training_run.sh

#!/usr/bin/env bash
# =============================================================================
# Fireworks RFT on tau2-bench banking — training launcher
#
# Sequenced to match FIREWORKS_TRAINING/OVERVIEW.md §7. It is GUARDED: by default
# it only prints/validates. Pass --go to actually launch the (paid) RFT job.
#
#   bash FIREWORKS_TRAINING/training_run.sh            # dry-run: checks + prints command
#   bash FIREWORKS_TRAINING/training_run.sh --smoke     # FAST local check: 1 Easy + 1 Medium
#   bash FIREWORKS_TRAINING/training_run.sh --localtest # full local eval (all subset tasks)
#   bash FIREWORKS_TRAINING/training_run.sh --go        # actually launch the RFT job
#   bash FIREWORKS_TRAINING/training_run.sh --skip      # --go, but skip local validation,
#                                                       # force evaluator re-upload, and
#                                                       # reuse the last uploaded dataset
#
# Fast repeat --go runs (reuse what's already uploaded, skip the ~15-min local
# validation). IDs are printed by a prior run:
#   SKIP_VALIDATION=1 \
#   EVALUATOR_ID=test-banking-test-banking \
#   DATASET_ID=test-banking-test-banking-dataset-20260702151138 \
#   bash FIREWORKS_TRAINING/training_run.sh --go
#
# Per repo convention, Claude does NOT run this — you do.
# =============================================================================
set -euo pipefail
cd "$(dirname "$0")/.."   # repo root

# ----------------------------- configuration --------------------------------
# Verify every value against `eval-protocol create rft --help` before --go.
#
# Confirmed via `firectl get deployment b2f04w5c` (2026-06-25):
#   Base Model : accounts/fireworks/models/qwen3p6-27b   (qwen3_5, 27.4B, served FP4)
#   Your acct  : accounts/dheerajmpaicmu-4f4056          (where the trained model lands)
# NOTE: 27B may sit above Fireworks' easy tunable tier — if `--go` is rejected on
# size/eligibility, that's why; check with Fireworks or try a ≤16B base first.
ACCOUNT="${ACCOUNT:-accounts/dheerajmpaicmu-4f4056}"                  # your account
BASE_MODEL="${BASE_MODEL:-accounts/fireworks/models/qwen3p6-27b}"     # the frozen base
OUTPUT_MODEL="${OUTPUT_MODEL:-qwen3p6-27b-banking-rft-v1}"            # name of the LoRA output
DATASET="${DATASET:-FIREWORKS_TRAINING/data/banking_easy_medium.jsonl}"
EVALUATOR="${EVALUATOR:-FIREWORKS_TRAINING/test_banking.py}"     # holds @evaluation_test
WANDB_PROJECT="${WANDB_PROJECT:-tau2-banking-rft}"

# --- knobs that directly address the two problems (see OVERVIEW §Problems) ---
# Problem 1 (long degenerate loops): keep rollouts SHORT + penalize repeats.
MAX_TOKENS="${MAX_TOKENS:-4096}"     # tight cap per response; small models ramble
#   (the hard turn cap lives in the MCP server's max_turns, and the early-stop
#    loop guard lives in score_from_reward_info.should_early_stop — both wired in
#    the rollout processor, not here.)
# Problem 2 (sparse end reward): density comes from the partial-credit reward
#   (param_partial), NOT from a flag here. EP still gives one scalar per rollout.
N_ROLLOUTS="${N:-16}"                # group size: P(>=1 success | p=0.1) ~ 81% at 16
RL_LOSS="${RL_LOSS:-DAPO}"           # DAPO drops all-zero AND all-one groups (saves compute)
# The Fireworks SDK Literal wants UPPERCASE with underscores:
#   {METHOD_UNSPECIFIED, GRPO, DAPO, DPO, ORPO, GSPO_TOKEN}
# (fireworks/types/shared_params/reinforcement_learning_loss_config.py). EP's own
# help text says "{grpo,dapo,gspo-token}" but that's stale — lowercase is rejected
# with "invalid Literal value". Normalize so RL_LOSS=dapo / gspo-token still work.
RL_LOSS="$(printf '%s' "$RL_LOSS" | tr '[:lower:]-' '[:upper:]_')"
TEMPERATURE="${TEMPERATURE:-0.8}"    # exploration -> in-group variance -> gradient exists
EPOCHS="${EPOCHS:-1}"                # smoke run: prove the loop moves, not SOTA
LORA_RANK="${LORA_RANK:-8}"

