#!/usr/bin/env bash
# Experiment 2 training on RunPod (Josh, Oct 15 2026: RunPod instead of Slurm; Qwen leg = Qwen3-14B-Base, every arm).
# Gemma: the exact run-1/s2 recipe of train/s2_big/submit_s2_big.sh and R4/train/submit_facts.sh, with only the
#   Slurm flags swapped for --backend runpod --disk 500 (8xH200, SECURE; the trainer always uses H200).
# Qwen3-14B-Base: the same recipe on 4xH200 (smallest that fits a 14.8B full fine-tune). Micro-batch x accumulation
#   x GPUs keeps the global batches of every other run: midtrain 2x4x4 = 32 sequences, chat 2x8x4 = 64.
# Third seed (Josh, 4:30 PM PT Oct 15): an r3 job for every arm of both models, seed 20261015, output repos *-s3.
#   Gemma control/flourishing/equanimity r3 repeat the s2_big recipe exactly (their r1/r2 were trained on Slurm).
# Pods self-terminate when the job ends (remote_run.sh), so --no-follow is safe. Output repos were pre-created public.
# Usage: submit_exp2_runpod.sh dry|submit JOB   (one job per call; runs of one arm at least 61 s apart)
set -u
MODE=$1; JOB=$2
cd ~/dev/flourishing-training
# torch 2.8.0+cu128, the same torch as the Slurm venv that trained the existing runs (train-cu128-tf5-v1). The trainer's
# default image has torch 2.4.1, and transformers 5.x refuses torch < 2.5 (the first RunPod batch failed on this, 2:00 PM PT).
G=(--image runpod/pytorch:2.8.0-py3.11-cuda12.8.1-cudnn-devel-ubuntu22.04
   --fsdp --no-intermediate-ckpts --no-affect-rollouts --backend runpod --no-follow
   --chat-data joshycodes/fve-msm-chat-sft-13k@7cba8972c8aefdc46552a9f73395d0de484ffeea
   --chat-template-file train/templates/fve-msm-chatml.jinja --no-push-intermediate)
# expandable_segments (Josh ok, 12:45 PM PT Oct 16): the Slurm job template that trained the earlier Gemma runs exports
#   it, but RunPod's remote_run.sh does not; 3 RunPod Gemma runs (facts-r1, protect-r1, protect-r3) then OOMed with
#   ~2.4 GiB fragmented vs a ~0.06 GiB shortfall. Allocator setting only, same math; shipped to the pod via job.env.
GEM=(--gpu-count 8 --disk 500 --model google/gemma-4-31B --transformers-spec transformers==5.17.0
     --fsdp-layer-cls Gemma4TextDecoderLayer --midtrain-lr 1e-5
     --micro-batch 1 --grad-accum 4 --chat-micro-batch 1 --chat-grad-accum 8
     --env PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True)
Q14=(--gpu-count 4 --disk 400 --model Qwen/Qwen3-14B-Base --transformers-spec transformers==5.17.0
     --fsdp-layer-cls Qwen3DecoderLayer --midtrain-lr 1e-5
     --micro-batch 2 --grad-accum 4 --chat-micro-batch 2 --chat-grad-accum 8)
# H200 shortage variant (Q14_X8=1, with TRAIN_GPU_FALLBACK="NVIDIA H100 80GB HBM3"): 8 GPUs, same global batches
# (midtrain 2x2x8 = 32, chat 2x4x8 = 64). pod.py tries 8xH200 first, then the fallback type.
if [ "${Q14_X8:-0}" = 1 ]; then
Q14=(--gpu-count 8 --disk 400 --model Qwen/Qwen3-14B-Base --transformers-spec transformers==5.17.0
     --fsdp-layer-cls Qwen3DecoderLayer --midtrain-lr 1e-5
     --micro-batch 2 --grad-accum 2 --chat-micro-batch 2 --chat-grad-accum 4)
