diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b9cfa7cc..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 @@ -29,3 +28,16 @@ jobs: run: uv run pre-commit run --all-files - name: Test with 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/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/runner.py b/judgearena/tuning/runner.py new file mode 100644 index 00000000..34ff0f87 --- /dev/null +++ b/judgearena/tuning/runner.py @@ -0,0 +1,346 @@ +"""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 asdict, dataclass +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np +import pandas as pd + +from judgearena.config import RunConfig, dump_config +from judgearena.log import get_logger +from judgearena.tuning.search_space import ( + apply_overrides, + axis_overrides, + build_neps_space, + decode_config, + parameter_specs, +) + +if TYPE_CHECKING: + from judgearena.tasks.schema import ResolvedTaskSpec + +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__) + +TrialExecutor = Callable[[Path], None] +"""Run the trial config at a path, raising ``CalledProcessError`` on failure.""" + +_SUMMARY_COLUMNS = ["config_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, + ) + + +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") + n_battles = annotations["battle_id"].nunique() + 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"], + "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": cost, + "result_dir": str(result_path.parent), + } + + +@dataclass(frozen=True) +class _TrialRunner: + """Materialize trial configs from one base config and record their results.""" + + base: dict + prices: dict[str, TokenPrice] + execute_trial: TrialExecutor + ignore_failed_trials: bool = False + + def config( + self, + overrides: dict[str, object], + *, + split: str, + battles_per_model: int, + folder: Path, + ) -> RunConfig: + values = apply_overrides( + self.base, + { + **overrides, + "meta_eval.split": split, + "meta_eval.battles_per_model": battles_per_model, + "run.result_folder": str(folder), + }, + ) + return RunConfig(**values) + + def run( + self, + overrides: dict[str, object], + *, + split: str, + battles_per_model: int, + folder: Path, + ) -> dict: + trial_cfg = self.config( + overrides, split=split, battles_per_model=battles_per_model, folder=folder + ) + encoded = json.dumps(overrides, sort_keys=True) + record = { + "config_id": hashlib.sha256(encoded.encode()).hexdigest()[:12], + "judge_model": trial_cfg.judge.model, + "overrides": encoded, + "battles_per_model": battles_per_model, + } + 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.get(trial_cfg.judge.model)), + } + 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) + return {**record, "status": "completed", **_read_trial(folder, price)} + + +def _objective(record: dict, objectives: list[str]) -> dict: + if record["status"] != "completed": + return { + "objective_to_minimize": ( + float("inf") + if len(objectives) == 1 + else [float("inf")] * len(objectives) + ), + "info_dict": record, + } + values = [ + 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, + "info_dict": record, + } + + +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 = 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): + 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, + # Only subprocess failures may be ignored; accounting errors stop the run. + ignore_errors=False, + **options, + ) + + +def run_tune_judge( + cfg: RunConfig, + task: ResolvedTaskSpec | 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 + 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={}, + execute_trial=execute_trial, + ) + 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) + tune_dir = prepare_session(cfg) + 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: + 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"] == max_battles) + ].drop_duplicates("config_id") + if final.empty: + raise RuntimeError( + "No configuration completed the highest fidelity; increase the NePS global budget" + ) + ranked = final.sort_values( + ["agreement", "cost_per_1k_battles", "config_id"], ascending=[False, True, True] + ) + from neps.optimizers.utils.multiobjective.epsnet import pareto_efficient + + priced = ranked.dropna(subset=["cost_per_1k_battles"]) + costs = priced[["agreement", "cost_per_1k_battles"]].to_numpy(dtype=float).copy() + costs[:, 0] *= -1 + 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( + [ + runner.run( + json.loads(pick.overrides), + split="test", + 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("\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( + index=False + ) + ) + print(f"Results: {tune_dir}") + return test_results 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_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 new file mode 100644 index 00000000..1bd6e4d1 --- /dev/null +++ b/tests/test_tuning.py @@ -0,0 +1,173 @@ +"""Representative tuning execution, accounting, and selection contracts.""" + +import json +import subprocess +from types import SimpleNamespace + +import pandas as pd +import pytest + +from judgearena.config import RunConfig, dump_config, load_config +from judgearena.tuning import runner + + +@pytest.fixture +def config(tmp_path): + return RunConfig( + task="tune-judge", + 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", + "search_space": {"judge": {"temperature": [0.0, 1.0]}}, + "fidelity": {"battles_per_model": {"lower": 1, "upper": 3}}, + "neps": {"total_evaluations_to_spend": 10, "ignore_errors": True}, + "price_per_million_tokens": {"Dummy/judge": 1.0}, + "run_dir": tmp_path / "tuning", + }, + ) + + +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=0.6 - cfg.judge.temperature / 10, + cohen_kappa=0.2, + coverage=1.0, + 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] * 2, "error": [None] * 2} + ).to_parquet(folder / "annotations.parquet") + + +def run(config, executor=execute): + pytest.importorskip("neps") + return runner.run_tune_judge(config, None, execute_trial=executor) + + +def test_search_failure_resume_and_held_out_evaluation(config, monkeypatch): + calls = [] + + def trial(path): + calls.append(path) + if len(calls) == 1: + raise subprocess.CalledProcessError(1, ["judge-trial"]) + execute(path) + + 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() == 1 + final = trials[(trials.status == "completed") & (trials.battles_per_model == 3)] + assert results.agreement.tolist() == [final.agreement.max()] + assert set(results.status) == {"completed"} + assert load_config(calls[-1]).meta_eval.battles_per_model == 3 + 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(config, monkeypatch): + config.tune_judge.optimizer = { + "name": "neps_random_search", + "ignore_fidelity": "highest_fidelity", + } + config.tune_judge.neps.update(total_evaluations_to_spend=2, sample_batch_size=3) + monkeypatch.setattr( + "judgearena.tuning.session.time.sleep", lambda _: pytest.fail("waited") + ) + assert not run(config).empty + assert len(pd.read_parquet(config.tune_judge.run_dir / "trials.parquet")) == 2 + + +@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"}, + ) + 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_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" + ) + + 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_excludes_failed_low_fidelity_and_dominated_trials(tmp_path): + pytest.importorskip("neps") + trials = pd.DataFrame( + [ + ("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="{}") + runner._evaluate_picks( + SimpleNamespace(run=lambda *a, **k: {"status": "completed"}), + SimpleNamespace(test_battles_per_model=None), + tmp_path, + trials, + 3, + ) + 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)