Skip to content

Instantly share code, notes, and snippets.

@ceramicwhite
Last active September 10, 2023 12:42
Show Gist options
  • Select an option

  • Save ceramicwhite/8f5435c4a655576d1b7ad5cbc14ee620 to your computer and use it in GitHub Desktop.

Select an option

Save ceramicwhite/8f5435c4a655576d1b7ad5cbc14ee620 to your computer and use it in GitHub Desktop.
lambda labs vm python diff version setup
#!/bin/bash
# This script is used to set up Python + CUDA-enabled PyTorch on a
# Lambda Labs VM.
PYTHON_VERSION=3.10
DEBIAN_FRONTEND=noninteractive
while getopts "v:" opt; do
case ${opt} in
v )
PYTHON_VERSION=$OPTARG
;;
* )
echo "Usage: $0 [-v version]"
exit 1
;;
esac
done
log() {
if [ -t 1 ]; then
echo -e "\n\e[1m$1\e[0m"
else
echo -e "\n$1"
fi
}
abort() {
log "An error occurred. Exiting..."
exit 1
}
trap 'abort' 0
log "Lambda Labs setup: Initiating installation of Python $PYTHON_VERSION and PyTorch with CUDA support..."
sudo timedatectl set-timezone UTC
sudo apt-get install software-properties-common -y
sudo add-apt-repository ppa:deadsnakes/ppa
sudo apt update && sudo apt install -y python$PYTHON_VERSION python${PYTHON_VERSION}-distutils python${PYTHON_VERSION}-dev python${PYTHON_VERSION}-venv python${PYTHON_VERSION}-tk
sudo update-alternatives --install /usr/bin/python python /usr/bin/python$PYTHON_VERSION 1
sudo update-alternatives --config python
log "Installing pip..."
wget https://bootstrap.pypa.io/get-pip.py
python get-pip.py && rm get-pip.py
log "Setting up Jupyter kernel for Python $PYTHON_VERSION..."
python -m pip install --upgrade --force-reinstall pyzmq ipykernel
python -m ipykernel install --user
log "Uninstalling system versions of PyTorch to install CUDA-enabled versions via Pip..."
sudo apt-get remove -y python3-torch-cuda python3-torchvision-cuda
log "Force reinstalling CUDA-enabled PyTorch packages via Pip..."
python -m pip install --upgrade --force-reinstall torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
for cudnn_so in /usr/lib/python3/dist-packages/tensorflow/libcudnn*; do
sudo ln -s "$cudnn_so" /usr/lib/x86_64-linux-gnu/
done
log "Setting path..."
export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH
trap : 0
log "Setup complete! Python $PYTHON_VERSION and CUDA-enabled PyTorch are now ready for use in the Jupyter environment."
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment