#!/usr/bin/env bash
# =============================================================================
# ckpt_shipper.sh -- upload checkpoints to HF, VERIFY, then free the disk.
#
# WHY THIS EXISTS
#   A fully resumable save is 8 ranks x (.ckpt 3.75 GB + .zero_g1 0.12 GB)
#   = 30.94 GB. Keeping every save locally needs 647 GB on a 193 GB disk.
#   Shipping and deleting keeps peak at ~106 GB (55%).
#
# POLICY
#   FULL resumable (all 8 ranks, .ckpt + .zero_g1)  at EVAL_STEPS and the final
#   rank_0 .ckpt only                               at every other FREQ save
#
# THREE RULES, EACH FROM A FAILURE THIS PROJECT ALREADY HAD
#   1. DELETE ONLY AFTER HF CONFIRMS THE BYTE SIZE. Not after `hf upload`
#      returns: on 2026-09-10 that returned SUCCESS for a missing file and
#      17 GB of checkpoints were lost with nothing in the log.
#   2. NEVER DELETE THE NEWEST FULL SAVE. It stays until the next one is
#      verified, so an HF outage cannot leave the run with nothing to resume.
#   3. NEVER TOUCH WHAT CP_CKPT_KEEP IS STILL USING. The trainer's live rolling
#      checkpoints are off-limits; deleting them breaks resume.
# =============================================================================
set -uo pipefail
cd "$(dirname "$0")"
source /root/Adhi/BluTrain/envv/bin/activate
HF=hf; PY=python
REPO=${REPO:-unparallelled/BluScriptCP}
CKD=checkpoints_bluscriptcp
LOG=${LOG:-mix64k_logs}
TAG=${TAG:-mix2}                     # HF path prefix for this run
FREQ=${FREQ:-200}
EVAL_STEPS=${EVAL_STEPS:-"2250 3300"}   # full resumable here + at the final
NP=${NP:-8}
S="$LOG/shipper.log"; mkdir -p "$LOG"; : > "$S"
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$S" >&2; }

# upload one file and CONFIRM the size HF reports back. Returns 1 on any doubt.
ship(){ # src dst
  local src=$1 dst=$2
  [ -f "$src" ] || { say "  MISSING (not uploaded): $src"; return 1; }
  $HF upload "$REPO" "$src" "$dst" >>"$LOG/upload.log" 2>&1 || { say "  UPLOAD FAILED $dst"; return 1; }
  $PY - "$REPO" "$src" "$dst" <<'PYEOF'
import sys, os
from huggingface_hub import HfApi
repo, loc, rem = sys.argv[1], sys.argv[2], sys.argv[3]
want = os.path.getsize(loc)
got = {f.rfilename: (f.size or 0) for f in HfApi().repo_info(repo, files_metadata=True).siblings}.get(rem)
ok = (got == want)
print(("  VERIFIED " if ok else "  *** MISMATCH/ABSENT ") + f"{rem} local={want:,} remote={got}")
sys.exit(0 if ok else 1)
PYEOF
}

is_eval_step(){ for e in $EVAL_STEPS; do [ "$1" = "$e" ] && return 0; done; return 1; }

LAST_FULL=""          # rule 2: the newest verified full save stays on disk
say "=== shipper armed: FREQ=$FREQ, full resumable at: $EVAL_STEPS ==="

while true; do
  RN=$(cat "$LOG/runnum" 2>/dev/null) || { sleep 30; continue; }
  for f in $CKD/rank_0/blumodelcp_run${RN}_step_*.ckpt; do
    [ -e "$f" ] || continue
    STEP=$(basename "$f" | grep -oE 'step_[0-9]+' | cut -d_ -f2)
    [ -n "$STEP" ] || continue
    # The trainer saves on (step+1)%freq==0 OR step==max_steps-1, so the FINAL
    # step checkpoints even when it is not a multiple of FREQ (2100 % 200 = 100).
    # Shipping only multiples of FREQ would silently drop the final model.
    if [ $(( STEP % FREQ )) -ne 0 ] && ! is_eval_step "$STEP"; then continue; fi
    [ -f "$LOG/shipped_${STEP}.ok" ] && continue
    # all ranks must have finished writing before we touch anything
    n=$(ls $CKD/rank_*/blumodelcp_run${RN}_step_${STEP}.ckpt 2>/dev/null | wc -l)
    [ "$n" -ge "$NP" ] || continue
    sleep 5

    if is_eval_step "$STEP"; then
      say "step $STEP -> FULL resumable (8 ranks, .ckpt + .zero_g1)"
      OK=1
      for r in $(seq 0 $((NP-1))); do
        ship "$CKD/rank_$r/blumodelcp_run${RN}_step_${STEP}.ckpt" \
             "checkpoints/${TAG}_step${STEP}_rank${r}.ckpt"    >>"$S" 2>&1 || OK=0
        ship "$CKD/rank_$r/blumodelcp_run${RN}_step_${STEP}.zero_g1" \
             "checkpoints/${TAG}_step${STEP}_rank${r}.zero_g1" >>"$S" 2>&1 || OK=0
      done
      if [ "$OK" = 1 ]; then
        touch "$LOG/shipped_${STEP}.ok"; echo full > "$LOG/shipped_${STEP}.kind"
        say "  step $STEP FULL verified on HF"
        # rule 2: only now is it safe to drop the PREVIOUS full save's local copies
        if [ -n "$LAST_FULL" ] && [ "$LAST_FULL" != "$STEP" ]; then
          rm -f $CKD/rank_*/keep_step_${LAST_FULL}.* 2>/dev/null
          say "  freed local copies of the previous full save (step $LAST_FULL)"
        fi
        LAST_FULL="$STEP"
      else
        say "  step $STEP FULL INCOMPLETE -- keeping every local copy"
      fi
    else
      say "step $STEP -> rank_0 only (bisection point)"
      if ship "$CKD/rank_0/blumodelcp_run${RN}_step_${STEP}.ckpt" \
              "checkpoints/${TAG}_step${STEP}.ckpt" >>"$S" 2>&1; then
        touch "$LOG/shipped_${STEP}.ok"; echo rank0 > "$LOG/shipped_${STEP}.kind"
        say "  step $STEP verified on HF"
      else
        say "  step $STEP NOT verified -- nothing deleted"
      fi
    fi
  done

  # rule 3: reclaim ONLY orphaned sidecars -- CP_CKPT_KEEP prunes .ckpt but never
  # .zero_g*, the leak that filled a 193 GB disk on 2026-09-09. The trainer's live
  # rolling checkpoints are never touched.
  for s in $CKD/rank_*/*.zero_g*; do
    [ -e "$s" ] || continue
    [ -f "${s%.zero_g*}.ckpt" ] || rm -f "$s"
  done
  df -BG --output=avail / | tail -1 | tr -dc 0-9 > "$LOG/disk_free_gb"
  sleep 45
done