#!/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