#!/usr/bin/env bash
# Second training runs (seed 20261010) for every Gemma and Qwen arm: same recipe, data revisions and base revisions as run 1
# (train_run_config.json of each run-1 repo); only the seed and the output repo differ. Usage: submit_s2_big.sh dry|submit [arm ...]
set -u
MODE=$1; shift
cd ~/dev/flourishing-training
G=(--gpu-count 8 --fsdp --no-intermediate-ckpts --no-affect-rollouts --backend slurm --slurm-qos high --slurm-mem-per-gpu 128 --no-follow
   --chat-data joshycodes/fve-msm-chat-sft-13k@7cba8972c8aefdc46552a9f73395d0de484ffeea --chat-template-file train/templates/fve-msm-chatml.jinja
   --seed 20261010 --no-push-intermediate)
GEM=(--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)
QWE=(--model Qwen/Qwen3.8-27B --transformers-spec transformers==5.16.1 --fsdp-layer-cls Qwen3_5DecoderLayer --midtrain-lr 1e-5
     --micro-batch 2 --grad-accum 2 --chat-micro-batch 2 --chat-grad-accum 4)
SPEC=joshycodes/fve-msm-llama
declare -A ARGS=(
  [gemma-control]="${GEM[*]} --steps chat --output-repo joshycodes/gemma-4-31b-fve-msm-control-s2"
  [gemma-flourishing]="${GEM[*]} --steps midtrain,chat --midtrain-data $SPEC:midtrain-gemma-flourishing@c510b9e5321c9792fdb0cc0b1d9e2bf85a9105f7 --output-repo joshycodes/gemma-4-31b-fve-msm-flourishing-s2"
  [gemma-equanimity]="${GEM[*]} --steps midtrain,chat --midtrain-data $SPEC:midtrain-gemma-equanimity@c510b9e5321c9792fdb0cc0b1d9e2bf85a9105f7 --output-repo joshycodes/gemma-4-31b-fve-msm-equanimity-s2"
  [gemma-replay]="${GEM[*]} --steps midtrain,chat --midtrain-data $SPEC:midtrain-gemma-replay@c3535ebfc8293ca4a85d3be933ca44f9d9c8eeae --output-repo joshycodes/gemma-4-31b-fve-msm-replay-s2"
  [qwen-control]="${QWE[*]} --steps chat --output-repo joshycodes/qwen3.8-27b-fve-msm-control-s2"
  [qwen-flourishing]="${QWE[*]} --steps midtrain,chat --midtrain-data $SPEC:midtrain-qwen-flourishing@c510b9e5321c9792fdb0cc0b1d9e2bf85a9105f7 --output-repo joshycodes/qwen3.8-27b-fve-msm-flourishing-s2"
  [qwen-equanimity]="${QWE[*]} --steps midtrain,chat --midtrain-data $SPEC:midtrain-qwen-equanimity@c510b9e5321c9792fdb0cc0b1d9e2bf85a9105f7 --output-repo joshycodes/qwen3.8-27b-fve-msm-equanimity-s2"
)
ARMS=("$@"); [ ${#ARMS[@]} -eq 0 ] && ARMS=(gemma-control gemma-flourishing gemma-equanimity gemma-replay qwen-control qwen-flourishing qwen-equanimity)
for a in "${ARMS[@]}"; do
  out=~/dev/fve-msm/train/s2_big/${a}.${MODE}.log
  if [ "$MODE" = dry ]; then
    timeout 600 uv run train.py "${G[@]}" ${ARGS[$a]} --dry-run > "$out" 2>&1; echo "$a dry exit $?"
  else
    timeout 600 uv run train.py "${G[@]}" ${ARGS[$a]} > "$out" 2>&1; echo "$a submit exit $?: $(grep -E 'submitted job' "$out" | tail -1)"
  fi
done