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

Question

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

Answer 1

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

使用 cgroup 时，这不会唯一标识 GPU。对于 cgroups，所有 GPU 的 CUDA_VISIBLE_DEVICES 将为 0，因为每个进程只能看到一个 GPU（其他进程被 cgroup 隐藏）。 (5认同)

Answer 2

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

echo ${SLURM_STEP_GPUS:-$SLURM_JOB_GPUS}

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

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