File size: 2,527 Bytes
2ae6691 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | #!/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"
|