set -x

echo "---------------------------------------------------------"
echo "9 disable ipv6"
content="
net.ipv6.conf.all.disable_ipv6=1
net.ipv6.conf.default.disable_ipv6=1
net.ipv6.conf.lo.disable_ipv6=1"

file="/etc/sysctl.conf"

if [ -f "$file" ]; then
   echo -e "$content" >> "$file"
fi
sysctl -p
echo "---------------------------------------------------------"


echo "---------------------------------------------------------"
echo "3. install gpu driver"
apt update
apt install linux-headers-$(uname -r) -y

# Download the GPU driver
rm -f NVIDIA-Linux-x86_64-535.171.04.run
wget "https://us.download.nvidia.com/XFree86/Linux-x86_64/535.171.04/NVIDIA-Linux-x86_64-535.171.04.run" -O NVIDIA-Linux-x86_64-535.171.04.run &>/dev/null
chmod 755 NVIDIA-Linux-x86_64-535.171.04.run
./NVIDIA-Linux-x86_64-535.171.04.run -s &>/dev/null

# Function to check if the driver is installed
check_driver() {
    # Use the nvidia-smi command to get the list of GPU devices and filter out lines containing UUID using grep
    gpu_info=$(nvidia-smi -L | grep -i uuid)

    # Check if GPU information was successfully obtained
    if [ -n "$gpu_info" ]; then
        echo "Nvidia GPU driver is installed successfully"
        return 0  # Return success status
    else
        echo "Nvidia GPU driver is not installed or installation failed"
        return 1  # Return failure status
    fi
}

check_driver
echo "---------------------------------------------------------"


echo "---------------------------------------------------------"
echo "4. install docker"
apt-get update
apt-get install ca-certificates curl
install -m 0755 -d /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg -o /etc/apt/keyrings/docker.asc
chmod a+r /etc/apt/keyrings/docker.asc
# Add the repository to Apt sources:
echo \
  "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.asc] https://download.docker.com/linux/ubuntu \
  $(. /etc/os-release && echo "$VERSION_CODENAME") stable" | \
  tee /etc/apt/sources.list.d/docker.list > /dev/null

apt-get update
apt-get install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin -y
echo "---------------------------------------------------------"


echo "---------------------------------------------------------"
echo "5. NVIDIA CUDA Toolkit"
rm -f cuda_12.2.0_535.54.03_linux.run
wget "https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda_12.2.0_535.54.03_linux.run"
sh cuda_12.2.0_535.54.03_linux.run  --toolkit --no-drm  --silent
echo 'export PATH=/usr/local/cuda-12.2/bin:$PATH' >> /root/.bashrc
echo 'export PATH=/usr/local/cuda-12.2/bin:$PATH' >> /home/administrator/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.2/lib64:$LD_LIBRARY_PATH' >> /root/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.2/lib64:$LD_LIBRARY_PATH' >> /home/administrator/.bashrc
/usr/local/cuda-12.2/bin/nvcc -V
echo "---------------------------------------------------------"

echo "---------------------------------------------------------"
echo "6. install NVIDIA Container Toolkit"
rm -f /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg nvidia-container-toolkit.list
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
  && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \
    sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
    tee /etc/apt/sources.list.d/nvidia-container-toolkit.list
sed -i -e '/experimental/ s/^#//g' /etc/apt/sources.list.d/nvidia-container-toolkit.list
apt-get update
apt-get install -y nvidia-container-toolkit
echo "---------------------------------------------------------"

echo "---------------------------------------------------------"
echo "7. Configuring NVIDIA Container Runtime for Docker"
nvidia-ctk runtime configure --runtime=docker
systemctl restart docker
systemctl enable docker
docker run --gpus all nvidia/cuda:12.0.0-base-ubuntu22.04  nvidia-smi
echo "---------------------------------------------------------"

echo "---------------------------------------------------------"
echo "8. Fixed kernel version"
cur_kernel_name=$(uname -r)
kernel_uuid=$(df -P / | awk '/\// {print $1}' | xargs blkid -s UUID -o value)
kernel_uuid_str1="gnulinux-advanced-${kernel_uuid}>"
kernel_uuid_str2=gnulinux-${cur_kernel_name}-advanced-${kernel_uuid}
sed -i "s/GRUB_DEFAULT=0/GRUB_DEFAULT=\"${kernel_uuid_str1}${kernel_uuid_str2}\"/" /etc/default/grub
update-grub
echo "---------------------------------------------------------"
