Skip to content

Instantly share code, notes, and snippets.

@190ikp
Last active Jun 24, 2022
Embed
What would you like to do?
#!/usr/bin/env bash
set -euxo pipefail
export NVIDIA_DRIVER_VERSION=450
export CUDA_VERSION=11-1
modify_bashrc() {
# don't variable expand in heredoc
cat << 'EOF' >> ~/.bashrc
export PATH=$PATH:/usr/local/cuda/bin
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64
EOF
}
setup_repos() {
# shellcheck disable=SC1091
source /etc/os-release
ARCH="$(arch)"
sudo apt update
sudo apt upgrade --yes
# install requirements
sudo apt install --yes \
"linux-headers-$(uname -r)" \
build-essential
export VERSION_ID_FIXED=${VERSION_ID//./}
# add CUDA and cuDNN repository
sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/$ID$VERSION_ID_FIXED/$ARCH/3bf863cc.pub
sudo apt-key adv --fetch-keys "http://developer.download.nvidia.com/compute/cuda/repos/$ID$VERSION_ID_FIXED/$ARCH/7fa2af80.pub"
echo "deb https://developer.download.nvidia.com/compute/cuda/repos/$ID$VERSION_ID_FIXED/$ARCH /" |
sudo tee /etc/apt/sources.list.d/cuda.list
echo "deb https://developer.download.nvidia.com/compute/machine-learning/repos/$ID$VERSION_ID_FIXED/$ARCH /" |
sudo tee /etc/apt/sources.list.d/nvidia-ml.list
# add libnvidia-container & nvidia-docker repository
curl -fsSL https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -fsSL https://nvidia.github.io/nvidia-docker/$ID$VERSION_ID/nvidia-docker.list |
sudo tee /etc/apt/sources.list.d/nvidia-container.list
sudo apt update
}
setup_driver() {
setup_repos
# NVIDIA driver only
sudo apt install --yes --no-install-recommends \
nvidia-driver-"$NVIDIA_DRIVER_VERSION" \
nvidia-container-toolkit
echo 'NVIDIA driver setup done.'
echo 'Reboot now.'
}
setup_cuda() {
# no GPU driver
sudo apt install --yes --no-install-recommends \
cuda-toolkit-"$CUDA_VERSION"
modify_bashrc
echo 'CUDA setup done.'
}
echo 'Setup CUDA...'
eval "$1"
@190ikp
Copy link
Author

190ikp commented Mar 2, 2022

NOTE: This script currently supports only for Ubuntu distro.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment