setup_gpu.sh

#!/bin/bash
# ─── STP GPU Environment Setup ──────────────────────────────────────────────
#
# Sets up the STP environment on an NVIDIA GPU node (e.g., Berzelius A100).
# Creates two virtualenvs:
#   1. ~/venv_vllm  — vLLM + Ray for inference (CUDA)
#   2. ~/venv310    — Levanter + JAX for training (CUDA)
# Also installs Lean 4 + Mathlib4 if not already present.
#
# Prerequisites:
#   - Python 3.10+
#   - CUDA 12.x drivers installed
#   - cmake and a C++17 compiler
#
# Usage:
#   bash setup_gpu.sh
# ─────────────────────────────────────────────────────────────────────────────

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

echo "═══════════════════════════════════════════════"
echo "  STP GPU Environment Setup"
echo "  CUDA: $(nvidia-smi --query-gpu=driver_version --format=csv,noheader 2>/dev/null | head -1 || echo 'not detected')"
echo "  Python: $(python3 --version 2>/dev/null || echo 'not found')"
echo "═══════════════════════════════════════════════"

# ─── 1. vLLM inference environment ──────────────────────────────────────────
echo ""
echo "[1/3] Setting up vLLM inference environment (~venv_vllm)..."

if [ ! -d ~/venv_vllm ]; then
    python3 -m venv ~/venv_vllm
fi
source ~/venv_vllm/bin/activate

pip install --upgrade pip setuptools wheel

# PyTorch with CUDA 12.x
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

# vLLM (CUDA target — the default)
pip install vllm

# Ray for distributed inference
pip install "ray[default]>=2.34"

# Other STP dependencies
pip install transformers huggingface_hub sentencepiece
pip install numpy pgzip psutil func_timeout pyyaml wandb
pip install pydantic openai

deactivate

# ─── 2. Levanter training environment ───────────────────────────────────────
echo ""
echo "[2/3] Setting up Levanter training environment (~venv310)..."

if [ ! -d ~/venv310 ]; then
    python3 -m venv ~/venv310
fi
source ~/venv310/bin/activate

pip install --upgrade pip setuptools wheel

# JAX with CUDA 12 support
pip install "jax[cuda12]"

# PyTorch CPU (Levanter uses it for tokenizer/model loading only)
pip install torch --index-url https://download.pytorch.org/whl/cpu

# Install Levanter from the STP fork
if [ -d "$SCRIPT_DIR/levanter" ]; then
    cd "$SCRIPT_DIR/levanter"
    pip install -e ".[dev]"
    cd "$SCRIPT_DIR"
fi

# Additional training deps
pip install wandb huggingface_hub transformers sentencepiece
pip install "ray[default]>=2.34"

deactivate

# ─── 3. Lean 4 + Mathlib4 ───────────────────────────────────────────────────
echo ""
echo "[3/3] Setting up Lean 4 + Mathlib4..."

if ! command -v elan &> /dev/null; then
    echo "Installing elan (Lean version manager)..."
    curl -sSf https://raw.githubusercontent.com/leanprover/elan/master/elan-init.sh | sh -s -- -y --default-toolchain none
    export PATH="$HOME/.elan/bin:$PATH"
fi

# Ensure the right version is installed
elan toolchain install leanprover/lean4:v4.9.0-rc1 2>/dev/null || true
elan default leanprover/lean4:v4.9.0-rc1

LEAN_DIR="$HOME/lean"
if [ ! -d "$LEAN_DIR/mathlib4" ]; then
    echo "Cloning Mathlib4..."
    mkdir -p "$LEAN_DIR"
    cd "$LEAN_DIR"
    git clone --depth 1 https://github.com/leanprover-community/mathlib4.git
    cd mathlib4
    echo "Building Mathlib4 (this may take a while)..."
    lake build
    cd "$SCRIPT_DIR"
else
    echo "Mathlib4 already present at $LEAN_DIR/mathlib4"
fi

# ─── Done ────────────────────────────────────────────────────────────────────
echo ""
echo "═══════════════════════════════════════════════"
echo "  Setup complete!"
echo ""
echo "  Inference: source ~/venv_vllm/bin/activate"
echo "  Training:  source ~/venv310/bin/activate"
echo "  Lean:      ~/.elan/bin/lean --version"
echo ""
echo "  Next steps:"
echo "    1. Copy RL/.bash_alias_gpu.sh.template to RL/.bash_alias.sh"
echo "       and fill in STORAGE, WANDB_API_KEY, HUGGING_FACE_HUB_TOKEN"
echo "    2. Prepare datasets:  cd RL && python prepare_datasets.py"
echo "    3. Run SFT:           bash run_SFT_gpu.sh"
echo "    4. Run RL:            bash run_RL_steps_gpu.sh"
echo "═══════════════════════════════════════════════"