inference.sh

#!/bin/bash

# 检查显卡状态的函数
check_gpu_idle() {
    # 使用nvidia-smi检查所有GPU的使用情况
    # 如果有任何GPU使用率超过20%或显存占用超过20%则返回1,否则返回0
    local is_busy=0
    
    # 检查GPU使用率
    while read -r usage; do
        if [ "$usage" -gt 20 ]; then
            is_busy=1
            break
        fi
    done < <(nvidia-smi --query-gpu=utilization.gpu --format=csv,noheader,nounits)
    
    # 检查显存占用率
    while read -r memory; do
        if [ "$memory" -gt 20480 ]; then
            is_busy=1
            break
        fi
    done < <(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits)
    
    if [ $is_busy -eq 1 ]; then
        # echo "GPU使用率或显存占用超过20%,等待120秒后重试..."
        return 1
    fi
    return 0
}

echo "开始监控GPU状态..."

while true; do
    if ! check_gpu_idle; then
        echo "GPU或显存正在使用中,等待300秒后重试..."
        sleep 300
    else
        echo "所有GPU空闲,开始训练..."
        cd src
        python inference_2.py
        break
    fi
done