#/bin/bash
########################################################
# 根据ibv_devices命令来判断机型：
#                             （1）IB设备名以mlx5_bond_开头则判断是A800、H800机器
#                             （2）否则如果以mlx5_开头则是DGX机器
#                             （3）如果没有该命令 且 存在命令 "npu-smi info"则判断其是华为NPU机器
#                             （4）以上条件均不满足， 则为未识别的设备
# usage:
#       集群检查： bash cluster_check.sh -f host.txt
#       单节点检查： bash cluster_check.sh -s
#       k8s模式集群检查： bash cluster_check.sh -f host.txt -k -kc kubectl.config -n namespace -c container_name
#                    kubectl.config： 凭证文件
#                    namespace： 命名空间
#                    container_name： 容器名（获取方式： kubectl get pods xxxx -o jsonpath={.spec.containers[*].name}）
#       host.txt： ip列表(每行一个 IP)
########################################################
script_name=""

if command -v ibv_devices &> /dev/null; then
  # 获取IB设备名
    ib_device_name=$(ibv_devices | awk '/mlx[0-7]_bond_/{print $1; exit} /mlx[0-7]_/{print $1; exit}')

    # 根据IB设备名判断机型
    if [[ $ib_device_name =~ mlx[0-7]_bond_.* ]]; then
        echo "Detected A800 or H800 machine, Using 'cluster_check.py'."
        script_name="cluster_check.py"
    elif [[ $ib_device_name =~ mlx[0-7]_.* ]]; then
        echo "Detected DGX machine, Using 'dgx_cluster_check.py'."
        script_name="dgx_cluster_check.py"
    else
        echo "Unrecognized device."
        exit 0
    fi
elif command -v npu-smi &> /dev/null; then
    # 检查npu-smi info命令是否存在
    if npu-smi info &> /dev/null; then
        echo "Detected Huawei NPU machine, Using 'ascend_cluster_check.py'."
        script_name="ascend_cluster_check.py"
    fi
else
    echo "Unrecognized device."
    exit 0
fi

# Only proceed if a script name has been set
if [[ -n $script_name ]]; then
    url="https://llm-diagnose-1258344699.cos.ap-guangzhou.myqcloud.com/tool/${script_name}"
    # Use wget or curl to download the script
    wget "${url}" -O "${script_name}" || curl -fLo "${script_name}" "${url}"
    # Check if the download was successful
    if [[ -f "${script_name}" ]]; then
        chmod +x "${script_name}"
        python "${script_name}" "$@"
    else
        echo "Failed to download the script."
        exit 1
    fi
else
    echo "No script name was determined."
    exit 0
fi
