setup.sh

#!/usr/bin/env bash
# ─────────────────────────────────────────────────────────────
#  GPU Onboarding Tool — First-time setup
#  Requires: uv  (install with: curl -LsSf https://astral.sh/uv/install.sh | sh)
# ─────────────────────────────────────────────────────────────
set -e

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

# ── 1. Check uv ─────────────────────────────────────────────
if ! command -v uv &>/dev/null; then
    echo "[ERROR] 'uv' is not installed."
    echo "        Install it with:"
    echo "          curl -LsSf https://astral.sh/uv/install.sh | sh"
    echo "        Then re-run this script."
    exit 1
fi

echo "[INFO] uv found: $(uv --version)"

# ── 2. Check NVIDIA drivers ──────────────────────────────────
if ! command -v nvidia-smi &>/dev/null; then
    echo "[ERROR] nvidia-smi not found. Install NVIDIA drivers first."
    echo "        Ubuntu:  sudo apt install nvidia-driver-<version>"
    exit 1
fi
echo "[INFO] nvidia-smi found. Driver: $(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -1)"

# ── 3. Create virtual environment ───────────────────────────
echo "[INFO] Creating uv virtual environment (.venv) …"
uv venv .venv

# ── 4. Detect CUDA version from driver ──────────────────────
# nvidia-smi header line contains e.g. "CUDA Version: 12.2"
CUDA_VER=$(nvidia-smi | grep -oP "CUDA Version: \K[0-9]+\.[0-9]+" || echo "")
if [[ -z "$CUDA_VER" ]]; then
    echo "[WARN] Could not detect CUDA version — defaulting to cu121 wheel"
    TORCH_INDEX="https://download.pytorch.org/whl/cu121"
else
    CUDA_MAJOR=$(echo "$CUDA_VER" | cut -d. -f1)
    CUDA_MINOR=$(echo "$CUDA_VER" | cut -d. -f2)
    echo "[INFO] Detected CUDA $CUDA_VER"

    # Map CUDA version to the closest available PyTorch wheel index
    if   [[ $CUDA_MAJOR -eq 11 ]]; then
        TORCH_INDEX="https://download.pytorch.org/whl/cu118"
    elif [[ $CUDA_MAJOR -eq 12 && $CUDA_MINOR -le 3 ]]; then
        TORCH_INDEX="https://download.pytorch.org/whl/cu121"
    elif [[ $CUDA_MAJOR -eq 12 && $CUDA_MINOR -le 5 ]]; then
        TORCH_INDEX="https://download.pytorch.org/whl/cu124"
    else
        TORCH_INDEX="https://download.pytorch.org/whl/cu126"
    fi
fi
echo "[INFO] Using PyTorch wheel index: $TORCH_INDEX"

# ── 5. Install dependencies ──────────────────────────────────
echo "[INFO] Installing Python dependencies (numpy + torch + colorama) …"
echo "       NOTE: torch is ~2 GB first download — this may take a few minutes."
uv pip install --python .venv/bin/python --index-url "$TORCH_INDEX" torch
uv pip install --python .venv/bin/python numpy colorama

echo ""
echo "[DONE] Setup complete!"
echo "       Run tests with:  ./run.sh"
echo "       Or directly:     .venv/bin/python gpu_onboard.py --help"