From 78b406b363f7d6019ec21e528ef242c64de46bbf Mon Sep 17 00:00:00 2001 From: bora kargi Date: Tue, 18 Aug 2026 15:06:42 +0200 Subject: [PATCH 1/2] feat: report model request usage --- judgearena/artifacts/metadata.py | 6 + .../benchmarks/mt_bench/pairwise_judging.py | 1 + judgearena/benchmarks/runner.py | 7 +- judgearena/evaluate.py | 1 + judgearena/generate.py | 17 +- judgearena/models.py | 116 ++++++++++- judgearena/usage.py | 194 ++++++++++++++++++ tests/test_usage.py | 182 ++++++++++++++++ 8 files changed, 510 insertions(+), 14 deletions(-) create mode 100644 judgearena/usage.py create mode 100644 tests/test_usage.py diff --git a/judgearena/artifacts/metadata.py b/judgearena/artifacts/metadata.py index 6a82e597..2cf802b3 100644 --- a/judgearena/artifacts/metadata.py +++ b/judgearena/artifacts/metadata.py @@ -19,6 +19,8 @@ from pathlib import Path from typing import Any +from judgearena.usage import current_run_usage + METADATA_FILENAME = "run-metadata.v1.json" METADATA_SCHEMA_VERSION = "judgearena-run-metadata/v1" _REQUIREMENT_NAME_RE = re.compile(r"^\s*([A-Za-z0-9][A-Za-z0-9_.-]*)") @@ -295,6 +297,10 @@ def write_run_metadata( if task_definition is not None: metadata["task_definition"] = task_definition + usage = current_run_usage() + if usage is not None: + metadata["usage"] = usage.to_dict() + git_hash = _get_git_hash(start_path=Path(__file__).resolve().parent) if git_hash: metadata["git_hash"] = git_hash diff --git a/judgearena/benchmarks/mt_bench/pairwise_judging.py b/judgearena/benchmarks/mt_bench/pairwise_judging.py index f1fc0072..fce48b8e 100644 --- a/judgearena/benchmarks/mt_bench/pairwise_judging.py +++ b/judgearena/benchmarks/mt_bench/pairwise_judging.py @@ -96,6 +96,7 @@ def infer_pairwise_judgments_by_prompt_groups( chat_model=judge_chat_model, inputs=prompt_inputs, use_tqdm=use_tqdm, + stage="judging", ) for item_index, output, prompt_kwargs in zip( idxs, outputs, batch_kwargs, strict=True diff --git a/judgearena/benchmarks/runner.py b/judgearena/benchmarks/runner.py index dddd9f63..ca155b25 100644 --- a/judgearena/benchmarks/runner.py +++ b/judgearena/benchmarks/runner.py @@ -6,6 +6,7 @@ from judgearena.benchmarks.registry import resolve_benchmark from judgearena.log import get_logger +from judgearena.usage import track_usage if TYPE_CHECKING: from judgearena.config import RunConfig @@ -17,4 +18,8 @@ def run_benchmark(cfg: RunConfig) -> object: """Run a task through its registered benchmark adapter.""" resolved = resolve_benchmark(cfg.task) logger.info("Using %s benchmark adapter for %s.", resolved.adapter.name, cfg.task) - return resolved.adapter.runner(cfg, resolved.task) + with track_usage() as usage_tracker: + try: + return resolved.adapter.runner(cfg, resolved.task) + finally: + usage_tracker.render_summary() diff --git a/judgearena/evaluate.py b/judgearena/evaluate.py index db73641c..f232a570 100644 --- a/judgearena/evaluate.py +++ b/judgearena/evaluate.py @@ -227,6 +227,7 @@ def annotate_battles( inputs=inputs, use_tqdm=use_tqdm, return_top_logprobs=collect_top_logprobs, + stage="judging", ) if not collect_top_logprobs: judge_results = [InferenceResult(text=text) for text in judge_results] diff --git a/judgearena/generate.py b/judgearena/generate.py index c38b2bc7..3e6185a7 100644 --- a/judgearena/generate.py +++ b/judgearena/generate.py @@ -37,7 +37,10 @@ def generate_instructions( completions = do_inference( chat_model=chat_model, inputs=inputs, - use_tqdm=use_tqdm, + # This path historically used one synchronous batch regardless of the + # progress setting; preserve that execution behavior here. + use_tqdm=False, + stage="generation", ) df_outputs = pd.DataFrame( data={ @@ -88,6 +91,7 @@ def _infer_grouped_by_temperature( chat_model=group_model, inputs=group_inputs, use_tqdm=use_tqdm, + stage="generation", ) for i, out in zip(idxs, group_outs, strict=True): outputs[i] = out @@ -152,6 +156,7 @@ def generate_multiturn( chat_model=chat_model, inputs=turn1_inputs, use_tqdm=use_tqdm, + stage="generation", ) turn2_inputs = [] @@ -206,6 +211,7 @@ def generate_multiturn( chat_model=chat_model, inputs=turn2_inputs, use_tqdm=use_tqdm, + stage="generation", ) return pd.DataFrame( @@ -225,18 +231,19 @@ def generate_base( use_tqdm: bool = False, **engine_kwargs, ) -> pd.DataFrame: - model = make_model(model, max_tokens=max_tokens, **engine_kwargs) + chat_model = make_model(model, max_tokens=max_tokens, **engine_kwargs) inputs = [ truncate(instruction, max_len=truncate_input_chars) for instruction in instructions ] - completions = model.batch( + completions = do_inference( + chat_model=chat_model, inputs=inputs, - max_tokens=max_tokens, + use_tqdm=use_tqdm, + stage="generation", ) - completions = [x.content if hasattr(x, "content") else x for x in completions] df_outputs = pd.DataFrame( data={ diff --git a/judgearena/models.py b/judgearena/models.py index 648cbb27..b3d3baee 100644 --- a/judgearena/models.py +++ b/judgearena/models.py @@ -7,7 +7,7 @@ import os import time import warnings -from dataclasses import dataclass +from dataclasses import dataclass, replace from langchain_community.llms import LlamaCpp from langchain_openai import ChatOpenAI @@ -16,6 +16,7 @@ from judgearena.constants import VLLM_REASONING_END_STR, VLLM_REASONING_START_STR from judgearena.log import get_logger +from judgearena.usage import RequestUsage, RunUsage, record_usage from judgearena.utils.io import safe_parse_int logger = get_logger(__name__) @@ -447,11 +448,11 @@ async def ainvoke(self, input_item, **invoke_kwargs): @dataclass(frozen=True) class InferenceResult: - """A text completion with the first generated token's top logprobs, when - the backend was asked for (and returned) them.""" + """A text completion and optional provider response details.""" text: str first_token_top_logprobs: dict[str, float] | None = None + usage: RequestUsage | None = None def _first_token_top_logprobs(response) -> dict[str, float] | None: @@ -466,24 +467,97 @@ def _first_token_top_logprobs(response) -> dict[str, float] | None: } -def _to_inference_result(response) -> InferenceResult: +def _optional_number(value, number_type): + if value is None: + return None + try: + return number_type(value) + except (TypeError, ValueError): + return None + + +def _request_usage( + response, + *, + stage: str, + default_model: str | None, +) -> RequestUsage: + response_metadata = getattr(response, "response_metadata", None) or {} + raw_usage = response_metadata.get("token_usage") or {} + prompt_details = raw_usage.get("prompt_tokens_details") or {} + completion_details = raw_usage.get("completion_tokens_details") or {} + + input_tokens = _optional_number(raw_usage.get("prompt_tokens"), int) + output_tokens = _optional_number(raw_usage.get("completion_tokens"), int) + total_tokens = _optional_number(raw_usage.get("total_tokens"), int) + if total_tokens is None and input_tokens is not None and output_tokens is not None: + total_tokens = input_tokens + output_tokens + + model = response_metadata.get("model_name") or default_model + return RequestUsage( + stage=stage, + model=str(model) if model is not None else None, + input_tokens=input_tokens, + output_tokens=output_tokens, + total_tokens=total_tokens, + reasoning_tokens=_optional_number( + completion_details.get("reasoning_tokens"), int + ), + cached_tokens=_optional_number(prompt_details.get("cached_tokens"), int), + cost_usd=_optional_number(raw_usage.get("cost"), float), + ) + + +def _model_name(chat_model) -> str | None: + for attribute in ("model_name", "model", "model_path", "name"): + value = getattr(chat_model, attribute, None) + if isinstance(value, str) and value: + return value + return None + + +def _to_inference_result( + response, + *, + stage: str, + default_model: str | None, +) -> InferenceResult: if isinstance(response, InferenceResult): - return response + if response.usage is not None: + return response + return replace( + response, + usage=RequestUsage(stage=stage, model=default_model), + ) if hasattr(response, "content"): return InferenceResult( text=response.content, first_token_top_logprobs=_first_token_top_logprobs(response), + usage=_request_usage( + response, + stage=stage, + default_model=default_model, + ), ) - return InferenceResult(text=response) + return InferenceResult( + text=response, + usage=RequestUsage(stage=stage, model=default_model), + ) def do_inference( - chat_model, inputs, use_tqdm: bool = False, return_top_logprobs: bool = False + chat_model, + inputs, + use_tqdm: bool = False, + return_top_logprobs: bool = False, + *, + stage: str = "unspecified", ): """Run inference over *inputs*, returning a list of text completions. With ``return_top_logprobs=True``, returns ``InferenceResult`` objects carrying the first token's top logprobs where the backend provided them. + ``stage`` labels request usage for the run-level summary. Retries on rate-limit/server errors with exponential backoff. The async path (``use_tqdm=True``) retries individual calls; the batch path splits @@ -573,7 +647,33 @@ def batch_with_retry(batch_inputs, max_retries=5, base_delay=1.0): # Langchain chat models return AIMessage objects, barebones models plain # strings, and ChatVLLM with logprobs enabled InferenceResult objects. - res = [_to_inference_result(x) for x in res] + res = [ + _to_inference_result( + response, + stage=stage, + default_model=_model_name(chat_model), + ) + for response in res + ] + request_usage = [result.usage for result in res if result.usage is not None] + record_usage(request_usage) + + batch_summary = RunUsage(tuple(request_usage)).summary() + if batch_summary["usage_reported_requests"]: + cost = batch_summary["cost_usd"] + input_tokens = batch_summary["input_tokens"] + output_tokens = batch_summary["output_tokens"] + cost_text = ( + f", ${float(cost):.6f}" if cost is not None else ", cost unavailable" + ) + logger.info( + "Model usage (%s): %d request(s), %s input / %s output tokens%s.", + stage, + batch_summary["requests"], + f"{int(input_tokens):,}" if input_tokens is not None else "unknown", + f"{int(output_tokens):,}" if output_tokens is not None else "unknown", + cost_text, + ) if return_top_logprobs: return res return [r.text for r in res] diff --git a/judgearena/usage.py b/judgearena/usage.py new file mode 100644 index 00000000..f8d171de --- /dev/null +++ b/judgearena/usage.py @@ -0,0 +1,194 @@ +"""Model request usage collected during one benchmark run.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +from contextvars import ContextVar +from dataclasses import dataclass + + +@dataclass(frozen=True) +class RequestUsage: + """Usage record for one model request.""" + + stage: str + model: str | None = None + input_tokens: int | None = None + output_tokens: int | None = None + total_tokens: int | None = None + reasoning_tokens: int | None = None + cached_tokens: int | None = None + cost_usd: float | None = None + + @property + def has_token_usage(self) -> bool: + return any( + value is not None + for value in (self.input_tokens, self.output_tokens, self.total_tokens) + ) + + +def _coverage_status(reported: int, total: int) -> str: + if total == 0: + return "no_requests" + if reported == 0: + return "unavailable" + if reported < total: + return "partial" + return "complete" + + +def _sum_optional(requests: tuple[RequestUsage, ...], field: str) -> int | None: + values = [getattr(request, field) for request in requests] + reported = [int(value) for value in values if value is not None] + return sum(reported) if reported else None + + +def _summarize(requests: tuple[RequestUsage, ...]) -> dict[str, object]: + token_reported = sum(request.has_token_usage for request in requests) + cost_values = [ + request.cost_usd for request in requests if request.cost_usd is not None + ] + cost_reported = len(cost_values) + return { + "requests": len(requests), + "input_tokens": _sum_optional(requests, "input_tokens"), + "output_tokens": _sum_optional(requests, "output_tokens"), + "total_tokens": _sum_optional(requests, "total_tokens"), + "reasoning_tokens": _sum_optional(requests, "reasoning_tokens"), + "cached_tokens": _sum_optional(requests, "cached_tokens"), + "cost_usd": sum(cost_values) if cost_values else None, + "usage_reported_requests": token_reported, + "cost_reported_requests": cost_reported, + "usage_status": _coverage_status(token_reported, len(requests)), + "cost_status": _coverage_status(cost_reported, len(requests)), + } + + +@dataclass(frozen=True) +class RunUsage: + """Immutable snapshot of the model requests made during a run.""" + + requests: tuple[RequestUsage, ...] + + def summary(self) -> dict[str, object]: + return _summarize(self.requests) + + def to_dict(self) -> dict[str, object]: + by_stage = { + stage: _summarize( + tuple(request for request in self.requests if request.stage == stage) + ) + for stage in sorted({request.stage for request in self.requests}) + } + by_model = { + model: _summarize( + tuple( + request + for request in self.requests + if (request.model or "unknown") == model + ) + ) + for model in sorted( + {request.model or "unknown" for request in self.requests} + ) + } + return { + "scope": "model_requests_made_during_this_run", + "source": "provider_response", + "total": self.summary(), + "by_stage": by_stage, + "by_model": by_model, + } + + def render(self) -> None: + data = self.to_dict() + total = data["total"] + assert isinstance(total, dict) + + print("\nModel usage:") + if not self.requests: + print(" No model requests were made during this run.") + return + + by_stage = data["by_stage"] + assert isinstance(by_stage, dict) + for stage, summary in by_stage.items(): + assert isinstance(summary, dict) + print(f" {stage.capitalize()}: {_format_summary(summary)}") + print(f" Total: {_format_summary(total)}") + + +def _format_summary(summary: dict[str, object]) -> str: + requests = int(summary["requests"]) + input_tokens = summary["input_tokens"] + output_tokens = summary["output_tokens"] + cost = summary["cost_usd"] + cost_status = str(summary["cost_status"]) + + if input_tokens is None or output_tokens is None: + tokens = "token usage unavailable" + else: + tokens = f"{int(input_tokens):,} input / {int(output_tokens):,} output tokens" + + if cost_status == "unavailable": + cost_text = "cost unavailable" + elif cost_status == "partial": + cost_text = f"${float(cost):.6f} reported (partial)" + elif cost_status == "no_requests": + cost_text = "$0.000000" + else: + cost_text = f"${float(cost):.6f}" + return f"{requests} request(s), {tokens}, {cost_text}" + + +class UsageTracker: + """Collect request usage without changing benchmark runner signatures.""" + + def __init__(self) -> None: + self._requests: list[RequestUsage] = [] + + def record(self, usage: RequestUsage) -> None: + self._requests.append(usage) + + def record_many(self, usage: list[RequestUsage]) -> None: + self._requests.extend(usage) + + def snapshot(self) -> RunUsage: + return RunUsage(requests=tuple(self._requests)) + + def render_summary(self) -> None: + self.snapshot().render() + + +_CURRENT_TRACKER: ContextVar[UsageTracker | None] = ContextVar( + "judgearena_usage_tracker", default=None +) + + +@contextmanager +def track_usage() -> Iterator[UsageTracker]: + """Collect model usage in the current benchmark execution context.""" + current = _CURRENT_TRACKER.get() + if current is not None: + yield current + return + + tracker = UsageTracker() + token = _CURRENT_TRACKER.set(tracker) + try: + yield tracker + finally: + _CURRENT_TRACKER.reset(token) + + +def record_usage(requests: list[RequestUsage]) -> None: + tracker = _CURRENT_TRACKER.get() + if tracker is not None: + tracker.record_many(requests) + + +def current_run_usage() -> RunUsage | None: + tracker = _CURRENT_TRACKER.get() + return tracker.snapshot() if tracker is not None else None diff --git a/tests/test_usage.py b/tests/test_usage.py new file mode 100644 index 00000000..63c5f8a8 --- /dev/null +++ b/tests/test_usage.py @@ -0,0 +1,182 @@ +import json +from types import SimpleNamespace + +import pytest +from langchain_core.messages import AIMessage + +import judgearena.artifacts.metadata as metadata_module +import judgearena.benchmarks.runner as benchmark_runner +from judgearena.models import InferenceResult, do_inference +from judgearena.usage import ( + RequestUsage, + RunUsage, + current_run_usage, + record_usage, + track_usage, +) + + +class FakeModel: + model_name = "google/test-model" + + def __init__(self, responses): + self.responses = responses + + def batch(self, inputs, **kwargs): + return self.responses[: len(inputs)] + + +def _provider_message() -> AIMessage: + return AIMessage( + content="answer", + usage_metadata={ + "input_tokens": 120, + "output_tokens": 30, + "total_tokens": 150, + "input_token_details": {"cache_read": 20}, + "output_token_details": {"reasoning": 10}, + }, + response_metadata={ + "id": "gen-test", + "model_name": "google/test-model", + "token_usage": { + "prompt_tokens": 120, + "completion_tokens": 30, + "total_tokens": 150, + "prompt_tokens_details": {"cached_tokens": 20}, + "completion_tokens_details": {"reasoning_tokens": 10}, + "cost": 0.00125, + }, + }, + ) + + +def test_do_inference_collects_provider_usage_without_changing_text_results(): + with track_usage() as tracker: + outputs = do_inference( + FakeModel([_provider_message()]), + ["prompt"], + stage="judging", + ) + snapshot = tracker.snapshot() + + assert outputs == ["answer"] + assert len(snapshot.requests) == 1 + usage = snapshot.requests[0] + assert usage.stage == "judging" + assert usage.model == "google/test-model" + assert usage.input_tokens == 120 + assert usage.output_tokens == 30 + assert usage.reasoning_tokens == 10 + assert usage.cached_tokens == 20 + assert usage.cost_usd == pytest.approx(0.00125) + + +def test_do_inference_keeps_usage_on_structured_results(): + with track_usage(): + outputs = do_inference( + FakeModel([_provider_message()]), + ["prompt"], + return_top_logprobs=True, + stage="judging", + ) + + assert isinstance(outputs[0], InferenceResult) + assert outputs[0].text == "answer" + assert outputs[0].usage is not None + assert outputs[0].usage.total_tokens == 150 + + +def test_run_usage_reports_partial_cost_without_presenting_it_as_complete(): + usage = RunUsage( + requests=( + RequestUsage( + stage="generation", + model="candidate", + input_tokens=10, + output_tokens=5, + total_tokens=15, + cost_usd=0.1, + ), + RequestUsage( + stage="judging", + model="local-judge", + input_tokens=20, + output_tokens=2, + total_tokens=22, + ), + ) + ).to_dict() + + assert usage["total"] == { + "requests": 2, + "input_tokens": 30, + "output_tokens": 7, + "total_tokens": 37, + "reasoning_tokens": None, + "cached_tokens": None, + "cost_usd": pytest.approx(0.1), + "usage_reported_requests": 2, + "cost_reported_requests": 1, + "usage_status": "complete", + "cost_status": "partial", + } + assert set(usage["by_stage"]) == {"generation", "judging"} + assert set(usage["by_model"]) == {"candidate", "local-judge"} + + +def test_write_run_metadata_includes_active_run_usage(tmp_path, monkeypatch): + monkeypatch.setattr(metadata_module, "_get_dependency_versions", lambda **_: {}) + monkeypatch.setattr(metadata_module, "_get_git_hash", lambda **_: None) + + with track_usage(): + record_usage( + [ + RequestUsage( + stage="judging", + model="judge", + input_tokens=10, + output_tokens=1, + total_tokens=11, + cost_usd=0.002, + ) + ] + ) + path = metadata_module.write_run_metadata( + output_dir=tmp_path, + entrypoint="test", + run={"task": "unknown"}, + ) + + metadata = json.loads(path.read_text()) + assert metadata["usage"]["total"]["cost_usd"] == pytest.approx(0.002) + assert metadata["usage"]["by_stage"]["judging"]["requests"] == 1 + + +def test_run_benchmark_scopes_and_prints_usage(monkeypatch, capsys): + def fake_runner(cfg, task): + record_usage( + [ + RequestUsage( + stage="judging", + model="judge", + input_tokens=4, + output_tokens=1, + total_tokens=5, + cost_usd=0.0001, + ) + ] + ) + return "done" + + resolved = SimpleNamespace( + adapter=SimpleNamespace(name="test", runner=fake_runner), + task=None, + ) + monkeypatch.setattr(benchmark_runner, "resolve_benchmark", lambda task: resolved) + + result = benchmark_runner.run_benchmark(SimpleNamespace(task="test")) + + assert result == "done" + assert "Model usage:" in capsys.readouterr().out + assert current_run_usage() is None From cb98f4321dd3fd705f0d9d4ef192668de48658a6 Mon Sep 17 00:00:00 2001 From: bora kargi Date: Tue, 18 Aug 2026 15:28:50 +0200 Subject: [PATCH 2/2] refactor: simplify usage metadata --- judgearena/models.py | 2 +- judgearena/usage.py | 29 ++++++++++------------------- tests/test_usage.py | 6 ++---- 3 files changed, 13 insertions(+), 24 deletions(-) diff --git a/judgearena/models.py b/judgearena/models.py index b3d3baee..9a6e00ef 100644 --- a/judgearena/models.py +++ b/judgearena/models.py @@ -659,7 +659,7 @@ def batch_with_retry(batch_inputs, max_retries=5, base_delay=1.0): record_usage(request_usage) batch_summary = RunUsage(tuple(request_usage)).summary() - if batch_summary["usage_reported_requests"]: + if batch_summary["requests_with_token_usage"]: cost = batch_summary["cost_usd"] input_tokens = batch_summary["input_tokens"] output_tokens = batch_summary["output_tokens"] diff --git a/judgearena/usage.py b/judgearena/usage.py index f8d171de..2f6b9e12 100644 --- a/judgearena/usage.py +++ b/judgearena/usage.py @@ -29,16 +29,6 @@ def has_token_usage(self) -> bool: ) -def _coverage_status(reported: int, total: int) -> str: - if total == 0: - return "no_requests" - if reported == 0: - return "unavailable" - if reported < total: - return "partial" - return "complete" - - def _sum_optional(requests: tuple[RequestUsage, ...], field: str) -> int | None: values = [getattr(request, field) for request in requests] reported = [int(value) for value in values if value is not None] @@ -59,10 +49,8 @@ def _summarize(requests: tuple[RequestUsage, ...]) -> dict[str, object]: "reasoning_tokens": _sum_optional(requests, "reasoning_tokens"), "cached_tokens": _sum_optional(requests, "cached_tokens"), "cost_usd": sum(cost_values) if cost_values else None, - "usage_reported_requests": token_reported, - "cost_reported_requests": cost_reported, - "usage_status": _coverage_status(token_reported, len(requests)), - "cost_status": _coverage_status(cost_reported, len(requests)), + "requests_with_token_usage": token_reported, + "requests_with_cost": cost_reported, } @@ -125,19 +113,22 @@ def _format_summary(summary: dict[str, object]) -> str: input_tokens = summary["input_tokens"] output_tokens = summary["output_tokens"] cost = summary["cost_usd"] - cost_status = str(summary["cost_status"]) + requests_with_tokens = int(summary["requests_with_token_usage"]) + requests_with_cost = int(summary["requests_with_cost"]) if input_tokens is None or output_tokens is None: tokens = "token usage unavailable" else: tokens = f"{int(input_tokens):,} input / {int(output_tokens):,} output tokens" + if requests_with_tokens < requests: + tokens += " reported (partial)" - if cost_status == "unavailable": + if requests == 0: + cost_text = "$0.000000" + elif requests_with_cost == 0: cost_text = "cost unavailable" - elif cost_status == "partial": + elif requests_with_cost < requests: cost_text = f"${float(cost):.6f} reported (partial)" - elif cost_status == "no_requests": - cost_text = "$0.000000" else: cost_text = f"${float(cost):.6f}" return f"{requests} request(s), {tokens}, {cost_text}" diff --git a/tests/test_usage.py b/tests/test_usage.py index 63c5f8a8..48f025b2 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -116,10 +116,8 @@ def test_run_usage_reports_partial_cost_without_presenting_it_as_complete(): "reasoning_tokens": None, "cached_tokens": None, "cost_usd": pytest.approx(0.1), - "usage_reported_requests": 2, - "cost_reported_requests": 1, - "usage_status": "complete", - "cost_status": "partial", + "requests_with_token_usage": 2, + "requests_with_cost": 1, } assert set(usage["by_stage"]) == {"generation", "judging"} assert set(usage["by_model"]) == {"candidate", "local-judge"}