Assumes 2xT4 GPUs are attached to an ubuntu based image.

gpu_dataproc_packages.sh

#!/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