#!/bin/bash

# 请根据诊断信息，在华为云官网自助诊断文档中找到对应问题解决办法，尝试自恢复gpu实例。
# 如果是硬件问题，请提工单

set -E -o pipefail

DIAGNOSE_DIR=/root/diagnose_gpu_$(date '+%y-%m-%d_%H-%M-%S')
mkdir -p "${DIAGNOSE_DIR}"
DIAGNOSE_LOG_PATH=${DIAGNOSE_DIR}/diagnose_gpu.log
NVIDIA_SMI_PATH=
GPU_NUMS=


_printf_end_split_line() {
    logger.info "$(printf "%0.s="{0..120} | sed 's/[0-9]//g')"
    echo ""
}

logger.info() {
    echo "  info  $1" 2>&1 | tee -a "${DIAGNOSE_LOG_PATH}"
}

logger.error() {
    echo "  error  $1" 2>&1 | tee -a "${DIAGNOSE_LOG_PATH}"
}

logger.echo() {
    echo "        $1" 2>&1 | tee -a "${DIAGNOSE_LOG_PATH}"
}

# 首先确定是否是GPU实例
_diagnose_gpu.is_gpu() {
    local pci_device_path="/sys/bus/pci/devices"
    if [[ ! -d "${pci_device_path}" ]];then
        return 1
    fi

    # shellcheck disable=SC2045
    for device in $(ls "${pci_device_path}");do
        if [[ ! -d "${pci_device_path}/${device}" ]];then
            continue
        fi
        vendor="${pci_device_path}/${device}"/vendor
        if [[ ! -f  "${vendor}" ]];then
            continue
        fi
        vendorId=$(cat "${vendor}")
        if [[ "${vendorId}" == "0x10de" ]];then
            logger.info "the instance is a gpu instance"
            _printf_end_split_line
            return 0
        fi
    done
    logger.error "no gpu cards found by pci devices, please confirm the instance is gpu instance"
    _printf_end_split_line
    return 1
}

# 检查是否存在gpu掉卡
_diagnose_gpu.lost_card() {
    local nvidia_smi_card_info
    local nvidia_smi_card_nums
    nvidia_smi_card_info=$(nvidia-smi -L)
    nvidia_smi_card_nums=$(echo "${nvidia_smi_card_info}" | wc -l)
    GPU_NUMS="${nvidia_smi_card_nums}"

    local lspci_card_info
    local lspci_card_num
    lspci_card_info=$(lspci -d 10de:)
    lspci_card_num=$(echo "${lspci_card_info}" | wc -l)
    logger.info "${nvidia_smi_card_info}"
    logger.info "${lspci_card_info}"
    if [[ "${nvidia_smi_card_nums}" -ne "${lspci_card_num}" ]]; then
        logger.error "gpu instance lost card, please refer to the self-service diagnostic document to troubleshoot the cause of the card loss."
        _printf_end_split_line
        return
    fi
    logger.info "gpu instance does not lose card"
    _printf_end_split_line
}

# 检查驱动安装版本是否和内核版本一致
_diagnose_gpu.check_os_and_driver_kernel_version() {
    os_kernel_version=$(uname -r)
    driver_kernel_version=$(modinfo nvidia | grep nvidia.ko | awk -F 'modules/' '{printf $2}' | awk -F '/' '{printf $1}')
    if [[ "${os_kernel_version}" != "${driver_kernel_version}" ]]; then
        logger.error "os kernel version ${os_kernel_version} is inconsistent with driver's kernel version ${driver_kernel_version} at installation"
        _printf_end_split_line
        return
    fi
    logger.info "os kernel version ${os_kernel_version} is consistent with driver install kernel version ${driver_kernel_version}"
    _printf_end_split_line
}

# 检查nouveau驱动是否被禁用
_diagnose_gpu.check_nouveau() {
    nouveau=$(lsmod | grep nouveau)
    if [[ -n "${nouveau}" ]]; then
        logger.error "nouveau driver hasn't been disabled, please refer to the self-service diagnostic document to disable it"
        _printf_end_split_line
        return
    fi
    logger.info "nouveau driver is disabled"
    _printf_end_split_line
}

# 检查是否存在xid错误
_diagnose_gpu.check_xid_error_in_dmesg() {
    local xid_info
    if xid_info=$(dmesg | grep -i xid); then
        logger.error "${xid_info}"
        logger.error "the instance has xid error, please refer to the self-service diagnostic document to solve the xid problem"
        _printf_end_split_line
        return
    fi
    logger.info "the instance doesn't have xid error"
    _printf_end_split_line
}

