Linux监听空闲GPU，定时运行多卡并行python程序

1.介绍

本脚本功能主要用于在linux服务器中检查空闲GPU，定时运行python程序。

2.脚本

2.1环境配置

下载安装gpustat

pip install gpustat

2.2运行程序脚本

train_run.sh这是我们需要运行的多卡程序脚本，其中$1,$2 对应显卡编号

# #!/bin/bash
export PYTHONPATH=$PYTHONPATH:$(pwd)
CUDA_VISIBLE_DEVICES=$1,$2  torchrun --nproc_per_node=2 --master_port=29509 URF/run_test.py --dataroot=./dataset/train/URF --n_ep=1000 --n_ep_decay=800 --resume=./checkpoint/initial_weights_128.pth --stage=RF --batch_size=2 --name=ttt

2.3监听空闲GPU脚本

gpu_monitor.sh是监听脚本

while true
do
  # shellcheck disable=SC1060
  # stat1-stat8 对应显卡数量
  # $09表示gpustat查询到的显卡占用显存
  # '2p-9p' 表示显卡号 '2p' 对应第一张显卡，依此类推
  stat1=$(gpustat | awk '{print $09}' | sed -n '2p')
  stat2=$(gpustat | awk '{print $09}' | sed -n '3p')
  stat3=$(gpustat | awk '{print $09}' | sed -n '4p')
  stat4=$(gpustat | awk '{print $09}' | sed -n '5p')
  stat5=$(gpustat | awk '{print $09}' | sed -n '6p')
  stat6=$(gpustat | awk '{print $09}' | sed -n '7p')
  stat7=$(gpustat | awk '{print $09}' | sed -n '8p')
  stat8=$(gpustat | awk '{print $09}' | sed -n '9p')
  stat_arr=($stat1 $stat2 $stat3 $stat4 $stat5 $stat6 $stat7 $stat8)
  gpu_available=0
  gpu_available_index_arr=()
  # 得到空闲GPU的数量和对应的序号
  for i in ${!stat_arr[@]}
  do
     # 如果显存占用小于阈值(单位M)，继续
    if [ "${stat_arr[$i]}" -lt 10000 ]
    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 the code...'
    # 传值操作,即需要运行的代码脚本
    sh ./train_run.sh ${gpu_available_index_arr[0]} ${gpu_available_index_arr[1]}
    break # 防止下一次循环又重复运行上一行命令
  fi
  sleep 30 # 单位秒
done

运行命令bash gpu_monitor.sh

2.4定时启动脚本

time_run.sh是定时启动脚本

#!/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监控脚本
    bash gpu_monitor.sh
    break
 fi
 sleep 30
done

参考资料

1. 【Linux shell命令】实验室监控空闲GPU，及时运行python代码程序
2. shell中 -eq，-ne，-gt，-lt，-ge，-le数字比较符