Linux下抢占GPU资源

示例脚本

将以下脚本保存为 .sh 文件，在有 GPU 上的服务器上运行即可实现 GPU 资源的实时查询，一旦发现有符合要求的 GPU 资源释放便立即抢占。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
#!/bin/bash
mem=${1}
cmd=${*:2}
var=1 #控制最外层while中断
while [ ${var} -eq 1 ]
  do
    count=0
    for i in $(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits)
    #nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits：查询GPU的内存使用信息，以CSV格式返回（无表头和单位）
    do
      if [ ${i} -gt $mem ]
      then
        echo 'GPU'$count' is avaiable'
        CUDA_VISIBLE_DEVICES=$count $cmd
        var=0
        break
      fi
      count=$(($count+1)) #当前显卡数+1
    done
  done