|
| 1 | +# @todo - better / more robust parsing of inputs from env vars. |
| 2 | +## ------------------- |
| 3 | +## Constants |
| 4 | +## ------------------- |
| 5 | + |
| 6 | +# List of sub-packages to install. |
| 7 | +# @todo - pass this in from outside the script? |
| 8 | +# @todo - check the specified subpackages exist via apt pre-install? apt-rdepends cuda-9-0 | grep "^cuda-"? |
| 9 | + |
| 10 | +# Ideally choose from the list of meta-packages to minimise variance between cuda versions (although it does change too). Some of these packages may not be availble in older CUDA releases |
| 11 | +CUDA_PACKAGES_IN=( |
| 12 | + "cuda-compiler" |
| 13 | + "cuda-cudart-dev" |
| 14 | + "cuda-nvtx" |
| 15 | + "cuda-nvrtc-dev" |
| 16 | + "libcurand-dev" # 11-0+ |
| 17 | + "cuda-cccl" # 11.4+, provides cub and thrust. On 11.3 knwon as cuda-thrust-11-3 |
| 18 | +) |
| 19 | + |
| 20 | +## ------------------- |
| 21 | +## Bash functions |
| 22 | +## ------------------- |
| 23 | +# returns 0 (true) if a >= b |
| 24 | +function version_ge() { |
| 25 | + [ "$#" != "2" ] && echo "${FUNCNAME[0]} requires exactly 2 arguments." && exit 1 |
| 26 | + [ "$(printf '%s\n' "$@" | sort -V | head -n 1)" == "$2" ] |
| 27 | +} |
| 28 | +# returns 0 (true) if a > b |
| 29 | +function version_gt() { |
| 30 | + [ "$#" != "2" ] && echo "${FUNCNAME[0]} requires exactly 2 arguments." && exit 1 |
| 31 | + [ "$1" = "$2" ] && return 1 || version_ge $1 $2 |
| 32 | +} |
| 33 | +# returns 0 (true) if a <= b |
| 34 | +function version_le() { |
| 35 | + [ "$#" != "2" ] && echo "${FUNCNAME[0]} requires exactly 2 arguments." && exit 1 |
| 36 | + [ "$(printf '%s\n' "$@" | sort -V | head -n 1)" == "$1" ] |
| 37 | +} |
| 38 | +# returns 0 (true) if a < b |
| 39 | +function version_lt() { |
| 40 | + [ "$#" != "2" ] && echo "${FUNCNAME[0]} requires exactly 2 arguments." && exit 1 |
| 41 | + [ "$1" = "$2" ] && return 1 || version_le $1 $2 |
| 42 | +} |
| 43 | + |
| 44 | +## ------------------- |
| 45 | +## Select CUDA version |
| 46 | +## ------------------- |
| 47 | + |
| 48 | +# Get the cuda version from the environment as $cuda. |
| 49 | +CUDA_VERSION_MAJOR_MINOR=${cuda} |
| 50 | + |
| 51 | +# Split the version. |
| 52 | +# We (might/probably) don't know PATCH at this point - it depends which version gets installed. |
| 53 | +CUDA_MAJOR=$(echo "${CUDA_VERSION_MAJOR_MINOR}" | cut -d. -f1) |
| 54 | +CUDA_MINOR=$(echo "${CUDA_VERSION_MAJOR_MINOR}" | cut -d. -f2) |
| 55 | +CUDA_PATCH=$(echo "${CUDA_VERSION_MAJOR_MINOR}" | cut -d. -f3) |
| 56 | +# use lsb_release to find the OS. |
| 57 | +UBUNTU_VERSION=$(lsb_release -sr) |
| 58 | +UBUNTU_VERSION="${UBUNTU_VERSION//.}" |
| 59 | + |
| 60 | +echo "CUDA_MAJOR: ${CUDA_MAJOR}" |
| 61 | +echo "CUDA_MINOR: ${CUDA_MINOR}" |
| 62 | +echo "CUDA_PATCH: ${CUDA_PATCH}" |
| 63 | +# echo "UBUNTU_NAME: ${UBUNTU_NAME}" |
| 64 | +echo "UBUNTU_VERSION: ${UBUNTU_VERSION}" |
| 65 | + |
| 66 | +# If we don't know the CUDA_MAJOR or MINOR, error. |
| 67 | +if [ -z "${CUDA_MAJOR}" ] ; then |
| 68 | + echo "Error: Unknown CUDA Major version. Aborting." |
| 69 | + exit 1 |
| 70 | +fi |
| 71 | +if [ -z "${CUDA_MINOR}" ] ; then |
| 72 | + echo "Error: Unknown CUDA Minor version. Aborting." |
| 73 | + exit 1 |
| 74 | +fi |
| 75 | +# If we don't know the Ubuntu version, error. |
| 76 | +if [ -z ${UBUNTU_VERSION} ]; then |
| 77 | + echo "Error: Unknown Ubuntu version. Aborting." |
| 78 | + exit 1 |
| 79 | +fi |
| 80 | + |
| 81 | + |
| 82 | +## ------------------------------- |
| 83 | +## Select CUDA packages to install |
| 84 | +## ------------------------------- |
| 85 | +CUDA_PACKAGES="" |
| 86 | +for package in "${CUDA_PACKAGES_IN[@]}" |
| 87 | +do : |
| 88 | + # @todo This is not perfect. Should probably provide a separate list for diff versions |
| 89 | + # cuda-compiler-X-Y if CUDA >= 9.1 else cuda-nvcc-X-Y |
| 90 | + if [[ "${package}" == "cuda-nvcc" ]] && version_ge "$CUDA_VERSION_MAJOR_MINOR" "9.1" ; then |
| 91 | + package="cuda-compiler" |
| 92 | + elif [[ "${package}" == "cuda-compiler" ]] && version_lt "$CUDA_VERSION_MAJOR_MINOR" "9.1" ; then |
| 93 | + package="cuda-nvcc" |
| 94 | + # CUB/Thrust are packages in cuda-thrust in 11.3, but cuda-cccl in 11.4+ |
| 95 | + elif [[ "${package}" == "cuda-thrust" || "${package}" == "cuda-cccl" ]]; then |
| 96 | + # CUDA cuda-thrust >= 11.4 |
| 97 | + if version_ge "$CUDA_VERSION_MAJOR_MINOR" "11.4" ; then |
| 98 | + package="cuda-cccl" |
| 99 | + # Use cuda-thrust > 11.2 |
| 100 | + elif version_ge "$CUDA_VERSION_MAJOR_MINOR" "11.3" ; then |
| 101 | + package="cuda-thrust" |
| 102 | + # Do not include this pacakge < 11.3 |
| 103 | + else |
| 104 | + continue |
| 105 | + fi |
| 106 | + fi |
| 107 | + # CUDA 11+ includes lib* / lib*-dev packages, which if they existed previously where cuda-cu*- / cuda-cu*-dev- |
| 108 | + if [[ ${package} == libcu* ]] && version_lt "$CUDA_VERSION_MAJOR_MINOR" "11.0" ; then |
| 109 | + package="${package/libcu/cuda-cu}" |
| 110 | + fi |
| 111 | + # Build the full package name and append to the string. |
| 112 | + CUDA_PACKAGES+=" ${package}-${CUDA_MAJOR}-${CUDA_MINOR}" |
| 113 | +done |
| 114 | +echo "CUDA_PACKAGES ${CUDA_PACKAGES}" |
| 115 | + |
| 116 | +## ----------------- |
| 117 | +## Prepare to install |
| 118 | +## ----------------- |
| 119 | +CPU_ARCH="x86_64" |
| 120 | +PIN_FILENAME="cuda-ubuntu${UBUNTU_VERSION}.pin" |
| 121 | +PIN_URL="https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/${CPU_ARCH}/${PIN_FILENAME}" |
| 122 | +# apt keyring package now available https://developer.nvidia.com/blog/updating-the-cuda-linux-gpg-repository-key/ |
| 123 | +KERYRING_PACKAGE_FILENAME="cuda-keyring_1.0-1_all.deb" |
| 124 | +KEYRING_PACKAGE_URL="https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/${CPU_ARCH}/${KERYRING_PACKAGE_FILENAME}" |
| 125 | +REPO_URL="https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/${CPU_ARCH}/" |
| 126 | + |
| 127 | +echo "PIN_FILENAME ${PIN_FILENAME}" |
| 128 | +echo "PIN_URL ${PIN_URL}" |
| 129 | +echo "KEYRING_PACKAGE_URL ${KEYRING_PACKAGE_URL}" |
| 130 | +echo "APT_KEY_URL ${APT_KEY_URL}" |
| 131 | + |
| 132 | +## ----------------- |
| 133 | +## Check for root/sudo |
| 134 | +## ----------------- |
| 135 | + |
| 136 | +# Detect if the script is being run as root, storing true/false in is_root. |
| 137 | +is_root=false |
| 138 | +if (( $EUID == 0)); then |
| 139 | + is_root=true |
| 140 | +fi |
| 141 | +# Find if sudo is available |
| 142 | +has_sudo=false |
| 143 | +if command -v sudo &> /dev/null ; then |
| 144 | + has_sudo=true |
| 145 | +fi |
| 146 | +# Decide if we can proceed or not (root or sudo is required) and if so store whether sudo should be used or not. |
| 147 | +if [ "$is_root" = false ] && [ "$has_sudo" = false ]; then |
| 148 | + echo "Root or sudo is required. Aborting." |
| 149 | + exit 1 |
| 150 | +elif [ "$is_root" = false ] ; then |
| 151 | + USE_SUDO=sudo |
| 152 | +else |
| 153 | + USE_SUDO= |
| 154 | +fi |
| 155 | + |
| 156 | +## ----------------- |
| 157 | +## Install |
| 158 | +## ----------------- |
| 159 | +echo "Adding CUDA Repository" |
| 160 | +wget ${PIN_URL} |
| 161 | +$USE_SUDO mv ${PIN_FILENAME} /etc/apt/preferences.d/cuda-repository-pin-600 |
| 162 | +wget ${KEYRING_PACKAGE_URL} && ${USE_SUDO} dpkg -i ${KERYRING_PACKAGE_FILENAME} && rm ${KERYRING_PACKAGE_FILENAME} |
| 163 | +$USE_SUDO add-apt-repository "deb ${REPO_URL} /" |
| 164 | +$USE_SUDO apt-get update |
| 165 | + |
| 166 | +echo "Installing CUDA packages ${CUDA_PACKAGES}" |
| 167 | +$USE_SUDO apt-get -y install ${CUDA_PACKAGES} |
| 168 | + |
| 169 | +if [[ $? -ne 0 ]]; then |
| 170 | + echo "CUDA Installation Error." |
| 171 | + exit 1 |
| 172 | +fi |
| 173 | + |
| 174 | +## ----------------- |
| 175 | +## Set environment vars / vars to be propagated |
| 176 | +## ----------------- |
| 177 | + |
| 178 | +CUDA_PATH=/usr/local/cuda-${CUDA_MAJOR}.${CUDA_MINOR} |
| 179 | +echo "CUDA_PATH=${CUDA_PATH}" |
| 180 | +export CUDA_PATH=${CUDA_PATH} |
| 181 | +export PATH="$CUDA_PATH/bin:$PATH" |
| 182 | +export LD_LIBRARY_PATH="$CUDA_PATH/lib:$LD_LIBRARY_PATH" |
| 183 | +export LD_LIBRARY_PATH="$CUDA_PATH/lib64:$LD_LIBRARY_PATH" |
| 184 | +# Check nvcc is now available. |
| 185 | +nvcc -V |
| 186 | + |
| 187 | + |
| 188 | +# If executed on github actions, make the appropriate echo statements to update the environment |
| 189 | +if [[ $GITHUB_ACTIONS ]]; then |
| 190 | + # Set paths for subsequent steps, using ${CUDA_PATH} |
| 191 | + echo "Adding CUDA to CUDA_PATH, PATH and LD_LIBRARY_PATH" |
| 192 | + echo "CUDA_PATH=${CUDA_PATH}" >> $GITHUB_ENV |
| 193 | + echo "${CUDA_PATH}/bin" >> $GITHUB_PATH |
| 194 | + echo "LD_LIBRARY_PATH=${CUDA_PATH}/lib:${LD_LIBRARY_PATH}" >> $GITHUB_ENV |
| 195 | + echo "LD_LIBRARY_PATH=${CUDA_PATH}/lib64:${LD_LIBRARY_PATH}" >> $GITHUB_ENV |
| 196 | +fi |
0 commit comments