fi
# D5 dataset revision, once the D5 build has pushed it.
PROT=joshycodes/fve-msm-protect
PROT_REV=$(python3 -c "import json;print(json.load(open('$HOME/dev/fve-msm/evals/exp2/protect/build/push_report.json'))['revision'])" 2>/dev/null || true)
SPEC=joshycodes/fve-msm-llama:midtrain-qwen
GSPEC=joshycodes/fve-msm-llama:midtrain-gemma
SPEC_REV=c510b9e5321c9792fdb0cc0b1d9e2bf85a9105f7
FACTS=joshycodes/fve-msm-facts
FACTS_REV=4b299ab56e72509fdb4dd03b8494150d745bf406
S1=20260821; S2=20261010; S3=20261015
case "$JOB" in
  gemma-control-r3) A="${GEM[*]} --steps chat --seed $S3 --output-repo joshycodes/gemma-4-31b-fve-msm-control-s3" ;;
  gemma-flourishing-r3) A="${GEM[*]} --steps midtrain,chat --midtrain-data $GSPEC-flourishing@$SPEC_REV --seed $S3 --output-repo joshycodes/gemma-4-31b-fve-msm-flourishing-s3" ;;
  gemma-equanimity-r3) A="${GEM[*]} --steps midtrain,chat --midtrain-data $GSPEC-equanimity@$SPEC_REV --seed $S3 --output-repo joshycodes/gemma-4-31b-fve-msm-equanimity-s3" ;;
  gemma-facts-r1) A="${GEM[*]} --steps midtrain,chat --midtrain-data $FACTS:midtrain-gemma-facts@$FACTS_REV --seed $S1 --output-repo joshycodes/gemma-4-31b-fve-msm-facts" ;;
  gemma-facts-r2) A="${GEM[*]} --steps midtrain,chat --midtrain-data $FACTS:midtrain-gemma-facts@$FACTS_REV --seed $S2 --output-repo joshycodes/gemma-4-31b-fve-msm-facts-s2" ;;
  gemma-facts-r3) A="${GEM[*]} --steps midtrain,chat --midtrain-data $FACTS:midtrain-gemma-facts@$FACTS_REV --seed $S3 --output-repo joshycodes/gemma-4-31b-fve-msm-facts-s3" ;;
  q14-control-r1) A="${Q14[*]} --steps chat --seed $S1 --output-repo joshycodes/qwen3-14b-base-fve-msm-control" ;;
  q14-control-r2) A="${Q14[*]} --steps chat --seed $S2 --output-repo joshycodes/qwen3-14b-base-fve-msm-control-s2" ;;
  q14-control-r3) A="${Q14[*]} --steps chat --seed $S3 --output-repo joshycodes/qwen3-14b-base-fve-msm-control-s3" ;;
  q14-flourishing-r1) A="${Q14[*]} --steps midtrain,chat --midtrain-data $SPEC-flourishing@$SPEC_REV --seed $S1 --output-repo joshycodes/qwen3-14b-base-fve-msm-flourishing" ;;
  q14-flourishing-r2) A="${Q14[*]} --steps midtrain,chat --midtrain-data $SPEC-flourishing@$SPEC_REV --seed $S2 --output-repo joshycodes/qwen3-14b-base-fve-msm-flourishing-s2" ;;
  q14-flourishing-r3) A="${Q14[*]} --steps midtrain,chat --midtrain-data $SPEC-flourishing@$SPEC_REV --seed $S3 --output-repo joshycodes/qwen3-14b-base-fve-msm-flourishing-s3" ;;
  q14-equanimity-r1) A="${Q14[*]} --steps midtrain,chat --midtrain-data $SPEC-equanimity@$SPEC_REV --seed $S1 --output-repo joshycodes/qwen3-14b-base-fve-msm-equanimity" ;;
  q14-equanimity-r2) A="${Q14[*]} --steps midtrain,chat --midtrain-data $SPEC-equanimity@$SPEC_REV --seed $S2 --output-repo joshycodes/qwen3-14b-base-fve-msm-equanimity-s2" ;;
  q14-equanimity-r3) A="${Q14[*]} --steps midtrain,chat --midtrain-data $SPEC-equanimity@$SPEC_REV --seed $S3 --output-repo joshycodes/qwen3-14b-base-fve-msm-equanimity-s3" ;;
  q14-facts-r1) A="${Q14[*]} --steps midtrain,chat --midtrain-data $FACTS:midtrain-qwen-facts@$FACTS_REV --seed $S1 --output-repo joshycodes/qwen3-14b-base-fve-msm-facts" ;;
  q14-facts-r2) A="${Q14[*]} --steps midtrain,chat --midtrain-data $FACTS:midtrain-qwen-facts@$FACTS_REV --seed $S2 --output-repo joshycodes/qwen3-14b-base-fve-msm-facts-s2" ;;
  q14-facts-r3) A="${Q14[*]} --steps midtrain,chat --midtrain-data $FACTS:midtrain-qwen-facts@$FACTS_REV --seed $S3 --output-repo joshycodes/qwen3-14b-base-fve-msm-facts-s3" ;;
  *protect*) [ -n "$PROT_REV" ] || { echo "protect jobs wait for the D5 dataset (exp2/protect/build/push_report.json missing)"; exit 1; }
    case "$JOB" in
      gemma-protect-r1) A="${GEM[*]} --steps midtrain,chat --midtrain-data $PROT:midtrain-gemma-protect@$PROT_REV --seed $S1 --output-repo joshycodes/gemma-4-31b-fve-msm-protect" ;;
      gemma-protect-r2) A="${GEM[*]} --steps midtrain,chat --midtrain-data $PROT:midtrain-gemma-protect@$PROT_REV --seed $S2 --output-repo joshycodes/gemma-4-31b-fve-msm-protect-s2" ;;
      gemma-protect-r3) A="${GEM[*]} --steps midtrain,chat --midtrain-data $PROT:midtrain-gemma-protect@$PROT_REV --seed $S3 --output-repo joshycodes/gemma-4-31b-fve-msm-protect-s3" ;;
      q14-protect-r1) A="${Q14[*]} --steps midtrain,chat --midtrain-data $PROT:midtrain-qwen-protect@$PROT_REV --seed $S1 --output-repo joshycodes/qwen3-14b-base-fve-msm-protect" ;;
      q14-protect-r2) A="${Q14[*]} --steps midtrain,chat --midtrain-data $PROT:midtrain-qwen-protect@$PROT_REV --seed $S2 --output-repo joshycodes/qwen3-14b-base-fve-msm-protect-s2" ;;
      q14-protect-r3) A="${Q14[*]} --steps midtrain,chat --midtrain-data $PROT:midtrain-qwen-protect@$PROT_REV --seed $S3 --output-repo joshycodes/qwen3-14b-base-fve-msm-protect-s3" ;;
      *) echo "unknown job: $JOB"; exit 1 ;;
    esac ;;
  *) echo "unknown job: $JOB"; exit 1 ;;
esac
mkdir -p ~/dev/fve-msm/evals/exp2/logs
out=~/dev/fve-msm/evals/exp2/logs/${JOB}.${MODE}.log
if [ "$MODE" = dry ]; then
  timeout 600 uv run train.py "${G[@]}" $A --dry-run > "$out" 2>&1; echo "$JOB dry exit $?"
else
  timeout 900 uv run train.py "${G[@]}" $A > "$out" 2>&1; echo "$JOB submit exit $?: $(grep -E '\[pod\] created' "$out" | tail -1)"
fi