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

在炼丹的时候我们会碰到下面两种情况:

• 想连续跑多个训练测试，可以新建一个sh文件，里面按行写上各条命令，然后直接运行这个sh文件就行。这种方法是最简单的，在自己固定有卡和各条命令之间时间逻辑先后关系考虑好，每条命令都能运行的情况下是可以的（这样不用往下看了）
• 还有一种情况，就是自己没有固定的卡，想在别人运行完自己第一时间跑上。或者在上一种情况下，第二条命令还没及时调试好，但是先想让第一条命令先运行着，后面再找时间无缝训练上。属于这些情况就可以往下看：

本篇参考这篇博客，对其进行解析和改了一个小bug

提前工具准备

后面的脚本需要在gpustat基础上运行，所以先装一下:

pip install gpustat

然后运行一下gpustat就可以看到自己的gpu信息了，就像nvidia-smi一样，比如会有显卡型号，温度，显存占用量，比如我的:

脚本编写

while true
do
  # 改动项1 text_include_in_your_cmd 这里改成你命令中包含的字符串
  count=$(ps -ef | grep python | grep -c text_include_in_your_cmd)
  if [ $count -lt 2 ]  # 改动项2， 根据之前被占用的显卡数调整
    then
     # 改动项3 查询第1块gpu的容量2p 第2块3p  第3块4p  第四块5p 依次类推
     stat2=$(gpustat | awk '{print $11}' | sed -n '5p')
     if [ "$stat2" -lt 1000 ] # 改动项4 显存小于多少表示前面的跑好了
       then
         echo 'run my code'
         #改动项5 前面输入占用的gpu id 后面是运行代码
         CUDA_VISIBLE_DEVICES=4,5,6,7 python train.py
         break # 防止下一次循环又重复运行上一行命令
     fi
  fi
  sleep 2
done

这里要将上面几点解释一下：

• 改动项1： ps -ef是查询进程命令的，|是在其中过滤的，grep python是匹配到包含python这个字符串的行，因为训练一般都会有python，所以我这里加了python这个条件，grep后面加-c参数就是返回符合条件的行数，比如有text_include_in_your_cmd这个字符串的进程名有多少个，这里需要填写现在正被占用的显卡上运行的命令的几个独特一点的字符串吧，比如你之前运行了pyhon train_siamrpn.py这样的命令，你这里就换成siamrpn。
• 改动项2：当多卡运行的时候，count值一般会有很多个，起码有显卡数个吧，还有可能一些副进程也包含那些字符串的，所以这里条件是小于多少的时候进入我们要运行的看之前的训练情况，单卡的话可能得改成1
• 改动项3：这里面的11和5p是需要改的，11就是把gpustat返回的各行的第11个字符串提取出来，我们这里要提取的就是显存占用量。因为gpustat在不同gpu型号的机器上返回的是不一样的，所以11这个位置得自己调整着来。比如在这台电脑上是11
  在我自己的笔记本上就是15
  
  注意这里awk后返回的第一行是空行，因为一般索引都超出gpustat返回的第一行的所有信息(时间，驱动型号啥的)了
  5p的意思就是把awk返回的第5行打印下来，因为有一行空行，所以2p就是第一块卡，也就是CUDA_VISIBLE_DEVICES中的0，3p就是第二块卡，以此类推，看自己想监控哪块卡的显存使用量了
• 改动项4：我这里写着1000，代表5p这张卡显存使用量小于1000M的时候就进入我们的命令，具体使用量看自己的情况
• 改动项5：就是你要跑的命令了，包括用哪些卡，最后运行好为了防止循环运行，break跳出循环

OK，结束，希望善意使用