From f338d64b21cd7fb3d215a06d9973e672381550e6 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 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. Paths in the sbatch and litellm helpers are now env-driven rather than hardcoded to one cluster home. --- .gitignore | 4 + litellm/config.yaml | 74 ++++ litellm/spend_logger.py | 87 +++++ litellm/spend_report.py | 43 +++ 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 ++- .../intelligence/orchestrator/eval_backend.py | 46 ++- .../intelligence/orchestrator/grpo_trainer.py | 12 +- .../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 ++++++ 47 files changed, 4181 insertions(+), 351 deletions(-) create mode 100644 litellm/config.yaml create mode 100644 litellm/spend_logger.py create mode 100644 litellm/spend_report.py 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 create mode 100644 tests/agents/test_anonymize_tools.py create mode 100644 tests/test_orchestrator_learning/test_final_answer_sanitize.py diff --git a/.gitignore b/.gitignore index 80d3cefd..673fc7b8 100644 --- a/.gitignore +++ b/.gitignore @@ -133,3 +133,7 @@ oj-debug.*.json # Generated orchestrator SFT/eval data artifacts data/ + +# Training run artifacts (model checkpoints + W&B run dirs — regenerated, large) +checkpoints/ +wandb/ diff --git a/litellm/config.yaml b/litellm/config.yaml new file mode 100644 index 00000000..70fa1c12 --- /dev/null +++ b/litellm/config.yaml @@ -0,0 +1,74 @@ +# LiteLLM proxy config for OpenJarvis ToolOrchestra expert spend tracking. +# Launch: litellm --config litellm/config.yaml --port 4000 --host 127.0.0.1 +# Keys are read from process env (source OpenJarvis/.env first). +# No Postgres on this box (docker socket denied) -> admin UI/virtual keys are +# disabled; spend is written to litellm/logs/spend.jsonl by spend_logger.py. + +model_list: + # ---- OpenAI ---- + - model_name: gpt-5.5 + litellm_params: + model: openai/gpt-5.5 + api_key: os.environ/OPENAI_API_KEY + - model_name: gpt-5-mini + litellm_params: + model: openai/gpt-5-mini + api_key: os.environ/OPENAI_API_KEY + - model_name: gpt-4o + litellm_params: + model: openai/gpt-4o + api_key: os.environ/OPENAI_API_KEY + + # ---- Anthropic ---- + - model_name: claude-opus-4-8 + litellm_params: + model: anthropic/claude-opus-4-8 + api_key: os.environ/ANTHROPIC_API_KEY + - model_name: claude-sonnet-4-6 + litellm_params: + model: anthropic/claude-sonnet-4-6 + api_key: os.environ/ANTHROPIC_API_KEY + + # ---- Gemini ---- + - model_name: gemini-2.5-pro + litellm_params: + model: gemini/gemini-2.5-pro + api_key: os.environ/GEMINI_API_KEY + - model_name: gemini-2.5-flash + litellm_params: + model: gemini/gemini-2.5-flash + api_key: os.environ/GEMINI_API_KEY + + # ---- OpenRouter (Qwen coder etc.) ---- + - model_name: qwen-coder + litellm_params: + model: openrouter/qwen/qwen-2.5-coder-32b-instruct + api_key: os.environ/OPENROUTER_API_KEY + - model_name: openrouter/* + litellm_params: + model: openrouter/* + api_key: os.environ/OPENROUTER_API_KEY + + # ---- Local vLLM (self-hosted students / experts) ---- + - model_name: qwen3.5-9b-local + litellm_params: + model: openai/Qwen/Qwen3.5-9B + api_base: http://localhost:8011/v1 + api_key: EMPTY + input_cost_per_token: 0.0 + output_cost_per_token: 0.0 + - model_name: qwen3.6-27b-local + litellm_params: + model: openai/Qwen/Qwen3.6-27B-FP8 + api_base: http://localhost:8002/v1 + api_key: EMPTY + input_cost_per_token: 0.0 + output_cost_per_token: 0.0 + +litellm_settings: + # JSONL spend logger (Postgres-free spend tracking). + callbacks: spend_logger.proxy_handler_instance + drop_params: true + +general_settings: + master_key: os.environ/LITELLM_MASTER_KEY diff --git a/litellm/spend_logger.py b/litellm/spend_logger.py new file mode 100644 index 00000000..6dfd5046 --- /dev/null +++ b/litellm/spend_logger.py @@ -0,0 +1,87 @@ +"""Custom LiteLLM callback that appends per-request cost/usage to a JSONL file. + +This is the spend-tracking fallback used when no Postgres is available for the +LiteLLM admin UI. Every successful (and failed) proxy request writes one line to +$OJ_LITELLM_SPEND_LOG. Aggregate with `litellm/spend_report.py`. +""" + +from __future__ import annotations + +import datetime +import json +import os +import threading + +from litellm.integrations.custom_logger import CustomLogger + +LOG_PATH = os.environ.get( + "OJ_LITELLM_SPEND_LOG", + os.path.join(os.path.dirname(os.path.abspath(__file__)), "logs", "spend.jsonl"), +) +os.makedirs(os.path.dirname(LOG_PATH), exist_ok=True) + +_lock = threading.Lock() + + +def _write(kwargs, response_obj, start_time, end_time, status: str) -> None: + try: + cost = kwargs.get("response_cost") + litellm_params = kwargs.get("litellm_params") or {} + metadata = litellm_params.get("metadata") or {} + usage = None + try: + usage = getattr(response_obj, "usage", None) or ( + response_obj.get("usage") if isinstance(response_obj, dict) else None + ) + except Exception: + usage = None + + def _u(field): + if usage is None: + return None + return getattr(usage, field, None) if not isinstance(usage, dict) else usage.get(field) + + dur = None + try: + if start_time and end_time: + dur = (end_time - start_time).total_seconds() + except Exception: + dur = None + + rec = { + "ts": datetime.datetime.utcnow().isoformat() + "Z", + "status": status, + "model": kwargs.get("model"), + "model_group": metadata.get("model_group"), + "call_type": kwargs.get("call_type"), + "cost_usd": cost, + "prompt_tokens": _u("prompt_tokens"), + "completion_tokens": _u("completion_tokens"), + "total_tokens": _u("total_tokens"), + "duration_s": dur, + "key_alias": metadata.get("user_api_key_alias"), + "exception": str(kwargs.get("exception")) if status == "failure" else None, + } + line = json.dumps(rec) + with _lock: + with open(LOG_PATH, "a") as f: + f.write(line + "\n") + except Exception as exc: # never let logging break a request + print(f"[spend_logger] error: {exc}") + + +class SpendLogger(CustomLogger): + def log_success_event(self, kwargs, response_obj, start_time, end_time): + _write(kwargs, response_obj, start_time, end_time, "success") + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + _write(kwargs, response_obj, start_time, end_time, "success") + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + _write(kwargs, response_obj, start_time, end_time, "failure") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + _write(kwargs, response_obj, start_time, end_time, "failure") + + +proxy_handler_instance = SpendLogger() diff --git a/litellm/spend_report.py b/litellm/spend_report.py new file mode 100644 index 00000000..c157a30b --- /dev/null +++ b/litellm/spend_report.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 +"""Aggregate litellm/logs/spend.jsonl into a per-model spend summary. + +Postgres-free stand-in for the LiteLLM admin UI. Usage: + .venv/bin/python litellm/spend_report.py [path-to-spend.jsonl] +""" +import json +import os +import sys +from collections import defaultdict + +path = sys.argv[1] if len(sys.argv) > 1 else os.environ.get( + "OJ_LITELLM_SPEND_LOG", + os.path.join(os.path.dirname(os.path.abspath(__file__)), "logs", "spend.jsonl"), +) + +agg = defaultdict(lambda: {"calls": 0, "ok": 0, "fail": 0, "cost": 0.0, + "ptok": 0, "ctok": 0}) +total_cost = 0.0 +with open(path) as f: + for line in f: + line = line.strip() + if not line: + continue + r = json.loads(line) + g = r.get("model_group") or r.get("model") or "?" + a = agg[g] + a["calls"] += 1 + a["ok"] += 1 if r.get("status") == "success" else 0 + a["fail"] += 1 if r.get("status") == "failure" else 0 + a["cost"] += float(r.get("cost_usd") or 0) + a["ptok"] += int(r.get("prompt_tokens") or 0) + a["ctok"] += int(r.get("completion_tokens") or 0) + total_cost += float(r.get("cost_usd") or 0) + +hdr = f"{'model':<24}{'calls':>7}{'ok':>5}{'fail':>6}{'p_tok':>9}{'c_tok':>9}{'cost_usd':>12}" +print(hdr) +print("-" * len(hdr)) +for g, a in sorted(agg.items(), key=lambda kv: -kv[1]["cost"]): + print(f"{g:<24}{a['calls']:>7}{a['ok']:>5}{a['fail']:>6}" + f"{a['ptok']:>9}{a['ctok']:>9}{a['cost']:>12.6f}") +print("-" * len(hdr)) +print(f"{'TOTAL':<24}{'':>7}{'':>5}{'':>6}{'':>9}{'':>9}{total_cost:>12.6f}") diff --git a/pyproject.toml b/pyproject.toml index 7c718d9b..c9cb148b 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: _