Upload scripts/ckpt_shipper.sh with huggingface_hub

932355a verified about 2 months ago

5.37 kB

#!/usr/bin/env bash
# =============================================================================
# ckpt_shipper.sh -- upload checkpoints to HF, VERIFY, then free the disk.
#
# WHY THIS EXISTS
#   A fully resumable save is 8 ranks x (.ckpt 3.75 GB + .zero_g1 0.12 GB)
#   = 30.94 GB. Keeping every save locally needs 647 GB on a 193 GB disk.
#   Shipping and deleting keeps peak at ~106 GB (55%).
#
# POLICY
#   FULL resumable (all 8 ranks, .ckpt + .zero_g1)  at EVAL_STEPS and the final
#   rank_0 .ckpt only                               at every other FREQ save
#
# THREE RULES, EACH FROM A FAILURE THIS PROJECT ALREADY HAD
#   1. DELETE ONLY AFTER HF CONFIRMS THE BYTE SIZE. Not after `hf upload`
#      returns: on 2026-09-10 that returned SUCCESS for a missing file and
#      17 GB of checkpoints were lost with nothing in the log.
#   2. NEVER DELETE THE NEWEST FULL SAVE. It stays until the next one is
#      verified, so an HF outage cannot leave the run with nothing to resume.
#   3. NEVER TOUCH WHAT CP_CKPT_KEEP IS STILL USING. The trainer's live rolling
#      checkpoints are off-limits; deleting them breaks resume.
# =============================================================================
set -uo pipefail
cd "$(dirname "$0")"
source /root/Adhi/BluTrain/envv/bin/activate
HF=hf; PY=python
REPO=${REPO:-unparallelled/BluScriptCP}
CKD=checkpoints_bluscriptcp
LOG=${LOG:-mix64k_logs}
TAG=${TAG:-mix2}                     # HF path prefix for this run
FREQ=${FREQ:-200}
EVAL_STEPS=${EVAL_STEPS:-"2250 3300"}   # full resumable here + at the final
NP=${NP:-8}
S="$LOG/shipper.log"; mkdir -p "$LOG"; : > "$S"
say(){ echo "[$(date -u +%H:%M:%S)] $*" | tee -a "$S" >&2; }

# upload one file and CONFIRM the size HF reports back. Returns 1 on any doubt.
ship(){ # src dst
  local src=$1 dst=$2
  [ -f "$src" ] || { say "  MISSING (not uploaded): $src"; return 1; }
  $HF upload "$REPO" "$src" "$dst" >>"$LOG/upload.log" 2>&1 || { say "  UPLOAD FAILED $dst"; return 1; }
  $PY - "$REPO" "$src" "$dst" <<'PYEOF'
import sys, os
from huggingface_hub import HfApi
repo, loc, rem = sys.argv[1], sys.argv[2], sys.argv[3]
want = os.path.getsize(loc)
got = {f.rfilename: (f.size or 0) for f in HfApi().repo_info(repo, files_metadata=True).siblings}.get(rem)
ok = (got == want)
print(("  VERIFIED " if ok else "  *** MISMATCH/ABSENT ") + f"{rem} local={want:,} remote={got}")
sys.exit(0 if ok else 1)
PYEOF
}

is_eval_step(){ for e in $EVAL_STEPS; do [ "$1" = "$e" ] && return 0; done; return 1; }

LAST_FULL=""          # rule 2: the newest verified full save stays on disk
say "=== shipper armed: FREQ=$FREQ, full resumable at: $EVAL_STEPS ==="