_diagnose_gpu.find_nvidia_smi_path() {
    local nvidia_smi_path
    nvidia_smi_path=$(find / -name nvidia-smi | grep nvidia-smi)
    logger.info "${nvidia_smi_path}"
    if [[ -z "${nvidia_smi_path}" ]]; then
        logger.error "driver is not install , please install it"
        return 1
    fi
    NVIDIA_SMI_PATH="${nvidia_smi_path}"
    logger.info "nvidia-smi path: ${nvidia_smi_path}"
}

# 检查nvidia-smi的返回码错误，判断gpu卡是否正常
_diagnose_gpu.nvidia_smi() {
    local smi_info
    smi_info=$("${NVIDIA_SMI_PATH}")
    ret_code=$?
    if [[ "${ret_code}" -ne 0 ]]; then
        logger.error "nvidia-smi returns error code ${ret_code}, please refer to the self-service diagnostic document for further disposal"
        _printf_end_split_line
        return
    fi
    logger.info "${smi_info}"
    _printf_end_split_line
}

# 收集 kernel_devel版本
_diagnose_gpu.collect_kernel_devel_version() {
    local kernel_devel_version
    if ! kernel_devel_version=$(cat /proc/version); then
        logger.error "kernel-devel is not existed, please check whether kernel-devel has been installed"
        _printf_end_split_line
        return
    fi
    logger.info "kernel-devel version is: "
    logger.echo "${kernel_devel_version}"
    _printf_end_split_line
}

# 收集cuda版本
_diagnose_gpu.collect_cuda_version() {
    local cuda_version
    if ! cuda_version=$(ls -l /usr/local/cuda); then
        logger.error "cuda is not installed, if need cuda, please refer to the self-service diagnostic document to install cuda"
        _printf_end_split_line
        return
    fi
    logger.info "${cuda_version}"
    _printf_end_split_line
}

# 收集dmesg中nvidia的报错
_diagnose_gpu.collect_nvidia_error_in_dmesg() {
    local nvidia_error_info
    # 过滤掉认证失败的错误
    nvidia_error_info=$(dmesg -T | grep -E "NVIDIA|nouveau|nvidia|nv" | grep -E "error|failed" | grep -v "nvidia: module verification failed")
    if [[ -n "${nvidia_error_info}" ]];then
        logger.error "${nvidia_error_info}"
        logger.error "please refer to the self-service diagnostic document to solve the problem"
        _printf_end_split_line
    fi
    logger.info "the instance doesn't have nvidia error in dmesg"
    _printf_end_split_line
}

# 收集nvidia-smi的详细信息，查看是否有ecc等报错，并根据自助诊断文档找到相应的解决办法
_diagnose_gpu.check_nvidia_smi_detail() {
    for((gpu_id=0;gpu_id<"${GPU_NUMS}";gpu_id++));
    do
        _diagnose_gpu.check_fan_speed "${gpu_id}"
        _diagnose_gpu.check_temperatures "${gpu_id}"
        _diagnose_gpu.check_pstate "${gpu_id}"
        local ecc_mode
        if ! _diagnose_gpu.check_ecc_mode "${gpu_id}"; then
            logger.error "ecc mode is not enabled, please enable it"
            _printf_end_split_line
            continue
        fi
        if _diagnose_gpu.is_exist_sram; then
            _diagnose_gpu.check_sram_ecc_count "${gpu_id}"
            _diagnose_gpu.check_dram_ecc_count "${gpu_id}"
        else
            _diagnose_gpu.check_volatile_ecc_count "${gpu_id}"
            _diagnose_gpu.check_aggregate_ecc_count "${gpu_id}"
            _diagnose_gpu.check_retired_pages_ecc_count "${gpu_id}"
        fi
    done
}

_diagnose_gpu.is_exist_sram() {
    local sram_info
    sram_info=$(nvidia-smi -q | grep -i sram)
    if [[ -n "${sram_info}" ]]; then
        logger.info "the gpu is after ampere architecture"
        return 0
    fi
    logger.info "the gpu is before ampere architecture"
    return 1
}

