Judge tuning (5/5): Add reference prices and usage docs

README.md

@@ -339,6 +339,84 @@ 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` task searches judge settings against human preferences, following [Salinas et al., 2025](https://arxiv.org/abs/2501.17178). Each trial runs the matching meta-eval task, so trials share its data and judgement cache. Install NePS with the `tune` extra, then run:
+
+```bash
+uv sync --extra tune
+uv run judgearena --config_path configs/tune_judge.yaml
+uv run judgearena tasks show tune-judge   # task defaults
+```
+
+By default the task tunes on `meta-eval-lmarena-140k-en`, optimizes agreement and cost with `mo_hyperband` (`eta: 3`), and uses 10 to 90 battles per model. Your config overrides these field by field. If you pick another optimizer, set its kwargs (such as `eta`) yourself. [configs/tune_judge.yaml](configs/tune_judge.yaml) is a minimal example:
+
+```yaml
+task: tune-judge
+judge:                      # reference judge, also the search prior
+  model: VLLM/Qwen/Qwen2.5-7B-Instruct
+  prompt_preset: meta-eval-pair-score
+  temperature: 0.0
+  max_out_tokens: 8192
+generation: {truncate_judge_input_chars: 8192}
+tune_judge:
+  search_space:
+    judge:
+      model: [VLLM/Qwen/Qwen2.5-7B-Instruct, VLLM/google/gemma-2-9b-it]
+      prompt_preset: [meta-eval-pair-score, arena-hard]
+      swap_mode: [fixed, both]
+      temperature: {lower: 0, upper: 1, prior_confidence: low}
+      max_out_tokens: {lower: 1024, upper: 8192, log: true}
+  price_per_million_tokens:  # your estimates in USD per 1M tokens
+    VLLM/Qwen/Qwen2.5-7B-Instruct: {input: 0.12, output: 0.12}
+    VLLM/google/gemma-2-9b-it: {input: 0.14, output: 0.14}
+  neps: {total_evaluations_to_spend: 40}
+```
+
+**Search space.** A list (or `{choices: [...]}`) is categorical. `{lower, upper}` is a range whose type follows the run-config field, so `temperature` is continuous and `max_out_tokens` is an integer; `log: true` samples on a log scale. Parameters are nested like the run config and can include `generation.truncate_judge_input_chars`.
+
+**Priors.** Your reference judge is the prior wherever its value lies in the domain. To override, set `prior` to a value (also for categoricals) and `prior_confidence` to `low`, `medium` (default) or `high`.
+
+**Cache reuse.** A continuous value such as temperature usually creates a new model descriptor per sample, so nothing is reused. Use a categorical list like `[0.0, 0.1, 1.0]` when reuse matters.
+
+**Optimizer and objectives.** These are the defaults; set them only to change something:
+
+```yaml
+tune_judge:
+  meta_eval_task: meta-eval-comparia-fr
+  fidelity: {battles_per_model: {lower: 10, upper: 90}}
+  objectives: [agreement, cost]
+  optimizer: {name: mo_hyperband, eta: 3, mo_selector: epsnet}
+```
+
+- `mo_hyperband` and `primo` take `[agreement, cost]`; `neps_priorband` and `neps_hyperband` take `[agreement]`.
+- Optimizer kwargs are native NePS settings. `primo` also takes per-objective `prior_centers` and `prior_confidences`.
+- Other `neps.run` settings go under `neps`, such as budgets and `ignore_errors` (default `true`). Directory, completion and NePS cost-budget settings are not supported.
+- Leave `sample_prior_first` unset for `neps_priorband` and `neps_hyperband`; NePS 0.17 fails on it.
+
+**Cost.** Expected cost is USD per 1,000 battles, computed from the native input and output token counts of the judge responses, including counts stored with cached judgements. It is separate from NePS's search budget.
+
+- Prices come from `tune_judge.price_per_million_tokens` (one number, or `{input, output}`). Without an override, the OpenRouter catalog supplies an estimate, saved to `{store_root}/openrouter_pricing.json`. Agreement-only runs never fetch it. For compute nodes without internet, prefetch it on the login node: `curl -fsS https://openrouter.ai/api/v1/models | jq '.data' > <store_root>/openrouter_pricing.json`.
+- A searched model with no price stops the run before evaluation. Resolved rates are frozen in the run's `prices.json`, which resume and helper workers read.
+- Every reused judgement needs native token counts. Caches written before usage was stored do not have them, so the run errors; use a fresh `store_root` or tune on agreement only. Counts are never estimated.
+- A judge call skipped for exceeding the context limit adds zero generation tokens, and the failed annotation lowers agreement and coverage.
+- With `ignore_errors: true`, failed trials get infinite objective values. Pricing, usage and configuration errors still abort the run.
+
+**Results.** The primary process waits for active trials, reports the agreement/cost Pareto front at full budget, and tests the best-agreement configuration per model on the held-out split (`test_battles_per_model`, default the highest fidelity). Outputs are `trials.parquet`, `pareto.parquet` and `test_results.parquet`; NePS state stays in `neps/`, apart from the inference caches.
+
+**Resume and parallel workers.** Workers share an explicit tuning directory:
+
+```bash
+# Resume the same session.
+uv run judgearena --config_path configs/tune_judge.yaml --tune_judge.run_dir /path/to/tune-dir
+# Join as a helper that runs validation trials only.
+uv run judgearena --config_path configs/tune_judge.yaml --tune_judge.run_dir /path/to/tune-dir --tune_judge.search_only true
+```
+
+The example keeps the paper's limits fixed across judges (8,192 characters per candidate completion, 8,192 judge output tokens, native context, no instruction truncation), so only searched settings vary. Overlong prompts for vLLM chat judges stay failed annotations and are not cached.
+
+Rollback-mode SQLite caching on a shared mount worked in a two-node test on kislurm (NFSv3). That is no guarantee for other filesystems, and WAL must never be used on NFS.
+
 ## 📈 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`).

