你能把 srun 的不同任务读入 Python 吗？答案

【问题标题】：Can you read different tasks of srun into Python?你能把 srun 的不同任务读入 Python 吗？
【发布时间】：2021-06-29 05:21:33
【问题描述】：

运行 srun 时，标志“-n”通常指定每个节点需要多少个任务。有没有办法将不同的任务 ID 读入 python。比如定义一个基本的python文件test.py：

#test.py

#somehow read the task ID as taskID
taskID = ...

print('task ID: ' + str(taskID))

现在，有没有办法定义 taskID 以便

srun -n2 python3 test.py 将打印以下内容：

task ID 0
task ID 1

我知道这可以使用 MPI 等其他并行框架来完成，因为您可以定义每个任务的等级，但想知道这是否可以仅使用 slurm 来完成。

【问题讨论】：

这能回答你的问题吗？ How to pass the SLURM-jobID as an input argument to python? .. 从未使用过 srun，但这似乎很接近？
@JeffUK 在发布我的问题之前，我查看了该链接。至少据我了解，jobID 与我正在寻找的不同。对于我上面的示例，jobID 将返回整个 slurm 提交“srun -n2 python3 test.py”的 ID，而不是作业中的各个任务。

【解决方案1】：

SLURM_PROCID 环境变量应该给出你想要的：

#test.py
import os
#somehow read the task ID as taskID
taskID = os.environ["SLURM_PROCID"]

print('task ID: ' + str(taskID))

【讨论】：