mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-27 21:05:34 +00:00
feat(orchestrator): ADP-derived SFT cold-start data pipeline
Train a single ~9B orchestrator (Qwen3-8B) for local<->cloud routing via SFT on synthetic traces derived from NeuLab ADP (agent-data-collection). - sft_data/: full ADP trajectory loader, tier model + pricing, 4 paradigm renderers, reward-ranked selection, conversations-JSONL serializer, CLI - wire sft_trainer._generate_traces() to the pipeline (no GPU / no API keys) - Qwen3-8B SFT config + scripts/orchestrator CLI wrapper - offline fixture tests (12) Heuristic competence labels for the cold-start; real paradigm execution + GRPO are the documented v2.
This commit is contained in:
Andrew Park
@@ -0,0 +1,19 @@
|
||||
#!/usr/bin/env python
|
||||
"""Thin CLI wrapper: build the orchestrator SFT dataset from NeuLab ADP.
|
||||
|
||||
Equivalent to::
|
||||
|
||||
python -m openjarvis.learning.intelligence.orchestrator.sft_data.build [...]
|
||||
|
||||
Usage:
|
||||
python scripts/orchestrator/build_sft_data.py \
|
||||
--out data/orchestrator_sft_traces.jsonl \
|
||||
--max-tasks 2000 --adp-configs codeactinstruct,code_feedback,openhands
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.build import _main
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(_main())
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
# Orchestrator SFT cold-start — Qwen3-8B (~9B, "easier to train").
|
||||
#
|
||||
# Fields map onto OrchestratorSFTConfig. Build the dataset first (no GPU):
|
||||
# python -m openjarvis.learning.intelligence.orchestrator.sft_data.build \
|
||||
# --out data/orchestrator_sft_traces.jsonl --max-tasks 2000
|
||||
# then train on a rented GPU with regenerate_traces = false.
|
||||
|
||||
[model]
|
||||
model_name = "Qwen/Qwen3-8B"
|
||||
max_seq_length = 4096
|
||||
|
||||
[training]
|
||||
num_epochs = 3
|
||||
batch_size = 8
|
||||
learning_rate = 2e-5
|
||||
weight_decay = 0.01
|
||||
warmup_ratio = 0.1
|
||||
gradient_checkpointing = true
|
||||
|
||||
[data]
|
||||
trace_cache_path = "data/orchestrator_sft_traces.jsonl"
|
||||
regenerate_traces = false
|
||||
# Coding / agentic / tool-use ADP sub-configs (per the meeting steer).
|
||||
adp_configs = "codeactinstruct,code_feedback,openhands,agenttuning_os,swe-smith"
|
||||
distill_max_tasks = 2000
|
||||
|
||||
[checkpoint]
|
||||
checkpoint_dir = "checkpoints/orchestrator_sft"
|
||||
save_every_n_epochs = 1
|
||||
@@ -0,0 +1,55 @@
|
||||
"""Synthetic SFT-data generation for the orchestrator cold-start.
|
||||
|
||||
Turns NeuLab ADP (``neulab/agent-data-collection``) agent trajectories into
|
||||
THOUGHT/TOOL/INPUT ``conversations`` JSONL that
|
||||
:class:`~openjarvis.learning.intelligence.orchestrator.sft_trainer.OrchestratorSFTDataset`
|
||||
loads directly.
|
||||
|
||||
Pipeline::
|
||||
|
||||
ADP trajectory -> canonical Episode -> per-paradigm tiered renderings
|
||||
-> reward-ranked best-correct rendering -> conversations JSONL
|
||||
|
||||
Everything here runs with **no GPU and no API keys** (the cold-start does not
|
||||
re-execute models; it re-tiers the demonstrated ADP traces).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.adp_loader import (
|
||||
CanonicalStep,
|
||||
iter_trajectories,
|
||||
trajectory_rows_to_episode,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.build import (
|
||||
build_sft_dataset,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import (
|
||||
PARADIGMS,
|
||||
RenderedEpisode,
|
||||
render_all,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.select import select_best
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.serialize import to_record
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.tiers import (
|
||||
Difficulty,
|
||||
Tier,
|
||||
step_difficulty,
|
||||
tier_telemetry,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CanonicalStep",
|
||||
"Difficulty",
|
||||
"PARADIGMS",
|
||||
"RenderedEpisode",
|
||||
"Tier",
|
||||
"build_sft_dataset",
|
||||
"iter_trajectories",
|
||||
"render_all",
|
||||
"select_best",
|
||||
"step_difficulty",
|
||||
"tier_telemetry",
|
||||
"to_record",
|
||||
"trajectory_rows_to_episode",
|
||||
]
|
||||
@@ -0,0 +1,198 @@
|
||||
"""Full-trajectory reader for NeuLab ADP (``neulab/agent-data-collection``).
|
||||
|
||||
The eval-side loader (:mod:`openjarvis.evals.datasets.adp`) keeps only
|
||||
``(problem, reference)`` per row — it drops every intermediate step, which is
|
||||
exactly the signal the orchestrator needs. This module instead transcribes
|
||||
**all** turns of a trajectory into a canonical
|
||||
:class:`~openjarvis.learning.intelligence.orchestrator.types.Episode`.
|
||||
|
||||
No GPU / no API keys: we read the demonstrated traces and re-tier them
|
||||
downstream; we never re-execute a model here.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ast
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable, Iterator, List, MutableMapping, Optional
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.types import (
|
||||
Episode,
|
||||
OrchestratorAction,
|
||||
OrchestratorObservation,
|
||||
)
|
||||
|
||||
HF_DATASET_ID = "neulab/agent-data-collection"
|
||||
HF_SPLIT = "std"
|
||||
|
||||
# Same configs the eval loader concatenates; the meeting steer favours the
|
||||
# coding / agentic / tool-use ones.
|
||||
DEFAULT_CONFIGS: tuple[str, ...] = (
|
||||
"codeactinstruct",
|
||||
"code_feedback",
|
||||
"openhands",
|
||||
"agenttuning_os",
|
||||
"agenttuning_db",
|
||||
"swe-smith",
|
||||
)
|
||||
|
||||
# ADP turn ``class_`` values that denote a tool/code action vs. plain message.
|
||||
_CODE_CLASSES = {"code_action", "ipython_action", "bash_action"}
|
||||
_SEARCH_CLASSES = {"search_action", "browse_action", "web_action"}
|
||||
_MESSAGE_CLASSES = {"message_action"}
|
||||
|
||||
|
||||
@dataclass
|
||||
class CanonicalStep:
|
||||
"""One transcribed ADP turn: what the demonstration did at this step."""
|
||||
|
||||
kind: str
|
||||
"""Normalised step kind: ``reason`` | ``code`` | ``search`` | ``message``."""
|
||||
|
||||
content: str
|
||||
"""The turn's text (the action the demonstration took)."""
|
||||
|
||||
observation: str = ""
|
||||
"""The environment/tool result that followed, if any."""
|
||||
|
||||
is_final: bool = False
|
||||
"""Whether this is the trajectory's final answer turn."""
|
||||
|
||||
|
||||
def _parse_content(raw: object) -> List[MutableMapping[str, object]]:
|
||||
"""Parse the ``content`` field (a list, or a string repr of one)."""
|
||||
if isinstance(raw, list):
|
||||
return raw # type: ignore[return-value]
|
||||
if isinstance(raw, str):
|
||||
try:
|
||||
parsed = ast.literal_eval(raw)
|
||||
if isinstance(parsed, list):
|
||||
return parsed # type: ignore[return-value]
|
||||
except (ValueError, SyntaxError):
|
||||
pass
|
||||
return []
|
||||
|
||||
|
||||
def _classify(turn: MutableMapping[str, object]) -> str:
|
||||
cls = str(turn.get("class_") or "").lower()
|
||||
if cls in _CODE_CLASSES:
|
||||
return "code"
|
||||
if cls in _SEARCH_CLASSES:
|
||||
return "search"
|
||||
if cls in _MESSAGE_CLASSES:
|
||||
return "message"
|
||||
return "reason"
|
||||
|
||||
|
||||
def trajectory_rows_to_episode(
|
||||
record_id: str,
|
||||
turns: List[MutableMapping[str, object]],
|
||||
) -> Optional[Episode]:
|
||||
"""Transcribe one ADP trajectory's turns into a canonical ``Episode``.
|
||||
|
||||
The orchestrator-relevant structure is preserved: the first user turn is the
|
||||
problem; every subsequent agent turn becomes a step (with its following
|
||||
observation, if the trace recorded one). Returns ``None`` if there is no
|
||||
usable problem or no agent steps.
|
||||
"""
|
||||
problem: Optional[str] = None
|
||||
steps: List[CanonicalStep] = []
|
||||
|
||||
# Group agent turns with the observation/user turn that follows them.
|
||||
pending: Optional[CanonicalStep] = None
|
||||
for turn in turns:
|
||||
source = str(turn.get("source") or "").lower()
|
||||
text = str(turn.get("content") or "").strip()
|
||||
if not text:
|
||||
continue
|
||||
|
||||
if source == "user":
|
||||
if problem is None:
|
||||
problem = text
|
||||
elif pending is not None:
|
||||
# A user/environment turn after an agent action = its observation.
|
||||
pending.observation = text[:2000]
|
||||
continue
|
||||
|
||||
# Agent-sourced turn -> a new step.
|
||||
if pending is not None:
|
||||
steps.append(pending)
|
||||
pending = CanonicalStep(kind=_classify(turn), content=text)
|
||||
|
||||
if pending is not None:
|
||||
steps.append(pending)
|
||||
|
||||
if not problem or not steps:
|
||||
return None
|
||||
|
||||
steps[-1].is_final = True
|
||||
steps[-1].kind = "message" if steps[-1].kind == "reason" else steps[-1].kind
|
||||
|
||||
episode = Episode(
|
||||
task_id=record_id,
|
||||
initial_prompt=problem,
|
||||
ground_truth=steps[-1].content[:2000],
|
||||
# ADP rows are demonstrated solutions; treat them as correct teachers.
|
||||
correct=True,
|
||||
)
|
||||
for st in steps:
|
||||
episode.add_step(
|
||||
OrchestratorAction(
|
||||
thought="", # filled in by the paradigm renderer
|
||||
tool_name=st.kind,
|
||||
tool_input=st.content,
|
||||
is_final_answer=st.is_final,
|
||||
),
|
||||
OrchestratorObservation(content=st.observation or st.content),
|
||||
)
|
||||
episode.final_answer = steps[-1].content[:2000]
|
||||
episode.metadata["source"] = "adp"
|
||||
return episode
|
||||
|
||||
|
||||
def iter_trajectories(
|
||||
*,
|
||||
max_tasks: Optional[int] = None,
|
||||
configs: Iterable[str] = DEFAULT_CONFIGS,
|
||||
min_steps: int = 1,
|
||||
max_steps: int = 24,
|
||||
) -> Iterator[Episode]:
|
||||
"""Stream canonical ``Episode`` objects from ADP.
|
||||
|
||||
Network is touched lazily inside the loop (``datasets.load_dataset`` with
|
||||
``streaming=True``) so importing this module stays free. Configs that fail
|
||||
to load (gated/missing) are skipped.
|
||||
"""
|
||||
from datasets import load_dataset
|
||||
|
||||
emitted = 0
|
||||
for cfg in configs:
|
||||
if max_tasks is not None and emitted >= max_tasks:
|
||||
break
|
||||
try:
|
||||
stream = load_dataset(HF_DATASET_ID, cfg, split=HF_SPLIT, streaming=True)
|
||||
except Exception:
|
||||
continue
|
||||
for i, row in enumerate(stream):
|
||||
if max_tasks is not None and emitted >= max_tasks:
|
||||
break
|
||||
row = dict(row) # type: ignore[arg-type]
|
||||
turns = _parse_content(row.get("content"))
|
||||
row_id = row.get("id")
|
||||
rec_id = str(row_id) if row_id is not None else f"{cfg}-{i}"
|
||||
episode = trajectory_rows_to_episode(rec_id, turns)
|
||||
if episode is None:
|
||||
continue
|
||||
n = episode.num_turns()
|
||||
if n < min_steps or n > max_steps:
|
||||
continue
|
||||
emitted += 1
|
||||
yield episode
|
||||
|
||||
|
||||
__all__ = [
|
||||
"CanonicalStep",
|
||||
"DEFAULT_CONFIGS",
|
||||
"iter_trajectories",
|
||||
"trajectory_rows_to_episode",
|
||||
]
|
||||
@@ -0,0 +1,112 @@
|
||||
"""End-to-end SFT dataset builder + CLI.
|
||||
|
||||
ADP trajectory -> canonical Episode -> render_all -> select_best
|
||||
-> to_record -> JSONL (+ a sidecar ``.stats.json``)
|
||||
|
||||
Runs with no GPU and no API keys (the cold-start re-tiers demonstrated traces;
|
||||
it does not execute models). Network is only touched to stream ADP rows.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import logging
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
from typing import Callable, Iterable, Iterator, Optional
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.reward import (
|
||||
MultiObjectiveReward,
|
||||
Normalizers,
|
||||
RewardWeights,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.adp_loader import (
|
||||
DEFAULT_CONFIGS,
|
||||
iter_trajectories,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import render_all
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.select import select_best
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.serialize import to_record
|
||||
from openjarvis.learning.intelligence.orchestrator.types import Episode
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def build_sft_dataset(
|
||||
out_path: str,
|
||||
*,
|
||||
max_tasks: Optional[int] = 2000,
|
||||
configs: Iterable[str] = DEFAULT_CONFIGS,
|
||||
source: Optional[Callable[..., Iterator[Episode]]] = None,
|
||||
) -> dict:
|
||||
"""Build the SFT JSONL at ``out_path`` and return stats.
|
||||
|
||||
``source`` overrides the ADP stream (used by tests to inject fixtures); it
|
||||
must be a callable returning an iterator of canonical :class:`Episode`.
|
||||
"""
|
||||
reward = MultiObjectiveReward(RewardWeights(), Normalizers())
|
||||
out = Path(out_path)
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
episodes = (
|
||||
source(max_tasks=max_tasks, configs=configs)
|
||||
if source is not None
|
||||
else iter_trajectories(max_tasks=max_tasks, configs=configs)
|
||||
)
|
||||
|
||||
seen = 0
|
||||
written = 0
|
||||
dropped = 0
|
||||
paradigm_counts: Counter[str] = Counter()
|
||||
|
||||
with out.open("w") as fh:
|
||||
for episode in episodes:
|
||||
seen += 1
|
||||
best = select_best(render_all(episode), reward=reward)
|
||||
if best is None:
|
||||
dropped += 1
|
||||
continue
|
||||
record = to_record(best, reward=reward.compute(best.episode))
|
||||
fh.write(json.dumps(record) + "\n")
|
||||
written += 1
|
||||
paradigm_counts[best.paradigm] += 1
|
||||
|
||||
stats = {
|
||||
"out_path": str(out),
|
||||
"tasks_seen": seen,
|
||||
"records_written": written,
|
||||
"tasks_dropped": dropped,
|
||||
"paradigm_distribution": dict(paradigm_counts),
|
||||
}
|
||||
stats_path = out.with_suffix(out.suffix + ".stats.json")
|
||||
stats_path.write_text(json.dumps(stats, indent=2))
|
||||
logger.info("Wrote %d SFT records to %s (%s)", written, out, dict(paradigm_counts))
|
||||
return stats
|
||||
|
||||
|
||||
def _main(argv: Optional[list[str]] = None) -> int:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Build orchestrator SFT data from ADP."
|
||||
)
|
||||
parser.add_argument("--out", default="data/orchestrator_sft_traces.jsonl")
|
||||
parser.add_argument("--max-tasks", type=int, default=2000)
|
||||
parser.add_argument(
|
||||
"--adp-configs",
|
||||
default=",".join(DEFAULT_CONFIGS),
|
||||
help="comma-separated ADP sub-configs to stream",
|
||||
)
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
||||
configs = [c.strip() for c in args.adp_configs.split(",") if c.strip()]
|
||||
stats = build_sft_dataset(args.out, max_tasks=args.max_tasks, configs=configs)
|
||||
print(json.dumps(stats, indent=2))
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__": # pragma: no cover
|
||||
raise SystemExit(_main())
|
||||
|
||||
|
||||
__all__ = ["build_sft_dataset"]
|
||||
@@ -0,0 +1,188 @@
|
||||
"""Render a canonical ADP ``Episode`` under each local↔cloud paradigm.
|
||||
|
||||
Each renderer assigns a model **tier** to every step according to that
|
||||
paradigm's policy, attaches the estimated telemetry (cost/energy/latency/power)
|
||||
to the step's observation, and predicts whether the rendering would still solve
|
||||
the task. Selection (``select.py``) then keeps the cheapest predicted-correct
|
||||
rendering per task.
|
||||
|
||||
These are the cold-start *tiering* scaffolds of the real
|
||||
``agents/hybrid/*`` paradigms — same names, same local-first spirit, no live
|
||||
execution. Swapping in real runs is the v2 item in the design doc.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Callable, List
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.tiers import (
|
||||
Difficulty,
|
||||
Tier,
|
||||
covers,
|
||||
min_covering_tier,
|
||||
step_difficulty,
|
||||
tier_name,
|
||||
tier_telemetry,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.types import Episode
|
||||
|
||||
# Detected in an observation -> the demonstration had to retry -> harder task.
|
||||
_RETRY_MARKERS = (
|
||||
"your answer is wrong",
|
||||
"incorrect",
|
||||
"traceback",
|
||||
"error:",
|
||||
"failed",
|
||||
"try again",
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RenderedEpisode:
|
||||
"""An ``Episode`` re-tiered under one paradigm, plus its prediction."""
|
||||
|
||||
paradigm: str
|
||||
episode: Episode
|
||||
"""Deep copy with per-step tier in ``action.tool_name`` and telemetry filled."""
|
||||
|
||||
step_tiers: List[str] = field(default_factory=list)
|
||||
"""Tier key (``local``/``mid``/``frontier``/``search``) chosen per step."""
|
||||
|
||||
predicted_correct: bool = False
|
||||
|
||||
|
||||
def _difficulties(episode: Episode) -> List[Difficulty]:
|
||||
"""Per-step difficulty, with a trajectory-wide retry bump."""
|
||||
retry = any(
|
||||
any(m in (step.observation.content or "").lower() for m in _RETRY_MARKERS)
|
||||
for step in episode.steps
|
||||
)
|
||||
out: List[Difficulty] = []
|
||||
for step in episode.steps:
|
||||
kind = step.action.tool_name
|
||||
out.append(step_difficulty(kind, step.action.tool_input, retry_signal=retry))
|
||||
return out
|
||||
|
||||
|
||||
def _apply(
|
||||
episode: Episode,
|
||||
paradigm: str,
|
||||
tier_keys: List[str],
|
||||
predicted_correct: bool,
|
||||
) -> RenderedEpisode:
|
||||
"""Build a RenderedEpisode: stamp each step with its tier + telemetry."""
|
||||
ep = copy.deepcopy(episode)
|
||||
ep.total_cost_usd = 0.0
|
||||
ep.total_energy_joules = 0.0
|
||||
ep.total_latency_seconds = 0.0
|
||||
ep.total_tokens = 0
|
||||
ep.max_power_watts = 0.0
|
||||
|
||||
for step, tier_key in zip(ep.steps, tier_keys):
|
||||
tel = tier_telemetry(tier_key)
|
||||
step.observation.cost_usd = tel["cost_usd"]
|
||||
step.observation.energy_joules = tel["energy_joules"]
|
||||
step.observation.latency_seconds = tel["latency_seconds"]
|
||||
step.observation.power_watts = tel["power_watts"]
|
||||
step.observation.tokens = int(tel["tokens"])
|
||||
# Record the routing decision on the action for the serializer.
|
||||
step.action.thought = _thought_for(tier_key, step.action.tool_name)
|
||||
ep.total_cost_usd += tel["cost_usd"]
|
||||
ep.total_energy_joules += tel["energy_joules"]
|
||||
ep.total_latency_seconds += tel["latency_seconds"]
|
||||
ep.total_tokens += int(tel["tokens"])
|
||||
ep.max_power_watts = max(ep.max_power_watts, tel["power_watts"])
|
||||
|
||||
ep.correct = predicted_correct
|
||||
ep.metadata["paradigm"] = paradigm
|
||||
return RenderedEpisode(
|
||||
paradigm=paradigm,
|
||||
episode=ep,
|
||||
step_tiers=tier_keys,
|
||||
predicted_correct=predicted_correct,
|
||||
)
|
||||
|
||||
|
||||
def _thought_for(tier_key: str, kind: str) -> str:
|
||||
if kind == "search" or tier_key == "search":
|
||||
return "This step needs retrieval; route to web_search."
|
||||
reason = {
|
||||
"local": "Cheap and on-device; the local model can handle this step.",
|
||||
"mid": "Beyond easy; escalate to a cheap mid-tier cloud model.",
|
||||
"frontier": "Hard step; escalate to the frontier cloud model.",
|
||||
}[tier_key]
|
||||
return reason
|
||||
|
||||
|
||||
def _tier_key_for_step(kind: str, tier: Tier) -> str:
|
||||
return "search" if kind == "search" else tier_name(tier)
|
||||
|
||||
|
||||
# --- paradigm renderers -----------------------------------------------------
|
||||
|
||||
|
||||
def render_baseline_local(episode: Episode) -> RenderedEpisode:
|
||||
diffs = _difficulties(episode)
|
||||
tiers = [
|
||||
_tier_key_for_step(s.action.tool_name, Tier.LOCAL) for s in episode.steps
|
||||
]
|
||||
predicted = all(d == Difficulty.EASY for d in diffs)
|
||||
return _apply(episode, "baseline_local", tiers, predicted)
|
||||
|
||||
|
||||
def render_baseline_cloud(episode: Episode) -> RenderedEpisode:
|
||||
tiers = [
|
||||
_tier_key_for_step(s.action.tool_name, Tier.FRONTIER) for s in episode.steps
|
||||
]
|
||||
return _apply(episode, "baseline_cloud", tiers, True)
|
||||
|
||||
|
||||
def render_advisor(episode: Episode) -> RenderedEpisode:
|
||||
"""Local executor, frontier rewrite of the final answer."""
|
||||
diffs = _difficulties(episode)
|
||||
tiers: List[str] = []
|
||||
n = len(episode.steps)
|
||||
for i, s in enumerate(episode.steps):
|
||||
tier = Tier.FRONTIER if i == n - 1 else Tier.LOCAL
|
||||
tiers.append(_tier_key_for_step(s.action.tool_name, tier))
|
||||
predicted = all(d == Difficulty.EASY for d in diffs[:-1]) if n > 1 else True
|
||||
return _apply(episode, "advisor", tiers, predicted)
|
||||
|
||||
|
||||
def render_toolorchestra(episode: Episode) -> RenderedEpisode:
|
||||
"""Local-first per-step: assign each step its minimum covering tier."""
|
||||
diffs = _difficulties(episode)
|
||||
tiers = [
|
||||
_tier_key_for_step(s.action.tool_name, min_covering_tier(d))
|
||||
for s, d in zip(episode.steps, diffs)
|
||||
]
|
||||
predicted = all(
|
||||
covers(min_covering_tier(d), d) for d in diffs
|
||||
) # True by construction
|
||||
return _apply(episode, "toolorchestra", tiers, predicted)
|
||||
|
||||
|
||||
PARADIGMS: dict[str, Callable[[Episode], RenderedEpisode]] = {
|
||||
"baseline_local": render_baseline_local,
|
||||
"baseline_cloud": render_baseline_cloud,
|
||||
"advisor": render_advisor,
|
||||
"toolorchestra": render_toolorchestra,
|
||||
}
|
||||
|
||||
|
||||
def render_all(episode: Episode) -> List[RenderedEpisode]:
|
||||
"""Render ``episode`` under every paradigm."""
|
||||
return [render(episode) for render in PARADIGMS.values()]
|
||||
|
||||
|
||||
__all__ = [
|
||||
"PARADIGMS",
|
||||
"RenderedEpisode",
|
||||
"render_advisor",
|
||||
"render_all",
|
||||
"render_baseline_cloud",
|
||||
"render_baseline_local",
|
||||
"render_toolorchestra",
|
||||
]
|
||||
@@ -0,0 +1,42 @@
|
||||
"""Pick the best paradigm rendering per task.
|
||||
|
||||
The label for SFT is the *cheapest strategy that still solves the task*: among
|
||||
the predicted-correct renderings, keep the one with the highest multi-objective
|
||||
reward (accuracy minus cost/energy/latency/power). This is the local-first
|
||||
thesis expressed as a training target.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List, Optional
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.reward import (
|
||||
MultiObjectiveReward,
|
||||
Normalizers,
|
||||
RewardWeights,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import (
|
||||
RenderedEpisode,
|
||||
)
|
||||
|
||||
_DEFAULT_REWARD = MultiObjectiveReward(RewardWeights(), Normalizers())
|
||||
|
||||
|
||||
def select_best(
|
||||
renderings: List[RenderedEpisode],
|
||||
*,
|
||||
reward: Optional[MultiObjectiveReward] = None,
|
||||
) -> Optional[RenderedEpisode]:
|
||||
"""Return the highest-reward predicted-correct rendering, or ``None``.
|
||||
|
||||
``None`` means no paradigm was predicted to solve the task — the task is
|
||||
dropped from the SFT set rather than teaching a wrong trajectory.
|
||||
"""
|
||||
scorer = reward or _DEFAULT_REWARD
|
||||
correct = [r for r in renderings if r.predicted_correct]
|
||||
if not correct:
|
||||
return None
|
||||
return max(correct, key=lambda r: scorer.compute(r.episode))
|
||||
|
||||
|
||||
__all__ = ["select_best"]
|
||||
@@ -0,0 +1,103 @@
|
||||
"""Serialize a winning :class:`RenderedEpisode` into a ``conversations`` record.
|
||||
|
||||
Output schema matches what
|
||||
:class:`~openjarvis.learning.intelligence.orchestrator.sft_trainer.OrchestratorSFTDataset`
|
||||
already consumes::
|
||||
|
||||
{"conversations": [{"role": "system"|"user"|"assistant"|"tool", "content": ...}],
|
||||
"paradigm": ..., "reward": ..., "metrics": {...}}
|
||||
|
||||
The assistant turns use the canonical THOUGHT/TOOL/INPUT format from
|
||||
``prompt_registry``; the routing tool encodes the tier the orchestrator chose
|
||||
(``local_model`` / ``mid_model`` / ``frontier_model`` / ``web_search``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.prompt_registry import (
|
||||
build_system_prompt,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import (
|
||||
RenderedEpisode,
|
||||
)
|
||||
|
||||
# Routing vocabulary the orchestrator learns to emit.
|
||||
TIER_TOOL = {
|
||||
"local": "local_model",
|
||||
"mid": "mid_model",
|
||||
"frontier": "frontier_model",
|
||||
"search": "web_search",
|
||||
}
|
||||
ROUTING_TOOLS = ["local_model", "mid_model", "frontier_model", "web_search"]
|
||||
|
||||
_SYSTEM_PROMPT = build_system_prompt(ROUTING_TOOLS)
|
||||
|
||||
|
||||
def to_record(rendered: RenderedEpisode, *, reward: float = 0.0) -> Dict[str, Any]:
|
||||
"""Convert a winning rendering into one SFT JSONL record."""
|
||||
ep = rendered.episode
|
||||
conversations: list[dict[str, str]] = [
|
||||
{"role": "system", "content": _SYSTEM_PROMPT},
|
||||
{"role": "user", "content": ep.initial_prompt},
|
||||
]
|
||||
|
||||
n = len(ep.steps)
|
||||
for i, (step, tier_key) in enumerate(zip(ep.steps, rendered.step_tiers)):
|
||||
tool = TIER_TOOL.get(tier_key, "local_model")
|
||||
if step.action.is_final_answer or i == n - 1:
|
||||
conversations.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": (
|
||||
f"THOUGHT: {step.action.thought}\n"
|
||||
f"TOOL: {tool}\n"
|
||||
f"INPUT: {step.action.tool_input}"
|
||||
),
|
||||
}
|
||||
)
|
||||
conversations.append(
|
||||
{"role": "tool", "name": tool, "content": step.observation.content}
|
||||
)
|
||||
conversations.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": (
|
||||
"THOUGHT: The result answers the task.\n"
|
||||
f"FINAL_ANSWER: {ep.final_answer}"
|
||||
),
|
||||
}
|
||||
)
|
||||
else:
|
||||
conversations.append(
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": (
|
||||
f"THOUGHT: {step.action.thought}\n"
|
||||
f"TOOL: {tool}\n"
|
||||
f"INPUT: {step.action.tool_input}"
|
||||
),
|
||||
}
|
||||
)
|
||||
conversations.append(
|
||||
{"role": "tool", "name": tool, "content": step.observation.content}
|
||||
)
|
||||
|
||||
return {
|
||||
"conversations": conversations,
|
||||
"task_id": ep.task_id,
|
||||
"paradigm": rendered.paradigm,
|
||||
"reward": reward,
|
||||
"metrics": {
|
||||
"cost_usd": ep.total_cost_usd,
|
||||
"energy_joules": ep.total_energy_joules,
|
||||
"latency_seconds": ep.total_latency_seconds,
|
||||
"tokens": ep.total_tokens,
|
||||
"max_power_watts": ep.max_power_watts,
|
||||
"num_steps": n,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
__all__ = ["ROUTING_TOOLS", "TIER_TOOL", "to_record"]
|
||||
@@ -0,0 +1,119 @@
|
||||
"""Model tiers, per-tier telemetry estimates, and step-difficulty heuristics.
|
||||
|
||||
This is the *one estimate* of the cold-start (see the design doc): because we
|
||||
do not execute models, we approximate "what would this step cost on tier X?"
|
||||
with flat per-tier telemetry, and "how hard is this step?" with a heuristic over
|
||||
the ADP step kind and the trajectory's retry signals.
|
||||
|
||||
Costs reuse the authoritative hybrid pricing table so the orchestrator is
|
||||
trained against the same dollars the paradigm harness charges.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import IntEnum
|
||||
|
||||
from openjarvis.agents.hybrid import _prices
|
||||
|
||||
# Representative model per tier (drives the $ side of telemetry).
|
||||
TIER_MODEL: dict[str, str] = {
|
||||
"local": "qwen3:8b", # local vLLM -> unknown to PRICES -> $0
|
||||
"mid": "gemini-2.5-flash",
|
||||
"frontier": "claude-sonnet-4-6",
|
||||
"search": "qwen3:8b", # retrieval is a tool call, not a model tier
|
||||
}
|
||||
|
||||
# Rough token budgets per step kind, for the cost estimate.
|
||||
_TOKENS_IN = 800
|
||||
_TOKENS_OUT = {"local": 400, "mid": 400, "frontier": 600, "search": 80}
|
||||
|
||||
# Non-$ telemetry per step (joules / seconds / watts). Local work burns some
|
||||
# on-device energy/power but no dollars; cloud work is ~free locally but costs
|
||||
# dollars and adds latency. Tuned (against reward.py's default weights and
|
||||
# normalizers) so the local-first ordering holds: on a step a tier can cover,
|
||||
# local beats cloud, and per-step escalation beats all-frontier. These are the
|
||||
# tunable knobs of the cold-start estimate — see the design doc.
|
||||
_ENERGY_J = {"local": 5.0, "mid": 0.5, "frontier": 0.5, "search": 0.2}
|
||||
_LATENCY_S = {"local": 1.5, "mid": 1.5, "frontier": 4.0, "search": 1.0}
|
||||
_POWER_W = {"local": 20.0, "mid": 5.0, "frontier": 5.0, "search": 5.0}
|
||||
_SEARCH_COST_USD = 0.001
|
||||
|
||||
|
||||
class Tier(IntEnum):
|
||||
"""Model capability tier; ordered so a higher tier covers a harder step."""
|
||||
|
||||
LOCAL = 0
|
||||
MID = 1
|
||||
FRONTIER = 2
|
||||
|
||||
|
||||
class Difficulty(IntEnum):
|
||||
"""Estimated step difficulty; compared against :class:`Tier` rank."""
|
||||
|
||||
EASY = 0
|
||||
MED = 1
|
||||
HARD = 2
|
||||
|
||||
|
||||
_TIER_NAME = {Tier.LOCAL: "local", Tier.MID: "mid", Tier.FRONTIER: "frontier"}
|
||||
|
||||
|
||||
def tier_name(tier: Tier) -> str:
|
||||
return _TIER_NAME[tier]
|
||||
|
||||
|
||||
def tier_telemetry(tier_key: str) -> dict[str, float]:
|
||||
"""Estimated ``(cost, energy, latency, power, tokens)`` for one step at a tier.
|
||||
|
||||
``tier_key`` is one of ``local``/``mid``/``frontier``/``search``.
|
||||
"""
|
||||
out_tokens = _TOKENS_OUT[tier_key]
|
||||
if tier_key == "search":
|
||||
cost = _SEARCH_COST_USD
|
||||
else:
|
||||
cost = _prices.cost(TIER_MODEL[tier_key], _TOKENS_IN, out_tokens)
|
||||
return {
|
||||
"cost_usd": cost,
|
||||
"energy_joules": _ENERGY_J[tier_key],
|
||||
"latency_seconds": _LATENCY_S[tier_key],
|
||||
"power_watts": _POWER_W[tier_key],
|
||||
"tokens": float(_TOKENS_IN + out_tokens),
|
||||
}
|
||||
|
||||
|
||||
def step_difficulty(
|
||||
kind: str, content: str, *, retry_signal: bool = False
|
||||
) -> Difficulty:
|
||||
"""Heuristic difficulty for a canonical step.
|
||||
|
||||
- ``code`` actions are at least MED (real execution / logic).
|
||||
- very long actions, or any step in a trajectory that hit a retry signal
|
||||
("your answer is wrong" etc.), bump to HARD.
|
||||
- everything else (plain messages / short reasoning) is EASY.
|
||||
"""
|
||||
base = Difficulty.MED if kind == "code" else Difficulty.EASY
|
||||
if retry_signal or len(content) > 1500:
|
||||
return Difficulty.HARD
|
||||
return base
|
||||
|
||||
|
||||
def covers(tier: Tier, difficulty: Difficulty) -> bool:
|
||||
"""True iff ``tier`` is capable enough for a step of this ``difficulty``."""
|
||||
return int(tier) >= int(difficulty)
|
||||
|
||||
|
||||
def min_covering_tier(difficulty: Difficulty) -> Tier:
|
||||
"""The cheapest tier that still covers ``difficulty`` (local-first optimum)."""
|
||||
return Tier(int(difficulty))
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Difficulty",
|
||||
"Tier",
|
||||
"TIER_MODEL",
|
||||
"covers",
|
||||
"min_covering_tier",
|
||||
"step_difficulty",
|
||||
"tier_name",
|
||||
"tier_telemetry",
|
||||
]
|
||||
@@ -74,6 +74,11 @@ class OrchestratorSFTConfig:
|
||||
trace_cache_path: str = "data/orchestrator_sft_traces.jsonl"
|
||||
regenerate_traces: bool = False
|
||||
|
||||
# ADP cold-start trace synthesis (see sft_data/build.py). Empty configs
|
||||
# falls back to sft_data.adp_loader.DEFAULT_CONFIGS.
|
||||
adp_configs: str = ""
|
||||
distill_max_tasks: int = 2000
|
||||
|
||||
# Checkpoint
|
||||
checkpoint_dir: str = "checkpoints/orchestrator_sft"
|
||||
save_every_n_epochs: int = 1
|
||||
@@ -267,11 +272,34 @@ class OrchestratorSFTTrainer:
|
||||
)
|
||||
|
||||
def _generate_traces(self) -> None:
|
||||
"""Generate SFT traces (placeholder — requires running engine)."""
|
||||
"""Generate SFT traces from the ADP corpus (no GPU / no API keys).
|
||||
|
||||
Builds the THOUGHT/TOOL/INPUT ``conversations`` JSONL by re-tiering
|
||||
NeuLab ADP trajectories through the local↔cloud paradigms and keeping
|
||||
the cheapest predicted-correct rendering per task. See
|
||||
``sft_data/build.py`` and the cold-start design doc.
|
||||
"""
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.build import (
|
||||
build_sft_dataset,
|
||||
)
|
||||
|
||||
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()
|
||||
configs = (
|
||||
[c.strip() for c in self.config.adp_configs.split(",") if c.strip()]
|
||||
or None
|
||||
)
|
||||
try:
|
||||
stats = build_sft_dataset(
|
||||
str(trace_path),
|
||||
max_tasks=self.config.distill_max_tasks,
|
||||
**({"configs": configs} if configs else {}),
|
||||
)
|
||||
logger.info("Generated SFT traces: %s", stats)
|
||||
except Exception as exc: # network/datasets optional — fail soft to empty
|
||||
logger.warning("ADP 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:
|
||||
|
||||
@@ -0,0 +1,187 @@
|
||||
"""Offline tests for the orchestrator SFT cold-start pipeline (no network/GPU).
|
||||
|
||||
Covers: ADP transcription -> canonical Episode, difficulty/tier heuristics,
|
||||
paradigm rendering, reward-ranked selection, serialization round-trip, and the
|
||||
end-to-end builder via an injected fixture source.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.adp_loader import (
|
||||
trajectory_rows_to_episode,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.build import (
|
||||
build_sft_dataset,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.paradigms import (
|
||||
PARADIGMS,
|
||||
render_all,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.select import select_best
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.serialize import to_record
|
||||
from openjarvis.learning.intelligence.orchestrator.sft_data.tiers import (
|
||||
Difficulty,
|
||||
Tier,
|
||||
covers,
|
||||
min_covering_tier,
|
||||
step_difficulty,
|
||||
tier_telemetry,
|
||||
)
|
||||
from openjarvis.learning.intelligence.orchestrator.types import Episode
|
||||
|
||||
# --- fixtures ---------------------------------------------------------------
|
||||
|
||||
_EASY_TURNS = [
|
||||
{"source": "user", "class_": "message_action", "content": "What is 2 + 2?"},
|
||||
{"source": "agent", "class_": "message_action", "content": "The answer is 4."},
|
||||
]
|
||||
|
||||
_CODE_TURNS = [
|
||||
{
|
||||
"source": "user",
|
||||
"class_": "message_action",
|
||||
"content": "Sort this list in Python.",
|
||||
},
|
||||
{"source": "agent", "class_": "code_action", "content": "sorted([3, 1, 2])"},
|
||||
{"source": "user", "class_": "message_action", "content": "[1, 2, 3]"},
|
||||
{"source": "agent", "class_": "message_action", "content": "Sorted: [1, 2, 3]"},
|
||||
]
|
||||
|
||||
|
||||
def _easy_episode() -> Episode:
|
||||
ep = trajectory_rows_to_episode("easy-1", _EASY_TURNS)
|
||||
assert ep is not None
|
||||
return ep
|
||||
|
||||
|
||||
def _code_episode() -> Episode:
|
||||
ep = trajectory_rows_to_episode("code-1", _CODE_TURNS)
|
||||
assert ep is not None
|
||||
return ep
|
||||
|
||||
|
||||
# --- adp_loader -------------------------------------------------------------
|
||||
|
||||
|
||||
def test_transcribe_keeps_all_steps_and_problem():
|
||||
ep = _code_episode()
|
||||
assert ep.initial_prompt == "Sort this list in Python."
|
||||
assert ep.num_turns() == 2 # two agent turns
|
||||
assert ep.steps[0].action.tool_name == "code"
|
||||
assert ep.steps[-1].action.is_final_answer is True
|
||||
assert ep.correct is True # ADP rows are demonstrated solutions
|
||||
|
||||
|
||||
def test_transcribe_rejects_empty():
|
||||
assert trajectory_rows_to_episode("x", []) is None
|
||||
assert trajectory_rows_to_episode("x", [{"source": "user", "content": ""}]) is None
|
||||
|
||||
|
||||
# --- tiers ------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_difficulty_heuristic():
|
||||
assert step_difficulty("message", "short") == Difficulty.EASY
|
||||
assert step_difficulty("code", "sorted([])") == Difficulty.MED
|
||||
assert step_difficulty("message", "x", retry_signal=True) == Difficulty.HARD
|
||||
|
||||
|
||||
def test_min_covering_tier_and_covers():
|
||||
assert min_covering_tier(Difficulty.EASY) == Tier.LOCAL
|
||||
assert min_covering_tier(Difficulty.MED) == Tier.MID
|
||||
assert covers(Tier.FRONTIER, Difficulty.HARD)
|
||||
assert not covers(Tier.LOCAL, Difficulty.HARD)
|
||||
|
||||
|
||||
def test_tier_telemetry_local_is_free_cloud_costs():
|
||||
assert tier_telemetry("local")["cost_usd"] == 0.0
|
||||
assert tier_telemetry("frontier")["cost_usd"] > 0.0
|
||||
local_e = tier_telemetry("local")["energy_joules"]
|
||||
mid_e = tier_telemetry("mid")["energy_joules"]
|
||||
assert local_e > mid_e
|
||||
|
||||
|
||||
# --- paradigms --------------------------------------------------------------
|
||||
|
||||
|
||||
def test_render_all_covers_every_paradigm():
|
||||
rendered = render_all(_code_episode())
|
||||
assert {r.paradigm for r in rendered} == set(PARADIGMS)
|
||||
|
||||
|
||||
def test_baseline_local_incorrect_on_code_task():
|
||||
rendered = {r.paradigm: r for r in render_all(_code_episode())}
|
||||
# code step is MED -> local can't cover it
|
||||
assert rendered["baseline_local"].predicted_correct is False
|
||||
assert rendered["baseline_cloud"].predicted_correct is True
|
||||
assert rendered["toolorchestra"].predicted_correct is True
|
||||
|
||||
|
||||
def test_toolorchestra_escalates_code_step():
|
||||
rendered = {r.paradigm: r for r in render_all(_code_episode())}
|
||||
tiers = rendered["toolorchestra"].step_tiers
|
||||
assert tiers[0] == "mid" # code step escalated off local
|
||||
assert rendered["baseline_local"].step_tiers[0] == "local"
|
||||
|
||||
|
||||
# --- select -----------------------------------------------------------------
|
||||
|
||||
|
||||
def test_select_prefers_local_on_easy_task():
|
||||
best = select_best(render_all(_easy_episode()))
|
||||
assert best is not None
|
||||
# all-easy -> baseline_local is cheapest-correct
|
||||
assert best.paradigm == "baseline_local"
|
||||
|
||||
|
||||
def test_select_drops_unsolved_task():
|
||||
ep = _easy_episode()
|
||||
# force every paradigm to be predicted-incorrect
|
||||
renderings = render_all(ep)
|
||||
for r in renderings:
|
||||
r.predicted_correct = False
|
||||
assert select_best(renderings) is None
|
||||
|
||||
|
||||
# --- serialize --------------------------------------------------------------
|
||||
|
||||
|
||||
def test_serialize_schema_and_final_answer():
|
||||
best = select_best(render_all(_code_episode()))
|
||||
assert best is not None
|
||||
rec = to_record(best, reward=0.5)
|
||||
convo = rec["conversations"]
|
||||
assert convo[0]["role"] == "system"
|
||||
assert convo[1]["role"] == "user"
|
||||
assert any(
|
||||
c["role"] == "assistant" and "FINAL_ANSWER:" in c["content"] for c in convo
|
||||
)
|
||||
assert rec["paradigm"] == best.paradigm
|
||||
assert rec["metrics"]["num_steps"] == best.episode.num_turns()
|
||||
|
||||
|
||||
# --- build (end-to-end, injected source) ------------------------------------
|
||||
|
||||
|
||||
def test_build_sft_dataset_end_to_end(tmp_path: Path):
|
||||
def fake_source(*, max_tasks=None, configs=None):
|
||||
yield _easy_episode()
|
||||
yield _code_episode()
|
||||
|
||||
out = tmp_path / "traces.jsonl"
|
||||
stats = build_sft_dataset(str(out), max_tasks=10, source=fake_source)
|
||||
|
||||
assert stats["records_written"] == 2
|
||||
assert stats["tasks_seen"] == 2
|
||||
lines = out.read_text().strip().splitlines()
|
||||
assert len(lines) == 2
|
||||
for line in lines:
|
||||
rec = json.loads(line)
|
||||
assert rec["conversations"][0]["role"] == "system"
|
||||
|
||||
stats_file = out.with_suffix(out.suffix + ".stats.json")
|
||||
assert stats_file.exists()
|
||||
assert sum(stats["paradigm_distribution"].values()) == 2
|
||||
Reference in New Issue
Block a user