代码改变世界

统计所有GPUs利用率，根据利用与自动分配GPU任务

2021-02-02 17:05  丨o聽乄雨o丨  阅读(621)  评论(0)  编辑  收藏  举报

统计所有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