#!/usr/bin/env bash
# =============================================================================
# finish_all.sh -- run the rest of the 2026-09-03 plan UNATTENDED:
#
#   arm C' -> arm A' -> eval both -> rebuild notebook -> upload to HF -> ready
#
# Detached with nohup, so it survives the chat session closing.
#
# ALREADY DONE before this script starts (do not redo):
#   run7' (run 6, step 750)      corrected-cache 16k base
#   arm B' (run 7, step 400)     IN2 needle at distance 0
#   eval of base_c + armB_c at 1024/2048/4096/8192/16360, N=100
#   cache_fix_results.ipynb      run7' vs run7 comparison
#
# RUN NUMBERING: bluscriptCP picks run_number as the first free
# CP_BluScreipt_Training_logs/CP_Training_log<N>.csv. Logs 1-26 exist, so arm C
# takes 27 and arm A takes 28. Neither collides with the run6/7/8 checkpoints --
# which matters, because exactly that collision (arm B taking run 7 while an old
# run7_step_750 sat in rank_0) let CP_CKPT_KEEP=1 prune arm B's step-200 shard on
# rank_0 and made it unresumable.
#
# EVERY checkpoint is verified across ALL FOUR ZeRO-1 rank shards. A zero exit
# code is not proof: bluscriptCP catches a failed save and only prints [WARN],
# and a 3-of-4 shard set loads silently wrong.
#
# THE EVAL REUSES the existing newck_len_* dirs on purpose. run_longeval.sh's
# prompt signature covers N, length, depths, seed and haystack -- NOT the model
# set -- so adding armC_c/armA_c to CKPT_EXTRA reuses the prompts, skips the
# finished base_c/armB_c cells, and rebuilds summary.csv with all four models.
# DO NOT change PROMPT_LEN/N/DEPTHS here or the signature flips and every
# completed cell is wiped.
# =============================================================================
set -uo pipefail
cd "$(dirname "$0")"

ROOT="$(cd ../.. && pwd)"
export LD_LIBRARY_PATH="$ROOT/Tensor-Implementations/lib:$ROOT/Profiler/lib:${LD_LIBRARY_PATH:-}"
export OMPI_ALLOW_RUN_AS_ROOT=1 OMPI_ALLOW_RUN_AS_ROOT_CONFIRM=1

PY=${PY:-/root/Adhi/BluTrain/envv/bin/python}
HFCLI=/root/Adhi/BluTrain/envv/bin/hf
EXE=./build/bluscriptCP_exec
CK=checkpoints_bluscriptcp/rank_0
LOG=finish_logs; mkdir -p "$LOG"
M="$LOG/finish.log"
REPO=unparallelled/BluScriptCP

say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }
die(){ say "FATAL: $*"; exit 1; }

T=16384; YARN=4; ORIG=4096; NP=4; GPUS=0,1,2,3
ARCH="CP_N_EMBD=1536 CP_N_LAYER=36 CP_N_HEAD=12 CP_N_KVHEAD=4 CP_FFN=4096 CP_WEIGHT_TYING=0"
COMMON="CP_SIZE=$NP CP_ATTN_MODE=ring CP_ATTN_FUSION=split CP_T=$T \
YARN_SCALE=$YARN YARN_ORIG_MAXPOS=$ORIG \
CP_ZERO=1 CP_ZERO_V2=1 CP_ZERO_SYMMETRIC=1 CP_BF16=1 CP_OPTIMIZER=muon \
CP_MAX_LR=1e-3 CP_MIN_LR=1e-4 CP_WARMUP=1 \
CP_CKPT=1 CP_CKPT_KEEP=1 CP_CKPT_NEW_RUN=1"

R7="$CK/blumodelcp_run6_step_750.ckpt"      # run7', every arm's parent
LENS="1024 2048 4096 8192 16360"
N=100
TASKS="niah_single val_backward"

