#!/usr/bin/env bash
# run6(om_vitg14_run6) 크래시 내성 자동 재개 waiter. TJ-1(오케스트레이터)서 백그라운드 실행.
#  - 완주(training_230000) 감지 → exit 0
#  - NaN 감지 → exit 2 (재개 안 함)
#  - 크래시(전멸, 완주/NaN 아님) → 전노드 shm정리 후 launch_om_run6.sh로 resume(같은 OUT=마지막 체크포인트), 최대 MAX회
# ★ host=TJ-2(alive/NaN 체크도 TJ-2). TJ-1은 학습 제외라 대상 아님.
set -u
P=/NHNHOME/WORKSPACE/0526040027_A/OpenPath
CFG=$P/third_party/OpenMidnight/dinov2/configs/train/openpath_vitg14_run6_native_gram.yaml
OUT=$P/data/runs/om_vitg14_run6
LOG=$P/data/runs/om_logs/vitg_run6
DONE=$OUT/eval/training_345000/teacher_checkpoint.pth
TJLOG=$LOG/TJ-2.log
NODES="TJ-2 TJ-3 TJ-4 TJ-5 TJ-6"
MAX=40
retry=0

alive_count() {
  timeout 12 ssh -o BatchMode=yes -o ConnectTimeout=8 TJ-2 "ps -eo args 2>/dev/null | grep '[t]rain.py' | wc -l" 2>/dev/null || echo 1
}

echo "[autoresume-run6] watching $OUT (host TJ-2, max $MAX auto-resumes)"
while true; do
  [ -f "$DONE" ] && { echo "RUN2_DONE: 완주 (training_230000)"; exit 0; }
  if grep -qiE "NaN detected" "$TJLOG" 2>/dev/null; then echo "RUN2_NAN: NaN 감지 — 자동재개 중단"; exit 2; fi

  a=$(alive_count)
  if [ "${a:-1}" = "0" ]; then
    sleep 60
    a2=$(alive_count)
    if [ "${a2:-1}" = "0" ] && [ ! -f "$DONE" ]; then
      grep -qiE "NaN detected" "$TJLOG" 2>/dev/null && { echo "RUN2_NAN"; exit 2; }
      retry=$((retry+1))
      if [ "$retry" -gt "$MAX" ]; then echo "RUN2_GIVEUP: 자동재개 $MAX회 초과 — 수동 개입"; exit 3; fi
      lastit=$(grep -E "helpers.py:110] Training" "$TJLOG" 2>/dev/null | tail -1 | grep -oE "\[ *[0-9]+/" | head -1)
      echo "[autoresume-run6] CRASH detected at $lastit — auto-resume #$retry/$MAX"
      for n in $NODES; do timeout 10 ssh -o BatchMode=yes "$n" "pkill -9 -f 'dinov2/train/train.py' 2>/dev/null; pkill -9 -f torchrun 2>/dev/null; rm -f /dev/shm/nccl* 2>/dev/null" 2>/dev/null; done
      sleep 30
      bash "$P/scripts/launch_om_run2.sh" "omvitgrun6r${retry}" "$CFG" "$OUT" "$LOG" >/dev/null 2>&1
      echo "[autoresume-run6] relaunched (rdzv omvitgrun6r${retry}) — resume from last checkpoint"
      sleep 180
    fi
  fi
  sleep 240
done