|
1 | | -# @todo - better / more robust parsing of inputs from env vars. |
| 1 | +#!/bin/bash |
| 2 | + |
2 | 3 | ## ------------------- |
3 | | -## Constants |
| 4 | +## Configuration |
4 | 5 | ## ------------------- |
5 | 6 |
|
6 | | -# @todo - apt repos/known supported versions? |
7 | | - |
8 | | -# @todo - GCC support matrix? |
9 | | - |
10 | | -# List of sub-packages to install. |
11 | | -# @todo - pass this in from outside the script? |
12 | | -# @todo - check the specified subpackages exist via apt pre-install? apt-rdepends cuda-9-0 | grep "^cuda-"? |
13 | | - |
14 | | -# Ideally choose from the list of meta-packages to minimise variance between cuda versions (although it does change too) |
15 | | -CUDA_PACKAGES_IN=( |
16 | | - "command-line-tools" |
17 | | - "libraries-dev" |
| 7 | +# CUDA sub-packages to install. |
| 8 | +declare -a CUDA_PACKAGES_IN=( |
| 9 | + "cuda-command-line-tools" |
| 10 | + "cuda-libraries-dev" |
| 11 | + "cuda-nvcc" |
18 | 12 | ) |
19 | 13 |
|
20 | 14 | ## ------------------- |
21 | | -## Bash functions |
| 15 | +## Helper Functions |
22 | 16 | ## ------------------- |
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 |
| 17 | + |
| 18 | +# Function to check if a command exists |
| 19 | +command_exists() { |
| 20 | + command -v "$1" >/dev/null 2>&1 |
42 | 21 | } |
43 | 22 |
|
44 | 23 | ## ------------------- |
45 | | -## Select CUDA version |
| 24 | +## Input Validation and Environment Setup |
46 | 25 | ## ------------------- |
47 | 26 |
|
48 | | -# Get the cuda version from the environment as $cuda. |
49 | | -CUDA_VERSION_MAJOR_MINOR=${cuda} |
| 27 | +# Get CUDA version from environment variable 'cuda' |
| 28 | +CUDA_VERSION_MAJOR_MINOR="${cuda}" |
50 | 29 |
|
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 |
| 30 | +# Validate CUDA version |
| 31 | +if [ -z "$CUDA_VERSION_MAJOR_MINOR" ]; then |
| 32 | + echo "Error: CUDA version not specified. Please set the 'cuda' environment variable (e.g., cuda=12.2)." |
| 33 | + exit 1 |
70 | 34 | fi |
71 | | -if [ -z "${CUDA_MINOR}" ] ; then |
72 | | - echo "Error: Unknown CUDA Minor version. Aborting." |
73 | | - exit 1 |
| 35 | + |
| 36 | +CUDA_MAJOR=$(echo "$CUDA_VERSION_MAJOR_MINOR" | cut -d. -f1) |
| 37 | +CUDA_MINOR=$(echo "$CUDA_VERSION_MAJOR_MINOR" | cut -d. -f2) |
| 38 | + |
| 39 | +# Check for root/sudo |
| 40 | +if ! command_exists sudo && [ "$EUID" -ne 0 ]; then |
| 41 | + echo "Error: This script requires root privileges. Please run with sudo." |
| 42 | + exit 1 |
74 | 43 | 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 |
| 44 | + |
| 45 | +SUDO_CMD="" |
| 46 | +if [ "$EUID" -ne 0 ]; then |
| 47 | + SUDO_CMD="sudo" |
79 | 48 | fi |
80 | 49 |
|
| 50 | +# Get Ubuntu version |
| 51 | +UBUNTU_VERSION=$(lsb_release -sr) |
81 | 52 |
|
82 | | -## --------------------------- |
83 | | -## GCC studio support check? |
84 | | -## --------------------------- |
| 53 | +# Validate Ubuntu version |
| 54 | +if [ -z "$UBUNTU_VERSION" ]; then |
| 55 | + echo "Error: Could not determine Ubuntu version." |
| 56 | + exit 1 |
| 57 | +fi |
85 | 58 |
|
86 | | -# @todo |
| 59 | +# Format Ubuntu version for URLs (e.g., 20.04 -> 2004) |
| 60 | +UBUNTU_VERSION_FORMATTED=$(echo "$UBUNTU_VERSION" | tr -d '.') |
87 | 61 |
|
88 | | -## ------------------------------- |
89 | | -## Select CUDA packages to install |
90 | | -## ------------------------------- |
91 | | -CUDA_PACKAGES="" |
92 | | -for package in "${CUDA_PACKAGES_IN[@]}" |
93 | | -do : |
94 | | - # @todo This is not perfect. Should probably provide a separate list for diff versions |
95 | | - # cuda-compiler-X-Y if CUDA >= 9.1 else cuda-nvcc-X-Y |
96 | | - if [[ "${package}" == "nvcc" ]] && version_ge "$CUDA_VERSION_MAJOR_MINOR" "9.1" ; then |
97 | | - package="compiler" |
98 | | - elif [[ "${package}" == "compiler" ]] && version_lt "$CUDA_VERSION_MAJOR_MINOR" "9.1" ; then |
99 | | - package="nvcc" |
100 | | - fi |
101 | | - # Build the full package name and append to the string. |
102 | | - CUDA_PACKAGES+=" cuda-${package}-${CUDA_MAJOR}-${CUDA_MINOR}" |
103 | | -done |
104 | | -echo "CUDA_PACKAGES ${CUDA_PACKAGES}" |
| 62 | +echo "CUDA Version: $CUDA_VERSION_MAJOR_MINOR" |
| 63 | +echo "Ubuntu Version: $UBUNTU_VERSION" |
| 64 | +echo "Ubuntu Version Formatted: $UBUNTU_VERSION_FORMATTED" |
105 | 65 |
|
106 | | -## ----------------- |
107 | | -## Prepare to install |
108 | | -## ----------------- |
| 66 | +## ------------------- |
| 67 | +## Install CUDA |
| 68 | +## ------------------- |
109 | 69 |
|
110 | | -PIN_FILENAME="cuda-ubuntu${UBUNTU_VERSION}.pin" |
111 | | -PIN_URL="https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/x86_64/${PIN_FILENAME}" |
112 | | -APT_KEY_URL="https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/x86_64/3bf863cc.pub" |
113 | | -REPO_URL="https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION}/x86_64/" |
| 70 | +# Download and install the CUDA keyring |
| 71 | +KEYRING_URL="https://developer.download.nvidia.com/compute/cuda/repos/ubuntu${UBUNTU_VERSION_FORMATTED}/x86_64/cuda-keyring_1.1-1_all.deb" |
| 72 | +echo "Downloading CUDA keyring from: $KEYRING_URL" |
| 73 | +wget -nv "$KEYRING_URL" -O cuda-keyring.deb |
| 74 | +$SUDO_CMD dpkg -i cuda-keyring.deb |
| 75 | +rm cuda-keyring.deb |
114 | 76 |
|
115 | | -echo "PIN_FILENAME ${PIN_FILENAME}" |
116 | | -echo "PIN_URL ${PIN_URL}" |
117 | | -echo "APT_KEY_URL ${APT_KEY_URL}" |
| 77 | +# Update package list |
| 78 | +echo "Updating package list..." |
| 79 | +$SUDO_CMD apt-get update |
118 | 80 |
|
119 | | -## ----------------- |
120 | | -## Check for root/sudo |
121 | | -## ----------------- |
| 81 | +# Construct the list of CUDA packages to install |
| 82 | +CUDA_PACKAGES="" |
| 83 | +for package in "${CUDA_PACKAGES_IN[@]}"; do |
| 84 | + CUDA_PACKAGES+=" ${package}-${CUDA_MAJOR}-${CUDA_MINOR}" |
| 85 | +done |
122 | 86 |
|
123 | | -# Detect if the script is being run as root, storing true/false in is_root. |
124 | | -is_root=false |
125 | | -if (( $EUID == 0)); then |
126 | | - is_root=true |
127 | | -fi |
128 | | -# Find if sudo is available |
129 | | -has_sudo=false |
130 | | -if command -v sudo &> /dev/null ; then |
131 | | - has_sudo=true |
132 | | -fi |
133 | | -# Decide if we can proceed or not (root or sudo is required) and if so store whether sudo should be used or not. |
134 | | -if [ "$is_root" = false ] && [ "$has_sudo" = false ]; then |
135 | | - echo "Root or sudo is required. Aborting." |
136 | | - exit 1 |
137 | | -elif [ "$is_root" = false ] ; then |
138 | | - USE_SUDO=sudo |
139 | | -else |
140 | | - USE_SUDO= |
| 87 | +# Special handling for nvcc in older versions |
| 88 | +if (( $(echo "$CUDA_MAJOR < 9" | bc -l) )); then |
| 89 | + CUDA_PACKAGES+=" cuda-nvcc-${CUDA_MAJOR}-${CUDA_MINOR}" |
141 | 90 | fi |
142 | 91 |
|
143 | | -## ----------------- |
144 | | -## Install |
145 | | -## ----------------- |
146 | | -echo "Adding CUDA Repository" |
147 | | -wget ${PIN_URL} |
148 | | -$USE_SUDO mv ${PIN_FILENAME} /etc/apt/preferences.d/cuda-repository-pin-600 |
149 | | -$USE_SUDO apt-key adv --fetch-keys ${APT_KEY_URL} |
150 | | -$USE_SUDO add-apt-repository "deb ${REPO_URL} /" |
151 | | -$USE_SUDO apt-get update |
152 | | - |
153 | | -echo "Installing CUDA packages ${CUDA_PACKAGES}" |
154 | | -$USE_SUDO apt-get -y install ${CUDA_PACKAGES} |
155 | | - |
156 | | -if [[ $? -ne 0 ]]; then |
157 | | - echo "CUDA Installation Error." |
158 | | - exit 1 |
| 92 | +# Install CUDA packages |
| 93 | +echo "Installing CUDA packages: $CUDA_PACKAGES" |
| 94 | +$SUDO_CMD apt-get install -y --no-install-recommends $CUDA_PACKAGES |
| 95 | + |
| 96 | +if [ $? -ne 0 ]; then |
| 97 | + echo "Error: Failed to install CUDA packages." |
| 98 | + exit 1 |
159 | 99 | fi |
160 | | -## ----------------- |
161 | | -## Set environment vars / vars to be propagated |
162 | | -## ----------------- |
163 | 100 |
|
164 | | -CUDA_PATH=/usr/local/cuda-${CUDA_MAJOR}.${CUDA_MINOR} |
165 | | -echo "CUDA_PATH=${CUDA_PATH}" |
166 | | -export CUDA_PATH=${CUDA_PATH} |
| 101 | +## ------------------- |
| 102 | +## Environment Variables |
| 103 | +## ------------------- |
167 | 104 |
|
| 105 | +CUDA_PATH="/usr/local/cuda-${CUDA_MAJOR}.${CUDA_MINOR}" |
| 106 | +echo "CUDA_PATH=$CUDA_PATH" |
168 | 107 |
|
169 | | -# Quick test. @temp |
| 108 | +# Update environment variables for the current shell |
| 109 | +export CUDA_PATH |
170 | 110 | export PATH="$CUDA_PATH/bin:$PATH" |
171 | | -export LD_LIBRARY_PATH="$CUDA_PATH/lib:$LD_LIBRARY_PATH" |
172 | | -nvcc -V |
173 | | - |
174 | | -# If executed on github actions, make the appropriate echo statements to update the environment |
175 | | -if [[ $GITHUB_ACTIONS ]]; then |
176 | | - # Set paths for subsequent steps, using ${CUDA_PATH} |
177 | | - echo "Adding CUDA to CUDA_PATH, PATH and LD_LIBRARY_PATH" |
178 | | - echo "CUDA_PATH=${CUDA_PATH}" >> $GITHUB_ENV |
179 | | - echo "${CUDA_PATH}/bin" >> $GITHUB_PATH |
180 | | - echo "LD_LIBRARY_PATH=${CUDA_PATH}/lib:${LD_LIBRARY_PATH}" >> $GITHUB_ENV |
| 111 | +export LD_LIBRARY_PATH="$CUDA_PATH/lib64:$LD_LIBRARY_PATH" |
| 112 | + |
| 113 | +# Verify installation (optional) |
| 114 | +echo "Verifying installation..." |
| 115 | +if command_exists nvcc; then |
| 116 | + nvcc --version |
| 117 | +else |
| 118 | + echo "Warning: nvcc not found. Installation might be incomplete." |
181 | 119 | fi |
| 120 | + |
| 121 | +# Update environment variables for GitHub Actions (if applicable) |
| 122 | +if [ -n "$GITHUB_ACTIONS" ]; then |
| 123 | + echo "Setting environment variables for GitHub Actions..." |
| 124 | + echo "CUDA_PATH=$CUDA_PATH" >> "$GITHUB_ENV" |
| 125 | + echo "$CUDA_PATH/bin" >> "$GITHUB_PATH" |
| 126 | + echo "LD_LIBRARY_PATH=$CUDA_PATH/lib64:$LD_LIBRARY_PATH" >> "$GITHUB_ENV" |
| 127 | +fi |
| 128 | + |
| 129 | +echo "CUDA installation complete." |
0 commit comments