统计所有GPU利用率:
1 unset GPUs 2 declare -A GPUs 3 unset i 4 for u in `nvidia-smi -q -d UTILIZATION | grep Gpu | awk '{ print $3 }'` 5 do 6 i=$(($i+1)) 7 id=$(($i-1)) 8 GPUs+=(["$id"]="$u") 9 done
设置阈值15%,判断如果特定[$1]GPU利用率小于15%,则可使用,否则提示不可用,并自动从GPU列表中寻找可用GPU,直至找到:
1 gpu_util=15 2 if [ "$1" ] ; then 3 if [[ ${GPUs["$1"]} -lt "$gpu_util" ]];then 4 gpu_id=$2 5 else 6 echo "Warnning: Your selected GPU ID:[$2] is busy and unavailable!" 7 echo "..." 8 for key in $(echo "${!GPUs[*]}") 9 do 10 if [[ ${GPUs["$key"]} -lt "$gpu_util" ]];then 11 gpu_id=$key 12 break 13 fi 14 done 15 fi16 fi
根据上述判断,得到可用GPU
1 unset CUDA_VISIBLE_DEVICES 2 export CUDA_VISIBLE_DEVICES=$gpu_id