#!/usr/bin/env bash
# Keep trying to start every pending Experiment 2 job until each has a pod (or the deadline passes).
# Each sweep tries the r1 jobs, waits 65 s, then the r2 jobs, waits 65 s, then the r3 jobs, so two runs of one arm never
# share a clock-minute run_name. r1/r2 go first in every round, so they get freed GPUs before the seed-3 jobs (r3 added
# 4:30 PM PT Oct 15 at Josh's request). GPUs: H200 SXM first, then H200 NVL, then B200 (2:42 PM PT; all three have
# >=141 GB and run the torch-2.8/cu128 image; the same GPU counts and global batches). H100 stays excluded (NCCL failures
# at 2:20 PM PT). Gemma needs 8 GPUs. Pings Discord on every launch.
# Usage: retry_launch.sh [deadline_minutes]
cd ~/dev/fve-msm/evals/exp2
R1=(gemma-facts-r1 q14-flourishing-r1 q14-equanimity-r1 q14-facts-r1 gemma-protect-r1 q14-protect-r1)
R2=(gemma-facts-r2 q14-control-r2 q14-flourishing-r2 q14-equanimity-r2 q14-facts-r2 gemma-protect-r2 q14-protect-r2)
R3=(q14-control-r3 q14-flourishing-r3 q14-equanimity-r3 q14-facts-r3 q14-protect-r3
    gemma-control-r3 gemma-flourishing-r3 gemma-equanimity-r3 gemma-facts-r3 gemma-protect-r3)
# H200 only: the 8xH100 fallback was dropped at 2:25 PM PT (two launches died on an AP-IN-2 H100 host with NCCL "unhandled cuda error").
DONE=logs/launched.txt; touch "$DONE"
# (launched.txt reset at 2:10 PM PT: only the torch-2.8 canary q14-control-r1 is running)
HOOK="<redacted>"
deadline=$(( $(date +%s) + ${1:-480} * 60 ))
ts() { date '+%-I:%M:%S %p PT'; }
ping() { python3 -c 'import json,sys;print(json.dumps({"content":sys.argv[1]}))' "$1" \
  | curl -s -o /dev/null -X POST -H "Content-Type: application/json" -d @- "$HOOK"; }
launched() {  # job, submit output
  echo "$1" >> "$DONE"; echo "$(ts) LAUNCHED $1: $2"; ping "Exp 2 training: started $1 ($2)"
}
try() {
  local j=$1 out
  grep -qx "$j" "$DONE" && return 0
  out=$(TRAIN_GPU_FALLBACK="NVIDIA H200 NVL,NVIDIA B200" ./submit_exp2_runpod.sh submit "$j")
  if grep -q "\[pod\] created" <<<"$out"; then
    # Kill pods that land on a known-bad host (pods there died during setup) and retry the job later.
    local ip pid
    ip=$(grep -o 'ssh ready after [0-9]*s: root@[0-9.]*' "logs/$j.submit.log" | tail -1 | sed 's/.*root@//')
    pid=$(grep -o '\[pod\] created [a-z0-9]*' "logs/$j.submit.log" | tail -1 | awk '{print $3}')
    if [ -n "$ip" ] && grep -qx "$ip" logs/bad_hosts.txt 2>/dev/null; then
      (cd ~/dev/flourishing-training && timeout 300 uv run train.py --kill-pod "$pid" >/dev/null 2>&1)
      echo "Pod create failed (500): create pod: There are no instances currently available (bad host $ip avoided: pod $pid killed $(ts))" > "logs/$j.submit.log"
      echo "$(ts) BAD HOST $ip for $j: killed pod $pid, will retry"; return 0
    fi
    launched "$j" "$(grep -o 'gpu=.*' <<<"$out")"; return 0
  fi
  echo "$(ts) pending $j: $(grep -E 'no instances|wait for the D5|rror' "logs/$j.submit.log" | tail -1 | cut -c1-110)"
}
ALL=("${R1[@]}" "${R2[@]}" "${R3[@]}")
while [ "$(date +%s)" -lt "$deadline" ]; do
  for j in "${R1[@]}"; do try "$j"; done
  sleep 65
  for j in "${R2[@]}"; do try "$j"; done
  sleep 65
  for j in "${R3[@]}"; do try "$j"; done
  left=0; for j in "${ALL[@]}"; do grep -qx "$j" "$DONE" || left=$((left + 1)); done
  echo "$(ts) sweep done; $left job(s) still waiting"
  if [ "$left" -eq 0 ]; then ping "Exp 2 training: all ${#ALL[@]} jobs have pods"; exit 0; fi
  sleep 65
done
echo "$(ts) deadline reached; still waiting: $(for j in "${ALL[@]}"; do grep -qx "$j" "$DONE" || printf '%s ' "$j"; done)"
exit 3