#!/usr/bin/env bash
# =============================================================================
# One command: 8x H200, 128k extension + NIAH eval. Setup ~17 min (build and
# downloads run CONCURRENTLY -- the build is the only critical path).
#
# WHY H200: 128k at cp=8 needs 16,384 tok/rank -> ~77,400 MB. Measured OOM on
# H100 (81,559 MB, 94.9%); on H200 (143,771 MB) that is 53.8%. cp=8 is the
# ceiling on 8 GPUs so there is no other lever.
#
# INVARIANT: SCALE 32 x ORIG 4096 == CP_T 131072. The 32k run broke this
# (32 x 4096 != 32768), trained on virtual 0..1023 instead of 0..4095, and went
# 90 -> 49. t128_train.sh hard-fails rather than repeat it.
# =============================================================================
set -uo pipefail
cd "$(dirname "$0")"
R=/root/Adhi/BluTrain/dist/Context_Parallelism
say(){ echo "[$(date -u +%H:%M:%S)] $*"; }

say "=== STAGE 1: build + downloads, concurrently ==="
( cd "$R"
  make libtensor SM_ARCH=90 -j$(nproc) > b_lib.log 2>&1 || { echo LIBFAIL; exit 1; }
  make CP_FUSED_ROPE=1 CP_ATTN_SPLIT=1 SM_ARCH=90 bluscript-cp -j$(nproc) > b_cp.log 2>&1
  [ -x build/bluscriptCP_exec ] && echo "BUILD_OK $(stat -c%s build/bluscriptCP_exec)" || { echo BUILDFAIL; tail -20 b_cp.log; }
) > "$R/_build.out" 2>&1 &
BPID=$!

( cd "$R"
  export PATH=/root/Adhi/BluTrain/envv/bin:$PATH
  pip install -q tiktoken pyarrow 2>&1 | tail -1
  mkdir -p checkpoints_680m/rank_0 checkpoints_bluscriptcp/rank_0 \
           Data_Loader/Data/b2_8192_16384 Data_Loader/Data/evalsrc Tests/bluscriptcp/qa_data
  R1=unparallelled/BluScriptCP; RD=unparallelled/BluScriptCP-data
  g(){ hf download "$1" ${3:+--repo-type dataset} --quiet --include "$2" --local-dir "$4"; }
  g $R1 'checkpoints/t64_680m_step1200_rank0.ckpt' "" dl1 && mv -f dl1/checkpoints/*.ckpt checkpoints_bluscriptcp/rank_0/t64.ckpt
  g $R1 'checkpoints/p1_680m_step1200_rank0.ckpt'  "" dl2 && mv -f dl2/checkpoints/*.ckpt checkpoints_bluscriptcp/rank_0/p1final.ckpt
  hf download navingv/varybatch --quiet --include 'varybatch.ckpt' --local-dir dl3 && mv -f dl3/varybatch.ckpt checkpoints_680m/rank_0/
  ln -sf "$PWD/checkpoints_680m/rank_0/varybatch.ckpt" checkpoints_bluscriptcp/rank_0/varybatch.ckpt
  g $RD 'cp_data_by_length/b2_8192_16384/val_shard_00000.bin' 1 dl4 && mv -f dl4/cp_data_by_length/b2_8192_16384/*.bin Data_Loader/Data/b2_8192_16384/
  for b in "b6_gte_131072 0 0" "b5_65536_131072 4 8" "b4_32768_65536 2 5" "b3_16384_32768 1 3" "b1_4096_8192 1 2"; do
    set -- $b; for i in $(seq -f '%05g' $2 $3); do
      hf download $RD --repo-type dataset --quiet --include "cp_data_by_length/$1/train_shard_$i.bin" --local-dir dl5 || true; done; done
  for d in dl5/cp_data_by_length/*; do n=$(basename $d); mkdir -p src128/$n; mv -n $d/*.bin src128/$n/ 2>/dev/null; done
  cp -f src128/b1_4096_8192/train_shard_*.bin Data_Loader/Data/evalsrc/ 2>/dev/null | head -1
  cp -f Data_Loader/Data/b2_8192_16384/val_shard_00000.bin Data_Loader/Data/evalsrc/
  g $R1 'data/ruler_qa/*' "" dl6; mv -f dl6/data/ruler_qa/*.parquet Tests/bluscriptcp/qa_data/ 2>/dev/null
  g $R1 'scripts/longeval_gen.py' "" dl7 && cp -f dl7/scripts/longeval_gen.py Tests/bluscriptcp/
  for s in t128_train.sh run_ruler.sh mix_buckets.py run_longeval.sh; do
    g $R1 "scripts/$s" "" dl8 && cp -f dl8/scripts/$s . 2>/dev/null; done
  chmod +x *.sh 2>/dev/null
  rm -rf dl1 dl2 dl3 dl4 dl5 dl6 dl7 dl8
  echo "DL_OK"
) > "$R/_dl.out" 2>&1 &
DPID=$!

wait $DPID; say "downloads: $(tail -1 $R/_dl.out)"
wait $BPID; say "build:     $(tail -1 $R/_build.out)"
grep -q BUILD_OK "$R/_build.out" || { say "FATAL build failed"; tail -20 "$R/_build.out"; exit 1; }
grep -q DL_OK    "$R/_dl.out"    || { say "FATAL downloads failed"; tail -20 "$R/_dl.out"; exit 1; }

say "=== STAGE 2: mix (b6 8.9 / b5 39 / b4 25 / b3 15 / b1 10) ==="
cd "$R"; export PATH=/root/Adhi/BluTrain/envv/bin:$PATH
python mix_buckets.py --out Data_Loader/Data/mix128 --data-root src128 \
  --src b6_gte_131072=0.11 --src b5_65536_131072=0.39 --src b4_32768_65536=0.25 \
  --src b3_16384_32768=0.15 --src b1_4096_8192=0.10 \
  --total-tokens 640e6 --shard-tokens 96e6 --window 131072 --grad-accum 4 2>&1 | tail -8
cp -f Data_Loader/Data/b2_8192_16384/val_shard_00000.bin Data_Loader/Data/mix128/

say "=== STAGE 3: training (probe is the go/no-go) ==="
./t128_train.sh || { say "training stopped -- see t128_logs/master.log"; exit 1; }

say "=== STAGE 4: NIAH eval, 3 models x 6 lengths, all 8 GPUs ==="
RN=$(cat t128_logs/runnum); FIN=$(ls -t checkpoints_bluscriptcp/rank_0/blumodelcp_run${RN}_step_*.ckpt | head -1)
EXTRA="t128=$(basename $FIN) t64=t64.ckpt base680=varybatch.ckpt" \
MODELS="t128 t64 base680" LENS="4096 8192 16384 32744 65512 131000" N=100 \
RULER_TASKS="niah_single" ./run_ruler.sh
say "=== DONE ==="