# --- fast-iteration knobs (skip the slow parts on repeat runs) ---
# SKIP_VALIDATION=1 : skip the local pytest validation before submission
#                     (that's the ~15-min rollout run you sit through every --go).
# EVALUATOR_ID=...  : reuse an evaluator already on Fireworks (e.g.
#                     test-banking-test-banking). Skips the interactive test
#                     selector; if it already exists remotely and FORCE isn't
#                     set, create-rft skips the re-upload automatically.
# DATASET_ID=...    : reuse a dataset already on Fireworks (printed by a prior
#                     run, e.g. test-banking-test-banking-dataset-20260702151138).
#                     Skips rebuilding + re-uploading the JSONL.
# FORCE=1          : delete + re-upload the evaluator even if it already exists
#                     remotely (REQUIRED after a failed build — without it the
#                     exists-check just re-polls the broken build and aborts).
SKIP_VALIDATION="${SKIP_VALIDATION:-0}"
EVALUATOR_ID="${EVALUATOR_ID:-}"
DATASET_ID="${DATASET_ID:-}"
FORCE="${FORCE:-0}"

# --- --skip mode: preset the fast knobs above, then behave exactly like --go ---
# One-word fast relaunch: skip the ~15-min local validation, force re-upload of
# the evaluator (safe even after a failed remote build), and reuse the uploaded
# dataset below (override with DATASET_ID=... if a newer one exists).
# Must happen HERE (before RFT_CMD is assembled), not in the case block.
MODE="${1:-}"
if [ "$MODE" = "--skip" ]; then
  SKIP_VALIDATION=1
  FORCE=1
  DATASET_ID="${DATASET_ID:-test-banking-test-banking-dataset-20260702153218}"
  # Required in --yes mode: the repo has TWO eval tests (test_banking + test_airline),
  # so non-interactive discovery aborts with "Multiple evaluation tests found"
  # unless the evaluator is named explicitly.
  EVALUATOR_ID="${EVALUATOR_ID:-test-banking-test-banking}"
  MODE="--go"
  echo "== --skip: SKIP_VALIDATION=1 FORCE=1 EVALUATOR_ID=$EVALUATOR_ID DATASET_ID=$DATASET_ID =="
fi

# ----------------------------- preflight ------------------------------------
echo "== preflight =="
command -v eval-protocol >/dev/null || { echo "ERROR: eval-protocol not installed (pip install eval-protocol)"; exit 1; }
# [ -n "${FIREWORKS_API_KEY:-}" ] || { echo "ERROR: FIREWORKS_API_KEY not set"; exit 1; }
[ -f "$DATASET" ] || { echo "ERROR: dataset $DATASET missing — build it first:"; echo "    python3 workspace/build_banking_easy.py"; exit 1; }
case "$BASE_MODEL" in
  *REPLACE_ME*) echo "ERROR: BASE_MODEL still a placeholder."; exit 1;;
esac
[ -f "$EVALUATOR" ] || { echo "WARN: evaluator $EVALUATOR not found — create it from OVERVIEW §2 (imports workspace/score_from_reward_info.py)"; }
echo "  eval-protocol: $(eval-protocol --version 2>/dev/null || echo '?')"
# Which binary actually runs for --preview/--go (PATH resolution at this moment).
# NOTE: --localtest prepends .venv/bin to PATH later, so ITS binary may differ —
# that mode prints its own "using:" line.
echo "  eval-protocol bin: $(command -v eval-protocol)"
echo "  python bin       : $(command -v python)"
echo "  rows in dataset: $(wc -l < "$DATASET")"
# inference routing: local eval can go through the LeanMCP gateway (logged), the
# on-Fireworks training rollouts cannot (trainee policy is internal to Fireworks).
if [ "${USE_LEANMCP_GATEWAY:-0}" = "1" ]; then
  [ -n "${LEANMCP_API_KEY:-}" ] || echo "  WARN: USE_LEANMCP_GATEWAY=1 but LEANMCP_API_KEY not set"
  echo "  local-eval inference: LeanMCP gateway (https://aigateway.leanmcp.com/v1/fireworks) -> Fireworks"
else
  echo "  local-eval inference: Fireworks direct"
fi
echo