_diagnose_gpu.check_dram_ecc_count() {
    local gpu_id=$1
    local value_list
    value_list=$(nvidia-smi -i "${gpu_id}" --query-remapped-rows=remapped_rows.correctable,remapped_rows.uncorrectable --format=csv,noheader,nounits)
    local dram_ecc_count_info
    local gpu_bus_id
    dram_ecc_count_info=$(echo "${value_list}" | sed s/[[:space:]]//g)
    IFS=","
    value_array=(${dram_ecc_count_info})
    local correctable="${value_array[0]}"
    local uncorrectable="${value_array[1]}"
    logger.info "correctable dram ecc count is ${correctable}"
    if [[ "${correctable}" != *"N/A"* && "${correctable}" -gt 1000 ]]; then
        logger.error "gpu ${gpu_id} instance need restart due to correctable dram ecc count more than 1000"
    fi
    logger.info "uncorrectable dram ecc count is ${uncorrectable}"
    if [[ "${uncorrectable}" != *"N/A"* && "${uncorrectable}" -gt 60 ]]; then
        logger.error "gpu ${gpu_id} instance need restart due to uncorrectable dram ecc count more than 60"
    fi
    _printf_end_split_line
}

_diagnose_gpu.check_sram_ecc_count() {
    local gpu_id=$1
    local value_list
    local sram_ecc_count_info
    local gpu_bus_id
    value_list=$(nvidia-smi -i "${gpu_id}" --query-gpu=pci.bus_id,ecc.errors.corrected.volatile.sram,ecc.errors.corrected.aggregate.sram --format=csv,noheader,nounits)
    sram_ecc_count_info=$(echo "${value_list}" | sed s/[[:space:]]//g)
    IFS=","
    value_array=(${sram_ecc_count_info})
    gpu_bus_id="${value_array[0]}"
    local volatile_sram_ecc_count="${value_array[1]}"
    local aggregate_sram_ecc_count="${value_array[2]}"
    if [[ "${volatile_sram_ecc_count}" -gt 4 ]]; then
        logger.error "gpu ${gpu_bus_id} have hardware malfunction due to volatile sram ecc count ${volatile_sram_ecc_count} more than 4"
        logger.info "please refer to the self-service diagnostic document to solve the problem"
    fi
    if [[ "${aggregate_sram_ecc_count}" -gt 4 ]]; then
        logger.error "gpu ${gpu_bus_id} have hardware malfunction due to aggregate sram ecc count ${aggregate_sram_ecc_count} more than 4"
        logger.info "please refer to the self-service diagnostic document to solve the problem"
    fi
    logger.info "volatile sram ecc count is ${volatile_sram_ecc_count}"
    logger.info "aggregate sram ecc count is ${aggregate_sram_ecc_count}"
    _printf_end_split_line
}

_diagnose_gpu.check_retired_pages_ecc_count() {
    local gpu_id=$1
    local value_list
    value_list=$(nvidia-smi -i "${gpu_id}" --query-gpu=pci.bus_id,retired_pages.single_bit_ecc.count, \
                                                      retired_pages.double_bit.count, \
                                                      retired_pages.pending --format=csv,noheader,nounits)
    local retired_pages_ecc_count_info
    local gpu_bus_id
    retired_pages_ecc_count_info=$(echo "${value_list}" | sed s/[[:space:]]//g)
    IFS=","
    value_array=(${retired_pages_ecc_count_info})
    gpu_bus_id="${value_array[0]}"
    local retired_pages_ecc_count
    retired_pages_ecc_count=$((value_array[1] + value_array[2]))
    if [[ "${retired_pages_ecc_count}" -gt 60 ]]; then
        logger.error "gpu ${gpu_bus_id} have hardware malfunction due to double_bit_ecc count and single_bit_ecc more than 60"
        logger.error "please refer to the self-service diagnostic document to solve the problem"
    fi
    if [[ "${value_array[2]}" == "yes" ]]; then
        logger.error "the instance need restart to recover retired pages "
    fi
    logger.info "retired pages ecc count is ${aggregate_ecc_count}"
    _printf_end_split_line
}

_diagnose_gpu.check_aggregate_ecc_count() {
    local gpu_id=$1
    local value_list
    value_list=$(nvidia-smi -i "${gpu_id}" --query-gpu=pci.bus_id,ecc.errors.corrected.aggregate.device_memory,ecc.errors.corrected.aggregate.register_file,ecc.errors.corrected.aggregate.l1_cache,ecc.errors.corrected.aggregate.l2_cache,ecc.errors.corrected.aggregate.texture_memory,ecc.errors.corrected.aggregate.total --format=csv,noheader,nounits)
    local aggregate_ecc_count_info
    local gpu_bus_id
    aggregate_ecc_count_info=$(echo "${value_list}" | sed s/[[:space:]]//g)
    IFS=","
    value_array=(${aggregate_ecc_count_info})
    gpu_bus_id="${value_array[0]}"
    local aggregate_ecc_count
    aggregate_ecc_count=$((value_array[6] - value_array[1]))
    if [[ "${aggregate_ecc_count}" -gt 5 ]]; then
        logger.error "gpu ${gpu_bus_id} have hardware malfunction due to that aggregate ecc count more than 5"
        logger.error "please refer to the self-service diagnostic document to solve the problem"
    fi
    logger.info "aggregate ecc count is ${aggregate_ecc_count}"
    _printf_end_split_line
}

_diagnose_gpu.check_volatile_ecc_count() {
    local gpu_id=$1
    local value_list
    value_list=$(nvidia-smi -i "${gpu_id}" --query-gpu=pci.bus_id,ecc.errors.corrected.volatile.device_memory,ecc.errors.corrected.volatile.register_file,ecc.errors.corrected.volatile.l1_cache,ecc.errors.corrected.volatile.l2_cache,ecc.errors.corrected.volatile.texture_memory,ecc.errors.corrected.volatile.total --format=csv,noheader,nounits)
    local volatile_ecc_count_info
    local gpu_bus_id
    volatile_ecc_count_info=$(echo "${value_list}" | sed s/[[:space:]]//g)
    IFS=","
    value_array=(${volatile_ecc_count_info})
    gpu_bus_id="${value_array[0]}"
    local volatile_ecc_count
    volatile_ecc_count=$((value_array[6] - value_array[1]))
    if [[ "${volatile_ecc_count}" -gt 5 ]]; then
        logger.error "gpu ${gpu_bus_id} have hardware malfunction due to volatile ecc count more than 5"
        logger.error "please refer to the self-service diagnostic document to solve the problem"
    fi
    logger.info "volatile ecc count is ${volatile_ecc_count}"
    _printf_end_split_line
}

_diagnose_gpu.check_ecc_mode() {
    local gpu_id=$1
    local value_list
    value_list=$(nvidia-smi -i "${gpu_id}" --query-gpu=pci.bus_id,ecc.mode.current --format=csv,noheader,nounits)
    # shellcheck disable=SC2001
    local ecc_mode_info
    local value_array
    local gpu_bus_id
    local ecc_mode
    ecc_mode_info=$(echo "${value_list}" | sed s/[[:space:]]//g)
    IFS=","
    value_array=(${ecc_mode_info})
    gpu_bus_id="${value_array[0]}"
    ecc_mode="${value_array[1]}"
    logger.info "gpu ${gpu_bus_id} ecc mode is ${ecc_mode}"
    if [[ ${ecc_mode} != "Enabled" ]]; then
        logger.error "ecc mode is ${ecc_mode}"
        return 1
    fi
    _printf_end_split_line
}

_diagnose_gpu.check_fan_speed() {
    local gpu_id=$1
    local value_list
    value_list=$(nvidia-smi -i "${gpu_id}" --query-gpu=pci.bus_id,fan.speed --format=csv,noheader,nounits)
    # shellcheck disable=SC2001
    local fan_info
    local value_array
    local gpu_bus_id
    fan_info=$(echo "${value_list}" | sed s/[[:space:]]//g)
    IFS=","
    value_array=(${fan_info})
    gpu_bus_id="${value_array[0]}"
    fan_speed="${value_array[1]}"
    if [[ "${fan_speed}" == *"N/A"* || "${fan_speed}" == *"Not Supported"* ]]; then
        logger.info "gpu ${gpu_bus_id} is not support fan"
        _printf_end_split_line
        return
    fi
    if [[ "${fan_speed}" == 0 || "${fan_speed}" == "ERR!" ]]; then
        logger.error "gpu ${gpu_bus_id} fan speed abnormal, please refer to the self-service diagnostic document for further disposal"
        _printf_end_split_line
        return
    fi
    logger.info "gpu ${gpu_bus_id} fan speed normal"
    _printf_end_split_line
}

_diagnose_gpu.check_temperatures() {
    local gpu_id=$1
    local value_list
    value_list=$(nvidia-smi -i "${gpu_id}" --query-gpu=pci.bus_id,temperature.gpu --format=csv,noheader,nounits)
    # shellcheck disable=SC2001
    local temp_info
    local value_array
    local gpu_bus_id
    local gpu_core_temp
    temp_info=$(echo "${value_list}" | sed s/[[:space:]]//g)
    IFS=","
    value_array=(${temp_info})
    gpu_bus_id="${value_array[0]}"
    gpu_core_temp="${value_array[1]}"
    local max_core_op_temp
    max_core_op_temp=$(nvidia-smi -q -i "${gpu_id}" | grep -i 'GPU Max Operating Temp' | awk -F ': ' '{print $2}' | awk -F ' C' '{print $1}')
    if [[ "${gpu_core_temp}" -ge "${max_core_op_temp}" ]]; then
        logger.info "gpu ${gpu_bus_id} core temperature is greater than max operating temperature, please stop working and stop instance"
        _printf_end_split_line
        return
    fi
    logger.info "gpu ${gpu_bus_id} temperature normal"
    _printf_end_split_line
}

_diagnose_gpu.check_pstate() {
    local gpu_id=$1
    local value_list
    value_list=$(nvidia-smi -i "${gpu_id}" --query-gpu=pci.bus_id,pstate --format=csv,noheader,nounits)
    # shellcheck disable=SC2001
    local temp_info
    local value_array
    local gpu_bus_id
    local gpu_pstate_level
    pstate_info=$(echo "${value_list}" | sed s/[[:space:]]//g)
    IFS=","
    value_array=(${pstate_info})
    gpu_bus_id="${value_array[0]}"
    gpu_pstate_level="${value_array[1]}"
    logger.info "gpu ${gpu_bus_id} pstate level ${gpu_pstate_level}"
    logger.info "if pstate is not in p0 ~ p15, please restart instance to recover it, please refer to the self-service diagnostic document for further disposal"
    _printf_end_split_line
}

_diagnose_gpu.collect_driver_install_log() {
    cp /var/log/nvidia-installer.log "${DIAGNOSE_DIR}"
    logger.info "success to collect nvidia driver install log"
    _printf_end_split_line
}

_diagnose_gpu.collect_nvidia_log() {
    (
        # shellcheck disable=SC2164
        cd /root
        bug_report_sh_path=$(find / -name nvidia-bug-report.sh)
        # shellcheck disable=SC2128
        if [[ -z "${bug_report_sh_path}" ]]; then
            return
        fi
        bash "${bug_report_sh_path}" > /dev/null
        mv nvidia-bug-report.log.gz "${DIAGNOSE_DIR}"
    )
    logger.info "success to collect nvidia bug report log"
    _printf_end_split_line
}

_diagnose_gpu.collect_log() {
    _diagnose_gpu.collect_cuda_version
    _diagnose_gpu.collect_kernel_devel_version
    _diagnose_gpu.collect_nvidia_error_in_dmesg
    _diagnose_gpu.collect_driver_install_log
    _diagnose_gpu.collect_nvidia_log
}

_diagnose_gpu.run() {
    if ! _diagnose_gpu.is_gpu; then
        logger.error "the instance is not a gpu instance, skip collecting gpu diagnose information"
        return 1
    fi
    if ! _diagnose_gpu.find_nvidia_smi_path; then
        return 1
    fi
    _diagnose_gpu.nvidia_smi
    _diagnose_gpu.lost_card
    _diagnose_gpu.check_os_and_driver_kernel_version
    _diagnose_gpu.check_nouveau
    _diagnose_gpu.check_xid_error_in_dmesg
    _diagnose_gpu.check_nvidia_smi_detail
    _diagnose_gpu.collect_log
}

_diagnose_gpu.tar() {
    tar -zcvPf "${DIAGNOSE_DIR}".tar.gz "${DIAGNOSE_DIR}" > /dev/null
    echo -e "\033[31mIf our suggestions cannot solve your problem, please upload ${DIAGNOSE_DIR}.tar.gz to oncall with ticket\033[0m"
}

_diagnose_gpu.clean() {
    if [[ -d "${DIAGNOSE_DIR}" ]]; then
        rm -rf "${DIAGNOSE_DIR}"
    fi
}

diagnose_gpu() {
    logger.info "now start diagnosing gpu"
    if ! _diagnose_gpu.run; then
        return 1
    fi
    _diagnose_gpu.tar
    _diagnose_gpu.clean
    echo "gpu diagnosis ended "
}

diagnose_gpu "$@"