#!/usr/bin/env bash
# =============================================================================
# eval_at_1000.sh -- capture the long b3 arm at step 1000 and score it.
#
# WHY A WATCHER AND NOT A TIMER: CP_CKPT_KEEP=2 with CP_CKPT_FREQ=25 means a
# checkpoint is pruned ~50 steps (~10 min) after it is written. The step-1000
# files must be secured the instant they appear on BOTH ranks, or they are gone.
#
# ORDER MATTERS: kill the trainer FIRST (stopping further saves, hence further
# pruning), then copy aside. Copying while the trainer runs risks reading a file
# that the pruner deletes mid-copy.
#
# EVAL REUSES decide_len_* ON PURPOSE. run_longeval.sh's prompt signature covers
# N/length/depths/seed/haystack but NOT the model set, so armA_c/l1/l2/l3 are
# already .done there and get skipped -- only l3long's 3 cells actually run.
# 15 min instead of 40. Directly comparable to the decision table.
#
# CAVEAT ON THE NUMBER THIS PRODUCES: step 1000 sits in the FLAT LR region
# (constant 3.0e-4 until step 3899). WSD's decay is where much of the loss drop
# lands, so this checkpoint UNDERSTATES what the same data/steps would give after
# a decay. It is a valid step-1000-vs-step-400 comparison (both undecayed at
# eval? no -- l3@400 WAS decayed, its 400-step cosine ran to completion). So
# l3long@1000 is undecayed and l3@400 is decayed: this comparison is BIASED
# AGAINST l3long. If it wins anyway, that is robust. If it loses, run the short
# decay before believing it:  STEPS=1100 STABLE=900 ./b3_long_run.sh resume 17
# =============================================================================
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=/root/Adhi/BluTrain/envv/bin/python
CK=checkpoints_bluscriptcp/rank_0
LOG=eval1000_logs; mkdir -p "$LOG"
M="$LOG/master.log"
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$M" >&2; }

RN=$(cat b3long_logs/runnum)
TARGET=${TARGET:-1000}
NAME="blumodelcp_run${RN}_step_${TARGET}.ckpt"
say "watching for $NAME on both ranks (run $RN, target step $TARGET)"

for i in $(seq 1 2400); do          # up to 20 h, 30 s poll
  n=$(ls checkpoints_bluscriptcp/rank_*/"$NAME" 2>/dev/null | wc -l)
  [ "$n" -ge 2 ] && break
  sleep 30
done
n=$(ls checkpoints_bluscriptcp/rank_*/"$NAME" 2>/dev/null | wc -l)
[ "$n" -ge 2 ] || { say "FATAL: $NAME never appeared on 2 ranks"; exit 1; }
say "step $TARGET present on $n ranks -- stopping trainer NOW to halt pruning"

for p in $(pgrep -f 'b3_long_run\.sh'); do kill "$p" 2>/dev/null; done
sleep 2
for p in $(pgrep -f 'mpirun.*bluscriptCP_exec') $(pgrep -f 'build/bluscriptCP_exec'); do kill "$p" 2>/dev/null; done
sleep 10
say "trainer stopped at step $(grep -acE '^step +[0-9]+ \|' b3long_logs/l3long.log)"

# preserve under a name the pruner will never touch (different run prefix)
SAFE="l3long1000_${NAME}"
for r in 0 1; do
  [ -f "checkpoints_bluscriptcp/rank_$r/$SAFE" ] && continue
  cp "checkpoints_bluscriptcp/rank_$r/$NAME" "checkpoints_bluscriptcp/rank_$r/$SAFE"
done
say "preserved as $SAFE ($(ls checkpoints_bluscriptcp/rank_*/$SAFE | wc -l) ranks, $(du -ch checkpoints_bluscriptcp/rank_*/$SAFE | tail -1 | cut -f1))"

EXTRA="armA_c=blumodelcp_run28_step_400.ckpt \
l1=lenbucket_l1_blumodelcp_run1_step_400.ckpt \
l2=lenbucket_l2_blumodelcp_run2_step_400.ckpt \
l3=blumodelcp_run1_step_400.ckpt \
l3long=$SAFE"
MODELS="armA_c l1 l2 l3 l3long"
say "eval models: $MODELS  (only l3long's cells are new)"

for LEN in 4096 8192 16360; do
  say "PROMPT_LEN=$LEN"
  GEN_DATA_ROOT=lendata/b3_16384_32768 \
  HAY="$PWD/lendata/haystack/b2val_haystack.bin" \
  PY=$PY OUT=decide_len_$LEN TASKS="niah_single" CKPTS="$MODELS" N=100 \
  PROMPT_LEN=$LEN MAX_NEW=24 T=16384 DEPTHS="0.0,0.25,0.5,0.75,1.0" \
  YARN_ORIG=4096 GPUS=0,1 JOBS_PER_GPU=2 BF16=1 GREEDY=1 \
  CKPT_EXTRA="$EXTRA" ./run_longeval.sh >> "$LOG/eval_$LEN.log" 2>&1 \
    || say "  WARN nonzero at $LEN"
  say "  -> $(sed -n 1p decide_len_$LEN/summary.csv) / $(sed -n 2p decide_len_$LEN/summary.csv)"
done
say "=== EVAL@1000 DONE ==="

# ---- CONTINUE the run from where it stopped -------------------------------
# resume mode pins CP_REWARMUP_ANCHOR=0, so the flat 3.0e-4 continues instead of
# a fresh 100-step warmup ramp re-anchoring at step 1000 (bluscriptCP.cpp:2105).
say "resuming run $RN to step 4000"
setsid nohup ./b3_long_run.sh resume "$RN" >> b3long_boot.out 2>&1 < /dev/null &
sleep 90
say "resumed: step $(grep -acE '^step +[0-9]+ \|' b3long_logs/l3long.log) logged, $(pgrep -c -f 'build/bluscriptCP_exec') trainer proc(s)"
say "LR after resume: $(grep -aE '^step ' b3long_logs/l3long.log | tail -1 | grep -oE 'lr [0-9.e-]+')  (must still be ~3.0e-4, NOT ramping)"