#!/bin/bash
# Stage 2: smoke test -> full LoRA instruction tuning -> evaluation. Safe to re-run (resumes).
set -uo pipefail
cd "$HOME/atlas"
export PYTHONUNBUFFERED=1 HF_HOME="$HOME/.cache/huggingface" HF_HUB_OFFLINE=1 PATH="$HOME/.local/bin:$PATH"
python3 -m pip install -q --user --break-system-packages peft pyarrow || exit 10
python3 -c "import peft, pyarrow; print('peft', peft.__version__, '| pyarrow', pyarrow.__version__)"
BS="${BATCH_SIZE:-8}"; GA="${GRAD_ACCUM:-4}"

if [ "${SKIP_SMOKE:-0}" != "1" ]; then
  echo "=== [1/3] smoke test (10 steps) ==="
  rm -rf "$HOME/checkpoints/stage2_smoke"
  MAX_STEPS=10 LOG_EVERY=2 SAVE_EVERY=1000000 CKPT_DIR="$HOME/checkpoints/stage2_smoke" BATCH_SIZE=$BS GRAD_ACCUM=$GA python3 train_stage2.py
  rc=$?
  if [ $rc -eq 3 ]; then
    echo "OOM -> retrying with BATCH_SIZE=4 GRAD_ACCUM=8"; BS=4; GA=8
    rm -rf "$HOME/checkpoints/stage2_smoke"
    MAX_STEPS=10 LOG_EVERY=2 SAVE_EVERY=1000000 CKPT_DIR="$HOME/checkpoints/stage2_smoke" BATCH_SIZE=$BS GRAD_ACCUM=$GA python3 train_stage2.py
    rc=$?
  fi
  if [ $rc -ne 0 ]; then echo "SMOKE TEST FAILED (exit $rc)"; exit 12; fi
fi

echo "=== [2/3] full stage-2 run (BATCH_SIZE=$BS GRAD_ACCUM=$GA) ==="
BATCH_SIZE=$BS GRAD_ACCUM=$GA python3 train_stage2.py
rc=$?
if [ $rc -eq 3 ] && [ "$BS" -gt 4 ]; then
  echo "OOM in full run -> continuing with BATCH_SIZE=4 GRAD_ACCUM=8"
  BATCH_SIZE=4 GRAD_ACCUM=8 python3 train_stage2.py
  rc=$?
fi
if [ $rc -ne 0 ]; then echo "TRAINING FAILED (exit $rc)"; exit $rc; fi

echo "=== [3/3] evaluation ==="
python3 eval_stage2.py
rc=$?
echo "=== run_stage2.sh finished (exit $rc) ==="
exit $rc