#!/bin/sh

AUTO_INSTALL="/root/auto_install"
CUDA_DIR="/root/auto_install/cuda"

mkdir -p ${AUTO_INSTALL}
mkdir -p ${CUDA_DIR}
log=${AUTO_INSTALL}"/auto_install.log"
PROCESS_NAME="$0"

DRIVER_PROCESS_NAME="${AUTO_INSTALL}/NVIDIA-Linux-x86_64"
CUDA_PROCESS_NAME="${AUTO_INSTALL}/cuda"
DOWNLOAD_DRIVER="NVIDIA-Linux"


CUDNN_PROCESS_NAME="${AUTO_INSTALL}/cudnn"
NCCL_PROCESS_NAME="${AUTO_INSTALL}/nccl"
DOWNLOAD_PROCESS_NAME="wget"

SUCCESS_STR="ALL INSTALL OK"
DOWNLOAD_SUCCESS_STR="Download OK"
REBOOT_STR="Please execute reboot command for the driver installation to take effect."

DRIVER_FAIL_STR="Driver INSTALL FAIL"
CUDA_FAIL_STR="CUDA INSTALL FAIL"
CUDNN_FAIL_STR="CUDNN INSTALL FAIL"
NCCL_FAIL_STR="NCCL INSTALL FAIL"
DOWNLOAD_FAIL_STR="Download FAIL"

install_notes="The script automatically downloads and installs a NVIDIA GPU driver and CUDA, CUDNN library.
1. The installation takes 15 to 20 minutes, depending on the intranet bandwidth and the quantity of vCPU cores of the instance. Please do not operate the GPU or install any GPU-related software until the GPU driver is installed successfully.
2. After the GPU is installed successfully, user must restart instance to make nvidia driver take effect."


check_install()
{
    b=''
    if [ "$1" = "NVIDIA" ]; then
        ProcessName=${DRIVER_PROCESS_NAME}
        t=2
    elif [ "$1" = "cuda" ]; then
        ProcessName=${CUDA_PROCESS_NAME}
        t=3
    elif [ "$1" = "cudnn" ]; then
        ProcessName=${CUDNN_PROCESS_NAME}
        t=0.5
    elif [ "$1" = "nccl" ]; then
        ProcessName=${NCCL_PROCESS_NAME}
        t=0.5
    fi
    i=0
    while true
    do
        pid_num=$(ps -ef | grep "${ProcessName}" | grep -v grep | wc -l)
        if [ $pid_num -eq 0 ]; then
            str=$(printf "%-100s" "#")
            b=$(echo "$str" | sed 's/ /#/g')
            printf "| %-100s | %d%% \r\n" "$b" "100";
            break
        fi
        i=$(($i+1))
        if [ $i -ge 90 ];then
            i=90
        fi
        str=$(printf "%-${i}s" "#")
        b=$(echo "$str" | sed 's/ /#/g')
        printf "| %-100s | %d%% \r" "$b" "$i";
        sleep $t
    done
    echo
    return 0
}

check_download()
{
    name=$1
    i=0
    b=''
    filesize=0
    percent=0

    sleep 0.5
    while true
    do
        pid_num=$(ps -ef | grep wget |grep ${name} |grep -v grep | wc -l)
        if [ $pid_num -eq 0 ]; then
            filesize=$(du -sk ${AUTO_INSTALL}/${name}* | awk '{print $1}')
            str=$(printf "%-100s" "#")
            b=$(echo "$str" | sed 's/ /#/g')
            printf "%-8s| %-100s | %d%% \r\n" "${filesize}K" "$b" "100";
            break
        fi
        line=$(tail -2 ${log})
        filesize=$(echo $line | awk -F ' ' '{print $1}')
        percent=$(echo $line | awk -F '%' '{print $1}' | awk -F ' ' '{print $NF}')
        if [ "$percent" -ge 0 ] 2>/dev/null ;then
           str=$(printf "%-${percent}s" "#")
           b=$(echo "$str" | sed 's/ /#/g')
           printf "%-8s| %-100s | %d%% \r" "${filesize}" "$b" "$percent";
        else
            continue
        fi
        sleep 0.5

    done
    return 0
}

