#!/usr/bin/env bash
# Stop the AdamW control arm cleanly once step 400 is checkpointed.
set -uo pipefail; cd "$(dirname "$0")"
say(){ echo "[$(date -u +%H:%M:%S)] $*"; }
LOG=p1_const_adamw_logs/train.log
CK=checkpoints_const_adamw/rank_0/blumodelcp_run2_step_400.ckpt
SCRIPT_PID=30792

say "waiting for step 400 checkpoint to be written safely..."
for i in $(seq 1 360); do
  grep -qa 'Saved: .*step_400\.ckpt (safely)' "$LOG" && { say "step 400 checkpoint confirmed safe"; break; }
  kill -0 "$SCRIPT_PID" 2>/dev/null || { say "chain already exited"; break; }
  sleep 10
done
grep -qa 'Saved: .*step_400\.ckpt (safely)' "$LOG" || { say "ABORT: no safe step-400 ckpt, leaving run alone"; exit 1; }

sleep 5   # let the .zero_g1 shard flush
say "sizes: $(stat -Lc%s "$CK" 2>/dev/null) + $(stat -Lc%s "${CK%.ckpt}.zero_g1" 2>/dev/null)"

# parent script first so it cannot advance to the next stage, then the trainer
kill "$SCRIPT_PID" 2>/dev/null && say "stopped abl_phase1.sh ($SCRIPT_PID)"
for p in $(nvidia-smi --query-compute-apps=pid --format=csv,noheader); do
  kill "$p" 2>/dev/null && say "TERM trainer $p"
done
for i in $(seq 1 30); do
  [ -z "$(nvidia-smi --query-compute-apps=pid --format=csv,noheader)" ] && break; sleep 2; done
for p in $(nvidia-smi --query-compute-apps=pid --format=csv,noheader); do
  kill -9 "$p" 2>/dev/null && say "KILL -9 $p"; done
sleep 3
say "GPU now: $(nvidia-smi --query-gpu=memory.used --format=csv,noheader)"
say "last step: $(grep '^step' "$LOG" | tail -1)"
say "CUT_DONE"