#!/bin/bash

apt update -y &>/dev/null
apt install linux-headers-$(uname -r) -y &>/dev/null

# Download the GPU driver
wget "https://download.cloudclusters.io/A40.run" -O A40.run &>/dev/null
chmod 755 A40.run
./A40.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
