Expose stable Trace-Bench task evaluation API (e.g. for recursive OpenTrace)

tests/m1/test_opto_integration.py

@@ -0,0 +1,87 @@
+"""Tests for the generic opto/OpenTrace bundle-evaluation integration."""
+from __future__ import annotations
+
+import pytest
+
+from trace_bench.integrations.opto import BundleEval, evaluate_bundle
+from trace_bench.registry import normalize_task_id
+
+
+class _Node:
+    """Node-like param: _extract_response reads .data (production shape)."""
+    def __init__(self, data="GOOD"): self.data = data
+
+
+class _Guide:
+    # Rewards the known-good response; mirrors a real guide scoring a fixed
+    # artifact response against each example.
+    def __call__(self, x, response, info):
+        return (0.5 if response == "GOOD" else 0.0), f"resp={response}"
+
+    def get_score_dict(self, x, response, info):
+        return {"accuracy": 1.0 if response == "GOOD" else 0.0, "cost": 0.1}
+
+
+def _bundle(**over):
+    base = {"param": _Node("GOOD"), "guide": _Guide(),
+            "train_dataset": {"inputs": ["a", "b"], "infos": [{}, {}]}}
+    base.update(over)
+    return base
+
+
+def test_evaluate_bundle_returns_mean_reward_and_score_dicts():
+    mean_reward, evals = evaluate_bundle(_bundle(), max_examples=2)
+    assert mean_reward == 0.5
+    assert [type(e) for e in evals] == [BundleEval, BundleEval]
+    assert evals[0].score_dict == {"accuracy": 1.0, "cost": 0.1}
+
+
+def test_evaluate_bundle_prefers_validation_split_and_caps_examples():
+    bundle = _bundle(validate_dataset={"inputs": ["a"], "infos": [{}]})
+    mean_reward, evals = evaluate_bundle(bundle, max_examples=5)
+    assert len(evals) == 1 and mean_reward == 0.5      # validate split used, capped
+
+
+def test_evaluate_bundle_unwraps_trainable_param():
+    mean_reward, evals = evaluate_bundle(_bundle(param=_Node("GOOD")), max_examples=1)
+    assert mean_reward == 0.5 and evals[0].feedback == "resp=GOOD"
+
+
+def test_evaluate_bundle_raises_on_empty_dataset():
+    with pytest.raises(ValueError, match="no evaluable examples"):
+        evaluate_bundle(_bundle(train_dataset={"inputs": [], "infos": []}))
+
+
+def test_normalize_task_id_is_public_and_stable():
+    assert normalize_task_id("example:foo") == "trace_examples:foo"
+    assert normalize_task_id("hf:GSM8K") == "hf:GSM8K"
+    assert normalize_task_id("online_bin_packing") == "llm4ad:online_bin_packing"
+
+
+def test_evaluate_bundle_strict_score_dict_reraises():
+    class _BadDict(_Guide):
+        def get_score_dict(self, x, response, info):
+            raise RuntimeError("scorer unavailable")
+
+    bundle = _bundle(guide=_BadDict())
+    # default: swallow -> score_dict None, reward still present
+    mean_reward, evals = evaluate_bundle(bundle, max_examples=1)
+    assert evals[0].score_dict is None and mean_reward == 0.5
+    # strict: re-raise so a silent None can't masquerade as a real vector
+    with pytest.raises(RuntimeError, match="scorer unavailable"):
+        evaluate_bundle(bundle, max_examples=1, strict_score_dict=True)
+
+
+def test_evaluate_bundle_mean_reward_matches_runner_mean():
+    # parity guard: aggregate reward equals the runner's own mean over the split
+    from trace_bench.runner import _score_dataset, _evaluation_dataset
+    bundle = _bundle()
+    name, dataset = _evaluation_dataset(bundle)
+    runner_summary = _score_dataset(bundle, dataset, name)
+    mean_reward, _ = evaluate_bundle(bundle, max_examples=len(dataset["inputs"]))
+    assert mean_reward == pytest.approx(runner_summary["score"])
+
+
+def test_integrations_package_exports_public_api():
+    import trace_bench.integrations as integ
+    assert set(integ.__all__) == {"BundleEval", "evaluate_bundle", "load_and_evaluate"}

trace_bench/integrations/__init__.py

@@ -1,4 +1,6 @@
 from __future__ import annotations
 
-__all__ = []
+# Public integration surface for external OpenTrace / opto consumers.
+from .opto import BundleEval, evaluate_bundle, load_and_evaluate
 
+__all__ = ["BundleEval", "evaluate_bundle", "load_and_evaluate"]

trace_bench/integrations/opto.py