check_install_log()
{
    if [ ! -f "$log" ];then
        echo "NVIDIA install log $log not exist! Install may fail!"
        echo
        exit 1
    fi

    if [ "$1" = "NVIDIA" ]; then
        succstr=$(cat $log |grep "${SUCCESS_STR}")
        str2=$(cat $log |grep "INSTALL_ERROR")
        if [ -n "${succstr}" ] && [ -z "${str2}" ]; then
            echo "${succstr} !!"
            echo "${REBOOT_STR} !!"
            echo
            return 0
        else
            echo "Install may have some INSTALL_ERROR, please execute 'nvidia-smi' to check driver installed, if driver not install please check log $log !"
            return 1
        fi
    fi

    if [ "$1" = "DRIVER" ]; then
        failstr=${DRIVER_FAIL_STR}
    elif [ "$1" = "CUDA" ]; then
        failstr=${CUDA_FAIL_STR}
    elif [ "$1" = "CUDNN" ]; then
        failstr=${CUDNN_FAIL_STR}
    elif [ "$1" = "NCCL" ]; then
        failstr=${NCCL_FAIL_STR}
    fi
    str1=$(cat $log |grep "${failstr}")
    if [ -n "${str1}" ] ;then
        echo
        echo "${failstr} ! please check install log ${log} !"
        return 1
    fi
}

check_install_process()
{
    if [ "$1" = "NULL" ]; then
        echo "CHECKING AUTO INSTALL,INSTALL RDMA=${6}, PLEASE WAIT ......"
    else
        echo "CHECKING AUTO INSTALL, DRIVER_VERSION=${1} CUDA_VERSION=${2} CUDNN_VERSION=${3}, PLEASE WAIT ......"
    fi
    echo "$install_notes"
    echo

    while true
    do
        pid_num=$(ps -ef | grep ${PROCESS_NAME} |grep -v grep | grep -v check | wc -l)
        if [ $pid_num -eq 0 ];then
            check_install_log "NVIDIA"
            return 0
        else
            pid_num=$(ps -ef | grep ${DOWNLOAD_PROCESS_NAME} |grep driver | grep -v rdma |grep -v grep | wc -l)
            if [ $pid_num -gt 0 ];then
                echo "Driver-${1} downloading, it takes 30 seconds or more. Remaining installation time 15 to 20 minutes!"
                check_download ${DOWNLOAD_DRIVER}
            fi

            pid_num=$(ps -ef | grep ${DOWNLOAD_PROCESS_NAME} |grep cuda |grep -v nccl |grep -v rapids |grep -v miniconda |grep -v Tensor |grep -v grep | wc -l)
            if [ $pid_num -gt 0 ];then
                echo "CUDA-${2} downloading, it takes 3 minutes or more. Remaining installation time 14 - 19 minutes!"
                while true
                do
                    check_download "cuda"
                    sleep 1
                    pid_num=$(ps -ef | grep ${DOWNLOAD_PROCESS_NAME} |grep cuda |grep -v nccl |grep -v rapids |grep -v miniconda |grep -v Tensor | grep -v grep | wc -l)
                    if [ $pid_num -eq 0 ];then
                        break
                    fi
                done
            fi

            pid_num=$(ps -ef | grep ${DOWNLOAD_PROCESS_NAME} |grep cudnn |grep -v Tensor |grep -v grep | wc -l)
            if [ $pid_num -gt 0 ];then
                echo "cuDNN-${3} downloading, it tasks 1 minutes or more. Remaining installation time 12 - 16 minutes!"
                check_download "cudnn"
            fi


            pid_num=$(ps -ef | grep "${DRIVER_PROCESS_NAME}" |grep -v grep | wc -l)
            if [ $pid_num -gt 0 ];then
                echo
                echo "Driver-${1} installing, it tasks 1 to 3 minutes. Remaining installation time 11 to 15 minutes!"
                check_install "NVIDIA"
                check_install_log "DRIVER"
            fi
            pid_num=$(ps -ef | grep "${CUDA_PROCESS_NAME}" |grep -v nccl |grep -v grep | wc -l)
            if [ $pid_num -gt 0 ];then
                echo "CUDA-${2} installing, it tasks 2 to 5 minutes. Remaining installation time 9 to 12 minutes!"
                check_install "cuda"
                check_install_log "CUDA"
            fi
            pid_num=$(ps -ef | grep ${CUDNN_PROCESS_NAME} |grep -v grep | wc -l)
            if [ $pid_num -gt 0 ];then
                echo "cuDNN-${3} installing, it takes about 10 seconds. Remaining installation time 6 to 9 minutes!"
                check_install "cudnn"
                check_install_log "CUDNN"
            fi
            pid_num=$(ps -ef | grep ${NCCL_PROCESS_NAME} |grep -v grep | wc -l)
            if [ $pid_num -gt 0 ];then
                echo "NCCL installing, it taskes about 10 seconds. "
                check_install "nccl"
                check_install_log "NCCL"
            fi

        fi
        sleep 1
    done
}

