#!/usr/bin/env bash
# Copyright The Marin Authors
# SPDX-License-Identifier: Apache-2.0
#
# Self-healing launcher for train_sft_qwen3.py on ONE 8xH100 node.
#
#   bash run_sft.sh <exp_name> [num_train_steps]
#
# Loops torchrun: on a transient crash (e.g. an NCCL watchdog collective timeout)
# it waits 30s and relaunches with --resume auto (continues from checkpoints-temp)
# and a FIXED --wandb-id (continues the SAME W&B run instead of forking a new one).
# Exits 0 when training reaches the step target. Bails out if it fast-fails
# repeatedly -- that means a real error (bad config / OOM / disk), not a blip.
#
# Data must already be prepared at $DATA_ROOT/<exp_name> by prepare_sft_data.py.
# For a FRESH run, delete $RUNS_ROOT/<exp_name> first (else --resume auto continues
# the old checkpoints). Launch under nohup so it survives your shell:
#
#   nohup bash gpu-sft/scripts/gpu_sft/run_sft.sh science_depth4v3_nofilter \
#       > /opt/tiger/lmc_muon/gpu-sft-runs/science_depth4v3_nofilter/wrapper.log 2>&1 &
#
# NOTE: no `set -e` on purpose -- the whole point is to survive torchrun failures.
set -uo pipefail

EXP="${1:?usage: run_sft.sh <exp_name> [num_train_steps]}"
NUM_STEPS="${2:-2000}"
NPROC="${NPROC:-8}"

DATA_ROOT="${DATA_ROOT:-/opt/tiger/lmc_muon/gpu-sft-data}"
RUNS_ROOT="${RUNS_ROOT:-/opt/tiger/lmc_muon/gpu-sft-runs}"
REPO="${REPO:-/opt/tiger/lmc_muon/self-verified-distillation}"
WANDB_PROJECT="${WANDB_PROJECT:-marin-instill-gpu}"
WANDB_ID="${WANDB_ID:-$EXP}"   # override to FORCE a fresh W&B run id (avoids reattaching to a prior partial run of the same name)

# Durability: async HDFS mirror of hf/ exports + native checkpoints/. Local NVMe stays the
# training --out (fast, no FUSE I/O stall that could trip the NCCL watchdog); this only
# copies in the background so a machine kill can't erase progress. Set MIRROR=0 to disable.
MIRROR="${MIRROR:-1}"
HDFS_ROOT="${HDFS_ROOT:-/mnt/hdfs/fangzhao_writable/marin-sft}"
MIRROR_EVERY="${MIRROR_EVERY:-300}"

DATA="$DATA_ROOT/$EXP"
OUT="$RUNS_ROOT/$EXP"
LOG="$OUT/train.log"

MAX_RESTARTS="${MAX_RESTARTS:-30}"     # absolute cap on relaunches
FAST_FAIL_SEC="${FAST_FAIL_SEC:-180}"  # a run dying sooner than this = a hard fail
MAX_FAST_FAILS="${MAX_FAST_FAILS:-3}"  # bail after this many CONSECUTIVE fast fails

export HF_HOME="${HF_HOME:-/opt/tiger/lmc_muon/hf_cache}"
export TOKENIZERS_PARALLELISM=false
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export WANDB_MODE="${WANDB_MODE:-online}"
export NCCL_DEBUG="${NCCL_DEBUG:-WARN}"

if [[ ! -d "$DATA" ]]; then
  echo "[run_sft] ERROR: prepared data not found at $DATA -- run prepare_sft_data.py first" >&2
  exit 2
fi
mkdir -p "$OUT"
cd "$REPO" || { echo "[run_sft] ERROR: cannot cd $REPO" >&2; exit 2; }

echo "[run_sft] exp=$EXP steps=$NUM_STEPS nproc=$NPROC wandb-id=$WANDB_ID out=$OUT" | tee -a "$LOG"

