From e6c47e95323dd09e49d024a60aa6ac0f6a30372f Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Sat, 11 Jul 2026 11:55:59 -0700 Subject: [PATCH] orchestrator: track the SFT training pipeline and fix ToolOrchestra tracing The orchestrator SFT pipeline was living entirely in the working tree and never got committed. One script in that lineage (clean_sft_data.py) has already been lost with no way to recover it, so track the rest before the same thing happens again: run_sft_fsdp.py full-parameter FSDP trainer (launched via accelerate) sft_tokenize.py conversation -> tokens + assistant-only masking render_sft_data.py record -> chat-template rendering make_splits.py train / holdout / overfit splits chat_orchestrator.py REPL for eyeballing a trained checkpoint's routing format_eval_sample.py, upload_to_braintrust.py scripts/train/ fsdp4.yaml, fsdp8.yaml, eval sbatch toolorchestra/tracing.py was also untracked despite being imported by unified.py and rollout.py, which left the branch broken for a fresh checkout. Add it along with the rest of the ToolOrchestra package split. Also here: mmlu_pro / supergpqa scorers, expert pricing for the OpenRouter Qwen builds (estimates, flagged in-file), an expanded calculator/web_search tool surface, and the reject-sampling clean gate that run_sft_fsdp reads via --require-clean. GRPO is left at its upstream state -- the RL stage hasn't been run yet, so it stays out of this change. Paths in the sbatch are env-driven rather than hardcoded to one cluster home. --- .gitignore | 7 + pyproject.toml | 1 + .../orchestrator/build_orchestrator_sft.py | 194 ++++++- scripts/orchestrator/build_unified_sft.py | 103 ---- scripts/orchestrator/chat_orchestrator.py | 109 ++++ scripts/orchestrator/eval_orchestrator.py | 21 + scripts/orchestrator/format_eval_sample.py | 239 ++++++++ scripts/orchestrator/make_splits.py | 120 ++++ scripts/orchestrator/render_sft_data.py | 290 ++++++++++ scripts/orchestrator/run_sft_fsdp.py | 338 +++++++++++ .../orchestrator/sft_configs/fsdp_modal.yaml | 13 + scripts/orchestrator/sft_tokenize.py | 148 +++++ scripts/orchestrator/upload_to_braintrust.py | 361 ++++++++++++ scripts/train/eval_orch_gemma_matx.sbatch | 90 +++ scripts/train/fsdp4.yaml | 25 + scripts/train/fsdp8.yaml | 25 + src/openjarvis/agents/hybrid/_base.py | 8 +- src/openjarvis/agents/hybrid/_prices.py | 12 + .../agents/hybrid/expert_registry.py | 274 +++++++-- .../agents/hybrid/toolorchestra/parsing.py | 51 +- .../agents/hybrid/toolorchestra/rollout.py | 188 +++++- .../agents/hybrid/toolorchestra/tracing.py | 161 ++++++ .../agents/hybrid/toolorchestra/unified.py | 139 ++++- .../agents/hybrid/toolorchestra/workers.py | 3 +- src/openjarvis/agents/orchestrator.py | 90 ++- src/openjarvis/evals/core/scorer.py | 70 ++- src/openjarvis/evals/scorers/gaia_exact.py | 9 +- src/openjarvis/evals/scorers/mmlu_pro_mcq.py | 43 +- src/openjarvis/evals/scorers/supergpqa_mcq.py | 43 +- .../orchestrator-gemma-4-12b-grpo.toml | 54 -- .../configs/orchestrator-qwen3-8b-grpo.toml | 44 -- .../configs/orchestrator-qwen3.5-9b-grpo.toml | 54 -- .../intelligence/orchestrator/eval_backend.py | 46 +- .../intelligence/orchestrator/grpo_trainer.py | 537 +++++------------- .../orchestrator/sft_data/datasets.py | 119 +++- .../orchestrator/sft_data/reject_sample.py | 363 +++++++++++- .../sft_data/unified_serialize.py | 181 +++++- .../orchestrator/sft_data/verify.py | 61 +- src/openjarvis/tools/calculator.py | 15 + src/openjarvis/tools/web_search.py | 26 +- tests/agents/test_anonymize_tools.py | 129 +++++ tests/agents/test_expert_registry.py | 42 +- .../sft_data/test_rejection_pipeline.py | 14 +- .../sft_data/test_verify.py | 10 +- .../test_build_orchestrator_sft.py | 16 +- .../test_eval_backend.py | 10 +- .../test_final_answer_sanitize.py | 112 ++++ .../test_grpo_epoch.py | 408 ------------- 48 files changed, 4125 insertions(+), 1291 deletions(-) delete mode 100644 scripts/orchestrator/build_unified_sft.py create mode 100644 scripts/orchestrator/chat_orchestrator.py create mode 100644 scripts/orchestrator/format_eval_sample.py create mode 100644 scripts/orchestrator/make_splits.py create mode 100644 scripts/orchestrator/render_sft_data.py create mode 100644 scripts/orchestrator/run_sft_fsdp.py create mode 100644 scripts/orchestrator/sft_configs/fsdp_modal.yaml create mode 100644 scripts/orchestrator/sft_tokenize.py create mode 100644 scripts/orchestrator/upload_to_braintrust.py create mode 100644 scripts/train/eval_orch_gemma_matx.sbatch create mode 100644 scripts/train/fsdp4.yaml create mode 100644 scripts/train/fsdp8.yaml create mode 100644 src/openjarvis/agents/hybrid/toolorchestra/tracing.py delete mode 100644 src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-gemma-4-12b-grpo.toml delete mode 100644 src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3-8b-grpo.toml delete mode 100644 src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3.5-9b-grpo.toml create mode 100644 tests/agents/test_anonymize_tools.py create mode 100644 tests/test_orchestrator_learning/test_final_answer_sanitize.py delete mode 100644 tests/test_orchestrator_learning/test_grpo_epoch.py diff --git a/.gitignore b/.gitignore index 80d3cefd..6e978eda 100644 --- a/.gitignore +++ b/.gitignore @@ -133,3 +133,10 @@ oj-debug.*.json # Generated orchestrator SFT/eval data artifacts data/ + +# Training run artifacts (model checkpoints + W&B run dirs — regenerated, large) +checkpoints/ +wandb/ + +# Local LiteLLM spend-tracking proxy (machine-specific, not part of the package) +litellm/ diff --git a/pyproject.toml b/pyproject.toml index 6e253ed7..03159a49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ + "braintrust>=0.0.150", "click>=8", "datasets>=4.5.0", "ddgs>=9.11.4", diff --git a/scripts/orchestrator/build_orchestrator_sft.py b/scripts/orchestrator/build_orchestrator_sft.py index b83c08a2..b6f8cd0a 100644 --- a/scripts/orchestrator/build_orchestrator_sft.py +++ b/scripts/orchestrator/build_orchestrator_sft.py @@ -42,6 +42,10 @@ from __future__ import annotations import argparse import json import logging +import os +import sys +import time +from pathlib import Path from typing import Optional from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog @@ -63,12 +67,25 @@ def main(argv: Optional[list[str]] = None) -> int: p = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter ) - p.add_argument("--out", default="data/orchestrator_sft_v1.jsonl") + # Default output is timestamped (e.g. data/orchestrator_sft_06-22-1230am.jsonl) + # so every run lands in its own file and concurrent runs never collide. Pass + # --out explicitly to override (sharded runs do this, one file per shard). + p.add_argument("--out", default=None, + help="Output JSONL. Default: data/orchestrator_sft_.jsonl") # The orchestrator == the local Qwen3-8B self-sampling over its own rollouts. p.add_argument("--orchestrator-endpoint", default="http://localhost:8001/v1", help="OpenAI-compatible vLLM base URL serving the orchestrator.") p.add_argument("--orchestrator-model", default="qwen3-8b") p.add_argument("--orchestrator-api-key", default="EMPTY") + # Provenance stamps: when set, every written record carries these fields so + # the JSONL is self-identifying once pooled across model families (gemma vs + # qwen). Default None -> no stamping (existing qwen runs unaffected). + p.add_argument("--gen-model", default=None, + help="Full HF id of the generating orchestrator, stamped as " + "record['gen_model'] (e.g. google/gemma-4-26B-A4B-it).") + p.add_argument("--orchestrator-label", default=None, + help="Short label stamped as record['orchestrator_model'] " + "(e.g. gemma-4-26b).") # Local OSS model endpoints; repeatable "model_id=base_url" (e.g. # "Qwen/Qwen3.5-9B=http://localhost:8001/v1"). Unmapped local models are # still listed but served unconfigured (base_url=None). @@ -77,13 +94,112 @@ def main(argv: Optional[list[str]] = None) -> int: help="Local model id -> vLLM base URL (repeatable).") p.add_argument("--max-tasks", type=int, default=None, help="Cap on tasks (default: all of load_sft_tasks()).") + p.add_argument("--skip-task-ids-from", action="append", default=[], + metavar="GLOB", + help="Glob of prior data.jsonl files; skip task_ids already " + "generated there so this run only does unseen prompts " + "(resume). Repeatable; applied before sharding.") p.add_argument("--samples-per-task", type=int, default=8) p.add_argument("--max-keep-per-task", type=int, default=1) p.add_argument("--max-turns", type=int, default=8) p.add_argument("--temperature", type=float, default=0.7) + # Orchestrator completion cap per turn. The library default (4096) intermittently + # truncates the final answer mid-sentence on longer reasoning turns; bump it so + # the trace ends cleanly. Raise further (e.g. 16384) if traces still cut off. + p.add_argument("--max-tokens", type=int, default=8192, + help="Per-turn completion cap for the orchestrator (default 8192).") + # Number of tasks rolled out concurrently against vLLM. Each in-flight task + # issues its samples sequentially, so ~concurrency requests hit the server at + # once. ~30-50 is comfortable on 1xL40S (prefix caching + continuous batching). + p.add_argument("--concurrency", type=int, default=32, + help="Tasks rolled out in parallel (default 32).") + # Balanced-smoke pull: cap//4 from each of GeneralThought + OpenThoughts + # code/math/science instead of the GeneralThought-only fast cap path. Use for a + # representative smoke; the real run omits --max-tasks for the full 8K balanced set. + p.add_argument("--balanced", action=argparse.BooleanOptionalAction, default=True, + help="Draw an EVEN cross-domain sample (GeneralThought + " + "OpenThoughts code/math/science). Default ON — pass " + "--no-balanced for the old GeneralThought-only skew.") + # Data-parallel sharding: split the (deterministic, seed-42) task list across + # N independent driver processes, each pointed at its own orchestrator vLLM + # replica. Shard i takes tasks[i::N] (a strided slice, so every shard stays + # domain-balanced). Each writes its own --out file; concatenate the shard + # JSONLs afterward. Lets the GPU-bound orchestrator scale across idle GPUs. + p.add_argument("--shard-index", type=int, default=0, + help="This shard's index in [0, shard-count).") + p.add_argument("--shard-count", type=int, default=1, + help="Total number of shards (default 1 = no sharding).") + # Throughput knob: stop sampling a task as soon as max-keep-per-task passing + # trajectories are found, instead of always running all --samples-per-task. + # ~3-4x faster when most tasks solve early, but keeps the *first* passers + # rather than the *cheapest* of N (drops the cost-optimisation signal). + p.add_argument("--stop-at-keep", action="store_true", + help="Short-circuit a task once max-keep passing samples found.") + # Anonymize model experts (opaque random labels, uniform description, no cost + # line, shuffled order) so the policy can't route on a model's name/position/ + # cost. The anon->real map is saved per record (metrics.anon_map) for analysis. + p.add_argument("--anonymize-experts", action="store_true", + help="Hide expert identity (random labels + shuffle) when routing.") + # By default we keep EVERY rolled-out trajectory (correct + incorrect), each + # tagged with ``correct`` (verifier verdict) and ``kept`` (the cheapest-correct + # sample the rejection sampler would pick). Lets you compute accuracy / inspect + # failures from the JSONL directly. --rejection-only restores the old + # drop-the-failures behaviour (only cheapest-correct written). + p.add_argument("--rejection-only", action="store_true", + help="Drop incorrect rollouts; write only the cheapest-correct " + "sample per task (the original behaviour).") + # Rollouts call shell_exec / file_write with model-chosen relative paths + # (e.g. ``solution.py``), which otherwise land in the repo root. Run them + # from a throwaway scratch dir so generated files never dirty the tree. + # gitignored via the existing ``scratch/`` rule. + p.add_argument("--scratch-dir", default="scratch/sft-rollouts", + help="CWD for rollouts; stray tool-written files go here.") args = p.parse_args(argv) logging.basicConfig(level=logging.INFO, format="%(message)s") + # Quiet the per-request HTTP / dataset-stream spam so the run log stays + # readable (rollout calls and dataset shards otherwise flood it). + for _noisy in ("httpx", "httpcore", "urllib3", "datasets", "fsspec", + "huggingface_hub", "openai"): + logging.getLogger(_noisy).setLevel(logging.WARNING) + + # Every run lands in its OWN datetimed folder, all collected flat under + # data/runs/ — so runs sort by time, never collide, and never litter data/. + # --out is treated as an optional LABEL for the folder, not a path; the data + # file inside is always data.jsonl (+ .stats.json/.pretty.json/.txt/.html). + # Folder name: _