SLURM环境下如何在多文件上并行运行同一Python脚本

问题根源分析 #

你遇到的两个问题本质是任务调度逻辑和SLURM资源分配不匹配：

• 用sem时：sem默认只会在提交脚本的节点上调度任务，不会自动分发到SLURM分配的其他节点，加上默认并发数受单节点核心数限制，所以只跑了约70个实例。
• 用srun循环时：bash循环变量展开时机问题，所有异步启动的srun会共享循环变量的最终值，导致所有实例处理同一个文件。

推荐解决方案 #

方案1：GNU Parallel + SLURM集成（高效灵活） #

利用GNU Parallel的SLURM适配能力，自动将任务分发到所有分配的节点，无需手动管理节点和并发数。

示例SLURM提交脚本：

#!/bin/bash
#SBATCH --nodes=8
#SBATCH --ntasks-per-node=128  # 8*128=1024，匹配你需要的任务数
#SBATCH --cpus-per-task=1
#SBATCH --time=01:00:00

# 让Parallel与SLURM集成，自动用srun提交任务
export PARALLEL_SLURM=1

# 定义单个文件处理逻辑
process_csv() {
    input="$1"
    output="${input%.csv}_summary.csv"
    python your_script.py --input "$input" --output "$output"
}

# 导出函数，让Parallel可以调用
export -f process_csv

# 用Parallel调度所有CSV文件，并发数等于SLURM总任务数
parallel --jobs $SLURM_NTASKS process_csv ::: /path/to/your/*.csv

方案2：SLURM数组任务（原生无依赖） #

用SLURM原生的数组任务功能，每个数组实例对应一个文件，无需额外工具。

步骤1：生成文件列表

ls /path/to/your/*.csv > file_list.txt

步骤2：编写SLURM数组脚本

#!/bin/bash
#SBATCH --nodes=8
#SBATCH --ntasks-per-node=128
#SBATCH --cpus-per-task=1
#SBATCH --time=01:00:00
#SBATCH --array=1-1024  # 数字需与文件总数一致，%1024限制并发数等于总任务数

# 根据数组索引获取对应文件
input_csv=$(sed -n "${SLURM_ARRAY_TASK_ID}p" file_list.txt)
output_file="${input_csv%.csv}_summary.csv"

# 运行Python脚本
python your_script.py --input "$input_csv" --output "$output_file"

修复你之前的srun循环问题（不推荐，但可应急） #

如果坚持用bash循环+srun，需要将文件参数传递给子shell，避免变量覆盖：

#!/bin/bash
#SBATCH --nodes=8
#SBATCH --ntasks=1024
#SBATCH --cpus-per-task=1
#SBATCH --time=01:00:00

for input_csv in /path/to/your/*.csv; do
    # 用--exclusive确保每个任务独占资源，通过子shell传递参数
    srun --exclusive --mem-per-cpu=100M bash -c '
        input="$1"
        output="${input%.csv}_summary.csv"
        python your_script.py --input "$input" --output "$output"
    ' _ "$input_csv" &
done
wait

内容的提问来源于stack exchange，提问作者camilo