#!/usr/bin/env bash
# One-click ZBP scaling ladder. Clone -> ./scripts/run_ladder.sh   (run on the H200 node, nohup-able)
#
# Configurable via env vars (defaults in brackets):
#   SIZES="s60 m124 l350"        model configs from configs/model/   [s60 m124 l350]
#   ARMS="bp zbp_n16 zbp_n64"    train configs from configs/train/   [bp zbp_n16 zbp_n64]
#   NPROC=8                      GPUs per run                        [auto: nvidia-smi -L]
#   TOKENS=2e9                   override token budget per run       [per train config]
#   DATA=data/fineweb            token shards (prepared if missing)  [data/fineweb]
#   PREP_TOKENS=3e9              tokens to tokenize if DATA missing  [3e9]
#   OUT=runs                     output root                        [runs]
#   MICRO_BS=8                   per-GPU micro batch                [8]
#   SET="k=v k2=v2"              extra --set overrides for every run (e.g. vocab=8192 seq_len=256)
#   DRY=1                        print the plan and exit
#
# Runs are sequential (each takes the whole node), resume-safe: a finished run leaves OUT/<name>/DONE
# and is skipped on re-invocation, so the script can be re-run after interruptions.
set -euo pipefail
cd "$(dirname "$0")/.."

SIZES=${SIZES:-"s60 m124 l350"}
ARMS=${ARMS:-"bp zbp_n16 zbp_n64"}
NPROC=${NPROC:-$(nvidia-smi -L 2>/dev/null | wc -l)}
DATA=${DATA:-data/fineweb}
PREP_TOKENS=${PREP_TOKENS:-3e9}
OUT=${OUT:-runs}
MICRO_BS=${MICRO_BS:-8}
EXTRA="--set ${TOKENS:+tokens=$TOKENS} ${SET:-}"
[ "$EXTRA" = "--set  " ] && EXTRA=""

echo "== zbp-scaling ladder: sizes=[$SIZES] arms=[$ARMS] nproc=$NPROC data=$DATA out=$OUT ${TOKENS:+tokens=$TOKENS}"
[ "$NPROC" -ge 1 ] || { echo "no GPUs detected; set NPROC"; exit 1; }
python -c "import zbp_scaling" 2>/dev/null || pip install -e . -q

if [ ! -f "$DATA/train.bin" ]; then
  echo "== preparing data ($PREP_TOKENS tokens -> $DATA) — run this on the training node"
  [ "${DRY:-0}" = 1 ] || python scripts/prepare_data.py --tokens "$PREP_TOKENS" --out "$DATA"
fi

for size in $SIZES; do
  for arm in $ARMS; do
    [ -f "configs/model/$size.yaml" ] || { echo "missing configs/model/$size.yaml"; exit 1; }
    [ -f "configs/train/$arm.yaml" ] || { echo "missing configs/train/$arm.yaml"; exit 1; }
    name="${size}_${arm}"
    dir="$OUT/$name"
    if [ -f "$dir/DONE" ]; then echo "== skip $name (done)"; continue; fi
    echo "== run $name  (model=configs/model/$size.yaml train=configs/train/$arm.yaml)"
    [ "${DRY:-0}" = 1 ] && continue
    mkdir -p "$dir"
    if [ "$NPROC" -gt 1 ]; then
      torchrun --standalone --nproc_per_node="$NPROC" scripts/train.py \
        --model "configs/model/$size.yaml" --train "configs/train/$arm.yaml" \
        --data "$DATA" --out "$dir" --micro_bs "$MICRO_BS" $EXTRA 2>&1 | tee -a "$dir/stdout.log"
    else
      python scripts/train.py --model "configs/model/$size.yaml" --train "configs/train/$arm.yaml" \
        --data "$DATA" --out "$dir" --micro_bs "$MICRO_BS" $EXTRA 2>&1 | tee -a "$dir/stdout.log"
    fi
    touch "$dir/DONE"
  done
done
echo "== ladder complete: $OUT"