[ -x "$EXE" ] || die "$EXE missing"
[ -f "$R7" ]  || die "run7' parent checkpoint missing: $R7"

say "=== finish_all.sh start ==="
say "free disk $(df -BG --output=avail /root | tail -1 | tr -dc '0-9')G"

# Warn about the HF token UP FRONT rather than 2 hours later at the upload step.
if [ -z "${HF_TOKEN:-}" ] && [ ! -f "$HOME/.cache/huggingface/token" ] && [ ! -f hf_token.txt ]; then
  say "WARNING: no HF token found. Training and eval will proceed; the UPLOAD"
  say "  step will wait for one. Provide it either way:"
  say "    echo '<token>' > $(pwd)/hf_token.txt"
  say "    or: export HF_TOKEN=<token>  (only helps if set before this script)"
fi

# ---------------------------------------------------------------------------
# train one arm and verify it landed on all four ranks
# ---------------------------------------------------------------------------
train_arm(){          # $1 = arm letter
  # Separate statements: in `local A=$1 lf="...${A}..."` bash expands every word
  # BEFORE the local builtin runs, so ${A} is still unset and `set -u` aborts.
  local A="$1"
  local lf="$LOG/arm_${A}.log"
  if [ -f "$LOG/arm_${A}.ckpt" ] && [ -f "$(cat "$LOG/arm_${A}.ckpt")" ]; then
    say "arm $A already done -> $(cat "$LOG/arm_${A}.ckpt")"; return 0
  fi
  [ -d "in2_data/arm_$A" ] || { say "arm $A: no in2_data/arm_$A -- SKIPPING"; return 1; }
  say "training arm $A' (400 steps from run7')"
  env $ARCH $COMMON CP_INIT_FROM="$PWD/$R7" CP_DATA_ROOT="$PWD/in2_data/arm_$A" \
      CP_B=4 CP_GLOBAL_BATCH=524288 CP_MAX_STEPS=400 \
      CP_REWARMUP=100 CP_REWARMUP_PEAK=0.3 CP_CKPT_FREQ=200 \
      CUDA_VISIBLE_DEVICES=$GPUS mpirun --allow-run-as-root -np $NP "$EXE" \
      > "$lf" 2>&1
  local rc=$?
  local RN; RN=$(grep -aoE 'CP_Training_log[0-9]+' "$lf" | head -1 | grep -oE '[0-9]+')
  [ -n "$RN" ] || { say "arm $A: no run number in $lf (rc=$rc) -- FAILED"; return 1; }
  local CKF; CKF=$(ls -t $CK/blumodelcp_run${RN}_step_*.ckpt 2>/dev/null | head -1)
  [ -n "$CKF" ] || { say "arm $A: run $RN wrote NO checkpoint (rc=$rc) -- FAILED"; return 1; }
  # all four ZeRO-1 shards, or the load is silently partial
  local base n; base=$(basename "$CKF")
  n=$(ls checkpoints_bluscriptcp/rank_*/"$base" 2>/dev/null | wc -l)
  [ "$n" -ge 4 ] || { say "arm $A: only $n/4 rank shards for $base -- FAILED"; return 1; }
  grep -qa 'checkpoint save.*failed' "$lf" && { say "arm $A: a checkpoint save FAILED"; return 1; }
  echo "$CKF" > "$LOG/arm_${A}.ckpt"
  say "arm $A' -> run $RN, $base, final loss $(grep -aE '^step ' "$lf" | tail -1 | grep -oE 'loss: [0-9.]+')"
  return 0
}

train_arm c || say "arm C did not complete -- continuing with whatever exists"
train_arm a || say "arm A did not complete -- continuing with whatever exists"