configs/tune_judge.yaml

@@ -0,0 +1,27 @@
+# Install the tune extra, then run: uv run judgearena --config_path configs/tune_judge.yaml
+task: tune-judge
+judge:
+  model: VLLM/Qwen/Qwen2.5-7B-Instruct
+  prompt_preset: meta-eval-pair-score
+  temperature: 0.0
+  max_out_tokens: 8192
+  max_model_len: null
+generation:
+  truncate_judge_input_chars: 8192
+tune_judge:
+  search_space:
+    judge:
+      model: [VLLM/Qwen/Qwen2.5-7B-Instruct, VLLM/google/gemma-2-9b-it]
+      prompt_preset: [meta-eval-pair-score, arena-hard, default_with_explanation]
+      swap_mode: [fixed, both]
+      temperature: {lower: 0.0, upper: 1.0, prior_confidence: low}
+      max_out_tokens: {lower: 1024, upper: 8192, log: true}
+  # User-provided estimates in USD per 1M tokens (input and output); these
+  # are explicit overrides, not live OpenRouter catalog prices.
+  price_per_million_tokens:
+    VLLM/Qwen/Qwen2.5-7B-Instruct: {input: 0.12, output: 0.12}
+    VLLM/google/gemma-2-9b-it: {input: 0.14, output: 0.14}
+  neps:
+    total_evaluations_to_spend: 40
+# Task defaults select meta-eval-lmarena-140k-en, agreement/cost objectives,
+# mo_hyperband, and 10-90 battles per model. Each can be overridden here.

judgearena/pricing.py

@@ -1,9 +1,15 @@
 """Shared reference token pricing and request cost calculation."""
 
+import json
 from dataclasses import dataclass
+from pathlib import Path
+from urllib.error import URLError
+from urllib.request import urlopen
 
 from judgearena.usage import RequestUsage
 
+_OPENROUTER_URL = "https://openrouter.ai/api/v1/models"
+
 
 @dataclass(frozen=True)
 class TokenPrice:
@@ -30,21 +36,89 @@ def reference_cost(usage: RequestUsage, price: TokenPrice) -> float:
     ) / 1_000_000
 
 
-def resolve_prices(models, overrides, *, require_cost=True) -> dict[str, TokenPrice]:
-    """Resolve explicit USD-per-million-token rates before searching."""
+def _fetch_openrouter(catalog_cache: Path | None) -> list[dict]:
+    try:
+        with urlopen(_OPENROUTER_URL, timeout=10) as response:
+            records = json.load(response)["data"]
+        if catalog_cache is not None:
+            catalog_cache.parent.mkdir(parents=True, exist_ok=True)
+            catalog_cache.write_text(json.dumps(records), encoding="utf-8")
+        return records
+    except (URLError, TimeoutError, OSError) as error:
+        if catalog_cache is not None and catalog_cache.exists():
+            cached = json.loads(catalog_cache.read_text(encoding="utf-8"))
+            if isinstance(cached, list):
+                return cached
+            raise ValueError(
+                "Cached OpenRouter pricing catalog uses an unsupported format; "
+                "refresh the catalog or set price_per_million_tokens overrides"
+            ) from error
+        raise ValueError(
+            "Could not fetch OpenRouter pricing and no cached catalog is available; "
+            "fetch the catalog on a login node or set "
+            "price_per_million_tokens overrides"
+        ) from error
+
+
+def _find_price(model: str, records: list[dict]) -> TokenPrice | None:
+    requested_id = model.split("/", 1)[-1].lower()
+    suffix = "/".join(requested_id.split("/")[-2:])
+    record = next(
+        (
+            item
+            for item in records
+            if str(item.get("hugging_face_id") or "").lower() == requested_id
+        ),
+        None,
+    )
+    if record is None:
+        record = next(
+            (item for item in records if str(item.get("id") or "").lower() == suffix),
+            None,
+        )
+    if record is None:
+        return None
+    rates = record.get("pricing", {})
+    if "prompt" not in rates or "completion" not in rates:
+        return None
+    return TokenPrice(
+        float(rates["prompt"]) * 1_000_000,
+        float(rates["completion"]) * 1_000_000,
+        "openrouter_reference",
+    )
+
+
+def resolve_prices(
+    models,
+    overrides,
+    *,
+    catalog_cache: Path | None = None,
+    require_cost=True,
+) -> dict[str, TokenPrice]:
+    """Resolve explicit or OpenRouter reference USD-per-million-token rates."""
     if not require_cost:
         return {}
     prices = {}