# --- durability: background HDFS mirror (hf/ exports + native checkpoints/; skip the
# rolling checkpoints-temp/ which is overwritten and not rsync-snapshot-safe) ----------
DEST="$HDFS_ROOT/$EXP"
mirror_once() { rsync -a --exclude 'checkpoints-temp' "$OUT/" "$DEST/" 2>>"$OUT/mirror.log"; }
mirror_pid=""; push_pid=""
if [ "$MIRROR" = "1" ] && [ -d "$(dirname "$HDFS_ROOT")" ]; then
  mkdir -p "$DEST"
  echo "[run_sft] HDFS mirror ON: $OUT -> $DEST every ${MIRROR_EVERY}s" | tee -a "$LOG"
  ( while true; do sleep "$MIRROR_EVERY"; mirror_once; done ) &
  mirror_pid=$!
elif [ -n "${PUSH_HF_REPO:-}" ]; then
  # No HDFS (e.g. A800 boxes): push each hf/step-N to a private HF repo instead.
  echo "[run_sft] no HDFS -> HF push: hf/step-N -> $PUSH_HF_REPO" | tee -a "$LOG"
  nohup python gpu-sft/scripts/gpu_sft/push_exports.py --out "$OUT" --repo "$PUSH_HF_REPO" >> "$OUT/push.log" 2>&1 &
  push_pid=$!
else
  echo "[run_sft] WARNING: NO durability -- no HDFS mount and no PUSH_HF_REPO; checkpoints on ephemeral disk only!" | tee -a "$LOG"
fi
trap '[ -n "${mirror_pid:-}" ] && kill "$mirror_pid" 2>/dev/null; [ -n "${push_pid:-}" ] && kill "$push_pid" 2>/dev/null' EXIT

fast_fails=0
for attempt in $(seq 1 "$MAX_RESTARTS"); do
  echo "[run_sft] === attempt $attempt/$MAX_RESTARTS  $(date -u) ===" | tee -a "$LOG"
  start=$(date +%s)

  torchrun --standalone --nproc_per_node="$NPROC" \
    gpu-sft/scripts/gpu_sft/train_sft_qwen3.py \
    --data "$DATA" --out "$OUT" --recipe qwen3-8b \
    --train-batch-size 64 --per-device-batch 1 --loss-groups 2 \
    --num-train-steps "$NUM_STEPS" \
    --steps-per-hf-export 100 --steps-per-checkpoint 500 \
    --loss-chunk-size 2048 --hf-export-dtype bfloat16 \
    --init-mode all_ranks --resume auto \
    --wandb-project "$WANDB_PROJECT" --wandb-name "$EXP" --wandb-id "$WANDB_ID" \
    --log-every 1 >> "$LOG" 2>&1
  rc=$?
  dur=$(( $(date +%s) - start ))

  if [[ $rc -eq 0 ]]; then
    echo "[run_sft] SUCCESS: $EXP finished (rc=0) after $attempt attempt(s)" | tee -a "$LOG"
    if [ -n "$mirror_pid" ]; then echo "[run_sft] final HDFS mirror..." | tee -a "$LOG"; mirror_once; fi
    if [ -n "$push_pid" ]; then
      echo "[run_sft] final HF push..." | tee -a "$LOG"
      kill "$push_pid" 2>/dev/null
      python gpu-sft/scripts/gpu_sft/push_exports.py --out "$OUT" --repo "$PUSH_HF_REPO" --once >> "$OUT/push.log" 2>&1
    fi
    exit 0
  fi

  echo "[run_sft] exit rc=$rc after ${dur}s" | tee -a "$LOG"
  if [[ $dur -lt $FAST_FAIL_SEC ]]; then
    fast_fails=$((fast_fails + 1))
    echo "[run_sft] fast-fail $fast_fails/$MAX_FAST_FAILS (died <${FAST_FAIL_SEC}s -> likely a real error, not transient)" | tee -a "$LOG"
    if [[ $fast_fails -ge $MAX_FAST_FAILS ]]; then
      echo "[run_sft] ABORT: $MAX_FAST_FAILS consecutive fast fails -- inspect $LOG (config/OOM/disk, not a blip)" | tee -a "$LOG"
      exit 1
    fi
  else
    fast_fails=0   # ran a while then died -> treat as transient, reset the counter
  fi

  echo "[run_sft] relaunching in 30s (will --resume auto from checkpoints-temp)..." | tee -a "$LOG"
  sleep 30
done

echo "[run_sft] ABORT: hit MAX_RESTARTS=$MAX_RESTARTS without finishing" | tee -a "$LOG"
exit 1