如何用 SLURM 数组任务高效并行处理上万 FITS 文件-Python教程-PHP中文网

如何用 SLURM 数组任务高效并行处理上万 FITS 文件

本文介绍使用 SLURM 数组作业（--array）替代逐文件提交 10,000 个独立作业的最优实践，显著降低调度开销、提升资源利用率，并支持 Conda 环境激活与参数化 Python 脚本调用。

本文介绍使用 slurm 数组作业（`--array`）替代逐文件提交 10,000 个独立作业的最优实践，显著降低调度开销、提升资源利用率，并支持 conda 环境激活与参数化 python 脚本调用。

在处理大规模同构任务（如对 10,000 个 FITS 文件逐一执行相同 Python 脚本）时，为每个文件单独提交一个 sbatch 作业（即“1 文件 = 1 作业”）会带来严重调度压力：SLURM 需维护上万个作业状态，排队延迟高、I/O 和元数据开销大，且难以动态调整并发粒度。更优解是采用 SLURM 作业数组（Job Array） —— 仅提交一个数组作业，由 SLURM 自动派生数千个子任务（task），每个子任务通过环境变量 $SLURM_ARRAY_TASK_ID 获取唯一索引，进而映射到对应输入文件。

✅ 推荐方案：基于数组索引的批量分发

首先，将所有 FITS 文件路径统一存入一个文本文件（按需排序），例如：

find "input_dir" -name "*.fits" | sort > file_list.txt

然后编写主 SLURM 脚本 run_array.sh：

#!/bin/bash
#SBATCH -c 1                    # 每个子任务使用 1 CPU 核心
#SBATCH --time=02:00:00          # 单任务最长运行 2 小时（按实际调整）
#SBATCH --mem=4G                 # 每任务分配 4GB 内存
#SBATCH --job-name=fits_proc_%a  # %a 自动替换为任务 ID，便于识别
#SBATCH --output=log/%A_%a.out   # %A 为数组作业主 ID，%a 为子任务 ID
#SBATCH --error=log/%A_%a.err
#SBATCH --array=1-10000          # 启动 1–10000 号子任务（总数 = 文件行数）

# 创建日志目录（仅首次运行时需确保存在，建议提前 mkdir -p log/）
mkdir -p log/

# 激活 Conda 环境（推荐使用 conda run 或显式 source，避免 init 副作用）
source $HOME/miniconda3/etc/profile.d/conda.sh
conda activate my_env

# 读取第 $SLURM_ARRAY_TASK_ID 行的文件路径（从 1 开始编号）
INPUT_FILE=$(sed -n "${SLURM_ARRAY_TASK_ID}p" file_list.txt)

# 执行 Python 脚本，传入文件路径及可选标志（如是否生成统计 CSV）
python3 -u my_python_code.py "$INPUT_FILE" "False" 3

? 关键说明：$SLURM_ARRAY_TASK_ID 是 SLURM 内置变量，值为当前子任务序号（1, 2, ..., 10000）。sed -n "${N}p" 精确提取第 N 行，避免 head -n $N | tail -1 的低效管道。

提交命令简洁明了：

sbatch run_array.sh

⚙️ 进阶优化建议

动态数组范围：若文件数量未知，可用 $(wc -l 替代硬编码 <code>1-10000：
```
sbatch --array=1-$(wc -l 
```
Conda 环境更稳健激活：
若 conda activate 在非交互 shell 中失效，改用 conda run（无需激活）：
```
conda run -n my_env python3 -u my_python_code.py "$INPUT_FILE" "False" 3
```

错误容错与重试：添加简单重试逻辑（最多 2 次）：

for i in {1..2}; do
    if python3 -u my_python_code.py "$INPUT_FILE" "False" 3; then
        exit 0
    elif [ $i -eq 2 ]; then
        echo "Failed after 2 attempts on $INPUT_FILE" >&2
        exit 1
    fi
    sleep 5
done