linux GPU查询和空闲时自动抢占脚本

以下脚本，是为了方便同学们跑实验用，禁止利用此脚本恶意抢占GPU资源

在实验室服务器中，GPU数量有限，但是经常出现有些同学有意无意的占用多块GPU，造成资源缺乏，不能及时跑实验。

工具

依赖工具：gpustat
安装方式：pip install gpustat
效果如下：

脚本

#!/bin/sh

while true
do
  count=$(ps -ef | grep -c kk_test)
  if [ $count -lt 2 ]
    then
     # 改动项 查询第1块gpu的容量--2p 第2块3p--2  第三块--4p  第四块--5p 
     stat2=$(gpustat | awk '{print $11}' | sed -n '5p')
     if [ "$stat2" -lt 1000 ]
       then
         echo 'run'
         #改动项 前面输入占用的gpu id 后面是运行代码
         CUDA_VISIBLE_DEVICES=3 python example.py  
     fi
  fi
  sleep 2
done

上面需要改动有两处

stat2=$(gpustat | awk '{print $11}' | sed -n '5p')

后面的’5p’根据需要查询的gpu改动，改动规则为
查询第1块gpu的容量–2p 第2块3p–2 第三块–4p 第四块–5p
同时还要在下面

CUDA_VISIBLE_DEVICES=3 python example.py  

改动到指定查询的gpu，example.py是自己需要运行的代码
这样就可以保证，在gpu出现空闲时自动抢占gpu资源
注：

  if [ "$stat2" -lt 1000 ]

这句话的意思是当查询的GPU容量使用小于1000M时，进行抢占