#!/bin/bash
# 检查 NVIDIA GPU 数量

INIT_LOG=/tmp/gpu-band-width-check.log
# Print All To LogFile
PRINT(){
    echo -n  "$*"'<br />' >> $INIT_LOG
}

# Print Error To LogFile
PRINT_ERR(){
    echo -n "<font color='red'>$*</font><br />" >> $INIT_LOG
    echo -e "\033[31m$*\033[0m"
}

# Print Info To LogFile
PRINT_INFO(){
    echo -n "<font color='green'>$*</font><br />" >> $INIT_LOG
    echo -e "\033[32m$*\033[0m"
}

_check() {
    _check="PASS"
    local gpu_conf gpu_type gpu_count real_count
    if [[ -z $gpu_conf ]] && ! lspci -n | grep -q 10de:;then
        return 0
    fi
    gpu_type=10de::0300
    gpu_count=8
    real_count=$(lspci -nn -d "$gpu_type" | wc -l)
    if [[ $real_count != "$gpu_count" ]];then
        _check="FAIL"
        PRINT_ERR "[ERROR] GPU count is $gpu_count, real count is $real_count"
    else
        PRINT_INFO "[INFO] real GPU count is $gpu_count"
    fi
    # 检测 GPU 带宽
    # Speed \d+GT/s, 会根据负载自动调整，无需检查
    for bdf in $(lspci -nD -d $gpu_type |awk '{print $1}');do
        gpu_width_type=$(lspci -vvv -s $bdf |grep 'LnkCap:\|LnkSta:'|grep -oP 'Width x\d+' |sort -u)
        if [[ $(wc -l <<< "$gpu_width_type" ) != 1 ]];then
            _check="FAIL"
            slot=$(dmidecode -t slot |grep -iE "add|des" |grep -B1 $bdf)
            PRINT_ERR "[ERROR] GPU Bandwidth $gpu_width_type ${slot//$'\t'/ }. "
        fi
    done
    if [[ $_check != 'FAIL' ]];then
        PRINT_INFO "[INFO] GPU Bandwidth has only one type; $gpu_width_type "
    fi
    gpu_width_type=$(lspci -vvv -d "$gpu_type" |grep 'LnkCap:\|LnkSta:' |grep -oP 'Width x\d+' |sort -u )
    if [[ $(wc -l <<< "$gpu_width_type" ) == 1 ]];then
        PRINT_INFO "[INFO] GPU Bandwidth has only one type; $gpu_width_type. "
    else
        _check="FAIL"
        PRINT_ERR "[ERROR] GPU Bandwidth has multiple type; $gpu_width_type."
    fi
}

_check


