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

时间:2017-05-14 18:22:51

标签: gpu nvidia slurm sbatch

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

3 个答案:

答案 0 :(得分:4)

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

答案 1 :(得分:1)

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

echo ${SLURM_STEP_GPUS:-$SLURM_JOB_GPUS}

注意 CUDA_VISIBLE_DEVICES 可能与实际值 (see @isarandi's comment) 不对应。

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

答案 2 :(得分:0)

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

一种跟踪此类信息的方法是在运行作业时记录所有与SLURM相关的变量,例如(遵循Kaldislurm.pl,这是包装Slurm作业的绝佳脚本)通过在sbatch运行的脚本中包含以下命令:

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