Setup script for Tinygrad with Nvidia GPU on Ubuntu native/wsl2

tinygrad_setup.sh

#!/bin/bash

# Setup script for Tinygrad, including all supported backends, made it for WSL2 Ubuntu 22.04 running on Windows 10/11 with an up-to-date Nvidia GPU driver (it should work for native Ubuntu too)

# To use this script, save it in a file: nano setup_environment.sh
# give it execute permissions:  chmod +x setup_environment.sh
# and then run it:  ./setup_environment.sh
# Make sure to run it in a terminal where you have administrative (sudo) access, as some steps require it.

# Initialize variables
update_command="sudo apt-get update && sudo apt-get dist-upgrade"
install_command="wget -nv -O- https://lambdalabs.com/install-lambda-stack.sh | I_AGREE_TO_THE_CUDNN_LICENSE=1 sh -"

# Check if running on WSL2
kernel_release=$(uname -r)
if [[ $kernel_release == *"-microsoft-standard"* ]]; then
  echo "You are running WSL2. Note: Install NVIDIA Drivers, CUDA, and cuDNN on your Windows host, not within WSL2."
else
  echo "You are running native Linux."
fi

# Step 1: Check if CUDA is installed and working
echo "Step 1: Verifying if CUDA is installed and working..."
nvidia-smi

# Capture the exit status of the last command
status=$?

# User decision variables
update_choice=""
install_choice=""

if [ $status -eq 0 ]; then
  echo "CUDA is installed and working."
  echo "Step 2: Asking for update choice..."
  read -p "Would you like to check for updates for NVIDIA/PyTorch/Tinygrad/Keras? (y/n): " update_choice
  if [ "$update_choice" == "y" ]; then
    echo "Step 3: Performing update..."
    eval $update_command
  fi
else
  echo "CUDA is not installed or not working."
  echo "Step 2: Asking for install choice..."
  read -p "Would you like to install the NVIDIA stack, including NVIDIA, PyTorch, Tinygrad, Keras, CUDA, cuDNN, and the NVIDIA driver? (y/n): " install_choice
  if [[ "$install_choice" == "y" && $kernel_release == *"-microsoft-standard"* ]]; then
    echo "Please install NVIDIA Drivers, CUDA, and cuDNN on your Windows host instead. Exiting."
    exit 1
  elif [ "$install_choice" == "y" ]; then
    echo "Step 3: Performing installation..."
    eval $install_command
  else
    echo "Exiting."
    exit 1
  fi
fi

if [ "$install_choice" == "y" ] || [ $status -eq 0 ]; then
  echo "Step 4: Installing Python Virtual Environment..."
  sudo apt install -y python3-venv
  python3 -m venv tiny
  source tiny/bin/activate
  
  echo "Step 5: Installing tinygrad..."
  sudo apt-get install -y python3-wheel
  sudo apt install -y git
  git clone https://github.com/tinygrad/tinygrad.git
  cd tinygrad
  python3 -m pip install -e .
  cd ..

  echo "Step 6: Installing compilers..."
  sudo apt install -y build-essential binutils gdb
  sudo apt-get install -y manpages-dev
  
  echo "Step 7: Installing Clang backend..."
  sudo apt install -y clang
  
  echo "Step 8: Installing LLVM backend..."
  pip install llvmlite
  
  echo "Step 9: Installing OpenCL and its dependencies..."
  pip install pyopencl
  sudo apt install -y pocl-opencl-icd
  
  echo "Step 10: Installing Triton and PyTorch backend..."
  pip install triton
  
  echo "Step 11: Installing pycuda dependencies..."
  sudo apt-get install -y freeglut3 freeglut3-dev libxi-dev libxmu-dev
  sudo apt install -y nvidia-cuda-toolkit gedit
  
  echo "Step 12: Configuring PATH and LD_LIBRARY_PATH for CUDA..."
  echo "" >> ~/.bashrc
  echo 'export PATH="/usr/local/cuda/bin:${PATH}"' >> ~/.bashrc
  echo "" >> ~/.bashrc
  echo 'export LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"' >> ~/.bashrc
  
  echo "Step 13: Testing CUDA Compiler..."
  nvcc --version
  
  echo "Step 14: Installing pycuda backend..."
  sudo apt-get install -y build-essential python-setuptools libboost-python-dev libboost-thread-dev
  pip install pycuda
  pip install nevergrad
  
  echo "Step 15: Installing WEBGPU backend..."
  pip install wgpu
  
  echo "Step 16: Installing EXTRA..."
  pip install pytest
  pip install graphviz
  pip install pydot
  pip install onnx
  pip install nevergrad
  pip install jupyterlab
  sudo apt install graphviz
  echo "All steps completed successfully!"
fi

# --Condensate script--
#!/bin/bash
# test nvidia
# nvidia-smi
# if [ $? -eq 0 ]; then sudo apt-get update && sudo apt-get dist-upgrade; fi 
# if [ $? -ne 0 ]; then wget -nv -O- https://lambdalabs.com/install-lambda-stack.sh | I_AGREE_TO_THE_CUDNN_LICENSE=1 sh -; fi
# sudo apt install -y python3-venv
# python3 -m venv tiny
# source tiny/bin/activate
# sudo apt-get install -y python3-wheel
# sudo apt install -y git
# git clone https://github.com/tinygrad/tinygrad.git
# cd tinygrad
# python3 -m pip install -e .
# cd ..
# sudo apt install -y build-essential binutils gdb
# sudo apt-get install -y manpages-dev
# pip install jupyterlab
# sudo apt install -y clang
# pip install llvmlite
# pip install pyopencl
# sudo apt install -y pocl-opencl-icd
# pip install triton
# sudo apt-get install -y freeglut3 freeglut3-dev libxi-dev libxmu-dev
# sudo apt install -y nvidia-cuda-toolkit gedit
# echo 'export PATH="/usr/local/cuda/bin:${PATH}"' >> ~/.bashrc
# echo 'export LD_LIBRARY_PATH="/usr/local/cuda/lib64:${LD_LIBRARY_PATH}"' >> ~/.bashrc
# sudo apt-get install -y build-essential python-setuptools libboost-python-dev libboost-thread-dev
# pip install pycuda
# pip install nevergrad
# pip install wgpu
# pip install pytest
#  pip install graphviz
# pip install pydot
# pip install onnx
# pip install nevergrad
# pip install jupyterlab
# sudo apt install graphviz