# ----------------------------- the command ----------------------------------
# NOTE (2026-07): the `eval-protocol create rft` CLI moved to namespaced flags.
# Most old flags kept aliases (--base-model, --output-model, --rl-loss-method,
# --temperature, --epochs, --lora-rank, --wandb-project all still resolve), but
# TWO were removed and now error:
#   --n           -> --response-candidates-count   (group size / candidates per prompt)
#   --max-tokens  -> --max-output-tokens           (per-response token cap)
# Verify against `eval-protocol create rft --help` if it drifts again.
# The evaluator + MCP server + dataset are auto-discovered from the EP test in the
# working dir (test_banking.py's @evaluation_test / input_dataset), so we do NOT
# pass --dataset/--evaluator here; point the test at a subset via EP_DATASET instead.
RFT_CMD=(eval-protocol create rft
  --base-model "$BASE_MODEL"
  --output-model "$OUTPUT_MODEL"
  --response-candidates-count "$N_ROLLOUTS"
  --rl-loss-method "$RL_LOSS"
  --temperature "$TEMPERATURE"
  --max-output-tokens "$MAX_TOKENS"
  --epochs "$EPOCHS"
  --lora-rank "$LORA_RANK"
  --wandb-project "$WANDB_PROJECT"
  --yes   # non-interactive: no test-selector / secrets prompts mid-run
)
[ "$SKIP_VALIDATION" = "1" ] && RFT_CMD+=(--skip-validation)
[ -n "$EVALUATOR_ID" ]       && RFT_CMD+=(--evaluator "$EVALUATOR_ID")
[ -n "$DATASET_ID" ]         && RFT_CMD+=(--dataset "$DATASET_ID")
[ "$FORCE" = "1" ]           && RFT_CMD+=(--force)

echo "== RFT command =="
printf '  %q ' "${RFT_CMD[@]}"; echo; echo

# ----------------------------- modes ----------------------------------------
run_localtest() {
  # Gateway OFF: call Fireworks directly with FIREWORKS_API_KEY. Override with
  # USE_LEANMCP_GATEWAY=1 if you ever want the LeanMCP-proxied (logged) path.
  export USE_LEANMCP_GATEWAY="${USE_LEANMCP_GATEWAY:-0}"
  # MUST use the env that has tau2 (the repo .venv), else EP's whole-repo
  # collection aborts on tests/conftest.py importing tau2. See
  # workspace/setup_localtest_env.sh (run that once to install ep into .venv).
  EP_BIN="eval-protocol"
  [ -x ".venv/bin/eval-protocol" ] && EP_BIN=".venv/bin/eval-protocol"
  # EP spawns the MCP server with a bare `python` resolved from PATH (see
  # eval_protocol/.../default_mcp_gym_rollout_processor.py: cmd = ["python", ...]).
  # If PATH points at miniconda base, the server crashes on imports it lacks
  # (rank_bm25 / tau2). Force the venv onto the front of PATH so the spawned
  # server uses the same interpreter that has every dep installed.
  [ -d ".venv/bin" ] && export PATH="$PWD/.venv/bin:$PATH"
  echo "   using: $EP_BIN  (PATH python: $(command -v python))"
  # Non-interactive: skip the "Select this test?" prompt.
  #   --yes           : don't prompt
  #   --entry a::b     : run exactly this test (so --yes has no selector to show)
  #   --ignore-docker  : run pytest on the host (don't try to build/run a container)
  "$EP_BIN" local-test --yes --entry "${EVALUATOR}::test_banking" --ignore-docker
}

case "$MODE" in
  --localtest)
    echo "== ep local-test (offline eval, no training) — inference direct to Fireworks =="
    run_localtest
    ;;
  --smoke)
    # Minimal local sanity check: 1 Easy + 1 Medium task only (test_banking.py's
    # dataset_adapter reads MAX_<BUCKET> env caps). Runs in minutes, not ~15.
    echo "== SMOKE local-test: 1 Easy + 1 Medium rollout only =="
    export MAX_EASY=1 MAX_MEDIUM=1
    run_localtest
    ;;
  --preview)
    # EP's OWN dry-run: actually invokes the CLI and prints the planned SDK call
    # (validating flags + evaluator/dataset discovery) WITHOUT submitting a paid job.
    echo "== EP dry-run (no submission) — validating flags + discovery =="
    "${RFT_CMD[@]}" --dry-run
    ;;
  --go)
    echo "== LAUNCHING RFT (this is a paid job) =="
    echo "NOTE: the trainee policy's rollouts run INSIDE Fireworks and are NOT routed"
    echo "      through the LeanMCP gateway (no client api_base to hijack during RFT)."
    echo "      LeanMCP logging applies to local eval (--localtest), not this job."
    "${RFT_CMD[@]}"
    echo
    echo "Track it:  eval-protocol list   (or W&B project '$WANDB_PROJECT')"
    echo "Deploy:    firectl deployment create --model $ACCOUNT/models/$OUTPUT_MODEL"
    echo "           (new firectl syntax is 'firectl <noun> <verb>'; the trained LoRA"
    echo "            lands in YOUR account: $ACCOUNT/models/$OUTPUT_MODEL)"
    ;;
  *)
    echo "Dry run only. Re-run with --smoke (1+1 fast check), --localtest (full local eval),"
    echo "or --go to launch. For fast repeat --go runs see the header (SKIP_VALIDATION/"
    echo "EVALUATOR_ID/DATASET_ID)."
    ;;
esac