ubuntu 20 / lambdalabs pyenv+poetry+cuda12 bring-up

setup.sh

#!/usr/bin/env bash

## PURPOSE: Fix problems with the currently outdated lambdalabs default image.

# Does:
#  - Installs latest Python 3.10 using pyenv
#  - Installs the current poetry environment (assuming this is running in a poetry project dir)
#  - Installs latest pytorch CUDA 12 nightly into current poetry environment

# Script for "normal-maxing" the default lambdalabs image as of June 2023.
# Their default image is providing Ubuntu 20, which is somewhat outdated with
# Ubuntu 22 being the new LTS.
# This script is for running in a python poetry project dir (with the "poetry install")
# and will give you a working pytorch CUDA 12 environment. 

set -e                                 
set -x                                                                              

# NOTE: avoid a big `apt upgrade` because it takes 10+ minutes for not much benefit.
                                                                                                                                 
# install packages we need for other things                                                                                      
sudo apt update                                                                                                                  
sudo apt remove man-db -y                                                                                                       
sudo apt install zlib1g zlib1g-dev libssl-dev libbz2-dev libsqlite3-dev libcap-dev libffi-dev fd-find jq -y                    
sudo snap install ripgrep --classic                                                                                            
# disable weird things we don't care about                                                                                     
sudo service cloudflared stop &                                                                                                
sudo service lambda-jupyter stop &                                                                                             
sudo service docker stop &                                                                                                     
sudo service containerd stop &                                                                     
sudo service unattended-upgrades stop &                                                            
# setup pyenv (if it doesn't already exist)                                                  
if [[ ! -d ~/.pyenv ]]; then                                                   
    curl https://pyenv.run | bash                                              
    #                                                                          
    echo 'export PYENV_ROOT="$HOME/.pyenv"' >> ~/.bashrc                       
    echo 'command -v pyenv >/dev/null || export PATH="$PYENV_ROOT/bin:$PATH"' >> ~/.bashrc
    echo 'eval "$(pyenv init -)"' >> ~/.bashrc                                 
    echo 'alias po=poetry' >> ~/.bashrc                                                   
    #                                                         
    echo 'export PYENV_ROOT="$HOME/.pyenv"' >> ~/.zshrc       
    echo 'command -v pyenv >/dev/null || export PATH="$PYENV_ROOT/bin:$PATH"' >> ~/.zshrc       
    echo 'eval "$(pyenv init -)"' >> ~/.zshrc
    echo 'alias po=poetry' >> ~/.zshrc                                                          
    # start again with the shell changes...                                                                             
    exec $0                                                                                                             
fi                                                                                                                      
# gather version metadata for all python versions                                                                       
pyenv update                                                                                                            
# install latest python 3.10 as fast as possible                                                                        
PYVER=$(pyenv install -l |grep '\s\+3.10' |tail -1)
MAKE_OPTS="-j $(getconf _NPROCESSORS_ONLN)" time pyenv install $PYVER                                                   
pyenv global $PYVER                                                                                                     
#                                                                                                                       
pip install pip poetry setuptools wheel -U                                                                              
# install local project                                                                                                 
poetry install                                                                                                          
# system image is CUDA 12, but pytorch doesn't support CUDA 12 without a 2 GB nightly build...                          
# https://github.com/pytorch/pytorch/issues/91122                                   
# TODO: bake this into the poetry config directly                                                                       
poetry run pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121   
# now this should work: po run python -m torch.utils.collect_env                                                        
# also we need a missing kernel module for some reason                                                    
sudo modprobe nvidia_uvm