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

_install() {
    if [[ $(lspci -n -d 10de:) ]];then
        PRINT_INFO "[INFO] 发现 NVIDIA 设备"
    elif [[ $(lspci -n -d 9999:4001) ]];then
        PRINT_INFO "[INFO] 发现 MATE-X 设备"
    else
        PRINT_INFO "[INFO] 未发现 NVIDIA 或 MATE-X 设备"
        return 0
    fi
}

get_script_env(){
    local BASEDIR SCRIPT_DIR
    SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
    BASEDIR=$SCRIPT_DIR/../../../
    python $BASEDIR/utils/cmdb_api.py Server
    MachineType=$(cat $BASEDIR/cmdb.info |grep -E ^MachineType: |awk -F ":" '{print $2}')
    if [[ $debug ]];then
        NR_FAIL=1
        echo [debug]: MachineType: $MachineType
        set -x
    fi
}

_check() {
    _check="PASS"
    local gpu_conf gpu_name gpu_pciid gpu_count real_count
    gpu_conf=$(/usr/bin/python "$BASEDIR"/utils/GetApolloConf.py "gpu_machinetype_name" "$MachineType" "gpu")
    # 配置中心获取不到机型GPU配置 并且 机器中没有VNIDIA的设备则不需要检查
    if [[ -z $gpu_conf ]] && _install;then
        return 0
    fi
    gpu_name=${gpu_conf%\**}
    gpu_count=${gpu_conf#*\*}
    if [[ -z $gpu_name || -z $gpu_count ]];then
        _check="FAIL"
        PRINT_ERR "[ERROR] get GPU Count failed -- GetApolloConf uops.gpu_machinetype_name $MachineType ;请硬件运维检查阿波罗配置 uops.gpu_machinetype_name $MachineType"
        return 1
    fi
    gpu_pciid=$(/usr/bin/python "$BASEDIR"/utils/GetApolloConf.py "gpu_name_pciid" "$gpu_name" "gpu" )
    # 检测 GPU 数量
    real_count=0
    for pciid in $gpu_pciid
    do
        ((real_count+=$(lspci -nn -d "$pciid" | wc -l)))
    done
    if [[ $real_count != "$gpu_count" ]];then
        _check="FAIL"
        PRINT_ERR "[ERROR] $MachineType GPU count is $gpu_count, real count is $real_count"
    else
        PRINT_INFO "[INFO] $MachineType real GPU count is $gpu_count"
    fi
    # 检测 GPU 带宽
    if [[ $MachineType == 'G1.U10.V8.34' ]];then
        return 0
    fi
    # Speed \d+GT/s, 会根据负载自动调整，无需检查
    for bdf in $(lspci -nD -d $gpu_pciid |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] $MachineType GPU Bandwidth $gpu_width_type ${slot//$'\t'/ }. "
        fi
    done
    if [[ $_check != 'FAIL' ]];then
        PRINT_INFO "[INFO] $MachineType GPU Bandwidth has only one type; $gpu_width_type "
    fi
    gpu_width_type=$(lspci -vvv -d "$gpu_pciid" |grep 'LnkCap:\|LnkSta:' |grep -oP 'Width x\d+' |sort -u )
    if [[ $(wc -l <<< "$gpu_width_type" ) == 1 ]];then
        PRINT_INFO "[INFO] $MachineType GPU Bandwidth has only one type; $gpu_width_type. "
    else
        _check="FAIL"
        PRINT_ERR "[ERROR] $MachineType GPU Bandwidth has multiple type; $gpu_width_type."
    fi
}
