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

Question

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

Answer 1

您可以使用环境变量获取 GPU id CUDA_VISIBLE_DEVICES。此变量是分配给作业的 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。