本文档旨在指导用户如何在SLURM环境下，并行运行同一个Python脚本处理多个输入文件。我们将详细解释如何配置SLURM脚本，以及如何使用srun命令有效地分配任务到多个节点，从而实现高效的并行处理。此外，我们还会简要介绍SLURM作业数组，作为另一种可选的解决方案。
理解SLURM任务分配
在使用SLURM时，理解任务分配至关重要。以下是几个关键参数的解释：
- --nodes: 指定请求的节点数量。每个节点拥有的CPU数量取决于集群配置。
- --ntasks: 指定请求的总任务数量。
- --ntasks-per-node: 指定每个节点上运行的最大任务数量。如果指定了--ntasks，则此参数表示每个节点的最大任务数；否则，表示每个节点的精确任务数。
如果已知每个节点的CPU数量，建议使用--ntasks和--cpus-per-task，避免过度订阅节点，从而影响性能。
使用srun进行并行执行
srun命令用于在SLURM分配的资源上立即运行任务。在sbatch脚本中，srun允许用户利用已分配的资源执行并行任务。
立即学习“Python免费学习笔记（深入）”；
以下是一个示例SLURM脚本，展示了如何使用srun并行处理多个文件：
#!/bin/bash
#SBATCH --nodes=8
#SBATCH --ntasks-per-node=128
INPUT_DIR='path/to/input/dir'
OUTPUT_DIR='/path/to/output/dir'
# 读取文件名到数组
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
# 循环处理每个文件
# 轮询分配节点 (0, 1, ..., 7, 0, 1, ...)
NODE_NUMBER=$((j % SLURM_NNODES))
# 动态生成文件名
INPUT_FILE_NAME="$INPUT_DIR/${INPUT_STEMS[$j]}.txt"
OUTPUT_FILE_NAME="$OUTPUT_DIR/${INPUT_STEMS[$j]}.txt"
# 在指定节点上运行任务
srun -N1 -n1 -w ${SLURM_NODELIST} --nodelist=$(hostname -s)$NODE_NUMBER python_script.py --input "$INPUT_FILE_NAME" > "$OUTPUT_FILE_NAME" &
done
wait代码解释：
- #SBATCH 指令: 设置SLURM作业的资源需求，包括节点数量和每个节点的任务数。
- 文件名读取: 从文件中读取输入文件名，存储到数组INPUT_STEMS中。
- 循环处理: 使用for循环遍历每个输入文件。
- 轮询节点分配: 使用取模运算(%)将任务轮询分配到各个节点。SLURM_NNODES是SLURM提供的环境变量，表示分配的节点总数。
- 动态文件名生成: 根据循环索引和输入目录，动态生成输入和输出文件名。
-
srun 命令:
- -N1 -n1: 指定每个任务使用1个节点和1个核心。
- -w ${SLURM_NODELIST}: 确保任务仅在分配的节点上运行。
- --nodelist=$(hostname -s)$NODE_NUMBER: 指定任务运行的具体节点。hostname -s 获取当前节点的主机名，然后添加轮询分配的节点编号。
- python_script.py --input "$INPUT_FILE_NAME" > "$OUTPUT_FILE_NAME": 运行Python脚本，并将输出重定向到输出文件。
- &: 将任务放入后台运行，实现并行执行。
- wait 命令: 等待所有后台任务完成。
注意事项：
- 确保INPUT_DIR、OUTPUT_DIR和INPUT_STEMS_FILE变量设置为正确的值。
- 根据实际情况调整--nodes和--ntasks-per-node参数。
- python_script.py 脚本需要能够处理单个输入文件，并生成相应的输出文件。
- --nodelist参数可能需要根据集群的具体配置进行调整。
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/${INPUT_STEMS[$TASK_ID]}.txt"
python_script.py --input "$INPUT_FILE_NAME" > "$OUTPUT_FILE_NAME"代码解释：
- #SBATCH --array=0-999: 创建一个包含1000个子作业的作业数组，索引从0到999。
- $SLURM_ARRAY_TASK_ID: SLURM提供的环境变量，表示当前子作业的索引。
- 动态文件名生成: 根据$SLURM_ARRAY_TASK_ID动态生成输入和输出文件名。
- 运行Python脚本: 运行Python脚本，处理对应的输入文件。
选择哪种方法？
- 如果需要更精细地控制任务分配到哪些节点，或者需要在任务之间进行通信，则使用srun方法。
- 如果任务是完全独立的，并且不需要复杂的任务分配，则使用作业数组。
总结
本文档介绍了两种在SLURM环境下并行运行Python脚本处理多个文件的方法：使用srun和使用作业数组。选择哪种方法取决于具体的应用场景和需求。理解SLURM的任务分配机制，并根据实际情况调整脚本参数，可以有效地提高并行处理的效率。记住根据你的集群配置调整节点和任务数量，并仔细检查文件路径，以确保脚本正确运行。
