#!/usr/bin/env bash
# Runs on the main pod (all GPUs): setup -> text cache -> [teacher samples] -> training.
#   DISC_REAL=adm|teacher TEACHER_COUNT=2400 bash scripts/main_run.sh [--set key=value ...]
# Expects code at /workspace/code and prompts at /workspace/prompts/{train,eval}.jsonl.
set -euo pipefail
cd /workspace/code
N=$(nvidia-smi -L | wc -l)
M=/workspace/models/MOSS-Image-SFT
echo "== setup $(date -u +%T)"
bash scripts/pod_setup.sh
echo "== text cache $(date -u +%T)"
torchrun --nproc-per-node "$N" -m rapid_moss.textenc --model $M \
    --prompts /workspace/prompts/train.jsonl /workspace/prompts/eval.jsonl --out /workspace/cache
EXTRA=(disc_real=adm)
if [[ "${DISC_REAL:-teacher}" == teacher ]]; then
    echo "== teacher samples $(date -u +%T)"
    torchrun --nproc-per-node "$N" -m rapid_moss.teachergen --model $M --cache /workspace/cache \
        --prompts /workspace/prompts/train.jsonl --out /workspace/teacher --count "${TEACHER_COUNT:-2000}"
    EXTRA=()
fi
echo "== train $(date -u +%T)"
torchrun --nproc-per-node "$N" -m rapid_moss.train --config configs/main.json --set "${EXTRA[@]}" "$@"
echo "== done $(date -u +%T)"