diff --git a/scripts/orchestrator/build_sft_data.py b/scripts/orchestrator/build_sft_data.py new file mode 100644 index 00000000..41ef73c7 --- /dev/null +++ b/scripts/orchestrator/build_sft_data.py @@ -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()) diff --git a/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3-8b-sft.toml b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3-8b-sft.toml new file mode 100644 index 00000000..57d8581b --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/configs/orchestrator-qwen3-8b-sft.toml @@ -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 diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/__init__.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/__init__.py new file mode 100644 index 00000000..5e7f5338 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/__init__.py @@ -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", +] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/adp_loader.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/adp_loader.py new file mode 100644 index 00000000..b37e5313 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/adp_loader.py @@ -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", +] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/build.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/build.py new file mode 100644 index 00000000..fe213be9 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/build.py @@ -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"] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/paradigms.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/paradigms.py new file mode 100644 index 00000000..86b3089b --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/paradigms.py @@ -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", +] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/select.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/select.py new file mode 100644 index 00000000..07097474 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/select.py @@ -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"] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/serialize.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/serialize.py new file mode 100644 index 00000000..f9bbd1ad --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/serialize.py @@ -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"] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_data/tiers.py b/src/openjarvis/learning/intelligence/orchestrator/sft_data/tiers.py new file mode 100644 index 00000000..4632a262 --- /dev/null +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_data/tiers.py @@ -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", +] diff --git a/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py b/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py index 2c2e0645..9cd39522 100644 --- a/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py +++ b/src/openjarvis/learning/intelligence/orchestrator/sft_trainer.py @@ -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: diff --git a/tests/test_orchestrator_learning/sft_data/__init__.py b/tests/test_orchestrator_learning/sft_data/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/test_orchestrator_learning/sft_data/test_sft_data.py b/tests/test_orchestrator_learning/sft_data/test_sft_data.py new file mode 100644 index 00000000..cc21df7c --- /dev/null +++ b/tests/test_orchestrator_learning/sft_data/test_sft_data.py @@ -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