create_nvidia_repo_centos()
{
    #baseurl_centos=$(curl http://100.100.100.200/latest/meta-data/source-address | head -1)
    #cudaurl=$baseurl_centos"/opsx/ecs/linux/rpm/cuda/${version}/\$basearch/"
    #driverurl=$baseurl_centos"/opsx/ecs/linux/rpm/driver/${version}/\$basearch/"
    driverurl="http://192.168.1.3/kernel"
    #echo "[ecs-cuda]" > /etc/yum.repos.d/nvidia.repo
    #echo "name=ecs cuda - \$basearch" >> /etc/yum.repos.d/nvidia.repo
    #echo "baseurl=$cudaurl" >> /etc/yum.repos.d/nvidia.repo
    #echo "enabled=1" >> /etc/yum.repos.d/nvidia.repo
    #echo "gpgcheck=0" >> /etc/yum.repos.d/nvidia.repo
    echo "[ecs-driver-kernel]" >> /etc/yum.repos.d/nvidia.repo
    echo "name=ecs driver kernel - \$basearch" >> /etc/yum.repos.d/nvidia.repo
    echo "baseurl=$driverurl" >> /etc/yum.repos.d/nvidia.repo
    echo "enabled=1" >> /etc/yum.repos.d/nvidia.repo
    echo "gpgcheck=0" >> /etc/yum.repos.d/nvidia.repo
    yum clean all >> $log 2>&1
    yum makecache >> $log 2>&1
}


