【Linux shell命令】实验室监控空闲GPU，及时运行python代码程序

1. 环境配置

下载安装 gpustat

pip install gpustat

2. 脚本编写

假设我有一段需要运行的python代码如下：

CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --master_port 10025 --nproc_per_node=2 tools/relation_train_net.py --config-file "configs/e2e_relation_X_101_32_8_FPN_1x.yaml" MODEL.ROI_RELATION_HEAD.USE_GT_BOX True MODEL.ROI_RELATION_HEAD.USE_GT_OBJECT_LABEL True MODEL.ROI_RELATION_HEAD.PREDICTOR MotifPredictor SOLVER.IMS_PER_BATCH 12 TEST.IMS_PER_BATCH 2 DTYPE "float16" SOLVER.MAX_ITER 50000 SOLVER.VAL_PERIOD 2000 SOLVER.CHECKPOINT_PERIOD 2000 GLOVE_DIR /home/kaihua/glove MODEL.PRETRAINED_DETECTOR_CKPT /home/kaihua/checkpoints/pretrained_faster_rcnn/model_final.pth OUTPUT_DIR /home/kaihua/checkpoints/motif-precls-exmp

现在我需要把它放到一个shell文件中，方便等下程序传入空闲的GPU序号。

新建一个shell脚本，例子如下：

vi run_code.sh

将我要运行的代码复制到这个文件中，这里需要做一点修改，将代码中的CUDA_VISIBLE_DEVICES=0,1改为CUDA_VISIBLE_DEVICES=$1,$2，这样等下传入run_code.sh的前两个GPU序号就会被获取到（不同数量的GPU需要灵活处理，以此类推很容易吧？）。

执行:wq保存退出。

---

再新建一个shell脚本，作为GPU监控脚本，例子如下：

vi my_gpu_shell.sh

将下面的代码复制到这个文件中：

while true
do
  # 改动项1 relation_train_net 这里改成你命令中包含的字符串
  # 比如我要运行的代码如下：
  # CUDA_VISIBLE_DEVICES=0,1 python ... --nproc_per_node=2 ./relation_train_net.py ...
  # 这个主要是控制不要同时占用太多显卡，免得被打
  count=$(ps -ef | grep python | grep -c relation_train_net)
  if [ $count -lt 2 ]  # 改动项2， 根据之前被占用的显卡数调整
    then
     # 改动项3 查询第1块gpu的容量2p 第2块3p  第3块4p  第四块5p 依次类推
     # 这里我有4张卡，获取四块GPU的显存占用情况
     stat1=$(gpustat | awk '{print $11}' | sed -n '2p')
     stat2=$(gpustat | awk '{print $11}' | sed -n '3p')
     stat3=$(gpustat | awk '{print $11}' | sed -n '4p')
     stat4=$(gpustat | awk '{print $11}' | sed -n '5p')
     stat_arr=($stat1 $stat2 $stat3 $stat4)
     gpu_available=0
     gpu_available_index_arr=()
     # 得到空闲GPU的数量和对应的序号
     for i in ${!stat_arr[@]}
     do
       # 如果显存占用小于100M，继续
       if [ "${stat_arr[$i]}" -lt 100 ]
       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 running my code...'
       # 传值操作，这是我们上面创建的另一个shell脚本
       sh ./run_code.sh ${gpu_available_index_arr[0]} ${gpu_available_index_arr[1]}
       break # 防止下一次循环又重复运行上一行命令
     fi
  fi
  sleep 30
done

:wq保存并退出。开始执行！

sh my_gpu_shell.sh

3. 扩展：定时检测GPU空闲

有时候我们不想搞得太明显，也不能做得太过分。偷偷摸摸的事情当然是深夜时分最好操作，我们就需要一个定时脚本。

再新建一个shell脚本，作为定时执行脚本，例子如下：

vi time.sh

将以下内容复制到该shell脚本中:wq保存并退出。

#!/bin/bash
# 2022年9月25日1点0分0秒
date="20220925-010000"

while true
do
 now="$(date "+%Y%m%d-%H%M%S")"
 echo 当前时间：$now
 if [[ $date > $now ]]; then
    echo "$date > $now"
 elif [[ $date -eq $now ]]; then
    echo "$date == $now"
 else
    echo "$date < $now"
    # 执行GPU监控脚本
    sh my_gpu_shell.sh
    break
 fi
 sleep 30
done

开始执行！

sh time.sh

本文参考下面的文章并做改进：
Ubuntu 查询gpu空闲自动运行下一个程序