#!/usr/bin/env bash
# One-time setup of a remote GPU box for batched ATOM training (docs/BATCHED_TRAINING.md §Runbook).
# Run from the repo root on the GPU box:   bash scripts/setup_gpu_box.sh
# Env: PYTHON=python3 (>= 3.10)  VENV=.venv  TORCH_INDEX="" (e.g. https://download.pytorch.org/whl/cu128 to
#      force a CUDA build; empty = PyPI default wheel, which is CUDA-enabled on Linux x86_64)
#      SKIP_TESTS=0  BENCH=1
# No HF / transformers / tokenizers are installed or needed: torch + numpy (+ pytest for the checks).
set -euo pipefail
cd "$(dirname "$0")/.."
PYTHON=${PYTHON:-python3}; VENV=${VENV:-.venv}; TORCH_INDEX=${TORCH_INDEX:-}; SKIP_TESTS=${SKIP_TESTS:-0}; BENCH=${BENCH:-1}
echo "== host: $(hostname)  $(date -Is)"
echo "== GPU"; if command -v nvidia-smi >/dev/null; then
  nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv,noheader
  nvidia-smi | sed -n '1,4p' | grep -i "cuda version" || true
else echo "nvidia-smi NOT found: no NVIDIA driver -> CUDA training impossible"; exit 1; fi
echo "== disk (repo fs)"; df -h . | tail -1
echo "== CPU/RAM"; nproc; free -g | sed -n '1,2p'
$PYTHON - <<'P'
import sys
assert sys.version_info >= (3, 10), f"python >= 3.10 needed, got {sys.version}"
print("python", sys.version.split()[0])
P
[ -x "$VENV/bin/python" ] || $PYTHON -m venv "$VENV"
"$VENV/bin/python" -m pip install -q --upgrade pip
if [ -n "$TORCH_INDEX" ]; then "$VENV/bin/pip" install -q torch --index-url "$TORCH_INDEX"; else "$VENV/bin/pip" install -q torch; fi
"$VENV/bin/pip" install -q numpy pytest
"$VENV/bin/python" - <<'P'
import torch
print("torch", torch.__version__, "cuda build", torch.version.cuda, "cuda available", torch.cuda.is_available())
assert torch.cuda.is_available(), "torch cannot see the GPU: set TORCH_INDEX to a CUDA wheel matching the driver"
for i in range(torch.cuda.device_count()):
    p = torch.cuda.get_device_properties(i)
    print(f"  gpu{i}: {p.name} {p.total_memory/2**30:.1f} GiB sm_{p.major}{p.minor}")
x = torch.randn(1024, 1024, device="cuda"); print("  matmul ok", float((x @ x).sum()) == float((x @ x).sum()))
P
if [ "$SKIP_TESTS" != "1" ]; then
  echo "== batched equivalence tests (CPU + CUDA)"
  PYTHONPATH=. "$VENV/bin/python" -m pytest -q test/test_batched_train.py
fi
if [ "$BENCH" = "1" ]; then
  echo "== short GPU benchmark (bytes/s vs B)"
  mkdir -p docs/artifacts/batched_training
  PYTHONPATH=. "$VENV/bin/python" tools/bench_batched.py --device cuda --threads 4 --nice 0 --k 128 --windows 2 \
    --batch 64 256 1024 --out docs/artifacts/batched_training/bench_gpu_$(hostname).json
fi
echo "== setup OK"