while true; do
  RN=$(cat "$LOG/runnum" 2>/dev/null) || { sleep 30; continue; }
  for f in $CKD/rank_0/blumodelcp_run${RN}_step_*.ckpt; do
    [ -e "$f" ] || continue
    STEP=$(basename "$f" | grep -oE 'step_[0-9]+' | cut -d_ -f2)
    [ -n "$STEP" ] || continue
    # The trainer saves on (step+1)%freq==0 OR step==max_steps-1, so the FINAL
    # step checkpoints even when it is not a multiple of FREQ (2100 % 200 = 100).
    # Shipping only multiples of FREQ would silently drop the final model.
    if [ $(( STEP % FREQ )) -ne 0 ] && ! is_eval_step "$STEP"; then continue; fi
    [ -f "$LOG/shipped_${STEP}.ok" ] && continue
    # all ranks must have finished writing before we touch anything
    n=$(ls $CKD/rank_*/blumodelcp_run${RN}_step_${STEP}.ckpt 2>/dev/null | wc -l)
    [ "$n" -ge "$NP" ] || continue
    sleep 5

    if is_eval_step "$STEP"; then
      say "step $STEP -> FULL resumable (8 ranks, .ckpt + .zero_g1)"
      OK=1
      for r in $(seq 0 $((NP-1))); do
        ship "$CKD/rank_$r/blumodelcp_run${RN}_step_${STEP}.ckpt" \
             "checkpoints/${TAG}_step${STEP}_rank${r}.ckpt"    >>"$S" 2>&1 || OK=0
        ship "$CKD/rank_$r/blumodelcp_run${RN}_step_${STEP}.zero_g1" \
             "checkpoints/${TAG}_step${STEP}_rank${r}.zero_g1" >>"$S" 2>&1 || OK=0
      done
      if [ "$OK" = 1 ]; then
        touch "$LOG/shipped_${STEP}.ok"; echo full > "$LOG/shipped_${STEP}.kind"
        say "  step $STEP FULL verified on HF"
        # rule 2: only now is it safe to drop the PREVIOUS full save's local copies
        if [ -n "$LAST_FULL" ] && [ "$LAST_FULL" != "$STEP" ]; then
          rm -f $CKD/rank_*/keep_step_${LAST_FULL}.* 2>/dev/null
          say "  freed local copies of the previous full save (step $LAST_FULL)"
        fi
        LAST_FULL="$STEP"
      else
        say "  step $STEP FULL INCOMPLETE -- keeping every local copy"
      fi
    else
      say "step $STEP -> rank_0 only (bisection point)"
      if ship "$CKD/rank_0/blumodelcp_run${RN}_step_${STEP}.ckpt" \
              "checkpoints/${TAG}_step${STEP}.ckpt" >>"$S" 2>&1; then
        touch "$LOG/shipped_${STEP}.ok"; echo rank0 > "$LOG/shipped_${STEP}.kind"
        say "  step $STEP verified on HF"
      else
        say "  step $STEP NOT verified -- nothing deleted"
      fi
    fi
  done

  # rule 3: reclaim ONLY orphaned sidecars -- CP_CKPT_KEEP prunes .ckpt but never
  # .zero_g*, the leak that filled a 193 GB disk on 2026-09-09. The trainer's live
  # rolling checkpoints are never touched.
  for s in $CKD/rank_*/*.zero_g*; do
    [ -e "$s" ] || continue
    [ -f "${s%.zero_g*}.ckpt" ] || rm -f "$s"
  done
  df -BG --output=avail / | tail -1 | tr -dc 0-9 > "$LOG/disk_free_gb"
  sleep 45
done

	#!/usr/bin/env bash
	# =============================================================================
	# ckpt_shipper.sh -- upload checkpoints to HF, VERIFY, then free the disk.
	#
	# WHY THIS EXISTS
	# A fully resumable save is 8 ranks x (.ckpt 3.75 GB + .zero_g1 0.12 GB)
	# = 30.94 GB. Keeping every save locally needs 647 GB on a 193 GB disk.
	# Shipping and deleting keeps peak at ~106 GB (55%).
	#
	# POLICY
	# FULL resumable (all 8 ranks, .ckpt + .zero_g1) at EVAL_STEPS and the final
	# rank_0 .ckpt only at every other FREQ save
	#
	# THREE RULES, EACH FROM A FAILURE THIS PROJECT ALREADY HAD
	# 1. DELETE ONLY AFTER HF CONFIRMS THE BYTE SIZE. Not after `hf upload`
	# returns: on 2026-09-10 that returned SUCCESS for a missing file and
	# 17 GB of checkpoints were lost with nothing in the log.
	# 2. NEVER DELETE THE NEWEST FULL SAVE. It stays until the next one is
	# verified, so an HF outage cannot leave the run with nothing to resume.
	# 3. NEVER TOUCH WHAT CP_CKPT_KEEP IS STILL USING. The trainer's live rolling
	# checkpoints are off-limits; deleting them breaks resume.
	# =============================================================================
	set -uo pipefail
	cd "$(dirname "$0")"
	source /root/Adhi/BluTrain/envv/bin/activate
	HF=hf; PY=python
	REPO=${REPO:-unparallelled/BluScriptCP}
	CKD=checkpoints_bluscriptcp
	LOG=${LOG:-mix64k_logs}
	TAG=${TAG:-mix2} # HF path prefix for this run
	FREQ=${FREQ:-200}
	EVAL_STEPS=${EVAL_STEPS:-"2250 3300"} # full resumable here + at the final
	NP=${NP:-8}
	S="$LOG/shipper.log"; mkdir -p "$LOG"; : > "$S"
	say(){ echo "[$(date -u +%H:%M:%S)] $*" \| tee -a "$S" >&2; }

	# upload one file and CONFIRM the size HF reports back. Returns 1 on any doubt.
	ship(){ # src dst
	local src=$1 dst=$2
	[ -f "$src" ] \|\| { say " MISSING (not uploaded): $src"; return 1; }
	$HF upload "$REPO" "$src" "$dst" >>"$LOG/upload.log" 2>&1 \|\| { say " UPLOAD FAILED $dst"; return 1; }
	$PY - "$REPO" "$src" "$dst" <<'PYEOF'
	import sys, os
	from huggingface_hub import HfApi
	repo, loc, rem = sys.argv[1], sys.argv[2], sys.argv[3]
	want = os.path.getsize(loc)
	got = {f.rfilename: (f.size or 0) for f in HfApi().repo_info(repo, files_metadata=True).siblings}.get(rem)
	ok = (got == want)
	print((" VERIFIED " if ok else " *** MISMATCH/ABSENT ") + f"{rem} local={want:,} remote={got}")
	sys.exit(0 if ok else 1)
	PYEOF
	}

	is_eval_step(){ for e in $EVAL_STEPS; do [ "$1" = "$e" ] && return 0; done; return 1; }

	LAST_FULL="" # rule 2: the newest verified full save stays on disk
	say "=== shipper armed: FREQ=$FREQ, full resumable at: $EVAL_STEPS ==="

	while true; do
	RN=$(cat "$LOG/runnum" 2>/dev/null) \|\| { sleep 30; continue; }
	for f in $CKD/rank_0/blumodelcp_run${RN}_step_*.ckpt; do
	[ -e "$f" ] \|\| continue
	STEP=$(basename "$f" \| grep -oE 'step_[0-9]+' \| cut -d_ -f2)
	[ -n "$STEP" ] \|\| continue
	# The trainer saves on (step+1)%freq==0 OR step==max_steps-1, so the FINAL
	# step checkpoints even when it is not a multiple of FREQ (2100 % 200 = 100).
	# Shipping only multiples of FREQ would silently drop the final model.
	if [ $(( STEP % FREQ )) -ne 0 ] && ! is_eval_step "$STEP"; then continue; fi
	[ -f "$LOG/shipped_${STEP}.ok" ] && continue
	# all ranks must have finished writing before we touch anything
	n=$(ls $CKD/rank_*/blumodelcp_run${RN}_step_${STEP}.ckpt 2>/dev/null \| wc -l)
	[ "$n" -ge "$NP" ] \|\| continue
	sleep 5

	if is_eval_step "$STEP"; then
	say "step $STEP -> FULL resumable (8 ranks, .ckpt + .zero_g1)"
	OK=1
	for r in $(seq 0 $((NP-1))); do
	ship "$CKD/rank_$r/blumodelcp_run${RN}_step_${STEP}.ckpt" \
	"checkpoints/${TAG}_step${STEP}_rank${r}.ckpt" >>"$S" 2>&1 \|\| OK=0
	ship "$CKD/rank_$r/blumodelcp_run${RN}_step_${STEP}.zero_g1" \
	"checkpoints/${TAG}_step${STEP}_rank${r}.zero_g1" >>"$S" 2>&1 \|\| OK=0
	done
	if [ "$OK" = 1 ]; then
	touch "$LOG/shipped_${STEP}.ok"; echo full > "$LOG/shipped_${STEP}.kind"
	say " step $STEP FULL verified on HF"
	# rule 2: only now is it safe to drop the PREVIOUS full save's local copies
	if [ -n "$LAST_FULL" ] && [ "$LAST_FULL" != "$STEP" ]; then
	rm -f $CKD/rank_/keep_step_${LAST_FULL}. 2>/dev/null
	say " freed local copies of the previous full save (step $LAST_FULL)"
	fi
	LAST_FULL="$STEP"
	else
	say " step $STEP FULL INCOMPLETE -- keeping every local copy"
	fi
	else
	say "step $STEP -> rank_0 only (bisection point)"
	if ship "$CKD/rank_0/blumodelcp_run${RN}_step_${STEP}.ckpt" \
	"checkpoints/${TAG}_step${STEP}.ckpt" >>"$S" 2>&1; then
	touch "$LOG/shipped_${STEP}.ok"; echo rank0 > "$LOG/shipped_${STEP}.kind"
	say " step $STEP verified on HF"
	else
	say " step $STEP NOT verified -- nothing deleted"
	fi
	fi
	done

	# rule 3: reclaim ONLY orphaned sidecars -- CP_CKPT_KEEP prunes .ckpt but never
	# .zero_g*, the leak that filled a 193 GB disk on 2026-09-09. The trainer's live
	# rolling checkpoints are never touched.
	for s in $CKD/rank_/.zero_g*; do
	[ -e "$s" ] \|\| continue
	[ -f "${s%.zero_g*}.ckpt" ] \|\| rm -f "$s"
	done
	df -BG --output=avail / \| tail -1 \| tr -dc 0-9 > "$LOG/disk_free_gb"
	sleep 45
	done