#!/bin/bash
# Waits for a final checkpoint, then:
#  1) keeps it forever as a hard-link copy (no extra disk space; survives deletion of the original)
#  2) exports a weights-only bf16 file (~8 GB) for downloading to a PC, plus its md5 fingerprint
# Settings (environment): NAME = checkpoint folder, OUT = target folder, FILE = export file name,
#   MAXSTEP (optional) = warn and stop if training passes this step without the checkpoint.
# Example, final Stage 2 model (login node):
#   NAME=ckpt_619699 OUT=/fsx/ajay-stu/ar-data/stage2_final FILE=model_stage2_final.pt \
#     nohup ./save_final.sh > logs/save_stage2.log 2>&1 &
set -uo pipefail
AR=${AR:-/fsx/ajay-stu/AR}
CKDIR=${CKDIR:-/fsx/ajay-stu/ar-data/checkpoints}
: "${NAME:?set NAME, e.g. ckpt_619699}" "${OUT:?set OUT, the target folder}" "${FILE:?set FILE, e.g. model_stage2_final.pt}"
MAXSTEP=${MAXSTEP:-0}
EXC="ip-10-0-136-14,ip-10-0-136-94,ip-10-0-136-253,ip-10-0-138-178,ip-10-0-139-113,ip-10-0-140-247"
cd "$AR" || exit 1
[ -f .venv/bin/activate ] && source .venv/bin/activate
export PYTHONPATH=$PWD

# latest training step from the running job's log (empty if no job is running)
latest_step() {
  local j
  j=$(squeue -u "${USER:-$(id -un)}" -h -n ar-real -t RUNNING -o %i 2>/dev/null | head -1)
  [ -n "$j" ] && grep "^step" "logs/real_$j.log" 2>/dev/null | tail -1 | awk '{print $2}' || true
}

# 1) wait until the checkpoint is completely written (.metadata is written last)
echo "$(date -u) waiting for $CKDIR/$NAME"
until [ -f "$CKDIR/$NAME/.metadata" ]; do
  s=$(latest_step)
  if [ "$MAXSTEP" -gt 0 ] && [ "${s:-0}" -gt "$MAXSTEP" ]; then
    echo "$(date -u) WARNING: training is at step $s but $NAME was never saved (disk full?)"
    exit 1
  fi
  sleep "${POLL:-300}"
done
sleep "${SETTLE:-60}"                                     # let the last file writes settle
echo "$(date -u) found complete $NAME"

# 2) permanent hard-link copy
mkdir -p "$OUT"
[ -d "$OUT/$NAME" ] || cp -al "$CKDIR/$NAME" "$OUT/$NAME" || { echo "hard-link copy failed"; exit 1; }
echo "$(date -u) permanent copy: $OUT/$NAME ($(ls "$OUT/$NAME" | wc -l) files)"

# 3) weights-only bf16 export for the PC (CPU job, ~5 min) and its fingerprint
M=$OUT/$FILE
if [ ! -f "$M" ]; then
  srun -p dev -N1 -n1 --gpus=0 --cpus-per-task=16 --exclude=$EXC \
    python -m ar.eval.export_ckpt --ckpt "$OUT/$NAME" --out "$M" || { echo "export failed"; exit 1; }
fi
( cd "$OUT" && md5sum "$FILE" > "${FILE%.pt}.md5" )
ls -lh "$M"
cat "$OUT/${FILE%.pt}.md5"
echo "$(date -u) DONE: $FILE saved"