Simple script for setting up NVIDIA software on bare-metal

setup_vmi_nvidia.sh

#!/usr/bin/env bash
# author: com/github/dualvtable

set -eo pipefail

check_root()
{
   if [[ $EUID -ne 0 ]]; then
      echo "This installer must be run as root."
      exit 1
   fi
}

usage()
{
cat << EOF
usage: $(basename "$0") [options]

Options:
 -h, --help             This help menu
EOF
}

#setup CUDA network repos
setup_repos()
{
    distribution=$(. /etc/os-release;echo $ID$VERSION_ID | sed -e 's/\.//g') \
    && wget https://developer.download.nvidia.com/compute/cuda/repos/$distribution/x86_64/cuda-keyring_1.1-1_all.deb \
    && dpkg -i cuda-keyring_1.1-1_all.deb
    #setup nvidia-docker2 repos

    distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
    && curl -s -L https://nvidia.github.io/libnvidia-container/gpgkey | sudo apt-key add - \
    && curl -s -L https://nvidia.github.io/libnvidia-container/stable/$distribution/libnvidia-container.list | sudo tee /etc/apt/sources.list.d/libnvidia-container.list

    echo "Setup CUDA and NVIDIA container package repos..." >&2
    sudo apt-get update
}

install_driver()
{
    driver_branch=$1
    if [ $driver_branch ]; then
       driver_branch="-$driver_branch"
    fi
    sudo apt-get install -y gcc make linux-headers-$(uname -r) golang-go \
    && sudo apt-get install -y cuda-drivers$driver_branch
}

check_docker()
{
    error_code=$(docker run --rm hello-world)
    if [ $error_code -ne 0 ]; then
       echo "Docker failed to install..." >&2
       exit $error_code
    fi
}

install_docker()
{
    curl https://get.docker.com | sh \
    && sudo systemctl start docker \
    && sudo systemctl enable docker

    check_docker
}

install_nvidia_runtime()
{
    sudo apt-get install -y nvidia-container-toolkit \
    && sudo systemctl restart docker
}

while true; do
   case "$1" in
      '-h'|'--help')
         usage
         shift
         continue
      ;;
      '--')
         shift
         break
      ;;
      *)
         echo "$@"
         usage
         exit 1
      ;;
   esac
done

check_root

setup_repos
install_docker
install_driver 550
install_nvidia_runtime