如何使用一个 slurm 批处理脚本并行运行作业？

Question

我正在尝试与一个 Slurm 批处理脚本并行运行多个 python 脚本。看一下下面的例子：

#!/bin/bash
#
#SBATCH --job-name=test
#SBATCH --output=/dev/null
#SBATCH --error=/dev/null
#SBATCH --ntasks=2
#SBATCH --cpus-per-task=1
#SBATCH --mem-per-cpu=1G
#SBATCH --partition=All
#SBATCH --time=5:00

srun sleep 60
srun sleep 60
wait

如何调整脚本以便执行仅需 60 秒（而不是 120 秒）？将脚本拆分为两个脚本不是一个选项。

Answer 1

正如所写，该脚本并行运行两个sleep命令，连续运行两次。

每个srun命令都会启动一个步骤，并且由于您设置了--ntasks=2每个步骤，所以会实例化两个任务（此处为sleep命令）。

如果你想并行运行两个 1-task 步骤，你应该这样写：

srun --exclusive -n 1 -c 1 sleep 60 &
srun --exclusive -n 1 -c 1 sleep 60 &
wait

然后，每个步骤仅实例化一个任务，并以分隔符为背景&，这意味着下一个任务srun可以立即开始。该wait命令确保脚本仅在两个步骤完成后才终止。

在这种情况下，xargs命令和GNU 并行命令可用于避免写入多个相同的srun行或避免for-循环。

例如，如果您有多个文件，则需要运行脚本：

find /path/to/data/*.csv -print0 | xargs -0 -n1 -P $SLURM_NTASKS srun -n1 --exclusive python my_python_script.py

这相当于写了很多

srun -n 1 -c 1 --exclusive python my_python_script.py /path/to/data/file1.csv &
srun -n 1 -c 1 --exclusive python my_python_script.py /path/to/data/file1.csv &
srun -n 1 -c 1 --exclusive python my_python_script.py /path/to/data/file1.csv &
[...]

GNU 并行对于迭代参数值很有用：

parallel -P $SLURM_NTASKS srun  -n1 --exclusive python my_python_script.py ::: {1..1000}

会跑

python my_python_script.py 1
python my_python_script.py 2
python my_python_script.py 3
...
python my_python_script.py 1000

另一种方法是运行

srun python my_python_script.py

并且，在 Python 脚本内查找SLURM_PROCID环境变量并根据其值拆分工作。该srun命令将启动脚本的多个实例，每个实例都会“看到”不同的值SLURM_PROCID。

import os
print(os.environ['SLURM_PROCID'])