create_repo_centos8()
{
    rename '.repo' '.repo.bak' /etc/yum.repos.d/*.repo
    wget https://mirrors.aliyun.com/repo/Centos-vault-8.5.2111.repo -O /etc/yum.repos.d/Centos-vault-8.5.2111.repo
    wget https://mirrors.aliyun.com/repo/epel-archive-8.repo -O /etc/yum.repos.d/epel-archive-8.repo
    
    sed -i 's/mirrors.cloud.aliyuncs.com/url_tmp/g'  /etc/yum.repos.d/Centos-vault-8.5.2111.repo &&  sed -i 's/mirrors.aliyun.com/mirrors.cloud.aliyuncs.com/g' /etc/yum.repos.d/Centos-vault-8.5.2111.repo && sed -i 's/url_tmp/mirrors.aliyun.com/g' /etc/yum.repos.d/Centos-vault-8.5.2111.repo
    sed -i 's/mirrors.aliyun.com/mirrors.cloud.aliyuncs.com/g' /etc/yum.repos.d/epel-archive-8.repo
}

create_repo_centos_vault()
{
    version_small=$(cat /etc/centos-release | grep -o "[[:digit:]].*[[:digit:]]\>")
    if [ "$version" = "7" -a "$version_small" != "7.9.2009" ]; then
       echo "no default kernel source, add centos-vault kernel updates"
       touch /etc/yum.repos.d/CentOS-Base-Vault.repo
       baseurl="http://mirrors-internal.cmecloud.cn/centos-vault/"$version_small"/os/\$basearch/"
       updatesurl="http://mirrors-internal.cmecloud.cn/centos-vault/"$version_small"/updates/\$basearch/"
       echo "[base]" >> /etc/yum.repos.d/CentOS-Base-Vault.repo
       echo "name=CentOS-Vault-Base" >> /etc/yum.repos.d/CentOS-Base-Vault.repo
       echo "baseurl=$baseurl" >> /etc/yum.repos.d/CentOS-Base-Vault.repo
       echo "enabled=1" >> /etc/yum.repos.d/CentOS-Base-Vault.repo
       echo "gpgcheck=0" >> /etc/yum.repos.d/CentOS-Base-Vault.repo
       echo "[updates]" >> /etc/yum.repos.d/CentOS-Base-Vault.repo
       echo "name=CentOS-Vault-Updates" >> /etc/yum.repos.d/CentOS-Base-Vault.repo
       echo "baseurl=$updatesurl" >> /etc/yum.repos.d/CentOS-Base-Vault.repo
       echo "enabled=1" >> /etc/yum.repos.d/CentOS-Base-Vault.repo
       echo "gpgcheck=0" >> /etc/yum.repos.d/CentOS-Base-Vault.repo
    fi
    yum clean all >> $log 2>&1
    yum makecache fast >> $log 2>&1
}

create_repo_bclinux_updates()
{
    touch /etc/yum.repos.d/BClinux-Base-Updates.repo
    if [ "$version" = "7" ] ; then
       echo "no default kernel source, add bclinux small version kernel updates"
       version_small=$(cat /etc/system-release | grep -o "[[:digit:]].*[[:digit:]]\>"|cut -b 1-3)
       baseurl="http://mirrors-internal.cmecloud.cn/bclinux/el"$version_small"/os/\$basearch/"
       updatesurl="http://mirrors-internal.cmecloud.cn/bclinux/el"$version_small"/updates/\$basearch/"
       echo "[bclinuxbase]" > /etc/yum.repos.d/BClinux-Base-Updates.repo
       echo "name=BClinux-Base" >> /etc/yum.repos.d/BClinux-Base-Updates.repo
       echo "baseurl=$baseurl" >> /etc/yum.repos.d/BClinux-Base-Updates.repo
       echo "enabled=1" >> /etc/yum.repos.d/BClinux-Base-Updates.repo
       echo "gpgcheck=0" >> /etc/yum.repos.d/BClinux-Base-Updates.repo
       echo "[bclinuxupdates]" >> /etc/yum.repos.d/BClinux-Base-Updates.repo
       echo "name=BClinux-Updates" >> /etc/yum.repos.d/BClinux-Base-Updates.repo
       echo "baseurl=$updatesurl" >> /etc/yum.repos.d/BClinux-Base-Updates.repo
       echo "enabled=1" >> /etc/yum.repos.d/BClinux-Base-Updates.repo
       echo "gpgcheck=0" >> /etc/yum.repos.d/BClinux-Base-Updates.repo
    fi
    yum clean all >> $log 2>&1
    yum makecache fast >> $log 2>&1
    if [ ! -f "/usr/bin/gcc" ]; then
        yum downgrade -y glibc glibc-common
        yum install -y gcc gcc-c++
    fi
}

disable_nouveau_centos()
{
    if  [ ! -f /etc/modprobe.d/blacklist-nouveau.conf ];then
        echo "blacklist nouveau" > /etc/modprobe.d/blacklist-nouveau.conf
        echo "options nouveau modeset=0" >> /etc/modprobe.d/blacklist-nouveau.conf
    fi
    content=$(lsmod |grep nouveau)
    if [ -n "$content" ];then
        rmmod nouveau
        echo "***exec \"dracut --force\" to regenerate the kernel initramfs"
        dracut --force
    fi
}

disable_nouveau_bclinux()
{
    if  [ ! -f /etc/modprobe.d/blacklist.conf ];then
        echo "blacklist nouveau" > /etc/modprobe.d/blacklist.conf
        echo "options nouveau modeset=0" >> /etc/modprobe.d/blacklist.conf
    fi
    content=$(lsmod |grep nouveau)
    if [ -n "$content" ];then
        rmmod nouveau
        echo "***exec \"dracut --force\" to regenerate the kernel initramfs"
        dracut --force
    fi
}

disable_nouveau_ubuntu()
{
    if  [ ! -f /etc/modprobe.d/blacklist-nouveau.conf ];then
        echo "blacklist nouveau" > /etc/modprobe.d/blacklist-nouveau.conf
        echo "blacklist lbm-nouveau" >> /etc/modprobe.d/blacklist-nouveau.conf
        echo "options nouveau modeset=0" >> /etc/modprobe.d/blacklist-nouveau.conf
    fi
    content=$(lsmod |grep nouveau)
    if [ -n "$content" ];then
        rmmod nouveau
        echo "***exec \"update-initramfs -u\" to regenerate the kernel initramfs"
        update-initramfs -u
    fi
}
install_kernel_centos()
{
    kernel_version=$(uname -r)
    kernel_devel_num=$(rpm -qa | grep kernel-devel | grep $kernel_version | wc -l)
    if [ $kernel_devel_num -eq 0 ];then
        echo "******exec \"yum install -y kernel-devel-$kernel_version\""
        yum install -y kernel-devel-$kernel_version
        if [ $? -ne 0 ]; then
            echo "INSTALL_ERROR: install kernel-devel fail!!!"
            return 1
        fi
    fi
    return 0
}
install_kernel_bclinux()
{
    kernel_version=$(uname -r)
    kernel_devel_num=$(rpm -qa | grep kernel-devel | grep $kernel_version | wc -l)
    kernel_headers_num=$(rpm -qa | grep kernel-devel | grep $kernel_version | wc -l)
    #remove kernel bek rpms installed default
    kernel_bek_devel_num=$(rpm -qa | grep kernel-bek-devel | wc -l)
    kernel_bek_headers_num=$(rpm -qa | grep kernel-bek-headers | wc -l)
    if [ $kernel_bek_devel_num -ne 0 ];then
        echo "******exec \"rpm -e kernel-bek-devel-* --nodeps\""
        rpm -e kernel-bek-devel-* --nodeps
        if [ $? -ne 0 ]; then
            echo "UNINSTALL_ERROR: uninstall kernel-bek-devel fail!!!"
            return 1
        fi
    fi
    if [ $kernel_bek_headers_num -ne 0 ];then
        echo "******exec \"rpm -e kernel-bek-headers-* --nodeps\""
        rpm -e kernel-bek-headers-* --nodeps
        if [ $? -ne 0 ]; then
            echo "UNINSTALL_ERROR: uninstall kernel-bek-headers fail!!!"
            return 1
        fi
    fi
    #Install default version of kernel-devel and kernel-headers
    if [ $kernel_devel_num -eq 0 ];then
        echo "******exec \"yum install -y kernel-devel-$kernel_version\""
        yum install -y kernel-devel-$kernel_version
        if [ $? -ne 0 ]; then
            echo "INSTALL_ERROR: install kernel-devel fail!!!"
            return 1
        fi
    fi
    if [ $kernel_headers_num -eq 0 ];then
        echo "******exec \"yum install -y kernel-headers-$kernel_version\""
        yum install -y kernel-headers-$kernel_version
        if [ $? -ne 0 ]; then
            echo "INSTALL_ERROR: install kernel-headers fail!!!"
            return 1
        fi
    fi
    return 0
}
install_kernel_sles()
{
    kernel_version=$(uname -r|awk -F'-' '{print $1"-"$2}')
    kernel_devel_num=$(rpm -qa | grep kernel-default-devel | wc -l)
    if [ $kernel_devel_num -eq 0 ];then
        echo "***exec \"zypper install -y kernel-default-devel=$kernel_version\""
        zypper install -y kernel-default-devel=$kernel_version
        if [ $? -ne 0 ]; then
            echo "INSTALL_ERROR: install kernel-default-devel fail!!!"
            return 1
        fi
    fi
}
install_kernel_ubuntu()
{
    kernel_version=$(uname -r)
    linux_headers_num=$(dpkg --list |grep linux-headers | grep $kernel_version | wc -l)
    if [ $linux_headers_num -eq 0 ];then
        echo "***exec \"apt-get install -y --allow-unauthenticated linux-headers-$kernel_version\""
        apt-get install -y --allow-unauthenticated linux-headers-$kernel_version
        if [ $? -ne 0 ]; then
            echo "INSTALL_ERROR: install linux-headers fail!!!"
            return 1
        fi
    fi
}


install_kernel_debian()
{
    apt-get update
    kernel_version=$(uname -r)
    if [ "${kernel_version}" = "4.19.0-17-amd64" ];then
	cd ${AUTO_INSTALL}
        wget -t 100 --timeout=10 ${download_url}/rdma/sccgn7ex/Debian10u10/linux-header.tar.gz
        tar zxf linux-header.tar.gz
        cd  linux-header
        dpkg -i linux-kbuild-4.19_4.19.194-3_amd64.deb
        dpkg -i linux-compiler-gcc-8-x86_4.19.194-3_amd64.deb
        dpkg -i linux-headers-4.19.0-17-common_4.19.194-3_all.deb
        dpkg -i linux-headers-4.19.0-17-amd64_4.19.194-3_amd64.deb
            if [ $? -ne 0 ]; then
                echo "INSTALL_ERROR: install linux-headers fail!!!"
                return 1
    	    fi

    else
        linux_headers_num=$(dpkg --list |grep linux-headers | grep $kernel_version | wc -l)
        if [ $linux_headers_num -eq 0 ];then
            echo "***exec \"apt-get install -y  linux-headers-$kernel_version\""
            apt-get install -y linux-headers-$kernel_version
            if [ $? -ne 0 ]; then
                echo "INSTALL_ERROR: install linux-headers fail!!!"
                return 1
    	    fi
        fi
    fi
}


download()
{
    cd ${AUTO_INSTALL}
    if [ ! -f ${driver_file} ]; then
        wget -t 100 --timeout=10 ${download_url}/driver/${driver_file}
    fi
    if [ $? -ne 0 ]; then
        echo "INSTALL_ERROR: Download driver fail!!! return: $?"
        return 1
    fi

    if [ "$os" = "centos" -a "$version" = "6" ];then
        if [ "${cuda_big_version}" = "8.0" -o "${cuda_big_version}" = "9.0" -o "${cuda_big_version}" = "9.2" \
             -o "${cuda_big_version}" = "10.0" ];then

            ar=$(curl ${download_url}/cuda/${cuda_version}/ > ./tmp)
            echo "${download_url}/cuda/${cuda_version}/"
            cudafilelist=$(cat ./tmp | perl -n -e'/>(cuda[^<]*)</ && print "$1 \n"' | grep -v ubuntu)
        else #cuda10.1 cuda10.2 , cuda11 not support rhel6
            ar=$(curl ${download_url}/cuda/${cuda_version}/rhel6/ > ./tmp)
            cudafilelist=$(cat ./tmp | perl -n -e'/"(cuda[^<]*)"/ && print "$1 \n"')
        fi
    else
        ar=$(curl ${download_url}/cuda/${cuda_version}/ > ./tmp)
        echo "${download_url}/cuda/${cuda_version}/"
        cudafilelist=$(cat ./tmp | perl -n -e'/"(cuda[^<]*)"/ && print "$1 \n"' | grep -v ubuntu)
    fi

    if [ -z "$cudafilelist" ]; then
        echo "INSTALL_ERROR: Download CUDA fail!!! get cuda-${cuda_version} filename fail!!"
        return 1
    fi

    cd ${CUDA_DIR}
    echo $cudafilelist
    for cudafile in $cudafilelist
    do
        sleep 1
        if [ "$os" = "centos" -a "$version" = "6" ];then
            if [ "${cuda_big_version}" = "8.0" -o "${cuda_big_version}" = "9.0" -o "${cuda_big_version}" = "9.2" \
                 -o "${cuda_big_version}" = "10.0" ];then
                wget -t 100 --timeout=10 ${download_url}/cuda/${cuda_version}/$cudafile
            else
                wget -t 100 --timeout=10 ${download_url}/cuda/${cuda_version}/rhel6/$cudafile
            fi
        else
            wget -t 100 --timeout=10 ${download_url}/cuda/${cuda_version}/$cudafile
        fi
        if [ $? -ne 0 ]; then
            echo "INSTALL_ERROR: Download CUDA fail!!! wget $cudafile fail! return: $?"
            return 1
        fi
    done
    chmod +x ${CUDA_DIR}/*

    cd ${AUTO_INSTALL}
    wget -t 100 --timeout=10 ${download_url}/cudnn/${cuda_big_version}/${cudnn_file}
    #if [ $? -ne 0 ]; then
    #    echo "INSTALL_ERROR: Download cuDNN fail!!! return :$?"
    #    return 1
    #fi

    chmod +x ${AUTO_INSTALL}/*
    echo "$DOWNLOAD_SUCCESS_STR !"
    return 0
}

install_driver()
{
    ${AUTO_INSTALL}/${driver_file} --silent
    if [ $? -ne 0 ]; then
        echo "INSTALL_ERROR: driver install fail!!!"
        return 1
    fi
    echo "DRIVER $driver_version install OK !"
    return 0
}

install_cuda()
{
    cd ${CUDA_DIR}
    cuda_file=$(ls -S | grep cuda | grep $cuda_version | head -1)
    echo "cuda file: "$cuda_file
    if [ -z "$cuda_file" ]
    then
        echo "INSTALL_ERROR: cuda file is null, cuda install fail!!!"
        return 1
    fi

    sh ${CUDA_DIR}/$cuda_file --silent --toolkit --samples --samplespath=/root
    if [ $? -ne 0 ]; then
        echo "INSTALL_ERROR: cuda install fail!!!"
        return 1
    fi

    cuda_patchfile=$(ls | grep cuda | grep $cuda_version | grep -v ${cuda_file})
    for cuda_patch in $cuda_patchfile
    do
        echo "install cuda patch file: "$cuda_patch
        sh ${CUDA_DIR}/${cuda_patch} --silent --installdir=/usr/local/cuda --accept-eula
        if [ $? -ne 0 ]; then
            echo "INSTALL_ERROR: cuda patch install fail!!!"
            return 1
        fi
    done
    echo "CUDA $cuda_version install OK !"
    return 0
}

install_cudnn()
{
    tar zxvf ${AUTO_INSTALL}/$cudnn_file -C /usr/local
    if [ $? -ne 0 ]; then
        echo "INSTALL_ERROR: CUDNN INSTALL FAIL !!!"
        return 1
    fi
    echo "CUDNN $cudnn_version install OK !"
    return 0
}

install_nccl()
{
    cd ${AUTO_INSTALL}
    #download nccl
    curl ${download_url}/nccl/${cuda_big_version}/ > ./tmp
    br=$(cat ./tmp | perl -n -e'/>nccl_(.*)-1\+cuda.*/ && print "$1 \n"')
    cr=$(echo "${br}" | sort -rV | head -n1)
    nccl_version=$(echo ${cr} | awk -F ' ' '{print $1}')
    echo "max nccl version:$nccl_version"
    nccl_dir="nccl_${nccl_version}-1+cuda${cuda_big_version}_x86_64"
    nccl_file="${nccl_dir}.txz"

    echo $nccl_file
    wget -t 100 --timeout=10 ${download_url}/nccl/${cuda_big_version}/${nccl_file}
    chmod +x $nccl_file
    tar xf ${AUTO_INSTALL}/${nccl_file} && cp -r ${nccl_dir} /usr/local/nccl
    if [ $? -ne 0 ]; then
        echo "INSTALL_ERROR: NCCL INSTALL FAIL !!!"
        return 1
    fi
    echo "NCCL $nccl_version install OK !"
    return 0
}

