如何在多個 GPU 節點上獲取分配給 SLURM 作業的 GPU ID？

Question

當我將帶有選項 --gres=gpu:1 的 SLURM 作業提交給具有兩個 GPU 的節點時，如何獲取分配給該作業的 GPU 的 ID？ 是否有用於此目的的環境變量？ 我使用的 GPU 都是 nvidia GPU。 謝謝。

Answer 1

您可以使用環境變量CUDA_VISIBLE_DEVICES獲取 GPU id。 此變量是分配給作業的 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。