本文档旨在指导用户如何在 SLURM 环境下，利用并行计算能力，高效地在多个输入文件上运行同一个 Python 脚本。我们将探讨如何正确配置 SLURM 脚本，利用 srun 命令分配任务，以及如何使用 Job Arrays 简化流程，从而充分利用集群资源，加速数据处理。
使用 srun 并行化 Python 脚本
当需要在 SLURM 集群上并行运行同一脚本处理多个文件时，srun 命令是关键。与 sem (GNU parallel) 不同，srun 与 SLURM 调度器紧密集成，能更好地利用分配的资源。
以下是一个示例 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'
# 将文件名读入数组
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, ..., 8, 0, 1, ...)
NODE_NUMBER=$((j % $SLURM_NNODES))
# 动态生成文件名
INPUT_FILE_NAME="$INPUT_DIR/${INPUT_STEMS[$j]}.txt"
OUTPUT_FILE_NAME="$OUTPUT_DIR/$j.txt"
# 在 1 个节点上运行 1 个任务，使用轮询分配。
# 这样，作业在不同的节点上运行
srun -N1 -n1 -w "$(( $(hostname -s) $NODE_NUMBER ))" python_script.py --input "$INPUT_FILE_NAME" > "$OUTPUT_FILE_NAME" &
done
wait代码解释：
立即学习“Python免费学习笔记（深入）”；
- #SBATCH 指令： 定义了 SLURM 作业的资源需求，例如节点数量 (--nodes) 和每个节点的任务数量 (--ntasks-per-node)。请务必根据集群的具体配置调整这些参数。
- 读取文件名： 从文件中读取输入文件名，并将其存储在数组 INPUT_STEMS 中。
- 循环迭代： 使用 for 循环遍历文件名数组，为每个文件创建一个任务。
- 轮询分配： 使用取模运算 (%) 将任务轮询分配到不同的节点。$SLURM_NNODES 变量包含了分配的节点总数。
- 动态生成文件名： 根据循环索引动态生成输入和输出文件名。
-
srun 命令： 使用 srun 命令在指定的节点上运行 Python 脚本。
- -N1 -n1 指定每个任务使用 1 个节点和 1 个核心。
- -w $(( $(hostname -s) $NODE_NUMBER )) 指定在特定节点上运行任务，其中 $NODE_NUMBER 是轮询分配的节点索引。
- python_script.py --input "$INPUT_FILE_NAME" > "$OUTPUT_FILE_NAME" 是实际执行的 Python 脚本，并将输出重定向到指定的输出文件。
- & 符号将任务放入后台运行，实现并行执行。
- wait 命令： 等待所有后台任务完成。
注意事项：
- 请根据实际情况修改 INPUT_DIR、OUTPUT_DIR 和 INPUT_STEMS_FILE 变量。
- 确保 Python 脚本 python_script.py 可执行，并且能够正确处理输入文件。
- 根据集群的节点配置和任务需求，调整 --nodes 和 --ntasks-per-node 参数。
- -w 参数可能需要根据集群的具体配置进行调整。某些集群可能需要使用不同的方式来指定节点。
- hostname -s 命令用于获取当前节点的主机名。
使用 Job Arrays 简化并行执行
SLURM 的 Job Arrays 提供了一种更简洁的方式来并行执行多个任务。通过 Job Arrays，可以将多个相似的任务作为一个整体提交，SLURM 会自动为每个任务分配资源。
以下是一个使用 Job Arrays 的示例 SLURM 脚本：
#!/bin/bash
#SBATCH --array=0-999
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
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)
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"代码解释：
立即学习“Python免费学习笔记（深入）”；
- #SBATCH --array=0-999： 定义了一个 Job Array，包含 1000 个任务，索引从 0 到 999。
- $SLURM_ARRAY_TASK_ID： SLURM 会为每个任务设置一个环境变量 $SLURM_ARRAY_TASK_ID，表示当前任务的索引。
- 动态生成文件名： 根据 $SLURM_ARRAY_TASK_ID 动态生成输入和输出文件名。
- 执行 Python 脚本： 使用 Python 脚本处理指定的输入文件，并将结果输出到指定的输出文件。
使用 Job Arrays 的优点：
- 简化了 SLURM 脚本，减少了代码量。
- 更方便地管理和监控多个任务。
- SLURM 可以更有效地分配资源，提高集群利用率。
注意事项：
- 请根据实际情况修改 --array 参数，指定 Job Array 的任务数量和索引范围。
- 确保 $SLURM_ARRAY_TASK_ID 变量在脚本中正确使用。
- Job Arrays 会创建大量的独立作业，可能会对 SLURM 调度器造成压力。
总结
本文介绍了两种在 SLURM 环境下并行执行 Python 脚本的方法：使用 srun 命令和使用 Job Arrays。srun 命令提供了更灵活的控制，可以手动分配任务到不同的节点。Job Arrays 则更加简洁，方便管理大量的相似任务。选择哪种方法取决于具体的应用场景和需求。理解 SLURM 的工作原理和正确配置 SLURM 脚本是充分利用集群资源的关键。 记住要根据集群的具体配置和任务需求，调整 SLURM 脚本中的参数，以获得最佳的性能。