enable_pm()
{
    if [ "$os" = "centos" -o "$os" = "bclinux" ];then
        yum install bzip2 -y
    fi
    cd /usr/share/doc/NVIDIA_GLX-1.0/sample*
    bunzip2 nvidia-persistenced-init.tar.bz2
    tar xvf nvidia-persistenced-init.tar
    cd nvidia-persistenced-init && sh install.sh -u root
}

set_env()
{
    env_path="/usr/local/cuda-${cuda_big_version}/bin:"
    #env_library="/usr/local/cuda-${cuda_big_version}/lib64:/usr/local/nccl/lib:"
    env_library="/usr/local/cuda-${cuda_big_version}/lib64:"
    env1="export PATH=${env_path}\$PATH"
    env2="export LD_LIBRARY_PATH=${env_library}\$LD_LIBRARY_PATH"

    echo $env1 >> ${env_file}
    echo $env2 >> ${env_file}

}


if [ -f "/etc/os-release" ];then
    #os=$(cat /etc/os-release |grep "^ID="|awk -F '=' '{print $2}'|sed 's/\"//g')
    os=$(. /etc/os-release;echo $ID)
    if [ "$os" = "ubuntu" ];then
        profile_file="/root/.profile"
        env_file="/root/.bashrc"
    elif [ "$os" = "centos" ];then
        profile_file="/root/.bash_profile"
        env_file="/root/.bashrc"
    elif [ "$os" = "alinux" ];then
        profile_file="/root/.bash_profile"
        env_file="/root/.bashrc"
    elif [ "$os" = "sles" ];then
        env_file="/root/.bash_profile"
        profile_file="/root/.bash_profile"
    elif [ "$os" = "debian" ]; then
        profile_file="/root/.profile"
        env_file="/root/.bashrc"
    elif [ "$os" = "bclinux" ];then
        profile_file="/root/.bash_profile"
        env_file="/root/.bashrc"
    fi
