#!/usr/bin/env bash
# Pod setup (classifier-plan §3.2 stage "setup"), run on the pod after `sync.sh up`; rerun after a pod restart.
#   bootstrap.sh            download google/gemma-4-12B and google/gemma-4-12B-it, check_env, selftest
#   BASE_ONLY=1 bootstrap.sh  base checkpoint only (screen dropped by the ledger replan rule)
# venv, HF cache and base weights live on the container disk (/root): the network volume writes small files slowly
# and every stage reloads ~24 GB. Runs live on /workspace (survives the pod). No credentials are needed: the base
# checkpoints are public and ungated (Apache-2.0); nothing is pushed from the pod.
set -euo pipefail
W=/workspace
CODE=$W/code
M=${MODELS:-/root/models}
V=${VENV:-/root/venv}
export HF_HOME=/root/hf HF_XET_HIGH_PERFORMANCE=1
cd "$CODE"
[ -d "$V" ] || python3 -m venv --system-site-packages "$V"
. "$V/bin/activate"
pip install -q --upgrade pip
# Keep the image's torch: nothing may replace it.
python -c "import torch; print('torch==' + torch.__version__.split('+')[0])" > "$W/constraints.txt"
before=$(python -c "import torch; print(torch.__version__)")
pip install -q -c "$W/constraints.txt" -r training/requirements-train.txt
after=$(python -c "import torch; print(torch.__version__)")
[ "$before" = "$after" ] || { echo "FATAL: pip replaced the image's torch ($before -> $after)"; exit 1; }
models=(gemma-4-12B)
[ "${BASE_ONLY:-0}" = 1 ] || models+=(gemma-4-12B-it)
df -h /root | tail -1
for m in "${models[@]}"; do
  [ -f "$M/$m/.complete" ] || { hf download "google/$m" --local-dir "$M/$m" && touch "$M/$m/.complete"; }
done
nvidia-smi --query-gpu=name,memory.total,driver_version --format=csv | tee "$W/nvidia-smi.txt"
pip freeze > "$W/pip-freeze.txt"
status=0
args=()
for m in "${models[@]}"; do args+=(--model "$M/$m"); done
python -m training.check_env "${args[@]}" --out "$W/check_env.json" || status=1
# Tiny random Gemma 4 with the real attention shapes of the base config (no 12B weights needed).
python -m training.selftest --device cuda --model "$M/gemma-4-12B" --out "$W/selftest.json" || status=1
if [ $status = 0 ]; then echo "BOOTSTRAP OK"; else
  echo "BOOTSTRAP CHECKS FAILED: read check_env.json and selftest.json; no 12B run until fixed (§4.3 fallbacks)"; fi
exit $status