admin/OpenJarvis
1
0
Fork 0
mirror of https://github.com/open-jarvis/OpenJarvis.git synced 2026-07-30 10:52:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

toolorchestra: carry provider-aware search through the package split

Browse Source 
The monolith->package split landed while #558 was adding provider-aware web
search to toolorchestra.py upstream. Rebasing resolved that as delete-vs-modify,
which would have silently dropped the feature: the package always routed search
to Anthropic. Port it -- search now follows cloud_endpoint (openai-web-search /
gemini-web-search / anthropic-web-search), with the matching worker types,
per-search cost surcharges, and pool validation.

_TOOLORCH_SEARCH_TYPES replaces the scattered == "anthropic-web-search" checks,
so the one-shot guard and the "needs a solver" rule now cover every search type
rather than just the Anthropic one -- previously a search worker could have been
picked as the synthesis fallback.

Also, to get back under the ruff gate CI now enforces (it landed in the 95
commits we were behind): formatting, import order, and one real bug -- calculator
had "abs" twice in the same dict literal, and eval_backend had four imports
stranded below a function body.

The build_orchestrator_sft driver test read OJ_DATA_ROOT from the ambient
environment, so it wrote to the real experiments tree instead of tmp_path for
anyone who had sourced .env. Drop the var in the test.
This commit is contained in:
Andrew Park
2026-07-11 12:53:52 -07:00
parent 4864ea7776
commit 9571a158a2
38 changed files with 1874 additions and 870 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/orchestrator/build_orchestrator_sft.py
+143 -62
View File
@@ -71,35 +71,58 @@ def main(argv: Optional[list[str]] = None) -> int:
# data/orchestrator/raw/{label}_{MMDD}[_{tag}]/ and the file inside is always
# data.jsonl — see the run_dir block below. Use --out only to distinguish two
# runs of the same orchestrator on the same day (e.g. --out balanced50).
p.add_argument("--out", default=None,
help="Optional tag appended to the run dir name (not a path).")
p.add_argument(
"--out",
default=None,
help="Optional tag appended to the run dir name (not a path).",
)
# 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-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).")
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).
p.add_argument("--local-endpoint", action="append", default=[],
metavar="MODEL_ID=URL",
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(
"--local-endpoint",
action="append",
default=[],
metavar="MODEL_ID=URL",
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)
@@ -107,68 +130,107 @@ def main(argv: Optional[list[str]] = None) -> int:
# 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).")
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).")
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.")
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).")
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.")
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.")
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).")
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.")
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"):
for _noisy in (
"httpx",
"httpcore",
"urllib3",
"datasets",
"fsspec",
"huggingface_hub",
"openai",
):
logging.getLogger(_noisy).setLevel(logging.WARNING)
# Every run lands in its OWN folder under data/orchestrator/raw/, named to
# MATCH the sft/ file it will eventually produce — only the stage word differs:
# Every run lands in its OWN folder under <data-root>/raw/, named to MATCH the
# sft/ file it will eventually produce — only the stage word differs:
#
# data/orchestrator/raw/qwen_0707/data.jsonl <- every rollout, incl. failures
# data/orchestrator/sft/qwen_clean_0707.jsonl <- reject-sampled from it
# raw/qwen_0707/data.jsonl <- every rollout, incl. failures
# sft/qwen_clean_0707.jsonl <- reject-sampled from it
#
# so any curated file traces back to its generation run by eye. raw/ sits ABOVE
# the sft/rl fork on purpose: SFT keeps only correct+clean rows, but GRPO needs
@@ -177,9 +239,13 @@ def main(argv: Optional[list[str]] = None) -> int:
# suffix (creation order is preserved, and a MM-DD-HHMMpm stamp is redundant
# with mtime). --out is an optional extra tag, NOT a path; the data file inside
# is always data.jsonl.
#
# OJ_DATA_ROOT keeps the data OUT of the git checkout (repo-relative default so
# a fresh clone still works); this workspace points it at the experiments tree.
data_root = Path(os.getenv("OJ_DATA_ROOT", "data/orchestrator"))
prefix = args.orchestrator_label or "orch"
tag = f"_{Path(args.out).stem}" if args.out else ""
base = Path("data/orchestrator/raw") / f"{prefix}_{time.strftime('%m%d')}{tag}"
base = data_root / "raw" / f"{prefix}_{time.strftime('%m%d')}{tag}"
run_dir = base.resolve()
n = 1
while run_dir.exists(): # same name+day (incl. parallel shards) -> disambiguate
@@ -192,6 +258,7 @@ def main(argv: Optional[list[str]] = None) -> int:
lock_p.write_text(f"{time.strftime('%m-%d-%I%M%p').lower()} pid={os.getpid()}")
args.out = str(out_p)
import atexit
atexit.register(lambda: lock_p.exists() and lock_p.unlink())
logging.info("Run dir: %s", run_dir)
@@ -199,8 +266,7 @@ def main(argv: Optional[list[str]] = None) -> int:
# tracing.run_context() reads these). setdefault so an explicit env wins.
os.environ.setdefault("OJ_RUN_LABEL", label)
os.environ.setdefault("OJ_GEN_MODEL", args.gen_model or args.orchestrator_model)
os.environ.setdefault("OJ_RUN_STAGE",
"smoke" if args.max_tasks else "prod")
os.environ.setdefault("OJ_RUN_STAGE", "smoke" if args.max_tasks else "prod")
os.environ["OJ_CFG_TEMPERATURE"] = str(args.temperature)
os.environ["OJ_CFG_MAX_TURNS"] = str(args.max_turns)
os.environ["OJ_CFG_ANONYMIZE"] = str(bool(args.anonymize_experts))
@@ -247,8 +313,10 @@ def main(argv: Optional[list[str]] = None) -> int:
def rollout_fn(task):
try:
return run_unified_rollout(
task.instruction, tools,
call_orchestrator=call_orch, dispatch=dispatch,
task.instruction,
tools,
call_orchestrator=call_orch,
dispatch=dispatch,
max_turns=args.max_turns,
anonymize=args.anonymize_experts,
)
@@ -265,6 +333,7 @@ def main(argv: Optional[list[str]] = None) -> int:
# remaining unseen tasks stride evenly and never re-cover finished prompts.
if args.skip_task_ids_from:
import glob as _glob
skip_ids = set()
for pattern in args.skip_task_ids_from:
for fp in _glob.glob(pattern):
@@ -285,18 +354,31 @@ def main(argv: Optional[list[str]] = None) -> int:
if skip_ids:
n_before = len(tasks)
tasks = [t for t in tasks if t.task_id not in skip_ids]
logging.info("Resume: %d done task_ids -> skipping; %d of %d tasks remain",
len(skip_ids), len(tasks), n_before)
logging.info(
"Resume: %d done task_ids -> skipping; %d of %d tasks remain",
len(skip_ids),
len(tasks),
n_before,
)
if args.shard_count > 1:
n_all = len(tasks)
tasks = tasks[args.shard_index::args.shard_count]
logging.info("Shard %d/%d -> %d of %d tasks",
args.shard_index, args.shard_count, len(tasks), n_all)
tasks = tasks[args.shard_index :: args.shard_count]
logging.info(
"Shard %d/%d -> %d of %d tasks",
args.shard_index,
args.shard_count,
len(tasks),
n_all,
)
from collections import Counter as _Counter
_dom = _Counter(getattr(t, "domain", "unknown") for t in tasks)
logging.info("Loaded %d SFT tasks | balanced=%s | domain split: %s",
len(tasks), args.balanced,
", ".join(f"{d}={n}" for d, n in sorted(_dom.items())))
logging.info(
"Loaded %d SFT tasks | balanced=%s | domain split: %s",
len(tasks),
args.balanced,
", ".join(f"{d}={n}" for d, n in sorted(_dom.items())),
)
stats = generate_sft_dataset(
args.out,
@@ -321,8 +403,7 @@ def main(argv: Optional[list[str]] = None) -> int:
from render_sft_data import render
if stats.get("records_written", 0) > 0:
rinfo = render(args.out)
logging.info("Rendered viewable traces: %s + %s",
rinfo["pretty"], rinfo["txt"])
logging.info("Rendered viewable traces: %s + %s", rinfo["pretty"], rinfo["txt"])
print(json.dumps(stats, indent=2))
return 0
scripts/orchestrator/chat_orchestrator.py
+42 -18
View File
@@ -23,25 +23,43 @@ Usage:
# point local expert tiers at their own vLLMs if you have them up (optional):
# --local-endpoint Qwen/Qwen3.6-27B=http://localhost:8002/v1
"""
import argparse
import sys
import time
def _parse_args(argv=None):
p = argparse.ArgumentParser(description=__doc__,
formatter_class=argparse.RawDescriptionHelpFormatter)
p.add_argument("--endpoint", default="http://localhost:8020/v1",
help="vLLM endpoint serving the orchestrator checkpoint.")
p.add_argument("--model", default="sft-qwen-8k",
help="Served-model-name of the checkpoint.")
p.add_argument("--api-key", default="EMPTY", help="API key for the endpoint (local vLLM = EMPTY).")
p = argparse.ArgumentParser(
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
)
p.add_argument(
"--endpoint",
default="http://localhost:8020/v1",
help="vLLM endpoint serving the orchestrator checkpoint.",
)
p.add_argument(
"--model", default="sft-qwen-8k", help="Served-model-name of the checkpoint."
)
p.add_argument(
"--api-key",
default="EMPTY",
help="API key for the endpoint (local vLLM = EMPTY).",
)
p.add_argument("--temperature", type=float, default=0.0)
p.add_argument("--max-turns", type=int, default=8, help="Same default as the eval.")
p.add_argument("--max-tokens", type=int, default=3000,
help="Keep high: the model reasons a lot before it emits the delegation.")
p.add_argument("--local-endpoint", action="append", default=[],
help="MODEL_ID=URL for a local expert tier. Repeatable.")
p.add_argument(
"--max-tokens",
type=int,
default=3000,
help="Keep high: the model reasons a lot before it emits the delegation.",
)
p.add_argument(
"--local-endpoint",
action="append",
default=[],
help="MODEL_ID=URL for a local expert tier. Repeatable.",
)
return p.parse_args(argv)
@@ -68,8 +86,10 @@ def main(argv=None):
temperature=args.temperature,
)
print(f"orchestrator: {args.model} @ {args.endpoint} "
f"(temp={args.temperature}, max_turns={args.max_turns})")
print(
f"orchestrator: {args.model} @ {args.endpoint} "
f"(temp={args.temperature}, max_turns={args.max_turns})"
)
print("type a question and hit enter. Ctrl-D or 'quit' to exit.\n")
try:
@@ -86,7 +106,9 @@ def main(argv=None):
started = time.time()
full = backend.generate_full(
prompt, model=args.model, temperature=args.temperature,
prompt,
model=args.model,
temperature=args.temperature,
max_tokens=args.max_tokens,
)
elapsed = time.time() - started
@@ -96,10 +118,12 @@ def main(argv=None):
continue
print(f"\norch> {full.get('content', '').strip()}\n")
print(f" [turns={full.get('turn_count', '?')} "
f"tool_calls={full.get('tool_calls', '?')} "
f"cost=${full.get('cost_usd', 0.0):.4f} "
f"{elapsed:.1f}s]\n")
print(
f" [turns={full.get('turn_count', '?')} "
f"tool_calls={full.get('tool_calls', '?')} "
f"cost=${full.get('cost_usd', 0.0):.4f} "
f"{elapsed:.1f}s]\n"
)
finally:
backend.close()
return 0
scripts/orchestrator/eval_orchestrator.py
+2 -3
View File
@@ -258,9 +258,7 @@ def main(argv: Optional[List[str]] = None) -> int:
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}"
)
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()
@@ -307,6 +305,7 @@ def main(argv: Optional[List[str]] = None) -> int:
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}"},
scripts/orchestrator/format_eval_sample.py
+48 -20
View File
@@ -30,6 +30,7 @@ Usage:
... --input <file> --all # one .txt per record
... --input <file> --all --only-wrong # only incorrect samples
"""
from __future__ import annotations
import argparse
@@ -81,8 +82,10 @@ def build_question_map(benchmark: str, n: int, seed: int = 42) -> dict:
ds.load(max_samples=n, seed=seed)
return {r.record_id: r.problem for r in ds.iter_records()}
except Exception as exc: # noqa: BLE001 - never fail the render over this
print(f" [warn] could not load dataset for {benchmark!r}: "
f"{type(exc).__name__}: {exc}")
print(
f" [warn] could not load dataset for {benchmark!r}: "
f"{type(exc).__name__}: {exc}"
)
return {}
@@ -134,7 +137,11 @@ def parse_scoring(meta: dict) -> dict:
def format_record(rec: dict, question: str | None) -> str:
parsed = parse_scoring(rec.get("scoring_metadata"))
is_correct = rec.get("is_correct")
verdict = "CORRECT" if is_correct else ("INCORRECT" if is_correct is False else "UNSCORED")
verdict = (
"CORRECT"
if is_correct
else ("INCORRECT" if is_correct is False else "UNSCORED")
)
def _fmt(v, fmt):
try:
@@ -142,14 +149,16 @@ def format_record(rec: dict, question: str | None) -> str:
except (ValueError, TypeError):
return str(v)
top = " · ".join([
str(rec.get("record_id", "?")),
str(rec.get("benchmark", "?")),
verdict,
f"score {_fmt(rec.get('score'), '{:.3f}')}",
f"lat {_fmt(rec.get('latency_seconds'), '{:.1f}')}s",
f"cost ${_fmt(rec.get('cost_usd'), '{:.4f}')}",
])
top = " · ".join(
[
str(rec.get("record_id", "?")),
str(rec.get("benchmark", "?")),
verdict,
f"score {_fmt(rec.get('score'), '{:.3f}')}",
f"lat {_fmt(rec.get('latency_seconds'), '{:.1f}')}s",
f"cost ${_fmt(rec.get('cost_usd'), '{:.4f}')}",
]
)
parts = [top]
err = rec.get("error")
@@ -164,7 +173,13 @@ def format_record(rec: dict, question: str | None) -> str:
parts.append(_banner("GOLD"))
gold = parsed["gold"]
parts.append(_indent(str(gold) if gold not in (None, "") else "(gold unavailable — see SCORING below)"))
parts.append(
_indent(
str(gold)
if gold not in (None, "")
else "(gold unavailable — see SCORING below)"
)
)
if parsed["kind"] == "judge" and parsed["judge"]:
parts.append(_banner("JUDGE"))
@@ -172,7 +187,9 @@ def format_record(rec: dict, question: str | None) -> str:
if not parsed["clean"]:
parts.append(_banner("SCORING (raw)"))
parts.append(_indent(json.dumps(rec.get("scoring_metadata"), indent=2, default=str)))
parts.append(
_indent(json.dumps(rec.get("scoring_metadata"), indent=2, default=str))
)
return "\n".join(parts).rstrip() + "\n"
@@ -182,10 +199,17 @@ def main(argv=None):
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
)
p.add_argument("--input", required=True, help="Graded *_orchestrator.jsonl file.")
p.add_argument("--out-dir", default="results/formatted", help="Where the .txt files go.")
p.add_argument("--seed", type=int, default=42, help="Subset seed used at eval time.")
p.add_argument("--only-wrong", action="store_true",
help="Render only incorrect/unscored samples.")
p.add_argument(
"--out-dir", default="results/formatted", help="Where the .txt files go."
)
p.add_argument(
"--seed", type=int, default=42, help="Subset seed used at eval time."
)
p.add_argument(
"--only-wrong",
action="store_true",
help="Render only incorrect/unscored samples.",
)
g = p.add_mutually_exclusive_group()
g.add_argument("--n", type=int, default=1, help="Format the first N records.")
g.add_argument("--lines", help="Comma-separated 1-indexed line numbers.")
@@ -213,7 +237,9 @@ def main(argv=None):
if args.only_wrong:
idxs = [i for i in idxs if not records[i].get("is_correct")]
benchmark = _normalize_benchmark(records[0].get("benchmark", src.stem.split("_")[0]))
benchmark = _normalize_benchmark(
records[0].get("benchmark", src.stem.split("_")[0])
)
# Load the whole subset so any selected id resolves (the eval-time subset is
# the first len(records) of seed=42).
qmap = build_question_map(benchmark, n=len(records), seed=args.seed)
@@ -230,8 +256,10 @@ def main(argv=None):
for w in written:
print(w)
print(f"\nwrote {len(written)} file(s) to {out_dir}/"
+ (" (dataset unavailable — questions omitted)" if not qmap else ""))
print(
f"\nwrote {len(written)} file(s) to {out_dir}/"
+ (" (dataset unavailable — questions omitted)" if not qmap else "")
)
return 0
scripts/orchestrator/make_splits.py
+42 -15
View File
@@ -27,6 +27,7 @@ Deterministic (seed 42).
--pool data/orchestrator/sft/qwen_clean_0711.jsonl \
--pool data/orchestrator/sft/gemma_clean_0711.jsonl
"""
import argparse
import json
import os
@@ -36,7 +37,12 @@ from collections import defaultdict
from datetime import datetime
from pathlib import Path
OUT = Path("data/orchestrator/sft")
# 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
@@ -46,16 +52,32 @@ 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("--date", default=datetime.now().strftime("%m%d"),
help="date tag in the filename (default: today, MMDD)")
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(
"--date",
default=datetime.now().strftime("%m%d"),
help="date tag in the filename (default: today, MMDD)",
)
ap.add_argument("--out", type=Path, default=OUT)
args = ap.parse_args()
@@ -94,8 +116,11 @@ def main() -> int:
args.out.mkdir(parents=True, exist_ok=True)
written = {}
for split, data in (("train", train), ("holdout", holdout),
(f"overfit{OVERFIT_N}", overfit)):
for split, data in (
("train", train),
("holdout", holdout),
(f"overfit{OVERFIT_N}", overfit),
):
f = args.out / f"{args.name}_{split}_{args.date}.jsonl"
f.write_text("".join(json.dumps(r) + "\n" for r in data))
written[split] = f
@@ -109,8 +134,10 @@ def main() -> int:
from upload_to_braintrust import autoupload
autoupload(
[f"{written['train']}={args.name}_train_{args.date}",
f"{written['holdout']}={args.name}_holdout_{args.date}"],
[
f"{written['train']}={args.name}_train_{args.date}",
f"{written['holdout']}={args.name}_holdout_{args.date}",
],
run_label=os.getenv("OJ_RUN_LABEL", f"{args.name}_splits_{args.date}"),
description=f"{args.name} orchestrator-SFT splits (make_splits.py)",
)
scripts/orchestrator/render_sft_data.py
+55 -23
View File
@@ -64,7 +64,11 @@ def _fmt_tool_calls(content: str, anon_map=None) -> str:
lines = [f"-> {_call_label(name, anon_map)}"]
if isinstance(args, dict):
for k, v in args.items():
val = v if isinstance(v, str) else json.dumps(v, ensure_ascii=False, indent=2)
val = (
v
if isinstance(v, str)
else json.dumps(v, ensure_ascii=False, indent=2)
)
indented = "\n".join(" " + ln for ln in val.splitlines())
lines.append(f" {k}:\n{indented}")
blocks.append("\n".join(lines))
@@ -136,7 +140,8 @@ def _system_header(records: List[dict]) -> str:
if str(t.get("role", "")).lower() == "system":
return (
"=" * 80 + "\nSHARED SYSTEM PROMPT + TOOL CATALOG "
"(identical for every record below)\n" + "=" * 80
"(identical for every record below)\n"
+ "=" * 80
+ f"\n{t.get('content', '').strip()}\n"
)
return ""
@@ -186,27 +191,42 @@ def _html_turns(record: dict) -> str:
if role == "system":
continue
if role == "user":
out.append(f'<div class="turn user"><div class="role">user</div>{_esc(content.strip())}</div>')
out.append(
f'<div class="turn user"><div class="role">user</div>{_esc(content.strip())}</div>'
)
elif role in ("tool", "function"):
name = _esc(t.get("name", "tool"))
out.append(
f'<div class="tool"><details><summary>tool result · {name} '
f'({len(content)} chars)</summary><pre>{_esc(content.strip())}</pre></details></div>'
f"({len(content)} chars)</summary><pre>{_esc(content.strip())}</pre></details></div>"
)
elif role == "assistant":
prose = _TOOL_CALL_RE.sub("", content).strip()
if prose:
out.append(f'<div class="turn assistant"><div class="role">assistant</div>{_esc(prose)}</div>')
out.append(
f'<div class="turn assistant"><div class="role">assistant</div>{_esc(prose)}</div>'
)
for name, args in _parse_tool_calls(content):
real = anon_map.get(name)
label = (f'{_esc(str(name))} <span class="real">&rarr; {_esc(str(real))}</span>'
if real else _esc(str(name)))
label = (
f'{_esc(str(name))} <span class="real">&rarr; {_esc(str(real))}</span>'
if real
else _esc(str(name))
)
rows = [f'<div class="cname">&rarr; {label}</div>']
if isinstance(args, dict):
for k, v in args.items():
val = v if isinstance(v, str) else json.dumps(v, ensure_ascii=False, indent=2)
rows.append(f'<div class="arg"><b>{_esc(str(k))}</b><pre>{_esc(val)}</pre></div>')
out.append(f'<div class="turn call"><div class="role">tool call</div>{"".join(rows)}</div>')
val = (
v
if isinstance(v, str)
else json.dumps(v, ensure_ascii=False, indent=2)
)
rows.append(
f'<div class="arg"><b>{_esc(str(k))}</b><pre>{_esc(val)}</pre></div>'
)
out.append(
f'<div class="turn call"><div class="role">tool call</div>{"".join(rows)}</div>'
)
return "".join(out)
@@ -215,16 +235,26 @@ def _html_record(record: dict, index: int) -> str:
pills = [f'<span class="pill dom">{_esc(str(record.get("domain")))}</span>']
correct = record.get("correct")
if correct is not None:
pills.append('<span class="pill ok">correct</span>' if correct
else '<span class="pill wrong">wrong</span>')
pills.append(
'<span class="pill ok">correct</span>'
if correct
else '<span class="pill wrong">wrong</span>'
)
if record.get("kept"):
pills.append('<span class="pill kept">kept</span>')
meta = (f'<span class="meta">cost ${m.get("cost_usd")} · {m.get("tokens")} tok · '
f'{m.get("num_tool_calls")} calls · {m.get("num_turns")} turns</span>')
summary = (f'<span>#{index}</span><span>{_esc(str(record.get("task_id")))}</span>'
+ "".join(pills) + meta)
return (f'<details class="rec"><summary>{summary}</summary>'
f'<div class="body">{_html_turns(record)}</div></details>')
meta = (
f'<span class="meta">cost ${m.get("cost_usd")} · {m.get("tokens")} tok · '
f"{m.get('num_tool_calls')} calls · {m.get('num_turns')} turns</span>"
)
summary = (
f"<span>#{index}</span><span>{_esc(str(record.get('task_id')))}</span>"
+ "".join(pills)
+ meta
)
return (
f'<details class="rec"><summary>{summary}</summary>'
f'<div class="body">{_html_turns(record)}</div></details>'
)
def _html_doc(records: List[dict], title: str) -> str:
@@ -262,9 +292,7 @@ def render(jsonl_path: str) -> dict:
pretty.write_text(json.dumps(records, indent=2, ensure_ascii=False))
txt = src.with_suffix(".txt")
banner = (
f"ORCHESTRATOR SFT TRANSCRIPT — {src.name}\n{len(records)} records\n\n"
)
banner = f"ORCHESTRATOR SFT TRANSCRIPT — {src.name}\n{len(records)} records\n\n"
sep = f"\n\n{_RULE}\n\n"
parts = [banner + _system_header(records)]
parts.append(sep.join(_transcript(r, i + 1) for i, r in enumerate(records)))
@@ -273,8 +301,12 @@ def render(jsonl_path: str) -> dict:
htmlp = src.with_suffix(".html")
htmlp.write_text(_html_doc(records, src.name))
return {"records": len(records), "pretty": str(pretty),
"txt": str(txt), "html": str(htmlp)}
return {
"records": len(records),
"pretty": str(pretty),
"txt": str(txt),
"html": str(htmlp),
}
def main(argv: List[str]) -> int:
scripts/orchestrator/run_sft_fsdp.py
+109 -40
View File
@@ -24,6 +24,7 @@ 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
@@ -39,8 +40,14 @@ from pathlib import Path
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"}
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:
@@ -53,7 +60,9 @@ def record_is_correct(r: dict) -> bool:
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", [])))
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:
@@ -84,7 +93,9 @@ def select(records, variant: str, require_clean: bool = True):
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)]
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:
@@ -94,14 +105,25 @@ def select(records, variant: str, require_clean: bool = True):
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(
"--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)
@@ -109,9 +131,12 @@ def main() -> int:
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(
"--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")
@@ -120,10 +145,11 @@ def main() -> int:
import random
from datetime import timedelta
import torch
from torch.utils.data import DataLoader
from accelerate import Accelerator
from accelerate.utils import set_seed, InitProcessGroupKwargs
from accelerate.utils import InitProcessGroupKwargs, set_seed
from torch.utils.data import DataLoader
from transformers import AutoModelForCausalLM, AutoTokenizer
set_seed(args.seed)
@@ -133,7 +159,10 @@ def main() -> int:
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))])
kwargs_handlers=[
InitProcessGroupKwargs(timeout=timedelta(minutes=pg_timeout_min))
],
)
is_main = accelerator.is_main_process
def log(m):
@@ -153,14 +182,20 @@ def main() -> int:
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")
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)
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")
@@ -179,7 +214,9 @@ def main() -> int:
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)
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
@@ -187,21 +224,32 @@ def main() -> int:
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)
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)
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")
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.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)
@@ -234,17 +282,25 @@ def main() -> int:
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)))
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"
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})
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})")
@@ -262,9 +318,11 @@ def main() -> int:
p.mkdir(parents=True, exist_ok=True)
tok.save_pretrained(str(p))
unwrapped.save_pretrained(
str(p), is_main_process=is_main,
str(p),
is_main_process=is_main,
save_function=accelerator.save,
state_dict=accelerator.get_state_dict(model))
state_dict=accelerator.get_state_dict(model),
)
if is_main:
log(f"checkpoint saved -> {p}")
@@ -319,9 +377,11 @@ def main() -> int:
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,
str(out_dir),
is_main_process=is_main,
save_function=accelerator.save,
state_dict=accelerator.get_state_dict(model))
state_dict=accelerator.get_state_dict(model),
)
if is_main:
log(f"saved -> {out_dir}")
if use_wandb:
@@ -329,8 +389,17 @@ def main() -> int:
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)}))
print(
"FSDP_SFT_DONE "
+ json.dumps(
{
"variant": args.variant,
"steps": gstep,
"wall_s": round(time.time() - t0, 1),
"out": str(out_dir),
}
)
)
return 0
scripts/orchestrator/sft_tokenize.py
+8 -5
View File
@@ -6,6 +6,7 @@ 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
@@ -102,8 +103,9 @@ def build_example(
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_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
@@ -114,7 +116,7 @@ def build_example(
if spans:
selected = spans if supervise_all_turns else spans[-1:]
labels = [-100] * len(full)
for (s, e) in selected:
for s, e in selected:
for k in range(s, e):
labels[k] = full[k]
last_start = selected[-1][0]
@@ -122,8 +124,9 @@ def build_example(
# 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_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
scripts/orchestrator/upload_to_braintrust.py
+118 -40
View File
@@ -34,6 +34,7 @@ Programmatic (used by the pipeline auto-upload hook in make_splits.py):
autoupload(["data/orchestrator/sft/qwen_train_0707.jsonl=qwen_train_0707"],
run_label="...")
"""
import argparse
import json
import logging
@@ -82,6 +83,7 @@ def _dataset_desc(dataset):
from openjarvis.learning.intelligence.orchestrator.sft_data.reject_sample import (
dataset_description,
)
return dataset_description(dataset)
except Exception:
return ""
@@ -89,10 +91,16 @@ def _dataset_desc(dataset):
def _git_sha():
try:
return subprocess.check_output(
["git", "rev-parse", "--short", "HEAD"],
cwd=Path(__file__).resolve().parent, stderr=subprocess.DEVNULL,
).decode().strip() or None
return (
subprocess.check_output(
["git", "rev-parse", "--short", "HEAD"],
cwd=Path(__file__).resolve().parent,
stderr=subprocess.DEVNULL,
)
.decode()
.strip()
or None
)
except Exception:
return None
@@ -113,7 +121,17 @@ def _model_short(rows, path):
def _split_of(path):
stem = Path(path).stem.lower()
for s in ("holdout", "train", "clean", "overfit", "partial", "8k", "4k", "2k", "1k"):
for s in (
"holdout",
"train",
"clean",
"overfit",
"partial",
"8k",
"4k",
"2k",
"1k",
):
if s in stem:
return s
return "data"
@@ -129,8 +147,16 @@ def _config_from_env():
for env_key, meta_key, cast in [
("OJ_CFG_TEMPERATURE", "temperature", float),
("OJ_CFG_MAX_TURNS", "max_turns", int),
("OJ_CFG_ANONYMIZE", "anonymize", lambda v: v.strip().lower() in ("1", "true", "yes")),
("OJ_CFG_REJECTION_ONLY", "rejection_only", lambda v: v.strip().lower() in ("1", "true", "yes")),
(
"OJ_CFG_ANONYMIZE",
"anonymize",
lambda v: v.strip().lower() in ("1", "true", "yes"),
),
(
"OJ_CFG_REJECTION_ONLY",
"rejection_only",
lambda v: v.strip().lower() in ("1", "true", "yes"),
),
]:
v = os.getenv(env_key)
if v not in (None, ""):
@@ -148,7 +174,7 @@ def _row(rec):
asst = [c["content"] for c in convs if c["role"] == "assistant"]
fa = asst[-1] if asst else ""
m = list(_FA.finditer(fa))
model_answer = fa[m[-1].end():].strip() if m else fa.strip()
model_answer = fa[m[-1].end() :].strip() if m else fa.strip()
# Gold reference the verifier graded against (stamped at generation time as
# `gold_answer`). `expected` in Braintrust is the GOLD so the UI shows
# gold-vs-model; the model's own answer goes to metadata.model_answer.
@@ -158,7 +184,9 @@ def _row(rec):
routed = []
for c in convs:
if c["role"] == "assistant":
for t in re.findall(r"<tool_call>(\{.*?\})</tool_call>", c["content"], re.DOTALL):
for t in re.findall(
r"<tool_call>(\{.*?\})</tool_call>", c["content"], re.DOTALL
):
try:
real = am.get(json.loads(t).get("name"))
if real:
@@ -171,14 +199,18 @@ def _row(rec):
correct = rec.get("correct")
clean = rec.get("clean")
kept = rec.get("kept")
tags = [t for t in [
gen_model,
f"domain:{domain}" if domain else None,
"correct" if correct else "incorrect",
"clean" if clean else "dirty",
"kept" if kept else "dropped",
"gold" if has_gold else "no_gold",
] if t]
tags = [
t
for t in [
gen_model,
f"domain:{domain}" if domain else None,
"correct" if correct else "incorrect",
"clean" if clean else "dirty",
"kept" if kept else "dropped",
"gold" if has_gold else "no_gold",
]
if t
]
return {
"input": question,
"expected": gold,
@@ -191,7 +223,8 @@ def _row(rec):
"area": rec.get("area"),
"difficulty": rec.get("difficulty") or None,
"dataset": rec.get("dataset"),
"dataset_description": rec.get("dataset_description") or _dataset_desc(rec.get("dataset")),
"dataset_description": rec.get("dataset_description")
or _dataset_desc(rec.get("dataset")),
"subsector": rec.get("subsector"),
"task_id": rec.get("task_id"),
"correct": correct,
@@ -222,7 +255,9 @@ def _routed_dist(rows):
for cc in r.get("conversations", []):
if cc.get("role") != "assistant":
continue
for t in re.findall(r"<tool_call>(\{.*?\})</tool_call>", cc.get("content", ""), re.DOTALL):
for t in re.findall(
r"<tool_call>(\{.*?\})</tool_call>", cc.get("content", ""), re.DOTALL
):
try:
real = am.get(json.loads(t).get("name"))
if real:
@@ -257,7 +292,7 @@ def _dataset_metadata(rows, path, run_label, gen_model):
"source_datasets": ["GeneralThought-430K-filtered", "OpenThoughts3-1.2M"],
"task_mix": "balanced (GeneralThought + OpenThoughts code/math/science)",
"filter": "correct + clean (rejects file-write echoes, garble, "
"reasoning-degeneration, truncated tails, essays/markdown dumps)",
"reasoning-degeneration, truncated tails, essays/markdown dumps)",
"leak_free": "train/holdout disjoint by task_id",
"n_total": n,
"n_correct": sum(1 for r in rows if r.get("correct")),
@@ -277,8 +312,16 @@ def _dataset_metadata(rows, path, run_label, gen_model):
return {k: v for k, v in meta.items() if v is not None}
def upload_dataset(path, name=None, *, project_id=None, project=None,
run_label=None, gen_model=None, description=""):
def upload_dataset(
path,
name=None,
*,
project_id=None,
project=None,
run_label=None,
gen_model=None,
description="",
):
"""Upload one JSONL file as a Braintrust dataset. Returns (name, n_rows, url).
Target project: ``project`` (name) if given, else ``project_id`` /
@@ -288,24 +331,33 @@ def upload_dataset(path, name=None, *, project_id=None, project=None,
rows = [json.loads(l) for l in open(path) if l.strip()]
name = name or _default_name(rows, path)
gen_model = gen_model or os.getenv("OJ_GEN_MODEL") or \
next((r.get("gen_model") for r in rows if r.get("gen_model")), None)
gen_model = (
gen_model
or os.getenv("OJ_GEN_MODEL")
or next((r.get("gen_model") for r in rows if r.get("gen_model")), None)
)
run_label = run_label or os.getenv("OJ_RUN_LABEL") or name
init_kwargs = {"name": name, "description": description or None,
"metadata": _dataset_metadata(rows, path, run_label, gen_model)}
init_kwargs = {
"name": name,
"description": description or None,
"metadata": _dataset_metadata(rows, path, run_label, gen_model),
}
if project:
init_kwargs["project"] = project
else:
init_kwargs["project_id"] = project_id or os.getenv(
"OJ_BRAINTRUST_PROJECT_ID", DEFAULT_PROJECT_ID)
"OJ_BRAINTRUST_PROJECT_ID", DEFAULT_PROJECT_ID
)
ds = braintrust.init_dataset(**init_kwargs)
for r in rows:
ds.insert(**_row(r))
ds.flush()
summ = ds.summarize()
url = getattr(summ, "dataset_url", None) or (project or init_kwargs.get("project_id"))
url = getattr(summ, "dataset_url", None) or (
project or init_kwargs.get("project_id")
)
return name, len(rows), url
@@ -313,7 +365,13 @@ def autoupload(specs, *, run_label=None, gen_model=None, description=""):
"""No-op-safe wrapper for pipeline hooks. Honors OJ_BRAINTRUST_AUTOUPLOAD
(default ON) and never raises: any failure (missing key/pkg, network) is
logged and swallowed so the data pipeline is never broken by telemetry."""
if os.getenv("OJ_BRAINTRUST_AUTOUPLOAD", "1").strip().lower() in ("0", "false", "no", "off", ""):
if os.getenv("OJ_BRAINTRUST_AUTOUPLOAD", "1").strip().lower() in (
"0",
"false",
"no",
"off",
"",
):
logger.info("[braintrust] autoupload disabled (OJ_BRAINTRUST_AUTOUPLOAD)")
return
if not os.getenv("BRAINTRUST_API_KEY"):
@@ -327,22 +385,37 @@ def autoupload(specs, *, run_label=None, gen_model=None, description=""):
for spec in specs:
path, _, name = spec.partition("=")
try:
nm, n, url = upload_dataset(path, name or None, run_label=run_label,
gen_model=gen_model, description=description)
nm, n, url = upload_dataset(
path,
name or None,
run_label=run_label,
gen_model=gen_model,
description=description,
)
logger.info("[braintrust] uploaded %s: %d rows -> %s", nm, n, url)
print(f"[braintrust] uploaded {nm}: {n} rows -> {url}")
except Exception as exc:
logger.warning("[braintrust] autoupload FAILED for %s (%s) — continuing", path, exc)
print(f"[braintrust] autoupload FAILED for {path} ({exc}) — pipeline unaffected")
logger.warning(
"[braintrust] autoupload FAILED for %s (%s) — continuing", path, exc
)
print(
f"[braintrust] autoupload FAILED for {path} ({exc}) — pipeline unaffected"
)
def main():
logging.basicConfig(level=logging.INFO, format="%(message)s")
p = argparse.ArgumentParser()
p.add_argument("--project", default=None,
help="Target project by NAME (overrides OJ_BRAINTRUST_PROJECT_ID).")
p.add_argument("--project-id", default=None,
help="Target project by id (default: OJ_BRAINTRUST_PROJECT_ID or research).")
p.add_argument(
"--project",
default=None,
help="Target project by NAME (overrides OJ_BRAINTRUST_PROJECT_ID).",
)
p.add_argument(
"--project-id",
default=None,
help="Target project by id (default: OJ_BRAINTRUST_PROJECT_ID or research).",
)
p.add_argument("--run-label", default=None)
p.add_argument("--gen-model", default=None, help="Specific gen model id override.")
p.add_argument("--description", default="")
@@ -351,9 +424,14 @@ def main():
for spec in args.specs:
path, _, name = spec.partition("=")
nm, n, url = upload_dataset(
path, name or None, project_id=args.project_id, project=args.project,
run_label=args.run_label, gen_model=args.gen_model,
description=args.description)
path,
name or None,
project_id=args.project_id,
project=args.project,
run_label=args.run_label,
gen_model=args.gen_model,
description=args.description,
)
print(f"[braintrust] {nm}: {n} rows -> {url}")