Setup script for docker + nvidia-container-toolkit for Ubuntu

docker-nvidia.sh

#!/bin/sh

command_exists() {
        command -v "$@" > /dev/null 2>&1
}

do_install() {
        set -e
        
        user="$(id -un 2>/dev/null || true)"

        sh_c='sh -c'
        if [ "$user" != 'root' ]; then
                if command_exists sudo; then
                        sh_c='sudo -E sh -c'
                elif command_exists su; then
                        sh_c='su -c'
                else
                        cat >&2 <<-'EOF'
                        Error: this installer needs the ability to run commands as root.
                        We are unable to find either "sudo" or "su" available to make this happen.
                        EOF
                        exit 1
                fi
        fi

        # Install docker
        echo '[*] Installing Docker using the install script'
        sh_c "curl -fsSL https://get.docker.com | sh"

        # Install drivers
        echo '[*] Installing GPU drivers using ubuntu-drivers'
        sh_c 'apt-get install -y ubuntu-drivers-common curl'
        sh_c 'ubuntu-drivers install'

        echo '[*] Installing NVIDIA Container Toolkit'
        sh_c $(cat <<EOF
                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
EOF
)
        sh_c 'apt-get update'
        sh_c 'apt-get install -y nvidia-container-toolkit'

        echo '[*] Configuring Docker container runtime'
        sh_c 'nvidia-ctk runtime configure --runtime=docker'

        echo '[*] Restarting Docker'
        sh_c 'systemctl restart docker'
        
        if [ "$user" != 'root' ]; then
                echo '[*] Adding current user to docker group'
                sh_c "usermod -aG docker $user"
        fi

        echo '[+] Done!'
}

do_install