+    unresolved = []
     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")
+        rate = overrides.get(model)
+        if rate is not None:
+            if isinstance(rate, (int, float)):
+                rate = {"input": rate, "output": rate}
+            prices[model] = TokenPrice(
+                float(rate["input"]), float(rate["output"]), "user_override"
+            )
+        else:
+            unresolved.append(model)
+    records = _fetch_openrouter(catalog_cache) if unresolved else []
+    for model in unresolved:
+        price = _find_price(model, records)
+        if price is not None:
+            prices[model] = price
     missing = sorted(set(models) - prices.keys())
     if missing:
         raise ValueError(
-            f"No token price for {missing}; set tune_judge.price_per_million_tokens"
+            f"No token price for {missing}; set price_per_million_tokens "
+            f"overrides for: {missing}"
         )
     return prices

judgearena/tuning/runner.py

@@ -276,6 +276,8 @@ def run_tune_judge(
             prices = resolve_prices(
                 models,
                 tuning.price_per_million_tokens,
+                catalog_cache=Path(cfg.run.store_root or cfg.run.result_folder)
+                / "openrouter_pricing.json",
                 require_cost=True,
             )
             if not tuning.search_only:

tests/test_tuning_pricing.py

@@ -0,0 +1,66 @@
+import json
+from contextlib import nullcontext
+from io import BytesIO
+
+import pytest
+
+from judgearena import pricing
+from judgearena.pricing import TokenPrice, reference_cost, resolve_prices
+from judgearena.usage import RequestUsage
+
+
+def test_override_and_agreement_only_skip_catalog(monkeypatch):
+    monkeypatch.setattr(pricing, "urlopen", lambda *a, **k: pytest.fail("fetched"))
+    assert resolve_prices(["VLLM/local"], {"VLLM/local": 1.2}) == {
+        "VLLM/local": TokenPrice(1.2, 1.2, "user_override")
+    }
+    assert resolve_prices(["VLLM/local"], {}, require_cost=False) == {}
+
+
+@pytest.mark.parametrize(
+    "provider,hf_match", [("VLLM", True), ("OpenRouter", True), ("VLLM", False)]
+)
+def test_reference_matching_and_catalog_snapshot(
+    provider, hf_match, monkeypatch, tmp_path
+):
+    records = [dict(id="org/model", pricing=dict(prompt="1e-6", completion="2e-6"))]
+    if hf_match:
+        records.append(
+            dict(
+                id="other/model",
+                hugging_face_id="ORG/Model",
+                pricing=dict(prompt="3e-6", completion="4e-6"),
+            )
+        )
+    response = json.dumps({"data": records}).encode()
+    monkeypatch.setattr(
+        pricing, "urlopen", lambda *a, **k: nullcontext(BytesIO(response))
+    )
+    model, cache = f"{provider}/org/model", tmp_path / "prices.json"
+    expected = (
+        TokenPrice(3, 4, "openrouter_reference")
+        if hf_match
+        else TokenPrice(1, 2, "openrouter_reference")
+    )
+    assert resolve_prices([model], {}, catalog_cache=cache)[model] == expected
+    assert json.loads(cache.read_text()) == records
+    with pytest.raises(ValueError, match="overrides"):
+        resolve_prices(["VLLM/unknown"], {}, catalog_cache=cache)
+
+
+@pytest.mark.parametrize(
+    "tokens,expected", [((0, 3), 6), ((2, 1), 4), ((None, 3), None)]
+)
+def test_native_counts_and_no_double_counted_reasoning(tokens, expected):
+    usage = RequestUsage(
+        stage="judging",
+        input_tokens=tokens[0],
+        output_tokens=tokens[1],
+        reasoning_tokens=100,
+    )
+    price = TokenPrice(1_000_000, 2_000_000, "ref")
+    if expected is None:
+        with pytest.raises(ValueError, match="missing native"):
+            reference_cost(usage, price)
+    else:
+        assert reference_cost(usage, price) == expected