当我使用选项 --gres=gpu:1 向具有两个 GPU 的节点提交 SLURM 作业时,如何获取为该作业分配的 GPU 的 ID?是否有用于此目的的环境变量?我使用的 GPU 都是 nvidia GPU。 谢谢。
您可以通过环境变量
CUDA_VISIBLE_DEVICES
获取GPU id。该变量是分配给作业的 GPU ID 的逗号分隔列表。
您可以检查给定节点的环境变量
SLURM_STEP_GPUS
或 SLURM_JOB_GPUS
:
echo ${SLURM_STEP_GPUS:-$SLURM_JOB_GPUS}
注意
CUDA_VISIBLE_DEVICES
可能与实际值不符(参见@isarandi的评论)。
另请注意,这也适用于非 Nvidia GPU。