else
    issue=$(cat /etc/issue | grep CentOS)
    if [ -n "$issue" ];then
        os="centos"
        env_file="/root/.bashrc"
        profile_file="/root/.bash_profile"
    fi
fi

if [ "$1" = "check" ];then
    check_install_process $2 $3 $4 $5 $6 $7
    sed -i '/auto_install/d' $profile_file
    exit 0
else
    driver_version="$1"
    cuda_version="$2"
    cudnn_version="$3"
    is_install_aiacc_train="$4"
    is_install_aiacc_inference="$5"
    is_install_rdma="$6"
    is_install_rapids="$7"

    echo "begin to install, driver: $driver_version, cuda: $cuda_version, cudnn: $cudnn_version " >> $log 2>&1
    driver_file="NVIDIA-Linux-x86_64-"${driver_version}".run"
    cuda_big_version=$(echo $cuda_version | awk -F'.' '{print $1"."$2}')
    cudnn_file="cudnn-"${cuda_big_version}"-linux-x64-v"${cudnn_version}".tgz"
    cudnn_big_version=$(echo $cudnn_version | awk -F'.' '{print $1}')

    echo "sh ${PROCESS_NAME} check $driver_version $cuda_version $cudnn_version ${is_install_aiacc_train} ${is_install_aiacc_inference} ${is_install_rdma} ${is_install_rapids}" | tee -a $profile_file
    #echo "sh ${PROCESS_NAME} check $driver_version $cuda_version $cudnn_version ${is_install_perseus}" | tee -a $profile_file
