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

Posted 2023-04-15

【中文标题】如何在多个 GPU 节点上获取分配给 SLURM 作业的 GPU ID？

【英文标题】：How to get the ID of GPU allocated to a SLURM job on a multiple GPUs node?

【发布时间】：2017-10-13 12:36:13

【问题描述】：

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

【参考方案1】：

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

【讨论】：

它有效。谢谢。似乎环境变量 GPU_DEVICE_ORDINAL 也有效。

这在使用 cgroups 时不能唯一识别 GPU。使用 cgroups，所有 GPU 的 CUDA_VISIBLE_DEVICES 将为 0，因为每个进程只能看到一个 GPU（其他进程被 cgroup 隐藏）。

【参考方案2】：

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

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

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

【参考方案3】：

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

echo $SLURM_STEP_GPUS:-$SLURM_JOB_GPUS

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

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