Set up Docker for use with NVIDIA GPU on Ubuntu 22.04

setup-docker.sh

#!/bin/bash

# This script automates the installation and setup of Docker, Docker Compose, and NVIDIA Docker extensions on Ubuntu 22.04.
# It is designed to run without user interaction. Run this script with root privileges.
# Script URL: https://gist.github.com/fadenb/5eb42f646706891abe2c6f12799f6e89/raw/setup-docker.sh

# Ensure the script is run as root
if [[ $EUID -ne 0 ]]; then
    echo "This script must be run as root. Use 'sudo' to run it."
    exit 1
fi

# Function to adjust needrestart configuration
adjust_needrestart_config() {
    if [[ -f /etc/needrestart/needrestart.conf ]]; then
        sed -i 's/^#\(.*\$nrconf{restart} = 'l';\)/\1/' /etc/needrestart/needrestart.conf
    fi
}

# Function to update system packages
update_packages() {
    apt update -y || { echo "Failed to update packages"; exit 1; }
    apt -o Dpkg::Options::="--force-confold" upgrade -y || { echo "Failed to upgrade packages"; exit 1; }
}

# Function to install prerequisites
install_prerequisites() {
    apt install -y apt-transport-https ca-certificates curl software-properties-common || { echo "Failed to install prerequisites"; exit 1; }
}

# Function to install Docker
install_docker() {
    curl -fsSL https://download.docker.com/linux/ubuntu/gpg | apt-key add - || { echo "Failed to add Docker GPG key"; exit 1; }
    echo "deb [arch=amd64] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | tee /etc/apt/sources.list.d/docker.list
    apt update -y || { echo "Failed to update packages for Docker"; exit 1; }
    apt install -y docker-ce || { echo "Failed to install Docker"; exit 1; }
}

# Function to add the current user to Docker group
add_user_to_docker_group() {
    usermod -aG docker "${USER:?}" || { echo "Failed to add user to Docker group"; exit 1; }
}

# Function to install NVIDIA Docker Extensions with updated URL
install_nvidia_docker_extensions() {
    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
    apt update -y || { echo "Failed to update packages for NVIDIA Docker"; exit 1; }
    apt install -y nvidia-container-toolkit || { echo "Failed to install NVIDIA container toolkit"; exit 1; }
}

# Function to install Docker Compose
install_docker_compose() {
    compose_version=$(curl -s https://api.github.com/repos/docker/compose/releases/latest | grep -oP '(?<=tag_name\": \")([^\"]+)')
    curl -L "https://github.com/docker/compose/releases/download/$compose_version/docker-compose-$(uname -s)-$(uname -m)" -o /usr/local/bin/docker-compose || { echo "Failed to download Docker Compose"; exit 1; }
    chmod +x /usr/local/bin/docker-compose || { echo "Failed to set executable permissions on Docker Compose"; exit 1; }
}

# Main script execution
adjust_needrestart_config
update_packages
install_prerequisites
install_docker
add_user_to_docker_group
install_nvidia_docker_extensions
install_docker_compose

# Ensure the script runs without rebooting
echo "Setup is complete. Please reboot your system manually."

# End of the script