Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Assumes 2xT4 GPUs are attached to an ubuntu based image.
#!/bin/bash
OS_NAME=$(lsb_release -is | tr '[:upper:]' '[:lower:]')
readonly OS_NAME
OS_DIST=$(lsb_release -cs)
readonly OS_DIST
CUDA_VERSION='10.2'
readonly CUDA_VERSION
readonly DEFAULT_NVIDIA_DEBIAN_GPU_DRIVER_VERSION='460.56'
readonly DEFAULT_NVIDIA_DEBIAN_GPU_DRIVER_URL="https://us.download.nvidia.com/XFree86/Linux-x86_64/${DEFAULT_NVIDIA_DEBIAN_GPU_DRIVER_VERSION}/NVIDIA-Linux-x86_64-${DEFAULT_NVIDIA_DEBIAN_GPU_DRIVER_VERSION}.run"
readonly NVIDIA_BASE_DL_URL='https://developer.download.nvidia.com/compute'
# Parameters for NVIDIA-provided Ubuntu GPU driver
readonly NVIDIA_UBUNTU_REPOSITORY_URL="${NVIDIA_BASE_DL_URL}/cuda/repos/ubuntu1804/x86_64"
readonly NVIDIA_UBUNTU_REPOSITORY_KEY="${NVIDIA_UBUNTU_REPOSITORY_URL}/7fa2af80.pub"
readonly NVIDIA_UBUNTU_REPOSITORY_CUDA_PIN="${NVIDIA_UBUNTU_REPOSITORY_URL}/cuda-ubuntu1804.pin"
function execute_with_retries() {
local -r cmd=$1
for ((i = 0; i < 10; i++)); do
if eval "$cmd"; then
return 0
fi
sleep 5
done
return 1
}
function install_nvidia_gpu_driver() {
curl -fsSL --retry-connrefused --retry 10 --retry-max-time 30 \
"${NVIDIA_UBUNTU_REPOSITORY_KEY}" | apt-key add -
curl -fsSL --retry-connrefused --retry 10 --retry-max-time 30 \
"${NVIDIA_UBUNTU_REPOSITORY_CUDA_PIN}" -o /etc/apt/preferences.d/cuda-repository-pin-600
add-apt-repository "deb ${NVIDIA_UBUNTU_REPOSITORY_URL} /"
execute_with_retries "apt-get update"
if [[ -n "${CUDA_VERSION}" ]]; then
local -r cuda_package=cuda-${CUDA_VERSION//./-}
else
local -r cuda_package=cuda
fi
# Without --no-install-recommends this takes a very long time.
execute_with_retries "apt-get install -y -q --no-install-recommends ${cuda_package}"
echo "NVIDIA GPU driver provided by NVIDIA was installed successfully"
}
function main() {
# updates
export DEBIAN_FRONTEND=noninteractive
execute_with_retries "apt-get update"
execute_with_retries "apt-get install -y -q pciutils"
execute_with_retries "apt-get install -y -q 'linux-headers-$(uname -r)'"
install_nvidia_gpu_driver
}
main
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment