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

Question

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

- Negelis

Answer 1

您可以通过环境变量CUDA_VISIBLE_DEVICES获取GPU id。这个变量是分配给作业的GPU id的逗号分隔列表。

- Carles Fenoy

---

它可以工作。谢谢。看起来环境变量GPU_DEVICE_ORDINAL也起作用。 - Negelis

使用cgroups时，这无法唯一标识GPU。因为每个进程只能看到单个GPU（其他GPU被cgroup隐藏），所以对于所有GPU来说，CUDA_VISIBLE_DEVICES的值都为0。 - isarandi

Answer 2

您可以检查给定节点的环境变量SLURM_STEP_GPUS或SLURM_JOB_GPUS：

echo ${SLURM_STEP_GPUS:-$SLURM_JOB_GPUS}

请注意CUDA_VISIBLE_DEVICES可能与实际值不对应（参见@isarandi的评论）。

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

- bryant1410

Answer 3

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

跟踪这些信息的一种方法是在运行作业时记录所有与SLURM相关的变量，例如（遵循Kaldi的slurm.pl，这是一个很棒的脚本来包装Slurm作业），通过在sbatch运行的脚本中包含以下命令：

Original Answer翻译成“最初的回答”

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

- leilu