#!/bin/bash
# Idempotent: safe to re-run after a stop/wake. Resumes training from ~/checkpoints/stage1/latest.pt.
set -uo pipefail
cd "$HOME/atlas"
export PYTHONUNBUFFERED=1
export HF_HOME="$HOME/.cache/huggingface"
export HF_HUB_ENABLE_HF_TRANSFER=1

echo "=== [1/4] python packages ==="
PIPFLAGS="--user --break-system-packages"
python3 -c "import sys; sys.exit(0 if sys.prefix != sys.base_prefix else 1)" && PIPFLAGS=""
python3 -m pip install -q $PIPFLAGS -U "transformers>=4.56" accelerate "huggingface_hub>=0.34" hf_transfer pillow sentencepiece protobuf || exit 10
export PATH="$HOME/.local/bin:$PATH"
python3 -c "import torch, transformers; print('torch', torch.__version__, '| transformers', transformers.__version__, '| cuda', torch.cuda.is_available())"
nvidia-smi --query-gpu=name,memory.total,memory.used --format=csv

echo "=== [2/4] downloads ==="
python3 prepare.py || exit 11
df -h "$HOME" | tail -1

export HF_HUB_OFFLINE=1   # everything is cached now; training never needs the token
GC="${GRAD_CKPT:-0}"
if [ "${SKIP_SMOKE:-0}" != "1" ]; then
  echo "=== [3/4] smoke test (20 steps) ==="
  rm -rf "$HOME/checkpoints/smoke"
  MAX_STEPS=20 LOG_EVERY=5 SAVE_EVERY=1000000 CKPT_DIR="$HOME/checkpoints/smoke" GRAD_CKPT=$GC python3 train.py
  rc=$?
  if [ $rc -eq 3 ] && [ "$GC" = "0" ]; then
    echo "OOM without gradient checkpointing -> retrying with GRAD_CKPT=1"
    GC=1
    rm -rf "$HOME/checkpoints/smoke"
    MAX_STEPS=20 LOG_EVERY=5 SAVE_EVERY=1000000 CKPT_DIR="$HOME/checkpoints/smoke" GRAD_CKPT=1 python3 train.py
    rc=$?
  fi
  if [ $rc -ne 0 ]; then echo "SMOKE TEST FAILED (exit $rc)"; exit 12; fi
fi

echo "=== [4/4] full stage-1 run (GRAD_CKPT=$GC) ==="
GRAD_CKPT=$GC python3 train.py
rc=$?
if [ $rc -eq 3 ] && [ "$GC" = "0" ]; then
  echo "OOM in full run -> resuming with GRAD_CKPT=1"
  GRAD_CKPT=1 python3 train.py
  rc=$?
fi
echo "=== run.sh finished (exit $rc) ==="
exit $rc