#!/bin/bash

BASE_DIR=/root/gpu_auto_install
PKG_DIR=$BASE_DIR/pkg
LOG_FILE=$BASE_DIR/gpu_install.log
GPU_DRIVER_VERSION=$1
CUDA_VERSION=$2
CUDNN_VERSION=$3
GPU_DRIVER="NVIDIA-Linux-x86_64-${GPU_DRIVER_VERSION}.run"
CUDA="cuda_${CUDA_VERSION}_linux.run"
CUDNN="cudnn-linux-x86_64-${CUDNN_VERSION}_cuda${CUDA_VERSION%%.*}.tar.gz"
DOWNLOAD_URL=100.100.100.101/nvidia
PROFILE="/root/.profile"
PID_FILE=$BASE_DIR/pid
PROCESS=`basename $0`
INSTALL_SCRIPT=`readlink -f $0`

function log_debug() {
    msg="$(date '+%Y-%m-%d %H:%M:%S') [DEBUG] $@"
    echo $msg >> $LOG_FILE && echo -e "\033[32m"  ${msg}  "\033[0m"
}

function log_info() {
    msg="$(date '+%Y-%m-%d %H:%M:%S') [INFO] $@"
    echo $msg >> $LOG_FILE && echo -e "\033[32m"  ${msg} "\033[0m" >> /dev/tty1
}

function log_warn() {
    msg="$(date '+%Y-%m-%d %H:%M:%S') [WARN] $@"
    echo $msg >> $LOG_FILE && echo -e "\033[33m" ${msg} "\033[0m" >> /dev/tty1
}

function log_error() {
    msg="$(date '+%Y-%m-%d %H:%M:%S') [ERROR] $@"
    echo $msg >> $LOG_FILE && echo -e "\033[31m" ${msg} "\033[0m" >> /dev/tty1
}

function init_env() {
    local reserved_size=20480

    free_size=$(df -BM / | awk 'NR==2{print $4}' | sed 's/[^0-9]//g')
    if [[ $free_size -lt $reserved_size ]]; then
        log_error "Insufficient disk space on root partition."
        return 1
    fi

    if [ -d $PKG_DIR ]; then
        log_error "The automatic installation has failed before. Please clean the installation directory."
        return 1
    fi

    if which nvidia-smi > /dev/null 2>&1; then
        log_debug "Gpu driver is already exist."
        return 1
    fi

    mkdir -p $PKG_DIR
    return 0
}

function download_package() {
    local save_path=$1
    local url=$2
    local retry_cnt=10
    local retry_interval=5

    log_debug "Download package $1 from $2"
    if [ -f $save_path ]; then
        rm -rf $save_path
    fi

    for i in $(seq 1 $retry_cnt); do
        wget -O "$save_path" "$url"
        if [ $? -eq 0 ]; then
            return 0
        fi
        log_warn "Download $save_path error, retry $i"
        sleep $retry_interval
    done
    log_error "Download package $1 from $2 failed."
    return 1
}

function disable_nouveau() {
    local black_list_nouveau="/etc/modprobe.d/blacklist-nouveau.conf"

    if [ -f "$black_list_nouveau" ]; then
        return
    fi
    log_info "Disable nouveau driver"
    cat <<EOF > $black_list_nouveau
blacklist nouveau
options nouveau modeset=0
EOF
    rmmod nouveau
}

function check_nvidia_driver_installed() {
    nvidia-smi >>$LOG_FILE 2>&1 || return 1
}

function install_gpu_driver() {
    if [ -z "$GPU_DRIVER_VERSION" ]; then
        return 0
    fi
    log_info "Download $GPU_DRIVER"
    download_package $PKG_DIR/$GPU_DRIVER $DOWNLOAD_URL/driver/$GPU_DRIVER
    if [ $? -ne 0 ]; then
        return 1
    fi
    log_info "Install $GPU_DRIVER"
    chmod +x $PKG_DIR/$GPU_DRIVER
    $PKG_DIR/$GPU_DRIVER -s >>$LOG_FILE 2>&1

    check_nvidia_driver_installed
    if [ $? -ne 0 ]; then
        log_error "Install $GPU_DRIVER failed"
        return 1
    fi
    return 0
}

function check_cuda_installed() {
    sed -i '$ a export PATH=/usr/local/cuda/bin:$PATH' /etc/profile
    sed -i '$ a export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' /etc/profile
    source /etc/profile
    /usr/local/cuda/bin/nvcc --version >>$LOG_FILE 2>&1 || return 1
}

function install_cuda() {
    if [ -z $CUDA_VERSION ]; then
        return 0
    fi
    log_info "Download $CUDA"
    download_package $PKG_DIR/$CUDA $DOWNLOAD_URL/cuda/$CUDA
    if [ $? -ne 0 ]; then
        return 1
    fi
    log_info "Install $CUDA"
    chmod +x $PKG_DIR/$CUDA
    $PKG_DIR/$CUDA --silent --toolkit >>$LOG_FILE 2>&1
    check_cuda_installed
    if [ $? -ne 0 ]; then
        log_error "Install $CUDA_DRIVER failed"
        return 1
    fi
    return 0
}

function install_cudnn() {
    if [ -z $CUDNN_VERSION ]; then
        return 0
    fi
    log_info "Download $CUDNN"
    download_package $PKG_DIR/$CUDNN $DOWNLOAD_URL/cudnn/$CUDNN
    if [ $? -ne 0 ]; then
        return 1
    fi
    tar -zxf $PKG_DIR/$CUDNN -C $PKG_DIR
    cp $PKG_DIR/cudnn-*-archive/include/cudnn*.h /usr/local/cuda/include
    cp -P $PKG_DIR/cudnn-*-archive/lib/libcudnn* /usr/local/cuda/lib64
    chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*
    return 0
}

function clean_env() {
    rm -rf $PKG_DIR
    reboot
}

function check_auto_install() {
    echo "Installing driver cuda cudnn, please wait 5-10 minitues. To avoid installation failures, do not operate the instance during installation. After the installation, the instance restarts automatically."

    while true
    do
        status=$(cloud-init status | grep running | wc -l)
        if [ $status -gt 0 ]; then
            echo "Installation running, you can check $LOG_FILE to get more infomation"
            sleep 10
        else
            break
        fi
    done

    if [ ! -f $LOG_FILE ]; then
        echo "Installation fail: $LOG_FILE not found"
        return 1
    fi

    succ_log_count=$(grep "Auto install succ" $LOG_FILE | wc -l)
    if [ $succ_log_count -gt 0 ]; then
        echo "Installation of DRIVER CUDA CUDNN already succeeded"
    else
        echo "Installation fail, please check log file $LOG_FILE"
    fi
}

if [ $# -eq 1 ] && [ $1 == "check" ]; then
    check_auto_install
    sed -i "/${PROCESS}/d" $PROFILE
else
    init_env || exit 1

    log_info "Auto install $GPU_DRIVER $CUDA $CUDNN"
    echo $$ > $PID_FILE
    echo "/bin/bash $INSTALL_SCRIPT check" | tee -a $PROFILE

    install_gpu_driver || exit 1
    install_cuda || exit 1
    install_cudnn || exit 1

    log_info "Auto install succ, reboot now"
    clean_env
fi

