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

_install() {
    if lspci -n | grep -q 10de:;then
        PRINT_INFO "[INFO] 发现 NVIDIA 设备；识别为 GPU 压测"
    elif lspci -n | grep -q 9999:4001;then
        PRINT_INFO "[INFO] 发现 MATE-X 设备；识别为 GPU 压测"
    else
        PRINT_INFO "[INFO] 未发现 NVIDIA 或 MATE-X 设备，识别为 CCCPU 压测"
    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_type gpu_count
    gpu_conf=$(/usr/bin/python "$BASEDIR"/utils/GetApolloConf.py "gpu_machinetype_name" "$MachineType" "gpu")
    # 配置中心获取不到机型GPU配置 并且 机器中没有VNIDIA的设备则不需要检查
    if [[ -z $gpu_conf ]] && ! lspci -n | grep -qE '10de:|9999:4001';then
        cpu_tag=1
        gpu_tag=0
    else
        cpu_tag=0
        gpu_tag=1
    fi
    # 检查 CPU 是 x86_64
    if [[ $(arch) != x86_64 ]] ;then
        PRINT_INFO "[ERROR]  GPU/CPU_MEMORY 压测 目前只支持 x86_64 架构"
        return 0
    fi
    gpu_type=${gpu_conf%\**}
    gpu_count=${gpu_conf#*\*}
    if [[ -z $gpu_type || -z $gpu_count ]] && [[ $gpu_tag = 1 ]];then
        _check="FAIL"
        PRINT_ERR "[ERROR] get GPU Count failed -- 请硬件运维检查阿波罗配置 uops.gpu_machinetype_name $MachineType"
        return 1
    fi
    gpu_test_success=$(/usr/bin/python "$BASEDIR"/utils/GetApolloConf.py "gpu_test_success" "$SN" "gpu" 2> /dev/null)
    if [[ -z $gpu_test_success ]];then
        :
    elif [[ $gpu_test_success == 'True' ]];then
        [[ -f /Meta-Llama-3-70B-Instruct-GPTQ-Int8.tar.gz ]] && rm -f /Meta-Llama-3-70B-Instruct-GPTQ-Int8.tar.gz
        [[ -d /Meta-Llama-3-70B-Instruct-GPTQ-Int8 ]] && rm -rf /Meta-Llama-3-70B-Instruct-GPTQ-Int8
        PRINT_INFO "[INFO]  GPU 压测 发现 Apollo 压测成功标志"
        return 0
    elif [[ $gpu_test_success == 'False' ]];then
        _check="FAIL"
        PRINT_ERR "[ERROR] get GPU Test Success False -- 请硬件运维检查阿波罗配置 uops.gpu_test_success $SN"
        return 1
    else
        _check="FAIL"
        PRINT_ERR "[ERROR] get GPU Test Success $gpu_test_success -- 请硬件运维检查阿波罗配置 uops.gpu_test_success $SN"
        return 1
    fi

    if grep -q 'UCLOUD_GPU_DEMAG=0' /ucloud/exports.ini &&  grep -q 'UCLOUD_TEST_SEL=0' /ucloud/exports.ini &&  grep -q 'UCLOUD_GPU_OK=1' /ucloud/exports.ini ;then
        [[ -f /Meta-Llama-3-70B-Instruct-GPTQ-Int8.tar.gz ]] && rm -f /Meta-Llama-3-70B-Instruct-GPTQ-Int8.tar.gz
        [[ -d /Meta-Llama-3-70B-Instruct-GPTQ-Int8 ]] && rm -rf /Meta-Llama-3-70B-Instruct-GPTQ-Int8
        PRINT_INFO "[INFO]  GPU 压测 发现压测结果 $gpu_type $gpu_count ok... $(date +%F_%T_%Z) "
        return 0
    fi

    # 从此处往后为进入压测环境
    if echo "$gpu_type" | grep -q '4090' ;then
        export KERNEL_VERSION=vmlinuz-5.15.0-43-generic
        export INITRD_VERSION=initrd-5.15.0-43-535.161.07-706.cpio.xz
        export SCRIPTS_VERSION=initrd-scripts-gpu.img
        export TIME=48h
        export GPU_TASK=all
        export info_630_check="GPU 压测 $gpu_type 4090 : $TIME $GPU_TASK ing..."
    elif echo "$gpu_type" | grep -qEw '5090|5090D' ;then
        export KERNEL_VERSION=vmlinuz-5.15.0-43-generic
        export INITRD_VERSION=initrd-5.15.0-43-575.64.05-cuda12.9-1202.cpio.xz
        export SCRIPTS_VERSION=initrd-scripts-gpu.img
        export TIME=48h
        export GPU_TASK=all
        export info_630_check="GPU 压测 $gpu_type 5090 : $TIME $GPU_TASK ing..."
    elif echo "$gpu_type" | grep -qEw 'H20|L20|A800|H800|A100|H100' ;then
        export KERNEL_VERSION=vmlinuz-5.15.0-94-generic
        export INITRD_VERSION=initrd-5.15.0-94-550.90.07-ofed23.10-0925.cpio.xz
        export SCRIPTS_VERSION=initrd-scripts-gpu.img
        export TIME=48h
        export GPU_TASK=all
        export info_630_check="GPU 压测 $gpu_type 550.90.07 : $TIME $GPU_TASK ing..."
    elif echo "$gpu_type" | grep -qE '2080|3080|3090|2080Ti|3080TI|Tesla|TeslaT4|TeslaP40' ;then
        export KERNEL_VERSION=vmlinuz-5.15.0-43-generic
        export INITRD_VERSION=initrd-5.15.0-43-535.161.07-706.cpio.xz
        export SCRIPTS_VERSION=initrd-scripts-gpu.img
        export TIME=60h
        export GPU_TASK=burn_stressapptest
        export info_630_check="GPU 压测 $gpu_type : $TIME $GPU_TASK ing..."
    elif echo "$gpu_type" | grep -q METAX-C500 ;then
        export KERNEL_VERSION=vmlinuz-5.15.0-58-generic
        export INITRD_VERSION=initrd-5.15.0-58-maca-mxc500-2.24-1112.cpio.xz
        export SCRIPTS_VERSION=initrd-scripts-gpu.img
        export TIME=72h
        # export GPU_TASK=mxvs
        export GPU_TASK=burn_stressapptest
        export info_630_check="GPU 压测 $gpu_type METAX-C500 : $TIME $GPU_TASK ing..."
        export DOWNLOAD_URL=10.81.40.4/stress_test/boot
    elif [[ $cpu_tag = 1 ]] ;then
        export KERNEL_VERSION=vmlinuz-5.15.0-43-generic
        export INITRD_VERSION=initrd-5.15.0-43-535.161.07-706.cpio.xz
        export SCRIPTS_VERSION=initrd-scripts-gpu.img
        export TIME=24h
        export GPU_TASK=stressapptest
        export info_630_check="CCCCPU 压测 MEM : $TIME $GPU_TASK ing..."
    else
        export KERNEL_VERSION=vmlinuz-5.15.0-43-generic
        export INITRD_VERSION=initrd-5.15.0-43-535.161.07-706.cpio.xz
        export SCRIPTS_VERSION=initrd-scripts-gpu.img
        export TIME=60h
        export GPU_TASK=burn_stressapptest
        export info_630_check="GPU 压测 $gpu_type non-matched : $TIME $GPU_TASK ing..."
    fi

    _check="FAIL"
    if [[ $CHANNEL == "reinstall" && $gpu_tag == 1 ]];then
        export TIME=12h
        export GPU_TASK=burn_stressapptest_reinstall
        export info_630_check="GPU 压测 $gpu_type reinstall : $TIME burn_stressapptest_reinstall ing..."
    fi
    if [[ $cpu_tag == 1 ]];then
        export TIME=24h
        export GPU_TASK=stressapptest
        export info_630_check="CCCCPU 压测 MEM : $TIME stressapptest ing..."
    fi
    gpu_test_time=$(/usr/bin/python "$BASEDIR"/utils/GetApolloConf.py "gpu_test_time" "$SN" "gpu")
    TIME=${gpu_test_time:-$TIME}

    PRINT_ERR "[ERROR]  $info_630_check $(date +%F_%T_%Z)"
    DOWNLOAD_URL=${DOWNLOAD_URL:-softlib.service.ucloud.cn/usa/stress_test}
    # 根据前方设置变量 在主程序退出后进入压测环境
    trap pressure_test EXIT
    pressure_test(){
        if [[ $NR_FAIL != 1 ]];then
            exit 1
        fi
        PRINT_ERR "test info_630_check $info_630_check $(date +%F_%T_%Z)"
        UPLOAD_LOG GPU
        curl -s $DOWNLOAD_URL/pangu-boot-stress_test.sh |bash -x
    }
}

