如何在Python中读取srun分配的任务ID 仅使用Slurm实现

结论 #

仅通过Slurm即可实现该需求，无需依赖MPI等额外并行框架。

实现逻辑 #

Slurm在通过srun启动每一个任务进程时，会自动为每个进程注入专属的环境变量，其中SLURM_PROCID的值就是当前任务在本次启动的所有任务中的排序编号，从0开始递增，正好匹配你需要的任务ID。

修改后的test.py代码 #

# test.py
import os

# 读取Slurm注入的环境变量获取任务ID
taskID = int(os.environ.get("SLURM_PROCID"))

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

效果验证 #

执行命令 srun -n2 python3 test.py 后，会直接输出你预期的结果：

task ID: 0
task ID: 1

可选优化：如果需要支持在非Slurm环境下本地运行脚本，可以给os.environ.get加默认值，例如int(os.environ.get("SLURM_PROCID", 0))，避免环境变量不存在时报错。

内容的提问来源于stack exchange，提问作者gkroiz