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

Question

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

Answer 1

您可以使用环境变量CUDA_VISIBLE_DEVICES获取 GPU id。 此变量是分配给作业的 GPU id 的逗号分隔列表。

Answer 2

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

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

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

Answer 3

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

echo ${SLURM_STEP_GPUS:-$SLURM_JOB_GPUS}

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

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