当前位置：网站首页>服务器自动抢占GPU运行程序

服务器自动抢占GPU运行程序

2022-07-19 18:56:00 【非线性光学元件】

其原理是通过nvidia-smi扫描每块显卡上的内存，然后查询已经使用的内存，若已经使用的显卡内存不大于一个阈值，则运行python脚本（你也可以替换成别的命令）

#!/bin/bash
var=0
ocp_memory=${2:-500}
while [ $var -eq 0 ]
echo 'waiting for available gpu...'
do
    count=0
    for i in $(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits)
    do
        if [ $i -lt ocp_memory ]
        then
            echo 'GPU'$count' is avaiable'
            CUDA_VISIBLE_DEVICES=$count python ${1}
            var=1
            break
        fi
        count=$(($count+1))    
    done    
done