SLURM教程：Python多文件并行处理方法

想要高效并行运行Python脚本处理多个文件？本文档为你提供详细的SLURM教程，助你轻松应对大规模数据处理任务。我们将深入讲解如何在SLURM环境下，利用`srun`命令和作业数组两种方式，实现Python脚本的并行执行。首先，我们会解析SLURM的任务分配机制，包括`--nodes`、`--ntasks`、`--ntasks-per-node`等关键参数的配置，避免资源浪费。接着，通过实例演示如何使用`srun`命令将任务分配到指定节点，实现细粒度的任务控制。此外，还会介绍SLURM作业数组，一种更简便的并行处理方案。无论你选择哪种方法，都能显著提升Python脚本处理多文件的效率，快来学习吧！

本文档旨在指导用户如何在SLURM环境下，并行运行同一个Python脚本处理多个输入文件。我们将详细解释如何配置SLURM脚本，以及如何使用srun命令有效地分配任务到多个节点，从而实现高效的并行处理。此外，我们还会简要介绍SLURM作业数组，作为另一种可选的解决方案。

理解SLURM任务分配

在使用SLURM时，理解任务分配至关重要。以下是几个关键参数的解释：

• --nodes: 指定请求的节点数量。每个节点拥有的CPU数量取决于集群配置。
• --ntasks: 指定请求的总任务数量。
• --ntasks-per-node: 指定每个节点上运行的最大任务数量。如果指定了--ntasks，则此参数表示每个节点的最大任务数；否则，表示每个节点的精确任务数。

如果已知每个节点的CPU数量，建议使用--ntasks和--cpus-per-task，避免过度订阅节点，从而影响性能。

使用srun进行并行执行

srun命令用于在SLURM分配的资源上立即运行任务。在sbatch脚本中，srun允许用户利用已分配的资源执行并行任务。

以下是一个示例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

代码解释：

1. #SBATCH 指令: 设置SLURM作业的资源需求，包括节点数量和每个节点的任务数。
2. 文件名读取: 从文件中读取输入文件名，存储到数组INPUT_STEMS中。
3. 循环处理: 使用for循环遍历每个输入文件。
4. 轮询节点分配: 使用取模运算(%)将任务轮询分配到各个节点。SLURM_NNODES是SLURM提供的环境变量，表示分配的节点总数。
5. 动态文件名生成: 根据循环索引和输入目录，动态生成输入和输出文件名。
6. 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脚本，并将输出重定向到输出文件。
   • &: 将任务放入后台运行，实现并行执行。
7. 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"

代码解释：

1. #SBATCH --array=0-999: 创建一个包含1000个子作业的作业数组，索引从0到999。
2. $SLURM_ARRAY_TASK_ID: SLURM提供的环境变量，表示当前子作业的索引。
3. 动态文件名生成: 根据$SLURM_ARRAY_TASK_ID动态生成输入和输出文件名。
4. 运行Python脚本: 运行Python脚本，处理对应的输入文件。

选择哪种方法？

• 如果需要更精细地控制任务分配到哪些节点，或者需要在任务之间进行通信，则使用srun方法。
• 如果任务是完全独立的，并且不需要复杂的任务分配，则使用作业数组。

总结

本文档介绍了两种在SLURM环境下并行运行Python脚本处理多个文件的方法：使用srun和使用作业数组。选择哪种方法取决于具体的应用场景和需求。理解SLURM的任务分配机制，并根据实际情况调整脚本参数，可以有效地提高并行处理的效率。记住根据你的集群配置调整节点和任务数量，并仔细检查文件路径，以确保脚本正确运行。

本篇关于《SLURM教程：Python多文件并行处理方法》的介绍就到此结束啦，但是学无止境，想要了解学习更多关于文章的相关知识，请关注golang学习网公众号！