如何获取分配给多 GPU 节点上的 SLURM 作业的 GPU ID？

Question

当我使用选项 --gres=gpu:1 向具有两个 GPU 的节点提交 SLURM 作业时，如何获取为该作业分配的 GPU 的 ID？是否有用于此目的的环境变量？我使用的 GPU 都是 nvidia GPU。 谢谢。

Answer 1

您可以通过环境变量

CUDA_VISIBLE_DEVICES

获取GPU id。该变量是分配给作业的 GPU ID 的逗号分隔列表。

Answer 2

您可以检查给定节点的环境变量

SLURM_STEP_GPUS

或

SLURM_JOB_GPUS

：

echo ${SLURM_STEP_GPUS:-$SLURM_JOB_GPUS}

注意

CUDA_VISIBLE_DEVICES

可能与实际值不符（参见@isarandi的评论）。

另请注意，这也适用于非 Nvidia GPU。

Answer 3

Slurm 将此信息存储在环境变量中，

SLURM_JOB_GPUS

或

SLURM_STEP_GPUS

。

跟踪此类信息的一种方法是在运行作业时记录所有 SLURM 相关变量，例如（遵循 Kaldi 的 slurm.pl，这是包装 Slurm 作业的一个很棒的脚本），包括以下内容

sbatch

运行的脚本中的命令:

set | grep SLURM | while read line; do echo "# $line"; done