其原理是通过nvidia-smi扫描每块显卡上的内存,然后查询已经使用的内存,若已经使用的显卡内存不大于一个阈值,则运行python脚本(你也可以替换成别的命令)
#!/bin/bash
var=0
ocp_memory=${2:-500}
while [ $var -eq 0 ]
echo 'waiting for available gpu...'
docount=0for i in $(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits)doif [ $i -lt ocp_memory ]thenecho 'GPU'$count' is avaiable'CUDA_VISIBLE_DEVICES=$count python ${1}var=1breakficount=$(($count+1)) done
done
使用时的命令为
sh ocp_gpu.sh test.py 500
test.py即需要运行的py文件,500代表显卡被占有内存小于等于500MB。