SLURM 并行执行：在多个文件上运行相同的 Python 脚本

本文档旨在指导用户如何在 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免费学习笔记（深入）”；

1. #SBATCH 指令： 定义了 SLURM 作业的资源需求，例如节点数量 (--nodes) 和每个节点的任务数量 (--ntasks-per-node)。请务必根据集群的具体配置调整这些参数。
2. 读取文件名： 从文件中读取输入文件名，并将其存储在数组 INPUT_STEMS 中。
3. 循环迭代： 使用 for 循环遍历文件名数组，为每个文件创建一个任务。
4. 轮询分配： 使用取模运算 (%) 将任务轮询分配到不同的节点。$SLURM_NNODES 变量包含了分配的节点总数。
5. 动态生成文件名： 根据循环索引动态生成输入和输出文件名。
6. 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 脚本，并将输出重定向到指定的输出文件。
   • & 符号将任务放入后台运行，实现并行执行。
7. 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 会自动为每个任务分配资源。

行者AI

行者AI绘图创作，唤醒新的灵感，创造更多可能

100

查看详情

以下是一个使用 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免费学习笔记（深入）”；

1. #SBATCH --array=0-999： 定义了一个 Job Array，包含 1000 个任务，索引从 0 到 999。
2. $SLURM_ARRAY_TASK_ID： SLURM 会为每个任务设置一个环境变量 $SLURM_ARRAY_TASK_ID，表示当前任务的索引。
3. 动态生成文件名： 根据 $SLURM_ARRAY_TASK_ID 动态生成输入和输出文件名。
4. 执行 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 脚本中的参数，以获得最佳的性能。

以上就是SLURM 并行执行：在多个文件上运行相同的 Python 脚本的详细内容，更多请关注php中文网其它相关文章！