# ---------------------------------------------------------------------------
# evaluate every model that exists, matched cache (YARN_ORIG=4096)
# ---------------------------------------------------------------------------
EXTRA="base_c=blumodelcp_run6_step_750.ckpt"
B=$(ls -t $CK/blumodelcp_run7_step_400.ckpt 2>/dev/null | head -1)
[ -n "$B" ] && EXTRA="$EXTRA armB_c=$(basename "$B")"
for A in c a; do
  f="$LOG/arm_${A}.ckpt"
  [ -f "$f" ] || continue
  p=$(cat "$f"); [ -f "$p" ] || continue
  EXTRA="$EXTRA arm$(echo $A | tr a-z A-Z)_c=$(basename "$p")"
done
MODELS=$(echo $EXTRA | tr ' ' '\n' | cut -d= -f1 | tr '\n' ' ')
say "eval models: $MODELS"

for LEN in $LENS; do
  say "eval PROMPT_LEN=$LEN (finished cells are skipped)"
  PY=$PY OUT=newck_len_$LEN TASKS="$TASKS" CKPTS="$MODELS" N=$N \
  PROMPT_LEN=$LEN MAX_NEW=24 T=$T DEPTHS="0.0,0.25,0.5,0.75,1.0" \
  YARN_ORIG=$ORIG GPUS=0,1,2,3 JOBS_PER_GPU=2 BF16=1 GREEDY=1 \
  CKPT_EXTRA="$EXTRA" ./run_longeval.sh >> "$LOG/eval_$LEN.log" 2>&1 \
    || say "  WARN: eval at $LEN returned nonzero"
  [ -f "newck_len_$LEN/summary.csv" ] && say "  -> $(head -1 newck_len_$LEN/summary.csv)"
done
say "eval done"

# ---------------------------------------------------------------------------
# rebuild the notebook and the markdown report with all four models
# ---------------------------------------------------------------------------
$PY Tests/bluscriptcp/mk_cache_notebook.py --out cache_fix_results.ipynb \
    >> "$LOG/notebook.log" 2>&1 \
  && say "notebook rebuilt: cache_fix_results.ipynb" \
  || say "WARN: notebook rebuild failed -- see $LOG/notebook.log"

{
  echo "# Corrected-cache retrain results -- 2026-09-03"
  echo
  echo "Models: run7' (base_c), arm B' (armB_c), arm C' (armC_c), arm A' (armA_c),"
  echo "all trained under YARN_ORIG_MAXPOS=4096 and evaluated under the same."
  echo "Generative substring-match scoring, greedy, bf16, N=$N per cell."
  echo
  for LEN in $LENS; do
    echo "## PROMPT_LEN=$LEN"; echo '```'
    cat "newck_len_$LEN/summary.csv" 2>/dev/null || echo "(missing)"
    echo '```'; echo
  done
  echo "## Training losses"; echo '```'
  echo "run7'  (run 6, step 750): $(grep -aE '^step ' retrain_logs/run7c.log | tail -1 | grep -oE 'loss: [0-9.]+')"
  echo "arm B' (run 7, step 400): $(grep -aE '^step ' retrain_logs/arm_b.log | tail -1 | grep -oE 'loss: [0-9.]+')"
  for A in c a; do
    [ -f "$LOG/arm_${A}.log" ] && \
      echo "arm ${A}' : $(grep -aE '^step ' "$LOG/arm_${A}.log" | tail -1 | grep -oE 'loss: [0-9.]+')"
  done
  echo '```'
} > "$LOG/RESULTS.md"
say "wrote $LOG/RESULTS.md"

# ---------------------------------------------------------------------------
# upload. Waits up to 2 h for a token so the chat can supply one late.
# ---------------------------------------------------------------------------
for i in $(seq 1 120); do
  if [ -f hf_token.txt ]; then export HF_TOKEN=$(tr -d ' \n' < hf_token.txt); fi
  [ -n "${HF_TOKEN:-}" ] && break
  [ -f "$HOME/.cache/huggingface/token" ] && break
  [ "$i" = 1 ] && say "waiting for an HF token (checking hf_token.txt every 60s, up to 2h)"
  sleep 60