fi
echo "os:$os" >> $log 2>&1
ubuntu_version=""

if [ "$os" = "ubuntu" ]; then
    disable_nouveau_ubuntu >> $log 2>&1
    apt-get update
    if [ ! -f "/usr/bin/gcc" ]; then
	  apt-get install -y gcc
    fi	  
    apt-get install -y make

#    version=$(cat /etc/os-release |grep "VERSION_ID=" | awk -F '=' '{print $2}'|sed 's/\"//g')
    version=$(. /etc/os-release;echo $VERSION_ID)
    if [ "$version" = "16.04" ]; then
        ubuntu_version="ubuntu1604"
    elif [ "$version" = "18.04" ];then
        ubuntu_version="ubuntu1804"
    elif [ "$version" = "20.04" ];then
        ubuntu_version="ubuntu2004"
    elif [ "$version" = "22.04" ];then
        ubuntu_version="ubuntu2204"
        apt-get install -y bzip2 
    else
        echo "ERROR: Ubuntu version $version is not supported!" >> $log 2>&1
        exit 1
    fi

elif [ "$os" = "centos" ]; then
    disable_nouveau_centos >> $log 2>&1

    if [ ! -f "/usr/bin/gcc" ]; then
        yum install -y gcc
    fi


    if [ -f "/etc/os-release" ];then
        #version=$(cat /etc/os-release |grep "VERSION_ID=" | awk -F '=' '{print $2}'|sed 's/\"//g')
        version=$(. /etc/os-release;echo $VERSION_ID)
    else
        if [ ! -f "/usr/bin/lsb_release" ]; then
            pkgname=$(yum provides /usr/bin/lsb_release |grep centos|grep x86_64 |head -1 |awk -F: '{print $1}')
            if [ -z "$pkgname" ]; then
                echo "INSTALL_ERROR: /usr/bin/lsb_release pkg not exists!" >> $log 2>&1
                exit 1
            fi
            yum install -y $pkgname >> $log 2>&1
        fi
        str=$(lsb_release -r | awk -F'[:.]' '{print $2}')
        version=$(echo $str | sed 's/ //g')

    fi

    if [ "$version" = "8" ]; then
       echo "no nvidia source, install elfutils-libelf-devel"
      
       create_repo_centos8

       yum install elfutils-libelf-devel -y
    fi
    #create_nvidia_repo_centos
    create_repo_centos_vault

