Linux——监控GPU集群显存并自动运行python训练脚本

可喜可贺，由于有时候模型的实验时间刚好结束在半夜，同时此时GPU闲置会造成极大的资源浪费，同时会降低实验的效率，因此！监控GPU显存并自动运行python训练脚本便是刚需神器！！！

在借鉴了网上各路大佬的博文之后，整理了一般更适用于自身情况的shell，并以此分享记录~
参考文章：

• Ubuntu查询gpu空闲自动运行下一个程序

依赖包

pip install gpustat

主脚本程序

• 功能：
  • GPU显存实时监控；
  • 满足指定显存条件自动运行指定命令；
• 使用场景：
  • 我的使用场景是8卡的GPU集群，需求是等当前8卡程序执行完后，立即开始执行新的脚本程序。

gpu_check.sh

while true
do
  # 改动项1 用来查看你前一个在执行的训练是否还存在，这里为(train.py)
  count=$(ps -ef | grep python | grep -c train.py)	# 
  echo "$count train.py are running."
  if [ $count -lt 8 ]  # 改动项2， 根据之前被占用的显卡数调整
    then
     # 改动项3 查询第1块gpu的容量2p 第2块3p  第3块4p  第四块5p 依次类推
     # 这里只需要检测0,1GPU的显存情况便可以知道上一个程序是否执行完
     # 这里的11代指gpustat第11个单元的内容，需要自己进行更改！！！每个人不一定一样，例如:
     
     # 我在8卡集群上就为9，则命令为 gpustat | awk '{print $9}' | sed -n '2p'才能打印出当前显存
     # 我在本地单卡机器上为11，则命令为 gpustat | awk '{print $11}' | sed -n '2p'才能打印出当前显存
     stat1=$(gpustat | awk '{print $9}' | sed -n '2p')	# 2p代表 GPU 0	
     stat2=$(gpustat | awk '{print $9}' | sed -n '3p') # 3p代表 GPU 1
     stat_arr=($stat1 $stat2)
     gpu_available=0
     gpu_available_index_arr=()
     # 得到空闲GPU的数量和对应的序号
     for i in ${!stat_arr[@]}
     do
       # 如果显存占用小于100M，继续
       if [ "${stat_arr[$i]}" -lt 5000 ]
       then
         gpu_available=$[gpu_available+1]
         gpu_available_index_arr[${#gpu_available_index_arr[@]}]=$i
       fi
     done
     echo '-可用GPU数:'$gpu_available', 第'${gpu_available_index_arr[@]}'块GPU可用'
     # 如果GPU数大于指定数量，取指定数量GPU开始训练
     if [ $gpu_available -ge 2 ]
     then
       echo 'start training.'
       # 需要执行的python shell脚本
       sh ./run_code.sh
       break # 防止下一次循环又重复运行上一行命令
     fi
  fi
  echo "waiting for training..."
  sleep 300	# 每300s执行一次，对于一般实验来说，每5分钟执行一次即可,可以自行调整
done

python命令脚本（示例）

• 功能：执行具体命令。

run_code.sh

bash tools/dist_train.sh projects/configs/maptr/maptr_tiny_r50_24e.py 2

实验证明可行性

• 当GPU空闲时，运行bash gpu_check.sh：

• 程序直接开始执行：

• 此时使用gpustat命令查看GPU显存：

• 重新开一个新的终端运行脚本：（一直处于等待状态直到GPU占用满足条件才会执行）