done

if [ -z "${HF_TOKEN:-}" ] && [ ! -f "$HOME/.cache/huggingface/token" ]; then
  say "UPLOAD SKIPPED: no HF token after 2h."
  cat > TEARDOWN_BLOCKED.txt <<EOF
NOT SAFE TO TEAR DOWN -- nothing was uploaded.

Everything below exists only on this box. Provide a token and re-run the upload:

    echo '<hf-token>' > $(pwd)/hf_token.txt
    ./finish_all.sh          # training/eval are idempotent, it jumps to upload

What still needs uploading:
    cache_fix_results.ipynb          the results notebook
    newck_len_*/                     corrected-cache eval data
    $LOG/                            arm C'/arm A' training logs + RESULTS.md
    retrain_logs/                    run7'/arm B' training logs
    checkpoints_bluscriptcp/rank_0/  run6 (run7'), run7 (arm B'), run27, run28
    checkpoints_old_run7/            yesterday's buggy-cache run7 base
EOF
  say "wrote TEARDOWN_BLOCKED.txt"
  exit 1
fi

say "uploading to $REPO"
up(){ # $1=local  $2=remote
  [ -e "$1" ] || { say "  skip $1 (absent)"; return 0; }
  if $HFCLI upload "$REPO" "$1" "$2" >>"$LOG/upload.log" 2>&1; then
    say "  ok   $1 -> $2"
  else
    say "  FAIL $1 -> $2 (see $LOG/upload.log)"; UPFAIL=1
  fi
}
UPFAIL=0
up cache_fix_results.ipynb                 docs/cache_fix_results_2026-09-03.ipynb
up "$LOG/RESULTS.md"                       docs/RETRAIN_RESULTS_2026-09-03.md
up EVAL_COMPARISON.md                      docs/EVAL_COMPARISON_2026-09-02.md
for LEN in $LENS; do up "newck_len_$LEN"   "logs/newck_len_$LEN/"; done
up retrain_logs                            logs/retrain_corrected_cache/
up "$LOG"                                  logs/finish_all/
up eval_logs                               logs/eval_run7c/
up eval_run7c.sh                           scripts/eval_run7c.sh
up finish_all.sh                           scripts/finish_all.sh
up retrain_corrected_cache.sh              scripts/retrain_corrected_cache.sh
up Tests/bluscriptcp/mk_cache_notebook.py  scripts/mk_cache_notebook.py
up Tests/bluscriptcp/longeval_gen.py       scripts/longeval_gen.py
up Tests/bluscriptcp/longeval_score.py     scripts/longeval_score.py
up run_longeval.sh                         scripts/run_longeval.sh

# checkpoints last: biggest, and least useful if something above failed
for f in $CK/blumodelcp_run6_step_750.ckpt $CK/blumodelcp_run7_step_400.ckpt; do
  [ -f "$f" ] && up "$f" "checkpoints/$(basename "$f")"
done
for A in c a; do
  [ -f "$LOG/arm_${A}.ckpt" ] && { p=$(cat "$LOG/arm_${A}.ckpt"); up "$p" "checkpoints/$(basename "$p")"; }
done

if [ "$UPFAIL" = 0 ]; then
  { echo "TEARDOWN READY -- $(date -u +'%Y-%m-%d %H:%M:%S UTC')"
    echo "All results, logs, scripts and checkpoints uploaded to $REPO."
    echo "Verify with: $HFCLI repo-files $REPO list | grep -E 'newck_len|cache_fix'"; } > TEARDOWN_READY.txt
  rm -f TEARDOWN_BLOCKED.txt
  say "=== TEARDOWN READY ==="
else
  say "=== SOME UPLOADS FAILED -- NOT safe to tear down (see $LOG/upload.log) ==="
  echo "Some uploads FAILED -- see $LOG/upload.log" > TEARDOWN_BLOCKED.txt
  exit 1
fi