From 1a843d1ece8701234befd751bfae07e53606d1b6 Mon Sep 17 00:00:00 2001 From: Andrew Park Date: Wed, 15 Jul 2026 12:50:03 -0700 Subject: [PATCH] orchestrator: SFT training + eval harness Tokenize the cleaned routing data (all-turns masking for ChatML), full-parameter FSDP SFT (fsdp4/8 configs for Lambda 8xH100), and an eval backend that drives the fine-tuned orchestrator as the scored model over GAIA/MMLU-Pro/SuperGPQA. Includes split-making, Braintrust upload, eval sample rendering, scorer fixes, and a boxed-tool-call salvage for SFT'd checkpoints that emit their action inside \boxed{}. --- .gitignore | 10 + scripts/orchestrator/eval_orchestrator.py | 358 +++++++++++++++ scripts/orchestrator/make_splits.py | 176 ++++++++ scripts/orchestrator/run_sft_fsdp.py | 407 ++++++++++++++++++ scripts/orchestrator/sft_tokenize.py | 151 +++++++ scripts/train/fsdp4.yaml | 25 ++ scripts/train/fsdp8.yaml | 25 ++ src/openjarvis/agents/orchestrator.py | 90 +++- src/openjarvis/evals/core/scorer.py | 75 +++- 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 +- .../configs/orchestrator-qwen3.5-9b-sft.toml | 34 ++ .../intelligence/orchestrator/eval_backend.py | 288 +++++++++++++ .../intelligence/orchestrator/sft_trainer.py | 103 ++++- .../test_eval_backend.py | 182 ++++++++ .../test_final_answer_sanitize.py | 120 ++++++ .../test_sft_trainer.py | 2 +- 18 files changed, 2105 insertions(+), 36 deletions(-) create mode 100644 scripts/orchestrator/eval_orchestrator.py create mode 100644 scripts/orchestrator/make_splits.py create mode 100644 scripts/orchestrator/run_sft_fsdp.py create mode 100644 scripts/orchestrator/sft_tokenize.py create mode 100644 scripts/train/fsdp4.yaml create mode 100644 scripts/train/fsdp8.yaml create mode 100644 src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3.5-9b-sft.toml create mode 100644 src/openjarvis/learning/intelligence/orchestrator/eval_backend.py create mode 100644 tests/test_orchestrator_learning/test_eval_backend.py create mode 100644 tests/test_orchestrator_learning/test_final_answer_sanitize.py diff --git a/.gitignore b/.gitignore index 75e2da3b..6e978eda 100644 --- a/.gitignore +++ b/.gitignore @@ -130,3 +130,13 @@ learning.db minion_logs/ *.oj-debug.json 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/scripts/orchestrator/eval_orchestrator.py b/scripts/orchestrator/eval_orchestrator.py new file mode 100644 index 00000000..b0b12759 --- /dev/null +++ b/scripts/orchestrator/eval_orchestrator.py @@ -0,0 +1,358 @@ +#!/usr/bin/env python +"""Run the ToolOrchestra orchestrator over N-sample subsets of several +benchmarks and score them with the existing eval infra. + +Reuses, verbatim: + * ``openjarvis.evals.cli._build_dataset`` / ``_build_scorer`` / + ``_build_judge_backend`` (loaders + scorers + judge wiring), + * ``openjarvis.evals.core.runner.EvalRunner`` (parallel run + scoring), + * ``openjarvis.evals.core.types.RunConfig`` (run config), + +and plugs in our orchestrator as the "model" via +``openjarvis.learning.intelligence.orchestrator.eval_backend.OrchestratorBackend``. + +For each benchmark it runs ``EvalRunner(config, dataset, backend, scorer).run()``, +collects accuracy / cost / latency, writes a combined ``summary.json``, and +prints a table. + +NOTE on the judge: the OpenAI and Gemini keys are both dead — Anthropic is the +only live provider, so the judge is ``claude-haiku-4-5-20251001`` (the same model +the generation-side verifier uses, see ``sft_data/verify.py:JUDGE_MODEL``). The +judge backend is still built via the ``cloud`` engine (``_build_judge_backend``) +— the model string selects the provider/route. Override with ``--judge-model`` / +``--judge-engine``. + +Example +------- + .venv/bin/python scripts/orchestrator/eval_orchestrator.py \\ + --benchmarks gaia,mmlu_pro --n 100 \\ + --orchestrator-endpoint http://localhost:8001/v1 \\ + --orchestrator-model qwen3-8b \\ + --output-dir ~/.openjarvis/experiments/hybrid/runs/orch-eval +""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path +from typing import Any, Dict, List, Optional + +# Friendly aliases (underscored / shorthand) -> the EXACT registry keys used by +# openjarvis.evals.cli.BENCHMARKS / _build_dataset / _build_scorer. +BENCHMARK_ALIASES: Dict[str, str] = { + "terminalbench_v2_1": "terminalbench-v2.1", + "terminalbench-v2_1": "terminalbench-v2.1", + "terminalbench_v21": "terminalbench-v2.1", + "terminalbench-v2.1": "terminalbench-v2.1", + "mmlu_pro": "mmlu-pro", + "mmlu-pro": "mmlu-pro", + "gaia": "gaia", + "taubench": "taubench", + "supergpqa": "supergpqa", +} + +DEFAULT_BENCHMARKS = "gaia,terminalbench_v2_1,taubench,mmlu_pro,supergpqa" + + +def _normalize_benchmark(name: str) -> str: + key = name.strip() + return BENCHMARK_ALIASES.get(key, BENCHMARK_ALIASES.get(key.lower(), key)) + + +def _parse_args(argv: Optional[List[str]] = None) -> argparse.Namespace: + p = argparse.ArgumentParser( + description="Run the orchestrator over benchmark subsets and score them.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + p.add_argument( + "--benchmarks", + default=DEFAULT_BENCHMARKS, + help="Comma-separated benchmark names (aliases normalized to registry keys).", + ) + p.add_argument("--n", type=int, default=100, help="Samples per benchmark.") + p.add_argument("--seed", type=int, default=42, help="Subset seed.") + p.add_argument( + "--orchestrator-endpoint", + default="http://localhost:8001/v1", + help="OpenAI-compatible base URL for the served orchestrator.", + ) + p.add_argument( + "--orchestrator-model", + default="qwen3-8b", + help="Served orchestrator model id.", + ) + p.add_argument( + "--orchestrator-api-key", + default="EMPTY", + help="API key for the orchestrator endpoint (EMPTY for local vLLM).", + ) + p.add_argument( + "--local-endpoint", + action="append", + default=[], + metavar="MODEL_ID=URL", + help=( + "Wire a local OSS tool: map a catalog model id to its vLLM base " + "URL, e.g. 'Qwen/Qwen3.5-9B=http://localhost:8003/v1'. Repeatable." + ), + ) + p.add_argument( + "--max-turns", + type=int, + default=8, + help="Max orchestrator turns per sample.", + ) + p.add_argument( + "--temperature", + type=float, + default=1.0, + help="Orchestrator sampling temperature.", + ) + p.add_argument( + "--max-workers", + type=int, + default=4, + help="Parallel samples per benchmark.", + ) + p.add_argument( + "--output-dir", + default="results/orchestrator-eval", + help="Directory for per-benchmark JSONL + combined summary.json.", + ) + # OpenAI AND Gemini keys are both dead — Anthropic is the only live provider, + # and this matches the generation-side verifier (verify.JUDGE_MODEL). + p.add_argument( + "--base-model", + action="store_true", + help="Orchestrator is an UNTRAINED base model: serve with native tools= " + "instead of the baked fine-tuned JSON prompt.", + ) + p.add_argument( + "--judge-model", + default="claude-haiku-4-5-20251001", + help="LLM-judge model.", + ) + p.add_argument( + "--judge-engine", + default="cloud", + help="Engine key for the judge backend.", + ) + return p.parse_args(argv) + + +def _run_one( + benchmark: str, + *, + n: int, + seed: int, + backend, + judge_model: str, + judge_engine: str, + max_workers: int, + orchestrator_model: str, + output_dir: Path, +) -> Dict[str, Any]: + """Run + score one benchmark, returning a row dict (or an error row).""" + # Imports here so --help works without importing the (heavy) eval stack. + from openjarvis.evals.cli import ( + _build_dataset, + _build_judge_backend, + _build_scorer, + ) + from openjarvis.evals.core.runner import EvalRunner + from openjarvis.evals.core.types import RunConfig + + output_path = output_dir / f"{benchmark.replace('.', '_')}_orchestrator.jsonl" + + dataset = _build_dataset(benchmark) + # The runner reloads internally; we also load here per spec so callers can + # inspect/size the subset up front. load() is idempotent for these sets. + dataset.load(max_samples=n, seed=seed) + + judge_backend = _build_judge_backend(judge_model, engine_key=judge_engine) + scorer = _build_scorer(benchmark, judge_backend, judge_model) + + config = RunConfig( + benchmark=benchmark, + backend="orchestrator", + model=orchestrator_model, + max_samples=n, + max_workers=max_workers, + seed=seed, + judge_model=judge_model, + judge_engine=judge_engine, + output_path=str(output_path), + ) + + runner = EvalRunner(config, dataset, backend, scorer) + started = time.time() + try: + summary = runner.run() + finally: + if judge_backend is not None: + judge_backend.close() + elapsed = time.time() - started + + return { + "benchmark": benchmark, + "accuracy": summary.accuracy, + "scored_samples": summary.scored_samples, + "correct": summary.correct, + "errors": summary.errors, + "total_samples": summary.total_samples, + "mean_latency_seconds": summary.mean_latency_seconds, + "total_cost_usd": summary.total_cost_usd, + "mean_continuous_score": summary.mean_continuous_score, + "wall_seconds": round(elapsed, 2), + "output_path": str(output_path), + } + + +def _print_table(rows: List[Dict[str, Any]]) -> None: + cols = [ + ("benchmark", 22, "{}"), + ("accuracy", 9, "{:.4f}"), + ("correct", 8, "{}"), + ("scored", 7, "{}"), + ("errors", 7, "{}"), + ("cost($)", 10, "{:.4f}"), + ("lat(s)", 9, "{:.2f}"), + ("wall(s)", 9, "{:.1f}"), + ] + header = " ".join(f"{name:<{w}}" for name, w, _ in cols) + print("\n" + header) + print("-" * len(header)) + for r in rows: + if r.get("error"): + print(f"{r['benchmark']:<22} ERROR: {r['error']}") + continue + vals = { + "benchmark": r["benchmark"], + "accuracy": r["accuracy"], + "correct": r["correct"], + "scored": r["scored_samples"], + "errors": r["errors"], + "cost($)": r["total_cost_usd"], + "lat(s)": r["mean_latency_seconds"], + "wall(s)": r["wall_seconds"], + } + cells = [] + for name, w, fmt in cols: + v = vals[name] + try: + s = fmt.format(v) + except (ValueError, TypeError): + s = str(v) + cells.append(f"{s:<{w}}") + print(" ".join(cells)) + + +def main(argv: Optional[List[str]] = None) -> int: + args = _parse_args(argv) + + benchmarks = [ + _normalize_benchmark(b) for b in args.benchmarks.split(",") if b.strip() + ] + output_dir = Path(args.output_dir).expanduser() + output_dir.mkdir(parents=True, exist_ok=True) + + # Build the orchestrator backend once (stateless across benchmarks). + from openjarvis.learning.intelligence.orchestrator.eval_backend import ( + OrchestratorBackend, + ) + + local_endpoints: Dict[str, str] = {} + for pair in args.local_endpoint: + if "=" not in pair: + raise SystemExit(f"--local-endpoint expects MODEL_ID=URL, got: {pair!r}") + model_id, url = pair.split("=", 1) + local_endpoints[model_id.strip()] = url.strip() + + backend = OrchestratorBackend( + orchestrator_endpoint=args.orchestrator_endpoint, + orchestrator_model=args.orchestrator_model, + api_key=args.orchestrator_api_key, + local_endpoints=local_endpoints, + max_turns=args.max_turns, + temperature=args.temperature, + # BASE models need native tools=; FINE-TUNED models need the baked + # JSON prompt they were trained on. Serving either in the other mode + # handicaps it (base-in-baked scored 0.184 on a rerun; the June 0.40 + # baseline used native mode). + finetuned=not args.base_model, + ) + + rows: List[Dict[str, Any]] = [] + try: + for benchmark in benchmarks: + print(f"\n=== {benchmark} (n={args.n}, seed={args.seed}) ===") + try: + row = _run_one( + benchmark, + n=args.n, + seed=args.seed, + backend=backend, + judge_model=args.judge_model, + judge_engine=args.judge_engine, + max_workers=args.max_workers, + orchestrator_model=args.orchestrator_model, + output_dir=output_dir, + ) + print( + f" accuracy={row['accuracy']:.4f} " + f"({row['correct']}/{row['scored_samples']}) " + f"errors={row['errors']} cost=${row['total_cost_usd']:.4f}" + ) + except Exception as exc: # noqa: BLE001 - one bad bench shouldn't kill all + print(f" FAILED: {type(exc).__name__}: {exc}", file=sys.stderr) + row = {"benchmark": benchmark, "error": f"{type(exc).__name__}: {exc}"} + rows.append(row) + finally: + backend.close() + + # Resolve the EXACT model behind the served alias (e.g. "gemma-sft" -> + # the real id + checkpoint path vLLM reports), so results are self-identifying + # instead of an opaque label. Best-effort: never fail the summary over it. + served_id, served_path = args.orchestrator_model, None + try: + import urllib.request + + req = urllib.request.Request( + args.orchestrator_endpoint.rstrip("/") + "/models", + headers={"Authorization": f"Bearer {args.orchestrator_api_key}"}, + ) + with urllib.request.urlopen(req, timeout=10) as resp: + data = json.load(resp).get("data", []) + if data: + served_id = data[0].get("id", served_id) + # vLLM reports the loaded checkpoint path in `root` (or `parent`). + served_path = data[0].get("root") or data[0].get("parent") + except Exception: + pass + + combined = { + "orchestrator_model": args.orchestrator_model, + "orchestrator_served_id": served_id, + "orchestrator_model_path": served_path, + "orchestrator_endpoint": args.orchestrator_endpoint, + "judge_model": args.judge_model, + "judge_engine": args.judge_engine, + "n": args.n, + "seed": args.seed, + "max_turns": args.max_turns, + "benchmarks": rows, + } + summary_path = output_dir / "summary.json" + with open(summary_path, "w") as f: + json.dump(combined, f, indent=2, default=str) + + _print_table(rows) + print(f"\nCombined summary written to {summary_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/orchestrator/make_splits.py b/scripts/orchestrator/make_splits.py new file mode 100644 index 00000000..64727219 --- /dev/null +++ b/scripts/orchestrator/make_splits.py @@ -0,0 +1,176 @@ +#!/usr/bin/env python +"""Carve train / holdout / overfit100 splits from a clean orchestrator-SFT pool. + +Replaces make_tranches.py + make_gemma_tranches.py, which carved 1k/2k/4k/8k +data-scaling tranches from the June pool. That lineage is dead: the tranches it +produced trained checkpoints that evaluated at base ~= 1k ~= 2k (the pool taught +the orchestrator to re-derive answers itself instead of routing), and the pool +was deleted. The lesson is in data/orchestrator/README.md — it was a data-QUALITY +ceiling, so re-running the old scaling ladder buys nothing. + +Reads a clean pool (reject-sampled from raw/) and writes the SFT splits to sft/. +Naming is uniform — ``{name}-{split}-{stamp}`` with name in {qwen, gemma, pooled}; +pass several pools to merge-and-restratify into `pooled`. The splits INHERIT the +pool's stamp, so they stay tied to the run that generated the data. + +The holdout is domain-stratified so val-loss is leak-free, and `overfit100` is a +strict prefix of train (it is a memorisation sanity check — the model must be +able to fit 100 rows it *has* seen, else the format/masking is broken). +Deterministic (seed 42). + + # one pool -> qwen-{train,holdout,overfit100}-july-7-2026-0553pm.jsonl + python scripts/orchestrator/make_splits.py --name qwen \ + --pool .../sft/qwen-clean-july-7-2026-0553pm.jsonl + + # merge both -> pooled-{train,holdout,overfit100}-july-7-2026-0553pm.jsonl + python scripts/orchestrator/make_splits.py --name pooled \ + --pool .../sft/qwen-clean-july-7-2026-0553pm.jsonl \ + --pool .../sft/gemma-clean-july-7-2026-0553pm.jsonl +""" + +import argparse +import json +import os +import random +import sys +from collections import defaultdict +from pathlib import Path + +from openjarvis.learning.intelligence.orchestrator.sft_data.naming import ( + dataset_name, + run_stamp, + stamp_from, +) + +# Where the orchestrator data tree lives. Repo-relative by default so a fresh +# clone works out of the box; set OJ_DATA_ROOT to keep the data OUT of the git +# checkout (this workspace points it at ~/experiments/orchestrator/data, so a +# stray `git reset` can't touch hundreds of GB of generations). +DATA_ROOT = Path(os.getenv("OJ_DATA_ROOT", "data/orchestrator")) +OUT = DATA_ROOT / "sft" +SEED = 42 +OVERFIT_N = 100 +# Orchestrator that generated each pool — stamped onto every row so provenance +# survives a merge (the filename encodes it too, but the field is cheap insurance +# and is what the `pooled` splits rely on to stay attributable). +ORCH_MODEL = {"qwen": "qwen3.5-9b", "gemma": "gemma-4-26b"} + + +def main() -> int: + ap = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + ap.add_argument( + "--name", + required=True, + help="split family: qwen | gemma | pooled (drives the filename)", + ) + ap.add_argument( + "--pool", + action="append", + required=True, + metavar="PATH", + help="clean pool jsonl; repeat to merge (use with --name pooled)", + ) + ap.add_argument( + "--holdout-frac", + type=float, + default=0.15, + help="fraction held out, domain-stratified (default 0.15)", + ) + ap.add_argument( + "--stamp", + default=None, + help="Stamp for the output names, e.g. july-7-2026-0553pm. Defaults to the " + "POOL's stamp (read off its filename) so a split stays tied to the run " + "that generated it, NOT to the day it happened to be carved. Falls back " + "to now only if no pool name carries a stamp.", + ) + ap.add_argument("--out", type=Path, default=OUT) + args = ap.parse_args() + + # A split belongs to the generation run, not to the day it was carved — so + # inherit the pool's stamp unless told otherwise. + stamp = args.stamp or next( + (s for p in args.pool if (s := stamp_from(Path(p).name))), None + ) or run_stamp() + + rows = [] + for p in args.pool: + path = Path(p) + if not path.exists(): + print(f"!! pool not found: {path}", file=sys.stderr) + return 1 + # Infer the source orchestrator from the filename so merged pools stay + # attributable; don't clobber a stamp the pool already carries. + stem = path.name.split("-", 1)[0] + loaded = [json.loads(line) for line in path.open() if line.strip()] + for r in loaded: + r.setdefault("orchestrator_model", ORCH_MODEL.get(stem, stem)) + rows.extend(loaded) + print(f"loaded {len(loaded):>5} rows from {path}") + print(f"pool total: {len(rows)}") + + # GROUPED stratified holdout: split by task_id, never by row. + # + # The pool keeps SEVERAL rollouts per question (different attempts at the same + # task, same gold answer). Splitting by row therefore leaks: attempt A of a + # question lands in train and attempt B of the SAME question lands in holdout, + # so val-loss is measured on questions the model trained on. That is exactly + # what happened to the 0707 splits — 80% of qwen's holdout tasks were also in + # train, making every val-loss number optimistically biased. + # + # So: group rows by task_id, stratify the TASKS by domain, and assign whole + # tasks to one side. A question is never split across train and holdout. + by_task = defaultdict(list) + for i, r in enumerate(rows): + by_task[r.get("task_id")].append(i) + + task_dom = {t: rows[idxs[0]].get("domain", "misc") for t, idxs in by_task.items()} + by_dom = defaultdict(list) + for t, dom in task_dom.items(): + by_dom[dom].append(t) + + rng = random.Random(SEED) + holdout_tasks: set = set() + for dom, tasks in sorted(by_dom.items()): + tasks = sorted(tasks) # deterministic before sampling + k = round(args.holdout_frac * len(tasks)) + pick = rng.sample(tasks, min(k, len(tasks))) + holdout_tasks.update(pick) + print(f" {dom:<9} tasks={len(tasks):>4} holdout={len(pick)}") + + holdout_idx = {i for t in holdout_tasks for i in by_task[t]} + holdout = [rows[i] for i in sorted(holdout_idx)] + train = [r for i, r in enumerate(rows) if i not in holdout_idx] # order preserved + overfit = train[:OVERFIT_N] + + # Assert the leak is actually gone — this is the whole point of the grouping. + tr_tasks = {r.get("task_id") for r in train} + ho_tasks = {r.get("task_id") for r in holdout} + overlap = tr_tasks & ho_tasks + if overlap: + raise SystemExit(f"!! LEAK: {len(overlap)} task_ids in BOTH train and holdout") + print( + f"train={len(train)} rows / {len(tr_tasks)} tasks " + f"holdout={len(holdout)} rows / {len(ho_tasks)} tasks " + f"overfit={len(overfit)} (0 tasks shared)" + ) + + args.out.mkdir(parents=True, exist_ok=True) + written = {} + for split, data in ( + ("train", train), + ("holdout", holdout), + (f"overfit{OVERFIT_N}", overfit), + ): + f = args.out / f"{dataset_name(args.name, split, stamp)}.jsonl" + f.write_text("".join(json.dumps(r) + "\n" for r in data)) + written[split] = f + print(f"wrote {f} ({len(data)})") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/orchestrator/run_sft_fsdp.py b/scripts/orchestrator/run_sft_fsdp.py new file mode 100644 index 00000000..af63e791 --- /dev/null +++ b/scripts/orchestrator/run_sft_fsdp.py @@ -0,0 +1,407 @@ +#!/usr/bin/env python +"""Full-parameter SFT for the orchestrator policy via FSDP (multi-GPU). + +Launch with accelerate so the model + optimizer shard across GPUs (an 8-9B full +fine-tune won't fit one L40S otherwise): + + accelerate launch --config_file \ + scripts/orchestrator/run_sft_fsdp.py \ + --data data/orchestrator/sft/qwen_train_0707.jsonl [--data ...] \ + --variant correct --model Qwen/Qwen3.5-9B \ + --out checkpoints/sft_correct --epochs 3 --batch-size 8 --grad-accum 8 \ + --wandb-project orchestrator-sft + +--variant selects which trajectories to train on: + all -> every record + correct -> only records the verifier marked correct + correct_routed -> only correct records that delegated to a model expert + +By default (--require-clean) records explicitly marked clean=False by the clean +gate (bloated / garbled / unrouted) are also dropped; rows missing the flag +(legacy data) are kept. Pass --no-require-clean to skip this filter. + +Reuses the conversation->tokens + assistant-only masking from sft_tokenize.py. +On these no-NVLink L40S, the accelerate/FSDP NCCL env (NCCL_P2P_DISABLE=1 etc.) +must be set by the launcher. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import re +import sys +import time +from pathlib import Path + +# reuse the data builder from the LoRA launcher +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from sft_tokenize import build_example # noqa: E402 + +MODEL_EXPERTS = { + "gpt_5_5", + "claude_opus_4_8", + "qwen3_5_9b", + "qwen3_6_27b_fp8", + "qwen3_5_122b_a10b_fp8", + "qwen3_5_397b_a17b_fp8", +} + + +def record_is_correct(r: dict) -> bool: + c = r.get("correct") + if c is None: + c = r.get("metrics", {}).get("correct") + return bool(c) + + +def record_routed_to_expert(r: dict) -> bool: + """True if the trajectory called a model expert (resolving anon labels).""" + amap = r.get("metrics", {}).get("anon_map", {}) + names = re.findall( + r'"name"\s*:\s*"([a-z0-9_]+)"', json.dumps(r.get("conversations", [])) + ) + for n in names: + real = amap.get(n, n) + if real in MODEL_EXPERTS: + return True + return False + + +def record_clean_flag(r: dict): + """The record's `clean` flag: True/False, or None if absent (legacy data). + + None means the clean gate never ran on this row, so we can't judge it -> we + keep it (legacy behavior) rather than silently dropping older datasets. + """ + c = r.get("clean") + if c is None: + c = r.get("metrics", {}).get("clean") + return c + + +def select(records, variant: str, require_clean: bool = True): + """Filter by --variant, then (default) drop records explicitly marked unclean. + + require_clean drops rows whose `clean` flag is False. Rows missing the flag + (legacy) are kept regardless, matching pre-clean-gate behavior. + """ + if variant == "all": + base = list(records) + elif variant == "correct": + base = [r for r in records if record_is_correct(r)] + elif variant == "correct_routed": + base = [ + r for r in records if record_is_correct(r) and record_routed_to_expert(r) + ] + else: + raise ValueError(f"unknown variant {variant}") + if not require_clean: + return base + return [r for r in base if record_clean_flag(r) is not False] + + +def main() -> int: + p = argparse.ArgumentParser() + p.add_argument( + "--data", action="append", required=True, help="JSONL(s) (repeatable)" + ) + p.add_argument( + "--val-data", + default=None, + help="Held-out JSONL for val-loss (excluded from --data). Eval'd per epoch.", + ) + p.add_argument( + "--variant", choices=["all", "correct", "correct_routed"], default="correct" + ) + p.add_argument( + "--require-clean", + action=argparse.BooleanOptionalAction, + default=True, + help="Drop records whose `clean` flag is False (bloated / garbled / " + "unrouted). Rows missing the flag (legacy data) are kept. " + "Use --no-require-clean to disable.", + ) + p.add_argument("--model", default="Qwen/Qwen3.5-9B") + p.add_argument("--out", required=True) + p.add_argument("--epochs", type=float, default=3.0) + p.add_argument("--batch-size", type=int, default=8) + p.add_argument("--grad-accum", type=int, default=8) + p.add_argument("--lr", type=float, default=1e-5) + p.add_argument("--max-seq", type=int, default=8192) + p.add_argument( + "--supervise-last-only", + action="store_true", + help="Legacy: supervise ONLY the final assistant turn. Default " + "(off) supervises every assistant turn incl. routing.", + ) + p.add_argument("--warmup-ratio", type=float, default=0.03) + p.add_argument("--seed", type=int, default=42) + p.add_argument("--wandb-project", default="orchestrator-sft") + p.add_argument("--wandb-name", default="") + args = p.parse_args() + + import random + from datetime import timedelta + + import torch + from accelerate import Accelerator + from accelerate.utils import InitProcessGroupKwargs, set_seed + from torch.utils.data import DataLoader + from transformers import AutoModelForCausalLM, AutoTokenizer + + set_seed(args.seed) + # 30-min collective timeout: a slow first-step shard gather / graph build on a + # contended node was tripping the default 10-min NCCL watchdog (job 3892 died + # at step 1). This only delays a *real* hang's crash; it fixes spurious ones. + pg_timeout_min = int(os.environ.get("SFT_PG_TIMEOUT_MIN", "30")) + accelerator = Accelerator( + gradient_accumulation_steps=args.grad_accum, + kwargs_handlers=[ + InitProcessGroupKwargs(timeout=timedelta(minutes=pg_timeout_min)) + ], + ) + is_main = accelerator.is_main_process + + def log(m): + if is_main: + print(f"[fsdp-sft] {time.strftime('%H:%M:%S')} {m}", flush=True) + + tok = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True) + if tok.pad_token_id is None: + tok.pad_token = tok.eos_token + + # ---- data ---- + records = [] + for f in args.data: + records += [json.loads(l) for l in open(f) if l.strip()] + sel = select(records, args.variant, require_clean=args.require_clean) + log(f"variant={args.variant}: {len(sel)}/{len(records)} records") + if args.require_clean: + pre_clean = select(records, args.variant, require_clean=False) + n_unclean = len(pre_clean) - len(sel) + log( + f"require_clean=True: dropped {n_unclean} unclean rows " + f"({len(pre_clean)}->{len(sel)}); use --no-require-clean to keep them" + ) + else: + log("require_clean=False: NOT filtering on `clean` flag") + examples = [] + for r in sel: + ex = build_example( + tok, + r.get("conversations", []), + args.max_seq, + supervise_all_turns=not args.supervise_last_only, + ) + if ex: + examples.append(ex) + log(f"built {len(examples)} training examples") + if not examples: + log("no examples; abort") + return 1 + random.Random(args.seed).shuffle(examples) + + def collate(batch): + maxlen = max(len(b["input_ids"]) for b in batch) + ii, lab, am = [], [], [] + for b in batch: + pad = maxlen - len(b["input_ids"]) + ii.append(b["input_ids"] + [tok.pad_token_id] * pad) + lab.append(b["labels"] + [-100] * pad) + am.append([1] * len(b["input_ids"]) + [0] * pad) + return (torch.tensor(ii), torch.tensor(lab), torch.tensor(am)) + + dl = DataLoader( + examples, batch_size=args.batch_size, shuffle=True, collate_fn=collate + ) + + # ---- held-out val set (leak-free; excluded from --data upstream) ---- + val_dl = None + if args.val_data: + vrecs = [json.loads(l) for l in open(args.val_data) if l.strip()] + vex = [] + for r in vrecs: + ex = build_example( + tok, + r.get("conversations", []), + args.max_seq, + supervise_all_turns=not args.supervise_last_only, + ) + if ex: + vex.append(ex) + log(f"val: {len(vex)} examples from {args.val_data}") + if vex: + val_dl = DataLoader( + vex, batch_size=args.batch_size, shuffle=False, collate_fn=collate + ) + + model = AutoModelForCausalLM.from_pretrained( + args.model, + dtype=torch.bfloat16, + trust_remote_code=True, + attn_implementation="sdpa", + ) + if os.environ.get("SFT_NO_GRAD_CKPT") == "1": + log("gradient checkpointing DISABLED (SFT_NO_GRAD_CKPT=1)") + else: + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + model.config.use_cache = False + opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=0.0) + + model, opt, dl = accelerator.prepare(model, opt, dl) + if val_dl is not None: + val_dl = accelerator.prepare(val_dl) + + @torch.no_grad() + def evaluate(): + # Corpus-level mean token loss over the holdout. FSDP needs every rank to + # run the forward collectively, so each rank evals its shard and we reduce + # token-weighted sums (not a mean-of-batch-means, which would mis-weight). + model.eval() + tot_loss = torch.zeros((), device=accelerator.device) + tot_tok = torch.zeros((), device=accelerator.device) + for ii, lab, am in val_dl: + out = model(input_ids=ii, attention_mask=am, labels=lab) + ntok = (lab != -100).sum() + tot_loss += out.loss.detach() * ntok + tot_tok += ntok + tot_loss = accelerator.reduce(tot_loss, reduction="sum") + tot_tok = accelerator.reduce(tot_tok, reduction="sum") + model.train() + return (tot_loss / tot_tok.clamp(min=1)).item() + + steps_per_epoch = math.ceil(len(dl) / args.grad_accum) + total_steps = max(1, int(steps_per_epoch * args.epochs)) + warmup = max(1, int(total_steps * args.warmup_ratio)) + + def lr_at(s): + if s < warmup: + return s / warmup + return 0.5 * ( + 1 + math.cos(math.pi * (s - warmup) / max(1, total_steps - warmup)) + ) + + use_wandb = False + if is_main: + try: + import wandb + + base_name = ( + args.wandb_name or f"{Path(args.model).name}-{args.variant}-fsdp" + ) + run_name = f"{base_name}-{time.strftime('%m%d-%H%M', time.gmtime())}" + wandb.init( + project=args.wandb_project, + name=run_name, + config=vars(args) + | {"n_examples": len(examples), "total_steps": total_steps}, + ) + use_wandb = True + except Exception as e: + log(f"wandb off ({e})") + + log(f"total_steps={total_steps} steps/epoch={steps_per_epoch} warmup={warmup}") + + def _save_ckpt(out_path): + # Full-model checkpoint (gathers FSDP shards to rank0). Called per-epoch + # so a long run leaves usable partial models + lets us eval intermediate + # checkpoints (e.g. epoch1/epoch2) for the data-scaling sweep. + accelerator.wait_for_everyone() + unwrapped = accelerator.unwrap_model(model) + p = Path(out_path) + if is_main: + p.mkdir(parents=True, exist_ok=True) + tok.save_pretrained(str(p)) + unwrapped.save_pretrained( + str(p), + is_main_process=is_main, + save_function=accelerator.save, + state_dict=accelerator.get_state_dict(model), + ) + if is_main: + log(f"checkpoint saved -> {p}") + + model.train() + gstep = 0 + t0 = time.time() + if val_dl is not None: + vloss = evaluate() + log(f"val/loss (baseline, step 0) = {vloss:.4f}") + if use_wandb: + wandb.log({"val/loss": vloss}, step=0) + for epoch in range(math.ceil(args.epochs)): + for ii, lab, am in dl: + with accelerator.accumulate(model): + out = model(input_ids=ii, attention_mask=am, labels=lab) + accelerator.backward(out.loss) + if accelerator.sync_gradients: + lr = args.lr * lr_at(gstep) + for g in opt.param_groups: + g["lr"] = lr + accelerator.clip_grad_norm_(model.parameters(), 1.0) + opt.step() + opt.zero_grad() + if accelerator.sync_gradients: + gstep += 1 + # accelerator.gather() is a COLLECTIVE — EVERY rank must call it. + # Guarding it behind `is_main` made only rank0 run the [1]-float + # all-gather while ranks 1-3 marched on to the next step's FSDP + # param all-gather, desyncing the process group and hanging the job + # at the step1->step2 boundary (NCCL collective-shape mismatch). + loss = accelerator.gather(out.loss.detach()).mean().item() + if is_main: + log(f"step {gstep}/{total_steps} loss={loss:.4f} lr={lr:.2e}") + if use_wandb: + wandb.log({"train/loss": loss, "train/lr": lr}, step=gstep) + if gstep >= total_steps: + break + _save_ckpt(Path(args.out) / f"epoch{epoch + 1}") + if val_dl is not None: + vloss = evaluate() + log(f"val/loss (epoch {epoch + 1}) = {vloss:.4f}") + if use_wandb: + wandb.log({"val/loss": vloss, "epoch": epoch + 1}, step=gstep) + if gstep >= total_steps: + break + + accelerator.wait_for_everyone() + log("saving full model...") + unwrapped = accelerator.unwrap_model(model) + out_dir = Path(args.out) + if is_main: + out_dir.mkdir(parents=True, exist_ok=True) + tok.save_pretrained(str(out_dir)) + unwrapped.save_pretrained( + str(out_dir), + is_main_process=is_main, + save_function=accelerator.save, + state_dict=accelerator.get_state_dict(model), + ) + if is_main: + log(f"saved -> {out_dir}") + if use_wandb: + try: + wandb.finish() + except Exception: + pass + print( + "FSDP_SFT_DONE " + + json.dumps( + { + "variant": args.variant, + "steps": gstep, + "wall_s": round(time.time() - t0, 1), + "out": str(out_dir), + } + ) + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/orchestrator/sft_tokenize.py b/scripts/orchestrator/sft_tokenize.py new file mode 100644 index 00000000..8b5e6666 --- /dev/null +++ b/scripts/orchestrator/sft_tokenize.py @@ -0,0 +1,151 @@ +#!/usr/bin/env python +"""Conversation->tokens + assistant-only masking shared by the SFT trainers. + +Extracted from the old LoRA trainer so ``run_sft_fsdp.py`` (full-parameter FSDP, +the path we actually use) doesn't depend on it. ``build_example`` tokenizes one +``conversations`` record and returns ``input_ids`` + ``labels`` with everything +but the supervised assistant turns masked to -100. +""" + +from __future__ import annotations + +from typing import Any, Dict, List, Optional + + +def normalize_messages(conversations: List[Dict[str, Any]]) -> List[Dict[str, str]]: + """Map raw conversation turns to {role, content} with role in + system/user/assistant. tool turns fold into user (matches the native + OrchestratorSFTDataset fallback).""" + msgs: List[Dict[str, str]] = [] + for turn in conversations: + role = turn.get("role") or turn.get("from", "") + content = turn.get("content") or turn.get("value", "") or "" + if role in ("human", "user"): + role = "user" + elif role in ("gpt", "assistant"): + role = "assistant" + elif role == "tool": + content = f"[Tool '{turn.get('name', 'tool')}' returned]: {content}" + role = "user" + if role in ("system", "user", "assistant"): + msgs.append({"role": role, "content": str(content)}) + return msgs + + +def _lcp_len(a: List[int], b: List[int]) -> int: + """Length of the longest common prefix of two token-id lists.""" + n = 0 + for x, y in zip(a, b): + if x != y: + break + n += 1 + return n + + +def _chatml_assistant_spans(ids: List[int], tok) -> Optional[List[tuple]]: + """``(start, end)`` token spans of each assistant turn's content in a + ChatML-rendered sequence, located by ``<|im_start|>assistant ... <|im_end|>`` + markers (the ``<|im_end|>`` is included so the model learns to stop). + + Operates on the FINAL token stream, so it is correct even when the chat + template strips prior-turn ```` reasoning from history (Qwen3.x does), + which breaks any incremental re-rendering / prefix approach. Returns None for + non-ChatML templates (e.g. gemma) so the caller can fall back.""" + im_start = tok.convert_tokens_to_ids("<|im_start|>") + im_end = tok.convert_tokens_to_ids("<|im_end|>") + asst = tok.convert_tokens_to_ids("assistant") + unk = tok.unk_token_id + if im_start is None or im_end is None or im_start == unk or im_end == unk: + return None + nl = tok.convert_tokens_to_ids("Ċ") # the '\n' after the role header + spans: List[tuple] = [] + i, n = 0, len(ids) + while i < n: + if ids[i] == im_start and i + 1 < n and ids[i + 1] == asst: + j = i + 2 + if j < n and ids[j] == nl: + j += 1 + k = j + while k < n and ids[k] != im_end: + k += 1 + end = min(k + 1, n) + if end > j: + spans.append((j, end)) + i = end + else: + i += 1 + return spans + + +def build_example( + tok, + conversations, + max_seq: int, + supervise_all_turns: bool = True, +) -> Optional[Dict[str, List[int]]]: + """Tokenize one record -> input_ids + labels. + + ``supervise_all_turns=True`` (default): supervise EVERY assistant turn — the + intermediate routing ```` turns AND the final answer; only + system/user/tool tokens are masked. This teaches the model to actually route, + not just synthesize the last answer given someone else's tool calls. + + ``supervise_all_turns=False`` (legacy): supervise only the final assistant + turn; everything before it is masked to -100. + + Assistant spans are found by the ChatML turn markers on the rendered stream + (robust to Qwen3's prior-turn think-stripping). For non-ChatML templates + (gemma) we fall back to the proven last-turn longest-common-prefix boundary; + all-turns supervision there needs per-template markers and is not yet wired, + so it degrades to last-turn. Records whose last turn isn't assistant, that + supervise nothing, or whose final turn alone overflows max_seq are skipped.""" + msgs = normalize_messages(conversations) + if len(msgs) < 2 or msgs[-1]["role"] != "assistant": + return None + try: + full_text = tok.apply_chat_template( + msgs, tokenize=False, add_generation_prompt=False + ) + full = tok(full_text, add_special_tokens=False)["input_ids"] + except Exception: + return None + if not full: + return None + + spans = _chatml_assistant_spans(full, tok) + if spans: + selected = spans if supervise_all_turns else spans[-1:] + labels = [-100] * len(full) + for s, e in selected: + for k in range(s, e): + labels[k] = full[k] + last_start = selected[-1][0] + else: + # Non-ChatML template: proven last-turn boundary via longest-common-prefix + # (robust to templates whose add_generation_prompt emits extra preamble). + try: + prompt_text = tok.apply_chat_template( + msgs[:-1], tokenize=False, add_generation_prompt=True + ) + prompt = tok(prompt_text, add_special_tokens=False)["input_ids"] + except Exception: + return None + boundary = _lcp_len(prompt, full) + if boundary >= len(full): + return None + labels = list(full) + for i in range(min(boundary, len(labels))): + labels[i] = -100 + last_start = boundary + + if all(v == -100 for v in labels): # nothing left to supervise + return None + + # Keep the final assistant turn if too long: left-truncate from the front. + if len(full) > max_seq: + if len(full) - last_start >= max_seq: # final answer alone overflows -> drop + return None + cut = len(full) - max_seq + full = full[cut:] + labels = labels[cut:] + return {"input_ids": full, "labels": labels} diff --git a/scripts/train/fsdp4.yaml b/scripts/train/fsdp4.yaml new file mode 100644 index 00000000..1620634a --- /dev/null +++ b/scripts/train/fsdp4.yaml @@ -0,0 +1,25 @@ +compute_environment: LOCAL_MACHINE +distributed_type: FSDP +downcast_bf16: 'no' +machine_rank: 0 +main_process_ip: null +main_process_port: null +main_training_function: main +mixed_precision: bf16 +num_machines: 1 +num_processes: 4 +rdzv_backend: static +same_network: true +use_cpu: false +fsdp_config: + fsdp_version: 1 + fsdp_auto_wrap_policy: SIZE_BASED_WRAP + fsdp_min_num_params: 100000000 + fsdp_sharding_strategy: FULL_SHARD + fsdp_state_dict_type: FULL_STATE_DICT + fsdp_offload_params: false + fsdp_backward_prefetch: BACKWARD_PRE + fsdp_forward_prefetch: false + fsdp_use_orig_params: true + fsdp_cpu_ram_efficient_loading: true + fsdp_sync_module_states: true diff --git a/scripts/train/fsdp8.yaml b/scripts/train/fsdp8.yaml new file mode 100644 index 00000000..95de8ce0 --- /dev/null +++ b/scripts/train/fsdp8.yaml @@ -0,0 +1,25 @@ +compute_environment: LOCAL_MACHINE +distributed_type: FSDP +downcast_bf16: 'no' +machine_rank: 0 +main_process_ip: null +main_process_port: null +main_training_function: main +mixed_precision: bf16 +num_machines: 1 +num_processes: 8 +rdzv_backend: static +same_network: true +use_cpu: false +fsdp_config: + fsdp_version: 1 + fsdp_auto_wrap_policy: SIZE_BASED_WRAP + fsdp_min_num_params: 100000000 + fsdp_sharding_strategy: FULL_SHARD + fsdp_state_dict_type: FULL_STATE_DICT + fsdp_offload_params: false + fsdp_backward_prefetch: BACKWARD_PRE + fsdp_forward_prefetch: false + fsdp_use_orig_params: true + fsdp_cpu_ram_efficient_loading: true + fsdp_sync_module_states: true diff --git a/src/openjarvis/agents/orchestrator.py b/src/openjarvis/agents/orchestrator.py index cca03d2c..05964f64 100644 --- a/src/openjarvis/agents/orchestrator.py +++ b/src/openjarvis/agents/orchestrator.py @@ -13,6 +13,7 @@ Supports two modes: from __future__ import annotations import concurrent.futures +import json import re from typing import Any, List, Optional @@ -201,6 +202,70 @@ class OrchestratorAgent(ToolUsingAgent): return result + @staticmethod + def _boxed_tool_call(content, openai_tools, call_id): + r"""Salvage a tool call the model wrote as ``\boxed{tool: query}``. + + Some SFT'd checkpoints (Qwen reasoning-mode habit) emit their action + as a boxed string instead of a real ```` — e.g. + ``\boxed{web_search: latest news}`` or ``\boxed{web_search query: ...}``. + The engine's tool-call parser sees no structured call, so without this + the agent would grade the half-finished thought as its final answer and + never actually act. Here we recover ``(tool, args)`` and build a real + ``ToolCall`` keyed on the tool's primary parameter. Returns ``None`` when + the boxed content isn't a recognisable tool action (leaving genuine + boxed final answers untouched). + """ + if not openai_tools or "\\boxed" not in content: + return None + + # tool name -> primary parameter (first required, else first property) + param_of: dict[str, str] = {} + for t in openai_tools: + fn = t.get("function", t) + name = fn.get("name") + if not name: + continue + params = fn.get("parameters") or {} + props = list((params.get("properties") or {}).keys()) + req = params.get("required") or [] + param_of[name] = req[0] if req else (props[0] if props else "query") + + # The action is the LAST boxed expression in the response. + boxed = re.findall(r"\\boxed\{([^{}]*)\}", content) + if not boxed: + return None + inner = boxed[-1].strip() + + # Match a leading tool name (longest first to avoid prefix collisions), + # then strip the separator between name and args. Handles + # ``tool: q`` / ``tool query: q`` / ``tool q`` / ``tool(q)``. + for name in sorted(param_of, key=len, reverse=True): + m = re.match( + r"^\s*" + re.escape(name) + r"\b(.*)$", + inner, + re.DOTALL | re.IGNORECASE, + ) + if not m: + continue + rest = m.group(1).strip() + if rest.startswith("(") and rest.endswith(")"): + rest = rest[1:-1] # tool(args) wrapper -> drop matching parens + else: + # optional "query" word, then a ":" / "-" separator + rest = re.sub( + r"^\s*(?:query\s*)?[:\-]\s*", "", rest, flags=re.IGNORECASE + ) + arg = rest.strip().strip("\"'").strip() + if not arg: + return None + return ToolCall( + id=call_id, + name=name, + arguments=json.dumps({param_of[name]: arg}), + ) + return None + # ------------------------------------------------------------------ # Function-calling mode (original behaviour) # ------------------------------------------------------------------ @@ -245,8 +310,31 @@ class OrchestratorAgent(ToolUsingAgent): content = result.get("content", "") raw_tool_calls = result.get("tool_calls", []) - # No tool calls -> check continuation, then final answer + # No tool calls -> try to salvage a \boxed{tool: query} action the + # model wrote instead of a real ; else final answer. if not raw_tool_calls: + boxed_call = self._boxed_tool_call( + content, openai_tools, f"boxed_{turns}" + ) + if boxed_call is not None: + messages.append( + Message( + role=Role.ASSISTANT, + content=content, + tool_calls=[boxed_call], + ) + ) + tool_result = self._executor.execute(boxed_call) + all_tool_results.append(tool_result) + messages.append( + Message( + role=Role.TOOL, + content=tool_result.content, + tool_call_id=boxed_call.id, + name=boxed_call.name, + ) + ) + continue content = self._check_continuation(result, messages) content = self._strip_think_tags(content) self._emit_turn_end(turns=turns, content_length=len(content)) diff --git a/src/openjarvis/evals/core/scorer.py b/src/openjarvis/evals/core/scorer.py index cb95690a..bc02bd96 100644 --- a/src/openjarvis/evals/core/scorer.py +++ b/src/openjarvis/evals/core/scorer.py @@ -2,12 +2,44 @@ from __future__ import annotations +import logging +import random +import time from abc import ABC, abstractmethod from typing import Any, Dict, Optional, Tuple from openjarvis.evals.core.backend import InferenceBackend from openjarvis.evals.core.types import EvalRecord +LOGGER = logging.getLogger(__name__) + +# Transient failures worth retrying the judge call on. A single rate-limit +# (429) used to fall straight through to `llm_fallback_error` and score the +# sample WRONG — e.g. 93/100 GAIA judge calls 429'd on one run and zeroed the +# whole bench. Retry with exponential backoff so a busy judge endpoint doesn't +# get mis-scored as a wrong answer. +_JUDGE_MAX_RETRIES = 6 +_JUDGE_BASE_DELAY_S = 2.0 +_JUDGE_MAX_DELAY_S = 60.0 +_RETRYABLE_MARKERS = ( + "429", + "rate_limit", + "rate limit", + "overloaded", + "timeout", + "timed out", + "503", + "502", + "500", + "connection", + "temporarily unavailable", +) + + +def _is_retryable_judge_error(exc: Exception) -> bool: + msg = str(exc).lower() + return any(marker in msg for marker in _RETRYABLE_MARKERS) + class Scorer(ABC): """Base class for all scorers.""" @@ -42,14 +74,41 @@ class LLMJudgeScorer(Scorer): temperature: float = 0.0, max_tokens: int = 2048, ) -> str: - """Send a prompt to the judge LLM and return the response text.""" - return self._judge_backend.generate( - prompt, - model=self._judge_model, - system=system, - temperature=temperature, - max_tokens=max_tokens, - ) + """Send a prompt to the judge LLM and return the response text. + + Retries transient failures (429 / rate-limit / 5xx / timeout) with + exponential backoff + jitter so a busy judge endpoint doesn't get + mis-scored as a wrong answer. Non-retryable errors, or exhaustion of + the retry budget, re-raise to the caller (which records the failure). + """ + last_exc: Optional[Exception] = None + for attempt in range(_JUDGE_MAX_RETRIES): + try: + return self._judge_backend.generate( + prompt, + model=self._judge_model, + system=system, + temperature=temperature, + max_tokens=max_tokens, + ) + except Exception as exc: # noqa: BLE001 - re-raised below + last_exc = exc + if attempt == _JUDGE_MAX_RETRIES - 1 or not _is_retryable_judge_error( + exc + ): + raise + delay = min(_JUDGE_BASE_DELAY_S * (2**attempt), _JUDGE_MAX_DELAY_S) + delay += random.uniform(0.0, delay * 0.25) # jitter + LOGGER.warning( + "judge call failed (attempt %d/%d): %s — retrying in %.1fs", + attempt + 1, + _JUDGE_MAX_RETRIES, + exc, + delay, + ) + time.sleep(delay) + # Unreachable (loop either returns or raises), but keeps type-checkers happy. + raise last_exc # type: ignore[misc] __all__ = ["LLMJudgeScorer", "Scorer"] diff --git a/src/openjarvis/evals/scorers/gaia_exact.py b/src/openjarvis/evals/scorers/gaia_exact.py index 5b1c23a5..9b152eef 100644 --- a/src/openjarvis/evals/scorers/gaia_exact.py +++ b/src/openjarvis/evals/scorers/gaia_exact.py @@ -140,7 +140,14 @@ class GAIAScorer(LLMJudgeScorer): if structured_match: is_correct = structured_match.group(1).lower() == "yes" else: - is_correct = "CORRECT" in raw.upper() and "INCORRECT" not in raw.upper() + up = raw.upper() + # "not correct" / "is not correct" contain "CORRECT" and lack + # "INCORRECT" — guard the negations explicitly or they read True. + is_correct = ( + "CORRECT" in up + and "INCORRECT" not in up + and "NOT CORRECT" not in up + ) meta: Dict[str, Any] = { "match_type": "llm_fallback", diff --git a/src/openjarvis/evals/scorers/mmlu_pro_mcq.py b/src/openjarvis/evals/scorers/mmlu_pro_mcq.py index 75d9bacf..6a6d9553 100644 --- a/src/openjarvis/evals/scorers/mmlu_pro_mcq.py +++ b/src/openjarvis/evals/scorers/mmlu_pro_mcq.py @@ -27,6 +27,25 @@ class MMLUProScorer(LLMJudgeScorer): return "".join(chr(ord("A") + i) for i in range(n)) return "ABCDEFGHIJ" + def _extract_answer_direct( + self, + model_answer: str, + valid_letters: str, + ) -> Optional[str]: + """High-confidence, unambiguous letter straight from the output. + + Only an explicit ``\\boxed{X}`` counts here. We deliberately do NOT regex + prose like "the answer is X": the letters "I" and "A" are real English + words, so "FINAL_ANSWER: I need ..." would grab "I" — the exact bug that + mislabels answers. Semantic extraction is left to the judge below. + """ + if not model_answer: + return None + for cand in reversed(re.findall(r"\\boxed\{\s*([A-Za-z])\s*\}", model_answer)): + if cand.upper() in valid_letters: + return cand.upper() + return None + def _extract_answer_with_llm( self, problem: str, @@ -56,10 +75,13 @@ class MMLUProScorer(LLMJudgeScorer): ) extracted = raw_response.strip().upper() + if "NONE" in extracted: + return None - # Handle "The answer is: A" etc. + # Accept only an ISOLATED letter — never the first cap of a word like + # "It"/"None". Prefer an explicit "the answer is X" phrasing. answer_match = re.search( - r"(?:THE ANSWER IS:?\s*)?([A-Z])", + r"(?:THE ANSWER IS:?\s*)?\b([A-Z])\b(?![A-Za-z'])", extracted, re.IGNORECASE, ) @@ -86,11 +108,17 @@ class MMLUProScorer(LLMJudgeScorer): valid_letters = self._valid_letters_from_options(record.metadata) - candidate = self._extract_answer_with_llm( - record.problem, - model_answer, - valid_letters, - ) + # Parse straight from the output first (fast, deterministic, no "I" bug); + # only fall back to the judge when no isolated letter is present. + method = "direct" + candidate = self._extract_answer_direct(model_answer, valid_letters) + if not candidate: + method = "llm" + candidate = self._extract_answer_with_llm( + record.problem, + model_answer, + valid_letters, + ) if not candidate: return None, {"reason": "no_choice_letter_extracted"} @@ -99,6 +127,7 @@ class MMLUProScorer(LLMJudgeScorer): "reference_letter": ref, "candidate_letter": candidate, "valid_letters": valid_letters, + "extract_method": method, } return is_correct, meta diff --git a/src/openjarvis/evals/scorers/supergpqa_mcq.py b/src/openjarvis/evals/scorers/supergpqa_mcq.py index da66a003..5595e91a 100644 --- a/src/openjarvis/evals/scorers/supergpqa_mcq.py +++ b/src/openjarvis/evals/scorers/supergpqa_mcq.py @@ -27,6 +27,25 @@ class SuperGPQAScorer(LLMJudgeScorer): return "".join(chr(ord("A") + i) for i in range(n)) return "ABCD" + def _extract_answer_direct( + self, + model_answer: str, + valid_letters: str, + ) -> Optional[str]: + """High-confidence, unambiguous letter straight from the output. + + Only an explicit ``\\boxed{X}`` counts here. We deliberately do NOT regex + prose like "the answer is X": the letters "I" and "A" are real English + words, so "FINAL_ANSWER: I need ..." would grab "I" — the exact bug that + mislabels answers. Semantic extraction is left to the judge below. + """ + if not model_answer: + return None + for cand in reversed(re.findall(r"\\boxed\{\s*([A-Za-z])\s*\}", model_answer)): + if cand.upper() in valid_letters: + return cand.upper() + return None + def _extract_answer_with_llm( self, problem: str, @@ -56,10 +75,13 @@ class SuperGPQAScorer(LLMJudgeScorer): ) extracted = raw_response.strip().upper() + if "NONE" in extracted: + return None - # Handle "The answer is: A" etc. + # Accept only an ISOLATED letter — never the first cap of a word like + # "It"/"None". Prefer an explicit "the answer is X" phrasing. answer_match = re.search( - r"(?:THE ANSWER IS:?\s*)?([A-Z])", + r"(?:THE ANSWER IS:?\s*)?\b([A-Z])\b(?![A-Za-z'])", extracted, re.IGNORECASE, ) @@ -86,11 +108,17 @@ class SuperGPQAScorer(LLMJudgeScorer): valid_letters = self._valid_letters_from_options(record.metadata) - candidate = self._extract_answer_with_llm( - record.problem, - model_answer, - valid_letters, - ) + # Parse straight from the output first (fast, deterministic, no "I" bug); + # only fall back to the judge when no isolated letter is present. + method = "direct" + candidate = self._extract_answer_direct(model_answer, valid_letters) + if not candidate: + method = "llm" + candidate = self._extract_answer_with_llm( + record.problem, + model_answer, + valid_letters, + ) if not candidate: return None, {"reason": "no_choice_letter_extracted"} @@ -99,6 +127,7 @@ class SuperGPQAScorer(LLMJudgeScorer): "reference_letter": ref, "candidate_letter": candidate, "valid_letters": valid_letters, + "extract_method": method, } return is_correct, meta diff --git a/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3.5-9b-sft.toml b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3.5-9b-sft.toml new file mode 100644 index 00000000..f3d325b0 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3.5-9b-sft.toml @@ -0,0 +1,34 @@ +# Orchestrator: Qwen3.5-9B. v1 = base self-sampling (the served orchestrator_model is this same base). +# +# v1 data = base Qwen3.5-9B self-sampling: roll the base model out over +# load_sft_tasks() (~8K reasoning tasks), keep correct trajectories. Build the +# dataset first (no GPU, needs the vLLM orchestrator endpoint up): +# .venv/bin/python scripts/orchestrator/build_orchestrator_sft.py \ +# --out data/orchestrator_qwen3.5-9b_sft_v1.jsonl --samples-per-task 8 +# then train on a rented GPU with regenerate_traces = false. + +[model] +model_name = "Qwen/Qwen3.5-9B" +max_seq_length = 4096 + +[training] +num_epochs = 3 +batch_size = 64 # 8000 traj / 64 ≈ 125 batches/epoch +learning_rate = 2e-5 +weight_decay = 0.01 +warmup_ratio = 0.1 +gradient_checkpointing = true + +[data] +trace_cache_path = "data/orchestrator_qwen3.5-9b_sft_v1.jsonl" +regenerate_traces = false +# Base Qwen3.5-9B self-sampling over load_sft_tasks() (see sft_data/reject_sample.py). +orchestrator_endpoint = "http://localhost:8001/v1" +orchestrator_model = "qwen3.5-9b" +samples_per_task = 8 +max_keep_per_task = 1 +distill_max_tasks = 0 # 0 -> all of load_sft_tasks() + +[checkpoint] +checkpoint_dir = "checkpoints/orchestrator_qwen3.5-9b_sft" +save_every_n_epochs = 1 diff --git a/src/openjarvis/learning/intelligence/orchestrator/eval_backend.py b/src/openjarvis/learning/intelligence/orchestrator/eval_backend.py new file mode 100644 index 00000000..bf682ab3 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/eval_backend.py @@ -0,0 +1,288 @@ +"""Eval backend that drives the ToolOrchestra orchestrator as the "model". + +This adapts our unified-tool orchestrator rollout +(:func:`openjarvis.agents.hybrid.toolorchestra.rollout.run_unified_rollout`) +to the eval framework's :class:`~openjarvis.evals.core.backend.InferenceBackend` +interface so it can be scored by the existing +:class:`~openjarvis.evals.core.runner.EvalRunner` over any registered benchmark. + +The orchestrator is an OpenAI-compatible chat model (served locally via vLLM, +default ``http://localhost:8001/v1``) that emits tool calls over the fixed +8-tool catalog from :func:`expert_registry.orchestrator_catalog`. For each eval +sample we run one full rollout and return ``rollout.final_answer`` as the model +answer, plus token/cost telemetry in the ``generate_full`` payload. + +Construction is import-safe and network-free; the OpenAI SDK is only touched +inside ``call_orchestrator`` at generate time (lazy import in ``unified.py``). +""" + +from __future__ import annotations + +import re +import time +from typing import Any, Dict, Optional + +from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog +from openjarvis.agents.hybrid.toolorchestra import rollout as _rollout_mod +from openjarvis.agents.hybrid.toolorchestra.unified import ( + make_call_orchestrator, + make_dispatch, +) +from openjarvis.evals.core.backend import InferenceBackend + +# The served fine-tuned model sometimes over-emits the answer marker, e.g. +# "FINAL_ANSWER: FINAL_ANSWER: 42" — a doubled prefix that breaks answer +# extraction and auto-scores the sample 0. Collapse any run of FINAL_ANSWER +# markers to a single one, keeping the text after the LAST marker as the answer. +_FA_MARK = re.compile(r"(?im)FINAL[_\s]?ANSWER\s*:?") + + +def _clean_final_answer(text: str) -> str: + t = (text or "").strip() + marks = list(_FA_MARK.finditer(t)) + if not marks: + return t + answer = t[marks[-1].end() :].strip() + return f"FINAL_ANSWER: {answer}" if answer else t + + +_OBS_CAP = 4000 # cap persisted observations so the eval JSONL stays readable + +DEFAULT_ENDPOINT = "http://localhost:8001/v1" +DEFAULT_MODEL = "qwen3-8b" + + +class OrchestratorBackend(InferenceBackend): + """Run the ToolOrchestra orchestrator rollout as an eval "model". + + Parameters + ---------- + orchestrator_endpoint: + OpenAI-compatible base URL for the served orchestrator (vLLM). Defaults + to ``http://localhost:8001/v1``. + orchestrator_model: + Served model id for the orchestrator. Defaults to ``"qwen3-8b"`` (will + be the Qwen3.5-9B checkpoint later). + api_key: + API key for the endpoint. ``"EMPTY"`` for a local vLLM server. + local_endpoints: + Optional dict mapping a local OSS model id (e.g. ``"Qwen/Qwen3.5-9B"``) + to its vLLM base URL. Unmapped local models are still listed in the + catalog but served as unconfigured (``base_url=None``). + max_turns: + Maximum orchestrator reasoning->action->observation turns per sample. + temperature: + Orchestrator sampling temperature. + dispatch_cfg: + Optional config dict passed to :func:`make_dispatch` (worker dispatch). + """ + + backend_id = "orchestrator" + framework_name = "openjarvis-orchestrator" + + + def _dump_trace(self, prompt, rollout, orch_model) -> None: + """Append one trajectory_to_record row to $OJ_EVAL_TRACE_DIR/traces.jsonl.""" + import hashlib + import json + import os + import threading + from pathlib import Path + + trace_dir = os.environ.get("OJ_EVAL_TRACE_DIR") + if not trace_dir: + return + try: + from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import ( + trajectory_to_record, + ) + + rec = trajectory_to_record( + hashlib.md5(prompt.encode()).hexdigest()[:12], + prompt, + self._tools, + rollout, + ) + rec["orchestrator_model"] = orch_model + Path(trace_dir).mkdir(parents=True, exist_ok=True) + if not hasattr(self, "_trace_lock"): + self._trace_lock = threading.Lock() + with self._trace_lock: + with open(Path(trace_dir) / "traces.jsonl", "a") as fh: + fh.write(json.dumps(rec, ensure_ascii=False) + "\n") + except Exception: # noqa: BLE001 — tracing must never fail an eval sample + pass + + + def __init__( + self, + *, + orchestrator_endpoint: str = DEFAULT_ENDPOINT, + orchestrator_model: str = DEFAULT_MODEL, + api_key: str = "EMPTY", + local_endpoints: Optional[Dict[str, str]] = None, + max_turns: int = 8, + temperature: float = 1.0, + finetuned: bool = True, + dispatch_cfg: Optional[Dict[str, Any]] = None, + ) -> None: + self.orchestrator_endpoint = orchestrator_endpoint + self.orchestrator_model = orchestrator_model + self.api_key = api_key + self.local_endpoints = dict(local_endpoints or {}) + self.max_turns = int(max_turns) + self.temperature = float(temperature) + # False -> serve with native tools= (a BASE model needs the server's tool + # template to emit reliable calls; the baked-JSON prompt is what a + # FINE-TUNED model was trained on). Getting this wrong handicaps the model: + # the first fixed-harness baseline rerun served the base model in baked + # mode and scored 0.184 — not comparable to anything. + self.finetuned = bool(finetuned) + self._dispatch_cfg = dict(dispatch_cfg or {}) + + # Build the catalog once; it is stateless. ``local_endpoints`` now maps + # full local model ids (e.g. ``"Qwen/Qwen3.5-9B"``) -> vLLM base_url. + self._tools = orchestrator_catalog(local_endpoints=self.local_endpoints) + + # ------------------------------------------------------------------ + # InferenceBackend abstract methods + # ------------------------------------------------------------------ + def generate( + self, + prompt: str, + *, + model: str, + system: str = "", + temperature: float = 0.0, + max_tokens: int = 2048, + ) -> str: + """Run a rollout and return just the final-answer text.""" + full = self.generate_full( + prompt, + model=model, + system=system, + temperature=temperature, + max_tokens=max_tokens, + ) + return full.get("content", "") or "" + + def generate_full( + self, + prompt: str, + *, + model: str, + system: str = "", + temperature: float = 0.0, + max_tokens: int = 2048, + ) -> Dict[str, Any]: + """Run one orchestrator rollout over ``prompt`` and return full details. + + Returns a dict with the keys the runner reads: ``content``, ``usage``, + ``model``, ``latency_seconds``, ``cost_usd``, plus ``tool_calls`` / + ``turn_count`` / ``framework``. On error, returns ``content=""`` and a + non-empty ``error`` field so the runner records it as a failed sample + rather than crashing the whole run. + """ + # The caller-supplied ``model`` is the RunConfig.model (the orchestrator + # model name). Fall back to the configured one if blank. + orch_model = model or self.orchestrator_model + # Prefer the explicit per-call temperature only when the caller set a + # non-default value; otherwise use the backend's configured temperature. + temp = temperature if temperature else self.temperature + + started = time.time() + try: + call_orch = make_call_orchestrator( + orch_model, + base_url=self.orchestrator_endpoint, + api_key=self.api_key, + temperature=temp, + max_tokens=max_tokens, + # Fine-tuned model: baked JSON prompt, no tools=. + # Base model (finetuned=False): native tools= — it can't emit + # reliable calls otherwise. + native_tools=not self.finetuned, + ) + dispatch = make_dispatch(self._dispatch_cfg) + rollout = _rollout_mod.run_unified_rollout( + prompt, + self._tools, + call_orchestrator=call_orch, + dispatch=dispatch, + max_turns=self.max_turns, + # Match TRAINING: the model was fine-tuned on anonymized expert + # labels (expert_xxxx), so it only picks valid labels / routes + # correctly when it sees the same anonymized names at inference. + anonymize=True, + ) + # Persist the FULL trajectory when asked (OJ_EVAL_TRACE_DIR). The + # scored JSONL keeps only the final answer — without this every eval + # conversation is discarded and failures can't be audited (e.g. + # "is GAIA 0.20 turn-exhaustion or wrong answers?" was unanswerable + # for the first sft879 run). Rows are the same `conversations` + # records the generation pipeline writes, so render_sft_data.py + # renders them as per-sample markdown. + self._dump_trace(prompt, rollout, orch_model) + except Exception as exc: # noqa: BLE001 - surface as a failed sample + return { + "content": "", + "error": f"{type(exc).__name__}: {exc}", + "usage": {"prompt_tokens": 0, "completion_tokens": 0}, + "model": orch_model, + "latency_seconds": time.time() - started, + "cost_usd": 0.0, + "tool_calls": 0, + "turn_count": 0, + "framework": self.framework_name, + } + + latency = time.time() - started + total_tokens = int(getattr(rollout, "tokens", 0) or 0) + anon_map = getattr(rollout, "anon_map", None) or {} + return { + "content": _clean_final_answer(rollout.final_answer or ""), + # The rollout reports a single combined token count; we surface it + # as completion_tokens so it still flows into total-token telemetry. + "usage": { + "prompt_tokens": 0, + "completion_tokens": total_tokens, + }, + "model": orch_model, + "latency_seconds": latency, + "cost_usd": float(getattr(rollout, "cost_usd", 0.0) or 0.0), + "tool_calls": int(getattr(rollout, "num_tool_calls", 0) or 0), + "turn_count": len(getattr(rollout, "turns", []) or []), + "framework": self.framework_name, + "trace_data": { + "parse_failures": int(getattr(rollout, "parse_failures", 0) or 0), + "tool_calls": [ + {"name": n, "arguments": a} for n, a in rollout.tool_calls() + ], + # Full step-by-step trace so eval samples are inspectable + # (format_eval_sample.py renders this). Observations capped. + "anon_map": anon_map, + "turns": [ + { + "reasoning": t.reasoning or "", + "tool_name": t.tool_name, + "real_model": anon_map.get(t.tool_name) + if t.tool_name + else None, + "arguments": t.arguments, + "observation": ( + (t.observation or "")[:_OBS_CAP] + + ( + "…[truncated]" + if t.observation and len(t.observation) > _OBS_CAP + else "" + ) + ), + } + for t in (getattr(rollout, "turns", []) or []) + ], + "final_answer": _clean_final_answer(rollout.final_answer or ""), + }, + } + + +__all__ = ["OrchestratorBackend", "DEFAULT_ENDPOINT", "DEFAULT_MODEL"] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py b/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py index 2c2e0645..4a41a832 100644 --- a/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py @@ -52,7 +52,7 @@ class OrchestratorSFTConfig: """Configuration for orchestrator SFT training.""" # Model - model_name: str = "Qwen/Qwen3-1.7B" + model_name: str = "Qwen/Qwen3.5-9B" max_seq_length: int = 4096 # Training @@ -63,17 +63,28 @@ class OrchestratorSFTConfig: warmup_ratio: float = 0.1 max_grad_norm: float = 1.0 - # Trace generation - teacher_engine_key: str = "" - teacher_model: str = "" - traces_per_query: int = 2 - max_attempts_per_trace: int = 3 + # Trace generation by base Qwen3-8B self-sampling (v1 cold-start). The + # orchestrator IS the local Qwen3-8B served over an OpenAI-compatible vLLM + # endpoint; we roll it out over load_sft_tasks() and keep correct trajectories. + # For v2, point orchestrator_endpoint/model at the v1 checkpoint's server. + orchestrator_endpoint: str = "http://localhost:8001/v1" + orchestrator_model: str = "qwen3-8b" + orchestrator_api_key: str = "EMPTY" generation_temperature: float = 0.7 # Data source - trace_cache_path: str = "data/orchestrator_sft_traces.jsonl" + trace_cache_path: str = "data/orchestrator_sft_v1.jsonl" regenerate_traces: bool = False + # Rejection-sampling cold-start over load_sft_tasks (see sft_data/reject_sample.py). + distill_max_tasks: int = 0 # 0 / None -> all of load_sft_tasks() + samples_per_task: int = 8 + max_keep_per_task: int = 1 + max_rollout_turns: int = 8 + # Optional local OSS model endpoints, mapping full model id (e.g. + # "Qwen/Qwen3.5-9B") -> vLLM base_url; omitted when unset. + local_endpoints: Dict[str, str] = field(default_factory=dict) + # Checkpoint checkpoint_dir: str = "checkpoints/orchestrator_sft" save_every_n_epochs: int = 1 @@ -267,11 +278,81 @@ class OrchestratorSFTTrainer: ) def _generate_traces(self) -> None: - """Generate SFT traces (placeholder — requires running engine).""" + """Generate SFT traces by base Qwen3-8B self-sampling (v1 cold-start). + + The orchestrator IS the local Qwen3-8B served over an OpenAI-compatible + vLLM endpoint (``orchestrator_endpoint``/``orchestrator_model``). We roll + it out over ``load_sft_tasks()`` on the fixed orchestrator tool catalog, + verify each trajectory's final answer against the gold + (``verify.make_verifier``), keep the cheapest-correct trajectory per task, + and serialize them into the ```` ``conversations`` JSONL the SFT + dataset consumes. For v2, point the endpoint/model at the v1 checkpoint's + server. See ``sft_data/reject_sample.py`` and the + ``scripts/orchestrator/build_orchestrator_sft.py`` driver. + """ + from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog + from openjarvis.agents.hybrid.toolorchestra.rollout import run_unified_rollout + from openjarvis.agents.hybrid.toolorchestra.unified import ( + make_call_orchestrator, + make_dispatch, + ) + from openjarvis.learning.intelligence.orchestrator.sft_data import ( + reject_sample as _reject_sample, + ) + from openjarvis.learning.intelligence.orchestrator.sft_data.datasets import ( + load_sft_tasks, + ) + + generate_sft_dataset = _reject_sample.generate_sft_dataset + from openjarvis.learning.intelligence.orchestrator.sft_data.verify import ( + make_verifier, + ) + trace_path = Path(self.config.trace_cache_path) - trace_path.parent.mkdir(parents=True, exist_ok=True) - if not trace_path.exists(): - trace_path.touch() + tools = orchestrator_catalog( + local_endpoints=self.config.local_endpoints or None, + ) + call_orch = make_call_orchestrator( + self.config.orchestrator_model, + base_url=self.config.orchestrator_endpoint, + api_key=self.config.orchestrator_api_key, + temperature=self.config.generation_temperature, + ) + dispatch = make_dispatch({}) + + def rollout_fn(task: Any) -> Any: + try: + return run_unified_rollout( + task.instruction, + tools, + call_orchestrator=call_orch, + dispatch=dispatch, + max_turns=self.config.max_rollout_turns, + ) + except Exception as exc: # network/key failures shouldn't kill the run + logger.warning("rollout failed for %s: %s", task.task_id, exc) + return None + + try: + tasks = load_sft_tasks() + if self.config.distill_max_tasks: + tasks = tasks[: self.config.distill_max_tasks] + stats = generate_sft_dataset( + str(trace_path), + tasks=tasks, + tools=tools, + rollout_fn=rollout_fn, + verify_fn=make_verifier(), + samples_per_task=self.config.samples_per_task, + max_keep_per_task=self.config.max_keep_per_task, + reward_fn=lambda r: -r.cost_usd, # cheapest-correct gets top reward + ) + logger.info("Generated SFT traces: %s", stats) + except Exception as exc: # network/datasets optional — fail soft to empty + logger.warning("trace generation failed (%s); writing empty set", exc) + trace_path.parent.mkdir(parents=True, exist_ok=True) + if not trace_path.exists(): + trace_path.touch() def _init_optimizer(self) -> None: if not HAS_TORCH or self.policy.model is None: diff --git a/tests/test_orchestrator_learning/test_eval_backend.py b/tests/test_orchestrator_learning/test_eval_backend.py new file mode 100644 index 00000000..d6bdac03 --- /dev/null +++ b/tests/test_orchestrator_learning/test_eval_backend.py @@ -0,0 +1,182 @@ +"""Offline tests for the orchestrator eval backend + the eval CLI's +benchmark-name normalization. + +No network / models: ``run_unified_rollout`` is monkeypatched to return a +canned rollout, and the dataset-key mapping is asserted against the real +``openjarvis.evals.cli`` registry without instantiating any (network-bound) +dataset. +""" + +from __future__ import annotations + +import pytest + +from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedTurn +from openjarvis.learning.intelligence.orchestrator import eval_backend as eb +from openjarvis.learning.intelligence.orchestrator.eval_backend import ( + OrchestratorBackend, +) + + +class _CannedRollout: + """Minimal stand-in for UnifiedRollout.""" + + def __init__(self, final_answer: str) -> None: + self.final_answer = final_answer + self.cost_usd = 0.123 + self.tokens = 42 + self.num_tool_calls = 1 + self.parse_failures = 0 + self.anon_map = {} + # generate_full now serializes a per-turn trace (reasoning / tool_name / + # arguments / observation), so turns must be real UnifiedTurn-shaped + # objects, not bare object() placeholders. + self.turns = [ + UnifiedTurn( + reasoning="let me search", + tool_name="web_search", + arguments={"query": "x"}, + observation="result", + ), + ] + + def tool_calls(self): + return [("web_search", {"query": "x"})] + + +def _patch_rollout(monkeypatch, rollout): + """Patch run_unified_rollout where the backend looks it up.""" + monkeypatch.setattr( + eb._rollout_mod, + "run_unified_rollout", + lambda *a, **k: rollout, + ) + + +def test_construct_is_network_free(): + """Constructing the backend builds the catalog but touches no network.""" + backend = OrchestratorBackend() + assert backend.backend_id == "orchestrator" + assert backend.orchestrator_endpoint == eb.DEFAULT_ENDPOINT + assert backend.orchestrator_model == eb.DEFAULT_MODEL + # Catalog is non-empty (cloud frontier + local OSS + basic tools). + assert len(backend._tools) > 0 + names = {t.name for t in backend._tools} + assert "web_search" in names + + +def test_local_endpoints_mapping(): + """local_endpoints (model-id -> base_url) wires the local OSS tool.""" + backend = OrchestratorBackend( + local_endpoints={"Qwen/Qwen3.5-9B": "http://m:9/v1"}, + ) + assert backend.local_endpoints["Qwen/Qwen3.5-9B"] == "http://m:9/v1" + # The corresponding local-OSS tool should carry that base_url. + tool = next(t for t in backend._tools if t.model == "Qwen/Qwen3.5-9B") + assert tool.base_url == "http://m:9/v1" + + +def test_generate_returns_final_answer(monkeypatch): + """generate() returns the rollout's final_answer string.""" + _patch_rollout(monkeypatch, _CannedRollout("42")) + backend = OrchestratorBackend() + out = backend.generate("what is 6*7?", model="qwen3-8b") + assert out == "42" + + +def test_generate_full_payload(monkeypatch): + """generate_full() returns the runner-expected payload shape.""" + _patch_rollout(monkeypatch, _CannedRollout("42")) + backend = OrchestratorBackend() + full = backend.generate_full("q", model="qwen3-8b") + assert full["content"] == "42" + assert full["cost_usd"] == 0.123 + assert full["usage"]["completion_tokens"] == 42 + assert full["tool_calls"] == 1 + assert full["turn_count"] == 1 + assert full["framework"] == "openjarvis-orchestrator" + assert "error" not in full + assert full["latency_seconds"] >= 0.0 + + +def test_generate_full_error_path(monkeypatch): + """An exception inside the rollout becomes a recorded error, not a crash.""" + + def _boom(*a, **k): + raise RuntimeError("vllm down") + + monkeypatch.setattr(eb._rollout_mod, "run_unified_rollout", _boom) + backend = OrchestratorBackend() + full = backend.generate_full("q", model="qwen3-8b") + assert full["content"] == "" + assert "vllm down" in full["error"] + assert full["usage"]["completion_tokens"] == 0 + + +# --------------------------------------------------------------------------- +# Eval-script benchmark name mapping +# --------------------------------------------------------------------------- + +# The 5 benchmarks the eval script targets, by the script's (alias) names. +SCRIPT_BENCHMARKS = [ + "gaia", + "terminalbench_v2_1", + "taubench", + "mmlu_pro", + "supergpqa", +] + + +def _load_script(): + import importlib.util + from pathlib import Path + + repo_root = Path(__file__).resolve().parents[2] + script_path = repo_root / "scripts" / "orchestrator" / "eval_orchestrator.py" + spec = importlib.util.spec_from_file_location("eval_orchestrator", script_path) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +def test_benchmark_aliases_map_to_registry_keys(): + """Each script benchmark normalizes to a real evals.cli BENCHMARKS key.""" + from openjarvis.evals.cli import BENCHMARKS + + script = _load_script() + for name in SCRIPT_BENCHMARKS: + key = script._normalize_benchmark(name) + assert key in BENCHMARKS, f"{name} -> {key} not in BENCHMARKS registry" + + +def test_normalize_specific_keys(): + """Spot-check the underscored aliases normalize to the dotted/dashed keys.""" + script = _load_script() + assert script._normalize_benchmark("terminalbench_v2_1") == "terminalbench-v2.1" + assert script._normalize_benchmark("mmlu_pro") == "mmlu-pro" + assert script._normalize_benchmark("gaia") == "gaia" + + +def test_build_dataset_accepts_keys(): + """_build_dataset must recognise the normalized keys. + + We only assert it doesn't raise the 'Unknown benchmark' ClickException; + datasets that need network/files at construction are xfail-skipped. + """ + import click + + from openjarvis.evals.cli import _build_dataset + + script = _load_script() + for name in SCRIPT_BENCHMARKS: + key = script._normalize_benchmark(name) + try: + _build_dataset(key) + except click.ClickException as exc: + if "Unknown benchmark" in str(exc): + pytest.fail(f"{key} not recognised by _build_dataset") + # Other ClickExceptions (missing files/creds) are acceptable here. + except Exception: + # Construction may require network/files — that's fine; the point + # is the key is recognised (didn't hit the unknown-benchmark else). + pytest.skip(f"{key} dataset construction needs resources") diff --git a/tests/test_orchestrator_learning/test_final_answer_sanitize.py b/tests/test_orchestrator_learning/test_final_answer_sanitize.py new file mode 100644 index 00000000..b1f4a844 --- /dev/null +++ b/tests/test_orchestrator_learning/test_final_answer_sanitize.py @@ -0,0 +1,120 @@ +"""Control-token sanitation in the SFT serializer + clean gate. + +Regression guard for the leak audit: leaked control/special tokens +(``<|im_end|>``, ``<|tool_call>``, gemma ````/````, +``<|"|>`` …) used to survive into the supervised final answer because the clean +gate only rejected the bare ```` form. Defense in depth now: + +1. ``_final_answer_block`` STRIPS residual control tokens so a good answer with a + stray token is salvaged (not dropped). +2. ``_target_is_clean`` REJECTS a rollout whose final answer is nothing but + control tokens, or where a token survives the strip. + +Both must leave legitimate ``<``/``>`` in math/code untouched. +""" + +from __future__ import annotations + +import pytest + +from openjarvis.agents.hybrid.toolorchestra.rollout import UnifiedRollout, UnifiedTurn +from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import ( + _target_is_clean, +) +from openjarvis.learning.intelligence.orchestrator.sft_data.unified_serialize import ( + _CONTROL_TOKEN_RE, + _final_answer_block, + _strip_control_tokens, +) + +# --- serializer strips control tokens, leaving a clean FINAL_ANSWER line ------ + + +@pytest.mark.parametrize( + "raw, expected_answer", + [ + ("FINAL_ANSWER: 42", "42"), + ("FINAL_ANSWER: 42<|end_of_turn|>", "42"), + ("FINAL_ANSWER: 42<|im_end|>", "42"), + ("FINAL_ANSWER: 42<|eot_id|>", "42"), + ("FINAL_ANSWER: The result<|tool_call>", "The result"), + ('FINAL_ANSWER: foo <|"|> bar', "foo bar"), + ("model", "model"), + ], +) +def test_final_answer_block_strips_tokens(raw: str, expected_answer: str) -> None: + out = _final_answer_block(raw) + assert out == f"FINAL_ANSWER: {expected_answer}" + # No control token survives the serializer. + assert not _CONTROL_TOKEN_RE.search(out), out + + +def test_bare_im_end_only_becomes_empty_answer() -> None: + # An answer that is nothing but a control token strips to empty; the clean + # gate (below) is what rejects such a rollout. + assert _final_answer_block("<|im_end|>") == "FINAL_ANSWER: " + + +def test_math_and_code_angle_brackets_survive() -> None: + # The narrow regex must NOT eat legitimate comparisons / generics. + for good in [ + "FINAL_ANSWER: x < 3 and y > 2", + "FINAL_ANSWER: List and ac", + "FINAL_ANSWER: if a < b: return a > 0", + ]: + out = _final_answer_block(good) + assert out == good, out + assert not _CONTROL_TOKEN_RE.search(out) + + +def test_strip_helper_handles_nested_and_whitespace() -> None: + assert _strip_control_tokens(" hi <|im_end|> ") == "hi" + assert _strip_control_tokens("a<|im_start|>b<|im_end|>c") == "abc" + + +# --- clean gate: salvage vs reject ------------------------------------------- + + +def _roll(final_answer: str) -> UnifiedRollout: + """A minimal well-formed rollout (one expert call, clean obs) whose only + variable is the final answer.""" + return UnifiedRollout( + turns=[ + UnifiedTurn( + reasoning="route it", + tool_name="expert_a", + arguments={"input": "q"}, + observation="valid observation", + ), + UnifiedTurn(reasoning=final_answer, tool_name=None), + ], + final_answer=final_answer, + cost_usd=0.01, + tokens=10, + num_tool_calls=1, + anon_map={"expert_a": "gpt"}, + ) + + +@pytest.mark.parametrize( + "final_answer, clean", + [ + # Salvageable: a good answer with a trailing stray token stays clean + # (the serializer strips the token from the emitted target). + ("42", True), + ("42<|end_of_turn|>", True), + ("The result<|tool_call>", True), + ('foo <|"|> bar', True), + ("model", True), + ("x < 3 and y > 2", True), + ("a normal plain answer", True), + # Unsalvageable: the answer is nothing but a control token -> rejected. + ("<|im_end|>", False), + ("<|eot_id|>", False), + ("", False), + # Bare tool-call tag: already rejected by the existing gate. + ("{}", False), + ], +) +def test_clean_gate_salvages_or_rejects(final_answer: str, clean: bool) -> None: + assert _target_is_clean(_roll(final_answer)) is clean diff --git a/tests/test_orchestrator_learning/test_sft_trainer.py b/tests/test_orchestrator_learning/test_sft_trainer.py index 941427f3..f05ab60d 100644 --- a/tests/test_orchestrator_learning/test_sft_trainer.py +++ b/tests/test_orchestrator_learning/test_sft_trainer.py @@ -15,7 +15,7 @@ from openjarvis.learning.intelligence.orchestrator.sft_trainer import ( class TestOrchestratorSFTConfig: def test_defaults(self): cfg = OrchestratorSFTConfig() - assert cfg.model_name == "Qwen/Qwen3-1.7B" + assert cfg.model_name == "Qwen/Qwen3.5-9B" assert cfg.num_epochs == 3 assert cfg.batch_size == 8 assert cfg.learning_rate == 2e-5