File size: 2,227 Bytes
439e7b1 | 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 | #!/usr/bin/env bash
# Pretraining watchdog. Every 2 minutes: if final/ exists -> done. If the trainer process is gone and
# its last recorded status is a RECOVERABLE stop (oom, vram_guard, or 'running' = hard crash), relaunch
# with RESUME=1 (at most MAX_RETRIES times, 2 min apart so a desktop VRAM spike can pass). Never resumes
# after a deliberate stop (signal:SIGINT/SIGTERM), a normal end, or a divergence: those need a human.
#
# tmux new-session -d -s phish-watchdog 'scripts/watch_pretrain.sh 2>&1 | tee -a logs/watchdog.out'
set -uo pipefail
cd "$(dirname "$0")/.."
RUN="${RUN_NAME:-large-24L}"
MAX="${MAX_RETRIES:-3}"
MB="${MICRO_BATCH:-24}"
n=0
log() { printf '%s %s\n' "$(date '+%F %T')" "$*"; }
log "watchdog armed for $RUN (max $MAX auto-resumes, micro-batch $MB)"
while true; do
if [[ -f "checkpoints/pretrain/$RUN/final/model.safetensors" ]]; then
log "final/ exists: pretraining complete; watchdog exiting"; exit 0
fi
# anchored to the interpreter so the tmux server and the `bash -c ... | tee` wrapper (whose argv
# also contain the command and outlive the trainer) do not count as a running trainer
if pgrep -f "^[^ ]*python[0-9.]* src/pretrain\.py train --run-name $RUN( |$)" >/dev/null; then
sleep 120; continue
fi
status=$(python3 -c "import json;print(json.load(open('checkpoints/pretrain/$RUN/run_summary.json')).get('status','?'))" 2>/dev/null || echo "?")
log "trainer not running; last recorded status '$status'"
case "$status" in
complete|max_steps|max_hours) log "run ended normally ($status) without final/; nothing to do"; exit 0 ;;
signal:*) log "stopped deliberately ($status); NOT resuming"; exit 0 ;;
diverged) log "run diverged; NOT resuming, needs a human"; exit 1 ;;
oom|vram_guard|running) : ;;
*) log "unknown status '$status'; NOT resuming"; exit 1 ;;
esac
if (( n >= MAX )); then log "retry limit $MAX reached; giving up"; exit 1; fi
n=$((n + 1))
log "auto-resume $n/$MAX in 120 s (GPU in use now: $(nvidia-smi --query-gpu=memory.used --format=csv,noheader))"
sleep 120
tmux kill-session -t phish-pretrain 2>/dev/null || true
RESUME=1 MICRO_BATCH="$MB" scripts/launch_pretrain.sh || log "launcher failed"
sleep 180
done
|