@@ -0,0 +1,112 @@
+"""Bundle evaluation helper for external OpenTrace / opto ``recursive_opt`` consumers.
+
+WHY THIS EXISTS
+---------------
+NewTrace's ``opto.features.recursive_opt`` adapter (``TraceBenchTaskAdapter``)
+needs to evaluate a loaded Trace-Bench bundle on its own dataset and read back
+BOTH the scalar reward and the optional multi-objective ``score_dict`` for its
+recursive scoring transforms (relative deltas, clipping, objective
+scalarization). The benchmark repo is the right home for "load a bundle and run
+it on its data"; the consumer keeps only the recursion-specific transforms.
+
+This module is therefore a thin, stable, PUBLIC API surface for that external
+consumer. It is intentionally not wired into ``runner.py``'s job pipeline (the
+runner has its own batched/staged evaluation with stub-LLM handling). To avoid
+the semantic drift the reviewer rightly flagged, the per-example evaluation here
+REUSES the runner's own ``_extract_response`` and ``_evaluation_dataset`` rather
+than re-deriving them — so "what counts as the response / which split is used"
+has exactly one definition in this repo.
+
+The only deliberate behavioral difference from ``runner._score_dataset`` is
+shape, by design for the consumer: this returns per-example ``BundleEval``
+records (reward + feedback + optional score_dict) instead of a single
+``{score, feedback}`` summary, because the recursive optimizer scores each
+example's objective vector. Aggregate reward matches the runner's mean.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+from trace_bench.registry import load_task_bundle, normalize_task_id
+
+__all__ = ["BundleEval", "evaluate_bundle", "load_and_evaluate"]
+
+
+@dataclass
+class BundleEval:
+    """One example's evaluation: official reward, feedback, optional score dict."""
+
+    reward: float
+    feedback: str
+    score_dict: Optional[Dict[str, float]] = None
+
+
+def evaluate_bundle(
+    bundle: Dict[str, Any],
+    *,
+    max_examples: int = 1,
+    strict_score_dict: bool = False,
+) -> Tuple[float, List[BundleEval]]:
+    """Evaluate a loaded bundle on up to ``max_examples`` of its own dataset.
+
+    Returns ``(mean_reward, per_example_evals)``. Split selection and response
+    extraction are delegated to the runner's helpers (single source of truth).
+
+    - ``strict_score_dict=False`` (default): a guide whose ``get_score_dict``
+      raises yields ``score_dict=None`` for that example (the scalar reward is
+      still recorded). This matches "score_dict is an optional enrichment".
+    - ``strict_score_dict=True``: re-raise instead of swallowing, for consumers
+      whose multi-objective scoring MUST have the dict (so a silent None can't
+      masquerade as a degenerate objective vector).
+
+    Raises ``ValueError`` when the chosen dataset is empty — a silent
+    zero-example "success" would hide a real loading/dataset failure.
+    """
+    # Reuse the runner's canonical helpers so this integration cannot drift from
+    # the repo's own response-extraction / split-selection logic. Imported lazily
+    # to avoid a circular import (runner -> integrations package -> this module).
+    from trace_bench.runner import _evaluation_dataset, _extract_response
+
+    dataset_name, dataset = _evaluation_dataset(bundle)
+    inputs = list(dataset.get("inputs") or [])
+    infos = list(dataset.get("infos") or dataset.get("info") or [])
+    limit = min(len(inputs), len(infos) or len(inputs),
+                max_examples if max_examples and max_examples > 0 else len(inputs))
+    if limit <= 0:
+        raise ValueError(f"bundle dataset {dataset_name!r} has no evaluable examples")
+
+    guide = bundle["guide"]
+    param = bundle["param"]
+    out: List[BundleEval] = []
+    for i in range(limit):
+        info = infos[i] if i < len(infos) else None
+        response = _extract_response(param, inputs[i])
+        reward, feedback = guide(inputs[i], response, info)
+        score_dict = None
+        if hasattr(guide, "get_score_dict"):
+            try:
+                score_dict = guide.get_score_dict(inputs[i], response, info)
+            except Exception:
+                if strict_score_dict:
+                    raise
+                score_dict = None
+        out.append(BundleEval(reward=float(reward), feedback=str(feedback),
+                              score_dict=score_dict))
+    return sum(ev.reward for ev in out) / len(out), out
+
+
+def load_and_evaluate(
+    task_id: str,
+    tasks_root: str,
+    *,
+    eval_kwargs: Optional[Dict[str, Any]] = None,
+    max_examples: int = 1,
+    strict_score_dict: bool = False,
+) -> Tuple[Dict[str, Any], float, List[BundleEval]]:
+    """One-call helper: load the bundle for ``task_id``, then evaluate it."""
+    bundle = load_task_bundle(normalize_task_id(task_id), tasks_root,
+                              eval_kwargs=eval_kwargs)
+    mean_reward, evals = evaluate_bundle(
+        bundle, max_examples=max_examples, strict_score_dict=strict_score_dict)
+    return bundle, mean_reward, evals

trace_bench/registry.py

@@ -401,6 +401,11 @@ def discover_tasks(tasks_root: str | Path, bench: Optional[str] = None) -> List[
     return specs
 
 
+def normalize_task_id(task_id: str) -> str:
+    """Public task-id normalization for external integrations (stable API)."""
+    return _normalize_task_id(task_id)
+
+
 def _normalize_task_id(task_id: str) -> str:
     if task_id.startswith("example:"):
         return task_id.replace("example:", "trace_examples:", 1)
@@ -516,5 +521,6 @@ def load_task_bundle(task_id: str, tasks_root: str | Path, eval_kwargs: Optional
     "expand_special_tasks",
     "load_task_bundle",
     "load_task_module",
+    "normalize_task_id",
     "priority_search_example_trainers_supported",
 ]