30_install_cuda

#!/usr/bin/env bash

function install_cuda_rhel_repo()
{
  # Optional argument $1 - RHEL Major version (e.g. 9)
  # Optional argument $2 - CPU Arch (default x86_64)

  local rhel_major_version=${1-}

  if [ "${rhel_major_version}" = "" ]; then
    rhel_major_version=$(source /etc/os-release; echo "${VERSION_ID%%.*}")
  fi

  cp /usr/local/share/just/info/cuda/keys/RPM-GPG-KEY-NVIDIA /etc/pki/rpm-gpg/RPM-GPG-KEY-NVIDIA
  cat - > /etc/yum.repos.d/cuda.repo << EOF
[cuda]
name=cuda
baseurl=https://developer.download.nvidia.com/compute/cuda/repos/rhel${rhel_major_version}/${2-x86_64}
enabled=1
gpgcheck=1
gpgkey=file:///etc/pki/rpm-gpg/RPM-GPG-KEY-NVIDIA
EOF
}

function install_cuda_ubuntu_repo()
{
  # Optional argument $1 - RHEL Major version (e.g. 9)
  # Optional argument $2 - CPU Arch (default x86_64)

  local ubuntu_major_version=${1-}

  if [ "${ubuntu_major_version}" = "" ]; then
    ubuntu_major_version=$(source /etc/os-release; echo "${VERSION_ID%%.*}")
  fi

  # The nvidia repo can't be downloaded without this.
  if ! dpkg -s ca-certificates &> /dev/null; then
    apt-get update
    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends ca-certificates
  fi

  # The "new" way to add a repo and key since at least 16.04
  cat /usr/local/share/just/info/cuda/keys/* > /usr/share/keyrings/cuda.asc
  echo "deb [signed-by=/usr/share/keyrings/cuda.asc] https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${ubuntu_major_version}04/${2-x86_64} /" > /etc/apt/sources.list.d/cuda.list
}

function install_cuda_repo_and_packages()
{
  source /usr/local/share/just/user_run_patch/10_load_cuda_env

  packages=()

  install_base=0
  install_runtime=0
  install_devel=0

  case "${CUDA_RECIPE_TARGET-runtime}" in
    runtime)
      install_base=1
      install_runtime=1
      ;;
    devel)
      install_base=1
      install_runtime=1
      install_devel=1
      ;;
    devel-only)
      install_devel=1
      ;;
  esac

  OLD_IFS="${IFS}"
  IFS=.
  cuda_version=(${CUDA_VERSION})
  IFS="${OLD_IFS}"

  if command -v dnf 2>&1 > /dev/null; then
    source /etc/os-release
    rhel_major_version=${VERSION_ID%%.*}

    ########
    # Base #
    ########

    if [ "${install_base}" = "1" ]; then
      install_cuda_rhel_repo "${rhel_major_version}" "${NVARCH}"
      # For libraries in the cuda-compat-* package: https://docs.nvidia.com/cuda/eula/index.html#attachment-a
      packages+=(cuda-cudart-${cuda_version[0]}-${cuda_version[1]}-${NV_CUDA_CUDART_VERSION}
                 cuda-compat-${cuda_version[0]}-${cuda_version[1]})

    fi

    ###########
    # Runtime # Should be accurate for 10-11, I didn't check every sub-version, but did some greps
    ###########

    if [ "${install_runtime}" = "1" ]; then
      packages+=(${NV_LIBCUBLAS_PACKAGE-} ${NV_LIBNCCL_PACKAGE-} ${NV_LIBNPP_PACKAGE-})


      packages+=(${NV_CUDA_LIB_VERSION+cuda-libraries-${cuda_version[0]}-${cuda_version[1]}-${NV_CUDA_LIB_VERSION}}
                 ${NV_NVTX_VERSION+cuda-nvtx-${cuda_version[0]}-${cuda_version[1]}-${NV_NVTX_VERSION}})

      if [ "${cuda_version[0]}" = "10" ]; then
        packages+=(${NV_LIBNPP_VERSION+cuda-npp-${cuda_version[0]}-${cuda_version[1]}-${NV_LIBNPP_VERSION}})
      fi

      if [ "${CUDA_VERSION}" = "10.0" ]; then
        packages+=(cuda-cublas-${cuda_version[0]}-${cuda_version[1]})
      fi

      #################
      # Runtime cudnn #
      #################

      if [ -n "${CUDNN_VERSION:+set}" ]; then
        packages+=(${NV_CUDNN_PACKAGE})
      fi
    fi

    #########
    # Devel #
    #########

    if [ "${install_devel}" = "1" ]; then
      packages+=(make)

      packages+=(cuda-command-line-tools-${cuda_version[0]}-${cuda_version[1]}-${NV_CUDA_LIB_VERSION}
                 cuda-minimal-build-${cuda_version[0]}-${cuda_version[1]}-${NV_CUDA_LIB_VERSION}
                 ${NV_NVPROF_VERSION+cuda-nvprof-${cuda_version[0]}-${cuda_version[1]}-${NV_NVPROF_VERSION}}
      )

      if [ "${cuda_version[0]}" = "10" ]; then
        packages+=(cuda-libraries-dev-${cuda_version[0]}-${cuda_version[1]}-${NV_CUDA_LIB_VERSION}
                   cuda-cudart-dev-${cuda_version[0]}-${cuda_version[1]}-${NV_CUDA_CUDART_DEV_VERSION}
                   ${NV_NVML_DEV_VERSION+cuda-nvml-dev-${cuda_version[0]}-${cuda_version[1]}-${NV_NVML_DEV_VERSION}}
                   ${NV_LIBNPP_DEV_VERSION+cuda-npp-dev-${cuda_version[0]}-${cuda_version[1]}-${NV_LIBNPP_DEV_VERSION}})
      else
        packages+=(cuda-libraries-devel-${cuda_version[0]}-${cuda_version[1]}-${NV_CUDA_LIB_VERSION}
                   cuda-cudart-devel-${cuda_version[0]}-${cuda_version[1]}-${NV_CUDA_CUDART_DEV_VERSION}
                   ${NV_NVML_DEV_VERSION+cuda-nvml-devel-${cuda_version[0]}-${cuda_version[1]}-${NV_NVML_DEV_VERSION}}
                   ${NV_LIBCUBLAS_DEV_VERSION+libcublas-devel-${cuda_version[0]}-${cuda_version[1]}-${NV_LIBCUBLAS_DEV_VERSION}})
      fi

      packages+=(${NV_LIBCUBLAS_DEV_PACKAGE-} ${NV_LIBNCCL_DEV_PACKAGE-} ${NV_LIBNPP_DEV_PACKAGE-} ${NV_NVPROF_DEV_PACKAGE-})

      ###############
      # Devel cudnn #
      ###############

      if [ -n "${CUDNN_VERSION:+set}" ]; then
        packages+=(${NV_CUDNN_PACKAGE} ${NV_CUDNN_PACKAGE_DEV})
      fi
    fi

    ###########
    # Install #
    ###########

    dnf install --setopt=obsoletes=0 -y "${packages[@]}"
    rm -rf /var/cache/yum/*
  elif command -v apt-get 2>&1 > /dev/null; then
    source /etc/os-release
    ubuntu_major_version=${VERSION_ID%%.*}

    ########
    # Base #
    ########

    if [ "${install_base}" = "1" ]; then
      packages+=(${NV_CUDA_COMPAT_PACKAGE-cuda-compat-${cuda_version[0]}-${cuda_version[1]}}
                 cuda-cudart-${cuda_version[0]}-${cuda_version[1]}=${NV_CUDA_CUDART_VERSION})
    fi

    ###########
    # Runtime #
    ###########

    if [ "${install_runtime}" = "1" ]; then
      if [ "${cuda_version[0]}" = "10" ]; then
        packages+=(${NV_LIBCUSPARSE_VERSION+cuda-cusparse-${cuda_version[0]}-${cuda_version[1]}=${NV_LIBCUSPARSE_VERSION}})
      else
        packages+=(${NV_LIBCUSPARSE_VERSION+libcusparse-${cuda_version[0]}-${cuda_version[1]}=${NV_LIBCUSPARSE_VERSION}})
      fi

      packages+=(${NV_CUDA_LIB_VERSION+cuda-libraries-${cuda_version[0]}-${cuda_version[1]}=${NV_CUDA_LIB_VERSION}}
                 ${NV_LIBNPP_PACKAGE-${NV_LIBNPP_VERSION+cuda-npp-10-0=${NV_LIBNPP_VERSION}}}
                 ${NV_NVTX_VERSION+cuda-nvtx-${cuda_version[0]}-${cuda_version[1]}=${NV_NVTX_VERSION}}
                 ${NV_LIBCUBLAS_PACKAGE-}
                 ${NV_LIBNCCL_PACKAGE-})

      #################
      # Runtime cudnn # Won't work on the older CUDAs
      #################

      if [ -n "${CUDNN_VERSION:+set}" ]; then
        packages+=(${NV_CUDNN_PACKAGE})
      fi
    fi

    #########
    # Devel # Not accurate for older CUDAs
    #########

    if [ "${install_devel}" = "1" ]; then
      packages+=(libtinfo5 libncursesw5
                 ${NV_CUDA_CUDART_DEV_VERSION+cuda-cudart-dev-${cuda_version[0]}-${cuda_version[1]}=${NV_CUDA_CUDART_DEV_VERSION}}
                 ${NV_CUDA_LIB_VERSION+cuda-command-line-tools-${cuda_version[0]}-${cuda_version[1]}=${NV_CUDA_LIB_VERSION}}
                 ${NV_CUDA_LIB_VERSION+cuda-minimal-build-${cuda_version[0]}-${cuda_version[1]}=${NV_CUDA_LIB_VERSION}}
                 ${NV_CUDA_LIB_VERSION+cuda-libraries-dev-${cuda_version[0]}-${cuda_version[1]}=${NV_CUDA_LIB_VERSION}}
                 ${NV_NVML_DEV_VERSION+cuda-nvml-dev-${cuda_version[0]}-${cuda_version[1]}=${NV_NVML_DEV_VERSION}}
                 ${NV_NVPROF_DEV_PACKAGE-}
                 ${NV_LIBNPP_DEV_PACKAGE-}
                 ${NV_LIBCUSPARSE_DEV_VERSION+libcusparse-dev-${cuda_version[0]}-${cuda_version[1]}=${NV_LIBCUSPARSE_DEV_VERSION}}
                 ${NV_LIBCUBLAS_DEV_PACKAGE-}
                 ${NV_LIBNCCL_DEV_PACKAGE-})

      ###############
      # Devel cudnn # Won't work on the older CUDAs
      ###############

      if [ -n "${CUDNN_VERSION:+set}" ]; then
        packages+=(${NV_CUDNN_PACKAGE} ${NV_CUDNN_PACKAGE_DEV})
      fi
    fi

    ###########
    # Install #
    ###########

    install_cuda_ubuntu_repo "${ubuntu_major_version}" "${NVARCH}"

    apt-get update

    DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends "${packages[@]}"
    apt-mark hold ${NV_LIBCUBLAS_DEV_PACKAGE_NAME-} ${NV_LIBNCCL_DEV_PACKAGE_NAME-} || :
    rm -rf /var/lib/apt/lists/*
  else
    echo "Unable to ID OS. Please amend this script" >&2
    exit 1
  fi

  if [ "${install_base}" = "1" ]; then
      # nvidia-docker 1.0
      echo "/usr/local/nvidia/lib" > /etc/ld.so.conf.d/nvidia.conf
      echo "/usr/local/nvidia/lib64" >> /etc/ld.so.conf.d/nvidia.conf
  fi
}

if [ -n "${BASH_SOURCE+set}" ]; then
  if [ "${BASH_SOURCE[0]}" = "${0}" ] || [ "$(basename "${BASH_SOURCE[0]}")" = "${0}" ]; then
    set -eu
    install_cuda_repo_and_packages ${@+"${@}"}

    # Self delete this file, so that cuda doesn't try to get installed twice
    exec bash -c "rm '${BASH_SOURCE[0]}'"
  fi
else
  case "${0}" in
    *30_install-cuda) # Ran
      set -eu
      install_cuda_repo_and_packages ${@+"${@}"}
      # sh doesn't have BASH_SOURCE or exact equivalent, this is closest, unless
      # they are sourcing it, then bad things can happen, hence the guards
      exec bash -c "rm '${0}'"
      ;;
    *sh) # Sourced for sh/dash/etc...
      ;;
  esac
fi