elif [ "$os" = "bclinux" ]; then
    disable_nouveau_bclinux >> $log 2>&1

    if [ ! -f "/usr/bin/gcc" ]; then
        yum install -y gcc gcc-c++
    fi


    if [ -f "/etc/os-release" ];then
        #version=$(cat /etc/os-release |grep "VERSION_ID=" | awk -F '=' '{print $2}'|sed 's/\"//g')
        version=$(. /etc/os-release;echo $VERSION_ID | cut -b 1)
    else
        if [ ! -f "/usr/bin/lsb_release" ]; then
            pkgname=$(yum provides /usr/bin/lsb_release |grep centos|grep x86_64 |head -1 |awk -F: '{print $1}')
            if [ -z "$pkgname" ]; then
                echo "INSTALL_ERROR: /usr/bin/lsb_release pkg not exists!" >> $log 2>&1
                exit 1
            fi
            yum install -y $pkgname >> $log 2>&1
        fi
        str=$(lsb_release -r | awk -F'[:.]' '{print $2}')
        version=$(echo $str | sed 's/ //g')

    fi
    if [ "$version" = "8" ]; then
        echo "no make tools in env, install it"
        yum install make elfutils-libelf-devel -y
    fi
    #Create repodata for bclinux default kernel
    create_repo_bclinux_updates
elif [ "$os" = "debian" ]; then
    version=$(. /etc/os-release;echo $VERSION_ID)
    echo "os is Debian, version:${version}" >> $log 2>&1
    apt-get update
else
    echo "INSTALL_ERROR: Invalid OS!! INSTALL FAIL!" >> $log 2>&1
    exit 1
fi


#baseurl=$(curl http://100.100.100.200/latest/meta-data/source-address | head -1)
#download_url="${baseurl}/opsx/ecs/linux/binary"
download_url="http://mirrors-internal.cmecloud.cn/bclinux/others/gpu"

install_kernel_${os} >> $log 2>&1
if [ $? -ne 0 ]; then
    echo "INSTALL_ERROR: kernel-devel install fail!!!" >> $log 2>&1
    exit 1
fi

begin_download=$(date '+%s')


download >> $log 2>&1
if [ $? -ne 0 ]; then
    exit 1
fi
end_download=$(date '+%s')
time_download=$((end_download-begin_download))
echo "NVIDIA download OK! Using time $time_download s !!" >> $log 2>&1

begin=$(date '+%s')
install_driver >> $log 2>&1
if [ $? -ne 0 ]; then
    exit 1
fi

enable_pm >> $log 2>&1

echo "NVIDIA install driver OK!!!" >> $log 2>&1

install_cuda >> $log 2>&1
if [ $? -ne 0 ]; then
    exit 1
fi
echo "NVIDIA install cuda OK!!"  >> $log 2>&1

install_cudnn >> $log 2>&1
if [ $? -ne 0 ]; then
    exit 1
fi
echo "NVIDIA install cudnn OK!!!" >> $log 2>&1

#install_nccl >> $log 2>&1
#if [ $? -ne 0 ]; then
#    exit 1
#fi


set_env
cd ${AUTO_INSTALL}
rm -f ${AUTO_INSTALL}/tmp
rm -rf ${AUTO_INSTALL}/cuda
rm -f ${AUTO_INSTALL}/NVIDIA*
rm -f ${AUTO_INSTALL}/nvidia*
rm -f ${AUTO_INSTALL}/cudnn*


end=$(date '+%s')
time_install=$((end-begin))
echo "Install using time $time_install !"
echo "Install using time $time_install !" >> $log 2>&1

lsmod |grep nvidia >> $log 2>&1
nvidia-smi >> $log 2>&1

echo  ${SUCCESS_STR} >> $log 2>&1
echo  ${REBOOT_STR} >> $log 2>&1
ldconfig
#sleep 60
#reboot

