#!/usr/bin/env bash
# Launch (or resume) the MLM pretraining run inside a detached tmux session.
#
#   scripts/launch_pretrain.sh                 # fresh run: EPOCHS=6 MICRO_BATCH=16 RUN_NAME=large-24L
#   EPOCHS=3 scripts/launch_pretrain.sh        # shorter plan (fixes the LR schedule length)
#   RESUME=1 scripts/launch_pretrain.sh        # continue the same run from its latest checkpoint
#   tmux attach -t phish-pretrain              # watch; Ctrl-b d to detach; Ctrl-C inside = save + clean exit
#
# The schedule length is fixed at the first launch (epochs x steps/epoch); a resume continues it.
# Measured on 2026-10-07: micro-batch 16 -> 13.65 GiB reserved, ~23.6k tokens/s, ~1.25 h per epoch.
set -euo pipefail
cd "$(dirname "$0")/.."
RUN_NAME="${RUN_NAME:-large-24L}"
EPOCHS="${EPOCHS:-6}"
MICRO_BATCH="${MICRO_BATCH:-16}"
SESSION="${SESSION:-phish-pretrain}"
PY=".venv/bin/python"
LOG="logs/pretrain_${RUN_NAME}_$(date +%Y%m%d-%H%M%S).out"

if tmux has-session -t "$SESSION" 2>/dev/null; then
  echo "tmux session '$SESSION' already exists; attach with: tmux attach -t $SESSION"; exit 1
fi
if [[ "${RESUME:-0}" == "1" ]]; then
  CMD="$PY src/pretrain.py train --run-name $RUN_NAME --resume --micro-batch $MICRO_BATCH"
else
  CMD="$PY src/pretrain.py train --run-name $RUN_NAME --epochs $EPOCHS --micro-batch $MICRO_BATCH"
fi
echo "launching in tmux '$SESSION': $CMD"
echo "stdout/stderr -> $LOG (the trainer also writes logs/pretrain-<timestamp>.log and checkpoints/pretrain/$RUN_NAME/metrics.csv)"
tmux new-session -d -s "$SESSION" -c "$PWD" "$CMD 2>&1 | tee '$LOG'; echo; echo '[pretrain exited with status '\${PIPESTATUS[0]}'] press Enter to close'; read"
sleep 2
tmux capture-pane -p -t "$SESSION" | tail -5
echo "attach: tmux attach -t $SESSION"