#!/usr/bin/env bash
# Pod setup, run on the pod after sync.sh up: venv, packages, the base checkpoint, then the environment
# check (classifier.check_env) and the tiny-model tests (classifier.selftest). Rerun after a pod restart.
#   bootstrap.sh            Qwen3.5-9B-Base
#   bootstrap.sh instruct   also Qwen3.5-9B (only for the optional extra arm E1)
set -euo pipefail
W=/workspace
cd "$W/code"
# The base checkpoint and the HF cache also live on the container disk: the network volume (FUSE) reads and
# writes about 200 MB/s, and every stage reloads the 19 GB model. Runs stay on /workspace, which survives the pod.
M=${MODELS:-/root/models}
export HF_HOME=/root/hf HF_XET_HIGH_PERFORMANCE=1
export PATH=/usr/local/cuda/bin:$PATH CUDA_HOME=${CUDA_HOME:-/usr/local/cuda}
# The venv lives on the container disk: the network volume writes small files very slowly (reference lesson).
V=${VENV:-/root/venv}
[ -d "$V" ] || python3 -m venv --system-site-packages "$V"
. "$V/bin/activate"
pip install -q --upgrade pip
# Keep the image's torch and Triton: nothing may replace them (the tilelang pin and selftest target them).
python -c "import torch, triton; print('torch==' + torch.__version__.split('+')[0]); print('triton==' + triton.__version__)" > "$W/constraints.txt"
before=$(python -c "import torch; print(torch.__version__)")
pip install -q -c "$W/constraints.txt" -r classifier/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; }
# Compiles CUDA code (devel image); without it transformers falls back to PyTorch and check_env fails.
MAX_JOBS=${MAX_JOBS:-16} pip install -q --no-build-isolation -c "$W/constraints.txt" causal-conv1d==1.7.0 \
  || echo "WARN: causal-conv1d did not build"
models=(Qwen3.5-9B-Base)
[ "${1:-}" = instruct ] && models+=(Qwen3.5-9B)
for m in "${models[@]}"; do
  [ -f "$M/$m/.complete" ] || { hf download "Qwen/$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"
status=0
python -m classifier.check_env --model "$M/Qwen3.5-9B-Base" --out "$W/check_env.json" || status=1
python -m classifier.selftest --model "$M/Qwen3.5-9B-Base" --out "$W/selftest.json" || status=1
[ $status = 0 ] && echo "BOOTSTRAP OK" || echo "BOOTSTRAP CHECKS FAILED: read check_env.json and selftest.json before any training"
exit $status