diff --git a/.gitignore b/.gitignore index 80d3cefd..6e978eda 100644 --- a/.gitignore +++ b/.gitignore @@ -133,3 +133,10 @@ oj-debug.*.json # Generated orchestrator SFT/eval data artifacts data/ + +# Training run artifacts (model checkpoints + W&B run dirs — regenerated, large) +checkpoints/ +wandb/ + +# Local LiteLLM spend-tracking proxy (machine-specific, not part of the package) +litellm/ diff --git a/pyproject.toml b/pyproject.toml index 6e253ed7..03159a49 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ + "braintrust>=0.0.150", "click>=8", "datasets>=4.5.0", "ddgs>=9.11.4", diff --git a/scripts/orchestrator/build_orchestrator_sft.py b/scripts/orchestrator/build_orchestrator_sft.py index b83c08a2..b6f8cd0a 100644 --- a/scripts/orchestrator/build_orchestrator_sft.py +++ b/scripts/orchestrator/build_orchestrator_sft.py @@ -42,6 +42,10 @@ from __future__ import annotations import argparse import json import logging +import os +import sys +import time +from pathlib import Path from typing import Optional from openjarvis.agents.hybrid.expert_registry import orchestrator_catalog @@ -63,12 +67,25 @@ def main(argv: Optional[list[str]] = None) -> int: p = argparse.ArgumentParser( description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter ) - p.add_argument("--out", default="data/orchestrator_sft_v1.jsonl") + # Default output is timestamped (e.g. data/orchestrator_sft_06-22-1230am.jsonl) + # so every run lands in its own file and concurrent runs never collide. Pass + # --out explicitly to override (sharded runs do this, one file per shard). + p.add_argument("--out", default=None, + help="Output JSONL. Default: data/orchestrator_sft_.jsonl") # The orchestrator == the local Qwen3-8B self-sampling over its own rollouts. p.add_argument("--orchestrator-endpoint", default="http://localhost:8001/v1", help="OpenAI-compatible vLLM base URL serving the orchestrator.") p.add_argument("--orchestrator-model", default="qwen3-8b") p.add_argument("--orchestrator-api-key", default="EMPTY") + # Provenance stamps: when set, every written record carries these fields so + # the JSONL is self-identifying once pooled across model families (gemma vs + # qwen). Default None -> no stamping (existing qwen runs unaffected). + p.add_argument("--gen-model", default=None, + help="Full HF id of the generating orchestrator, stamped as " + "record['gen_model'] (e.g. google/gemma-4-26B-A4B-it).") + p.add_argument("--orchestrator-label", default=None, + help="Short label stamped as record['orchestrator_model'] " + "(e.g. gemma-4-26b).") # Local OSS model endpoints; repeatable "model_id=base_url" (e.g. # "Qwen/Qwen3.5-9B=http://localhost:8001/v1"). Unmapped local models are # still listed but served unconfigured (base_url=None). @@ -77,13 +94,112 @@ def main(argv: Optional[list[str]] = None) -> int: help="Local model id -> vLLM base URL (repeatable).") p.add_argument("--max-tasks", type=int, default=None, help="Cap on tasks (default: all of load_sft_tasks()).") + p.add_argument("--skip-task-ids-from", action="append", default=[], + metavar="GLOB", + help="Glob of prior data.jsonl files; skip task_ids already " + "generated there so this run only does unseen prompts " + "(resume). Repeatable; applied before sharding.") p.add_argument("--samples-per-task", type=int, default=8) p.add_argument("--max-keep-per-task", type=int, default=1) p.add_argument("--max-turns", type=int, default=8) p.add_argument("--temperature", type=float, default=0.7) + # Orchestrator completion cap per turn. The library default (4096) intermittently + # truncates the final answer mid-sentence on longer reasoning turns; bump it so + # the trace ends cleanly. Raise further (e.g. 16384) if traces still cut off. + p.add_argument("--max-tokens", type=int, default=8192, + help="Per-turn completion cap for the orchestrator (default 8192).") + # Number of tasks rolled out concurrently against vLLM. Each in-flight task + # issues its samples sequentially, so ~concurrency requests hit the server at + # once. ~30-50 is comfortable on 1xL40S (prefix caching + continuous batching). + p.add_argument("--concurrency", type=int, default=32, + help="Tasks rolled out in parallel (default 32).") + # Balanced-smoke pull: cap//4 from each of GeneralThought + OpenThoughts + # code/math/science instead of the GeneralThought-only fast cap path. Use for a + # representative smoke; the real run omits --max-tasks for the full 8K balanced set. + p.add_argument("--balanced", action=argparse.BooleanOptionalAction, default=True, + help="Draw an EVEN cross-domain sample (GeneralThought + " + "OpenThoughts code/math/science). Default ON — pass " + "--no-balanced for the old GeneralThought-only skew.") + # Data-parallel sharding: split the (deterministic, seed-42) task list across + # N independent driver processes, each pointed at its own orchestrator vLLM + # replica. Shard i takes tasks[i::N] (a strided slice, so every shard stays + # domain-balanced). Each writes its own --out file; concatenate the shard + # JSONLs afterward. Lets the GPU-bound orchestrator scale across idle GPUs. + p.add_argument("--shard-index", type=int, default=0, + help="This shard's index in [0, shard-count).") + p.add_argument("--shard-count", type=int, default=1, + help="Total number of shards (default 1 = no sharding).") + # Throughput knob: stop sampling a task as soon as max-keep-per-task passing + # trajectories are found, instead of always running all --samples-per-task. + # ~3-4x faster when most tasks solve early, but keeps the *first* passers + # rather than the *cheapest* of N (drops the cost-optimisation signal). + p.add_argument("--stop-at-keep", action="store_true", + help="Short-circuit a task once max-keep passing samples found.") + # Anonymize model experts (opaque random labels, uniform description, no cost + # line, shuffled order) so the policy can't route on a model's name/position/ + # cost. The anon->real map is saved per record (metrics.anon_map) for analysis. + p.add_argument("--anonymize-experts", action="store_true", + help="Hide expert identity (random labels + shuffle) when routing.") + # By default we keep EVERY rolled-out trajectory (correct + incorrect), each + # tagged with ``correct`` (verifier verdict) and ``kept`` (the cheapest-correct + # sample the rejection sampler would pick). Lets you compute accuracy / inspect + # failures from the JSONL directly. --rejection-only restores the old + # drop-the-failures behaviour (only cheapest-correct written). + p.add_argument("--rejection-only", action="store_true", + help="Drop incorrect rollouts; write only the cheapest-correct " + "sample per task (the original behaviour).") + # Rollouts call shell_exec / file_write with model-chosen relative paths + # (e.g. ``solution.py``), which otherwise land in the repo root. Run them + # from a throwaway scratch dir so generated files never dirty the tree. + # gitignored via the existing ``scratch/`` rule. + p.add_argument("--scratch-dir", default="scratch/sft-rollouts", + help="CWD for rollouts; stray tool-written files go here.") args = p.parse_args(argv) logging.basicConfig(level=logging.INFO, format="%(message)s") + # Quiet the per-request HTTP / dataset-stream spam so the run log stays + # readable (rollout calls and dataset shards otherwise flood it). + for _noisy in ("httpx", "httpcore", "urllib3", "datasets", "fsspec", + "huggingface_hub", "openai"): + logging.getLogger(_noisy).setLevel(logging.WARNING) + + # Every run lands in its OWN datetimed folder, all collected flat under + # data/runs/ — so runs sort by time, never collide, and never litter data/. + # --out is treated as an optional LABEL for the folder, not a path; the data + # file inside is always data.jsonl (+ .stats.json/.pretty.json/.txt/.html). + # Folder name: _