SLURM并行运行Python脚本教程

本文档旨在指导用户如何使用 SLURM 作业调度器在多个输入文件上并行运行同一个 Python 脚本。文章详细解释了 SLURM 脚本的编写，着重讲解了如何正确配置节点和任务数量，以及如何使用 srun 命令有效地分配任务到各个节点，以实现最大程度的并行化。此外，还介绍了使用 SLURM 作业数组的替代方案，并比较了两种方法的优缺点。

SLURM 脚本配置

在使用 SLURM 并行运行脚本时，正确配置 SLURM 脚本至关重要。以下是一些关键参数的解释：

• --nodes: 指定要请求的节点数量。集群中每个节点的 CPU 数量取决于具体的硬件配置。
• --ntasks: 指定要请求的总任务数量。
• --ntasks-per-node: 指定每个节点上允许运行的最大任务数量。如果指定了 --ntasks，则此参数表示每个节点上允许的最大任务数；否则，它表示每个节点上运行的任务的确切数量。

为了避免节点超额订阅，建议使用 --ntasks 和 --cpus-per-task 参数，尤其是在同构集群中。

使用 srun 命令

srun 命令用于在 SLURM 分配的资源上立即运行作业。在 sbatch 脚本中，srun 允许用户从已分配的资源中选择用于运行任务。

以下是一个示例 SLURM 脚本，展示了如何使用 srun 在多个文件上并行运行 Python 脚本：

#!/bin/bash
#SBATCH --nodes=8
#SBATCH --ntasks-per-node=128

INPUT_DIR='path/to/input/dir'
OUTPUT_DIR='/path/to/output/dir'

# Read the file names into an array
INPUT_STEMS_FILE='/some/path/to/list/of/inputs.txt'
INPUT_STEMS=()
while IFS= read -r line; do
  INPUT_STEMS+=("$line")
done < <(tr -d '\r' < INPUT_STEMS_FILE)

for j in `seq 0 $(( ${#INPUT_STEMS[@]} - 1 ))`; do
  # Iterate over the indices for each of the N files

  # Round-robin allocation to nodes (0, 1, ..., 8, 0, 1, ...)
  NODE_NUMBER=$(($j % $SLURM_NNODES))

  # Dynamically generate filename
  INPUT_FILE_NAME="$INPUT_DIR/${INPUT_STEMS[$j]}.txt"
  OUTPUT_FILE_NAME="$OUTPUT_DIR/$j.txt"

  # Run a job on 1 task on 1 node, using the round-robin allocation.
  # The jobs run on different nodes, this way
  srun -N1 -n1 -w ${SLURM_NODELIST} --nodelist=$(hostname -s)$((NODE_NUMBER)) python_script.py --input $INPUT_FILE_NAME > $OUTPUT_FILE_NAME &
done

wait

代码解释：

1. 读取输入文件列表： 脚本首先从 INPUT_STEMS_FILE 读取输入文件名，并将它们存储在 INPUT_STEMS 数组中。
2. 循环遍历文件： 使用 for 循环遍历 INPUT_STEMS 数组中的每个文件。
3. 循环分配节点： 使用取模运算 (%) 将任务以循环方式分配给各个节点。$SLURM_NNODES 变量包含分配的节点总数。
4. 动态生成文件名： 根据循环索引 j 动态生成输入和输出文件名。
5. 使用 srun 运行任务： srun 命令用于在指定的节点上运行 Python 脚本。
   • -N1: 指定每个任务使用 1 个节点。
   • -n1: 指定每个任务运行在 1 个核心上。
   • -w ${SLURM_NODELIST}： 限制任务只能在分配的节点上运行。
   • --nodelist=$(hostname -s)$((NODE_NUMBER)): 指定任务运行的节点。 hostname -s 获取当前节点的主机名，然后加上循环分配的节点编号。
   • python_script.py --input $INPUT_FILE_NAME > $OUTPUT_FILE_NAME: 运行 Python 脚本，并将输出重定向到指定的文件。
   • &: 将任务放入后台运行。
6. wait 命令： wait 命令确保所有后台任务完成后脚本才会退出。

注意事项：

• 确保 INPUT_DIR 和 OUTPUT_DIR 路径正确。
• python_script.py 脚本必须能够处理单个输入文件并生成相应的输出文件。
• 根据实际情况调整 --nodes 和 --ntasks-per-node 参数。
• 这个脚本假设节点名称的格式是 -。如果你的集群节点命名方式不同，你需要相应地修改 --nodelist 参数。

使用 SLURM 作业数组

SLURM 作业数组是另一种并行处理多个任务的方法。使用作业数组，可以创建多个作业，每个作业处理一个输入文件。

优点：

• 易于使用。
• SLURM 会自动管理任务分配。

缺点：

• 会创建大量的作业，可能对调度器造成压力。

以下是一个使用 SLURM 作业数组的示例脚本：

#!/bin/bash
#SBATCH --array=0-999

INPUT_DIR='path/to/input/dir'
OUTPUT_DIR='/path/to/output/dir'

INPUT_STEMS_FILE='/some/path/to/list/of/inputs.txt'

# Read the file names into an array
INPUT_STEMS=()
while IFS= read -r line; do
  INPUT_STEMS+=("$line")
done < <(tr -d '\r' < INPUT_STEMS_FILE)


TASK_ID=$SLURM_ARRAY_TASK_ID

INPUT_FILE_NAME="$INPUT_DIR/${INPUT_STEMS[$TASK_ID]}.txt"
OUTPUT_FILE_NAME="$OUTPUT_DIR/$TASK_ID.txt"

python_script.py --input $INPUT_FILE_NAME > $OUTPUT_FILE_NAME

代码解释：

1. #SBATCH --array=0-999: 创建 1000 个作业，作业 ID 从 0 到 999。
2. TASK_ID=$SLURM_ARRAY_TASK_ID: 获取当前作业的 ID。
3. 根据 TASK_ID 动态生成输入和输出文件名。
4. 运行 Python 脚本，并将输出重定向到指定的文件。

总结

本文档介绍了两种使用 SLURM 在多个文件上并行运行 Python 脚本的方法：使用 srun 命令和使用 SLURM 作业数组。选择哪种方法取决于具体的需求和集群配置。使用 srun 命令可以更精细地控制任务分配，而作业数组则更易于使用。无论选择哪种方法，都需要仔细配置 SLURM 脚本，以确保任务能够有效地并行运行。