From 6ac2a2639cbdade563268eddbfcb352f20648963 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 30 Sep 2026 00:28:34 +0200 Subject: [PATCH 01/20] Add prompt-hashed validation/test split to meta-eval Judge tuning needs to select configurations on one partition of the human-labeled arena and report on another. Battles are assigned by a salted hash of their first user prompt, so repeated prompts never cross partitions and changing run.seed cannot move battles between them. --- judgearena/benchmarks/meta_eval/runner.py | 12 +++++++++++- judgearena/benchmarks/meta_eval/sampling.py | 20 ++++++++++++++++++++ judgearena/config.py | 7 +++++++ judgearena/reports.py | 3 ++- tests/test_meta_eval_foundation.py | 14 ++++++++++++++ 5 files changed, 54 insertions(+), 2 deletions(-) diff --git a/judgearena/benchmarks/meta_eval/runner.py b/judgearena/benchmarks/meta_eval/runner.py index 5c052d48..4d57d12d 100644 --- a/judgearena/benchmarks/meta_eval/runner.py +++ b/judgearena/benchmarks/meta_eval/runner.py @@ -10,6 +10,7 @@ import numpy as np import pandas as pd +from judgearena.arenas_utils import extract_turn_text from judgearena.artifacts import ( prepare_run_directory, safe_filename, @@ -24,6 +25,7 @@ ) from judgearena.benchmarks.meta_eval.sampling import ( MetaEvalSamplingError, + filter_split, sample_battles_per_model, select_top_models, ) @@ -133,6 +135,13 @@ def run_meta_eval( top_models, top_pool = select_top_models( arena_battles, top_models=cfg.meta_eval.top_models ) + if cfg.meta_eval.split != "all": + top_pool = filter_split( + top_pool, + top_pool["conversation_a"].map(lambda turns: extract_turn_text(turns[0])), + split=cfg.meta_eval.split, + validation_fraction=cfg.meta_eval.validation_fraction, + ) sample = sample_battles_per_model( top_pool, top_models, @@ -155,7 +164,7 @@ def run_meta_eval( timestamp = run_started_at.strftime("%Y%m%d_%H%M%S") result_name = ( f"{safe_filename(cfg.task)}-{safe_filename(cfg.judge.model)}-" - f"{cfg.judge.swap_mode}-{timestamp}" + f"{cfg.judge.swap_mode}-{cfg.meta_eval.split}-{timestamp}" ) result_dir = prepare_run_directory(cfg, Path(cfg.run.result_folder) / result_name) sample.loc[ @@ -198,6 +207,7 @@ def run_meta_eval( top_models=top_models, n_sampled_battles=len(sample), swap_mode=cfg.judge.swap_mode, + split=cfg.meta_eval.split, metrics=metric_results, ) results = report.to_dict() diff --git a/judgearena/benchmarks/meta_eval/sampling.py b/judgearena/benchmarks/meta_eval/sampling.py index 19e36e61..efb51062 100644 --- a/judgearena/benchmarks/meta_eval/sampling.py +++ b/judgearena/benchmarks/meta_eval/sampling.py @@ -7,6 +7,8 @@ import pandas as pd +_SPLIT_SALT = "judgearena-meta-eval-split" + class MetaEvalSamplingError(ValueError): """Raised when filtering or sampling yields an unusable battle subset.""" @@ -85,6 +87,24 @@ def select_top_models( return top, df_top +def _validation_bucket(prompt: str) -> float: + digest = hashlib.sha256(f"{_SPLIT_SALT}\0{prompt}".encode()).digest() + return int.from_bytes(digest[:8], "big") / 2**64 + + +def filter_split( + df: pd.DataFrame, prompts: pd.Series, *, split: str, validation_fraction: float +) -> pd.DataFrame: + """Keep battles whose prompt hashes into ``split``. + + Keying on the prompt keeps repeated prompts on one side and makes the + partition independent of the sampling seed. + """ + is_validation = prompts.map(_validation_bucket) < validation_fraction + keep = is_validation if split == "validation" else ~is_validation + return df.loc[keep.to_numpy()].reset_index(drop=True) + + def _stable_priority(battle_id: object, *, seed: int) -> str: payload = f"{seed}\0{battle_id}".encode() return hashlib.sha256(payload).hexdigest() diff --git a/judgearena/config.py b/judgearena/config.py index 74f68fbb..f0ac302e 100644 --- a/judgearena/config.py +++ b/judgearena/config.py @@ -403,6 +403,13 @@ class MetaEvalArgs(BaseModel): languages: list[str] | None = None """Restrict arena battles to these language codes. Defaults to all languages.""" + split: Literal["all", "validation", "test"] = "all" + """Prompt-hashed arena partition to sample from. Judge tuning selects on + ``validation`` and reports on ``test``.""" + + validation_fraction: float = Field(default=0.5, gt=0, lt=1) + """Share of arena prompts assigned to the validation partition.""" + class RunArgs(BaseModel): """Run-level settings: seed, output location, caching, and logging.""" diff --git a/judgearena/reports.py b/judgearena/reports.py index 783555f5..ea03982f 100644 --- a/judgearena/reports.py +++ b/judgearena/reports.py @@ -101,12 +101,13 @@ class MetaEvalReport(Report): top_models: list[str] n_sampled_battles: int swap_mode: str + split: str metrics: dict[str, dict[str, object]] def render(self) -> None: from judgearena.benchmarks.scoring import render_metrics - print(f"\n=== Meta-eval: {self.task} ===") + print(f"\n=== Meta-eval: {self.task} ({self.split} split) ===") print(f"Arena: {self.arena} | Judge: {self.judge_model}") print( f"Models: {len(self.top_models)} | Sampled battles: {self.n_sampled_battles} | " diff --git a/tests/test_meta_eval_foundation.py b/tests/test_meta_eval_foundation.py index 4233c900..798cbbac 100644 --- a/tests/test_meta_eval_foundation.py +++ b/tests/test_meta_eval_foundation.py @@ -11,6 +11,7 @@ import judgearena.evaluate as evaluate_module from judgearena.benchmarks.meta_eval.sampling import ( MetaEvalSamplingError, + filter_split, sample_battles_per_model, select_top_models, ) @@ -55,6 +56,19 @@ def test_sampling_rejects_insufficient_quota(): sample_battles_per_model(battles, ["a", "b"], battles_per_model=2, seed=0) +def test_split_partitions_battles_by_prompt(): + battles = pd.DataFrame({"battle_id": range(200)}) + prompts = battles["battle_id"].map(lambda i: f"prompt {i % 40}") + validation, test = ( + filter_split(battles, prompts, split=split, validation_fraction=0.5) + for split in ("validation", "test") + ) + + assert 0 < len(validation) < len(battles) + assert sorted([*validation["battle_id"], *test["battle_id"]]) == list(range(200)) + assert not set(prompts[validation["battle_id"]]) & set(prompts[test["battle_id"]]) + + def test_meta_eval_protocol_requires_no_baseline_and_compatible_metrics(): definition = get_packaged_task("meta-eval-comparia").spec.model_dump() protocol = definition["protocol"] From f7c0816f9607fd965972ec5748cbe28f0105bf2d Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 30 Sep 2026 08:48:05 +0200 Subject: [PATCH 02/20] Add successive-halving judge tuning on meta-eval tasks A tune_judge section on a meta-eval task expands a grid of judge overrides and runs each configuration as its own meta-eval subprocess on the validation split. After each rung, survivors are chosen by the non-dominated sort on judge cost and human agreement from Salinas et al. (ICML 2025). The best configuration per judge model is then scored once on the held-out test split. --- judgearena/benchmarks/runner.py | 4 + judgearena/config.py | 40 ++++- judgearena/tuning/__init__.py | 1 + judgearena/tuning/runner.py | 234 ++++++++++++++++++++++++++++ judgearena/tuning/search_space.py | 68 ++++++++ judgearena/tuning/selection.py | 88 +++++++++++ tests/test_generate_and_evaluate.py | 2 +- tests/test_tuning.py | 132 ++++++++++++++++ tests/test_usage.py | 4 +- 9 files changed, 569 insertions(+), 4 deletions(-) create mode 100644 judgearena/tuning/__init__.py create mode 100644 judgearena/tuning/runner.py create mode 100644 judgearena/tuning/search_space.py create mode 100644 judgearena/tuning/selection.py create mode 100644 tests/test_tuning.py diff --git a/judgearena/benchmarks/runner.py b/judgearena/benchmarks/runner.py index ca155b25..8f4b3c51 100644 --- a/judgearena/benchmarks/runner.py +++ b/judgearena/benchmarks/runner.py @@ -20,6 +20,10 @@ def run_benchmark(cfg: RunConfig) -> object: logger.info("Using %s benchmark adapter for %s.", resolved.adapter.name, cfg.task) with track_usage() as usage_tracker: try: + if cfg.tune_judge is not None: + from judgearena.tuning.runner import run_tune_judge + + return run_tune_judge(cfg, resolved.task) return resolved.adapter.runner(cfg, resolved.task) finally: usage_tracker.render_summary() diff --git a/judgearena/config.py b/judgearena/config.py index f0ac302e..b9dc99d6 100644 --- a/judgearena/config.py +++ b/judgearena/config.py @@ -5,7 +5,7 @@ import argparse import os from pathlib import Path -from typing import Literal +from typing import Any, Literal import yaml from pydantic import BaseModel, ConfigDict, Field, model_validator @@ -411,6 +411,37 @@ class MetaEvalArgs(BaseModel): """Share of arena prompts assigned to the validation partition.""" +class TuneJudgeArgs(BaseModel): + """Successive-halving search over judge settings on meta-eval battles.""" + + model_config = ConfigDict(use_attribute_docstrings=True, extra="forbid") + + search_space: dict[str, list[Any]] = Field(min_length=1) + """Axes of the grid, keyed by dotted config path (``judge.temperature``). An + axis whose values are mappings of dotted overrides changes them together.""" + + price_per_million_tokens: dict[str, float] + """USD per million judge tokens for every searched ``judge.model``.""" + + rungs: list[int] = Field(default=[10, 30, 90], min_length=1) + """Validation battles per model at each successive-halving rung.""" + + keep_fraction: float = Field(default=1 / 3, gt=0, le=1) + """Share of configurations kept after each rung but the last.""" + + min_agreement: float | None = None + """Drop configurations whose human agreement does not exceed this value.""" + + test_battles_per_model: int | None = Field(default=None, gt=0) + """Test battles per model for the selected judges. Defaults to the last rung.""" + + @model_validator(mode="after") + def _validate_rungs(self) -> TuneJudgeArgs: + if any(n <= 0 for n in self.rungs) or self.rungs != sorted(set(self.rungs)): + raise ValueError("tune_judge.rungs must be positive and increasing.") + return self + + class RunArgs(BaseModel): """Run-level settings: seed, output location, caching, and logging.""" @@ -467,6 +498,9 @@ class RunConfig(BaseSettings): meta_eval: MetaEvalArgs | None = None """Runtime settings used only by tasks with a meta-evaluation protocol.""" + tune_judge: TuneJudgeArgs | None = None + """Tune the judge on a meta-evaluation task instead of scoring it once.""" + run: RunArgs = Field(default_factory=RunArgs) """Run-level settings (seed, output, caching, logging).""" @@ -556,6 +590,10 @@ def _validate(self) -> RunConfig: raise ValueError( "meta_eval config is only valid for meta-evaluation tasks." ) + if self.tune_judge is not None and not is_meta_eval: + raise ValueError( + "tune_judge config is only valid for meta-evaluation tasks." + ) if is_elo: if self.elo is None: self.elo = EloArgs() diff --git a/judgearena/tuning/__init__.py b/judgearena/tuning/__init__.py new file mode 100644 index 00000000..a2fa5b7b --- /dev/null +++ b/judgearena/tuning/__init__.py @@ -0,0 +1 @@ +"""Judge configuration tuning against human-labeled arena battles.""" diff --git a/judgearena/tuning/runner.py b/judgearena/tuning/runner.py new file mode 100644 index 00000000..8a065b27 --- /dev/null +++ b/judgearena/tuning/runner.py @@ -0,0 +1,234 @@ +"""Successive-halving search over judge configurations on meta-eval battles. + +Follows Salinas et al., "Tuning LLM Judge Design Decisions for 1/1000 of the +Cost" (ICML 2025). Each rung judges more battles per model with the surviving +configurations and keeps the best by a non-dominated sort on judge cost and +human agreement. The search only sees the validation split; the best +configuration per judge model is then scored once on the held-out test split. +""" + +from __future__ import annotations + +import json +import math +import subprocess +import sys +from collections.abc import Callable +from dataclasses import dataclass +from datetime import UTC, datetime +from functools import cache +from pathlib import Path +from typing import TYPE_CHECKING + +import pandas as pd +import tiktoken + +from judgearena.artifacts import prepare_run_directory, safe_filename +from judgearena.config import RunConfig, TuneJudgeArgs, dump_config +from judgearena.log import get_logger +from judgearena.tuning.search_space import apply_overrides, expand_grid +from judgearena.tuning.selection import select_survivors + +if TYPE_CHECKING: + from judgearena.tasks.schema import ResolvedTaskSpec + +logger = get_logger(__name__) + +TrialExecutor = Callable[[Path], None] +"""Run the trial config at a path, raising ``CalledProcessError`` on failure.""" + +_SUMMARY_COLUMNS = ["trial_id", "judge_model", "agreement", "cost_per_1k_battles"] + + +def run_trial_subprocess(config_path: Path) -> None: + """Run one trial in a fresh process so each judge releases its GPU memory.""" + subprocess.run( + [sys.executable, "-m", "judgearena", "--config_path", str(config_path)], + check=True, + ) + + +@cache +def _judge_token_encoding() -> tiktoken.Encoding: + return tiktoken.encoding_for_model("gpt-4o") + + +def _read_trial(trial_dir: Path, price_per_million_tokens: float) -> dict: + (result_path,) = trial_dir.glob("*/results.json") + metrics = json.loads(result_path.read_text())["metrics"] + agreement = metrics["meta_eval_agreement"]["all"] + annotations = pd.read_parquet( + result_path.parent / "annotations.parquet", + columns=["battle_id", "judge_input", "judge_completion"], + ) + texts = [*annotations["judge_input"], *annotations["judge_completion"].fillna("")] + encoding = _judge_token_encoding() + tokens = sum(len(encoding.encode(text, disallowed_special=())) for text in texts) + tokens_per_battle = tokens / annotations["battle_id"].nunique() + return { + "agreement": agreement["accuracy_attempted"], + "cohen_kappa": agreement["cohen_kappa"], + "coverage": agreement["coverage"], + "n_battles": agreement["n_attempted"], + "spearman": metrics.get("meta_eval_ranking", {}) + .get("soft", {}) + .get("spearman"), + "tokens_per_battle": tokens_per_battle, + "cost_per_1k_battles": tokens_per_battle * price_per_million_tokens / 1000, + "result_dir": str(result_path.parent), + } + + +@dataclass(frozen=True) +class _TrialRunner: + """Materialize trial configs from one base config and record their results.""" + + base: dict + overrides: dict[str, dict[str, object]] + prices: dict[str, float] + execute_trial: TrialExecutor + + def config( + self, trial_id: str, *, split: str, battles_per_model: int, folder: Path + ) -> RunConfig: + values = apply_overrides( + self.base, + { + **self.overrides[trial_id], + "meta_eval.split": split, + "meta_eval.battles_per_model": battles_per_model, + "run.result_folder": str(folder), + }, + ) + return RunConfig(**values) + + def run( + self, trial_id: str, *, split: str, battles_per_model: int, folder: Path + ) -> dict: + trial_cfg = self.config( + trial_id, split=split, battles_per_model=battles_per_model, folder=folder + ) + record = { + "trial_id": trial_id, + "judge_model": trial_cfg.judge.model, + "overrides": json.dumps(self.overrides[trial_id], sort_keys=True), + "battles_per_model": battles_per_model, + } + folder.mkdir(parents=True, exist_ok=True) + config_path = folder / "config.yaml" + dump_config(trial_cfg, config_path) + try: + self.execute_trial(config_path) + except subprocess.CalledProcessError as exc: + logger.warning("Trial %s failed: %s", folder, exc) + return {**record, "status": "failed"} + price = self.prices[trial_cfg.judge.model] + return {**record, "status": "completed", **_read_trial(folder, price)} + + +def _successive_halving( + runner: _TrialRunner, tuning: TuneJudgeArgs, tune_dir: Path +) -> tuple[pd.DataFrame, pd.DataFrame]: + """Return every rung record and the completed survivors of the last rung.""" + alive = list(runner.overrides) + records = [] + for rung, battles_per_model in enumerate(tuning.rungs): + rung_records = [ + { + "rung": rung, + **runner.run( + trial_id, + split="validation", + battles_per_model=battles_per_model, + folder=tune_dir / f"rung-{rung}" / trial_id, + ), + } + for trial_id in alive + ] + records.extend(rung_records) + completed = pd.DataFrame( + [record for record in rung_records if record["status"] == "completed"] + ) + if completed.empty: + raise RuntimeError(f"Every judge configuration failed at rung {rung}.") + is_last = rung == len(tuning.rungs) - 1 + survivors = select_survivors( + completed["agreement"].to_numpy(dtype=float), + completed["cost_per_1k_battles"].to_numpy(dtype=float), + n_keep=( + len(completed) + if is_last + else math.ceil(len(alive) * tuning.keep_fraction) + ), + min_agreement=tuning.min_agreement, + ) + if not survivors: + raise RuntimeError( + f"No configuration exceeded min_agreement at rung {rung}." + ) + alive = completed["trial_id"].iloc[survivors].tolist() + logger.info("Rung %d kept %d of %d.", rung, len(alive), len(rung_records)) + return pd.DataFrame(records), completed.iloc[survivors] + + +def run_tune_judge( + cfg: RunConfig, + task: ResolvedTaskSpec | None = None, + *, + execute_trial: TrialExecutor = run_trial_subprocess, +) -> pd.DataFrame: + """Tune judge settings on validation battles and score the picks on test.""" + tuning = cfg.tune_judge + started_at = datetime.now(UTC) + tune_dir = prepare_run_directory( + cfg, + Path(cfg.run.result_folder) + / f"tune-{safe_filename(cfg.task)}-{started_at:%Y%m%d_%H%M%S}", + ) + runner = _TrialRunner( + base=cfg.model_dump(mode="json", exclude={"tune_judge"}), + overrides={ + trial.id: trial.overrides for trial in expand_grid(tuning.search_space) + }, + prices=tuning.price_per_million_tokens, + execute_trial=execute_trial, + ) + # Build every trial before judging so invalid settings fail up front. + judge_models = { + runner.config( + trial_id, split="validation", battles_per_model=1, folder=tune_dir + ).judge.model + for trial_id in runner.overrides + } + missing_prices = sorted(judge_models - set(tuning.price_per_million_tokens)) + if missing_prices: + raise ValueError( + f"tune_judge.price_per_million_tokens is missing models: {missing_prices}" + ) + logger.info("Tuning %d judge configurations in %s", len(runner.overrides), tune_dir) + + trials, final = _successive_halving(runner, tuning, tune_dir) + trials.to_parquet(tune_dir / "trials.parquet", index=False) + picks = final.loc[final.groupby("judge_model")["agreement"].idxmax(), "trial_id"] + test_results = pd.DataFrame( + [ + runner.run( + trial_id, + split="test", + battles_per_model=tuning.test_battles_per_model or tuning.rungs[-1], + folder=tune_dir / "test" / trial_id, + ) + for trial_id in picks + ] + ) + test_results.to_parquet(tune_dir / "test_results.parquet", index=False) + + print(f"\n=== Judge tuning: {cfg.task} ===") + print("Validation (last rung):") + ranked = final.sort_values("agreement", ascending=False) + print(ranked[_SUMMARY_COLUMNS].to_string(index=False)) + print("Test:") + summary = test_results.reindex(columns=[*_SUMMARY_COLUMNS, "status"]) + print(summary.to_string(index=False)) + print(f"Results: {tune_dir}") + return test_results diff --git a/judgearena/tuning/search_space.py b/judgearena/tuning/search_space.py new file mode 100644 index 00000000..7485f281 --- /dev/null +++ b/judgearena/tuning/search_space.py @@ -0,0 +1,68 @@ +"""Expand a judge-tuning search space into concrete trial overrides.""" + +from __future__ import annotations + +import copy +import hashlib +import itertools +import json +from collections.abc import Mapping, Sequence +from dataclasses import dataclass + +TUNER_OWNED_KEYS = frozenset( + {"task", "meta_eval.battles_per_model", "meta_eval.split", "run.result_folder"} +) + + +@dataclass(frozen=True) +class Trial: + """One judge configuration, given as dotted ``RunConfig`` overrides.""" + + overrides: dict[str, object] + + @property + def id(self) -> str: + payload = json.dumps(self.overrides, sort_keys=True) + return hashlib.sha256(payload.encode()).hexdigest()[:12] + + +def expand_grid(search_space: Mapping[str, Sequence[object]]) -> list[Trial]: + """Return the cartesian product of the search-space axes. + + An axis name is the dotted ``RunConfig`` path its values override. An axis + whose values are mappings is grouped instead: each value is a set of dotted + overrides that change together, and the axis name is only a label. + """ + axes = [ + [ + dict(value) if isinstance(value, Mapping) else {name: value} + for value in values + ] + for name, values in search_space.items() + ] + keys = {key for choices in axes for choice in choices for key in choice} + owned = sorted( + key for key in keys if key in TUNER_OWNED_KEYS or key.startswith("tune_judge") + ) + if owned: + raise ValueError(f"search_space must not override tuner-owned keys: {owned}") + return [ + Trial({key: value for choice in combo for key, value in choice.items()}) + for combo in itertools.product(*axes) + ] + + +def apply_overrides( + base: Mapping[str, object], overrides: Mapping[str, object] +) -> dict[str, object]: + """Return a copy of ``base`` with dotted-path ``overrides`` assigned.""" + values = copy.deepcopy(dict(base)) + for path, value in overrides.items(): + *parents, leaf = path.split(".") + node = values + for key in parents: + if node.get(key) is None: + node[key] = {} + node = node[key] + node[leaf] = value + return values diff --git a/judgearena/tuning/selection.py b/judgearena/tuning/selection.py new file mode 100644 index 00000000..65cd354d --- /dev/null +++ b/judgearena/tuning/selection.py @@ -0,0 +1,88 @@ +"""Multi-objective survivor selection used by successive halving. + +``pareto_efficiency_mask``, ``epsilon_net_indices`` and ``argsort_nondominated`` +are adapted from judgetuning (github.com/geoalgo/judgetuning, Apache-2.0), +which took them from TSBench by Oliver Borchert. +""" + +from __future__ import annotations + +import numpy as np +import scipy.stats as st + + +def pareto_efficiency_mask(X: np.ndarray) -> np.ndarray: + """Mark rows of ``X`` not dominated by another row, all columns minimized.""" + mask = np.ones(X.shape[0], dtype=bool) + for i, point in enumerate(X): + if mask[i]: + dominated = np.all(point <= X[mask], axis=1) & np.any( + point < X[mask], axis=1 + ) + mask[mask] = ~dominated + return mask + + +def epsilon_net_indices(X: np.ndarray, dim: int) -> np.ndarray: + """Order rows so each next row is farthest from those already chosen. + + The seed is the row minimizing column ``dim`` (Clarkson, 2005, p. 17). + """ + remaining = set(range(X.shape[0])) + order = [int(np.argmin(X[:, dim]))] + remaining.remove(order[0]) + while remaining: + candidates = list(remaining) + diff = X[candidates][:, None, :] - X[order][None, :, :] + choice = candidates[np.linalg.norm(diff, axis=-1).min(-1).argmax()] + order.append(choice) + remaining.remove(choice) + return np.array(order) + + +def argsort_nondominated( + X: np.ndarray, *, minimize: list[bool], dim: int, max_items: int | None = None +) -> list[int]: + """Sort rows by successive Pareto fronts, each ordered by an epsilon-net. + + Columns are quantile-normalized, so objectives on different scales weigh + equally. + """ + if X.shape[0] == 1: + return [0] + X = np.where(minimize, X, -X) + X = (st.rankdata(X, axis=0) - 1) / (X.shape[0] - 1) + remaining = np.arange(X.shape[0]) + indices: list[int] = [] + while remaining.size and (max_items is None or len(indices) < max_items): + pareto_mask = pareto_efficiency_mask(X[remaining]) + front = remaining[pareto_mask] + indices.extend(front[epsilon_net_indices(X[front], dim=dim)].tolist()) + remaining = remaining[~pareto_mask] + return indices[:max_items] + + +def select_survivors( + agreement: np.ndarray, + cost: np.ndarray, + *, + n_keep: int, + min_agreement: float | None = None, +) -> list[int]: + """Return indices of the ``n_keep`` best trials by (agreement ↑, cost ↓). + + Trials at or below ``min_agreement`` are dropped first; the paper used a + threshold just above a length-based baseline. + """ + candidates = np.arange(len(agreement)) + if min_agreement is not None: + candidates = candidates[agreement > min_agreement] + if not candidates.size: + return [] + order = argsort_nondominated( + np.stack([cost[candidates], agreement[candidates]], axis=1), + minimize=[True, False], + dim=1, + max_items=n_keep, + ) + return candidates[order].tolist() diff --git a/tests/test_generate_and_evaluate.py b/tests/test_generate_and_evaluate.py index b790fe67..9988ec92 100644 --- a/tests/test_generate_and_evaluate.py +++ b/tests/test_generate_and_evaluate.py @@ -218,7 +218,7 @@ def test_benchmark_dispatch_passes_the_resolved_task(monkeypatch): ) monkeypatch.setattr(benchmark_registry, "benchmark_adapters", lambda: (pairwise,)) monkeypatch.setattr(benchmark_registry, "get_packaged_task", lambda _task: resolved) - cfg = SimpleNamespace(task="yaml-task") + cfg = SimpleNamespace(task="yaml-task", tune_judge=None) result = benchmark_runner.run_benchmark(cfg) diff --git a/tests/test_tuning.py b/tests/test_tuning.py new file mode 100644 index 00000000..89eb96e8 --- /dev/null +++ b/tests/test_tuning.py @@ -0,0 +1,132 @@ +"""Tests for judge configuration tuning.""" + +from __future__ import annotations + +import json +import subprocess + +import numpy as np +import pandas as pd +import pytest +from pydantic import ValidationError + +from judgearena.config import RunConfig, load_config +from judgearena.tuning.runner import run_tune_judge +from judgearena.tuning.search_space import apply_overrides, expand_grid +from judgearena.tuning.selection import select_survivors + + +def test_expand_grid_combines_plain_and_grouped_axes(): + trials = expand_grid( + { + "judge.temperature": [0.0, 1.0], + "prompt": [ + {"judge.prompt_preset": "meta-eval-pair-score"}, + {"judge.prompt_preset": "alpaca-eval", "judge.top_logprobs": 20}, + ], + } + ) + + assert len({trial.id for trial in trials}) == 4 + assert { + "judge.temperature": 1.0, + "judge.prompt_preset": "alpaca-eval", + "judge.top_logprobs": 20, + } in [trial.overrides for trial in trials] + assert apply_overrides({"judge": {"model": "m"}}, trials[0].overrides)["judge"] == { + "model": "m", + "temperature": 0.0, + "prompt_preset": "meta-eval-pair-score", + } + with pytest.raises(ValueError, match="tuner-owned"): + expand_grid({"meta_eval.split": ["test"]}) + + +def test_select_survivors_sorts_pareto_fronts_by_agreement_first(): + agreement = np.array([0.50, 0.60, 0.55, 0.40]) + cost = np.array([1.0, 2.0, 3.0, 0.5]) + + assert select_survivors(agreement, cost, n_keep=2, min_agreement=0.45) == [1, 0] + assert select_survivors(agreement, cost, n_keep=3, min_agreement=0.45) == [1, 0, 2] + + +def _fake_trial(config_path): + cfg = load_config(config_path) + if cfg.judge.temperature == 1.0 and cfg.judge.swap_mode == "both": + raise subprocess.CalledProcessError(1, "judgearena") + run_dir = config_path.parent / "run" + run_dir.mkdir() + agreement = 0.6 - 0.1 * cfg.judge.temperature + metrics = { + "meta_eval_agreement": { + "all": { + "accuracy_attempted": agreement, + "cohen_kappa": agreement, + "coverage": 1.0, + "n_attempted": cfg.meta_eval.battles_per_model, + } + } + } + (run_dir / "results.json").write_text(json.dumps({"metrics": metrics})) + completion = "x " * (1 + 10 * (cfg.judge.swap_mode == "both")) + pd.DataFrame( + { + "battle_id": ["b1"], + "judge_input": ["prompt"], + "judge_completion": [completion], + } + ).to_parquet(run_dir / "annotations.parquet") + + +def _tuning_config(tmp_path, **tune_judge): + return RunConfig( + task="meta-eval-comparia", + judge={"model": "Dummy/judge"}, + run={"result_folder": str(tmp_path), "no_log_file": True}, + tune_judge={ + "search_space": { + "judge.temperature": [0.0, 0.5, 1.0], + "judge.swap_mode": ["fixed", "both"], + }, + "rungs": [2, 4], + "price_per_million_tokens": {"Dummy/judge": 1.0}, + **tune_judge, + }, + ) + + +def test_run_tune_judge_halves_on_validation_and_scores_pick_on_test(tmp_path): + test_results = run_tune_judge(_tuning_config(tmp_path), execute_trial=_fake_trial) + + (tune_dir,) = tmp_path.glob("tune-*") + trials = pd.read_parquet(tune_dir / "trials.parquet") + assert trials.groupby("rung").size().tolist() == [6, 2] + assert (trials["status"] == "failed").sum() == 1 + (pick,) = test_results["trial_id"] + pick_cfg = load_config(tune_dir / "test" / pick / "config.yaml") + assert pick_cfg.meta_eval.split == "test" + assert (pick_cfg.judge.temperature, pick_cfg.judge.swap_mode) == (0.0, "fixed") + for rung in trials.itertuples(): + rung_cfg = load_config( + tune_dir / f"rung-{rung.rung}" / rung.trial_id / "config.yaml" + ) + assert rung_cfg.meta_eval.split == "validation" + assert rung_cfg.tune_judge is None + + +def test_tune_judge_requires_prices_and_meta_eval_task(tmp_path): + with pytest.raises(ValueError, match="missing models"): + run_tune_judge( + _tuning_config(tmp_path, price_per_million_tokens={}), + execute_trial=_fake_trial, + ) + with pytest.raises(ValidationError, match="only valid for meta-evaluation"): + RunConfig( + task="alpaca-eval", + model={"name": "Dummy/a"}, + judge={"model": "Dummy/judge"}, + tune_judge={ + "search_space": {"judge.temperature": [0.0]}, + "price_per_million_tokens": {}, + }, + ) diff --git a/tests/test_usage.py b/tests/test_usage.py index 51e7a0a7..12e84a9f 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -77,7 +77,7 @@ def fake_runner(cfg, task): adapter=SimpleNamespace(name="test", runner=fake_runner), task=None ) monkeypatch.setattr(benchmark_runner, "resolve_benchmark", lambda task: resolved) - path = benchmark_runner.run_benchmark(SimpleNamespace(task="test")) + path = benchmark_runner.run_benchmark(SimpleNamespace(task="test", tune_judge=None)) usage = json.loads(path.read_text())["usage"] assert usage["total"]["requests"] == 2 assert usage["total"]["cost_usd"] == pytest.approx(0.002) @@ -99,7 +99,7 @@ def fake_runner(cfg, task): ) monkeypatch.setattr(benchmark_runner, "resolve_benchmark", lambda task: resolved) with pytest.raises(ValueError, match="backend failed"): - benchmark_runner.run_benchmark(SimpleNamespace(task="test")) + benchmark_runner.run_benchmark(SimpleNamespace(task="test", tune_judge=None)) assert current_run_usage() is None assert "No successful model responses" in capsys.readouterr().out From 548148788f7c8de4342aa8ee7d92514294a6978f Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 30 Sep 2026 08:49:14 +0200 Subject: [PATCH 03/20] Document judge tuning and add an example search space --- README.md | 22 ++++++++++++++++++++++ configs/tune_judge.yaml | 30 ++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+) create mode 100644 configs/tune_judge.yaml diff --git a/README.md b/README.md index 2e73395d..ece9e8e5 100644 --- a/README.md +++ b/README.md @@ -332,6 +332,28 @@ judgearena \ Runs save the selected sample, judge evidence, metric battles, configuration, and results under `--run.result_folder`. +`--meta_eval.split validation|test` restricts sampling to one half of the arena prompts (`--meta_eval.validation_fraction`, default 0.5). The split hashes each battle's prompt, so it does not depend on `--run.seed`. + +### Tuning a judge + +A `tune_judge` section turns a meta-eval task into a judge configuration search, following [Salinas et al., 2025](https://arxiv.org/abs/2501.17178): + +- Each entry in `search_space` is a dotted config path with candidate values. An axis whose values are mappings changes several settings together, e.g. `alpaca-eval` together with `judge.top_logprobs`. +- Every configuration runs as its own meta-eval subprocess on the validation split. +- Successive halving: each rung in `rungs` judges more battles per model. After each rung, the configurations are ranked by a non-dominated sort on human agreement and judge cost, and the best `keep_fraction` survive. +- Cost is judge tokens (input plus output, counted with tiktoken) times `price_per_million_tokens[judge.model]`. +- After the last rung, the best configuration for each judge model is scored once on the held-out test split. + +```bash +judgearena --config_path configs/tune_judge.yaml +``` + +Outputs go to `tune--/` under `--run.result_folder`: + +- `trials.parquet`: one row per trial and rung. +- `test_results.parquet`: test-split results for the selected configurations. +- One meta-eval run directory per trial. + ## 📈 Estimating ELO Ratings JudgeArena can estimate the ELO rating of a model by running it against opponents sampled from a human preference arena (`LMArena-100k`, `LMArena-140k`, or `ComparIA`). diff --git a/configs/tune_judge.yaml b/configs/tune_judge.yaml new file mode 100644 index 00000000..25ab4b38 --- /dev/null +++ b/configs/tune_judge.yaml @@ -0,0 +1,30 @@ +# Example judge-tuning run (successive halving, Salinas et al., ICML 2025). +# Run with: judgearena --config_path configs/tune_judge.yaml +# Each configuration runs as its own meta-eval on the validation split; the +# best configuration per judge model is then scored on the test split. +task: meta-eval-lmarena-140k-en +judge: + model: VLLM/Qwen/Qwen2.5-7B-Instruct + max_out_tokens: 2048 + max_model_len: 16384 +meta_eval: + top_models: 20 +tune_judge: + search_space: + judge.model: + - VLLM/Qwen/Qwen2.5-7B-Instruct + - VLLM/google/gemma-2-9b-it + # Grouped axis: settings in one entry change together. + prompt: + - {judge.prompt_preset: meta-eval-pair-score} + - {judge.prompt_preset: arena-hard} + - {judge.prompt_preset: default_with_explanation} + - {judge.prompt_preset: alpaca-eval, judge.top_logprobs: 20, judge.max_out_tokens: 1} + judge.temperature: [0.0, 0.1, 1.0] + judge.swap_mode: [fixed, both] + rungs: [10, 30, 90] + # USD per million tokens, from judgetuning's llm_specs.py (measured vLLM + # throughput). Every searched judge.model needs a price. + price_per_million_tokens: + VLLM/Qwen/Qwen2.5-7B-Instruct: 0.12 + VLLM/google/gemma-2-9b-it: 0.14 From cb18726fe186473285e271e8ed72651088317cc3 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 2 Oct 2026 15:08:35 +0200 Subject: [PATCH 04/20] Split judge tuning into tune-judge tasks searched with neps tune-judge-* tasks reuse the meta-eval arena definitions and run each trial as the matching meta-eval task, so both share one judgement cache. The custom successive halving and Pareto selection give way to neps (priorband by default, with the base config as prior; hyperband and mo_hyperband), run via AskAndTell with validation battles per model as the fidelity. --- README.md | 17 +- configs/tune_judge.yaml | 20 +- judgearena/benchmarks/registry.py | 2 + judgearena/benchmarks/runner.py | 4 - judgearena/config.py | 58 ++++-- .../definitions/meta_eval/_comparia.yaml | 19 ++ .../definitions/meta_eval/_lmarena-100k.yaml | 19 ++ .../definitions/meta_eval/_lmarena-140k.yaml | 19 ++ .../meta_eval/meta-eval-comparia.yaml | 20 +- .../meta_eval/meta-eval-lmarena-100k.yaml | 20 +- .../meta_eval/meta-eval-lmarena-140k.yaml | 20 +- .../meta_eval/tune-judge-comparia.yaml | 7 + .../meta_eval/tune-judge-lmarena-100k.yaml | 7 + .../meta_eval/tune-judge-lmarena-140k.yaml | 7 + judgearena/tasks/schema/__init__.py | 3 +- judgearena/tasks/schema/meta_eval.py | 8 + judgearena/tasks/schema/task.py | 5 +- judgearena/tuning/runner.py | 176 ++++++++++-------- judgearena/tuning/search_space.py | 122 ++++++++---- judgearena/tuning/selection.py | 88 --------- pyproject.toml | 3 + tests/test_mt_bench_downloads.py | 3 + tests/test_tuning.py | 164 ++++++++-------- 23 files changed, 443 insertions(+), 368 deletions(-) create mode 100644 judgearena/tasks/definitions/meta_eval/_comparia.yaml create mode 100644 judgearena/tasks/definitions/meta_eval/_lmarena-100k.yaml create mode 100644 judgearena/tasks/definitions/meta_eval/_lmarena-140k.yaml create mode 100644 judgearena/tasks/definitions/meta_eval/tune-judge-comparia.yaml create mode 100644 judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-100k.yaml create mode 100644 judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-140k.yaml delete mode 100644 judgearena/tuning/selection.py diff --git a/README.md b/README.md index ece9e8e5..03c0332d 100644 --- a/README.md +++ b/README.md @@ -336,13 +336,16 @@ Runs save the selected sample, judge evidence, metric battles, configuration, an ### Tuning a judge -A `tune_judge` section turns a meta-eval task into a judge configuration search, following [Salinas et al., 2025](https://arxiv.org/abs/2501.17178): - -- Each entry in `search_space` is a dotted config path with candidate values. An axis whose values are mappings changes several settings together, e.g. `alpaca-eval` together with `judge.top_logprobs`. -- Every configuration runs as its own meta-eval subprocess on the validation split. -- Successive halving: each rung in `rungs` judges more battles per model. After each rung, the configurations are ranked by a non-dominated sort on human agreement and judge cost, and the best `keep_fraction` survive. +The `tune-judge-*` tasks search judge configurations against the same human battles as their `meta-eval-*` counterparts, following [Salinas et al., 2025](https://arxiv.org/abs/2501.17178). They need the `tune` extra (`pip install 'judgearena[tune]'`), which installs [neps](https://github.com/automl/neps). + +- Each entry in `search_space` is a dotted config path with a list of choices or a `{lower, upper}` range. Choices that are mappings change several settings together, e.g. `alpaca-eval` together with `judge.top_logprobs`. +- A neps optimizer proposes each trial. The fidelity is validation battles per model, between `min_battles_per_model` and `max_battles_per_model`. + - `algorithm: priorband` (the default) uses the base config as a prior, so a known-good judge setup speeds up the search. + - `hyperband` samples uniformly. + - `mo_hyperband` also minimizes judge cost. +- Every trial runs as a meta-eval subprocess of the matching `meta-eval-*` task on the validation split, so tuning and meta-eval share the judgement cache. - Cost is judge tokens (input plus output, counted with tiktoken) times `price_per_million_tokens[judge.model]`. -- After the last rung, the best configuration for each judge model is scored once on the held-out test split. +- At the end, the best configuration for each judge model at the highest fidelity is scored once on the held-out test split. ```bash judgearena --config_path configs/tune_judge.yaml @@ -350,7 +353,7 @@ judgearena --config_path configs/tune_judge.yaml Outputs go to `tune--/` under `--run.result_folder`: -- `trials.parquet`: one row per trial and rung. +- `trials.parquet`: one row per neps trial. - `test_results.parquet`: test-split results for the selected configurations. - One meta-eval run directory per trial. diff --git a/configs/tune_judge.yaml b/configs/tune_judge.yaml index 25ab4b38..63341bc3 100644 --- a/configs/tune_judge.yaml +++ b/configs/tune_judge.yaml @@ -1,10 +1,14 @@ -# Example judge-tuning run (successive halving, Salinas et al., ICML 2025). -# Run with: judgearena --config_path configs/tune_judge.yaml -# Each configuration runs as its own meta-eval on the validation split; the -# best configuration per judge model is then scored on the test split. -task: meta-eval-lmarena-140k-en +# Example judge-tuning run (Salinas et al., ICML 2025) with neps PriorBand. +# Run with: pip install 'judgearena[tune]' +# judgearena --config_path configs/tune_judge.yaml +# Each configuration runs as a meta-eval-lmarena-140k-en trial on the +# validation split; the settings below (Qwen2.5-7B, meta-eval-pair-score, +# temperature 0, fixed swap) are the prior. The best configuration per judge +# model is then scored on the test split. +task: tune-judge-lmarena-140k-en judge: model: VLLM/Qwen/Qwen2.5-7B-Instruct + temperature: 0.0 max_out_tokens: 2048 max_model_len: 16384 meta_eval: @@ -22,7 +26,11 @@ tune_judge: - {judge.prompt_preset: alpaca-eval, judge.top_logprobs: 20, judge.max_out_tokens: 1} judge.temperature: [0.0, 0.1, 1.0] judge.swap_mode: [fixed, both] - rungs: [10, 30, 90] + algorithm: priorband + # One hyperband iteration with eta 3 over 10-90 battles is 20 trials. + max_evaluations: 40 + min_battles_per_model: 10 + max_battles_per_model: 90 # USD per million tokens, from judgetuning's llm_specs.py (measured vLLM # throughput). Every searched judge.model needs a price. price_per_million_tokens: diff --git a/judgearena/benchmarks/registry.py b/judgearena/benchmarks/registry.py index 7bab0b06..ff88ab2f 100644 --- a/judgearena/benchmarks/registry.py +++ b/judgearena/benchmarks/registry.py @@ -49,12 +49,14 @@ def benchmark_adapters() -> tuple[BenchmarkAdapter, ...]: from judgearena.benchmarks.mt_bench.runner import run_mt_bench_benchmark from judgearena.benchmarks.mt_bench_101.runner import run_mt_bench_101_benchmark from judgearena.benchmarks.pairwise.runner import run_pairwise + from judgearena.tuning.runner import run_tune_judge return ( BenchmarkAdapter("elo", frozenset(), run_elo), BenchmarkAdapter("meta_eval", frozenset(), run_meta_eval), BenchmarkAdapter("mt_bench", frozenset(), run_mt_bench_benchmark), BenchmarkAdapter("mt_bench_101", frozenset(), run_mt_bench_101_benchmark), + BenchmarkAdapter("tune_judge", frozenset(), run_tune_judge), BenchmarkAdapter("pairwise", None, run_pairwise), ) diff --git a/judgearena/benchmarks/runner.py b/judgearena/benchmarks/runner.py index 8f4b3c51..ca155b25 100644 --- a/judgearena/benchmarks/runner.py +++ b/judgearena/benchmarks/runner.py @@ -20,10 +20,6 @@ def run_benchmark(cfg: RunConfig) -> object: logger.info("Using %s benchmark adapter for %s.", resolved.adapter.name, cfg.task) with track_usage() as usage_tracker: try: - if cfg.tune_judge is not None: - from judgearena.tuning.runner import run_tune_judge - - return run_tune_judge(cfg, resolved.task) return resolved.adapter.runner(cfg, resolved.task) finally: usage_tracker.render_summary() diff --git a/judgearena/config.py b/judgearena/config.py index b9dc99d6..6c5f88c1 100644 --- a/judgearena/config.py +++ b/judgearena/config.py @@ -26,6 +26,7 @@ EloScoringSpec, MetaEvalProtocol, MTBenchProtocol, + TuneJudgeProtocol, ) # Set by build_run_config() for the duration of RunConfig() construction. @@ -411,34 +412,57 @@ class MetaEvalArgs(BaseModel): """Share of arena prompts assigned to the validation partition.""" +class FloatRange(BaseModel): + """Continuous search-space axis.""" + + model_config = ConfigDict(extra="forbid") + + lower: float + upper: float + log: bool = False + + class TuneJudgeArgs(BaseModel): - """Successive-halving search over judge settings on meta-eval battles.""" + """Multi-fidelity neps search over judge settings on meta-eval battles.""" model_config = ConfigDict(use_attribute_docstrings=True, extra="forbid") - search_space: dict[str, list[Any]] = Field(min_length=1) - """Axes of the grid, keyed by dotted config path (``judge.temperature``). An - axis whose values are mappings of dotted overrides changes them together.""" + search_space: dict[str, list[Any] | FloatRange] = Field(min_length=1) + """Axes keyed by dotted config path (``judge.temperature``), each a list of + choices or a ``{lower, upper}`` range. Choices that are mappings of dotted + overrides change those settings together. The base config is the prior.""" price_per_million_tokens: dict[str, float] """USD per million judge tokens for every searched ``judge.model``.""" - rungs: list[int] = Field(default=[10, 30, 90], min_length=1) - """Validation battles per model at each successive-halving rung.""" + algorithm: Literal["priorband", "hyperband", "mo_hyperband"] = "priorband" + """neps optimizer; ``mo_hyperband`` also minimizes judge cost.""" + + max_evaluations: int = Field(gt=0) + """Validation trials to run, counting each promotion to more battles.""" - keep_fraction: float = Field(default=1 / 3, gt=0, le=1) - """Share of configurations kept after each rung but the last.""" + min_battles_per_model: int = Field(default=10, gt=0) + """Lowest fidelity, in validation battles per model.""" - min_agreement: float | None = None - """Drop configurations whose human agreement does not exceed this value.""" + max_battles_per_model: int = 90 + """Highest fidelity, in validation battles per model.""" + + eta: int = Field(default=3, ge=2) + """Hyperband reduction factor between fidelities.""" + + prior_confidence: Literal["low", "medium", "high"] = "medium" + """How strongly priorband trusts the base config.""" test_battles_per_model: int | None = Field(default=None, gt=0) - """Test battles per model for the selected judges. Defaults to the last rung.""" + """Test battles per model for the selected judges. Defaults to the highest + fidelity.""" @model_validator(mode="after") - def _validate_rungs(self) -> TuneJudgeArgs: - if any(n <= 0 for n in self.rungs) or self.rungs != sorted(set(self.rungs)): - raise ValueError("tune_judge.rungs must be positive and increasing.") + def _validate_fidelity(self) -> TuneJudgeArgs: + if self.min_battles_per_model >= self.max_battles_per_model: + raise ValueError( + "tune_judge.min_battles_per_model must be below max_battles_per_model." + ) return self @@ -590,9 +614,11 @@ def _validate(self) -> RunConfig: raise ValueError( "meta_eval config is only valid for meta-evaluation tasks." ) - if self.tune_judge is not None and not is_meta_eval: + is_tune_judge = isinstance(protocol, TuneJudgeProtocol) + if (self.tune_judge is not None) != is_tune_judge: raise ValueError( - "tune_judge config is only valid for meta-evaluation tasks." + "tune_judge config is required by, and only valid for, " + "tune-judge tasks." ) if is_elo: if self.elo is None: diff --git a/judgearena/tasks/definitions/meta_eval/_comparia.yaml b/judgearena/tasks/definitions/meta_eval/_comparia.yaml new file mode 100644 index 00000000..6947fd5e --- /dev/null +++ b/judgearena/tasks/definitions/meta_eval/_comparia.yaml @@ -0,0 +1,19 @@ +extends: _base.yaml + +variants: + selector: language + values: [fr, en] + +dataset: + sources: + comparia: + type: huggingface_dataset + repo_id: ministere-culture/comparia-votes + revision: "7a40bce496c1f2aa3be4001da85a49cb4743042b" + allow_patterns: [votes.parquet] + +protocol: + arena: ComparIA + +metadata: + reference_implementation: https://huggingface.co/datasets/ministere-culture/comparia-votes diff --git a/judgearena/tasks/definitions/meta_eval/_lmarena-100k.yaml b/judgearena/tasks/definitions/meta_eval/_lmarena-100k.yaml new file mode 100644 index 00000000..dafbf9c3 --- /dev/null +++ b/judgearena/tasks/definitions/meta_eval/_lmarena-100k.yaml @@ -0,0 +1,19 @@ +extends: _base.yaml + +variants: + selector: language + values: [en, pl, ru, zh, de, ja, es, fr, ko, pt, fa, it, tr, vi, cs, ar, uk] + +dataset: + sources: + lmarena_100k: + type: huggingface_dataset + repo_id: lmarena-ai/arena-human-preference-100k + revision: "72e85b3ddc9c81bf7b659d6b03d4126dfd8fb34a" + allow_patterns: [data/*.parquet] + +protocol: + arena: LMArena-100k + +metadata: + reference_implementation: https://huggingface.co/datasets/lmarena-ai/arena-human-preference-100k diff --git a/judgearena/tasks/definitions/meta_eval/_lmarena-140k.yaml b/judgearena/tasks/definitions/meta_eval/_lmarena-140k.yaml new file mode 100644 index 00000000..36ea359e --- /dev/null +++ b/judgearena/tasks/definitions/meta_eval/_lmarena-140k.yaml @@ -0,0 +1,19 @@ +extends: _base.yaml + +variants: + selector: language + values: [en, pl, ru, zh, de, ja, es, fr, ko, pt, fa, it, tr, vi, cs, ar, uk] + +dataset: + sources: + lmarena_140k: + type: huggingface_dataset + repo_id: lmarena-ai/arena-human-preference-140k + revision: "6322995ab34d7c2693e3f47dd13fa5caa0789a74" + allow_patterns: [data/*.parquet] + +protocol: + arena: LMArena-140k + +metadata: + reference_implementation: https://huggingface.co/datasets/lmarena-ai/arena-human-preference-140k diff --git a/judgearena/tasks/definitions/meta_eval/meta-eval-comparia.yaml b/judgearena/tasks/definitions/meta_eval/meta-eval-comparia.yaml index f9a34f46..d0d35a00 100644 --- a/judgearena/tasks/definitions/meta_eval/meta-eval-comparia.yaml +++ b/judgearena/tasks/definitions/meta_eval/meta-eval-comparia.yaml @@ -1,21 +1,3 @@ -extends: _base.yaml +extends: _comparia.yaml task: meta-eval-comparia description: Score a judge against human votes from the ComparIA arena. - -variants: - selector: language - values: [fr, en] - -dataset: - sources: - comparia: - type: huggingface_dataset - repo_id: ministere-culture/comparia-votes - revision: "7a40bce496c1f2aa3be4001da85a49cb4743042b" - allow_patterns: [votes.parquet] - -protocol: - arena: ComparIA - -metadata: - reference_implementation: https://huggingface.co/datasets/ministere-culture/comparia-votes diff --git a/judgearena/tasks/definitions/meta_eval/meta-eval-lmarena-100k.yaml b/judgearena/tasks/definitions/meta_eval/meta-eval-lmarena-100k.yaml index 4fa21998..679d3db0 100644 --- a/judgearena/tasks/definitions/meta_eval/meta-eval-lmarena-100k.yaml +++ b/judgearena/tasks/definitions/meta_eval/meta-eval-lmarena-100k.yaml @@ -1,21 +1,3 @@ -extends: _base.yaml +extends: _lmarena-100k.yaml task: meta-eval-lmarena-100k description: Score a judge against human votes from LMSYS Chatbot Arena 100k. - -variants: - selector: language - values: [en, pl, ru, zh, de, ja, es, fr, ko, pt, fa, it, tr, vi, cs, ar, uk] - -dataset: - sources: - lmarena_100k: - type: huggingface_dataset - repo_id: lmarena-ai/arena-human-preference-100k - revision: "72e85b3ddc9c81bf7b659d6b03d4126dfd8fb34a" - allow_patterns: [data/*.parquet] - -protocol: - arena: LMArena-100k - -metadata: - reference_implementation: https://huggingface.co/datasets/lmarena-ai/arena-human-preference-100k diff --git a/judgearena/tasks/definitions/meta_eval/meta-eval-lmarena-140k.yaml b/judgearena/tasks/definitions/meta_eval/meta-eval-lmarena-140k.yaml index 52ba8059..2ffddfd5 100644 --- a/judgearena/tasks/definitions/meta_eval/meta-eval-lmarena-140k.yaml +++ b/judgearena/tasks/definitions/meta_eval/meta-eval-lmarena-140k.yaml @@ -1,21 +1,3 @@ -extends: _base.yaml +extends: _lmarena-140k.yaml task: meta-eval-lmarena-140k description: Score a judge against human votes from LMSYS Chatbot Arena 140k. - -variants: - selector: language - values: [en, pl, ru, zh, de, ja, es, fr, ko, pt, fa, it, tr, vi, cs, ar, uk] - -dataset: - sources: - lmarena_140k: - type: huggingface_dataset - repo_id: lmarena-ai/arena-human-preference-140k - revision: "6322995ab34d7c2693e3f47dd13fa5caa0789a74" - allow_patterns: [data/*.parquet] - -protocol: - arena: LMArena-140k - -metadata: - reference_implementation: https://huggingface.co/datasets/lmarena-ai/arena-human-preference-140k diff --git a/judgearena/tasks/definitions/meta_eval/tune-judge-comparia.yaml b/judgearena/tasks/definitions/meta_eval/tune-judge-comparia.yaml new file mode 100644 index 00000000..b17d9874 --- /dev/null +++ b/judgearena/tasks/definitions/meta_eval/tune-judge-comparia.yaml @@ -0,0 +1,7 @@ +extends: _comparia.yaml +task: tune-judge-comparia +description: Tune judge settings on human votes from the ComparIA arena. + +protocol: + runner: tune_judge + meta_eval_task: meta-eval-comparia diff --git a/judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-100k.yaml b/judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-100k.yaml new file mode 100644 index 00000000..7269fc6f --- /dev/null +++ b/judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-100k.yaml @@ -0,0 +1,7 @@ +extends: _lmarena-100k.yaml +task: tune-judge-lmarena-100k +description: Tune judge settings on human votes from LMSYS Chatbot Arena 100k. + +protocol: + runner: tune_judge + meta_eval_task: meta-eval-lmarena-100k diff --git a/judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-140k.yaml b/judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-140k.yaml new file mode 100644 index 00000000..f7ff84dc --- /dev/null +++ b/judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-140k.yaml @@ -0,0 +1,7 @@ +extends: _lmarena-140k.yaml +task: tune-judge-lmarena-140k +description: Tune judge settings on human votes from LMSYS Chatbot Arena 140k. + +protocol: + runner: tune_judge + meta_eval_task: meta-eval-lmarena-140k diff --git a/judgearena/tasks/schema/__init__.py b/judgearena/tasks/schema/__init__.py index 2c82226f..90039f45 100644 --- a/judgearena/tasks/schema/__init__.py +++ b/judgearena/tasks/schema/__init__.py @@ -15,7 +15,7 @@ ) from judgearena.tasks.schema.dataset import DatasetFields, DatasetSpec from judgearena.tasks.schema.elo import EloProtocol, EloScoringSpec -from judgearena.tasks.schema.meta_eval import MetaEvalProtocol +from judgearena.tasks.schema.meta_eval import MetaEvalProtocol, TuneJudgeProtocol from judgearena.tasks.schema.metrics import MetricSpec, ScoringSpec from judgearena.tasks.schema.mt_bench import ( MTBenchJudgeSpec, @@ -88,4 +88,5 @@ "TaskProvenance", "TaskSelection", "TaskSpec", + "TuneJudgeProtocol", ] diff --git a/judgearena/tasks/schema/meta_eval.py b/judgearena/tasks/schema/meta_eval.py index 3c01d120..6f0b52d6 100644 --- a/judgearena/tasks/schema/meta_eval.py +++ b/judgearena/tasks/schema/meta_eval.py @@ -39,3 +39,11 @@ def _validate_scoring(self) -> MetaEvalProtocol: f"meta-evaluation metric {request.metric} does not support breakdown_by" ) return self + + +class TuneJudgeProtocol(MetaEvalProtocol): + """Search judge settings by running trials of a meta-evaluation task.""" + + runner: Literal["tune_judge"] + meta_eval_task: str = Field(min_length=1) + """Task each trial runs as, so trials share its judgement cache.""" diff --git a/judgearena/tasks/schema/task.py b/judgearena/tasks/schema/task.py index 6f5ab80e..a9f6e61e 100644 --- a/judgearena/tasks/schema/task.py +++ b/judgearena/tasks/schema/task.py @@ -13,7 +13,7 @@ ) from judgearena.tasks.schema.dataset import DatasetSpec from judgearena.tasks.schema.elo import EloProtocol -from judgearena.tasks.schema.meta_eval import MetaEvalProtocol +from judgearena.tasks.schema.meta_eval import MetaEvalProtocol, TuneJudgeProtocol from judgearena.tasks.schema.mt_bench import MTBenchProtocol from judgearena.tasks.schema.mt_bench_101 import MTBench101Protocol from judgearena.tasks.schema.pairwise import PairwiseProtocol @@ -23,7 +23,8 @@ | MTBenchProtocol | MTBench101Protocol | EloProtocol - | MetaEvalProtocol, + | MetaEvalProtocol + | TuneJudgeProtocol, Field(discriminator="runner"), ] diff --git a/judgearena/tuning/runner.py b/judgearena/tuning/runner.py index 8a065b27..71942f82 100644 --- a/judgearena/tuning/runner.py +++ b/judgearena/tuning/runner.py @@ -1,16 +1,17 @@ -"""Successive-halving search over judge configurations on meta-eval battles. +"""Multi-fidelity search over judge configurations on meta-eval battles. Follows Salinas et al., "Tuning LLM Judge Design Decisions for 1/1000 of the -Cost" (ICML 2025). Each rung judges more battles per model with the surviving -configurations and keeps the best by a non-dominated sort on judge cost and -human agreement. The search only sees the validation split; the best -configuration per judge model is then scored once on the held-out test split. +Cost" (ICML 2025), with the fidelity being validation battles per model. A neps +optimizer (priorband by default, seeded with the base config as prior) proposes +each trial, which runs as the meta-eval task so trials share its judgement +cache. The best configuration per judge model at the highest fidelity is then +scored once on the held-out test split. """ from __future__ import annotations +import hashlib import json -import math import subprocess import sys from collections.abc import Callable @@ -26,8 +27,12 @@ from judgearena.artifacts import prepare_run_directory, safe_filename from judgearena.config import RunConfig, TuneJudgeArgs, dump_config from judgearena.log import get_logger -from judgearena.tuning.search_space import apply_overrides, expand_grid -from judgearena.tuning.selection import select_survivors +from judgearena.tuning.search_space import ( + apply_overrides, + axis_overrides, + build_neps_space, + decode_config, +) if TYPE_CHECKING: from judgearena.tasks.schema import ResolvedTaskSpec @@ -37,7 +42,7 @@ TrialExecutor = Callable[[Path], None] """Run the trial config at a path, raising ``CalledProcessError`` on failure.""" -_SUMMARY_COLUMNS = ["trial_id", "judge_model", "agreement", "cost_per_1k_battles"] +_SUMMARY_COLUMNS = ["config_id", "judge_model", "agreement", "cost_per_1k_battles"] def run_trial_subprocess(config_path: Path) -> None: @@ -84,17 +89,21 @@ class _TrialRunner: """Materialize trial configs from one base config and record their results.""" base: dict - overrides: dict[str, dict[str, object]] prices: dict[str, float] execute_trial: TrialExecutor def config( - self, trial_id: str, *, split: str, battles_per_model: int, folder: Path + self, + overrides: dict[str, object], + *, + split: str, + battles_per_model: int, + folder: Path, ) -> RunConfig: values = apply_overrides( self.base, { - **self.overrides[trial_id], + **overrides, "meta_eval.split": split, "meta_eval.battles_per_model": battles_per_model, "run.result_folder": str(folder), @@ -103,15 +112,21 @@ def config( return RunConfig(**values) def run( - self, trial_id: str, *, split: str, battles_per_model: int, folder: Path + self, + overrides: dict[str, object], + *, + split: str, + battles_per_model: int, + folder: Path, ) -> dict: trial_cfg = self.config( - trial_id, split=split, battles_per_model=battles_per_model, folder=folder + overrides, split=split, battles_per_model=battles_per_model, folder=folder ) + encoded = json.dumps(overrides, sort_keys=True) record = { - "trial_id": trial_id, + "config_id": hashlib.sha256(encoded.encode()).hexdigest()[:12], "judge_model": trial_cfg.judge.model, - "overrides": json.dumps(self.overrides[trial_id], sort_keys=True), + "overrides": encoded, "battles_per_model": battles_per_model, } folder.mkdir(parents=True, exist_ok=True) @@ -126,54 +141,54 @@ def run( return {**record, "status": "completed", **_read_trial(folder, price)} -def _successive_halving( +def _objective(record: dict, algorithm: str) -> dict | Exception: + if record["status"] != "completed": + return RuntimeError(f"Trial {record['config_id']} failed.") + disagreement = 1 - record["agreement"] + if algorithm == "mo_hyperband": + return {"objective_to_minimize": [disagreement, record["cost_per_1k_battles"]]} + return {"objective_to_minimize": disagreement} + + +def _search( runner: _TrialRunner, tuning: TuneJudgeArgs, tune_dir: Path -) -> tuple[pd.DataFrame, pd.DataFrame]: - """Return every rung record and the completed survivors of the last rung.""" - alive = list(runner.overrides) +) -> pd.DataFrame: + """Run the neps ask-and-tell loop on the validation split.""" + from neps import AskAndTell + from neps.optimizers import algorithms + + space = build_neps_space(tuning, runner.base) + optimizer = AskAndTell(getattr(algorithms, tuning.algorithm)(space, eta=tuning.eta)) + # neps resamples identical configs on small categorical spaces. + results: dict[tuple[str, int], dict] = {} records = [] - for rung, battles_per_model in enumerate(tuning.rungs): - rung_records = [ - { - "rung": rung, - **runner.run( - trial_id, - split="validation", - battles_per_model=battles_per_model, - folder=tune_dir / f"rung-{rung}" / trial_id, - ), - } - for trial_id in alive - ] - records.extend(rung_records) - completed = pd.DataFrame( - [record for record in rung_records if record["status"] == "completed"] - ) - if completed.empty: - raise RuntimeError(f"Every judge configuration failed at rung {rung}.") - is_last = rung == len(tuning.rungs) - 1 - survivors = select_survivors( - completed["agreement"].to_numpy(dtype=float), - completed["cost_per_1k_battles"].to_numpy(dtype=float), - n_keep=( - len(completed) - if is_last - else math.ceil(len(alive) * tuning.keep_fraction) - ), - min_agreement=tuning.min_agreement, - ) - if not survivors: - raise RuntimeError( - f"No configuration exceeded min_agreement at rung {rung}." + for _ in range(tuning.max_evaluations): + trial = optimizer.ask() + overrides, battles_per_model = decode_config(trial.config) + key = (json.dumps(overrides, sort_keys=True), battles_per_model) + if key not in results: + results[key] = runner.run( + overrides, + split="validation", + battles_per_model=battles_per_model, + folder=tune_dir / "trials" / trial.metadata.id, ) - alive = completed["trial_id"].iloc[survivors].tolist() - logger.info("Rung %d kept %d of %d.", rung, len(alive), len(rung_records)) - return pd.DataFrame(records), completed.iloc[survivors] + records.append({"neps_trial_id": trial.metadata.id, **results[key]}) + optimizer.tell(trial, _objective(results[key], tuning.algorithm)) + return pd.DataFrame(records) + + +def _meta_eval_task(task: ResolvedTaskSpec) -> str: + """Return the meta-eval task ID trials run as, keeping the variant suffix.""" + meta_eval_task = task.spec.protocol.meta_eval_task + if task.selection is None: + return meta_eval_task + return f"{meta_eval_task}-{task.selection.name}" def run_tune_judge( cfg: RunConfig, - task: ResolvedTaskSpec | None = None, + task: ResolvedTaskSpec, *, execute_trial: TrialExecutor = run_trial_subprocess, ) -> pd.DataFrame: @@ -185,46 +200,59 @@ def run_tune_judge( Path(cfg.run.result_folder) / f"tune-{safe_filename(cfg.task)}-{started_at:%Y%m%d_%H%M%S}", ) + base = cfg.model_dump(mode="json", exclude={"tune_judge"}) + if base["judge"]["prompt"] is None and base["judge"]["prompt_preset"] is None: + # Make the task's default preset explicit so it can be the prior. + default_preset = task.spec.protocol.judge.default_prompt_preset + base["judge"]["prompt_preset"] = default_preset runner = _TrialRunner( - base=cfg.model_dump(mode="json", exclude={"tune_judge"}), - overrides={ - trial.id: trial.overrides for trial in expand_grid(tuning.search_space) - }, + base={**base, "task": _meta_eval_task(task)}, prices=tuning.price_per_million_tokens, execute_trial=execute_trial, ) - # Build every trial before judging so invalid settings fail up front. - judge_models = { + # Build every choice before judging so invalid settings fail up front. + judge_models = {cfg.judge.model} | { runner.config( - trial_id, split="validation", battles_per_model=1, folder=tune_dir + overrides, split="validation", battles_per_model=1, folder=tune_dir ).judge.model - for trial_id in runner.overrides + for overrides in axis_overrides(tuning) } missing_prices = sorted(judge_models - set(tuning.price_per_million_tokens)) if missing_prices: raise ValueError( f"tune_judge.price_per_million_tokens is missing models: {missing_prices}" ) - logger.info("Tuning %d judge configurations in %s", len(runner.overrides), tune_dir) + logger.info("Tuning %s with neps %s in %s", cfg.task, tuning.algorithm, tune_dir) - trials, final = _successive_halving(runner, tuning, tune_dir) + trials = _search(runner, tuning, tune_dir) trials.to_parquet(tune_dir / "trials.parquet", index=False) - picks = final.loc[final.groupby("judge_model")["agreement"].idxmax(), "trial_id"] + final = trials[ + (trials["status"] == "completed") + & (trials["battles_per_model"] == tuning.max_battles_per_model) + ].drop_duplicates("config_id") + if final.empty: + raise RuntimeError( + "No configuration completed the highest fidelity; " + "increase tune_judge.max_evaluations." + ) + picks = final.loc[final.groupby("judge_model")["agreement"].idxmax()] test_results = pd.DataFrame( [ runner.run( - trial_id, + json.loads(pick.overrides), split="test", - battles_per_model=tuning.test_battles_per_model or tuning.rungs[-1], - folder=tune_dir / "test" / trial_id, + battles_per_model=( + tuning.test_battles_per_model or tuning.max_battles_per_model + ), + folder=tune_dir / "test" / pick.config_id, ) - for trial_id in picks + for pick in picks.itertuples() ] ) test_results.to_parquet(tune_dir / "test_results.parquet", index=False) - print(f"\n=== Judge tuning: {cfg.task} ===") - print("Validation (last rung):") + print(f"\n=== Judge tuning: {cfg.task} ({tuning.algorithm}) ===") + print(f"Validation ({tuning.max_battles_per_model} battles per model):") ranked = final.sort_values("agreement", ascending=False) print(ranked[_SUMMARY_COLUMNS].to_string(index=False)) print("Test:") diff --git a/judgearena/tuning/search_space.py b/judgearena/tuning/search_space.py index 7485f281..30fd8acb 100644 --- a/judgearena/tuning/search_space.py +++ b/judgearena/tuning/search_space.py @@ -1,55 +1,109 @@ -"""Expand a judge-tuning search space into concrete trial overrides.""" +"""Translate a judge-tuning search space to and from a neps search space.""" from __future__ import annotations import copy -import hashlib -import itertools import json -from collections.abc import Mapping, Sequence -from dataclasses import dataclass +from collections.abc import Iterator, Mapping +from typing import TYPE_CHECKING, Any + +from judgearena.config import FloatRange, TuneJudgeArgs + +if TYPE_CHECKING: + import neps TUNER_OWNED_KEYS = frozenset( {"task", "meta_eval.battles_per_model", "meta_eval.split", "run.result_folder"} ) +FIDELITY = "battles_per_model" + + +def choice_overrides(name: str, value: object) -> dict[str, object]: + """Return the dotted overrides one choice of an axis applies.""" + return dict(value) if isinstance(value, Mapping) else {name: value} -@dataclass(frozen=True) -class Trial: - """One judge configuration, given as dotted ``RunConfig`` overrides.""" +def axis_overrides(tuning: TuneJudgeArgs) -> Iterator[dict[str, object]]: + """Yield every choice and range bound, rejecting tuner-owned keys.""" + for name, axis in tuning.search_space.items(): + values = [axis.lower, axis.upper] if isinstance(axis, FloatRange) else axis + for value in values: + overrides = choice_overrides(name, value) + owned = sorted( + key + for key in overrides + if key in TUNER_OWNED_KEYS or key.startswith("tune_judge") + ) + if owned: + raise ValueError( + f"search_space must not override tuner-owned keys: {owned}" + ) + yield overrides - overrides: dict[str, object] - @property - def id(self) -> str: - payload = json.dumps(self.overrides, sort_keys=True) - return hashlib.sha256(payload.encode()).hexdigest()[:12] +def _lookup(values: Mapping[str, Any], path: str) -> object: + for key in path.split("."): + if not isinstance(values, Mapping) or key not in values: + return None + values = values[key] + return values -def expand_grid(search_space: Mapping[str, Sequence[object]]) -> list[Trial]: - """Return the cartesian product of the search-space axes. +def build_neps_space( + tuning: TuneJudgeArgs, base: Mapping[str, Any] +) -> neps.SearchSpace: + """Return the neps space whose prior is the base config. - An axis name is the dotted ``RunConfig`` path its values override. An axis - whose values are mappings is grouped instead: each value is a set of dotted - overrides that change together, and the axis name is only a label. + Choices are JSON-encoded override mappings so grouped settings stay one + categorical; the fidelity is the number of validation battles per model. """ - axes = [ - [ - dict(value) if isinstance(value, Mapping) else {name: value} - for value in values - ] - for name, values in search_space.items() - ] - keys = {key for choices in axes for choice in choices for key in choice} - owned = sorted( - key for key in keys if key in TUNER_OWNED_KEYS or key.startswith("tune_judge") + import neps + + confidence = tuning.prior_confidence + parameters = {} + for name, axis in tuning.search_space.items(): + if isinstance(axis, FloatRange): + prior = _lookup(base, name) + inside = ( + isinstance(prior, int | float) and axis.lower <= prior <= axis.upper + ) + parameters[name] = neps.HPOFloat( + lower=axis.lower, + upper=axis.upper, + log=axis.log, + prior=prior if inside else None, + prior_confidence=confidence, + ) + continue + overrides = [choice_overrides(name, value) for value in axis] + choices = [json.dumps(choice, sort_keys=True) for choice in overrides] + prior = next( + ( + encoded + for encoded, choice in zip(choices, overrides, strict=True) + if all(_lookup(base, key) == value for key, value in choice.items()) + ), + None, + ) + parameters[name] = neps.HPOCategorical( + choices=choices, prior=prior, prior_confidence=confidence + ) + parameters[FIDELITY] = neps.HPOInteger( + lower=tuning.min_battles_per_model, + upper=tuning.max_battles_per_model, + is_fidelity=True, ) - if owned: - raise ValueError(f"search_space must not override tuner-owned keys: {owned}") - return [ - Trial({key: value for choice in combo for key, value in choice.items()}) - for combo in itertools.product(*axes) - ] + return neps.SearchSpace(parameters) + + +def decode_config(config: Mapping[str, Any]) -> tuple[dict[str, object], int]: + """Return the dotted overrides and battles per model of a neps config.""" + overrides: dict[str, object] = {} + for name, value in config.items(): + if name == FIDELITY: + continue + overrides.update(json.loads(value) if isinstance(value, str) else {name: value}) + return overrides, int(config[FIDELITY]) def apply_overrides( diff --git a/judgearena/tuning/selection.py b/judgearena/tuning/selection.py deleted file mode 100644 index 65cd354d..00000000 --- a/judgearena/tuning/selection.py +++ /dev/null @@ -1,88 +0,0 @@ -"""Multi-objective survivor selection used by successive halving. - -``pareto_efficiency_mask``, ``epsilon_net_indices`` and ``argsort_nondominated`` -are adapted from judgetuning (github.com/geoalgo/judgetuning, Apache-2.0), -which took them from TSBench by Oliver Borchert. -""" - -from __future__ import annotations - -import numpy as np -import scipy.stats as st - - -def pareto_efficiency_mask(X: np.ndarray) -> np.ndarray: - """Mark rows of ``X`` not dominated by another row, all columns minimized.""" - mask = np.ones(X.shape[0], dtype=bool) - for i, point in enumerate(X): - if mask[i]: - dominated = np.all(point <= X[mask], axis=1) & np.any( - point < X[mask], axis=1 - ) - mask[mask] = ~dominated - return mask - - -def epsilon_net_indices(X: np.ndarray, dim: int) -> np.ndarray: - """Order rows so each next row is farthest from those already chosen. - - The seed is the row minimizing column ``dim`` (Clarkson, 2005, p. 17). - """ - remaining = set(range(X.shape[0])) - order = [int(np.argmin(X[:, dim]))] - remaining.remove(order[0]) - while remaining: - candidates = list(remaining) - diff = X[candidates][:, None, :] - X[order][None, :, :] - choice = candidates[np.linalg.norm(diff, axis=-1).min(-1).argmax()] - order.append(choice) - remaining.remove(choice) - return np.array(order) - - -def argsort_nondominated( - X: np.ndarray, *, minimize: list[bool], dim: int, max_items: int | None = None -) -> list[int]: - """Sort rows by successive Pareto fronts, each ordered by an epsilon-net. - - Columns are quantile-normalized, so objectives on different scales weigh - equally. - """ - if X.shape[0] == 1: - return [0] - X = np.where(minimize, X, -X) - X = (st.rankdata(X, axis=0) - 1) / (X.shape[0] - 1) - remaining = np.arange(X.shape[0]) - indices: list[int] = [] - while remaining.size and (max_items is None or len(indices) < max_items): - pareto_mask = pareto_efficiency_mask(X[remaining]) - front = remaining[pareto_mask] - indices.extend(front[epsilon_net_indices(X[front], dim=dim)].tolist()) - remaining = remaining[~pareto_mask] - return indices[:max_items] - - -def select_survivors( - agreement: np.ndarray, - cost: np.ndarray, - *, - n_keep: int, - min_agreement: float | None = None, -) -> list[int]: - """Return indices of the ``n_keep`` best trials by (agreement ↑, cost ↓). - - Trials at or below ``min_agreement`` are dropped first; the paper used a - threshold just above a length-based baseline. - """ - candidates = np.arange(len(agreement)) - if min_agreement is not None: - candidates = candidates[agreement > min_agreement] - if not candidates.size: - return [] - order = argsort_nondominated( - np.stack([cost[candidates], agreement[candidates]], axis=1), - minimize=[True, False], - dim=1, - max_items=n_keep, - ) - return candidates[order].tolist() diff --git a/pyproject.toml b/pyproject.toml index 2f3586b3..ee6e29ca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -113,3 +113,6 @@ vllm = [ "transformers>=5.5.1,<6.0.0", ] llamacpp = ["llama-cpp-python>=0.3.0"] +tune = [ + "neural-pipeline-search>=0.16.0", +] diff --git a/tests/test_mt_bench_downloads.py b/tests/test_mt_bench_downloads.py index 5ea1522b..f85a8230 100644 --- a/tests/test_mt_bench_downloads.py +++ b/tests/test_mt_bench_downloads.py @@ -127,6 +127,9 @@ def test_download_all_includes_mt_bench(tmp_path, monkeypatch): "meta-eval-lmarena-140k", "mt-bench", "mt-bench-101", + "tune-judge-comparia", + "tune-judge-lmarena-100k", + "tune-judge-lmarena-140k", ] assert all(path == tables_dir for _, path in hf_datasets) diff --git a/tests/test_tuning.py b/tests/test_tuning.py index 89eb96e8..6902a3fc 100644 --- a/tests/test_tuning.py +++ b/tests/test_tuning.py @@ -5,49 +5,73 @@ import json import subprocess -import numpy as np import pandas as pd import pytest from pydantic import ValidationError from judgearena.config import RunConfig, load_config +from judgearena.tasks.registry import get_packaged_task from judgearena.tuning.runner import run_tune_judge -from judgearena.tuning.search_space import apply_overrides, expand_grid -from judgearena.tuning.selection import select_survivors +from judgearena.tuning.search_space import ( + axis_overrides, + build_neps_space, + decode_config, +) +_SEARCH_SPACE = { + "judge.temperature": [0.0, 0.5, 1.0], + "prompt": [ + {"judge.prompt_preset": "meta-eval-pair-score", "judge.swap_mode": "fixed"}, + {"judge.prompt_preset": "meta-eval-pair-score", "judge.swap_mode": "both"}, + ], +} -def test_expand_grid_combines_plain_and_grouped_axes(): - trials = expand_grid( - { - "judge.temperature": [0.0, 1.0], - "prompt": [ - {"judge.prompt_preset": "meta-eval-pair-score"}, - {"judge.prompt_preset": "alpaca-eval", "judge.top_logprobs": 20}, - ], - } - ) - assert len({trial.id for trial in trials}) == 4 - assert { - "judge.temperature": 1.0, - "judge.prompt_preset": "alpaca-eval", - "judge.top_logprobs": 20, - } in [trial.overrides for trial in trials] - assert apply_overrides({"judge": {"model": "m"}}, trials[0].overrides)["judge"] == { - "model": "m", - "temperature": 0.0, - "prompt_preset": "meta-eval-pair-score", - } - with pytest.raises(ValueError, match="tuner-owned"): - expand_grid({"meta_eval.split": ["test"]}) +def _tuning_config(tmp_path, **tune_judge): + return RunConfig( + task="tune-judge-comparia-fr", + judge={"model": "Dummy/judge", "temperature": 0.0}, + run={"result_folder": str(tmp_path), "no_log_file": True}, + tune_judge={ + "search_space": _SEARCH_SPACE, + "price_per_million_tokens": {"Dummy/judge": 1.0}, + "max_evaluations": 6, + "min_battles_per_model": 2, + "max_battles_per_model": 6, + **tune_judge, + }, + ) -def test_select_survivors_sorts_pareto_fronts_by_agreement_first(): - agreement = np.array([0.50, 0.60, 0.55, 0.40]) - cost = np.array([1.0, 2.0, 3.0, 0.5]) +def test_neps_space_uses_base_config_as_prior(tmp_path): + pytest.importorskip("neps") + cfg = _tuning_config(tmp_path) + base = cfg.model_dump(mode="json") + base["judge"]["prompt_preset"] = "meta-eval-pair-score" + space = build_neps_space(cfg.tune_judge, base) - assert select_survivors(agreement, cost, n_keep=2, min_agreement=0.45) == [1, 0] - assert select_survivors(agreement, cost, n_keep=3, min_agreement=0.45) == [1, 0, 2] + temperature, prompt = ( + space.searchables[k] for k in ("judge.temperature", "prompt") + ) + assert json.loads(temperature.prior) == {"judge.temperature": 0.0} + assert json.loads(prompt.prior) == _SEARCH_SPACE["prompt"][0] + overrides, battles = decode_config( + { + "judge.temperature": temperature.choices[2], + "prompt": prompt.choices[1], + "battles_per_model": 6, + } + ) + assert overrides == {"judge.temperature": 1.0, **_SEARCH_SPACE["prompt"][1]} + assert battles == 6 + with pytest.raises(ValueError, match="tuner-owned"): + list( + axis_overrides( + _tuning_config(tmp_path).tune_judge.model_copy( + update={"search_space": {"meta_eval.split": ["test"]}} + ) + ) + ) def _fake_trial(config_path): @@ -68,65 +92,47 @@ def _fake_trial(config_path): } } (run_dir / "results.json").write_text(json.dumps({"metrics": metrics})) - completion = "x " * (1 + 10 * (cfg.judge.swap_mode == "both")) pd.DataFrame( - { - "battle_id": ["b1"], - "judge_input": ["prompt"], - "judge_completion": [completion], - } + {"battle_id": ["b1"], "judge_input": ["prompt"], "judge_completion": ["x"]} ).to_parquet(run_dir / "annotations.parquet") -def _tuning_config(tmp_path, **tune_judge): - return RunConfig( - task="meta-eval-comparia", - judge={"model": "Dummy/judge"}, - run={"result_folder": str(tmp_path), "no_log_file": True}, - tune_judge={ - "search_space": { - "judge.temperature": [0.0, 0.5, 1.0], - "judge.swap_mode": ["fixed", "both"], - }, - "rungs": [2, 4], - "price_per_million_tokens": {"Dummy/judge": 1.0}, - **tune_judge, - }, +@pytest.mark.parametrize("algorithm", ["priorband", "mo_hyperband"]) +def test_run_tune_judge_searches_validation_and_scores_pick_on_test( + tmp_path, algorithm +): + pytest.importorskip("neps") + cfg = _tuning_config(tmp_path, algorithm=algorithm) + test_results = run_tune_judge( + cfg, get_packaged_task(cfg.task), execute_trial=_fake_trial ) - -def test_run_tune_judge_halves_on_validation_and_scores_pick_on_test(tmp_path): - test_results = run_tune_judge(_tuning_config(tmp_path), execute_trial=_fake_trial) - (tune_dir,) = tmp_path.glob("tune-*") trials = pd.read_parquet(tune_dir / "trials.parquet") - assert trials.groupby("rung").size().tolist() == [6, 2] - assert (trials["status"] == "failed").sum() == 1 - (pick,) = test_results["trial_id"] - pick_cfg = load_config(tune_dir / "test" / pick / "config.yaml") - assert pick_cfg.meta_eval.split == "test" - assert (pick_cfg.judge.temperature, pick_cfg.judge.swap_mode) == (0.0, "fixed") - for rung in trials.itertuples(): - rung_cfg = load_config( - tune_dir / f"rung-{rung.rung}" / rung.trial_id / "config.yaml" - ) - assert rung_cfg.meta_eval.split == "validation" - assert rung_cfg.tune_judge is None + assert len(trials) == 6 + assert set(trials["battles_per_model"]) <= {2, 6} + (pick,) = test_results.itertuples() + pick_cfg = load_config(tune_dir / "test" / pick.config_id / "config.yaml") + assert (pick_cfg.task, pick_cfg.meta_eval.split) == ( + "meta-eval-comparia-fr", + "test", + ) + for trial_dir in (tune_dir / "trials").iterdir(): + trial_cfg = load_config(trial_dir / "config.yaml") + assert trial_cfg.task == "meta-eval-comparia-fr" + assert trial_cfg.meta_eval.split == "validation" + assert trial_cfg.tune_judge is None -def test_tune_judge_requires_prices_and_meta_eval_task(tmp_path): +def test_tune_judge_requires_prices_and_tune_task(tmp_path): + cfg = _tuning_config(tmp_path, price_per_million_tokens={}) with pytest.raises(ValueError, match="missing models"): - run_tune_judge( - _tuning_config(tmp_path, price_per_million_tokens={}), - execute_trial=_fake_trial, - ) - with pytest.raises(ValidationError, match="only valid for meta-evaluation"): + run_tune_judge(cfg, get_packaged_task(cfg.task), execute_trial=_fake_trial) + with pytest.raises(ValidationError, match="tune-judge tasks"): RunConfig( - task="alpaca-eval", - model={"name": "Dummy/a"}, + task="meta-eval-comparia-fr", judge={"model": "Dummy/judge"}, - tune_judge={ - "search_space": {"judge.temperature": [0.0]}, - "price_per_million_tokens": {}, - }, + tune_judge=cfg.tune_judge.model_dump(), ) + with pytest.raises(ValidationError, match="tune-judge tasks"): + RunConfig(task="tune-judge-comparia-fr", judge={"model": "Dummy/judge"}) From 7e911646309e6197d4f9db02715698633e10d844 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 2 Oct 2026 15:52:22 +0200 Subject: [PATCH 05/20] Name tuning run directories after the tune-judge task --- judgearena/tuning/runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/judgearena/tuning/runner.py b/judgearena/tuning/runner.py index 71942f82..641c775c 100644 --- a/judgearena/tuning/runner.py +++ b/judgearena/tuning/runner.py @@ -198,7 +198,7 @@ def run_tune_judge( tune_dir = prepare_run_directory( cfg, Path(cfg.run.result_folder) - / f"tune-{safe_filename(cfg.task)}-{started_at:%Y%m%d_%H%M%S}", + / f"{safe_filename(cfg.task)}-{started_at:%Y%m%d_%H%M%S}", ) base = cfg.model_dump(mode="json", exclude={"tune_judge"}) if base["judge"]["prompt"] is None and base["judge"]["prompt_preset"] is None: From baccfac617440d5dedb03e93bd7625c4feea7f8f Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 7 Oct 2026 00:58:52 +0200 Subject: [PATCH 06/20] Use NePS 0.17 runtime for resumable judge tuning --- .github/workflows/ci.yml | 4 +- judgearena/config.py | 93 ++++++---- judgearena/tuning/runner.py | 190 +++++++++++--------- judgearena/tuning/search_space.py | 128 +++++--------- judgearena/tuning/session.py | 95 ++++++++++ pyproject.toml | 2 +- tests/test_tuning.py | 278 +++++++++++++++++++----------- 7 files changed, 485 insertions(+), 305 deletions(-) create mode 100644 judgearena/tuning/session.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b9cfa7cc..5c7e6977 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -24,8 +24,8 @@ jobs: enable-cache: true python-version: "3.12" - name: Install dependencies - run: uv sync + run: uv sync --extra tune - name: Run pre-commit run: uv run pre-commit run --all-files - name: Test with pytest - run: uv run pytest + run: uv run --extra tune pytest diff --git a/judgearena/config.py b/judgearena/config.py index 6c5f88c1..0719332d 100644 --- a/judgearena/config.py +++ b/judgearena/config.py @@ -412,57 +412,74 @@ class MetaEvalArgs(BaseModel): """Share of arena prompts assigned to the validation partition.""" -class FloatRange(BaseModel): - """Continuous search-space axis.""" - - model_config = ConfigDict(extra="forbid") - - lower: float - upper: float - log: bool = False - - class TuneJudgeArgs(BaseModel): - """Multi-fidelity neps search over judge settings on meta-eval battles.""" + """NePS search settings and judge-evaluation objectives.""" model_config = ConfigDict(use_attribute_docstrings=True, extra="forbid") - search_space: dict[str, list[Any] | FloatRange] = Field(min_length=1) - """Axes keyed by dotted config path (``judge.temperature``), each a list of - choices or a ``{lower, upper}`` range. Choices that are mappings of dotted - overrides change those settings together. The base config is the prior.""" - - price_per_million_tokens: dict[str, float] - """USD per million judge tokens for every searched ``judge.model``.""" - - algorithm: Literal["priorband", "hyperband", "mo_hyperband"] = "priorband" - """neps optimizer; ``mo_hyperband`` also minimizes judge cost.""" - - max_evaluations: int = Field(gt=0) - """Validation trials to run, counting each promotion to more battles.""" + neps: dict[str, Any] + """NePS constructor specs in pipeline_space and native neps.run arguments.""" - min_battles_per_model: int = Field(default=10, gt=0) - """Lowest fidelity, in validation battles per model.""" + objectives: list[Literal["agreement", "cost_per_1k_battles"]] = Field( + default=["agreement", "cost_per_1k_battles"], min_length=1 + ) + """Ordered objectives; agreement is maximized and expected cost minimized.""" - max_battles_per_model: int = 90 - """Highest fidelity, in validation battles per model.""" + price_per_million_tokens: dict[str, float] + """USD per million judge tokens, including cached judgements.""" - eta: int = Field(default=3, ge=2) - """Hyperband reduction factor between fidelities.""" + test_battles_per_model: int | None = Field(default=None, gt=0) + """Test budget; defaults to the highest fidelity.""" - prior_confidence: Literal["low", "medium", "high"] = "medium" - """How strongly priorband trusts the base config.""" + run_dir: Path | None = None + """Explicit directory to resume or share; otherwise create a fresh run.""" - test_battles_per_model: int | None = Field(default=None, gt=0) - """Test battles per model for the selected judges. Defaults to the highest - fidelity.""" + search_only: bool = False + """Join an initialized run as a validation-only worker.""" @model_validator(mode="after") - def _validate_fidelity(self) -> TuneJudgeArgs: - if self.min_battles_per_model >= self.max_battles_per_model: + def _validate_runtime(self) -> TuneJudgeArgs: + reserved = { + "evaluate_pipeline", + "root_directory", + "overwrite_root_directory", + "continue_until_max_evaluation_completed", + "total_cost_to_spend", + "worker_cost_to_spend", + } + if forbidden := reserved.intersection(self.neps): raise ValueError( - "tune_judge.min_battles_per_model must be below max_battles_per_model." + f"Unsupported or managed NePS options: {sorted(forbidden)}" ) + if self.search_only and self.run_dir is None: + raise ValueError("search_only requires an explicit tune_judge.run_dir") + if not self.search_only: + if not any( + self.neps.get(k) + for k in ("total_evaluations_to_spend", "total_fidelities_to_spend") + ): + raise ValueError("The primary worker requires a finite global budget") + if any( + self.neps.get(k) is not None + for k in ("worker_evaluations_to_spend", "worker_fidelities_to_spend") + ): + raise ValueError("Worker-local budgets require search_only") + optimizer = self.neps.get("optimizer", "auto") + name = ( + optimizer.get("name") + if isinstance(optimizer, dict) + else optimizer[0] + if isinstance(optimizer, (list, tuple)) + else optimizer + ) + if name in {"neps_priorband", "neps_hyperband"} and len(self.objectives) != 1: + raise ValueError(f"{name} requires a single objective") + if name in {"mo_hyperband", "primo"} and len(self.objectives) != 2: + raise ValueError(f"{name} requires agreement and cost objectives") + if len(set(self.objectives)) != len(self.objectives): + raise ValueError("objectives must not contain duplicates") + if any(price < 0 for price in self.price_per_million_tokens.values()): + raise ValueError("Judge token prices must be nonnegative") return self diff --git a/judgearena/tuning/runner.py b/judgearena/tuning/runner.py index 641c775c..334030b1 100644 --- a/judgearena/tuning/runner.py +++ b/judgearena/tuning/runner.py @@ -1,33 +1,26 @@ -"""Multi-fidelity search over judge configurations on meta-eval battles. - -Follows Salinas et al., "Tuning LLM Judge Design Decisions for 1/1000 of the -Cost" (ICML 2025), with the fidelity being validation battles per model. A neps -optimizer (priorband by default, seeded with the base config as prior) proposes -each trial, which runs as the meta-eval task so trials share its judgement -cache. The best configuration per judge model at the highest fidelity is then -scored once on the held-out test split. -""" +"""Run NePS searches using meta-evaluation subprocesses.""" from __future__ import annotations import hashlib import json +import random import subprocess import sys from collections.abc import Callable from dataclasses import dataclass -from datetime import UTC, datetime from functools import cache from pathlib import Path from typing import TYPE_CHECKING +import numpy as np import pandas as pd import tiktoken -from judgearena.artifacts import prepare_run_directory, safe_filename -from judgearena.config import RunConfig, TuneJudgeArgs, dump_config +from judgearena.config import RunConfig, dump_config from judgearena.log import get_logger from judgearena.tuning.search_space import ( + FIDELITY, apply_overrides, axis_overrides, build_neps_space, @@ -37,6 +30,8 @@ if TYPE_CHECKING: from judgearena.tasks.schema import ResolvedTaskSpec +from judgearena.tuning.session import collect_trials, prepare_session + logger = get_logger(__name__) TrialExecutor = Callable[[Path], None] @@ -132,6 +127,12 @@ def run( folder.mkdir(parents=True, exist_ok=True) config_path = folder / "config.yaml" dump_config(trial_cfg, config_path) + if split == "test" and any(folder.glob("*/results.json")): + return { + **record, + "status": "completed", + **_read_trial(folder, self.prices[trial_cfg.judge.model]), + } try: self.execute_trial(config_path) except subprocess.CalledProcessError as exc: @@ -141,41 +142,64 @@ def run( return {**record, "status": "completed", **_read_trial(folder, price)} -def _objective(record: dict, algorithm: str) -> dict | Exception: +def _objective(record: dict, objectives: list[str]) -> dict: if record["status"] != "completed": - return RuntimeError(f"Trial {record['config_id']} failed.") - disagreement = 1 - record["agreement"] - if algorithm == "mo_hyperband": - return {"objective_to_minimize": [disagreement, record["cost_per_1k_battles"]]} - return {"objective_to_minimize": disagreement} + return { + "exception": RuntimeError(f"Trial {record['config_id']} failed"), + "info_dict": record, + } + values = [ + 1 - record[key] if key == "agreement" else record[key] for key in objectives + ] + return { + "objective_to_minimize": values[0] if len(values) == 1 else values, + "info_dict": record, + } -def _search( - runner: _TrialRunner, tuning: TuneJudgeArgs, tune_dir: Path -) -> pd.DataFrame: - """Run the neps ask-and-tell loop on the validation split.""" - from neps import AskAndTell - from neps.optimizers import algorithms - - space = build_neps_space(tuning, runner.base) - optimizer = AskAndTell(getattr(algorithms, tuning.algorithm)(space, eta=tuning.eta)) - # neps resamples identical configs on small categorical spaces. - results: dict[tuple[str, int], dict] = {} - records = [] - for _ in range(tuning.max_evaluations): - trial = optimizer.ask() - overrides, battles_per_model = decode_config(trial.config) - key = (json.dumps(overrides, sort_keys=True), battles_per_model) - if key not in results: - results[key] = runner.run( - overrides, - split="validation", - battles_per_model=battles_per_model, - folder=tune_dir / "trials" / trial.metadata.id, - ) - records.append({"neps_trial_id": trial.metadata.id, **results[key]}) - optimizer.tell(trial, _objective(results[key], tuning.algorithm)) - return pd.DataFrame(records) +def _search(runner: _TrialRunner, cfg: RunConfig, tune_dir: Path, space) -> None: + """Let NePS persist and coordinate the search across workers.""" + import neps + + def evaluate(pipeline_directory: Path, **config): + overrides, battles = decode_config(config) + record = runner.run( + overrides, + split="validation", + battles_per_model=battles, + folder=tune_dir / "trials" / pipeline_directory.name, + ) + return _objective(record, cfg.tune_judge.objectives) + + options = { + key: value + for key, value in cfg.tune_judge.neps.items() + if key != "pipeline_space" + } + # NePS 0.17's space compatibility check does not inspect optimizer mappings. + if isinstance(optimizer := options.get("optimizer"), dict): + options["optimizer"] = ( + optimizer["name"], + {key: value for key, value in optimizer.items() if key != "name"}, + ) + root = tune_dir / "neps" + continuing = (root / "pipeline_space.pkl").exists() + if cfg.tune_judge.search_only: + options.pop("total_evaluations_to_spend", None) + options.pop("total_fidelities_to_spend", None) + elif not continuing: + import torch + + random.seed(cfg.run.seed) + np.random.seed(cfg.run.seed) + torch.manual_seed(cfg.run.seed) + neps.run( + evaluate_pipeline=evaluate, + pipeline_space=None if continuing else space, + root_directory=root, + continue_until_max_evaluation_completed=False, + **options, + ) def _meta_eval_task(task: ResolvedTaskSpec) -> str: @@ -194,69 +218,73 @@ def run_tune_judge( ) -> pd.DataFrame: """Tune judge settings on validation battles and score the picks on test.""" tuning = cfg.tune_judge - started_at = datetime.now(UTC) - tune_dir = prepare_run_directory( - cfg, - Path(cfg.run.result_folder) - / f"{safe_filename(cfg.task)}-{started_at:%Y%m%d_%H%M%S}", - ) + specs = tuning.neps["pipeline_space"] + space = build_neps_space(specs) + max_battles = specs[FIDELITY]["upper"] base = cfg.model_dump(mode="json", exclude={"tune_judge"}) - if base["judge"]["prompt"] is None and base["judge"]["prompt_preset"] is None: - # Make the task's default preset explicit so it can be the prior. - default_preset = task.spec.protocol.judge.default_prompt_preset - base["judge"]["prompt_preset"] = default_preset runner = _TrialRunner( base={**base, "task": _meta_eval_task(task)}, prices=tuning.price_per_million_tokens, execute_trial=execute_trial, ) - # Build every choice before judging so invalid settings fail up front. - judge_models = {cfg.judge.model} | { - runner.config( - overrides, split="validation", battles_per_model=1, folder=tune_dir - ).judge.model - for overrides in axis_overrides(tuning) - } - missing_prices = sorted(judge_models - set(tuning.price_per_million_tokens)) - if missing_prices: + models = set() + for overrides in axis_overrides(specs): + trial_cfg = runner.config( + overrides, + split="validation", + battles_per_model=1, + folder=Path(cfg.run.result_folder), + ) + if "judge.model" not in specs or "judge.model" in overrides: + models.add(trial_cfg.judge.model) + if "judge.model" not in specs: + models.add(cfg.judge.model) + if missing := sorted(models - set(tuning.price_per_million_tokens)): raise ValueError( - f"tune_judge.price_per_million_tokens is missing models: {missing_prices}" + f"tune_judge.price_per_million_tokens is missing models: {missing}" ) - logger.info("Tuning %s with neps %s in %s", cfg.task, tuning.algorithm, tune_dir) - - trials = _search(runner, tuning, tune_dir) + tune_dir = prepare_session(cfg) + logger.info("Tuning %s in %s", cfg.task, tune_dir) + _search(runner, cfg, tune_dir, space) + if tuning.search_only: + return pd.DataFrame() + trials = collect_trials(tune_dir / "neps", wait=True) trials.to_parquet(tune_dir / "trials.parquet", index=False) + return _evaluate_picks(runner, tuning, tune_dir, trials, max_battles) + + +def _evaluate_picks(runner, tuning, tune_dir, trials, max_battles) -> pd.DataFrame: final = trials[ - (trials["status"] == "completed") - & (trials["battles_per_model"] == tuning.max_battles_per_model) + (trials["status"] == "completed") & (trials["battles_per_model"] == max_battles) ].drop_duplicates("config_id") if final.empty: raise RuntimeError( - "No configuration completed the highest fidelity; " - "increase tune_judge.max_evaluations." + "No configuration completed the highest fidelity; increase the NePS global budget" ) - picks = final.loc[final.groupby("judge_model")["agreement"].idxmax()] + ranked = final.sort_values( + ["agreement", "cost_per_1k_battles", "config_id"], ascending=[False, True, True] + ) + picks = ranked.groupby("judge_model", sort=False).head(1) test_results = pd.DataFrame( [ runner.run( json.loads(pick.overrides), split="test", - battles_per_model=( - tuning.test_battles_per_model or tuning.max_battles_per_model - ), + battles_per_model=tuning.test_battles_per_model or max_battles, folder=tune_dir / "test" / pick.config_id, ) for pick in picks.itertuples() ] ) test_results.to_parquet(tune_dir / "test_results.parquet", index=False) - - print(f"\n=== Judge tuning: {cfg.task} ({tuning.algorithm}) ===") - print(f"Validation ({tuning.max_battles_per_model} battles per model):") - ranked = final.sort_values("agreement", ascending=False) + print("\n=== Judge tuning ===") + print("Validation:") print(ranked[_SUMMARY_COLUMNS].to_string(index=False)) print("Test:") - summary = test_results.reindex(columns=[*_SUMMARY_COLUMNS, "status"]) - print(summary.to_string(index=False)) + print( + test_results.reindex(columns=[*_SUMMARY_COLUMNS, "status"]).to_string( + index=False + ) + ) print(f"Results: {tune_dir}") return test_results diff --git a/judgearena/tuning/search_space.py b/judgearena/tuning/search_space.py index 30fd8acb..3ce6720b 100644 --- a/judgearena/tuning/search_space.py +++ b/judgearena/tuning/search_space.py @@ -1,115 +1,69 @@ -"""Translate a judge-tuning search space to and from a neps search space.""" +"""Construct NePS domains using its public constructor arguments.""" from __future__ import annotations import copy -import json from collections.abc import Iterator, Mapping from typing import TYPE_CHECKING, Any -from judgearena.config import FloatRange, TuneJudgeArgs - if TYPE_CHECKING: import neps -TUNER_OWNED_KEYS = frozenset( - {"task", "meta_eval.battles_per_model", "meta_eval.split", "run.result_folder"} -) -FIDELITY = "battles_per_model" - - -def choice_overrides(name: str, value: object) -> dict[str, object]: - """Return the dotted overrides one choice of an axis applies.""" - return dict(value) if isinstance(value, Mapping) else {name: value} - - -def axis_overrides(tuning: TuneJudgeArgs) -> Iterator[dict[str, object]]: - """Yield every choice and range bound, rejecting tuner-owned keys.""" - for name, axis in tuning.search_space.items(): - values = [axis.lower, axis.upper] if isinstance(axis, FloatRange) else axis - for value in values: - overrides = choice_overrides(name, value) - owned = sorted( - key - for key in overrides - if key in TUNER_OWNED_KEYS or key.startswith("tune_judge") - ) - if owned: - raise ValueError( - f"search_space must not override tuner-owned keys: {owned}" - ) - yield overrides - - -def _lookup(values: Mapping[str, Any], path: str) -> object: - for key in path.split("."): - if not isinstance(values, Mapping) or key not in values: - return None - values = values[key] - return values - +FIDELITY = "meta_eval.battles_per_model" +DOMAIN_TYPES = frozenset({"Categorical", "Float", "Integer", "IntegerFidelity"}) -def build_neps_space( - tuning: TuneJudgeArgs, base: Mapping[str, Any] -) -> neps.SearchSpace: - """Return the neps space whose prior is the base config. - Choices are JSON-encoded override mappings so grouped settings stay one - categorical; the fidelity is the number of validation battles per model. - """ +def build_neps_space(specs: Mapping[str, dict]) -> neps.PipelineSpace: + """Translate constructor tags to domains without changing their arguments.""" import neps - confidence = tuning.prior_confidence parameters = {} - for name, axis in tuning.search_space.items(): - if isinstance(axis, FloatRange): - prior = _lookup(base, name) - inside = ( - isinstance(prior, int | float) and axis.lower <= prior <= axis.upper - ) - parameters[name] = neps.HPOFloat( - lower=axis.lower, - upper=axis.upper, - log=axis.log, - prior=prior if inside else None, - prior_confidence=confidence, - ) + for name, spec in specs.items(): + kind = spec["type"] + if kind not in DOMAIN_TYPES: + raise ValueError(f"Unsupported NePS domain: {kind}") + if name == FIDELITY: + if kind != "IntegerFidelity": + raise ValueError(f"{FIDELITY} must use IntegerFidelity") + elif not ( + name.startswith("judge.") or name == "generation.truncate_judge_input_chars" + ): + raise ValueError(f"Search space must not override tuner-owned keys: {name}") + elif kind == "IntegerFidelity": + raise ValueError(f"Only {FIDELITY} may be a fidelity") + kwargs = {key: value for key, value in spec.items() if key != "type"} + if kind == "Categorical": + kwargs["choices"] = tuple(kwargs["choices"]) + parameters[name] = getattr(neps, kind)(**kwargs) + if FIDELITY not in parameters: + raise ValueError(f"Search space requires {FIDELITY}") + return type("JudgeSpace", (neps.PipelineSpace,), parameters)() + + +def axis_overrides(specs: Mapping[str, dict]) -> Iterator[dict[str, object]]: + """Yield candidate values for preflight RunConfig validation.""" + for name, spec in specs.items(): + if name == FIDELITY: continue - overrides = [choice_overrides(name, value) for value in axis] - choices = [json.dumps(choice, sort_keys=True) for choice in overrides] - prior = next( - ( - encoded - for encoded, choice in zip(choices, overrides, strict=True) - if all(_lookup(base, key) == value for key, value in choice.items()) - ), - None, + values = ( + spec["choices"] + if spec["type"] == "Categorical" + else [spec["lower"], spec["upper"]] ) - parameters[name] = neps.HPOCategorical( - choices=choices, prior=prior, prior_confidence=confidence - ) - parameters[FIDELITY] = neps.HPOInteger( - lower=tuning.min_battles_per_model, - upper=tuning.max_battles_per_model, - is_fidelity=True, - ) - return neps.SearchSpace(parameters) + for value in values: + yield {name: value} def decode_config(config: Mapping[str, Any]) -> tuple[dict[str, object], int]: - """Return the dotted overrides and battles per model of a neps config.""" - overrides: dict[str, object] = {} - for name, value in config.items(): - if name == FIDELITY: - continue - overrides.update(json.loads(value) if isinstance(value, str) else {name: value}) - return overrides, int(config[FIDELITY]) + """Split the sampled judge settings from the evaluation fidelity.""" + overrides = dict(config) + return overrides, int(overrides.pop(FIDELITY)) def apply_overrides( base: Mapping[str, object], overrides: Mapping[str, object] ) -> dict[str, object]: - """Return a copy of ``base`` with dotted-path ``overrides`` assigned.""" + """Return a copy of base with dotted-path overrides assigned.""" values = copy.deepcopy(dict(base)) for path, value in overrides.items(): *parents, leaf = path.split(".") diff --git a/judgearena/tuning/session.py b/judgearena/tuning/session.py new file mode 100644 index 00000000..d6c6ae02 --- /dev/null +++ b/judgearena/tuning/session.py @@ -0,0 +1,95 @@ +"""Explicit tuning sessions and collection of shared NePS state.""" + +from __future__ import annotations + +import hashlib +import json +import time +from datetime import UTC, datetime +from importlib.metadata import PackageNotFoundError, version +from pathlib import Path + +import pandas as pd + +from judgearena.artifacts import prepare_run_directory, safe_filename +from judgearena.config import RunConfig + + +def _config_hash(cfg: RunConfig) -> str: + payload = cfg.model_dump(mode="json") + payload["run"] = {"seed": cfg.run.seed} + tuning = payload["tune_judge"] + tuning.pop("run_dir") + tuning.pop("search_only") + tuning["neps"] = { + key: value + for key, value in tuning["neps"].items() + if not key.startswith("worker_") + } + if cfg.judge.prompt is not None: + payload["prompt_contents"] = [ + path.read_text() + for path in (cfg.judge.prompt.system_file, cfg.judge.prompt.user_file) + ] + return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest() + + +def prepare_session(cfg: RunConfig) -> Path: + """Create a new run or verify an explicitly selected existing run.""" + tuning = cfg.tune_judge + tune_dir = tuning.run_dir or ( + Path(cfg.run.result_folder) + / f"{safe_filename(cfg.task)}-{datetime.now(UTC):%Y%m%d_%H%M%S_%f}" + ) + tune_dir = tune_dir.resolve() + metadata_path = tune_dir / "tuning-metadata.json" + fingerprint = _config_hash(cfg) + if metadata_path.exists(): + metadata = json.loads(metadata_path.read_text()) + if metadata["config_hash"] != fingerprint: + raise ValueError(f"Incompatible configuration for tuning run {tune_dir}") + if not tuning.search_only: + prepare_run_directory(cfg, tune_dir) + elif tuning.search_only: + raise ValueError("Start the primary process before joining its run_dir") + else: + prepare_run_directory(cfg, tune_dir) + versions = {"neural-pipeline-search": version("neural-pipeline-search")} + try: + versions["vllm"] = version("vllm") + except PackageNotFoundError: + pass + metadata = {"config_hash": fingerprint, "versions": versions} + temporary = metadata_path.with_suffix(".tmp") + temporary.write_text(json.dumps(metadata, indent=2)) + temporary.replace(metadata_path) + if tuning.search_only: + if not (tune_dir / "neps" / "pipeline_space.pkl").exists(): + raise ValueError("The primary process has not initialized NePS yet") + return tune_dir + + +def collect_trials(root: Path, *, wait: bool) -> pd.DataFrame: + """Wait for active evaluations and read all workers' trial records.""" + from neps.state import NePSState + + state = NePSState.create_or_load(root, load_only=True) + while True: + trials = state.lock_and_read_trials() + active = [ + t for t in trials.values() if t.metadata.state in {"pending", "evaluating"} + ] + if not wait or not active: + break + time.sleep(1) + records = [] + for trial in trials.values(): + if trial.metadata.state in {"success", "failed", "crashed"}: + records.append( + { + "neps_trial_id": trial.id, + "status": "failed", + **trial.report.extra, + } + ) + return pd.DataFrame(records) diff --git a/pyproject.toml b/pyproject.toml index ee6e29ca..a1887c56 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -114,5 +114,5 @@ vllm = [ ] llamacpp = ["llama-cpp-python>=0.3.0"] tune = [ - "neural-pipeline-search>=0.16.0", + "neural-pipeline-search==0.17.0", ] diff --git a/tests/test_tuning.py b/tests/test_tuning.py index 6902a3fc..6c1f9d18 100644 --- a/tests/test_tuning.py +++ b/tests/test_tuning.py @@ -1,138 +1,224 @@ -"""Tests for judge configuration tuning.""" +"""Contracts for native NePS judge tuning and persisted trial artifacts.""" from __future__ import annotations import json -import subprocess +from types import SimpleNamespace import pandas as pd import pytest from pydantic import ValidationError -from judgearena.config import RunConfig, load_config +from judgearena.config import RunConfig, dump_config, load_config from judgearena.tasks.registry import get_packaged_task -from judgearena.tuning.runner import run_tune_judge +from judgearena.tuning import runner from judgearena.tuning.search_space import ( + FIDELITY, + apply_overrides, axis_overrides, build_neps_space, decode_config, ) -_SEARCH_SPACE = { - "judge.temperature": [0.0, 0.5, 1.0], - "prompt": [ - {"judge.prompt_preset": "meta-eval-pair-score", "judge.swap_mode": "fixed"}, - {"judge.prompt_preset": "meta-eval-pair-score", "judge.swap_mode": "both"}, - ], -} - -def _tuning_config(tmp_path, **tune_judge): +def _tuning_config(tmp_path, algorithm="neps_priorband", **settings): + optimizer = {"name": algorithm, "eta": 3} + objectives = ["agreement"] + if algorithm in {"mo_hyperband", "primo"}: + objectives.append("cost_per_1k_battles") + if algorithm == "primo": + optimizer.update( + initial_design_size=2, + prior_confidences={ + f"objective_{i}": {"judge.temperature": 0.25, "judge.model": 0.25} + for i in range(2) + }, + prior_centers={ + f"objective_{i}": { + "judge.temperature": value, + "judge.model": "Dummy/judge", + } + for i, value in enumerate((0.2, 0.8)) + }, + ) + specs = { + "judge.temperature": dict(type="Float", lower=0.0, upper=1.0), + "judge.model": dict(type="Categorical", choices=["Dummy/judge", "Dummy/other"]), + FIDELITY: dict(type="IntegerFidelity", lower=1, upper=3), + } + for key, value in [("judge.temperature", 0.2), ("judge.model", 1)]: + specs[key].update(prior=value, prior_confidence="medium") return RunConfig( task="tune-judge-comparia-fr", - judge={"model": "Dummy/judge", "temperature": 0.0}, + judge={"model": "Dummy/judge", "temperature": 0.2, "swap_mode": "both"}, run={"result_folder": str(tmp_path), "no_log_file": True}, tune_judge={ - "search_space": _SEARCH_SPACE, - "price_per_million_tokens": {"Dummy/judge": 1.0}, - "max_evaluations": 6, - "min_battles_per_model": 2, - "max_battles_per_model": 6, - **tune_judge, + "neps": { + "pipeline_space": specs, + "optimizer": optimizer, + "total_evaluations_to_spend": 10, + }, + "objectives": objectives, + "price_per_million_tokens": {"Dummy/judge": 1.0, "Dummy/other": 1.0}, + "run_dir": tmp_path / "tuning", + **settings, }, ) -def test_neps_space_uses_base_config_as_prior(tmp_path): - pytest.importorskip("neps") - cfg = _tuning_config(tmp_path) - base = cfg.model_dump(mode="json") - base["judge"]["prompt_preset"] = "meta-eval-pair-score" - space = build_neps_space(cfg.tune_judge, base) - - temperature, prompt = ( - space.searchables[k] for k in ("judge.temperature", "prompt") +@pytest.fixture +def token_counter(monkeypatch): + monkeypatch.setattr( + runner, + "_judge_token_encoding", + lambda: SimpleNamespace(encode=lambda text, **kwargs: text.split()), ) - assert json.loads(temperature.prior) == {"judge.temperature": 0.0} - assert json.loads(prompt.prior) == _SEARCH_SPACE["prompt"][0] - overrides, battles = decode_config( - { - "judge.temperature": temperature.choices[2], - "prompt": prompt.choices[1], - "battles_per_model": 6, - } - ) - assert overrides == {"judge.temperature": 1.0, **_SEARCH_SPACE["prompt"][1]} - assert battles == 6 - with pytest.raises(ValueError, match="tuner-owned"): - list( - axis_overrides( - _tuning_config(tmp_path).tune_judge.model_copy( - update={"search_space": {"meta_eval.split": ["test"]}} - ) - ) - ) def _fake_trial(config_path): cfg = load_config(config_path) - if cfg.judge.temperature == 1.0 and cfg.judge.swap_mode == "both": - raise subprocess.CalledProcessError(1, "judgearena") - run_dir = config_path.parent / "run" - run_dir.mkdir() + folder = config_path.parent / "run" + folder.mkdir() agreement = 0.6 - 0.1 * cfg.judge.temperature - metrics = { - "meta_eval_agreement": { - "all": { - "accuracy_attempted": agreement, - "cohen_kappa": agreement, - "coverage": 1.0, - "n_attempted": cfg.meta_eval.battles_per_model, - } - } - } - (run_dir / "results.json").write_text(json.dumps({"metrics": metrics})) + agreement_metrics = dict( + accuracy_attempted=agreement, + cohen_kappa=agreement, + coverage=1.0, + n_attempted=cfg.meta_eval.battles_per_model, + ) + metrics = {"meta_eval_agreement": {"all": agreement_metrics}} + (folder / "results.json").write_text(json.dumps({"metrics": metrics})) pd.DataFrame( - {"battle_id": ["b1"], "judge_input": ["prompt"], "judge_completion": ["x"]} - ).to_parquet(run_dir / "annotations.parquet") + { + "battle_id": ["b1", "b1"], + "judge_input": ["system forward", "system reverse"], + "judge_completion": ["answer", "answer"], + "cache_hit": [True, True], + } + ).to_parquet(folder / "annotations.parquet") -@pytest.mark.parametrize("algorithm", ["priorband", "mo_hyperband"]) -def test_run_tune_judge_searches_validation_and_scores_pick_on_test( - tmp_path, algorithm -): +@pytest.mark.parametrize( + "algorithm", ["neps_priorband", "neps_hyperband", "mo_hyperband", "primo"] +) +def test_native_optimizer_smoke(tmp_path, token_counter, algorithm): pytest.importorskip("neps") - cfg = _tuning_config(tmp_path, algorithm=algorithm) - test_results = run_tune_judge( - cfg, get_packaged_task(cfg.task), execute_trial=_fake_trial + cfg = _tuning_config(tmp_path, algorithm) + calls = [] + + def execute(path): + calls.append(path) + _fake_trial(path) + + task = get_packaged_task(cfg.task) + results = runner.run_tune_judge(cfg, task, execute_trial=execute) + assert not results.empty and set(results.status) == {"completed"} + if algorithm != "neps_priorband": + return + folder = cfg.tune_judge.run_dir + trials = pd.read_parquet(folder / "trials.parquet") + assert trials.battles_per_model.max() == 3 + assert trials.neps_trial_id.is_unique + best = trials[trials.battles_per_model == 3].groupby("judge_model").agreement.max() + assert results.set_index("judge_model").agreement.to_dict() == best.to_dict() + for path in calls: + trial = load_config(path) + assert trial.task == "meta-eval-comparia-fr" and trial.tune_judge is None + assert trial.meta_eval.split == ( + "test" if "test" in path.parts else "validation" + ) + assert load_config(calls[-1]).meta_eval.battles_per_model == 3 + before = len(calls) + (folder / "test_results.parquet").unlink() + resumed = runner.run_tune_judge(cfg, task, execute_trial=execute) + pd.testing.assert_frame_equal(results, resumed) + assert len(calls) == before # Completed search and held-out artifacts are reused. + helper = _tuning_config(tmp_path, search_only=True) + assert runner.run_tune_judge(helper, task, execute_trial=execute).empty + assert len(calls) == before + incompatible = _tuning_config( + tmp_path, price_per_million_tokens={"Dummy/judge": 2.0, "Dummy/other": 2.0} ) - - (tune_dir,) = tmp_path.glob("tune-*") - trials = pd.read_parquet(tune_dir / "trials.parquet") - assert len(trials) == 6 - assert set(trials["battles_per_model"]) <= {2, 6} - (pick,) = test_results.itertuples() - pick_cfg = load_config(tune_dir / "test" / pick.config_id / "config.yaml") - assert (pick_cfg.task, pick_cfg.meta_eval.split) == ( - "meta-eval-comparia-fr", - "test", + with pytest.raises(ValueError, match="Incompatible configuration"): + runner.run_tune_judge(incompatible, task, execute_trial=execute) + assert len(calls) == before + fresh = _tuning_config(tmp_path, run_dir=tmp_path / "fresh") + runner.run_tune_judge(fresh, task, execute_trial=execute) + repeated = pd.read_parquet(fresh.tune_judge.run_dir / "trials.parquet") + assert repeated.overrides.tolist() == trials.overrides.tolist() + + +def test_dotted_domains_and_index_priors(tmp_path): + neps = pytest.importorskip("neps") + specs = _tuning_config(tmp_path).tune_judge.neps["pipeline_space"] + specs["judge.max_out_tokens"] = dict(type="Integer", lower=16, upper=32) + attrs = build_neps_space(specs).get_attrs() + assert attrs["judge.model"].prior == 1 + assert attrs["judge.model"].choices == ("Dummy/judge", "Dummy/other") + assert isinstance(attrs["judge.temperature"], neps.Float) + assert attrs["judge.temperature"].prior == 0.2 + assert isinstance(attrs["judge.max_out_tokens"], neps.Integer) + assert isinstance(attrs[FIDELITY], neps.IntegerFidelity) + assert {"judge.temperature": 1.0} in list(axis_overrides(specs)) + overrides, battles = decode_config( + {"judge.temperature": 0.37, "judge.model": "Dummy/other", FIDELITY: 3} + ) + base = {"judge": {"temperature": 0.2}} + assert apply_overrides(base, overrides)["judge"] == { + "temperature": 0.37, + "model": "Dummy/other", + } + assert base["judge"]["temperature"] == 0.2 and battles == 3 + for name, spec in [ + ("meta_eval.split", {"type": "Categorical", "choices": ["test"]}), + ("judge.temperature", {"type": "IntegerFidelity", "lower": 1, "upper": 3}), + ]: + with pytest.raises(ValueError, match="tuner-owned|Only"): + build_neps_space({**specs, name: spec}) + + +def test_cost_counts_both_cached_orientations(tmp_path, token_counter): + path = tmp_path / "config.yaml" + dump_config(_tuning_config(tmp_path), path) + _fake_trial(path) + record = runner._read_trial(tmp_path, price_per_million_tokens=2.0) + assert record["tokens_per_battle"] == 6 + assert record["cost_per_1k_battles"] == pytest.approx(0.012) + objective = runner._objective( + {**record, "status": "completed"}, ["agreement", "cost_per_1k_battles"] ) - for trial_dir in (tune_dir / "trials").iterdir(): - trial_cfg = load_config(trial_dir / "config.yaml") - assert trial_cfg.task == "meta-eval-comparia-fr" - assert trial_cfg.meta_eval.split == "validation" - assert trial_cfg.tune_judge is None + assert objective["objective_to_minimize"] == pytest.approx([0.42, 0.012]) + assert "cost" not in objective + + +@pytest.mark.parametrize( + "change,match", + [ + ({"neps": {"total_cost_to_spend": 1}}, "managed"), + ({"neps": {"root_directory": "elsewhere"}}, "managed"), + ({"neps": {}}, "finite global budget"), + ({"search_only": True, "run_dir": None}, "explicit"), + ({"neps": {"worker_evaluations_to_spend": 1}}, "Worker-local"), + ({"objectives": ["agreement", "cost_per_1k_battles"]}, "single objective"), + ({"neps": {"optimizer": {"name": "mo_hyperband"}}}, "agreement and cost"), + ({"price_per_million_tokens": {"Dummy/judge": -1}}, "nonnegative"), + ], +) +def test_tuning_config_validation(tmp_path, change, match): + values = _tuning_config(tmp_path).model_dump() + if "neps" in change and change["neps"]: + change = {**change, "neps": {**values["tune_judge"]["neps"], **change["neps"]}} + values["tune_judge"].update(change) + with pytest.raises(ValidationError, match=match): + RunConfig(**values) -def test_tune_judge_requires_prices_and_tune_task(tmp_path): +def test_prices_required_before_execution(tmp_path): + pytest.importorskip("neps") cfg = _tuning_config(tmp_path, price_per_million_tokens={}) with pytest.raises(ValueError, match="missing models"): - run_tune_judge(cfg, get_packaged_task(cfg.task), execute_trial=_fake_trial) - with pytest.raises(ValidationError, match="tune-judge tasks"): - RunConfig( - task="meta-eval-comparia-fr", - judge={"model": "Dummy/judge"}, - tune_judge=cfg.tune_judge.model_dump(), + runner.run_tune_judge( + cfg, + get_packaged_task(cfg.task), + execute_trial=lambda _: pytest.fail("executed without prices"), ) - with pytest.raises(ValidationError, match="tune-judge tasks"): - RunConfig(task="tune-judge-comparia-fr", judge={"model": "Dummy/judge"}) From 141457787c3feced64db5535eac59ab15e89095a Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 7 Oct 2026 00:58:58 +0200 Subject: [PATCH 07/20] Document native NePS search configuration and worker usage --- README.md | 62 +++++++++++++++++++++++++++-------- configs/tune_judge.yaml | 71 ++++++++++++++++++++++++++--------------- 2 files changed, 93 insertions(+), 40 deletions(-) diff --git a/README.md b/README.md index 03c0332d..0685ea67 100644 --- a/README.md +++ b/README.md @@ -336,26 +336,60 @@ Runs save the selected sample, judge evidence, metric battles, configuration, an ### Tuning a judge -The `tune-judge-*` tasks search judge configurations against the same human battles as their `meta-eval-*` counterparts, following [Salinas et al., 2025](https://arxiv.org/abs/2501.17178). They need the `tune` extra (`pip install 'judgearena[tune]'`), which installs [neps](https://github.com/automl/neps). +The `tune-judge-*` tasks search judge configurations against the same human battles as their `meta-eval-*` counterparts, following [Salinas et al., 2025](https://arxiv.org/abs/2501.17178). Install `pip install 'judgearena[tune]'` for [NePS](https://github.com/automl/neps) 0.17 (`neural-pipeline-search`). -- Each entry in `search_space` is a dotted config path with a list of choices or a `{lower, upper}` range. Choices that are mappings change several settings together, e.g. `alpaca-eval` together with `judge.top_logprobs`. -- A neps optimizer proposes each trial. The fidelity is validation battles per model, between `min_battles_per_model` and `max_battles_per_model`. - - `algorithm: priorband` (the default) uses the base config as a prior, so a known-good judge setup speeds up the search. - - `hyperband` samples uniformly. - - `mo_hyperband` also minimizes judge cost. -- Every trial runs as a meta-eval subprocess of the matching `meta-eval-*` task on the validation split, so tuning and meta-eval share the judgement cache. -- Cost is judge tokens (input plus output, counted with tiktoken) times `price_per_million_tokens[judge.model]`. -- At the end, the best configuration for each judge model at the highest fidelity is scored once on the held-out test split. +[configs/tune_judge.yaml](configs/tune_judge.yaml) uses native NePS settings under `tune_judge.neps`: `pipeline_space` has dotted config names, a `type` (`Categorical`, `Float`, `Integer`, or `IntegerFidelity`), and native constructor kwargs such as `choices`, `lower`, `upper`, `prior`, `prior_confidence`, and `log` where supported. Categorical `prior` is a choice **index**. Every parameter with a `prior` must also set `prior_confidence` (the example uses `medium`). Prompt presets are scalar choices (`meta-eval-pair-score`, `arena-hard`, `default_with_explanation`); grouped overrides are unsupported. The example searches temperature continuously and uses `meta_eval.battles_per_model` as the validation fidelity. + +For example, a continuous temperature parameter under `tune_judge.neps.pipeline_space`: + +```yaml +judge.temperature: + type: Float + lower: 0.0 + upper: 1.0 + prior: 0.0 + prior_confidence: medium +``` + +- `optimizer: {name: neps_priorband, eta: 3}` uses parameter priors; `neps_hyperband` samples uniformly. Both require `objectives: [agreement]` outside `neps`. +- `optimizer: {name: mo_hyperband, eta: 3, mo_selector: epsnet}` and `primo` support `objectives: [agreement, cost_per_1k_battles]`. Agreement is maximized (NePS minimizes `1 - agreement`); cost is minimized. +- For `primo`, put native `prior_centers` settings in the optimizer mapping, keyed by `objective_0` (disagreement) and `objective_1` (cost). Each contains dotted parameter names mapped to native setting **values**, for example: + +```yaml +optimizer: + name: primo + eta: 3 + mo_selector: epsnet + prior_centers: + objective_0: + judge.model: VLLM/Qwen/Qwen2.5-7B-Instruct + judge.prompt_preset: meta-eval-pair-score + judge.temperature: 0.0 + judge.max_out_tokens: 2048 + judge.swap_mode: both + objective_1: + judge.model: VLLM/Qwen/Qwen2.5-7B-Instruct + judge.prompt_preset: meta-eval-pair-score + judge.temperature: 0.0 + judge.max_out_tokens: 1024 + judge.swap_mode: fixed +``` + +Each trial runs the matching `meta-eval-*` task on the validation split and shares its judgement cache. Cost uses input plus output tokens (counted with tiktoken), including cached outputs, and `price_per_million_tokens[judge.model]` for every searched model. The Qwen/Gemma prices are illustrative estimates; this objective estimates USD per 1,000 battles, not actual API spend. NePS cost spend limits are unsupported and rejected; use `neps.total_evaluations_to_spend` (40 in the example). + +The main process waits for active trials after `neps.run` returns under normal budget accounting, then selects the best agreement configuration per judge model at the highest fidelity for held-out testing. The native NePS completion flag is not a worker barrier. Optional `test_battles_per_model` defaults to the highest fidelity. + +With NePS 0.17, leave `sample_prior_first` at its default for `neps_priorband` and `neps_hyperband`: enabling it fails in upstream trial-ID parsing. Priors still guide sampling without that option. For `primo`, its initial design must reach the highest fidelity before it switches to Bayesian optimization. ```bash judgearena --config_path configs/tune_judge.yaml +# Resume with the same explicit tuning directory and settings. +judgearena --config_path configs/tune_judge.yaml --tune_judge.run_dir /path/to/tune-dir +# Helper workers join that directory and run only the search. +judgearena --config_path configs/tune_judge.yaml --tune_judge.run_dir /path/to/tune-dir --tune_judge.search_only true ``` -Outputs go to `tune--/` under `--run.result_folder`: - -- `trials.parquet`: one row per neps trial. -- `test_results.parquet`: test-split results for the selected configurations. -- One meta-eval run directory per trial. +`run_dir: null` creates a tuning directory under `--run.result_folder`; set an explicit `tune_judge.run_dir` to resume or share workers. NePS state lives in `/neps/`, with meta-eval trial artifacts separately in `/trials/`. `trials.parquet` summarizes trials and `test_results.parquet` holds the selected configurations' test results. For multiple workers on the same node, share a node-local `store_root` for SQLite judgement caching; multi-node NFS cache sharing is not supported. ## 📈 Estimating ELO Ratings diff --git a/configs/tune_judge.yaml b/configs/tune_judge.yaml index 63341bc3..31c55945 100644 --- a/configs/tune_judge.yaml +++ b/configs/tune_judge.yaml @@ -1,38 +1,57 @@ -# Example judge-tuning run (Salinas et al., ICML 2025) with neps PriorBand. -# Run with: pip install 'judgearena[tune]' -# judgearena --config_path configs/tune_judge.yaml -# Each configuration runs as a meta-eval-lmarena-140k-en trial on the -# validation split; the settings below (Qwen2.5-7B, meta-eval-pair-score, -# temperature 0, fixed swap) are the prior. The best configuration per judge -# model is then scored on the test split. +# NePS 0.17 judge-tuning example; install with pip install 'judgearena[tune]'. +# Run: judgearena --config_path configs/tune_judge.yaml +# Validation trials use the matching meta-eval task and its judgement cache. task: tune-judge-lmarena-140k-en judge: model: VLLM/Qwen/Qwen2.5-7B-Instruct + prompt_preset: meta-eval-pair-score temperature: 0.0 max_out_tokens: 2048 max_model_len: 16384 meta_eval: top_models: 20 tune_judge: - search_space: - judge.model: - - VLLM/Qwen/Qwen2.5-7B-Instruct - - VLLM/google/gemma-2-9b-it - # Grouped axis: settings in one entry change together. - prompt: - - {judge.prompt_preset: meta-eval-pair-score} - - {judge.prompt_preset: arena-hard} - - {judge.prompt_preset: default_with_explanation} - - {judge.prompt_preset: alpaca-eval, judge.top_logprobs: 20, judge.max_out_tokens: 1} - judge.temperature: [0.0, 0.1, 1.0] - judge.swap_mode: [fixed, both] - algorithm: priorband - # One hyperband iteration with eta 3 over 10-90 battles is 20 trials. - max_evaluations: 40 - min_battles_per_model: 10 - max_battles_per_model: 90 - # USD per million tokens, from judgetuning's llm_specs.py (measured vLLM - # throughput). Every searched judge.model needs a price. + neps: + pipeline_space: + judge.model: + type: Categorical + choices: [VLLM/Qwen/Qwen2.5-7B-Instruct, VLLM/google/gemma-2-9b-it] + prior: 0 # Categorical priors are choice indices. + prior_confidence: medium + judge.prompt_preset: + type: Categorical + choices: [meta-eval-pair-score, arena-hard, default_with_explanation] + prior: 0 + prior_confidence: medium + judge.temperature: + type: Float + lower: 0.0 + upper: 1.0 + prior: 0.0 + prior_confidence: medium + judge.max_out_tokens: + type: Integer + lower: 1024 + upper: 2048 + prior: 2048 + prior_confidence: medium + judge.swap_mode: + type: Categorical + choices: [fixed, both] + prior: 0 + prior_confidence: medium + meta_eval.battles_per_model: + type: IntegerFidelity + lower: 10 + upper: 90 + optimizer: {name: mo_hyperband, eta: 3, mo_selector: epsnet} + total_evaluations_to_spend: 40 + objectives: [agreement, cost_per_1k_battles] + # Estimated USD per million tokens, not measured spend or provider quotes. price_per_million_tokens: VLLM/Qwen/Qwen2.5-7B-Instruct: 0.12 VLLM/google/gemma-2-9b-it: 0.14 + # Optional; null uses the highest validation fidelity for held-out testing. + test_battles_per_model: null + run_dir: null + search_only: false From 871df16342570cce55cf1bae2ffc9141a449e081 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:12:56 +0200 Subject: [PATCH 08/20] Drain running NePS evaluations without waiting for unused samples --- README.md | 2 ++ judgearena/tuning/session.py | 6 ++---- tests/test_tuning.py | 34 ++++++++++++++++++++++++++++++++++ 3 files changed, 38 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 0685ea67..641972f1 100644 --- a/README.md +++ b/README.md @@ -391,6 +391,8 @@ judgearena --config_path configs/tune_judge.yaml --tune_judge.run_dir /path/to/t `run_dir: null` creates a tuning directory under `--run.result_folder`; set an explicit `tune_judge.run_dir` to resume or share workers. NePS state lives in `/neps/`, with meta-eval trial artifacts separately in `/trials/`. `trials.parquet` summarizes trials and `test_results.parquet` holds the selected configurations' test results. For multiple workers on the same node, share a node-local `store_root` for SQLite judgement caching; multi-node NFS cache sharing is not supported. +The primary process can extend the global evaluation or fidelity budget when resuming. Helpers inherit the persisted global limits. Search-space, optimizer, objective, and pricing changes require a new run directory. + ## 📈 Estimating ELO Ratings JudgeArena can estimate the ELO rating of a model by running it against opponents sampled from a human preference arena (`LMArena-100k`, `LMArena-140k`, or `ComparIA`). diff --git a/judgearena/tuning/session.py b/judgearena/tuning/session.py index d6c6ae02..931d388f 100644 --- a/judgearena/tuning/session.py +++ b/judgearena/tuning/session.py @@ -24,7 +24,7 @@ def _config_hash(cfg: RunConfig) -> str: tuning["neps"] = { key: value for key, value in tuning["neps"].items() - if not key.startswith("worker_") + if not key.endswith("_to_spend") and key != "worker_id" } if cfg.judge.prompt is not None: payload["prompt_contents"] = [ @@ -76,9 +76,7 @@ def collect_trials(root: Path, *, wait: bool) -> pd.DataFrame: state = NePSState.create_or_load(root, load_only=True) while True: trials = state.lock_and_read_trials() - active = [ - t for t in trials.values() if t.metadata.state in {"pending", "evaluating"} - ] + active = [t for t in trials.values() if t.metadata.state == "evaluating"] if not wait or not active: break time.sleep(1) diff --git a/tests/test_tuning.py b/tests/test_tuning.py index 6c1f9d18..c93f44df 100644 --- a/tests/test_tuning.py +++ b/tests/test_tuning.py @@ -146,6 +146,19 @@ def execute(path): runner.run_tune_judge(fresh, task, execute_trial=execute) repeated = pd.read_parquet(fresh.tune_judge.run_dir / "trials.parquet") assert repeated.overrides.tolist() == trials.overrides.tolist() + extended = _tuning_config(tmp_path) + extended.tune_judge.neps["total_evaluations_to_spend"] = 11 + runner.run_tune_judge(extended, task, execute_trial=execute) + assert len(pd.read_parquet(folder / "trials.parquet")) == 11 + runner.run_tune_judge(helper, task, execute_trial=execute) + from neps.state import NePSState + + assert ( + NePSState.create_or_load( + folder / "neps", load_only=True + ).lock_and_get_global_budgets()[0] + == 11 + ) def test_dotted_domains_and_index_priors(tmp_path): @@ -177,6 +190,27 @@ def test_dotted_domains_and_index_priors(tmp_path): build_neps_space({**specs, name: spec}) +def test_budget_end_does_not_wait_for_unstarted_batch( + tmp_path, token_counter, monkeypatch +): + pytest.importorskip("neps") + cfg = _tuning_config(tmp_path) + cfg.tune_judge.neps.update( + optimizer={"name": "neps_random_search", "ignore_fidelity": "highest_fidelity"}, + total_evaluations_to_spend=2, + sample_batch_size=3, + ) + monkeypatch.setattr( + "judgearena.tuning.session.time.sleep", + lambda _: pytest.fail("Waiting for an unstarted trial"), + ) + results = runner.run_tune_judge( + cfg, get_packaged_task(cfg.task), execute_trial=_fake_trial + ) + assert not results.empty + assert len(pd.read_parquet(cfg.tune_judge.run_dir / "trials.parquet")) == 2 + + def test_cost_counts_both_cached_orientations(tmp_path, token_counter): path = tmp_path / "config.yaml" dump_config(_tuning_config(tmp_path), path) From b95b6600cc2933e06997a994840ade506d8b2bd8 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 7 Oct 2026 01:15:30 +0200 Subject: [PATCH 09/20] Remove obsolete meta-eval tuning dispatch test setup --- tests/test_generate_and_evaluate.py | 2 +- tests/test_usage.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_generate_and_evaluate.py b/tests/test_generate_and_evaluate.py index 9988ec92..b790fe67 100644 --- a/tests/test_generate_and_evaluate.py +++ b/tests/test_generate_and_evaluate.py @@ -218,7 +218,7 @@ def test_benchmark_dispatch_passes_the_resolved_task(monkeypatch): ) monkeypatch.setattr(benchmark_registry, "benchmark_adapters", lambda: (pairwise,)) monkeypatch.setattr(benchmark_registry, "get_packaged_task", lambda _task: resolved) - cfg = SimpleNamespace(task="yaml-task", tune_judge=None) + cfg = SimpleNamespace(task="yaml-task") result = benchmark_runner.run_benchmark(cfg) diff --git a/tests/test_usage.py b/tests/test_usage.py index 12e84a9f..51e7a0a7 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -77,7 +77,7 @@ def fake_runner(cfg, task): adapter=SimpleNamespace(name="test", runner=fake_runner), task=None ) monkeypatch.setattr(benchmark_runner, "resolve_benchmark", lambda task: resolved) - path = benchmark_runner.run_benchmark(SimpleNamespace(task="test", tune_judge=None)) + path = benchmark_runner.run_benchmark(SimpleNamespace(task="test")) usage = json.loads(path.read_text())["usage"] assert usage["total"]["requests"] == 2 assert usage["total"]["cost_usd"] == pytest.approx(0.002) @@ -99,7 +99,7 @@ def fake_runner(cfg, task): ) monkeypatch.setattr(benchmark_runner, "resolve_benchmark", lambda task: resolved) with pytest.raises(ValueError, match="backend failed"): - benchmark_runner.run_benchmark(SimpleNamespace(task="test", tune_judge=None)) + benchmark_runner.run_benchmark(SimpleNamespace(task="test")) assert current_run_usage() is None assert "No successful model responses" in capsys.readouterr().out From f5d864d2c14ea527fd2e765bd3bc08cc04611578 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 7 Oct 2026 09:24:26 +0200 Subject: [PATCH 10/20] Require OpenAI SDK types imported by vLLM 0.25 --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index a1887c56..81228bf8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -111,6 +111,7 @@ indent-style = "space" vllm = [ "vllm>=0.19.1,<1.0.0", "transformers>=5.5.1,<6.0.0", + "openai>=2.30.0", # vLLM 0.25 imports responses.NamespaceTool. ] llamacpp = ["llama-cpp-python>=0.3.0"] tune = [ From 99300e9df1adc3ce1f18dba1e8a2f6d3fb78b040 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 7 Oct 2026 15:13:09 +0200 Subject: [PATCH 11/20] Separate tuning execution from configuration contracts and usage docs --- README.md | 59 --------------------------------------- configs/tune_judge.yaml | 57 ------------------------------------- tests/test_tuning.py | 62 ++--------------------------------------- 3 files changed, 2 insertions(+), 176 deletions(-) delete mode 100644 configs/tune_judge.yaml diff --git a/README.md b/README.md index d4ca6ae0..2367ca42 100644 --- a/README.md +++ b/README.md @@ -339,65 +339,6 @@ Runs save the selected sample, judge evidence, metric battles, configuration, an `--meta_eval.split validation|test` restricts sampling to one half of the arena prompts (`--meta_eval.validation_fraction`, default 0.5). The split hashes each battle's prompt, so it does not depend on `--run.seed`. -### Tuning a judge - -The `tune-judge-*` tasks search judge configurations against the same human battles as their `meta-eval-*` counterparts, following [Salinas et al., 2025](https://arxiv.org/abs/2501.17178). Install `pip install 'judgearena[tune]'` for [NePS](https://github.com/automl/neps) 0.17 (`neural-pipeline-search`). - -[configs/tune_judge.yaml](configs/tune_judge.yaml) uses native NePS settings under `tune_judge.neps`: `pipeline_space` has dotted config names, a `type` (`Categorical`, `Float`, `Integer`, or `IntegerFidelity`), and native constructor kwargs such as `choices`, `lower`, `upper`, `prior`, `prior_confidence`, and `log` where supported. Categorical `prior` is a choice **index**. Every parameter with a `prior` must also set `prior_confidence` (the example uses `medium`). Prompt presets are scalar choices (`meta-eval-pair-score`, `arena-hard`, `default_with_explanation`); grouped overrides are unsupported. The example searches temperature continuously and uses `meta_eval.battles_per_model` as the validation fidelity. - -For example, a continuous temperature parameter under `tune_judge.neps.pipeline_space`: - -```yaml -judge.temperature: - type: Float - lower: 0.0 - upper: 1.0 - prior: 0.0 - prior_confidence: medium -``` - -- `optimizer: {name: neps_priorband, eta: 3}` uses parameter priors; `neps_hyperband` samples uniformly. Both require `objectives: [agreement]` outside `neps`. -- `optimizer: {name: mo_hyperband, eta: 3, mo_selector: epsnet}` and `primo` support `objectives: [agreement, cost_per_1k_battles]`. Agreement is maximized (NePS minimizes `1 - agreement`); cost is minimized. -- For `primo`, put native `prior_centers` settings in the optimizer mapping, keyed by `objective_0` (disagreement) and `objective_1` (cost). Each contains dotted parameter names mapped to native setting **values**, for example: - -```yaml -optimizer: - name: primo - eta: 3 - mo_selector: epsnet - prior_centers: - objective_0: - judge.model: VLLM/Qwen/Qwen2.5-7B-Instruct - judge.prompt_preset: meta-eval-pair-score - judge.temperature: 0.0 - judge.max_out_tokens: 2048 - judge.swap_mode: both - objective_1: - judge.model: VLLM/Qwen/Qwen2.5-7B-Instruct - judge.prompt_preset: meta-eval-pair-score - judge.temperature: 0.0 - judge.max_out_tokens: 1024 - judge.swap_mode: fixed -``` - -Each trial runs the matching `meta-eval-*` task on the validation split and shares its judgement cache. Cost uses input plus output tokens (counted with tiktoken), including cached outputs, and `price_per_million_tokens[judge.model]` for every searched model. The Qwen/Gemma prices are illustrative estimates; this objective estimates USD per 1,000 battles, not actual API spend. NePS cost spend limits are unsupported and rejected; use `neps.total_evaluations_to_spend` (40 in the example). - -The main process waits for active trials after `neps.run` returns under normal budget accounting, then selects the best agreement configuration per judge model at the highest fidelity for held-out testing. The native NePS completion flag is not a worker barrier. Optional `test_battles_per_model` defaults to the highest fidelity. - -With NePS 0.17, leave `sample_prior_first` at its default for `neps_priorband` and `neps_hyperband`: enabling it fails in upstream trial-ID parsing. Priors still guide sampling without that option. For `primo`, its initial design must reach the highest fidelity before it switches to Bayesian optimization. - -```bash -judgearena --config_path configs/tune_judge.yaml -# Resume with the same explicit tuning directory and settings. -judgearena --config_path configs/tune_judge.yaml --tune_judge.run_dir /path/to/tune-dir -# Helper workers join that directory and run only the search. -judgearena --config_path configs/tune_judge.yaml --tune_judge.run_dir /path/to/tune-dir --tune_judge.search_only true -``` - -`run_dir: null` creates a tuning directory under `--run.result_folder`; set an explicit `tune_judge.run_dir` to resume or share workers. NePS state lives in `/neps/`, with meta-eval trial artifacts separately in `/trials/`. `trials.parquet` summarizes trials and `test_results.parquet` holds the selected configurations' test results. For multiple workers on the same node, share a node-local `store_root` for SQLite judgement caching; multi-node NFS cache sharing is not supported. - -The primary process can extend the global evaluation or fidelity budget when resuming. Helpers inherit the persisted global limits. Search-space, optimizer, objective, and pricing changes require a new run directory. - ## 📈 Estimating ELO Ratings JudgeArena can estimate the ELO rating of a model by running it against opponents sampled from a human preference arena (`LMArena-100k`, `LMArena-140k`, or `ComparIA`). diff --git a/configs/tune_judge.yaml b/configs/tune_judge.yaml deleted file mode 100644 index 31c55945..00000000 --- a/configs/tune_judge.yaml +++ /dev/null @@ -1,57 +0,0 @@ -# NePS 0.17 judge-tuning example; install with pip install 'judgearena[tune]'. -# Run: judgearena --config_path configs/tune_judge.yaml -# Validation trials use the matching meta-eval task and its judgement cache. -task: tune-judge-lmarena-140k-en -judge: - model: VLLM/Qwen/Qwen2.5-7B-Instruct - prompt_preset: meta-eval-pair-score - temperature: 0.0 - max_out_tokens: 2048 - max_model_len: 16384 -meta_eval: - top_models: 20 -tune_judge: - neps: - pipeline_space: - judge.model: - type: Categorical - choices: [VLLM/Qwen/Qwen2.5-7B-Instruct, VLLM/google/gemma-2-9b-it] - prior: 0 # Categorical priors are choice indices. - prior_confidence: medium - judge.prompt_preset: - type: Categorical - choices: [meta-eval-pair-score, arena-hard, default_with_explanation] - prior: 0 - prior_confidence: medium - judge.temperature: - type: Float - lower: 0.0 - upper: 1.0 - prior: 0.0 - prior_confidence: medium - judge.max_out_tokens: - type: Integer - lower: 1024 - upper: 2048 - prior: 2048 - prior_confidence: medium - judge.swap_mode: - type: Categorical - choices: [fixed, both] - prior: 0 - prior_confidence: medium - meta_eval.battles_per_model: - type: IntegerFidelity - lower: 10 - upper: 90 - optimizer: {name: mo_hyperband, eta: 3, mo_selector: epsnet} - total_evaluations_to_spend: 40 - objectives: [agreement, cost_per_1k_battles] - # Estimated USD per million tokens, not measured spend or provider quotes. - price_per_million_tokens: - VLLM/Qwen/Qwen2.5-7B-Instruct: 0.12 - VLLM/google/gemma-2-9b-it: 0.14 - # Optional; null uses the highest validation fidelity for held-out testing. - test_battles_per_model: null - run_dir: null - search_only: false diff --git a/tests/test_tuning.py b/tests/test_tuning.py index c93f44df..73c44811 100644 --- a/tests/test_tuning.py +++ b/tests/test_tuning.py @@ -1,4 +1,4 @@ -"""Contracts for native NePS judge tuning and persisted trial artifacts.""" +"""Native NePS execution and persisted trial artifacts.""" from __future__ import annotations @@ -7,18 +7,11 @@ import pandas as pd import pytest -from pydantic import ValidationError from judgearena.config import RunConfig, dump_config, load_config from judgearena.tasks.registry import get_packaged_task from judgearena.tuning import runner -from judgearena.tuning.search_space import ( - FIDELITY, - apply_overrides, - axis_overrides, - build_neps_space, - decode_config, -) +from judgearena.tuning.search_space import FIDELITY def _tuning_config(tmp_path, algorithm="neps_priorband", **settings): @@ -161,35 +154,6 @@ def execute(path): ) -def test_dotted_domains_and_index_priors(tmp_path): - neps = pytest.importorskip("neps") - specs = _tuning_config(tmp_path).tune_judge.neps["pipeline_space"] - specs["judge.max_out_tokens"] = dict(type="Integer", lower=16, upper=32) - attrs = build_neps_space(specs).get_attrs() - assert attrs["judge.model"].prior == 1 - assert attrs["judge.model"].choices == ("Dummy/judge", "Dummy/other") - assert isinstance(attrs["judge.temperature"], neps.Float) - assert attrs["judge.temperature"].prior == 0.2 - assert isinstance(attrs["judge.max_out_tokens"], neps.Integer) - assert isinstance(attrs[FIDELITY], neps.IntegerFidelity) - assert {"judge.temperature": 1.0} in list(axis_overrides(specs)) - overrides, battles = decode_config( - {"judge.temperature": 0.37, "judge.model": "Dummy/other", FIDELITY: 3} - ) - base = {"judge": {"temperature": 0.2}} - assert apply_overrides(base, overrides)["judge"] == { - "temperature": 0.37, - "model": "Dummy/other", - } - assert base["judge"]["temperature"] == 0.2 and battles == 3 - for name, spec in [ - ("meta_eval.split", {"type": "Categorical", "choices": ["test"]}), - ("judge.temperature", {"type": "IntegerFidelity", "lower": 1, "upper": 3}), - ]: - with pytest.raises(ValueError, match="tuner-owned|Only"): - build_neps_space({**specs, name: spec}) - - def test_budget_end_does_not_wait_for_unstarted_batch( tmp_path, token_counter, monkeypatch ): @@ -225,28 +189,6 @@ def test_cost_counts_both_cached_orientations(tmp_path, token_counter): assert "cost" not in objective -@pytest.mark.parametrize( - "change,match", - [ - ({"neps": {"total_cost_to_spend": 1}}, "managed"), - ({"neps": {"root_directory": "elsewhere"}}, "managed"), - ({"neps": {}}, "finite global budget"), - ({"search_only": True, "run_dir": None}, "explicit"), - ({"neps": {"worker_evaluations_to_spend": 1}}, "Worker-local"), - ({"objectives": ["agreement", "cost_per_1k_battles"]}, "single objective"), - ({"neps": {"optimizer": {"name": "mo_hyperband"}}}, "agreement and cost"), - ({"price_per_million_tokens": {"Dummy/judge": -1}}, "nonnegative"), - ], -) -def test_tuning_config_validation(tmp_path, change, match): - values = _tuning_config(tmp_path).model_dump() - if "neps" in change and change["neps"]: - change = {**change, "neps": {**values["tune_judge"]["neps"], **change["neps"]}} - values["tune_judge"].update(change) - with pytest.raises(ValidationError, match=match): - RunConfig(**values) - - def test_prices_required_before_execution(tmp_path): pytest.importorskip("neps") cfg = _tuning_config(tmp_path, price_per_million_tokens={}) From 83ed0c2a26523dd375753c035cfa388288dd3ab7 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 7 Oct 2026 15:48:12 +0200 Subject: [PATCH 12/20] Report judge cost Pareto front and continue multi-objective trial failures --- .github/workflows/ci.yml | 18 ++++++++++-- judgearena/tuning/runner.py | 13 +++++++++ tests/test_tuning.py | 57 +++++++++++++++++++++++++++++++++++++ 3 files changed, 85 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 5c7e6977..8f3119be 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -7,7 +7,6 @@ on: push: branches: [ "main" ] pull_request: - branches: [ "main" ] permissions: contents: read @@ -24,8 +23,21 @@ jobs: enable-cache: true python-version: "3.12" - name: Install dependencies - run: uv sync --extra tune + run: uv sync - name: Run pre-commit run: uv run pre-commit run --all-files - name: Test with pytest - run: uv run --extra tune pytest + run: uv run pytest + + tuning: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v3 + - uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + python-version: "3.12" + - name: Install tuning dependencies + run: uv sync --extra tune + - name: Test tuning integration + run: uv run --extra tune pytest tests/test_tuning.py tests/test_tuning_config.py diff --git a/judgearena/tuning/runner.py b/judgearena/tuning/runner.py index 334030b1..bd97e1f8 100644 --- a/judgearena/tuning/runner.py +++ b/judgearena/tuning/runner.py @@ -146,6 +146,11 @@ def _objective(record: dict, objectives: list[str]) -> dict: if record["status"] != "completed": return { "exception": RuntimeError(f"Trial {record['config_id']} failed"), + "objective_to_minimize": ( + float("inf") + if len(objectives) == 1 + else [float("inf")] * len(objectives) + ), "info_dict": record, } values = [ @@ -264,6 +269,12 @@ def _evaluate_picks(runner, tuning, tune_dir, trials, max_battles) -> pd.DataFra ranked = final.sort_values( ["agreement", "cost_per_1k_battles", "config_id"], ascending=[False, True, True] ) + from neps.optimizers.utils.multiobjective.epsnet import pareto_efficient + + costs = ranked[["agreement", "cost_per_1k_battles"]].to_numpy().copy() + costs[:, 0] *= -1 + pareto = ranked.loc[pareto_efficient(costs)] + pareto.to_parquet(tune_dir / "pareto.parquet", index=False) picks = ranked.groupby("judge_model", sort=False).head(1) test_results = pd.DataFrame( [ @@ -280,6 +291,8 @@ def _evaluate_picks(runner, tuning, tune_dir, trials, max_battles) -> pd.DataFra print("\n=== Judge tuning ===") print("Validation:") print(ranked[_SUMMARY_COLUMNS].to_string(index=False)) + print("Agreement/cost Pareto front (validation):") + print(pareto[_SUMMARY_COLUMNS].to_string(index=False)) print("Test:") print( test_results.reindex(columns=[*_SUMMARY_COLUMNS, "status"]).to_string( diff --git a/tests/test_tuning.py b/tests/test_tuning.py index 73c44811..bef63aed 100644 --- a/tests/test_tuning.py +++ b/tests/test_tuning.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import subprocess from types import SimpleNamespace import pandas as pd @@ -198,3 +199,59 @@ def test_prices_required_before_execution(tmp_path): get_packaged_task(cfg.task), execute_trial=lambda _: pytest.fail("executed without prices"), ) + + +def test_failed_trial_does_not_stop_search(tmp_path, token_counter): + pytest.importorskip("neps") + cfg = _tuning_config(tmp_path, "mo_hyperband") + cfg.tune_judge.neps["ignore_errors"] = True + calls = [] + + def execute(path): + calls.append(path) + if len(calls) == 1: + raise subprocess.CalledProcessError(1, ["judge-trial"]) + _fake_trial(path) + + results = runner.run_tune_judge( + cfg, get_packaged_task(cfg.task), execute_trial=execute + ) + trials = pd.read_parquet(cfg.tune_judge.run_dir / "trials.parquet") + assert len(trials) == 10 + assert (trials.status == "failed").sum() == 1 + assert not results.empty and set(results.status) == {"completed"} + pareto = pd.read_parquet(cfg.tune_judge.run_dir / "pareto.parquet") + final = trials[(trials.status == "completed") & (trials.battles_per_model == 3)] + assert pareto.agreement.tolist() == [final.agreement.max()] + assert set(pareto.status) == {"completed"} + assert set(pareto.battles_per_model) == {3} + + +def test_pareto_report_uses_completed_full_fidelity_validation(tmp_path): + pytest.importorskip("neps") + trials = pd.DataFrame( + [ + dict( + config_id=name, + judge_model=name, + agreement=agreement, + cost_per_1k_battles=cost, + battles_per_model=fidelity, + status=status, + overrides="{}", + ) + for name, agreement, cost, fidelity, status in [ + ("cheap", 0.5, 1.0, 3, "completed"), + ("accurate", 0.8, 2.0, 3, "completed"), + ("dominated", 0.4, 3.0, 3, "completed"), + ("low-fidelity", 0.9, 0.1, 1, "completed"), + ("failed", 1.0, 0.0, 3, "failed"), + ] + ] + ) + executor = SimpleNamespace(run=lambda *args, **kwargs: {"status": "completed"}) + runner._evaluate_picks( + executor, SimpleNamespace(test_battles_per_model=None), tmp_path, trials, 3 + ) + front = pd.read_parquet(tmp_path / "pareto.parquet") + assert set(front.config_id) == {"cheap", "accurate"} From ce9d509bf05021c45de95369f34d7f4cd38faf05 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 7 Oct 2026 15:53:25 +0200 Subject: [PATCH 13/20] Cover failed trials in single and multi-objective searches --- tests/test_tuning.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/test_tuning.py b/tests/test_tuning.py index bef63aed..4bbca703 100644 --- a/tests/test_tuning.py +++ b/tests/test_tuning.py @@ -201,9 +201,10 @@ def test_prices_required_before_execution(tmp_path): ) -def test_failed_trial_does_not_stop_search(tmp_path, token_counter): +@pytest.mark.parametrize("algorithm", ["neps_priorband", "mo_hyperband"]) +def test_failed_trial_does_not_stop_search(tmp_path, token_counter, algorithm): pytest.importorskip("neps") - cfg = _tuning_config(tmp_path, "mo_hyperband") + cfg = _tuning_config(tmp_path, algorithm) cfg.tune_judge.neps["ignore_errors"] = True calls = [] @@ -222,7 +223,7 @@ def execute(path): assert not results.empty and set(results.status) == {"completed"} pareto = pd.read_parquet(cfg.tune_judge.run_dir / "pareto.parquet") final = trials[(trials.status == "completed") & (trials.battles_per_model == 3)] - assert pareto.agreement.tolist() == [final.agreement.max()] + assert set(pareto.agreement) == {final.agreement.max()} assert set(pareto.status) == {"completed"} assert set(pareto.battles_per_model) == {3} From 021baa8cb7e293b39167f808b695f90e3035a25e Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Thu, 8 Oct 2026 16:32:50 +0200 Subject: [PATCH 14/20] Route judge tuning through a configured meta-evaluation target --- judgearena/benchmarks/registry.py | 3 +- .../meta_eval/tune-judge-comparia.yaml | 7 ---- .../meta_eval/tune-judge-lmarena-100k.yaml | 7 ---- .../meta_eval/tune-judge-lmarena-140k.yaml | 7 ---- judgearena/tuning/runner.py | 12 ++----- tests/test_mt_bench_downloads.py | 3 -- tests/test_tuning.py | 33 ++++++++++++++++++- 7 files changed, 36 insertions(+), 36 deletions(-) delete mode 100644 judgearena/tasks/definitions/meta_eval/tune-judge-comparia.yaml delete mode 100644 judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-100k.yaml delete mode 100644 judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-140k.yaml diff --git a/judgearena/benchmarks/registry.py b/judgearena/benchmarks/registry.py index ff88ab2f..ba79e40f 100644 --- a/judgearena/benchmarks/registry.py +++ b/judgearena/benchmarks/registry.py @@ -7,6 +7,7 @@ from judgearena.tasks.registry import get_packaged_task from judgearena.tasks.schema import ResolvedTaskSpec +from judgearena.tuning import TUNE_JUDGE_TASK if TYPE_CHECKING: from judgearena.config import RunConfig @@ -56,7 +57,7 @@ def benchmark_adapters() -> tuple[BenchmarkAdapter, ...]: BenchmarkAdapter("meta_eval", frozenset(), run_meta_eval), BenchmarkAdapter("mt_bench", frozenset(), run_mt_bench_benchmark), BenchmarkAdapter("mt_bench_101", frozenset(), run_mt_bench_101_benchmark), - BenchmarkAdapter("tune_judge", frozenset(), run_tune_judge), + BenchmarkAdapter("tune_judge", frozenset({TUNE_JUDGE_TASK}), run_tune_judge), BenchmarkAdapter("pairwise", None, run_pairwise), ) diff --git a/judgearena/tasks/definitions/meta_eval/tune-judge-comparia.yaml b/judgearena/tasks/definitions/meta_eval/tune-judge-comparia.yaml deleted file mode 100644 index b17d9874..00000000 --- a/judgearena/tasks/definitions/meta_eval/tune-judge-comparia.yaml +++ /dev/null @@ -1,7 +0,0 @@ -extends: _comparia.yaml -task: tune-judge-comparia -description: Tune judge settings on human votes from the ComparIA arena. - -protocol: - runner: tune_judge - meta_eval_task: meta-eval-comparia diff --git a/judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-100k.yaml b/judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-100k.yaml deleted file mode 100644 index 7269fc6f..00000000 --- a/judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-100k.yaml +++ /dev/null @@ -1,7 +0,0 @@ -extends: _lmarena-100k.yaml -task: tune-judge-lmarena-100k -description: Tune judge settings on human votes from LMSYS Chatbot Arena 100k. - -protocol: - runner: tune_judge - meta_eval_task: meta-eval-lmarena-100k diff --git a/judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-140k.yaml b/judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-140k.yaml deleted file mode 100644 index f7ff84dc..00000000 --- a/judgearena/tasks/definitions/meta_eval/tune-judge-lmarena-140k.yaml +++ /dev/null @@ -1,7 +0,0 @@ -extends: _lmarena-140k.yaml -task: tune-judge-lmarena-140k -description: Tune judge settings on human votes from LMSYS Chatbot Arena 140k. - -protocol: - runner: tune_judge - meta_eval_task: meta-eval-lmarena-140k diff --git a/judgearena/tuning/runner.py b/judgearena/tuning/runner.py index bd97e1f8..25b25f84 100644 --- a/judgearena/tuning/runner.py +++ b/judgearena/tuning/runner.py @@ -207,17 +207,9 @@ def evaluate(pipeline_directory: Path, **config): ) -def _meta_eval_task(task: ResolvedTaskSpec) -> str: - """Return the meta-eval task ID trials run as, keeping the variant suffix.""" - meta_eval_task = task.spec.protocol.meta_eval_task - if task.selection is None: - return meta_eval_task - return f"{meta_eval_task}-{task.selection.name}" - - def run_tune_judge( cfg: RunConfig, - task: ResolvedTaskSpec, + task: ResolvedTaskSpec | None, *, execute_trial: TrialExecutor = run_trial_subprocess, ) -> pd.DataFrame: @@ -228,7 +220,7 @@ def run_tune_judge( max_battles = specs[FIDELITY]["upper"] base = cfg.model_dump(mode="json", exclude={"tune_judge"}) runner = _TrialRunner( - base={**base, "task": _meta_eval_task(task)}, + base={**base, "task": tuning.meta_eval_task}, prices=tuning.price_per_million_tokens, execute_trial=execute_trial, ) diff --git a/tests/test_mt_bench_downloads.py b/tests/test_mt_bench_downloads.py index f85a8230..5ea1522b 100644 --- a/tests/test_mt_bench_downloads.py +++ b/tests/test_mt_bench_downloads.py @@ -127,9 +127,6 @@ def test_download_all_includes_mt_bench(tmp_path, monkeypatch): "meta-eval-lmarena-140k", "mt-bench", "mt-bench-101", - "tune-judge-comparia", - "tune-judge-lmarena-100k", - "tune-judge-lmarena-140k", ] assert all(path == tables_dir for _, path in hf_datasets) diff --git a/tests/test_tuning.py b/tests/test_tuning.py index 4bbca703..2b9b5286 100644 --- a/tests/test_tuning.py +++ b/tests/test_tuning.py @@ -43,10 +43,11 @@ def _tuning_config(tmp_path, algorithm="neps_priorband", **settings): for key, value in [("judge.temperature", 0.2), ("judge.model", 1)]: specs[key].update(prior=value, prior_confidence="medium") return RunConfig( - task="tune-judge-comparia-fr", + task="tune-judge", judge={"model": "Dummy/judge", "temperature": 0.2, "swap_mode": "both"}, run={"result_folder": str(tmp_path), "no_log_file": True}, tune_judge={ + "meta_eval_task": "meta-eval-comparia-fr", "neps": { "pipeline_space": specs, "optimizer": optimizer, @@ -135,6 +136,9 @@ def execute(path): ) with pytest.raises(ValueError, match="Incompatible configuration"): runner.run_tune_judge(incompatible, task, execute_trial=execute) + changed_target = _tuning_config(tmp_path, meta_eval_task="meta-eval-comparia-en") + with pytest.raises(ValueError, match="Incompatible configuration"): + runner.run_tune_judge(changed_target, task, execute_trial=execute) assert len(calls) == before fresh = _tuning_config(tmp_path, run_dir=tmp_path / "fresh") runner.run_tune_judge(fresh, task, execute_trial=execute) @@ -256,3 +260,30 @@ def test_pareto_report_uses_completed_full_fidelity_validation(tmp_path): ) front = pd.read_parquet(tmp_path / "pareto.parquet") assert set(front.config_id) == {"cheap", "accurate"} + + +def test_generic_route_uses_target_defaults(tmp_path): + from judgearena.benchmarks.registry import resolve_benchmark + + tuning = _tuning_config(tmp_path).tune_judge.model_dump() + direct = RunConfig(task=tuning["meta_eval_task"], judge={"model": "Dummy/judge"}) + outer = RunConfig( + task="tune-judge", judge={"model": "Dummy/judge"}, tune_judge=tuning + ) + assert outer.judge == direct.judge + assert outer.meta_eval == direct.meta_eval + resolved = resolve_benchmark(outer.task) + assert resolved.adapter.name == "tune_judge" and resolved.task is None + for target in ("unknown-task", "mt-bench"): + with pytest.raises(ValueError): + RunConfig( + task="tune-judge", + judge={"model": "Dummy/judge"}, + tune_judge={**tuning, "meta_eval_task": target}, + ) + with pytest.raises(ValueError): + RunConfig( + task="tune-judge", + judge={"model": "Dummy/judge", "swap_mode": "random"}, + tune_judge=tuning, + ) From c809d4a7178079e92ce7d288932e4e41901e7c91 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Thu, 8 Oct 2026 18:38:17 +0200 Subject: [PATCH 15/20] Resolve automatic judge prices and use the short tuning configuration --- judgearena/tuning/prices.json | 40 ++++++++ judgearena/tuning/pricing.py | 180 ++++++++++++++++++++++++++++++++++ judgearena/tuning/runner.py | 83 +++++++++++----- pyproject.toml | 1 + tests/test_tuning.py | 39 ++++---- tests/test_tuning_pricing.py | 133 +++++++++++++++++++++++++ 6 files changed, 432 insertions(+), 44 deletions(-) create mode 100644 judgearena/tuning/prices.json create mode 100644 judgearena/tuning/pricing.py create mode 100644 tests/test_tuning_pricing.py diff --git a/judgearena/tuning/prices.json b/judgearena/tuning/prices.json new file mode 100644 index 00000000..5b57fa53 --- /dev/null +++ b/judgearena/tuning/prices.json @@ -0,0 +1,40 @@ +{ + "source": { + "url": "https://github.com/geoalgo/judgetuning/blob/main/judgetuning/llm_client/llm_specs.py", + "revision": "main (commit lookup unavailable during implementation)", + "fetched_file_sha256": "ff23df03a293c1ded9c521aad1f2f6812947dbda4b16b1808c4b2d3a4ab9aed4" + }, + "models": { + "meta-llama/Meta-Llama-3-8B-instruct": [0.20, 0.20, "historical-together-reference"], + "meta-llama/Meta-Llama-3-70B-instruct": [0.90, 0.90, "historical-together-reference"], + "meta-llama/Meta-Llama-3.1-8B-Instruct": [0.11, 0.11, "measured-runtime"], + "meta-llama/Meta-Llama-3.1-70B-Instruct": [0.90, 0.90, "historical-together-reference"], + "neuralmagic/Meta-Llama-3.1-70B-Instruct-FP8": [0.35, 0.35, "measured-runtime"], + "HuggingFaceTB/SmolLM2-1.7B-Instruct": [0.06, 0.06, "historical-size-tier-reference"], + "Qwen/Qwen2.5-0.5B-Instruct": [0.06, 0.06, "historical-size-tier-reference"], + "Qwen/Qwen2.5-1.5B-Instruct": [0.06, 0.06, "historical-size-tier-reference"], + "Qwen/Qwen2.5-3B-Instruct": [0.10, 0.10, "historical-size-tier-reference"], + "Qwen/Qwen2.5-7B-Instruct": [0.12, 0.12, "measured-runtime"], + "Qwen/Qwen2.5-14B-Instruct": [0.30, 0.30, "historical-size-tier-reference"], + "Qwen/Qwen2.5-32B-Instruct": [0.80, 0.80, "historical-size-tier-reference"], + "Qwen/Qwen2.5-32B-Instruct-GPTQ-Int8": [0.36, 0.36, "measured-runtime"], + "Qwen/Qwen2.5-72B-Instruct": [1.20, 1.20, "historical-size-tier-reference"], + "Qwen/Qwen2.5-72B-Instruct-GPTQ-Int8": [0.58, 0.58, "measured-runtime"], + "google/gemma-2-2b-it": [0.06, 0.06, "historical-size-tier-reference"], + "google/gemma-2-9b-it": [0.14, 0.14, "measured-runtime"], + "google/gemma-2-27b-it": [0.30, 0.30, "measured-runtime"], + "mistralai/Mixtral-8x7B-Instruct-v0.1": [1.20, 1.20, "historical-together-reference"], + "BAAI/JudgeLM-7B-v1.0": [0.20, 0.20, "historical-size-tier-reference"], + "WeOpenML/PandaLM-7B-v1": [0.20, 0.20, "historical-size-tier-reference"] + }, + "tiers": [ + [3, 0.06], + [4, 0.10], + [8, 0.20], + [21, 0.30], + [41, 0.80], + [70, 0.90], + [null, 1.20] + ], + "tier_note": "Historical references document tiers through 41B and examples at 70B and 72B. The 41-70B interval uses the 70B reference; above 70B, 1.20 is an explicit extrapolation from Qwen2.5-72B, not a source-stated tier." +} diff --git a/judgearena/tuning/pricing.py b/judgearena/tuning/pricing.py new file mode 100644 index 00000000..acefb491 --- /dev/null +++ b/judgearena/tuning/pricing.py @@ -0,0 +1,180 @@ +"""Resolve estimated judge-token prices from explicit reference sources.""" + +from __future__ import annotations + +import json +import urllib.request +from dataclasses import dataclass +from pathlib import Path + +from huggingface_hub import try_to_load_from_cache + +_TABLE = Path(__file__).with_name("prices.json") +_OPENROUTER_URL = "https://openrouter.ai/api/v1/models" +_LOCAL = {"vllm", "llamacpp", "dummy"} + + +@dataclass(frozen=True) +class TokenPrice: + input: float + output: float + source: str + + +def _model_id(model: str) -> str: + return model.split("/", 1)[-1] if "/" in model else model + + +def _hosted(model: str) -> bool: + provider = model.split("/", 1)[0].lower() if "/" in model else "" + return bool(provider and provider not in _LOCAL) + + +def _repo_snapshot(model_id: str) -> Path | None: + path = Path(model_id).expanduser() + if path.is_dir(): + return path + try: + config = try_to_load_from_cache(model_id, "config.json", repo_type="model") + except (OSError, ValueError): + config = None + if isinstance(config, (str, Path)): + return Path(config).parent + return None + + +def _parameter_count(model_id: str) -> float | None: + path = _repo_snapshot(model_id) + if path is None: + return None + quant_config = path / "config.json" + if quant_config.exists(): + config = json.loads(quant_config.read_text(encoding="utf-8")) + if config.get("quantization_config"): + return None + index = path / "model.safetensors.index.json" + if index.exists(): + shards = sorted(set(json.loads(index.read_text())["weight_map"].values())) + else: + shards = [p.name for p in path.glob("*.safetensors")] + total = 0 + for shard in shards: + with (path / shard).open("rb") as stream: + size = int.from_bytes(stream.read(8), "little") + header = json.loads(stream.read(size)) + total += sum( + _shape_size(tensor["shape"]) + for name, tensor in header.items() + if name != "__metadata__" + ) + return total / 1e9 if total else None + + +def _shape_size(shape: list[int]) -> int: + product = 1 + for dim in shape: + product *= dim + return product + + +def _tier_price(params_b: float, tiers: list) -> float | None: + for upper, price in tiers: + if upper is None or params_b <= upper: + return float(price) + return None + + +def _fetch_openrouter(cache_file: Path) -> list[dict]: + if cache_file.exists(): + cached = json.loads(cache_file.read_text(encoding="utf-8")) + if isinstance(cached, dict): + return [ + { + "id": model_id, + "pricing": {"prompt": prices[0], "completion": prices[1]}, + } + for model_id, prices in cached.items() + ] + return cached + with urllib.request.urlopen(_OPENROUTER_URL, timeout=30) as response: + records = json.load(response)["data"] + cache_file.parent.mkdir(parents=True, exist_ok=True) + cache_file.write_text(json.dumps(records), encoding="utf-8") + return records + + +def _hosted_price(model_id: str, records: list[dict]) -> TokenPrice | None: + target = model_id.lower() + suffix = "/".join(target.split("/")[-2:]) + match = next( + (r for r in records if str(r.get("hugging_face_id") or "").lower() == target), + None, + ) + if match is None: + match = next( + (r for r in records if str(r.get("id") or "").lower() == suffix), None + ) + if match is None: + return None + pricing = match.get("pricing", {}) + if "prompt" not in pricing or "completion" not in pricing: + return None + return TokenPrice( + float(pricing["prompt"]) * 1e6, + float(pricing["completion"]) * 1e6, + "openrouter_reference", + ) + + +def resolve_prices( + models: list[str] | set[str] | tuple[str, ...], + overrides: dict[str, float | dict[str, float]], + store_root: str | Path | None, + require_cost: bool = True, +) -> dict[str, TokenPrice]: + """Resolve USD/M-token estimates; never reports billed or measured cost.""" + table = json.loads(_TABLE.read_text(encoding="utf-8")) + specs = table["models"] + resolved: dict[str, TokenPrice] = {} + pending = [] + for model in sorted(set(models)): + override = overrides.get(model) + if override is not None: + pair = ( + {"input": override, "output": override} + if isinstance(override, (int, float)) + else override + ) + if "input" in pair and "output" in pair: + resolved[model] = TokenPrice( + float(pair["input"]), float(pair["output"]), "user_override" + ) + continue + provider = model.split("/", 1)[0].lower() if "/" in model else "" + model_id = _model_id(model) + spec = specs.get(model_id) + if spec and provider == "vllm": + resolved[model] = TokenPrice(float(spec[0]), float(spec[1]), spec[2]) + elif _hosted(model): + pending.append((model, model_id)) + elif provider == "vllm": + count = _parameter_count(model_id) + price = _tier_price(count, table["tiers"]) if count is not None else None + if price is not None: + resolved[model] = TokenPrice( + price, price, "historical-size-tier-reference" + ) + if pending: + if store_root is None: + raise ValueError( + "store_root is required to cache hosted model reference pricing" + ) + records = _fetch_openrouter(Path(store_root) / "openrouter_pricing.json") + for model, model_id in pending: + price = _hosted_price(model_id, records) + if price: + resolved[model] = price + missing = sorted(set(models) - set(resolved)) + if missing and require_cost: + raise ValueError(f"No token price could be resolved for: {missing}") + return resolved diff --git a/judgearena/tuning/runner.py b/judgearena/tuning/runner.py index 25b25f84..a9568aa4 100644 --- a/judgearena/tuning/runner.py +++ b/judgearena/tuning/runner.py @@ -8,7 +8,7 @@ import subprocess import sys from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import asdict, dataclass from functools import cache from pathlib import Path from typing import TYPE_CHECKING @@ -20,16 +20,17 @@ from judgearena.config import RunConfig, dump_config from judgearena.log import get_logger from judgearena.tuning.search_space import ( - FIDELITY, apply_overrides, axis_overrides, build_neps_space, decode_config, + parameter_specs, ) if TYPE_CHECKING: from judgearena.tasks.schema import ResolvedTaskSpec +from judgearena.tuning.pricing import TokenPrice, resolve_prices from judgearena.tuning.session import collect_trials, prepare_session logger = get_logger(__name__) @@ -53,7 +54,7 @@ def _judge_token_encoding() -> tiktoken.Encoding: return tiktoken.encoding_for_model("gpt-4o") -def _read_trial(trial_dir: Path, price_per_million_tokens: float) -> dict: +def _read_trial(trial_dir: Path, price: TokenPrice | None) -> dict: (result_path,) = trial_dir.glob("*/results.json") metrics = json.loads(result_path.read_text())["metrics"] agreement = metrics["meta_eval_agreement"]["all"] @@ -61,10 +62,22 @@ def _read_trial(trial_dir: Path, price_per_million_tokens: float) -> dict: result_path.parent / "annotations.parquet", columns=["battle_id", "judge_input", "judge_completion"], ) - texts = [*annotations["judge_input"], *annotations["judge_completion"].fillna("")] encoding = _judge_token_encoding() - tokens = sum(len(encoding.encode(text, disallowed_special=())) for text in texts) - tokens_per_battle = tokens / annotations["battle_id"].nunique() + input_tokens = sum( + len(encoding.encode(text, disallowed_special=())) + for text in annotations["judge_input"] + ) + output_tokens = sum( + len(encoding.encode(text, disallowed_special=())) + for text in annotations["judge_completion"].fillna("") + ) + n_battles = annotations["battle_id"].nunique() + tokens_per_battle = (input_tokens + output_tokens) / n_battles + cost = ( + (input_tokens * price.input + output_tokens * price.output) / n_battles / 1000 + if price + else None + ) return { "agreement": agreement["accuracy_attempted"], "cohen_kappa": agreement["cohen_kappa"], @@ -74,7 +87,7 @@ def _read_trial(trial_dir: Path, price_per_million_tokens: float) -> dict: .get("soft", {}) .get("spearman"), "tokens_per_battle": tokens_per_battle, - "cost_per_1k_battles": tokens_per_battle * price_per_million_tokens / 1000, + "cost_per_1k_battles": cost, "result_dir": str(result_path.parent), } @@ -84,7 +97,7 @@ class _TrialRunner: """Materialize trial configs from one base config and record their results.""" base: dict - prices: dict[str, float] + prices: dict[str, TokenPrice] execute_trial: TrialExecutor def config( @@ -131,14 +144,14 @@ def run( return { **record, "status": "completed", - **_read_trial(folder, self.prices[trial_cfg.judge.model]), + **_read_trial(folder, self.prices.get(trial_cfg.judge.model)), } try: self.execute_trial(config_path) except subprocess.CalledProcessError as exc: logger.warning("Trial %s failed: %s", folder, exc) return {**record, "status": "failed"} - price = self.prices[trial_cfg.judge.model] + price = self.prices.get(trial_cfg.judge.model) return {**record, "status": "completed", **_read_trial(folder, price)} @@ -154,7 +167,8 @@ def _objective(record: dict, objectives: list[str]) -> dict: "info_dict": record, } values = [ - 1 - record[key] if key == "agreement" else record[key] for key in objectives + 1 - record[key] if key == "agreement" else record["cost_per_1k_battles"] + for key in objectives ] return { "objective_to_minimize": values[0] if len(values) == 1 else values, @@ -176,13 +190,10 @@ def evaluate(pipeline_directory: Path, **config): ) return _objective(record, cfg.tune_judge.objectives) - options = { - key: value - for key, value in cfg.tune_judge.neps.items() - if key != "pipeline_space" - } + options = dict(cfg.tune_judge.neps) # NePS 0.17's space compatibility check does not inspect optimizer mappings. - if isinstance(optimizer := options.get("optimizer"), dict): + optimizer = cfg.tune_judge.optimizer + if isinstance(optimizer, dict): options["optimizer"] = ( optimizer["name"], {key: value for key, value in optimizer.items() if key != "name"}, @@ -215,13 +226,19 @@ def run_tune_judge( ) -> pd.DataFrame: """Tune judge settings on validation battles and score the picks on test.""" tuning = cfg.tune_judge - specs = tuning.neps["pipeline_space"] - space = build_neps_space(specs) - max_battles = specs[FIDELITY]["upper"] base = cfg.model_dump(mode="json", exclude={"tune_judge"}) + if base["judge"]["prompt"] is None and base["judge"]["prompt_preset"] is None: + from judgearena.tasks.registry import get_packaged_task + + base["judge"]["prompt_preset"] = get_packaged_task( + tuning.meta_eval_task + ).spec.protocol.judge.default_prompt_preset + specs = parameter_specs(tuning.search_space, base, tuning.fidelity) + space = build_neps_space(specs) + max_battles = tuning.fidelity["battles_per_model"]["upper"] runner = _TrialRunner( base={**base, "task": tuning.meta_eval_task}, - prices=tuning.price_per_million_tokens, + prices={}, execute_trial=execute_trial, ) models = set() @@ -236,11 +253,22 @@ def run_tune_judge( models.add(trial_cfg.judge.model) if "judge.model" not in specs: models.add(cfg.judge.model) - if missing := sorted(models - set(tuning.price_per_million_tokens)): - raise ValueError( - f"tune_judge.price_per_million_tokens is missing models: {missing}" - ) + prices = resolve_prices( + models, + tuning.price_per_million_tokens, + cfg.run.store_root or Path(cfg.run.result_folder), + require_cost="cost" in tuning.objectives, + ) + runner = _TrialRunner(runner.base, prices, execute_trial) tune_dir = prepare_session(cfg) + price_path = tune_dir / "prices.json" + snapshot = {model: asdict(price) for model, price in prices.items()} + if price_path.exists(): + stored = json.loads(price_path.read_text()) + if stored != snapshot: + raise ValueError("Resolved prices changed for this tuning session") + elif not tuning.search_only: + price_path.write_text(json.dumps(snapshot, indent=2)) logger.info("Tuning %s in %s", cfg.task, tune_dir) _search(runner, cfg, tune_dir, space) if tuning.search_only: @@ -263,9 +291,10 @@ def _evaluate_picks(runner, tuning, tune_dir, trials, max_battles) -> pd.DataFra ) from neps.optimizers.utils.multiobjective.epsnet import pareto_efficient - costs = ranked[["agreement", "cost_per_1k_battles"]].to_numpy().copy() + priced = ranked.dropna(subset=["cost_per_1k_battles"]) + costs = priced[["agreement", "cost_per_1k_battles"]].to_numpy(dtype=float).copy() costs[:, 0] *= -1 - pareto = ranked.loc[pareto_efficient(costs)] + pareto = priced.loc[pareto_efficient(costs)] if len(priced) else priced pareto.to_parquet(tune_dir / "pareto.parquet", index=False) picks = ranked.groupby("judge_model", sort=False).head(1) test_results = pd.DataFrame( diff --git a/pyproject.toml b/pyproject.toml index 81228bf8..ad77ff0e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,7 @@ exclude = ["slurmpilot_scripts*"] "templates/*/*.txt", ] "judgearena.tasks" = ["definitions/*/*.yaml"] +"judgearena.tuning" = ["prices.json"] "judgearena.benchmarks.pairwise.scoring" = ["*.csv.gz"] [dependency-groups] diff --git a/tests/test_tuning.py b/tests/test_tuning.py index 2b9b5286..dff53083 100644 --- a/tests/test_tuning.py +++ b/tests/test_tuning.py @@ -12,14 +12,13 @@ from judgearena.config import RunConfig, dump_config, load_config from judgearena.tasks.registry import get_packaged_task from judgearena.tuning import runner -from judgearena.tuning.search_space import FIDELITY def _tuning_config(tmp_path, algorithm="neps_priorband", **settings): optimizer = {"name": algorithm, "eta": 3} objectives = ["agreement"] if algorithm in {"mo_hyperband", "primo"}: - objectives.append("cost_per_1k_battles") + objectives.append("cost") if algorithm == "primo": optimizer.update( initial_design_size=2, @@ -36,12 +35,14 @@ def _tuning_config(tmp_path, algorithm="neps_priorband", **settings): }, ) specs = { - "judge.temperature": dict(type="Float", lower=0.0, upper=1.0), - "judge.model": dict(type="Categorical", choices=["Dummy/judge", "Dummy/other"]), - FIDELITY: dict(type="IntegerFidelity", lower=1, upper=3), + "judge": { + "temperature": {"lower": 0.0, "upper": 1.0}, + "model": { + "choices": ["Dummy/judge", "Dummy/other"], + "prior": "Dummy/other", + }, + } } - for key, value in [("judge.temperature", 0.2), ("judge.model", 1)]: - specs[key].update(prior=value, prior_confidence="medium") return RunConfig( task="tune-judge", judge={"model": "Dummy/judge", "temperature": 0.2, "swap_mode": "both"}, @@ -49,10 +50,11 @@ def _tuning_config(tmp_path, algorithm="neps_priorband", **settings): tune_judge={ "meta_eval_task": "meta-eval-comparia-fr", "neps": { - "pipeline_space": specs, - "optimizer": optimizer, "total_evaluations_to_spend": 10, }, + "search_space": specs, + "fidelity": {"battles_per_model": {"lower": 1, "upper": 3}}, + "optimizer": optimizer, "objectives": objectives, "price_per_million_tokens": {"Dummy/judge": 1.0, "Dummy/other": 1.0}, "run_dir": tmp_path / "tuning", @@ -164,8 +166,11 @@ def test_budget_end_does_not_wait_for_unstarted_batch( ): pytest.importorskip("neps") cfg = _tuning_config(tmp_path) + cfg.tune_judge.optimizer = { + "name": "neps_random_search", + "ignore_fidelity": "highest_fidelity", + } cfg.tune_judge.neps.update( - optimizer={"name": "neps_random_search", "ignore_fidelity": "highest_fidelity"}, total_evaluations_to_spend=2, sample_batch_size=3, ) @@ -184,20 +189,20 @@ def test_cost_counts_both_cached_orientations(tmp_path, token_counter): path = tmp_path / "config.yaml" dump_config(_tuning_config(tmp_path), path) _fake_trial(path) - record = runner._read_trial(tmp_path, price_per_million_tokens=2.0) + record = runner._read_trial(tmp_path, price=runner.TokenPrice(2.0, 4.0, "test")) assert record["tokens_per_battle"] == 6 - assert record["cost_per_1k_battles"] == pytest.approx(0.012) + assert record["cost_per_1k_battles"] == pytest.approx(0.016) objective = runner._objective( - {**record, "status": "completed"}, ["agreement", "cost_per_1k_battles"] + {**record, "status": "completed"}, ["agreement", "cost"] ) - assert objective["objective_to_minimize"] == pytest.approx([0.42, 0.012]) + assert objective["objective_to_minimize"] == pytest.approx([0.42, 0.016]) assert "cost" not in objective def test_prices_required_before_execution(tmp_path): pytest.importorskip("neps") - cfg = _tuning_config(tmp_path, price_per_million_tokens={}) - with pytest.raises(ValueError, match="missing models"): + cfg = _tuning_config(tmp_path, "mo_hyperband", price_per_million_tokens={}) + with pytest.raises(ValueError, match="No token price"): runner.run_tune_judge( cfg, get_packaged_task(cfg.task), @@ -273,7 +278,7 @@ def test_generic_route_uses_target_defaults(tmp_path): assert outer.judge == direct.judge assert outer.meta_eval == direct.meta_eval resolved = resolve_benchmark(outer.task) - assert resolved.adapter.name == "tune_judge" and resolved.task is None + assert resolved.adapter.name == "tune_judge" and resolved.task.spec.dataset is None for target in ("unknown-task", "mt-bench"): with pytest.raises(ValueError): RunConfig( diff --git a/tests/test_tuning_pricing.py b/tests/test_tuning_pricing.py new file mode 100644 index 00000000..0f809d54 --- /dev/null +++ b/tests/test_tuning_pricing.py @@ -0,0 +1,133 @@ +import json + +import pytest + +from judgearena.tuning import pricing +from judgearena.tuning.pricing import TokenPrice, resolve_prices + + +def test_source_prices_exact_model_identifiers(tmp_path): + prices = resolve_prices( + [ + "VLLM/Qwen/Qwen2.5-7B-Instruct", + "VLLM/google/gemma-2-9b-it", + "VLLM/google/gemma-2-27b-it", + "VLLM/Qwen/Qwen2.5-32B-Instruct-GPTQ-Int8", + ], + {}, + tmp_path, + ) + assert prices["VLLM/Qwen/Qwen2.5-7B-Instruct"] == TokenPrice( + 0.12, 0.12, "measured-runtime" + ) + assert prices["VLLM/google/gemma-2-9b-it"].input == 0.14 + assert prices["VLLM/google/gemma-2-27b-it"].input == 0.30 + assert prices["VLLM/Qwen/Qwen2.5-32B-Instruct-GPTQ-Int8"].input == 0.36 + + +@pytest.mark.parametrize( + "count, expected", + [(3, 0.06), (4, 0.10), (8, 0.20), (21, 0.30), (41, 0.80), (70, 0.90), (72, 1.20)], +) +def test_parameter_tiers(count, expected): + tiers = json.loads(pricing._TABLE.read_text())["tiers"] + assert pricing._tier_price(count, tiers) == expected + + +def _write_safetensors(path, shape): + header = {"weight": {"dtype": "I8", "shape": shape, "data_offsets": [0, 1]}} + encoded = json.dumps(header).encode() + (path / "model.safetensors").write_bytes( + len(encoded).to_bytes(8, "little") + encoded + ) + + +def test_offline_hf_cache_snapshot_counts_shape_products(tmp_path, monkeypatch): + snapshot = tmp_path / "models--org--model" / "snapshots" / "abc" + snapshot.mkdir(parents=True) + _write_safetensors(snapshot, [2_000_000_000]) + monkeypatch.setattr( + pricing, "try_to_load_from_cache", lambda *a, **k: snapshot / "config.json" + ) + result = resolve_prices(["VLLM/org/model"], {}, tmp_path / "store") + assert result["VLLM/org/model"].input == 0.06 + + +def test_quantized_model_without_exact_entry_has_no_tier_price(tmp_path, monkeypatch): + snapshot = tmp_path / "quantized" + snapshot.mkdir() + (snapshot / "config.json").write_text( + json.dumps({"quantization_config": {"bits": 4}}) + ) + monkeypatch.setattr( + pricing, "try_to_load_from_cache", lambda *a, **k: snapshot / "config.json" + ) + with pytest.raises(ValueError, match="quantized-model"): + resolve_prices(["VLLM/quantized-model"], {}, tmp_path) + + +def test_override_numeric_and_asymmetric(tmp_path, monkeypatch): + def fail(_): + raise AssertionError("override must avoid hosted pricing fetch") + + monkeypatch.setattr(pricing, "_fetch_openrouter", fail) + result = resolve_prices( + ["OpenRouter/org/a", "OpenAI/b"], + {"OpenRouter/org/a": 0.4, "OpenAI/b": {"input": 0.2, "output": 0.9}}, + tmp_path, + ) + assert result["OpenRouter/org/a"] == TokenPrice(0.4, 0.4, "user_override") + assert result["OpenAI/b"] == TokenPrice(0.2, 0.9, "user_override") + + +def test_unknown_required_fails_but_agreement_only_omits(tmp_path): + with pytest.raises(ValueError, match="unlisted"): + resolve_prices(["VLLM/unlisted"], {}, tmp_path) + assert resolve_prices(["VLLM/unlisted"], {}, tmp_path, require_cost=False) == {} + + +def test_openrouter_matches_hf_id_first_then_model_id(tmp_path, monkeypatch): + records = [ + { + "id": "provider/other", + "hugging_face_id": "org/model", + "pricing": {"prompt": "0.000002", "completion": "0.000007"}, + }, + { + "id": "org/fallback", + "pricing": {"prompt": "0.000003", "completion": "0.000009"}, + }, + ] + monkeypatch.setattr(pricing, "_fetch_openrouter", lambda path: records) + result = resolve_prices( + ["OpenRouter/org/model", "Provider/org/fallback"], {}, tmp_path + ) + assert result["OpenRouter/org/model"] == TokenPrice( + 2.0, 7.0, "openrouter_reference" + ) + assert result["Provider/org/fallback"] == TokenPrice( + 3.0, 9.0, "openrouter_reference" + ) + + +def test_nonhosted_unlisted_model_does_not_fetch(tmp_path, monkeypatch): + def fail(_): + raise AssertionError("unexpected hosted pricing fetch") + + monkeypatch.setattr(pricing, "_fetch_openrouter", fail) + assert resolve_prices(["VLLM/unlisted"], {}, tmp_path, require_cost=False) == {} + + +def test_legacy_openrouter_cache_mapping_is_supported(tmp_path): + cache_file = tmp_path / "openrouter_pricing.json" + cache_file.write_text(json.dumps({"org/model": [0.000002, 0.000007]})) + records = pricing._fetch_openrouter(cache_file) + assert pricing._hosted_price("org/model", records) == TokenPrice( + 2.0, 7.0, "openrouter_reference" + ) + + +def test_hf_cache_missing_sentinel_is_not_a_path(tmp_path, monkeypatch): + sentinel = object() + monkeypatch.setattr(pricing, "try_to_load_from_cache", lambda *a, **k: sentinel) + assert pricing._repo_snapshot("org/missing") is None From a9b1a862ad954daa7144e505f451b48e9a248003 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Thu, 8 Oct 2026 18:39:51 +0200 Subject: [PATCH 16/20] Pin judge pricing table source revision --- judgearena/tuning/prices.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/judgearena/tuning/prices.json b/judgearena/tuning/prices.json index 5b57fa53..49db40b2 100644 --- a/judgearena/tuning/prices.json +++ b/judgearena/tuning/prices.json @@ -1,7 +1,7 @@ { "source": { - "url": "https://github.com/geoalgo/judgetuning/blob/main/judgetuning/llm_client/llm_specs.py", - "revision": "main (commit lookup unavailable during implementation)", + "url": "https://github.com/geoalgo/judgetuning/blob/c5202286238d6af45c9652320ce5aeb62158e671/judgetuning/llm_client/llm_specs.py", + "revision": "c5202286238d6af45c9652320ce5aeb62158e671", "fetched_file_sha256": "ff23df03a293c1ded9c521aad1f2f6812947dbda4b16b1808c4b2d3a4ab9aed4" }, "models": { From 1d8d0750cab0f76714334d9e6bdb24a9021d613a Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 9 Oct 2026 09:40:42 +0200 Subject: [PATCH 17/20] Keep tuning execution cost override-only and move reference pricing out --- judgearena/tuning/prices.json | 40 -------- judgearena/tuning/pricing.py | 184 +++------------------------------- judgearena/tuning/runner.py | 1 - pyproject.toml | 1 - tests/test_tuning_pricing.py | 133 ------------------------ 5 files changed, 16 insertions(+), 343 deletions(-) delete mode 100644 judgearena/tuning/prices.json delete mode 100644 tests/test_tuning_pricing.py diff --git a/judgearena/tuning/prices.json b/judgearena/tuning/prices.json deleted file mode 100644 index 49db40b2..00000000 --- a/judgearena/tuning/prices.json +++ /dev/null @@ -1,40 +0,0 @@ -{ - "source": { - "url": "https://github.com/geoalgo/judgetuning/blob/c5202286238d6af45c9652320ce5aeb62158e671/judgetuning/llm_client/llm_specs.py", - "revision": "c5202286238d6af45c9652320ce5aeb62158e671", - "fetched_file_sha256": "ff23df03a293c1ded9c521aad1f2f6812947dbda4b16b1808c4b2d3a4ab9aed4" - }, - "models": { - "meta-llama/Meta-Llama-3-8B-instruct": [0.20, 0.20, "historical-together-reference"], - "meta-llama/Meta-Llama-3-70B-instruct": [0.90, 0.90, "historical-together-reference"], - "meta-llama/Meta-Llama-3.1-8B-Instruct": [0.11, 0.11, "measured-runtime"], - "meta-llama/Meta-Llama-3.1-70B-Instruct": [0.90, 0.90, "historical-together-reference"], - "neuralmagic/Meta-Llama-3.1-70B-Instruct-FP8": [0.35, 0.35, "measured-runtime"], - "HuggingFaceTB/SmolLM2-1.7B-Instruct": [0.06, 0.06, "historical-size-tier-reference"], - "Qwen/Qwen2.5-0.5B-Instruct": [0.06, 0.06, "historical-size-tier-reference"], - "Qwen/Qwen2.5-1.5B-Instruct": [0.06, 0.06, "historical-size-tier-reference"], - "Qwen/Qwen2.5-3B-Instruct": [0.10, 0.10, "historical-size-tier-reference"], - "Qwen/Qwen2.5-7B-Instruct": [0.12, 0.12, "measured-runtime"], - "Qwen/Qwen2.5-14B-Instruct": [0.30, 0.30, "historical-size-tier-reference"], - "Qwen/Qwen2.5-32B-Instruct": [0.80, 0.80, "historical-size-tier-reference"], - "Qwen/Qwen2.5-32B-Instruct-GPTQ-Int8": [0.36, 0.36, "measured-runtime"], - "Qwen/Qwen2.5-72B-Instruct": [1.20, 1.20, "historical-size-tier-reference"], - "Qwen/Qwen2.5-72B-Instruct-GPTQ-Int8": [0.58, 0.58, "measured-runtime"], - "google/gemma-2-2b-it": [0.06, 0.06, "historical-size-tier-reference"], - "google/gemma-2-9b-it": [0.14, 0.14, "measured-runtime"], - "google/gemma-2-27b-it": [0.30, 0.30, "measured-runtime"], - "mistralai/Mixtral-8x7B-Instruct-v0.1": [1.20, 1.20, "historical-together-reference"], - "BAAI/JudgeLM-7B-v1.0": [0.20, 0.20, "historical-size-tier-reference"], - "WeOpenML/PandaLM-7B-v1": [0.20, 0.20, "historical-size-tier-reference"] - }, - "tiers": [ - [3, 0.06], - [4, 0.10], - [8, 0.20], - [21, 0.30], - [41, 0.80], - [70, 0.90], - [null, 1.20] - ], - "tier_note": "Historical references document tiers through 41B and examples at 70B and 72B. The 41-70B interval uses the 70B reference; above 70B, 1.20 is an explicit extrapolation from Qwen2.5-72B, not a source-stated tier." -} diff --git a/judgearena/tuning/pricing.py b/judgearena/tuning/pricing.py index acefb491..be671c95 100644 --- a/judgearena/tuning/pricing.py +++ b/judgearena/tuning/pricing.py @@ -1,17 +1,6 @@ -"""Resolve estimated judge-token prices from explicit reference sources.""" +"""Explicit judge-token prices used by the tuning cost objective.""" -from __future__ import annotations - -import json -import urllib.request from dataclasses import dataclass -from pathlib import Path - -from huggingface_hub import try_to_load_from_cache - -_TABLE = Path(__file__).with_name("prices.json") -_OPENROUTER_URL = "https://openrouter.ai/api/v1/models" -_LOCAL = {"vllm", "llamacpp", "dummy"} @dataclass(frozen=True) @@ -21,160 +10,19 @@ class TokenPrice: source: str -def _model_id(model: str) -> str: - return model.split("/", 1)[-1] if "/" in model else model - - -def _hosted(model: str) -> bool: - provider = model.split("/", 1)[0].lower() if "/" in model else "" - return bool(provider and provider not in _LOCAL) - - -def _repo_snapshot(model_id: str) -> Path | None: - path = Path(model_id).expanduser() - if path.is_dir(): - return path - try: - config = try_to_load_from_cache(model_id, "config.json", repo_type="model") - except (OSError, ValueError): - config = None - if isinstance(config, (str, Path)): - return Path(config).parent - return None - - -def _parameter_count(model_id: str) -> float | None: - path = _repo_snapshot(model_id) - if path is None: - return None - quant_config = path / "config.json" - if quant_config.exists(): - config = json.loads(quant_config.read_text(encoding="utf-8")) - if config.get("quantization_config"): - return None - index = path / "model.safetensors.index.json" - if index.exists(): - shards = sorted(set(json.loads(index.read_text())["weight_map"].values())) - else: - shards = [p.name for p in path.glob("*.safetensors")] - total = 0 - for shard in shards: - with (path / shard).open("rb") as stream: - size = int.from_bytes(stream.read(8), "little") - header = json.loads(stream.read(size)) - total += sum( - _shape_size(tensor["shape"]) - for name, tensor in header.items() - if name != "__metadata__" - ) - return total / 1e9 if total else None - - -def _shape_size(shape: list[int]) -> int: - product = 1 - for dim in shape: - product *= dim - return product - - -def _tier_price(params_b: float, tiers: list) -> float | None: - for upper, price in tiers: - if upper is None or params_b <= upper: - return float(price) - return None - - -def _fetch_openrouter(cache_file: Path) -> list[dict]: - if cache_file.exists(): - cached = json.loads(cache_file.read_text(encoding="utf-8")) - if isinstance(cached, dict): - return [ - { - "id": model_id, - "pricing": {"prompt": prices[0], "completion": prices[1]}, - } - for model_id, prices in cached.items() - ] - return cached - with urllib.request.urlopen(_OPENROUTER_URL, timeout=30) as response: - records = json.load(response)["data"] - cache_file.parent.mkdir(parents=True, exist_ok=True) - cache_file.write_text(json.dumps(records), encoding="utf-8") - return records - - -def _hosted_price(model_id: str, records: list[dict]) -> TokenPrice | None: - target = model_id.lower() - suffix = "/".join(target.split("/")[-2:]) - match = next( - (r for r in records if str(r.get("hugging_face_id") or "").lower() == target), - None, - ) - if match is None: - match = next( - (r for r in records if str(r.get("id") or "").lower() == suffix), None - ) - if match is None: - return None - pricing = match.get("pricing", {}) - if "prompt" not in pricing or "completion" not in pricing: - return None - return TokenPrice( - float(pricing["prompt"]) * 1e6, - float(pricing["completion"]) * 1e6, - "openrouter_reference", - ) - - -def resolve_prices( - models: list[str] | set[str] | tuple[str, ...], - overrides: dict[str, float | dict[str, float]], - store_root: str | Path | None, - require_cost: bool = True, -) -> dict[str, TokenPrice]: - """Resolve USD/M-token estimates; never reports billed or measured cost.""" - table = json.loads(_TABLE.read_text(encoding="utf-8")) - specs = table["models"] - resolved: dict[str, TokenPrice] = {} - pending = [] - for model in sorted(set(models)): - override = overrides.get(model) - if override is not None: - pair = ( - {"input": override, "output": override} - if isinstance(override, (int, float)) - else override - ) - if "input" in pair and "output" in pair: - resolved[model] = TokenPrice( - float(pair["input"]), float(pair["output"]), "user_override" - ) - continue - provider = model.split("/", 1)[0].lower() if "/" in model else "" - model_id = _model_id(model) - spec = specs.get(model_id) - if spec and provider == "vllm": - resolved[model] = TokenPrice(float(spec[0]), float(spec[1]), spec[2]) - elif _hosted(model): - pending.append((model, model_id)) - elif provider == "vllm": - count = _parameter_count(model_id) - price = _tier_price(count, table["tiers"]) if count is not None else None - if price is not None: - resolved[model] = TokenPrice( - price, price, "historical-size-tier-reference" - ) - if pending: - if store_root is None: - raise ValueError( - "store_root is required to cache hosted model reference pricing" - ) - records = _fetch_openrouter(Path(store_root) / "openrouter_pricing.json") - for model, model_id in pending: - price = _hosted_price(model_id, records) - if price: - resolved[model] = price - missing = sorted(set(models) - set(resolved)) +def resolve_prices(models, overrides, *, require_cost=True) -> dict[str, TokenPrice]: + """Resolve explicit USD-per-million-token rates before searching.""" + prices = {} + for model in sorted(models): + if model not in overrides: + continue + rate = overrides[model] + if isinstance(rate, (int, float)): + rate = {"input": rate, "output": rate} + prices[model] = TokenPrice(rate["input"], rate["output"], "user_override") + missing = sorted(set(models) - prices.keys()) if missing and require_cost: - raise ValueError(f"No token price could be resolved for: {missing}") - return resolved + raise ValueError( + f"No token price for {missing}; set tune_judge.price_per_million_tokens" + ) + return prices diff --git a/judgearena/tuning/runner.py b/judgearena/tuning/runner.py index a9568aa4..d6e59565 100644 --- a/judgearena/tuning/runner.py +++ b/judgearena/tuning/runner.py @@ -256,7 +256,6 @@ def run_tune_judge( prices = resolve_prices( models, tuning.price_per_million_tokens, - cfg.run.store_root or Path(cfg.run.result_folder), require_cost="cost" in tuning.objectives, ) runner = _TrialRunner(runner.base, prices, execute_trial) diff --git a/pyproject.toml b/pyproject.toml index ad77ff0e..81228bf8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,6 @@ exclude = ["slurmpilot_scripts*"] "templates/*/*.txt", ] "judgearena.tasks" = ["definitions/*/*.yaml"] -"judgearena.tuning" = ["prices.json"] "judgearena.benchmarks.pairwise.scoring" = ["*.csv.gz"] [dependency-groups] diff --git a/tests/test_tuning_pricing.py b/tests/test_tuning_pricing.py deleted file mode 100644 index 0f809d54..00000000 --- a/tests/test_tuning_pricing.py +++ /dev/null @@ -1,133 +0,0 @@ -import json - -import pytest - -from judgearena.tuning import pricing -from judgearena.tuning.pricing import TokenPrice, resolve_prices - - -def test_source_prices_exact_model_identifiers(tmp_path): - prices = resolve_prices( - [ - "VLLM/Qwen/Qwen2.5-7B-Instruct", - "VLLM/google/gemma-2-9b-it", - "VLLM/google/gemma-2-27b-it", - "VLLM/Qwen/Qwen2.5-32B-Instruct-GPTQ-Int8", - ], - {}, - tmp_path, - ) - assert prices["VLLM/Qwen/Qwen2.5-7B-Instruct"] == TokenPrice( - 0.12, 0.12, "measured-runtime" - ) - assert prices["VLLM/google/gemma-2-9b-it"].input == 0.14 - assert prices["VLLM/google/gemma-2-27b-it"].input == 0.30 - assert prices["VLLM/Qwen/Qwen2.5-32B-Instruct-GPTQ-Int8"].input == 0.36 - - -@pytest.mark.parametrize( - "count, expected", - [(3, 0.06), (4, 0.10), (8, 0.20), (21, 0.30), (41, 0.80), (70, 0.90), (72, 1.20)], -) -def test_parameter_tiers(count, expected): - tiers = json.loads(pricing._TABLE.read_text())["tiers"] - assert pricing._tier_price(count, tiers) == expected - - -def _write_safetensors(path, shape): - header = {"weight": {"dtype": "I8", "shape": shape, "data_offsets": [0, 1]}} - encoded = json.dumps(header).encode() - (path / "model.safetensors").write_bytes( - len(encoded).to_bytes(8, "little") + encoded - ) - - -def test_offline_hf_cache_snapshot_counts_shape_products(tmp_path, monkeypatch): - snapshot = tmp_path / "models--org--model" / "snapshots" / "abc" - snapshot.mkdir(parents=True) - _write_safetensors(snapshot, [2_000_000_000]) - monkeypatch.setattr( - pricing, "try_to_load_from_cache", lambda *a, **k: snapshot / "config.json" - ) - result = resolve_prices(["VLLM/org/model"], {}, tmp_path / "store") - assert result["VLLM/org/model"].input == 0.06 - - -def test_quantized_model_without_exact_entry_has_no_tier_price(tmp_path, monkeypatch): - snapshot = tmp_path / "quantized" - snapshot.mkdir() - (snapshot / "config.json").write_text( - json.dumps({"quantization_config": {"bits": 4}}) - ) - monkeypatch.setattr( - pricing, "try_to_load_from_cache", lambda *a, **k: snapshot / "config.json" - ) - with pytest.raises(ValueError, match="quantized-model"): - resolve_prices(["VLLM/quantized-model"], {}, tmp_path) - - -def test_override_numeric_and_asymmetric(tmp_path, monkeypatch): - def fail(_): - raise AssertionError("override must avoid hosted pricing fetch") - - monkeypatch.setattr(pricing, "_fetch_openrouter", fail) - result = resolve_prices( - ["OpenRouter/org/a", "OpenAI/b"], - {"OpenRouter/org/a": 0.4, "OpenAI/b": {"input": 0.2, "output": 0.9}}, - tmp_path, - ) - assert result["OpenRouter/org/a"] == TokenPrice(0.4, 0.4, "user_override") - assert result["OpenAI/b"] == TokenPrice(0.2, 0.9, "user_override") - - -def test_unknown_required_fails_but_agreement_only_omits(tmp_path): - with pytest.raises(ValueError, match="unlisted"): - resolve_prices(["VLLM/unlisted"], {}, tmp_path) - assert resolve_prices(["VLLM/unlisted"], {}, tmp_path, require_cost=False) == {} - - -def test_openrouter_matches_hf_id_first_then_model_id(tmp_path, monkeypatch): - records = [ - { - "id": "provider/other", - "hugging_face_id": "org/model", - "pricing": {"prompt": "0.000002", "completion": "0.000007"}, - }, - { - "id": "org/fallback", - "pricing": {"prompt": "0.000003", "completion": "0.000009"}, - }, - ] - monkeypatch.setattr(pricing, "_fetch_openrouter", lambda path: records) - result = resolve_prices( - ["OpenRouter/org/model", "Provider/org/fallback"], {}, tmp_path - ) - assert result["OpenRouter/org/model"] == TokenPrice( - 2.0, 7.0, "openrouter_reference" - ) - assert result["Provider/org/fallback"] == TokenPrice( - 3.0, 9.0, "openrouter_reference" - ) - - -def test_nonhosted_unlisted_model_does_not_fetch(tmp_path, monkeypatch): - def fail(_): - raise AssertionError("unexpected hosted pricing fetch") - - monkeypatch.setattr(pricing, "_fetch_openrouter", fail) - assert resolve_prices(["VLLM/unlisted"], {}, tmp_path, require_cost=False) == {} - - -def test_legacy_openrouter_cache_mapping_is_supported(tmp_path): - cache_file = tmp_path / "openrouter_pricing.json" - cache_file.write_text(json.dumps({"org/model": [0.000002, 0.000007]})) - records = pricing._fetch_openrouter(cache_file) - assert pricing._hosted_price("org/model", records) == TokenPrice( - 2.0, 7.0, "openrouter_reference" - ) - - -def test_hf_cache_missing_sentinel_is_not_a_path(tmp_path, monkeypatch): - sentinel = object() - monkeypatch.setattr(pricing, "try_to_load_from_cache", lambda *a, **k: sentinel) - assert pricing._repo_snapshot("org/missing") is None From 6c1008400a07865a1eccd83594c4c5df71f6a4f7 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 9 Oct 2026 10:57:12 +0200 Subject: [PATCH 18/20] Use native request usage for cached and fresh judge cost objectives --- judgearena/pricing.py | 50 +++++++++++++++++ judgearena/tuning/pricing.py | 28 ---------- judgearena/tuning/runner.py | 105 +++++++++++++++++++++-------------- judgearena/tuning/session.py | 2 + tests/test_tuning.py | 94 ++++++++++++++++++++++++------- 5 files changed, 190 insertions(+), 89 deletions(-) create mode 100644 judgearena/pricing.py delete mode 100644 judgearena/tuning/pricing.py diff --git a/judgearena/pricing.py b/judgearena/pricing.py new file mode 100644 index 00000000..bd01eec4 --- /dev/null +++ b/judgearena/pricing.py @@ -0,0 +1,50 @@ +"""Shared reference token pricing and request cost calculation.""" + +from dataclasses import dataclass + +from judgearena.usage import RequestUsage + + +@dataclass(frozen=True) +class TokenPrice: + input: float + output: float + source: str + + +def reference_cost(usage: RequestUsage, price: TokenPrice) -> float: + """Calculate USD from native input/output token counts and per-million rates.""" + missing = [ + name + for name in ("input_tokens", "output_tokens") + if getattr(usage, name) is None + ] + if missing: + fields = " and ".join(missing) + raise ValueError( + f"Cannot calculate reference cost: missing native {fields}; " + "provide provider-reported input_tokens and output_tokens usage" + ) + return ( + usage.input_tokens * price.input + usage.output_tokens * price.output + ) / 1_000_000 + + +def resolve_prices(models, overrides, *, require_cost=True) -> dict[str, TokenPrice]: + """Resolve explicit USD-per-million-token rates before searching.""" + if not require_cost: + return {} + prices = {} + for model in sorted(models): + if model not in overrides: + continue + rate = overrides[model] + if isinstance(rate, (int, float)): + rate = {"input": rate, "output": rate} + prices[model] = TokenPrice(rate["input"], rate["output"], "user_override") + missing = sorted(set(models) - prices.keys()) + if missing: + raise ValueError( + f"No token price for {missing}; set tune_judge.price_per_million_tokens" + ) + return prices diff --git a/judgearena/tuning/pricing.py b/judgearena/tuning/pricing.py deleted file mode 100644 index be671c95..00000000 --- a/judgearena/tuning/pricing.py +++ /dev/null @@ -1,28 +0,0 @@ -"""Explicit judge-token prices used by the tuning cost objective.""" - -from dataclasses import dataclass - - -@dataclass(frozen=True) -class TokenPrice: - input: float - output: float - source: str - - -def resolve_prices(models, overrides, *, require_cost=True) -> dict[str, TokenPrice]: - """Resolve explicit USD-per-million-token rates before searching.""" - prices = {} - for model in sorted(models): - if model not in overrides: - continue - rate = overrides[model] - if isinstance(rate, (int, float)): - rate = {"input": rate, "output": rate} - prices[model] = TokenPrice(rate["input"], rate["output"], "user_override") - missing = sorted(set(models) - prices.keys()) - if missing and require_cost: - raise ValueError( - f"No token price for {missing}; set tune_judge.price_per_million_tokens" - ) - return prices diff --git a/judgearena/tuning/runner.py b/judgearena/tuning/runner.py index d6e59565..34ff0f87 100644 --- a/judgearena/tuning/runner.py +++ b/judgearena/tuning/runner.py @@ -9,13 +9,11 @@ import sys from collections.abc import Callable from dataclasses import asdict, dataclass -from functools import cache from pathlib import Path from typing import TYPE_CHECKING import numpy as np import pandas as pd -import tiktoken from judgearena.config import RunConfig, dump_config from judgearena.log import get_logger @@ -30,8 +28,9 @@ if TYPE_CHECKING: from judgearena.tasks.schema import ResolvedTaskSpec -from judgearena.tuning.pricing import TokenPrice, resolve_prices +from judgearena.pricing import TokenPrice, reference_cost, resolve_prices from judgearena.tuning.session import collect_trials, prepare_session +from judgearena.usage import request_usage_from_json logger = get_logger(__name__) @@ -49,35 +48,41 @@ def run_trial_subprocess(config_path: Path) -> None: ) -@cache -def _judge_token_encoding() -> tiktoken.Encoding: - return tiktoken.encoding_for_model("gpt-4o") - - def _read_trial(trial_dir: Path, price: TokenPrice | None) -> dict: (result_path,) = trial_dir.glob("*/results.json") metrics = json.loads(result_path.read_text())["metrics"] agreement = metrics["meta_eval_agreement"]["all"] - annotations = pd.read_parquet( - result_path.parent / "annotations.parquet", - columns=["battle_id", "judge_input", "judge_completion"], - ) - encoding = _judge_token_encoding() - input_tokens = sum( - len(encoding.encode(text, disallowed_special=())) - for text in annotations["judge_input"] - ) - output_tokens = sum( - len(encoding.encode(text, disallowed_special=())) - for text in annotations["judge_completion"].fillna("") - ) + annotations = pd.read_parquet(result_path.parent / "annotations.parquet") n_battles = annotations["battle_id"].nunique() - tokens_per_battle = (input_tokens + output_tokens) / n_battles - cost = ( - (input_tokens * price.input + output_tokens * price.output) / n_battles / 1000 - if price - else None + total_input = total_output = 0 + total_cost = 0.0 + has_usage_columns = {"usage_json", "error"}.issubset(annotations.columns) + has_complete_usage = has_usage_columns + if has_usage_columns: + for row in annotations.itertuples(index=False): + request_usage = request_usage_from_json(row.usage_json) + if request_usage is None and row.error == "context_length": + continue + input_tokens = request_usage.input_tokens if request_usage else None + output_tokens = request_usage.output_tokens if request_usage else None + if input_tokens is None or output_tokens is None: + has_complete_usage = False + continue + total_input += input_tokens + total_output += output_tokens + if price is not None: + total_cost += reference_cost(request_usage, price) + if price is not None and not has_complete_usage: + raise ValueError( + "Cannot calculate tuning cost: annotations are missing native input/output " + "token usage in " + f"{result_path.parent / 'annotations.parquet'}. Rerun with a fresh " + "store_root so usage is recorded, or use agreement-only tuning." + ) + tokens_per_battle = ( + (total_input + total_output) / n_battles if has_complete_usage else None ) + cost = total_cost / n_battles * 1000 if price is not None else None return { "agreement": agreement["accuracy_attempted"], "cohen_kappa": agreement["cohen_kappa"], @@ -99,6 +104,7 @@ class _TrialRunner: base: dict prices: dict[str, TokenPrice] execute_trial: TrialExecutor + ignore_failed_trials: bool = False def config( self, @@ -149,6 +155,8 @@ def run( try: self.execute_trial(config_path) except subprocess.CalledProcessError as exc: + if not self.ignore_failed_trials: + raise logger.warning("Trial %s failed: %s", folder, exc) return {**record, "status": "failed"} price = self.prices.get(trial_cfg.judge.model) @@ -158,7 +166,6 @@ def run( def _objective(record: dict, objectives: list[str]) -> dict: if record["status"] != "completed": return { - "exception": RuntimeError(f"Trial {record['config_id']} failed"), "objective_to_minimize": ( float("inf") if len(objectives) == 1 @@ -191,6 +198,7 @@ def evaluate(pipeline_directory: Path, **config): return _objective(record, cfg.tune_judge.objectives) options = dict(cfg.tune_judge.neps) + options.pop("ignore_errors", None) # NePS 0.17's space compatibility check does not inspect optimizer mappings. optimizer = cfg.tune_judge.optimizer if isinstance(optimizer, dict): @@ -214,6 +222,8 @@ def evaluate(pipeline_directory: Path, **config): pipeline_space=None if continuing else space, root_directory=root, continue_until_max_evaluation_completed=False, + # Only subprocess failures may be ignored; accounting errors stop the run. + ignore_errors=False, **options, ) @@ -253,21 +263,34 @@ def run_tune_judge( models.add(trial_cfg.judge.model) if "judge.model" not in specs: models.add(cfg.judge.model) - prices = resolve_prices( - models, - tuning.price_per_million_tokens, - require_cost="cost" in tuning.objectives, - ) - runner = _TrialRunner(runner.base, prices, execute_trial) tune_dir = prepare_session(cfg) - price_path = tune_dir / "prices.json" - snapshot = {model: asdict(price) for model, price in prices.items()} - if price_path.exists(): - stored = json.loads(price_path.read_text()) - if stored != snapshot: - raise ValueError("Resolved prices changed for this tuning session") - elif not tuning.search_only: - price_path.write_text(json.dumps(snapshot, indent=2)) + prices = {} + if "cost" in tuning.objectives: + price_path = tune_dir / "prices.json" + if price_path.exists(): + prices = { + model: TokenPrice(**price) + for model, price in json.loads(price_path.read_text()).items() + } + else: + prices = resolve_prices( + models, + tuning.price_per_million_tokens, + require_cost=True, + ) + if not tuning.search_only: + price_path.write_text( + json.dumps( + {model: asdict(price) for model, price in prices.items()}, + indent=2, + ) + ) + runner = _TrialRunner( + runner.base, + prices, + execute_trial, + ignore_failed_trials=tuning.neps.get("ignore_errors", False), + ) logger.info("Tuning %s in %s", cfg.task, tune_dir) _search(runner, cfg, tune_dir, space) if tuning.search_only: diff --git a/judgearena/tuning/session.py b/judgearena/tuning/session.py index 931d388f..b2fda494 100644 --- a/judgearena/tuning/session.py +++ b/judgearena/tuning/session.py @@ -19,6 +19,8 @@ def _config_hash(cfg: RunConfig) -> str: payload = cfg.model_dump(mode="json") payload["run"] = {"seed": cfg.run.seed} tuning = payload["tune_judge"] + if "cost" in tuning["objectives"]: + payload["cost_accounting"] = "native_usage_v1" tuning.pop("run_dir") tuning.pop("search_only") tuning["neps"] = { diff --git a/tests/test_tuning.py b/tests/test_tuning.py index dff53083..9e1788d9 100644 --- a/tests/test_tuning.py +++ b/tests/test_tuning.py @@ -63,15 +63,6 @@ def _tuning_config(tmp_path, algorithm="neps_priorband", **settings): ) -@pytest.fixture -def token_counter(monkeypatch): - monkeypatch.setattr( - runner, - "_judge_token_encoding", - lambda: SimpleNamespace(encode=lambda text, **kwargs: text.split()), - ) - - def _fake_trial(config_path): cfg = load_config(config_path) folder = config_path.parent / "run" @@ -90,6 +81,11 @@ def _fake_trial(config_path): "battle_id": ["b1", "b1"], "judge_input": ["system forward", "system reverse"], "judge_completion": ["answer", "answer"], + "usage_json": [ + json.dumps({"stage": "judging", "input_tokens": 2, "output_tokens": 1}), + json.dumps({"stage": "judging", "input_tokens": 2, "output_tokens": 1}), + ], + "error": [None, None], "cache_hit": [True, True], } ).to_parquet(folder / "annotations.parquet") @@ -98,7 +94,7 @@ def _fake_trial(config_path): @pytest.mark.parametrize( "algorithm", ["neps_priorband", "neps_hyperband", "mo_hyperband", "primo"] ) -def test_native_optimizer_smoke(tmp_path, token_counter, algorithm): +def test_native_optimizer_smoke(tmp_path, algorithm, monkeypatch): pytest.importorskip("neps") cfg = _tuning_config(tmp_path, algorithm) calls = [] @@ -127,11 +123,16 @@ def execute(path): assert load_config(calls[-1]).meta_eval.battles_per_model == 3 before = len(calls) (folder / "test_results.parquet").unlink() - resumed = runner.run_tune_judge(cfg, task, execute_trial=execute) - pd.testing.assert_frame_equal(results, resumed) - assert len(calls) == before # Completed search and held-out artifacts are reused. helper = _tuning_config(tmp_path, search_only=True) - assert runner.run_tune_judge(helper, task, execute_trial=execute).empty + with monkeypatch.context() as patch: + patch.setattr( + runner, + "resolve_prices", + lambda *a, **kw: pytest.fail("Refetched session prices"), + ) + resumed = runner.run_tune_judge(cfg, task, execute_trial=execute) + assert runner.run_tune_judge(helper, task, execute_trial=execute).empty + pd.testing.assert_frame_equal(results, resumed) assert len(calls) == before incompatible = _tuning_config( tmp_path, price_per_million_tokens={"Dummy/judge": 2.0, "Dummy/other": 2.0} @@ -161,9 +162,7 @@ def execute(path): ) -def test_budget_end_does_not_wait_for_unstarted_batch( - tmp_path, token_counter, monkeypatch -): +def test_budget_end_does_not_wait_for_unstarted_batch(tmp_path, monkeypatch): pytest.importorskip("neps") cfg = _tuning_config(tmp_path) cfg.tune_judge.optimizer = { @@ -185,7 +184,7 @@ def test_budget_end_does_not_wait_for_unstarted_batch( assert len(pd.read_parquet(cfg.tune_judge.run_dir / "trials.parquet")) == 2 -def test_cost_counts_both_cached_orientations(tmp_path, token_counter): +def test_cost_counts_both_cached_orientations(tmp_path): path = tmp_path / "config.yaml" dump_config(_tuning_config(tmp_path), path) _fake_trial(path) @@ -199,7 +198,58 @@ def test_cost_counts_both_cached_orientations(tmp_path, token_counter): assert "cost" not in objective -def test_prices_required_before_execution(tmp_path): +def test_agreement_only_reads_old_cached_artifacts(tmp_path): + path = tmp_path / "config.yaml" + dump_config(_tuning_config(tmp_path), path) + _fake_trial(path) + annotations_path = tmp_path / "run" / "annotations.parquet" + pd.read_parquet(annotations_path).drop(columns=["usage_json", "error"]).to_parquet( + annotations_path, index=False + ) + + result = runner._read_trial(tmp_path, price=None) + assert result["agreement"] == pytest.approx(0.58) + assert result["tokens_per_battle"] is None + assert result["cost_per_1k_battles"] is None + + +def test_missing_native_usage_aborts_even_when_failed_trials_are_ignored(tmp_path): + pytest.importorskip("neps") + cfg = _tuning_config(tmp_path, "mo_hyperband") + cfg.tune_judge.neps.update(ignore_errors=True, total_evaluations_to_spend=1) + + def execute(path): + _fake_trial(path) + annotations_path = path.parent / "run" / "annotations.parquet" + annotations = pd.read_parquet(annotations_path) + annotations["usage_json"] = None + annotations.to_parquet(annotations_path, index=False) + + from neps.exceptions import WorkerRaiseError + + with pytest.raises(WorkerRaiseError) as exc_info: + runner.run_tune_judge(cfg, get_packaged_task(cfg.task), execute_trial=execute) + assert isinstance(exc_info.value.__cause__, ValueError) + assert "fresh store_root" in str(exc_info.value.__cause__) + assert "agreement-only" in str(exc_info.value.__cause__) + + +def test_skipped_request_is_free_but_unparsed_response_keeps_cost(tmp_path): + path = tmp_path / "config.yaml" + dump_config(_tuning_config(tmp_path), path) + _fake_trial(path) + annotations_path = tmp_path / "run" / "annotations.parquet" + annotations = pd.read_parquet(annotations_path) + annotations.loc[0, ["usage_json", "error"]] = [None, "context_length"] + annotations.loc[1, "judge_completion"] = "unparseable" + annotations.to_parquet(annotations_path, index=False) + record = runner._read_trial(tmp_path, runner.TokenPrice(2, 4, "test")) + assert record["tokens_per_battle"] == 3 + assert record["cost_per_1k_battles"] == pytest.approx(0.008) + + +def test_prices_required_before_execution(tmp_path, monkeypatch): + pytest.importorskip("neps") cfg = _tuning_config(tmp_path, "mo_hyperband", price_per_million_tokens={}) with pytest.raises(ValueError, match="No token price"): @@ -211,7 +261,7 @@ def test_prices_required_before_execution(tmp_path): @pytest.mark.parametrize("algorithm", ["neps_priorband", "mo_hyperband"]) -def test_failed_trial_does_not_stop_search(tmp_path, token_counter, algorithm): +def test_failed_trial_does_not_stop_search(tmp_path, algorithm): pytest.importorskip("neps") cfg = _tuning_config(tmp_path, algorithm) cfg.tune_judge.neps["ignore_errors"] = True @@ -231,6 +281,10 @@ def execute(path): assert (trials.status == "failed").sum() == 1 assert not results.empty and set(results.status) == {"completed"} pareto = pd.read_parquet(cfg.tune_judge.run_dir / "pareto.parquet") + if algorithm == "neps_priorband": + assert pareto.empty + assert "cost" not in cfg.tune_judge.objectives + return final = trials[(trials.status == "completed") & (trials.battles_per_model == 3)] assert set(pareto.agreement) == {final.agreement.max()} assert set(pareto.status) == {"completed"} From 5abc6a8d4e8b888b479dd69c9d19c8eb65d29cbe Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 9 Oct 2026 16:48:31 +0200 Subject: [PATCH 19/20] Consolidate tuning integration tests and isolate session contracts --- tests/test_tuning.py | 283 ++++++++++++++++++++----------------------- 1 file changed, 129 insertions(+), 154 deletions(-) diff --git a/tests/test_tuning.py b/tests/test_tuning.py index 9e1788d9..eb737de9 100644 --- a/tests/test_tuning.py +++ b/tests/test_tuning.py @@ -9,16 +9,14 @@ import pandas as pd import pytest -from judgearena.config import RunConfig, dump_config, load_config +from judgearena.config import RunConfig, load_config from judgearena.tasks.registry import get_packaged_task from judgearena.tuning import runner def _tuning_config(tmp_path, algorithm="neps_priorband", **settings): + multi = algorithm in {"mo_hyperband", "primo"} optimizer = {"name": algorithm, "eta": 3} - objectives = ["agreement"] - if algorithm in {"mo_hyperband", "primo"}: - objectives.append("cost") if algorithm == "primo": optimizer.update( initial_design_size=2, @@ -27,100 +25,135 @@ def _tuning_config(tmp_path, algorithm="neps_priorband", **settings): for i in range(2) }, prior_centers={ - f"objective_{i}": { - "judge.temperature": value, - "judge.model": "Dummy/judge", - } - for i, value in enumerate((0.2, 0.8)) + f"objective_{i}": {"judge.temperature": v, "judge.model": "Dummy/judge"} + for i, v in enumerate((0.2, 0.8)) }, ) - specs = { - "judge": { - "temperature": {"lower": 0.0, "upper": 1.0}, - "model": { - "choices": ["Dummy/judge", "Dummy/other"], - "prior": "Dummy/other", - }, - } - } return RunConfig( task="tune-judge", judge={"model": "Dummy/judge", "temperature": 0.2, "swap_mode": "both"}, run={"result_folder": str(tmp_path), "no_log_file": True}, tune_judge={ "meta_eval_task": "meta-eval-comparia-fr", - "neps": { - "total_evaluations_to_spend": 10, + "neps": {"total_evaluations_to_spend": 10}, + "search_space": { + "judge": { + "temperature": {"lower": 0.0, "upper": 1.0}, + "model": ["Dummy/judge", "Dummy/other"], + } }, - "search_space": specs, "fidelity": {"battles_per_model": {"lower": 1, "upper": 3}}, "optimizer": optimizer, - "objectives": objectives, - "price_per_million_tokens": {"Dummy/judge": 1.0, "Dummy/other": 1.0}, + "objectives": ["agreement", "cost"] if multi else ["agreement"], + "price_per_million_tokens": dict.fromkeys( + ["Dummy/judge", "Dummy/other"], 1.0 + ), "run_dir": tmp_path / "tuning", **settings, }, ) -def _fake_trial(config_path): - cfg = load_config(config_path) - folder = config_path.parent / "run" +def _write_artifacts(folder, agreement, battles): folder.mkdir() - agreement = 0.6 - 0.1 * cfg.judge.temperature - agreement_metrics = dict( + scores = dict( accuracy_attempted=agreement, cohen_kappa=agreement, coverage=1.0, - n_attempted=cfg.meta_eval.battles_per_model, + n_attempted=battles, + ) + (folder / "results.json").write_text( + json.dumps({"metrics": {"meta_eval_agreement": {"all": scores}}}) ) - metrics = {"meta_eval_agreement": {"all": agreement_metrics}} - (folder / "results.json").write_text(json.dumps({"metrics": metrics})) + usage = json.dumps({"stage": "judging", "input_tokens": 2, "output_tokens": 1}) pd.DataFrame( - { - "battle_id": ["b1", "b1"], - "judge_input": ["system forward", "system reverse"], - "judge_completion": ["answer", "answer"], - "usage_json": [ - json.dumps({"stage": "judging", "input_tokens": 2, "output_tokens": 1}), - json.dumps({"stage": "judging", "input_tokens": 2, "output_tokens": 1}), - ], - "error": [None, None], - "cache_hit": [True, True], - } + {"battle_id": ["b1", "b1"], "usage_json": [usage, usage], "error": [None, None]} ).to_parquet(folder / "annotations.parquet") +def _fake_trial(config_path): + cfg = load_config(config_path) + _write_artifacts( + config_path.parent / "run", + 0.6 - 0.1 * cfg.judge.temperature, + cfg.meta_eval.battles_per_model, + ) + + +def _run(cfg, execute=_fake_trial): + return runner.run_tune_judge( + cfg, get_packaged_task(cfg.task), execute_trial=execute + ) + + +@pytest.fixture +def artifacts(tmp_path): + _write_artifacts(tmp_path / "run", 0.58, 1) + return tmp_path + + @pytest.mark.parametrize( - "algorithm", ["neps_priorband", "neps_hyperband", "mo_hyperband", "primo"] + "algorithm,fail_first", + [ + ("neps_priorband", False), + ("neps_hyperband", False), + ("mo_hyperband", False), + ("primo", False), + ("neps_priorband", True), + ("mo_hyperband", True), + ], ) -def test_native_optimizer_smoke(tmp_path, algorithm, monkeypatch): +def test_optimizer_searches_validation_then_tests_full_fidelity( + tmp_path, algorithm, fail_first +): pytest.importorskip("neps") cfg = _tuning_config(tmp_path, algorithm) + cfg.tune_judge.neps["ignore_errors"] = True calls = [] def execute(path): - calls.append(path) - _fake_trial(path) - - task = get_packaged_task(cfg.task) - results = runner.run_tune_judge(cfg, task, execute_trial=execute) - assert not results.empty and set(results.status) == {"completed"} - if algorithm != "neps_priorband": - return - folder = cfg.tune_judge.run_dir - trials = pd.read_parquet(folder / "trials.parquet") - assert trials.battles_per_model.max() == 3 - assert trials.neps_trial_id.is_unique - best = trials[trials.battles_per_model == 3].groupby("judge_model").agreement.max() - assert results.set_index("judge_model").agreement.to_dict() == best.to_dict() - for path in calls: trial = load_config(path) assert trial.task == "meta-eval-comparia-fr" and trial.tune_judge is None assert trial.meta_eval.split == ( "test" if "test" in path.parts else "validation" ) + calls.append(path) + if fail_first and len(calls) == 1: + raise subprocess.CalledProcessError(1, ["judge-trial"]) + _fake_trial(path) + + results = _run(cfg, execute) + folder = cfg.tune_judge.run_dir + trials = pd.read_parquet(folder / "trials.parquet") + assert len(trials) == 10 and trials.neps_trial_id.is_unique + assert (trials.status == "failed").sum() == int(fail_first) + final = trials[(trials.status == "completed") & (trials.battles_per_model == 3)] + assert not results.empty and set(results.status) == {"completed"} + assert ( + results.set_index("judge_model").agreement.to_dict() + == final.groupby("judge_model").agreement.max().to_dict() + ) assert load_config(calls[-1]).meta_eval.battles_per_model == 3 + pareto = pd.read_parquet(folder / "pareto.parquet") + if "cost" in cfg.tune_judge.objectives: + assert set(pareto.agreement) == {final.agreement.max()} + assert set(pareto.battles_per_model) == {3} + else: + assert pareto.empty + + +def test_session_resume_reuses_trials_and_prices(tmp_path, monkeypatch): + pytest.importorskip("neps") + cfg = _tuning_config(tmp_path) + calls = [] + + def execute(path): + calls.append(path) + _fake_trial(path) + + results = _run(cfg, execute) + folder = cfg.tune_judge.run_dir + trials = pd.read_parquet(folder / "trials.parquet") before = len(calls) (folder / "test_results.parquet").unlink() helper = _tuning_config(tmp_path, search_only=True) @@ -130,28 +163,26 @@ def execute(path): "resolve_prices", lambda *a, **kw: pytest.fail("Refetched session prices"), ) - resumed = runner.run_tune_judge(cfg, task, execute_trial=execute) - assert runner.run_tune_judge(helper, task, execute_trial=execute).empty + resumed = _run(cfg, execute) + assert _run(helper, execute).empty pd.testing.assert_frame_equal(results, resumed) assert len(calls) == before - incompatible = _tuning_config( - tmp_path, price_per_million_tokens={"Dummy/judge": 2.0, "Dummy/other": 2.0} - ) - with pytest.raises(ValueError, match="Incompatible configuration"): - runner.run_tune_judge(incompatible, task, execute_trial=execute) - changed_target = _tuning_config(tmp_path, meta_eval_task="meta-eval-comparia-en") - with pytest.raises(ValueError, match="Incompatible configuration"): - runner.run_tune_judge(changed_target, task, execute_trial=execute) + for change in ( + {"price_per_million_tokens": {"Dummy/judge": 2.0, "Dummy/other": 2.0}}, + {"meta_eval_task": "meta-eval-comparia-en"}, + ): + with pytest.raises(ValueError, match="Incompatible configuration"): + _run(_tuning_config(tmp_path, **change), execute) assert len(calls) == before fresh = _tuning_config(tmp_path, run_dir=tmp_path / "fresh") - runner.run_tune_judge(fresh, task, execute_trial=execute) + _run(fresh, execute) repeated = pd.read_parquet(fresh.tune_judge.run_dir / "trials.parquet") assert repeated.overrides.tolist() == trials.overrides.tolist() extended = _tuning_config(tmp_path) extended.tune_judge.neps["total_evaluations_to_spend"] = 11 - runner.run_tune_judge(extended, task, execute_trial=execute) + _run(extended, execute) assert len(pd.read_parquet(folder / "trials.parquet")) == 11 - runner.run_tune_judge(helper, task, execute_trial=execute) + _run(helper, execute) from neps.state import NePSState assert ( @@ -177,17 +208,13 @@ def test_budget_end_does_not_wait_for_unstarted_batch(tmp_path, monkeypatch): "judgearena.tuning.session.time.sleep", lambda _: pytest.fail("Waiting for an unstarted trial"), ) - results = runner.run_tune_judge( - cfg, get_packaged_task(cfg.task), execute_trial=_fake_trial - ) + results = _run(cfg) assert not results.empty assert len(pd.read_parquet(cfg.tune_judge.run_dir / "trials.parquet")) == 2 -def test_cost_counts_both_cached_orientations(tmp_path): - path = tmp_path / "config.yaml" - dump_config(_tuning_config(tmp_path), path) - _fake_trial(path) +def test_cost_counts_both_cached_orientations(artifacts): + tmp_path = artifacts record = runner._read_trial(tmp_path, price=runner.TokenPrice(2.0, 4.0, "test")) assert record["tokens_per_battle"] == 6 assert record["cost_per_1k_battles"] == pytest.approx(0.016) @@ -198,13 +225,10 @@ def test_cost_counts_both_cached_orientations(tmp_path): assert "cost" not in objective -def test_agreement_only_reads_old_cached_artifacts(tmp_path): - path = tmp_path / "config.yaml" - dump_config(_tuning_config(tmp_path), path) - _fake_trial(path) - annotations_path = tmp_path / "run" / "annotations.parquet" - pd.read_parquet(annotations_path).drop(columns=["usage_json", "error"]).to_parquet( - annotations_path, index=False +def test_agreement_only_reads_old_cached_artifacts(artifacts): + tmp_path = artifacts + pd.DataFrame({"battle_id": ["b1"]}).to_parquet( + tmp_path / "run" / "annotations.parquet" ) result = runner._read_trial(tmp_path, price=None) @@ -228,16 +252,14 @@ def execute(path): from neps.exceptions import WorkerRaiseError with pytest.raises(WorkerRaiseError) as exc_info: - runner.run_tune_judge(cfg, get_packaged_task(cfg.task), execute_trial=execute) + _run(cfg, execute) assert isinstance(exc_info.value.__cause__, ValueError) assert "fresh store_root" in str(exc_info.value.__cause__) assert "agreement-only" in str(exc_info.value.__cause__) -def test_skipped_request_is_free_but_unparsed_response_keeps_cost(tmp_path): - path = tmp_path / "config.yaml" - dump_config(_tuning_config(tmp_path), path) - _fake_trial(path) +def test_skipped_request_is_free_but_unparsed_response_keeps_cost(artifacts): + tmp_path = artifacts annotations_path = tmp_path / "run" / "annotations.parquet" annotations = pd.read_parquet(annotations_path) annotations.loc[0, ["usage_json", "error"]] = [None, "context_length"] @@ -248,71 +270,26 @@ def test_skipped_request_is_free_but_unparsed_response_keeps_cost(tmp_path): assert record["cost_per_1k_battles"] == pytest.approx(0.008) -def test_prices_required_before_execution(tmp_path, monkeypatch): - +def test_prices_required_before_execution(tmp_path): pytest.importorskip("neps") cfg = _tuning_config(tmp_path, "mo_hyperband", price_per_million_tokens={}) with pytest.raises(ValueError, match="No token price"): - runner.run_tune_judge( - cfg, - get_packaged_task(cfg.task), - execute_trial=lambda _: pytest.fail("executed without prices"), - ) - - -@pytest.mark.parametrize("algorithm", ["neps_priorband", "mo_hyperband"]) -def test_failed_trial_does_not_stop_search(tmp_path, algorithm): - pytest.importorskip("neps") - cfg = _tuning_config(tmp_path, algorithm) - cfg.tune_judge.neps["ignore_errors"] = True - calls = [] - - def execute(path): - calls.append(path) - if len(calls) == 1: - raise subprocess.CalledProcessError(1, ["judge-trial"]) - _fake_trial(path) - - results = runner.run_tune_judge( - cfg, get_packaged_task(cfg.task), execute_trial=execute - ) - trials = pd.read_parquet(cfg.tune_judge.run_dir / "trials.parquet") - assert len(trials) == 10 - assert (trials.status == "failed").sum() == 1 - assert not results.empty and set(results.status) == {"completed"} - pareto = pd.read_parquet(cfg.tune_judge.run_dir / "pareto.parquet") - if algorithm == "neps_priorband": - assert pareto.empty - assert "cost" not in cfg.tune_judge.objectives - return - final = trials[(trials.status == "completed") & (trials.battles_per_model == 3)] - assert set(pareto.agreement) == {final.agreement.max()} - assert set(pareto.status) == {"completed"} - assert set(pareto.battles_per_model) == {3} + _run(cfg, lambda _: pytest.fail("executed without prices")) def test_pareto_report_uses_completed_full_fidelity_validation(tmp_path): pytest.importorskip("neps") trials = pd.DataFrame( [ - dict( - config_id=name, - judge_model=name, - agreement=agreement, - cost_per_1k_battles=cost, - battles_per_model=fidelity, - status=status, - overrides="{}", - ) - for name, agreement, cost, fidelity, status in [ - ("cheap", 0.5, 1.0, 3, "completed"), - ("accurate", 0.8, 2.0, 3, "completed"), - ("dominated", 0.4, 3.0, 3, "completed"), - ("low-fidelity", 0.9, 0.1, 1, "completed"), - ("failed", 1.0, 0.0, 3, "failed"), - ] - ] + ("cheap", 0.5, 1.0, 3, "completed"), + ("accurate", 0.8, 2.0, 3, "completed"), + ("dominated", 0.4, 3.0, 3, "completed"), + ("low-fidelity", 0.9, 0.1, 1, "completed"), + ("failed", 1.0, 0.0, 3, "failed"), + ], + columns="config_id agreement cost_per_1k_battles battles_per_model status".split(), ) + trials = trials.assign(judge_model=trials.config_id, overrides="{}") executor = SimpleNamespace(run=lambda *args, **kwargs: {"status": "completed"}) runner._evaluate_picks( executor, SimpleNamespace(test_battles_per_model=None), tmp_path, trials, 3 @@ -333,16 +310,14 @@ def test_generic_route_uses_target_defaults(tmp_path): assert outer.meta_eval == direct.meta_eval resolved = resolve_benchmark(outer.task) assert resolved.adapter.name == "tune_judge" and resolved.task.spec.dataset is None - for target in ("unknown-task", "mt-bench"): + for target, judge in ( + ("unknown-task", {"model": "Dummy/judge"}), + ("mt-bench", {"model": "Dummy/judge"}), + (tuning["meta_eval_task"], {"model": "Dummy/judge", "swap_mode": "random"}), + ): with pytest.raises(ValueError): RunConfig( task="tune-judge", - judge={"model": "Dummy/judge"}, + judge=judge, tune_judge={**tuning, "meta_eval_task": target}, ) - with pytest.raises(ValueError): - RunConfig( - task="tune-judge", - judge={"model": "Dummy/judge", "swap_mode": "random"}, - tune_judge=tuning, - ) From befae240608e53701cdc741dabb917decdc6a7d2 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:09:10 +0200 Subject: [PATCH 20/20] Halve tuning test scope with representative integration coverage --- tests/test_chat_vllm.py | 13 +- tests/test_tuning.py | 356 ++++++++++++---------------------------- tests/test_usage.py | 22 +-- 3 files changed, 110 insertions(+), 281 deletions(-) diff --git a/tests/test_chat_vllm.py b/tests/test_chat_vllm.py index 42ed8dec..f6c202f5 100644 --- a/tests/test_chat_vllm.py +++ b/tests/test_chat_vllm.py @@ -342,15 +342,7 @@ def test_chat_vllm_skips_overlong_requests_without_caching(tmp_path, monkeypatch backend = models.ChatVLLM(model="test/model", max_tokens=128) monkeypatch.setattr(models, "make_model", lambda *_args, **_kwargs: backend) cache = JudgementInferenceCache(tmp_path, "meta-eval-test") - metadata = [ - { - "instruction_id": str(i), - "model_a": "a", - "model_b": "b", - "orientation": "direct", - } - for i in range(2) - ] + metadata = [dict(instruction_id=str(i), model_a="a", model_b="b") for i in range(2)] prepared = models.prepare_model("Dummy/judge", cache=cache) with track_usage() as tracker: assert models.do_inference( @@ -362,9 +354,6 @@ def test_chat_vllm_skips_overlong_requests_without_caching(tmp_path, monkeypatch None, ] assert captured["chat_call"]["messages"] == [[{"role": "user", "content": "valid"}]] - assert models.do_inference( - prepared, ["overlong", "valid"], cache_row_metadata=metadata - ) == ["", "ok"] import sqlite3 db = next(tmp_path.rglob("judgements.db")) diff --git a/tests/test_tuning.py b/tests/test_tuning.py index eb737de9..1bd6e4d1 100644 --- a/tests/test_tuning.py +++ b/tests/test_tuning.py @@ -1,6 +1,4 @@ -"""Native NePS execution and persisted trial artifacts.""" - -from __future__ import annotations +"""Representative tuning execution, accounting, and selection contracts.""" import json import subprocess @@ -9,315 +7,167 @@ import pandas as pd import pytest -from judgearena.config import RunConfig, load_config -from judgearena.tasks.registry import get_packaged_task +from judgearena.config import RunConfig, dump_config, load_config from judgearena.tuning import runner -def _tuning_config(tmp_path, algorithm="neps_priorband", **settings): - multi = algorithm in {"mo_hyperband", "primo"} - optimizer = {"name": algorithm, "eta": 3} - if algorithm == "primo": - optimizer.update( - initial_design_size=2, - prior_confidences={ - f"objective_{i}": {"judge.temperature": 0.25, "judge.model": 0.25} - for i in range(2) - }, - prior_centers={ - f"objective_{i}": {"judge.temperature": v, "judge.model": "Dummy/judge"} - for i, v in enumerate((0.2, 0.8)) - }, - ) +@pytest.fixture +def config(tmp_path): return RunConfig( task="tune-judge", - judge={"model": "Dummy/judge", "temperature": 0.2, "swap_mode": "both"}, + judge={"model": "Dummy/judge", "temperature": 0.0}, run={"result_folder": str(tmp_path), "no_log_file": True}, tune_judge={ "meta_eval_task": "meta-eval-comparia-fr", - "neps": {"total_evaluations_to_spend": 10}, - "search_space": { - "judge": { - "temperature": {"lower": 0.0, "upper": 1.0}, - "model": ["Dummy/judge", "Dummy/other"], - } - }, + "search_space": {"judge": {"temperature": [0.0, 1.0]}}, "fidelity": {"battles_per_model": {"lower": 1, "upper": 3}}, - "optimizer": optimizer, - "objectives": ["agreement", "cost"] if multi else ["agreement"], - "price_per_million_tokens": dict.fromkeys( - ["Dummy/judge", "Dummy/other"], 1.0 - ), + "neps": {"total_evaluations_to_spend": 10, "ignore_errors": True}, + "price_per_million_tokens": {"Dummy/judge": 1.0}, "run_dir": tmp_path / "tuning", - **settings, }, ) -def _write_artifacts(folder, agreement, battles): +def execute(path): + cfg = load_config(path) + assert cfg.task == "meta-eval-comparia-fr" and cfg.tune_judge is None + assert cfg.meta_eval.split == ("test" if "test" in path.parts else "validation") + folder = path.parent / "run" folder.mkdir() scores = dict( - accuracy_attempted=agreement, - cohen_kappa=agreement, + accuracy_attempted=0.6 - cfg.judge.temperature / 10, + cohen_kappa=0.2, coverage=1.0, - n_attempted=battles, + n_attempted=cfg.meta_eval.battles_per_model, ) (folder / "results.json").write_text( json.dumps({"metrics": {"meta_eval_agreement": {"all": scores}}}) ) usage = json.dumps({"stage": "judging", "input_tokens": 2, "output_tokens": 1}) pd.DataFrame( - {"battle_id": ["b1", "b1"], "usage_json": [usage, usage], "error": [None, None]} + {"battle_id": ["b1", "b1"], "usage_json": [usage] * 2, "error": [None] * 2} ).to_parquet(folder / "annotations.parquet") -def _fake_trial(config_path): - cfg = load_config(config_path) - _write_artifacts( - config_path.parent / "run", - 0.6 - 0.1 * cfg.judge.temperature, - cfg.meta_eval.battles_per_model, - ) - - -def _run(cfg, execute=_fake_trial): - return runner.run_tune_judge( - cfg, get_packaged_task(cfg.task), execute_trial=execute - ) - - -@pytest.fixture -def artifacts(tmp_path): - _write_artifacts(tmp_path / "run", 0.58, 1) - return tmp_path +def run(config, executor=execute): + pytest.importorskip("neps") + return runner.run_tune_judge(config, None, execute_trial=executor) -@pytest.mark.parametrize( - "algorithm,fail_first", - [ - ("neps_priorband", False), - ("neps_hyperband", False), - ("mo_hyperband", False), - ("primo", False), - ("neps_priorband", True), - ("mo_hyperband", True), - ], -) -def test_optimizer_searches_validation_then_tests_full_fidelity( - tmp_path, algorithm, fail_first -): - pytest.importorskip("neps") - cfg = _tuning_config(tmp_path, algorithm) - cfg.tune_judge.neps["ignore_errors"] = True +def test_search_failure_resume_and_held_out_evaluation(config, monkeypatch): calls = [] - def execute(path): - trial = load_config(path) - assert trial.task == "meta-eval-comparia-fr" and trial.tune_judge is None - assert trial.meta_eval.split == ( - "test" if "test" in path.parts else "validation" - ) + def trial(path): calls.append(path) - if fail_first and len(calls) == 1: + if len(calls) == 1: raise subprocess.CalledProcessError(1, ["judge-trial"]) - _fake_trial(path) + execute(path) - results = _run(cfg, execute) - folder = cfg.tune_judge.run_dir - trials = pd.read_parquet(folder / "trials.parquet") + results = run(config, trial) + root = config.tune_judge.run_dir + trials = pd.read_parquet(root / "trials.parquet") assert len(trials) == 10 and trials.neps_trial_id.is_unique - assert (trials.status == "failed").sum() == int(fail_first) + assert (trials.status == "failed").sum() == 1 final = trials[(trials.status == "completed") & (trials.battles_per_model == 3)] - assert not results.empty and set(results.status) == {"completed"} - assert ( - results.set_index("judge_model").agreement.to_dict() - == final.groupby("judge_model").agreement.max().to_dict() - ) + assert results.agreement.tolist() == [final.agreement.max()] + assert set(results.status) == {"completed"} assert load_config(calls[-1]).meta_eval.battles_per_model == 3 - pareto = pd.read_parquet(folder / "pareto.parquet") - if "cost" in cfg.tune_judge.objectives: - assert set(pareto.agreement) == {final.agreement.max()} - assert set(pareto.battles_per_model) == {3} - else: - assert pareto.empty - - -def test_session_resume_reuses_trials_and_prices(tmp_path, monkeypatch): - pytest.importorskip("neps") - cfg = _tuning_config(tmp_path) - calls = [] - - def execute(path): - calls.append(path) - _fake_trial(path) - - results = _run(cfg, execute) - folder = cfg.tune_judge.run_dir - trials = pd.read_parquet(folder / "trials.parquet") - before = len(calls) - (folder / "test_results.parquet").unlink() - helper = _tuning_config(tmp_path, search_only=True) - with monkeypatch.context() as patch: - patch.setattr( - runner, - "resolve_prices", - lambda *a, **kw: pytest.fail("Refetched session prices"), - ) - resumed = _run(cfg, execute) - assert _run(helper, execute).empty - pd.testing.assert_frame_equal(results, resumed) - assert len(calls) == before - for change in ( - {"price_per_million_tokens": {"Dummy/judge": 2.0, "Dummy/other": 2.0}}, - {"meta_eval_task": "meta-eval-comparia-en"}, - ): - with pytest.raises(ValueError, match="Incompatible configuration"): - _run(_tuning_config(tmp_path, **change), execute) - assert len(calls) == before - fresh = _tuning_config(tmp_path, run_dir=tmp_path / "fresh") - _run(fresh, execute) - repeated = pd.read_parquet(fresh.tune_judge.run_dir / "trials.parquet") - assert repeated.overrides.tolist() == trials.overrides.tolist() - extended = _tuning_config(tmp_path) - extended.tune_judge.neps["total_evaluations_to_spend"] = 11 - _run(extended, execute) - assert len(pd.read_parquet(folder / "trials.parquet")) == 11 - _run(helper, execute) - from neps.state import NePSState - - assert ( - NePSState.create_or_load( - folder / "neps", load_only=True - ).lock_and_get_global_budgets()[0] - == 11 + front = pd.read_parquet(root / "pareto.parquet") + assert set(front.agreement) == {final.agreement.max()} + monkeypatch.setattr( + runner, "resolve_prices", lambda *a, **k: pytest.fail("repriced") ) + pd.testing.assert_frame_equal(results, run(config, lambda _: pytest.fail("reran"))) + config.tune_judge.search_only = True + assert run(config, lambda _: pytest.fail("helper reran")).empty -def test_budget_end_does_not_wait_for_unstarted_batch(tmp_path, monkeypatch): - pytest.importorskip("neps") - cfg = _tuning_config(tmp_path) - cfg.tune_judge.optimizer = { +def test_budget_end_does_not_wait_for_unstarted_batch(config, monkeypatch): + config.tune_judge.optimizer = { "name": "neps_random_search", "ignore_fidelity": "highest_fidelity", } - cfg.tune_judge.neps.update( - total_evaluations_to_spend=2, - sample_batch_size=3, - ) + config.tune_judge.neps.update(total_evaluations_to_spend=2, sample_batch_size=3) monkeypatch.setattr( - "judgearena.tuning.session.time.sleep", - lambda _: pytest.fail("Waiting for an unstarted trial"), + "judgearena.tuning.session.time.sleep", lambda _: pytest.fail("waited") ) - results = _run(cfg) - assert not results.empty - assert len(pd.read_parquet(cfg.tune_judge.run_dir / "trials.parquet")) == 2 + assert not run(config).empty + assert len(pd.read_parquet(config.tune_judge.run_dir / "trials.parquet")) == 2 -def test_cost_counts_both_cached_orientations(artifacts): - tmp_path = artifacts - record = runner._read_trial(tmp_path, price=runner.TokenPrice(2.0, 4.0, "test")) - assert record["tokens_per_battle"] == 6 - assert record["cost_per_1k_battles"] == pytest.approx(0.016) - objective = runner._objective( - {**record, "status": "completed"}, ["agreement", "cost"] - ) - assert objective["objective_to_minimize"] == pytest.approx([0.42, 0.016]) - assert "cost" not in objective - - -def test_agreement_only_reads_old_cached_artifacts(artifacts): - tmp_path = artifacts - pd.DataFrame({"battle_id": ["b1"]}).to_parquet( - tmp_path / "run" / "annotations.parquet" +@pytest.mark.parametrize( + "case,expected", + [("cached", 0.016), ("skipped", 0.008), ("missing", None), ("agreement", None)], +) +def test_native_cost_accounting(config, tmp_path, case, expected): + path = tmp_path / "config.yaml" + direct = RunConfig( + task=config.tune_judge.meta_eval_task, + judge={"model": "Dummy/judge", "temperature": 0.0}, + meta_eval={"split": "validation"}, ) - - result = runner._read_trial(tmp_path, price=None) - assert result["agreement"] == pytest.approx(0.58) - assert result["tokens_per_battle"] is None - assert result["cost_per_1k_battles"] is None - - -def test_missing_native_usage_aborts_even_when_failed_trials_are_ignored(tmp_path): - pytest.importorskip("neps") - cfg = _tuning_config(tmp_path, "mo_hyperband") - cfg.tune_judge.neps.update(ignore_errors=True, total_evaluations_to_spend=1) - - def execute(path): - _fake_trial(path) - annotations_path = path.parent / "run" / "annotations.parquet" - annotations = pd.read_parquet(annotations_path) - annotations["usage_json"] = None - annotations.to_parquet(annotations_path, index=False) - - from neps.exceptions import WorkerRaiseError - - with pytest.raises(WorkerRaiseError) as exc_info: - _run(cfg, execute) - assert isinstance(exc_info.value.__cause__, ValueError) - assert "fresh store_root" in str(exc_info.value.__cause__) - assert "agreement-only" in str(exc_info.value.__cause__) + dump_config(direct, path) + execute(path) + annotations = tmp_path / "run" / "annotations.parquet" + rows = pd.read_parquet(annotations) + if case == "skipped": + rows.loc[0, ["usage_json", "error"]] = [None, "context_length"] + elif case in {"missing", "agreement"}: + rows = rows.drop(columns=["usage_json", "error"]) + rows.to_parquet(annotations) + if case == "missing": + with pytest.raises(ValueError, match="fresh store_root"): + runner._read_trial(tmp_path, runner.TokenPrice(2, 4, "test")) + else: + price = None if case == "agreement" else runner.TokenPrice(2, 4, "test") + record = runner._read_trial(tmp_path, price) + assert record["agreement"] == 0.6 + assert record["cost_per_1k_battles"] == ( + pytest.approx(expected) if expected else None + ) + objective = runner._objective( + {**record, "status": "completed"}, + ["agreement", "cost"] if price else ["agreement"], + ) + assert "cost" not in objective -def test_skipped_request_is_free_but_unparsed_response_keeps_cost(artifacts): - tmp_path = artifacts - annotations_path = tmp_path / "run" / "annotations.parquet" - annotations = pd.read_parquet(annotations_path) - annotations.loc[0, ["usage_json", "error"]] = [None, "context_length"] - annotations.loc[1, "judge_completion"] = "unparseable" - annotations.to_parquet(annotations_path, index=False) - record = runner._read_trial(tmp_path, runner.TokenPrice(2, 4, "test")) - assert record["tokens_per_battle"] == 3 - assert record["cost_per_1k_battles"] == pytest.approx(0.008) +def test_accounting_errors_are_not_ignored(config): + config.tune_judge.neps["total_evaluations_to_spend"] = 1 + def missing_usage(path): + execute(path) + pd.DataFrame({"battle_id": ["b1"]}).to_parquet( + path.parent / "run" / "annotations.parquet" + ) -def test_prices_required_before_execution(tmp_path): - pytest.importorskip("neps") - cfg = _tuning_config(tmp_path, "mo_hyperband", price_per_million_tokens={}) - with pytest.raises(ValueError, match="No token price"): - _run(cfg, lambda _: pytest.fail("executed without prices")) + neps = pytest.importorskip("neps") + with pytest.raises(neps.exceptions.WorkerRaiseError) as error: + run(config, missing_usage) + assert isinstance(error.value.__cause__, ValueError) -def test_pareto_report_uses_completed_full_fidelity_validation(tmp_path): +def test_pareto_excludes_failed_low_fidelity_and_dominated_trials(tmp_path): pytest.importorskip("neps") trials = pd.DataFrame( [ - ("cheap", 0.5, 1.0, 3, "completed"), - ("accurate", 0.8, 2.0, 3, "completed"), - ("dominated", 0.4, 3.0, 3, "completed"), - ("low-fidelity", 0.9, 0.1, 1, "completed"), - ("failed", 1.0, 0.0, 3, "failed"), + ("cheap", 0.5, 1, 3, "completed"), + ("accurate", 0.8, 2, 3, "completed"), + ("dominated", 0.4, 3, 3, "completed"), + ("low", 0.9, 0.1, 1, "completed"), + ("failed", 1, 0, 3, "failed"), ], columns="config_id agreement cost_per_1k_battles battles_per_model status".split(), ) trials = trials.assign(judge_model=trials.config_id, overrides="{}") - executor = SimpleNamespace(run=lambda *args, **kwargs: {"status": "completed"}) runner._evaluate_picks( - executor, SimpleNamespace(test_battles_per_model=None), tmp_path, trials, 3 - ) - front = pd.read_parquet(tmp_path / "pareto.parquet") - assert set(front.config_id) == {"cheap", "accurate"} - - -def test_generic_route_uses_target_defaults(tmp_path): - from judgearena.benchmarks.registry import resolve_benchmark - - tuning = _tuning_config(tmp_path).tune_judge.model_dump() - direct = RunConfig(task=tuning["meta_eval_task"], judge={"model": "Dummy/judge"}) - outer = RunConfig( - task="tune-judge", judge={"model": "Dummy/judge"}, tune_judge=tuning + SimpleNamespace(run=lambda *a, **k: {"status": "completed"}), + SimpleNamespace(test_battles_per_model=None), + tmp_path, + trials, + 3, ) - assert outer.judge == direct.judge - assert outer.meta_eval == direct.meta_eval - resolved = resolve_benchmark(outer.task) - assert resolved.adapter.name == "tune_judge" and resolved.task.spec.dataset is None - for target, judge in ( - ("unknown-task", {"model": "Dummy/judge"}), - ("mt-bench", {"model": "Dummy/judge"}), - (tuning["meta_eval_task"], {"model": "Dummy/judge", "swap_mode": "random"}), - ): - with pytest.raises(ValueError): - RunConfig( - task="tune-judge", - judge=judge, - tune_judge={**tuning, "meta_eval_task": target}, - ) + assert set(pd.read_parquet(tmp_path / "pareto.parquet").config_id) == { + "cheap", + "accurate", + } diff --git a/tests/test_usage.py b/tests/test_usage.py index c3726995..1e152fda 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -34,9 +34,12 @@ def test_do_inference_prefers_canonical_usage(): response_metadata={"token_usage": {"prompt_tokens": 999}, "cost": 0.125}, ) with track_usage() as tracker: - assert do_inference(FakeModel([message]), ["prompt"], stage="judging") == [ - "answer" - ] + (result,) = do_inference( + FakeModel([message]), ["prompt"], stage="judging", return_results=True + ) + assert ( + result.text == "answer" and result.usage == tracker.snapshot().requests[0] + ) usage = tracker.snapshot().requests[0] assert (usage.stage, usage.model) == ("judging", FakeModel.model_name) assert usage.input_tokens == 10 @@ -59,19 +62,6 @@ def test_structured_response_preserves_partial_fallback_usage(capsys): assert "input/output token usage unavailable" in capsys.readouterr().out -def test_return_results_exposes_native_usage_without_logprobs(): - message = AIMessage( - content="answer", - usage_metadata={"input_tokens": 10, "output_tokens": 2, "total_tokens": 12}, - ) - with track_usage() as tracker: - (result,) = do_inference(FakeModel([message]), ["prompt"], return_results=True) - assert result.text == "answer" - assert result.first_token_top_logprobs is None - assert result.usage == tracker.snapshot().requests[0] - assert result.usage.total_tokens == 12 - - def test_run_benchmark_saves_nested_usage_and_cleans_up(tmp_path, monkeypatch, capsys): monkeypatch.setattr(metadata_module, "_get_dependency_versions", lambda **_: {}) monkeypatch.setattr(metadata_module, "_get_git_hash", lambda **_: None)