From 52a4620f3fa8f5a88694dbafb1afbdd27a02d1c2 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:11:04 +0200 Subject: [PATCH 1/8] Add content-addressed local completion and judgement SQLite caches. --- judgearena/cache_sqlite.py | 341 +++++++++++++++++++++++++++++++++++++ tests/test_cache_sqlite.py | 117 +++++++++++++ 2 files changed, 458 insertions(+) create mode 100644 judgearena/cache_sqlite.py create mode 100644 tests/test_cache_sqlite.py diff --git a/judgearena/cache_sqlite.py b/judgearena/cache_sqlite.py new file mode 100644 index 00000000..0325ff39 --- /dev/null +++ b/judgearena/cache_sqlite.py @@ -0,0 +1,341 @@ +"""Content-addressed local SQLite caches.""" + +from __future__ import annotations + +import hashlib +import json +import os +import sqlite3 +import uuid +from datetime import UTC, datetime +from pathlib import Path +from typing import Any, Literal +from urllib.parse import quote + +import pandas as pd + +COMPLETION_DB_NAME = "completions.db" +JUDGEMENT_DB_NAME = "judgements.db" +DESCRIPTOR_FILENAME = "metadata.json" + +CacheKind = Literal["completions", "judgements"] + + +def stable_json_dumps(value: Any) -> str: + return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + + +def descriptor_hash(descriptor: dict[str, Any]) -> str: + return hashlib.sha256(stable_json_dumps(descriptor).encode()).hexdigest()[:16] + + +def input_hash(input_text: str) -> str: + return hashlib.sha256(input_text.encode()).hexdigest() + + +def cache_folder( + store_root: Path | str, + kind: CacheKind, + task: str, + model_spec: str, + descriptor: dict[str, Any], +) -> Path: + return cache_model_folder(store_root, kind, task, model_spec) / descriptor_hash( + descriptor + ) + + +def cache_model_folder( + store_root: Path | str, + kind: CacheKind, + task: str, + model_spec: str, +) -> Path: + provider, model = model_spec.split("/", 1) + return ( + Path(store_root) + / kind + / quote(task, safe="") + / quote(provider, safe="") + / quote(model, safe="") + ) + + +def write_descriptor(folder: Path, descriptor: dict[str, Any]) -> Path: + folder.mkdir(parents=True, exist_ok=True) + path = folder / DESCRIPTOR_FILENAME + if path.exists(): + if json.loads(path.read_text()) != descriptor: + raise ValueError(f"Descriptor does not match existing metadata at {path}.") + return path + + path.write_text(json.dumps(descriptor, indent=2, sort_keys=True) + "\n") + return path + + +def read_descriptor(folder: Path) -> dict[str, Any]: + descriptor = json.loads((folder / DESCRIPTOR_FILENAME).read_text()) + if folder.name != descriptor_hash(descriptor): + raise ValueError(f"Descriptor hash does not match cache folder {folder}.") + return descriptor + + +class _SQLiteCache: + table: str + schema: str + + def __init__(self, db_path: Path | str) -> None: + self.db_path = Path(db_path) + self._conn: sqlite3.Connection | None = None + + def _connect(self) -> sqlite3.Connection: + if self._conn is None: + self.db_path.parent.mkdir(parents=True, exist_ok=True) + self._conn = sqlite3.connect(self.db_path) + self._conn.execute(self.schema) + self._conn.commit() + return self._conn + + def _query( + self, + input_hashes: list[str] | None, + conditions: list[str], + params: list[Any], + ) -> pd.DataFrame: + if input_hashes is not None: + if not input_hashes: + return pd.read_sql( + f"SELECT * FROM {self.table} WHERE 0", self._connect() + ) + placeholders = ",".join("?" * len(input_hashes)) + conditions.append(f"input_hash IN ({placeholders})") + params.extend(input_hashes) + + where = f" WHERE {' AND '.join(conditions)}" if conditions else "" + return pd.read_sql( + f"SELECT * FROM {self.table}{where} ORDER BY instruction_id", + self._connect(), + params=params, + ) + + def _delete(self, conditions: list[str], params: list[Any]) -> int: + if not conditions: + raise ValueError("Delete requires at least one filter.") + where = f" WHERE {' AND '.join(conditions)}" if conditions else "" + with self._connect() as conn: + cursor = conn.execute(f"DELETE FROM {self.table}{where}", params) + return cursor.rowcount + + def merge_from(self, other_db: Path) -> int: + """Atomically merge another cache database using pushed_at last-write-wins.""" + self.close() + frames = [] + for db_path in (other_db, self.db_path): + if db_path.exists(): + with sqlite3.connect(db_path) as conn: + frames.append(pd.read_sql(f"SELECT * FROM {self.table}", conn)) + rows = ( + pd.concat(frames, ignore_index=True) + .sort_values("pushed_at", kind="stable") + .drop_duplicates("input_hash", keep="last") + ) + + self.db_path.parent.mkdir(parents=True, exist_ok=True) + temporary_db = self.db_path.with_name(f".{self.db_path.name}.{uuid.uuid4()}") + try: + with sqlite3.connect(temporary_db) as conn: + conn.execute(self.schema) + rows.to_sql(self.table, conn, if_exists="append", index=False) + os.replace(temporary_db, self.db_path) + finally: + temporary_db.unlink(missing_ok=True) + return len(rows) + + def close(self) -> None: + if self._conn is not None: + self._conn.close() + self._conn = None + + def __enter__(self) -> _SQLiteCache: + return self + + def __exit__(self, *_: object) -> None: + self.close() + + +class CompletionCache(_SQLiteCache): + """Completion rows keyed by the exact rendered model input.""" + + table = "completions" + schema = """ + CREATE TABLE IF NOT EXISTS completions ( + input_hash TEXT PRIMARY KEY, + input_text TEXT NOT NULL, + completion TEXT NOT NULL, + benchmark TEXT NOT NULL, + instruction_id TEXT NOT NULL, + model TEXT NOT NULL, + pushed_by TEXT NOT NULL, + pushed_at TEXT NOT NULL, + run_id TEXT NOT NULL + ) + """ + + def save( + self, + rows: pd.DataFrame, + *, + pushed_by: str, + run_id: str | None = None, + ) -> int: + now = datetime.now(UTC).isoformat() + resolved_run_id = run_id or str(uuid.uuid4()) + values = [ + ( + input_hash(str(row["input_text"])), + str(row["input_text"]), + str(row["completion"]), + str(row["benchmark"]), + str(row["instruction_id"]), + str(row["model"]), + pushed_by, + now, + resolved_run_id, + ) + for _, row in rows.iterrows() + ] + with self._connect() as conn: + conn.executemany( + "INSERT OR REPLACE INTO completions VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", + values, + ) + return len(values) + + def query( + self, + input_hashes: list[str] | None = None, + *, + instruction_id: str | None = None, + model: str | None = None, + ) -> pd.DataFrame: + conditions: list[str] = [] + params: list[Any] = [] + if instruction_id is not None: + conditions.append("instruction_id = ?") + params.append(str(instruction_id)) + if model is not None: + conditions.append("model = ?") + params.append(model) + return self._query(input_hashes, conditions, params) + + def delete( + self, + *, + instruction_id: str | None = None, + model: str | None = None, + ) -> int: + conditions: list[str] = [] + params: list[Any] = [] + if instruction_id is not None: + conditions.append("instruction_id = ?") + params.append(str(instruction_id)) + if model is not None: + conditions.append("model = ?") + params.append(model) + return self._delete(conditions, params) + + +class JudgementCache(_SQLiteCache): + """Raw judge completions keyed by the exact rendered judge input.""" + + table = "judgements" + schema = """ + CREATE TABLE IF NOT EXISTS judgements ( + input_hash TEXT PRIMARY KEY, + judge_input TEXT NOT NULL, + judge_completion TEXT NOT NULL, + benchmark TEXT NOT NULL, + instruction_id TEXT NOT NULL, + model_a TEXT NOT NULL, + model_b TEXT NOT NULL, + judge TEXT NOT NULL, + top_logprobs TEXT, + -- direct/reversed relative to the source model order, when applicable + orientation TEXT, + pushed_by TEXT NOT NULL, + pushed_at TEXT NOT NULL, + run_id TEXT NOT NULL + ) + """ + + def save( + self, + rows: pd.DataFrame, + *, + pushed_by: str, + run_id: str | None = None, + ) -> int: + now = datetime.now(UTC).isoformat() + resolved_run_id = run_id or str(uuid.uuid4()) + values = [ + ( + input_hash(str(row["judge_input"])), + str(row["judge_input"]), + str(row["judge_completion"]), + str(row["benchmark"]), + str(row["instruction_id"]), + str(row["model_a"]), + str(row["model_b"]), + str(row["judge"]), + ( + stable_json_dumps(row["top_logprobs"]) + if row.get("top_logprobs") is not None + else None + ), + row.get("orientation"), + pushed_by, + now, + resolved_run_id, + ) + for _, row in rows.iterrows() + ] + with self._connect() as conn: + conn.executemany( + "INSERT OR REPLACE INTO judgements " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", + values, + ) + return len(values) + + def query( + self, + input_hashes: list[str] | None = None, + *, + instruction_id: str | None = None, + model: str | None = None, + ) -> pd.DataFrame: + conditions: list[str] = [] + params: list[Any] = [] + if instruction_id is not None: + conditions.append("instruction_id = ?") + params.append(str(instruction_id)) + if model is not None: + conditions.append("(model_a = ? OR model_b = ?)") + params.extend((model, model)) + return self._query(input_hashes, conditions, params) + + def delete( + self, + *, + instruction_id: str | None = None, + model: str | None = None, + ) -> int: + conditions: list[str] = [] + params: list[Any] = [] + if instruction_id is not None: + conditions.append("instruction_id = ?") + params.append(str(instruction_id)) + if model is not None: + conditions.append("(model_a = ? OR model_b = ?)") + params.extend((model, model)) + return self._delete(conditions, params) diff --git a/tests/test_cache_sqlite.py b/tests/test_cache_sqlite.py new file mode 100644 index 00000000..28dc693e --- /dev/null +++ b/tests/test_cache_sqlite.py @@ -0,0 +1,117 @@ +import json + +import pandas as pd +import pytest + +from judgearena.cache_sqlite import ( + COMPLETION_DB_NAME, + JUDGEMENT_DB_NAME, + CompletionCache, + JudgementCache, + cache_folder, + input_hash, + write_descriptor, +) + +DESCRIPTOR = { + "model": "Qwen/Qwen3-8B", + "provider": "VLLM", + "sampling": {"max_tokens": 1024, "temperature": 0.0}, +} + + +def test_descriptor_paths_separate_completion_and_judgement_caches(tmp_path): + model = "VLLM/Qwen/Qwen3-8B" + completion_folder = cache_folder( + tmp_path, "completions", "arena-hard", model, DESCRIPTOR + ) + judgement_folder = cache_folder( + tmp_path, "judgements", "arena-hard", model, DESCRIPTOR + ) + + assert completion_folder != judgement_folder + assert completion_folder.name == judgement_folder.name + assert completion_folder / COMPLETION_DB_NAME != ( + judgement_folder / JUDGEMENT_DB_NAME + ) + + metadata_path = write_descriptor(completion_folder, DESCRIPTOR) + assert json.loads(metadata_path.read_text()) == DESCRIPTOR + write_descriptor(completion_folder, DESCRIPTOR) + + with pytest.raises(ValueError, match="does not match"): + write_descriptor(completion_folder, {**DESCRIPTOR, "provider": "OpenAI"}) + + +def test_completion_cache_uses_content_key_and_last_write(tmp_path): + db_path = tmp_path / COMPLETION_DB_NAME + first = pd.DataFrame( + [ + { + "input_text": "rendered prompt", + "completion": "first", + "benchmark": "arena-hard", + "instruction_id": "12", + "model": "VLLM/Qwen/Qwen3-8B", + } + ] + ) + second = first.assign(completion="second") + + with CompletionCache(db_path) as cache: + cache.save(first, pushed_by="alice") + cache.save(second, pushed_by="bob") + result = cache.query([input_hash("rendered prompt")]) + + assert result["completion"].tolist() == ["second"] + assert result["pushed_by"].tolist() == ["bob"] + + +def test_completion_cache_filters_and_deletes_by_instruction(tmp_path): + rows = pd.DataFrame( + [ + { + "input_text": f"prompt-{index}", + "completion": f"completion-{index}", + "benchmark": "arena-hard", + "instruction_id": str(index), + "model": "VLLM/Qwen/Qwen3-8B", + } + for index in range(2) + ] + ) + + with CompletionCache(tmp_path / COMPLETION_DB_NAME) as cache: + cache.save(rows, pushed_by="alice") + assert cache.query(instruction_id="1")["completion"].tolist() == [ + "completion-1" + ] + assert cache.delete(instruction_id="1") == 1 + assert cache.query()["instruction_id"].tolist() == ["0"] + + +def test_judgement_cache_filters_and_deletes_by_candidate_model(tmp_path): + rows = pd.DataFrame( + [ + { + "judge_input": f"judge prompt {index}", + "judge_completion": f"scores {index}", + "benchmark": "arena-hard", + "instruction_id": str(index), + "model_a": "candidate" if index == 0 else "baseline", + "model_b": "baseline" if index == 0 else "other", + "judge": "VLLM/Qwen/Qwen3-8B", + "top_logprobs": {"m": -0.1, "M": -2.0} if index == 0 else None, + "orientation": "direct", + } + for index in range(2) + ] + ) + + with JudgementCache(tmp_path / JUDGEMENT_DB_NAME) as cache: + cache.save(rows, pushed_by="alice") + result = cache.query(model="candidate") + assert result["judge_completion"].tolist() == ["scores 0"] + assert json.loads(result.iloc[0]["top_logprobs"]) == {"M": -2.0, "m": -0.1} + assert cache.delete(model="candidate") == 1 + assert cache.query()["instruction_id"].tolist() == ["1"] From 969347c50e399af296f9a6dc81fd1691ac31296a Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:11:16 +0200 Subject: [PATCH 2/8] Cache inference at do_inference with lazy PreparedModel. --- judgearena/inference.py | 210 ++++++++++++++++++++++++++++++++++ judgearena/models.py | 109 +++++++++++++++++- tests/test_inference_cache.py | 106 +++++++++++++++++ 3 files changed, 424 insertions(+), 1 deletion(-) create mode 100644 judgearena/inference.py create mode 100644 tests/test_inference_cache.py diff --git a/judgearena/inference.py b/judgearena/inference.py new file mode 100644 index 00000000..2a2e24d1 --- /dev/null +++ b/judgearena/inference.py @@ -0,0 +1,210 @@ +"""Lazy model preparation and inference-cache context.""" + +from __future__ import annotations + +import json +from abc import ABC, abstractmethod +from collections.abc import Callable +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, ClassVar + +import pandas as pd + +from judgearena.cache_sqlite import ( + COMPLETION_DB_NAME, + JUDGEMENT_DB_NAME, + CacheKind, + CompletionCache, + JudgementCache, + cache_folder, + stable_json_dumps, + write_descriptor, +) + +_ROLE_MAP = {"human": "user", "ai": "assistant", "system": "system"} + + +def canonicalize_chat_input(input_item: Any) -> str: + """Serialize a logical model input for content-addressed cache lookup.""" + if isinstance(input_item, str): + payload = {"type": "text", "text": input_item} + elif hasattr(input_item, "to_messages"): + payload = { + "type": "messages", + "messages": [ + { + "role": _ROLE_MAP.get(message.type, message.type), + "content": message.content, + } + for message in input_item.to_messages() + ], + } + else: + raise TypeError(f"Unsupported inference input: {type(input_item)!r}") + return stable_json_dumps(payload) + + +def build_model_descriptor( + provider: str, + model_name: str, + resolved_kwargs: dict[str, Any], +) -> dict[str, Any] | None: + """Describe output-affecting settings without constructing the backend.""" + if provider != "Dummy": + return None + return { + "schema_version": "judgearena-inference-cache/v1", + "provider": provider, + "model": model_name, + "input_mode": "chat", + "model_kwargs": resolved_kwargs, + } + + +@dataclass(frozen=True) +class CachedInferenceResult: + """Provider output fields required by downstream parsing.""" + + text: str + first_token_top_logprobs: dict[str, float] | None = None + + +@dataclass +class PreparedModel: + """Carry cache identity while deferring backend construction until a miss.""" + + model_spec: str + descriptor: dict[str, Any] | None + factory: Callable[[], Any] + cache: InferenceCache | None = None + _model: Any = field(default=None, init=False, repr=False) + + def materialize(self) -> Any: + if self._model is None: + self._model = self.factory() + return self._model + + +@dataclass(frozen=True) +class InferenceCache(ABC): + """Share cache lifecycle while subclasses define role-specific rows.""" + + store_root: Path + task: str + pushed_by: str = "judgearena" + + kind: ClassVar[CacheKind] + db_name: ClassVar[str] + store_type: ClassVar[type[CompletionCache] | type[JudgementCache]] + + def open_store(self, model: PreparedModel) -> CompletionCache | JudgementCache: + assert model.descriptor is not None + folder = cache_folder( + self.store_root, + self.kind, + self.task, + model.model_spec, + model.descriptor, + ) + write_descriptor(folder, model.descriptor) + return self.store_type(folder / self.db_name) + + def save_outputs( + self, + store: CompletionCache | JudgementCache, + model: PreparedModel, + input_texts: list[str], + outputs: list[Any], + metadata: list[dict[str, Any]], + indices: list[int], + ) -> None: + rows = [ + self.make_row( + model=model, + input_text=input_texts[index], + output=output, + metadata=metadata[index], + ) + for index, output in zip(indices, outputs, strict=True) + ] + store.save(pd.DataFrame(rows), pushed_by=self.pushed_by) + + @abstractmethod + def make_row( + self, + *, + model: PreparedModel, + input_text: str, + output: Any, + metadata: dict[str, Any], + ) -> dict[str, Any]: + """Convert one inference output to its role-specific storage row.""" + + @abstractmethod + def cached_result(self, row: pd.Series) -> CachedInferenceResult: + """Restore output fields from a stored row.""" + + +class CompletionInferenceCache(InferenceCache): + """Cache generated model completions.""" + + kind = "completions" + db_name = COMPLETION_DB_NAME + store_type = CompletionCache + + def make_row( + self, + *, + model: PreparedModel, + input_text: str, + output: Any, + metadata: dict[str, Any], + ) -> dict[str, Any]: + return { + "input_text": input_text, + "completion": output.text, + "benchmark": self.task, + "instruction_id": metadata["instruction_id"], + "model": model.model_spec, + } + + def cached_result(self, row: pd.Series) -> CachedInferenceResult: + return CachedInferenceResult(text=str(row["completion"])) + + +class JudgementInferenceCache(InferenceCache): + """Cache raw judge completions.""" + + kind = "judgements" + db_name = JUDGEMENT_DB_NAME + store_type = JudgementCache + + def make_row( + self, + *, + model: PreparedModel, + input_text: str, + output: Any, + metadata: dict[str, Any], + ) -> dict[str, Any]: + return { + "judge_input": input_text, + "judge_completion": output.text, + "benchmark": self.task, + "instruction_id": metadata["instruction_id"], + "model_a": metadata["model_a"], + "model_b": metadata["model_b"], + "judge": model.model_spec, + "top_logprobs": output.first_token_top_logprobs, + "orientation": metadata.get("orientation"), + } + + def cached_result(self, row: pd.Series) -> CachedInferenceResult: + top_logprobs = row["top_logprobs"] + return CachedInferenceResult( + text=str(row["judge_completion"]), + first_token_top_logprobs=( + json.loads(top_logprobs) if pd.notna(top_logprobs) else None + ), + ) diff --git a/judgearena/models.py b/judgearena/models.py index f949d5a9..8e3adc57 100644 --- a/judgearena/models.py +++ b/judgearena/models.py @@ -16,7 +16,14 @@ from tqdm.asyncio import tqdm from tqdm.contrib.logging import logging_redirect_tqdm +from judgearena.cache_sqlite import input_hash from judgearena.constants import VLLM_REASONING_END_STR, VLLM_REASONING_START_STR +from judgearena.inference import ( + InferenceCache, + PreparedModel, + build_model_descriptor, + canonicalize_chat_input, +) from judgearena.log import get_logger from judgearena.usage import RequestUsage, RunUsage, record_usage from judgearena.utils.io import safe_parse_int @@ -663,7 +670,7 @@ def batch_inference_once( return [result.text for result in results] -def do_inference( +def _do_inference_uncached( chat_model, inputs, use_tqdm: bool = False, @@ -772,6 +779,78 @@ def batch_with_retry(batch_inputs, max_retries=5, base_delay=1.0): return [result.text for result in results] +def do_inference( + chat_model, + inputs, + use_tqdm: bool = False, + return_top_logprobs: bool = False, + *, + stage: str = "unspecified", + cache_metadata: list[dict] | None = None, +): + """Reuse raw provider outputs and invoke the backend only for cache misses.""" + inputs = list(inputs) + if not isinstance(chat_model, PreparedModel): + return _do_inference_uncached( + chat_model, + inputs, + use_tqdm, + return_top_logprobs, + stage=stage, + ) + + cache = chat_model.cache + if cache is None or chat_model.descriptor is None: + return _do_inference_uncached( + chat_model.materialize(), + inputs, + use_tqdm, + return_top_logprobs, + stage=stage, + ) + if cache_metadata is None or len(cache_metadata) != len(inputs): + raise ValueError("cache_metadata must contain one row per inference input.") + + input_texts = [canonicalize_chat_input(item) for item in inputs] + input_hashes = [input_hash(input_text) for input_text in input_texts] + with cache.open_store(chat_model) as store: + cached_rows = store.query(input_hashes).set_index("input_hash") + results: list[InferenceResult | None] = [ + ( + InferenceResult(**cache.cached_result(cached_rows.loc[key]).__dict__) + if key in cached_rows.index + else None + ) + for key in input_hashes + ] + missing_indices = [ + index for index, result in enumerate(results) if result is None + ] + if missing_indices: + generated = _do_inference_uncached( + chat_model.materialize(), + [inputs[index] for index in missing_indices], + use_tqdm, + True, + stage=stage, + ) + for index, result in zip(missing_indices, generated, strict=True): + results[index] = result + cache.save_outputs( + store, + chat_model, + input_texts, + generated, + cache_metadata, + missing_indices, + ) + + resolved_results = [result for result in results if result is not None] + if return_top_logprobs: + return resolved_results + return [result.text for result in resolved_results] + + def _route_sampling_params( engine_kwargs: dict, *, @@ -814,6 +893,34 @@ def _route_sampling_params( return engine_kwargs +def prepare_model( + model: str, + max_tokens: int | None = 8192, + *, + cache: InferenceCache | None = None, + **engine_kwargs, +) -> PreparedModel: + """Prepare cache identity without constructing the provider backend.""" + provider, model_name = _split_model_spec(model) + resolved_kwargs = {**engine_kwargs, "max_tokens": max_tokens or 8192} + descriptor = ( + build_model_descriptor(provider, model_name, resolved_kwargs) + if cache is not None + else None + ) + factory_kwargs = engine_kwargs.copy() + return PreparedModel( + model_spec=model, + descriptor=descriptor, + factory=lambda: make_model( + model, + max_tokens=max_tokens, + **factory_kwargs, + ), + cache=cache, + ) + + def make_model(model: str, max_tokens: int | None = 8192, **engine_kwargs): """Instantiate a model wrapper from a provider/model-name string. diff --git a/tests/test_inference_cache.py b/tests/test_inference_cache.py new file mode 100644 index 00000000..3343fb97 --- /dev/null +++ b/tests/test_inference_cache.py @@ -0,0 +1,106 @@ +from langchain_core.messages import AIMessage + +import judgearena.models as models +from judgearena.inference import CompletionInferenceCache, JudgementInferenceCache +from judgearena.models import InferenceResult, do_inference, prepare_model +from judgearena.usage import track_usage + + +class EchoModel: + def __init__(self): + self.calls = [] + + def batch(self, inputs, **_kwargs): + self.calls.append(inputs) + return [AIMessage(content=f"generated:{item}") for item in inputs] + + +def test_full_hit_does_not_materialize_model(tmp_path, monkeypatch): + cache = CompletionInferenceCache(tmp_path, "arena-hard") + monkeypatch.setattr(models, "make_model", lambda *_args, **_kwargs: EchoModel()) + metadata = [{"instruction_id": "1"}] + do_inference( + prepare_model("Dummy/test-model", cache=cache), + ["prompt"], + cache_metadata=metadata, + ) + + def fail_if_materialized(*_args, **_kwargs): + raise AssertionError("cache hit materialized the model") + + monkeypatch.setattr(models, "make_model", fail_if_materialized) + with track_usage() as tracker: + outputs = do_inference( + prepare_model("Dummy/test-model", cache=cache), + ["prompt"], + cache_metadata=metadata, + ) + + assert outputs == ["generated:prompt"] + assert tracker.snapshot().requests == () + + +def test_mixed_hits_and_misses_preserve_order(tmp_path, monkeypatch): + cache = CompletionInferenceCache(tmp_path, "arena-hard") + first_backend = EchoModel() + monkeypatch.setattr(models, "make_model", lambda *_args, **_kwargs: first_backend) + do_inference( + prepare_model("Dummy/test-model", cache=cache), + ["hit"], + cache_metadata=[{"instruction_id": "hit"}], + ) + + backend = EchoModel() + monkeypatch.setattr(models, "make_model", lambda *_args, **_kwargs: backend) + outputs = do_inference( + prepare_model("Dummy/test-model", cache=cache), + ["miss-a", "hit", "miss-b"], + cache_metadata=[ + {"instruction_id": "a"}, + {"instruction_id": "hit"}, + {"instruction_id": "b"}, + ], + ) + + assert outputs == ["generated:miss-a", "generated:hit", "generated:miss-b"] + assert backend.calls == [["miss-a", "miss-b"]] + + +def test_judgement_hit_preserves_top_logprobs(tmp_path, monkeypatch): + cache = JudgementInferenceCache(tmp_path, "arena-hard") + + class LogprobModel: + def batch(self, inputs, **_kwargs): + return [ + InferenceResult( + text="m", + first_token_top_logprobs={"m": -0.1, "M": -2.0}, + ) + for _ in inputs + ] + + monkeypatch.setattr(models, "make_model", lambda *_args, **_kwargs: LogprobModel()) + metadata = [ + { + "instruction_id": "1", + "model_a": "candidate", + "model_b": "baseline", + "orientation": "direct", + } + ] + first = do_inference( + prepare_model("Dummy/judge", cache=cache), + ["judge prompt"], + return_top_logprobs=True, + cache_metadata=metadata, + ) + second = do_inference( + prepare_model("Dummy/judge", cache=cache), + ["judge prompt"], + return_top_logprobs=True, + cache_metadata=metadata, + ) + + assert second[0].text == first[0].text + assert second[0].first_token_top_logprobs == first[0].first_token_top_logprobs + assert second[0].usage is None From 13898409912321ba887cb91261b4b453dc64ca8d Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:11:24 +0200 Subject: [PATCH 3/8] Cache inference across model providers. --- judgearena/inference.py | 103 ++++++++++++++++++++++++------ judgearena/models.py | 116 ++++++++++++++++++++++------------ tests/test_inference_cache.py | 107 ++++++++++++++++++++++++++++++- 3 files changed, 265 insertions(+), 61 deletions(-) diff --git a/judgearena/inference.py b/judgearena/inference.py index 2a2e24d1..0005e285 100644 --- a/judgearena/inference.py +++ b/judgearena/inference.py @@ -6,8 +6,9 @@ from abc import ABC, abstractmethod from collections.abc import Callable from dataclasses import dataclass, field +from importlib import metadata as importlib_metadata from pathlib import Path -from typing import Any, ClassVar +from typing import Any, ClassVar, Literal import pandas as pd @@ -23,25 +24,51 @@ ) _ROLE_MAP = {"human": "user", "ai": "assistant", "system": "system"} +_CHAT_PROVIDERS = {"ChatOpenAI", "Dummy", "OpenRouter", "VLLM"} +_TEXT_PROVIDERS = {"LlamaCpp", "OpenAI", "Together"} +VLLM_EXECUTION_ONLY_KWARGS = { + "enforce_eager", + "gpu_memory_utilization", + "tensor_parallel_size", +} +InputMode = Literal["chat", "text"] -def canonicalize_chat_input(input_item: Any) -> str: - """Serialize a logical model input for content-addressed cache lookup.""" + +def provider_input_mode(provider: str) -> InputMode | None: + if provider in _CHAT_PROVIDERS: + return "chat" + if provider in _TEXT_PROVIDERS: + return "text" + return None + + +def _canonical_messages(input_item: Any) -> list[dict[str, Any]]: if isinstance(input_item, str): - payload = {"type": "text", "text": input_item} - elif hasattr(input_item, "to_messages"): - payload = { - "type": "messages", - "messages": [ - { - "role": _ROLE_MAP.get(message.type, message.type), - "content": message.content, - } - for message in input_item.to_messages() - ], - } + return [{"role": "user", "content": input_item}] + if hasattr(input_item, "to_messages"): + return [ + { + "role": _ROLE_MAP.get(message.type, message.type), + "content": message.content, + } + for message in input_item.to_messages() + ] + raise TypeError(f"Unsupported inference input: {type(input_item)!r}") + + +def canonicalize_model_input(input_item: Any, input_mode: InputMode) -> str: + """Serialize the chat messages or flattened text sent to a provider.""" + if input_mode == "text": + if isinstance(input_item, str): + text = input_item + elif hasattr(input_item, "to_string"): + text = input_item.to_string() + else: + raise TypeError(f"Unsupported inference input: {type(input_item)!r}") + payload = {"type": "text", "text": text} else: - raise TypeError(f"Unsupported inference input: {type(input_item)!r}") + payload = {"type": "messages", "messages": _canonical_messages(input_item)} return stable_json_dumps(payload) @@ -49,17 +76,51 @@ def build_model_descriptor( provider: str, model_name: str, resolved_kwargs: dict[str, Any], + endpoint: str | None = None, ) -> dict[str, Any] | None: - """Describe output-affecting settings without constructing the backend.""" - if provider != "Dummy": + """Describe output-affecting settings without constructing the backend. + + Omitted hosted settings retain provider defaults. Local-engine versions are + part of the key. VLLM tokenizer-template changes require a distinct model + revision or an explicit ``chat_template``. + """ + input_mode = provider_input_mode(provider) + if input_mode is None: return None - return { + + descriptor_kwargs = resolved_kwargs.copy() + if provider == "VLLM": + descriptor_kwargs["temperature"] = ( + 0.6 + if descriptor_kwargs.get("temperature") is None + else descriptor_kwargs["temperature"] + ) + descriptor_kwargs["top_p"] = ( + 0.95 + if descriptor_kwargs.get("top_p") is None + else descriptor_kwargs["top_p"] + ) + descriptor_kwargs["chat_template"] = descriptor_kwargs.get("chat_template") + descriptor_kwargs = { + key: value + for key, value in descriptor_kwargs.items() + if key not in VLLM_EXECUTION_ONLY_KWARGS + } + + descriptor = { "schema_version": "judgearena-inference-cache/v1", "provider": provider, "model": model_name, - "input_mode": "chat", - "model_kwargs": resolved_kwargs, + "input_mode": input_mode, + "model_kwargs": descriptor_kwargs, } + if provider == "VLLM": + descriptor["backend_version"] = importlib_metadata.version("vllm") + elif provider == "LlamaCpp": + descriptor["backend_version"] = importlib_metadata.version("llama-cpp-python") + if endpoint is not None: + descriptor["endpoint"] = endpoint.rstrip("/") + return descriptor @dataclass(frozen=True) diff --git a/judgearena/models.py b/judgearena/models.py index 8e3adc57..f370c003 100644 --- a/judgearena/models.py +++ b/judgearena/models.py @@ -22,7 +22,7 @@ InferenceCache, PreparedModel, build_model_descriptor, - canonicalize_chat_input, + canonicalize_model_input, ) from judgearena.log import get_logger from judgearena.usage import RequestUsage, RunUsage, record_usage @@ -811,7 +811,8 @@ def do_inference( if cache_metadata is None or len(cache_metadata) != len(inputs): raise ValueError("cache_metadata must contain one row per inference input.") - input_texts = [canonicalize_chat_input(item) for item in inputs] + input_mode = chat_model.descriptor["input_mode"] + input_texts = [canonicalize_model_input(item, input_mode) for item in inputs] input_hashes = [input_hash(input_text) for input_text in input_texts] with cache.open_store(chat_model) as store: cached_rows = store.query(input_hashes).set_index("input_hash") @@ -893,6 +894,61 @@ def _route_sampling_params( return engine_kwargs +_VLLM_ONLY_KWARGS = ( + "max_model_len", + "chat_template", + "language_model_only", + "gpu_memory_utilization", + "enforce_eager", + "tensor_parallel_size", + "quantization", + "kv_cache_dtype", + "reasoning_parser", + "reasoning_config", + "trust_remote_code", +) +_REMOTE_ENDPOINTS = { + "ChatOpenAI": "https://api.openai.com/v1", + "OpenAI": "https://api.openai.com/v1", + "OpenRouter": "https://openrouter.ai/api/v1", + "Together": "https://api.together.xyz/v1/completions", +} + + +def _resolve_model_config( + model: str, + max_tokens: int | None, + engine_kwargs: dict, +) -> tuple[str, str, dict, dict]: + """Resolve shared provider and sampling inputs once.""" + resolved_kwargs = engine_kwargs.copy() + resolved_kwargs["max_tokens"] = max_tokens or 8192 + sampling = { + key: resolved_kwargs.pop(key, None) + for key in ("temperature", "top_p", "top_k", "seed", "top_logprobs") + } + provider, model_name = _split_model_spec(model) + if provider != "VLLM": + for key in _VLLM_ONLY_KWARGS: + resolved_kwargs.pop(key, None) + return provider, model_name, resolved_kwargs, sampling + + +def _resolve_model_endpoint(provider: str, resolved_kwargs: dict) -> str | None: + if provider == "OpenRouter": + return _REMOTE_ENDPOINTS[provider] + for key in ("base_url", "openai_api_base", "together_api_base"): + if resolved_kwargs.get(key): + return str(resolved_kwargs[key]) + if provider in {"ChatOpenAI", "OpenAI"}: + return ( + os.getenv("OPENAI_BASE_URL") + or os.getenv("OPENAI_API_BASE") + or _REMOTE_ENDPOINTS[provider] + ) + return _REMOTE_ENDPOINTS.get(provider) + + def prepare_model( model: str, max_tokens: int | None = 8192, @@ -901,13 +957,21 @@ def prepare_model( **engine_kwargs, ) -> PreparedModel: """Prepare cache identity without constructing the provider backend.""" - provider, model_name = _split_model_spec(model) - resolved_kwargs = {**engine_kwargs, "max_tokens": max_tokens or 8192} + provider, model_name, resolved_kwargs, sampling = _resolve_model_config( + model, max_tokens, engine_kwargs + ) descriptor = ( - build_model_descriptor(provider, model_name, resolved_kwargs) + build_model_descriptor( + provider, + model_name, + {**resolved_kwargs, **sampling}, + endpoint=_resolve_model_endpoint(provider, resolved_kwargs), + ) if cache is not None else None ) + if cache is not None and descriptor is None: + logger.warning("Caching is not supported for %s; running uncached.", model) factory_kwargs = engine_kwargs.copy() return PreparedModel( model_spec=model, @@ -933,40 +997,14 @@ def make_model(model: str, max_tokens: int | None = 8192, **engine_kwargs): ``top_k``, ``seed``. vLLM-only keys (``max_model_len``, ``chat_template``) are stripped before reaching hosted providers. """ - # Avoid mutating the original engine_kwargs dictionary - # NOTE: this is a shallow copy since we are not modifying any - # mutable objects in the dictionary. - engine_kwargs = engine_kwargs.copy() - - # Dedicated arguments like max_tokens always win over engine_kwargs. - engine_kwargs["max_tokens"] = max_tokens or 8192 - - temperature = engine_kwargs.pop("temperature", None) - top_p = engine_kwargs.pop("top_p", None) - top_k = engine_kwargs.pop("top_k", None) - seed = engine_kwargs.pop("seed", None) - top_logprobs = engine_kwargs.pop("top_logprobs", None) - - model_provider, model_name = _split_model_spec(model) - - # vLLM-engine-only kwargs must not leak to remote-API providers - # (OpenRouter, OpenAI, Together): langchain-openai forwards unknown - # kwargs via model_kwargs into chat.completions.create, which rejects them. - if model_provider != "VLLM": - for key in ( - "max_model_len", - "chat_template", - "language_model_only", - "gpu_memory_utilization", - "enforce_eager", - "tensor_parallel_size", - "quantization", - "kv_cache_dtype", - "reasoning_parser", - "reasoning_config", - "trust_remote_code", - ): - engine_kwargs.pop(key, None) + model_provider, model_name, engine_kwargs, sampling = _resolve_model_config( + model, max_tokens, engine_kwargs + ) + temperature = sampling["temperature"] + top_p = sampling["top_p"] + top_k = sampling["top_k"] + seed = sampling["seed"] + top_logprobs = sampling["top_logprobs"] if model_provider == "Dummy": if top_logprobs is not None: diff --git a/tests/test_inference_cache.py b/tests/test_inference_cache.py index 3343fb97..1d4cf3db 100644 --- a/tests/test_inference_cache.py +++ b/tests/test_inference_cache.py @@ -1,7 +1,17 @@ +import json + +import pytest from langchain_core.messages import AIMessage +from langchain_core.prompts import ChatPromptTemplate +import judgearena.inference as inference import judgearena.models as models -from judgearena.inference import CompletionInferenceCache, JudgementInferenceCache +from judgearena.inference import ( + CompletionInferenceCache, + JudgementInferenceCache, + canonicalize_model_input, + provider_input_mode, +) from judgearena.models import InferenceResult, do_inference, prepare_model from judgearena.usage import track_usage @@ -104,3 +114,98 @@ def batch(self, inputs, **_kwargs): assert second[0].text == first[0].text assert second[0].first_token_top_logprobs == first[0].first_token_top_logprobs assert second[0].usage is None + + +def test_vllm_descriptor_contains_output_configuration(tmp_path, monkeypatch): + monkeypatch.setattr(inference.importlib_metadata, "version", lambda _name: "0.10.2") + cache = CompletionInferenceCache(tmp_path, "arena-hard") + + descriptor = prepare_model( + "VLLM/Qwen/Qwen3-8B", + max_tokens=32, + cache=cache, + temperature=0.2, + enforce_eager=True, + gpu_memory_utilization=0.9, + max_model_len=4096, + tensor_parallel_size=2, + ).descriptor + + assert descriptor["backend_version"] == "0.10.2" + assert descriptor["model_kwargs"]["temperature"] == 0.2 + assert descriptor["model_kwargs"]["top_p"] == 0.95 + assert descriptor["model_kwargs"]["max_model_len"] == 4096 + assert "tensor_parallel_size" not in descriptor["model_kwargs"] + + +@pytest.mark.parametrize( + ("model_spec", "expected_mode", "expected_endpoint"), + [ + ("Dummy/model", "chat", None), + ("VLLM/org/model", "chat", None), + ("OpenRouter/org/model", "chat", "https://openrouter.ai/api/v1"), + ("ChatOpenAI/model", "chat", "https://api.openai.com/v1"), + ("OpenAI/model", "text", "https://api.openai.com/v1"), + ("Together/org/model", "text", "https://api.together.xyz/v1/completions"), + ("LlamaCpp/./models/model.gguf", "text", None), + ], +) +def test_supported_provider_descriptors( + tmp_path, + monkeypatch, + model_spec, + expected_mode, + expected_endpoint, +): + versions = {"vllm": "0.10.2", "llama-cpp-python": "0.3.0"} + monkeypatch.setattr(inference.importlib_metadata, "version", versions.__getitem__) + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + cache = CompletionInferenceCache(tmp_path, "arena-hard") + + descriptor = prepare_model(model_spec, cache=cache).descriptor + + assert descriptor["input_mode"] == expected_mode + assert descriptor.get("endpoint") == expected_endpoint + + +def test_hosted_descriptor_hashes_routing_and_endpoint(tmp_path, monkeypatch): + cache = CompletionInferenceCache(tmp_path, "arena-hard") + unpinned = prepare_model("OpenRouter/org/model", cache=cache).descriptor + + pinned = prepare_model( + "OpenRouter/org/model", + cache=cache, + extra_body={"provider": {"order": ["Together"], "allow_fallbacks": False}}, + ).descriptor + assert unpinned != pinned + + monkeypatch.setenv("OPENAI_BASE_URL", "https://ambient.example/v1") + ambient = prepare_model("ChatOpenAI/model", cache=cache).descriptor + assert ambient["endpoint"] == "https://ambient.example/v1" + + +def test_unsupported_provider_runs_uncached(tmp_path, caplog): + model = prepare_model( + "Unsupported/model", + cache=CompletionInferenceCache(tmp_path, "arena-hard"), + ) + + assert model.descriptor is None + assert "Caching is not supported" in caplog.text + + +@pytest.mark.parametrize( + ("provider", "expected_type"), + [("OpenRouter", "messages"), ("Together", "text"), ("VLLM", "messages")], +) +def test_input_canonicalization_matches_provider_mode(provider, expected_type): + prompt = ChatPromptTemplate.from_messages( + [("system", "System"), ("user", "Question")] + ).invoke({}) + + payload = json.loads( + canonicalize_model_input(prompt, provider_input_mode(provider)) + ) + + assert payload["type"] == expected_type From b5ca140e6b0a55c77e98094f55ffc0e6766ff5ff Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Wed, 9 Sep 2026 14:11:29 +0200 Subject: [PATCH 4/8] Adopt inference caching in task runners. --- judgearena/benchmarks/elo/calibration.py | 21 +++- judgearena/benchmarks/elo/runner.py | 96 +++++++---------- judgearena/benchmarks/execution.py | 19 +++- judgearena/benchmarks/meta_eval/annotate.py | 9 ++ .../benchmarks/mt_bench/fastchat_compat.py | 4 + .../benchmarks/mt_bench/pairwise_judging.py | 13 +++ .../benchmarks/mt_bench/preset_judging.py | 4 + judgearena/benchmarks/mt_bench/runner.py | 32 +++--- judgearena/benchmarks/pairwise/runner.py | 46 ++++---- judgearena/config.py | 4 +- judgearena/evaluate.py | 22 ++++ judgearena/generate.py | 91 +++++++++++----- judgearena/utils/__init__.py | 6 -- judgearena/utils/io.py | 102 ------------------ tests/test_config.py | 4 +- tests/test_elo_calibration.py | 3 +- tests/test_estimate_elo_ratings.py | 21 ++-- tests/test_generate_and_evaluate.py | 15 +-- tests/test_meta_eval_annotations.py | 2 + tests/test_mt_bench_downloads.py | 15 +-- tests/test_seed_plumbing.py | 12 --- tests/test_usage.py | 10 +- 22 files changed, 272 insertions(+), 279 deletions(-) diff --git a/judgearena/benchmarks/elo/calibration.py b/judgearena/benchmarks/elo/calibration.py index dd3ebb9d..10575f35 100644 --- a/judgearena/benchmarks/elo/calibration.py +++ b/judgearena/benchmarks/elo/calibration.py @@ -11,8 +11,9 @@ from judgearena.arenas_utils import extract_turn_text from judgearena.benchmarks.elo.rating import winner_to_pref from judgearena.evaluate import judge_and_parse_prefs +from judgearena.inference import JudgementInferenceCache from judgearena.log import get_logger -from judgearena.models import make_model +from judgearena.models import prepare_model from judgearena.prompts.parsing import PairScore from judgearena.prompts.registry import ResolvedJudgePrompt @@ -62,6 +63,8 @@ def calibrate_pairscore_temperature( prompt: ResolvedJudgePrompt, truncate_input_chars: int | None, default_temperature: float, + arena: str, + inference_cache: JudgementInferenceCache | None = None, ) -> float | None: """Judge sampled human battles and return a fitted PairScore temperature.""" if not enabled: @@ -103,7 +106,12 @@ def calibrate_pairscore_temperature( for index in calibration_battles.index ] - calibration_judge = make_model(model=judge_model, **dict(judge_model_kwargs)) + source_rows = source_battles.loc[calibration_battles.index] + calibration_judge = prepare_model( + model=judge_model, + cache=inference_cache, + **dict(judge_model_kwargs), + ) annotations, _, _ = judge_and_parse_prefs( judge_chat_model=calibration_judge, instructions=instructions, @@ -115,6 +123,15 @@ def calibrate_pairscore_temperature( prompt_preset=prompt.preset_name, parse=prompt.parser, truncate_input_chars=truncate_input_chars, + cache_metadata=[ + { + "instruction_id": f"{arena}:{row.question_id}", + "model_a": row.model_a, + "model_b": row.model_b, + "orientation": "direct", + } + for row in source_rows.itertuples() + ], ) score_differences: list[float] = [] diff --git a/judgearena/benchmarks/elo/runner.py b/judgearena/benchmarks/elo/runner.py index 64b15460..7815784c 100644 --- a/judgearena/benchmarks/elo/runner.py +++ b/judgearena/benchmarks/elo/runner.py @@ -1,6 +1,4 @@ -import hashlib from datetime import UTC, datetime -from functools import partial from pathlib import Path from typing import TYPE_CHECKING @@ -19,10 +17,13 @@ from judgearena.benchmarks.elo.rating import ( arena_anchor_battles, prefs_to_battle_results, - sampling_cache_token, select_seeded_random_arena_battles, ) -from judgearena.benchmarks.execution import build_generation_kwargs +from judgearena.benchmarks.execution import ( + build_completion_cache, + build_generation_kwargs, + build_judgement_cache, +) from judgearena.benchmarks.scoring import build_metrics, calculate_metrics from judgearena.datasets import load_battles from judgearena.evaluate import ( @@ -33,10 +34,9 @@ ) from judgearena.generate import generate_instructions from judgearena.log import get_logger -from judgearena.models import build_default_judge_model_kwargs, make_model +from judgearena.models import build_default_judge_model_kwargs, prepare_model from judgearena.reports import EloReport from judgearena.tasks.schema import EloProtocol, ResolvedTaskSpec -from judgearena.utils import cache_function_dataframe if TYPE_CHECKING: from judgearena.config import RunConfig @@ -108,6 +108,11 @@ def run_elo(cfg: "RunConfig", task: ResolvedTaskSpec | None = None) -> dict: extract_turn_text(row["conversation_a"][0]) for _, row in df_battles.iterrows() ], + index=( + arena + ":" + df_battles["question_id"].astype(str) + if "question_id" in df_battles + else df_battles.index.astype(str) + ), name="instruction", ) logger.debug("First instruction:\n%s", instructions.iloc[0][:300]) @@ -121,48 +126,13 @@ def run_elo(cfg: "RunConfig", task: ResolvedTaskSpec | None = None) -> dict: # dropped battle_thinking_token_budget). extra_kwargs = build_generation_kwargs(cfg, cfg.model.name, role="A") use_tqdm = False - gen_fun = partial( - generate_instructions, + completions_df = generate_instructions( + instructions=instructions, + model=cfg.model.name, truncate_input_chars=cfg.generation.truncate_all_input_chars, use_tqdm=use_tqdm, + inference_cache=build_completion_cache(cfg), **extra_kwargs, - ) - - def replace_slash(s: str) -> str: - return s.replace("/", "_") - - languages_str = ( - "-".join(sorted(selected_languages)) if selected_languages else "all" - ) - extra_kwargs_str = ( - "_".join(f"{k}={v}" for k, v in sorted(extra_kwargs.items())) - if extra_kwargs - else "" - ) - cache_token = sampling_cache_token( - sampling_metadata, - n_instructions=cfg.generation.n_instructions, - n_instructions_per_language=cfg.elo.n_instructions_per_language, - ) - cache_suffix = ( - f"{arena}_{replace_slash(cfg.model.name)}_" - f"{cache_token}_" - f"{languages_str}_{cfg.generation.truncate_all_input_chars}_{extra_kwargs['max_tokens']}" - + (f"_{extra_kwargs_str}" if extra_kwargs_str else "") - ) - if len(cache_suffix) > 100: - cache_hash = hashlib.sha256(cache_suffix.encode()).hexdigest()[:16] - logger.debug( - "Cache suffix too long (%d chars), using hash: %s (full: %s)", - len(cache_suffix), - cache_hash, - cache_suffix, - ) - cache_suffix = cache_hash - completions_df = cache_function_dataframe( - lambda: gen_fun(instructions=instructions, model=cfg.model.name), - ignore_cache=cfg.run.ignore_cache, - cache_name=f"elo/{cache_suffix}", ).set_index("instruction_index") completions = completions_df.loc[:, "completion"] @@ -210,8 +180,9 @@ def replace_slash(s: str) -> str: ) def run_judge() -> pd.DataFrame: - judge_chat_model = make_model( + judge_chat_model = prepare_model( model=cfg.judge.model, + cache=build_judgement_cache(cfg), **judge_extra_kwargs, ) annotations, annotations_reversed, prefs = judge_and_parse_prefs( @@ -227,6 +198,25 @@ def run_judge() -> pd.DataFrame: parse=resolved_prompt.parser, truncate_input_chars=cfg.generation.truncate_judge_input_chars, use_tqdm=use_tqdm, + cache_metadata=[ + { + "instruction_id": str(instructions.index[index]), + "model_a": ( + cfg.model.name + if our_model_is_position_a[index] + else opponent_models[index] + ), + "model_b": ( + opponent_models[index] + if our_model_is_position_a[index] + else cfg.model.name + ), + "orientation": ( + "direct" if our_model_is_position_a[index] else "reversed" + ), + } + for index in range(n) + ], ) if annotations_reversed is None: row_annotations = list(annotations) @@ -259,17 +249,7 @@ def run_judge() -> pd.DataFrame: ) return frame - # Stripping reasoning traces changes the judged text but not the cached - # completions, so it must be part of the judge cache key. Only append when - # enabled so prior (non-stripped) runs keep their existing cache hashes. - judge_cache_suffix = f"judge_{cache_suffix}" - if cfg.judge.strip_thinking_before_judging: - judge_cache_suffix += "_stripthinking" - df_judge = cache_function_dataframe( - run_judge, - ignore_cache=cfg.run.ignore_cache, - cache_name=f"elo/{judge_cache_suffix}", - ) + df_judge = run_judge() # Restore position arrays and prefs from cache (in case loaded from disk) use_model_a_as_opponent = df_judge["use_model_a_as_opponent"].to_numpy() @@ -308,6 +288,8 @@ def run_judge() -> pd.DataFrame: prompt=resolved_prompt, truncate_input_chars=cfg.generation.truncate_judge_input_chars, default_temperature=cfg.elo.soft_elo_temperature, + arena=arena, + inference_cache=build_judgement_cache(cfg), ) # Reparse cached responses at this run's temperature, keeping the selected diff --git a/judgearena/benchmarks/execution.py b/judgearena/benchmarks/execution.py index 3ec96581..8a173ca5 100644 --- a/judgearena/benchmarks/execution.py +++ b/judgearena/benchmarks/execution.py @@ -2,12 +2,14 @@ from __future__ import annotations +from pathlib import Path from typing import TYPE_CHECKING, Literal +from judgearena.inference import CompletionInferenceCache, JudgementInferenceCache from judgearena.models import ( build_default_judge_model_kwargs, is_thinking_model, - make_model, + prepare_model, ) if TYPE_CHECKING: @@ -35,8 +37,9 @@ def build_generation_kwargs( def build_judge(cfg: RunConfig): """Construct the configured judge consistently across benchmark runners.""" - return make_model( + return prepare_model( model=cfg.judge.model, + cache=build_judgement_cache(cfg), **build_default_judge_model_kwargs( cfg.judge.model, cfg.model.engine_kwargs, @@ -45,3 +48,15 @@ def build_judge(cfg: RunConfig): ), ), ) + + +def build_completion_cache(cfg: RunConfig) -> CompletionInferenceCache | None: + if cfg.run.store_root is None: + return None + return CompletionInferenceCache(Path(cfg.run.store_root), cfg.task) + + +def build_judgement_cache(cfg: RunConfig) -> JudgementInferenceCache | None: + if cfg.run.store_root is None: + return None + return JudgementInferenceCache(Path(cfg.run.store_root), cfg.task) diff --git a/judgearena/benchmarks/meta_eval/annotate.py b/judgearena/benchmarks/meta_eval/annotate.py index be47be30..843ed6ff 100644 --- a/judgearena/benchmarks/meta_eval/annotate.py +++ b/judgearena/benchmarks/meta_eval/annotate.py @@ -161,6 +161,15 @@ def annotate_sample( parse=parser, truncate_input_chars=cfg.generation.truncate_judge_input_chars, use_tqdm=cfg.run.use_tqdm, + cache_metadata=[ + { + "instruction_id": battle["battle_id"], + "model_a": battle["model_a"], + "model_b": battle["model_b"], + "orientation": "direct", + } + for _, battle in df_sample.iterrows() + ], ) n_battles = len(df_sample) diff --git a/judgearena/benchmarks/mt_bench/fastchat_compat.py b/judgearena/benchmarks/mt_bench/fastchat_compat.py index 45aa7424..b20b22ae 100644 --- a/judgearena/benchmarks/mt_bench/fastchat_compat.py +++ b/judgearena/benchmarks/mt_bench/fastchat_compat.py @@ -362,6 +362,8 @@ def judge_mt_bench_pairwise_fastchat( items=items, use_tqdm=use_tqdm, swap_answers=False, + model_a=model_a, + model_b=model_b, ) g2_judgments: list[str] | None = None @@ -371,6 +373,8 @@ def judge_mt_bench_pairwise_fastchat( items=items, use_tqdm=use_tqdm, swap_answers=True, + model_a=model_a, + model_b=model_b, ) annotations: list[dict[str, Any]] = [] diff --git a/judgearena/benchmarks/mt_bench/pairwise_judging.py b/judgearena/benchmarks/mt_bench/pairwise_judging.py index fce48b8e..f07669f3 100644 --- a/judgearena/benchmarks/mt_bench/pairwise_judging.py +++ b/judgearena/benchmarks/mt_bench/pairwise_judging.py @@ -76,6 +76,8 @@ def infer_pairwise_judgments_by_prompt_groups( items: list[MTBenchJudgeItem], use_tqdm: bool, swap_answers: bool, + model_a: str, + model_b: str, ) -> tuple[list[str], list[dict[str, str]]]: judgments: list[str] = [""] * len(items) used_prompt_kwargs: list[dict[str, str]] = [{} for _ in items] @@ -97,6 +99,17 @@ def infer_pairwise_judgments_by_prompt_groups( inputs=prompt_inputs, use_tqdm=use_tqdm, stage="judging", + cache_metadata=[ + { + "instruction_id": ( + f"{items[item_index].question_id}:turn-{items[item_index].turn}" + ), + "model_a": model_b if swap_answers else model_a, + "model_b": model_a if swap_answers else model_b, + "orientation": "reversed" if swap_answers else "direct", + } + for item_index in idxs + ], ) for item_index, output, prompt_kwargs in zip( idxs, outputs, batch_kwargs, strict=True diff --git a/judgearena/benchmarks/mt_bench/preset_judging.py b/judgearena/benchmarks/mt_bench/preset_judging.py index fc0de802..efe34891 100644 --- a/judgearena/benchmarks/mt_bench/preset_judging.py +++ b/judgearena/benchmarks/mt_bench/preset_judging.py @@ -178,6 +178,8 @@ def judge_mt_bench_with_preset( items=items, use_tqdm=use_tqdm, swap_answers=False, + model_a=model_a, + model_b=model_b, ) annotations: list[dict[str, Any]] = [] @@ -246,6 +248,8 @@ def _append_results( items=items, use_tqdm=use_tqdm, swap_answers=True, + model_a=model_a, + model_b=model_b, ) ) _append_results(swapped_judgments, swapped_prompt_kwargs, swapped=True) diff --git a/judgearena/benchmarks/mt_bench/runner.py b/judgearena/benchmarks/mt_bench/runner.py index 140f9916..4c05252b 100644 --- a/judgearena/benchmarks/mt_bench/runner.py +++ b/judgearena/benchmarks/mt_bench/runner.py @@ -15,6 +15,10 @@ import pandas as pd from judgearena.artifacts import prepare_run_directory, write_run_metadata_safely +from judgearena.benchmarks.execution import ( + build_completion_cache, + build_judgement_cache, +) from judgearena.benchmarks.mt_bench.fastchat_compat import ( judge_mt_bench_pairwise_fastchat, ) @@ -27,14 +31,10 @@ ) from judgearena.generate import generate_multiturn from judgearena.log import get_logger -from judgearena.models import is_thinking_model, make_model +from judgearena.models import is_thinking_model, prepare_model from judgearena.prompts.registry import ResolvedJudgePrompt, resolve_run_judge_prompt from judgearena.reports import BattleReport from judgearena.tasks.schema import MTBenchProtocol -from judgearena.utils import ( - cache_function_dataframe, - generation_cache_token, -) logger = get_logger(__name__) @@ -87,8 +87,6 @@ def _generate_mt_bench_completions( protocol: MTBenchProtocol, questions_df: pd.DataFrame, ) -> tuple[pd.DataFrame, pd.DataFrame]: - cache_prefix = cfg.task - def _run_generation( model_name: str, *, generation_kwargs: dict[str, object] ) -> pd.DataFrame: @@ -107,6 +105,7 @@ def _run_generation( use_tqdm=cfg.run.use_tqdm, temperature_config=temperature_config, strip_thinking_before_turn_2_prompt=cfg.judge.strip_thinking_before_judging, + inference_cache=build_completion_cache(cfg), **generation_kwargs, ) @@ -120,19 +119,12 @@ def _load_or_generate(model_name: str, *, role: str) -> pd.DataFrame: completions=loaded_answers, model_name=model_name, ) - # Fold the resolved generation kwargs into the cache key so changing any - # sampling param busts cached completions instead of reusing a stale run. generation_kwargs = _build_mt_bench_generation_kwargs( cfg=cfg, model_spec=model_name, role=role ) - sampling_token = generation_cache_token(generation_kwargs) - generated_answers = cache_function_dataframe( - lambda: _run_generation(model_name, generation_kwargs=generation_kwargs), - ignore_cache=cfg.run.ignore_cache, - cache_name=( - f"{cache_prefix}_{model_name}_{cfg.generation.n_instructions}_" - f"{sampling_token}" - ), + generated_answers = _run_generation( + model_name, + generation_kwargs=generation_kwargs, ) return _align_mt_bench_completions( questions_df=questions_df, @@ -461,7 +453,11 @@ def run_mt_bench_benchmark(cfg: RunConfig, task: ResolvedTaskSpec | None = None) judge_model_kwargs.setdefault( "temperature", protocol.judge.fastchat_temperature ) - judge_chat_model = make_model(model=cfg.judge.model, **judge_model_kwargs) + judge_chat_model = prepare_model( + model=cfg.judge.model, + cache=build_judgement_cache(cfg), + **judge_model_kwargs, + ) if resolved_prompt.delegated: return _run_mt_bench_fastchat( cfg=cfg, diff --git a/judgearena/benchmarks/pairwise/runner.py b/judgearena/benchmarks/pairwise/runner.py index 0382205d..63980e88 100644 --- a/judgearena/benchmarks/pairwise/runner.py +++ b/judgearena/benchmarks/pairwise/runner.py @@ -12,7 +12,11 @@ import pandas as pd from judgearena.artifacts import prepare_run_directory, write_run_metadata_safely -from judgearena.benchmarks.execution import build_generation_kwargs, build_judge +from judgearena.benchmarks.execution import ( + build_completion_cache, + build_generation_kwargs, + build_judge, +) from judgearena.benchmarks.pairwise.baselines import resolve_baseline_plan from judgearena.benchmarks.scoring import build_metrics, calculate_metrics from judgearena.datasets.pairwise import load_pairwise_task_data @@ -22,7 +26,6 @@ from judgearena.reports import BattleReport from judgearena.tasks.registry import get_packaged_task from judgearena.tasks.schema import ResolvedTaskSpec -from judgearena.utils import cache_function_dataframe, generation_cache_token if TYPE_CHECKING: from judgearena.config import RunConfig @@ -99,11 +102,6 @@ def run_pairwise(cfg: "RunConfig", resolved_task: ResolvedTaskSpec | None = None run_started_at = datetime.now(UTC) - # Not working with vllm, not detecting model changes and serving the same cache for two different models... - # if not cfg.run.ignore_cache: - # set_langchain_cache() - ignore_cache = cfg.run.ignore_cache - resolved_task = resolved_task or get_packaged_task(cfg.task) if resolved_task is None: raise ValueError(f"Unknown task {cfg.task!r}.") @@ -201,6 +199,7 @@ def _run_generation( model=model_spec, truncate_input_chars=cfg.generation.truncate_all_input_chars, use_tqdm=cfg.run.use_tqdm, + inference_cache=build_completion_cache(cfg), **generation_kwargs, ) @@ -219,18 +218,10 @@ def _load_or_generate_completions( ) if preloaded is not None: return preloaded - # Fold the resolved generation kwargs into the cache key so that changing - # any sampling param (temperature, seed, top_p/k, max_tokens, ...) busts - # the cached completions instead of silently reusing a stale run. generation_kwargs = build_generation_kwargs(cfg, model_spec, role=role) - sampling_token = generation_cache_token(generation_kwargs) - generated = cache_function_dataframe( - lambda: _run_generation(model_spec, generation_kwargs=generation_kwargs), - ignore_cache=ignore_cache, - cache_name=( - f"{cfg.task}_{model_spec}_{cfg.generation.n_instructions}_" - f"{sampling_token}" - ), + generated = _run_generation( + model_spec, + generation_kwargs=generation_kwargs, ) completions = _align_completion_series(generated) return ( @@ -297,6 +288,12 @@ def _load_or_generate_completions( judged_B = completions_A.mask(swap_mask, completions_B) else: judged_A, judged_B = completions_A, completions_B + candidate_models = pd.Series(cfg.model.name, index=instructions.index) + if swap_mask is not None: + judged_model_a = baseline_per_index.mask(swap_mask, cfg.model.name) + judged_model_b = candidate_models.mask(swap_mask, baseline_per_index) + else: + judged_model_a, judged_model_b = baseline_per_index, candidate_models annotations = [] annotations_reversed = [] if cfg.judge.swap_mode == "both" else None @@ -315,6 +312,19 @@ def _load_or_generate_completions( parse=group_prompt.parser, truncate_input_chars=cfg.generation.truncate_judge_input_chars, use_tqdm=cfg.run.use_tqdm, + cache_metadata=[ + { + "instruction_id": str(index), + "model_a": judged_model_a.loc[index], + "model_b": judged_model_b.loc[index], + "orientation": ( + "direct" + if judged_model_a.loc[index] == cfg.model.name + else "reversed" + ), + } + for index in group_index + ], ) annotations.extend(group_annotations) if group_reversed is not None: diff --git a/judgearena/config.py b/judgearena/config.py index 5f82f8c9..8e2280e5 100644 --- a/judgearena/config.py +++ b/judgearena/config.py @@ -410,8 +410,8 @@ class RunArgs(BaseModel): """Directory where annotations, results, and the resolved ``config.yaml`` are written (under a per-run subfolder).""" - ignore_cache: bool = False - """If set, ignore cached completions and regenerate them.""" + store_root: str | None = None + """Root directory for content-addressed completion and judgement caches.""" use_tqdm: bool = False """Show a tqdm progress bar (not compatible with vLLM).""" diff --git a/judgearena/evaluate.py b/judgearena/evaluate.py index fbf81054..b5b4c6cd 100644 --- a/judgearena/evaluate.py +++ b/judgearena/evaluate.py @@ -111,6 +111,7 @@ def annotate_battles( prompt_preset: str = DEFAULT_JUDGE_PROMPT_PRESET, strip_thinking_before_judging: bool = False, collect_top_logprobs: bool = False, + cache_metadata: list[dict] | None = None, ) -> list[JudgeAnnotation]: """ Directly evaluate from list of instructions and completions @@ -182,6 +183,7 @@ def annotate_battles( use_tqdm=use_tqdm, return_top_logprobs=collect_top_logprobs, stage="judging", + cache_metadata=cache_metadata, ) if not collect_top_logprobs: judge_results = [InferenceResult(text=text) for text in judge_results] @@ -229,6 +231,7 @@ def judge_and_parse_prefs( truncate_input_chars: int = 8192, use_tqdm: bool = False, parse: JudgeParser | None = None, + cache_metadata: list[dict] | None = None, ) -> tuple[list[JudgeAnnotation], list[JudgeAnnotation] | None, pd.Series]: """Run judge annotation and parse preferences, handling swap_mode='both'. @@ -266,10 +269,28 @@ def judge_and_parse_prefs( truncate_input_chars=truncate_input_chars, use_tqdm=use_tqdm, collect_top_logprobs=parse.requires_top_logprobs, + cache_metadata=cache_metadata, ) annotations_reversed = None if swap_mode == "both": + reversed_cache_metadata = ( + [ + { + **metadata, + "model_a": metadata["model_b"], + "model_b": metadata["model_a"], + "orientation": ( + "reversed" + if metadata.get("orientation") == "direct" + else "direct" + ), + } + for metadata in cache_metadata + ] + if cache_metadata is not None + else None + ) annotations_reversed = annotate_battles( judge_chat_model=judge_chat_model, instructions=instructions, @@ -282,6 +303,7 @@ def judge_and_parse_prefs( truncate_input_chars=truncate_input_chars, use_tqdm=use_tqdm, collect_top_logprobs=parse.requires_top_logprobs, + cache_metadata=reversed_cache_metadata, ) def _none_to_nan(x): diff --git a/judgearena/generate.py b/judgearena/generate.py index edc4e60a..4ca968f5 100644 --- a/judgearena/generate.py +++ b/judgearena/generate.py @@ -1,7 +1,8 @@ import pandas as pd from langchain_core.prompts import ChatPromptTemplate -from judgearena.models import batch_inference_once, do_inference, make_model +from judgearena.inference import CompletionInferenceCache +from judgearena.models import PreparedModel, do_inference, prepare_model from judgearena.utils import strip_thinking_tags, truncate @@ -12,9 +13,15 @@ def generate_instructions( max_tokens: int | None = 32768, use_tqdm: bool = True, system_prompt: str | None = None, + inference_cache: CompletionInferenceCache | None = None, **engine_kwargs, ) -> pd.DataFrame: - chat_model = make_model(model, max_tokens=max_tokens, **engine_kwargs) + chat_model = prepare_model( + model, + max_tokens=max_tokens, + cache=inference_cache, + **engine_kwargs, + ) # TODO improve prompt to generate instructions if system_prompt is None: @@ -39,6 +46,7 @@ def generate_instructions( inputs=inputs, use_tqdm=use_tqdm, stage="generation", + cache_metadata=[{"instruction_id": str(index)} for index in instructions.index], ) df_outputs = pd.DataFrame( data={ @@ -57,16 +65,23 @@ def _set_temperature_on_model(chat_model, temperature: float) -> None: chat_model.temperature = temperature +def _materialize_at_temperature(base_model: PreparedModel, temperature: float): + model = base_model.materialize() + _set_temperature_on_model(model, temperature) + return model + + def _infer_grouped_by_temperature( *, model_spec: str, - provider: str, max_tokens: int | None, model_kwargs: dict, - base_model, + base_model: PreparedModel | None, inputs: list, + instruction_ids: list[str], temperatures: list[float], use_tqdm: bool, + inference_cache: CompletionInferenceCache | None, ) -> list[str]: outputs: list[str] = [""] * len(inputs) groups: dict[float, list[int]] = {} @@ -77,12 +92,16 @@ def _infer_grouped_by_temperature( idxs = groups[temp] group_inputs = [inputs[i] for i in idxs] - if provider in {"VLLM", "LlamaCpp"}: - _set_temperature_on_model(base_model, temp) - group_model = base_model - else: - group_model = make_model( - model_spec, max_tokens=max_tokens, temperature=temp, **model_kwargs + group_model = prepare_model( + model_spec, + max_tokens=max_tokens, + temperature=temp, + cache=inference_cache, + **model_kwargs, + ) + if base_model is not None: + group_model.factory = lambda temperature=temp: _materialize_at_temperature( + base_model, temperature ) group_outs = do_inference( @@ -90,6 +109,9 @@ def _infer_grouped_by_temperature( inputs=group_inputs, use_tqdm=use_tqdm, stage="generation", + cache_metadata=[ + {"instruction_id": instruction_ids[index]} for index in idxs + ], ) for i, out in zip(idxs, group_outs, strict=True): outputs[i] = out @@ -105,6 +127,7 @@ def generate_multiturn( use_tqdm: bool = True, temperature_config: dict[str, float] | None = None, strip_thinking_before_turn_2_prompt: bool = False, + inference_cache: CompletionInferenceCache | None = None, **model_kwargs, ) -> pd.DataFrame: """Generate two-turn completions for MT-Bench style questions.""" @@ -112,12 +135,22 @@ def generate_multiturn( use_category_temperatures = temperature_config is not None local_provider = provider in {"VLLM", "LlamaCpp"} - if use_category_temperatures and local_provider: - chat_model = make_model( - model, max_tokens=max_tokens, temperature=0.0, **model_kwargs + base_model = ( + prepare_model( + model, + max_tokens=max_tokens, + temperature=0.0, + **model_kwargs, ) - else: - chat_model = make_model(model, max_tokens=max_tokens, **model_kwargs) + if use_category_temperatures and local_provider + else None + ) + chat_model = prepare_model( + model, + max_tokens=max_tokens, + cache=inference_cache, + **model_kwargs, + ) system_prompt = "You are a helpful assistant." idxs = questions.index.tolist() @@ -141,13 +174,14 @@ def generate_multiturn( if use_category_temperatures: completions_turn_1 = _infer_grouped_by_temperature( model_spec=model, - provider=provider, max_tokens=max_tokens, model_kwargs=model_kwargs, - base_model=chat_model, + base_model=base_model, inputs=turn1_inputs, + instruction_ids=[f"{index}:turn-1" for index in idxs], temperatures=temperatures, use_tqdm=use_tqdm, + inference_cache=inference_cache, ) else: completions_turn_1 = do_inference( @@ -155,6 +189,7 @@ def generate_multiturn( inputs=turn1_inputs, use_tqdm=use_tqdm, stage="generation", + cache_metadata=[{"instruction_id": f"{index}:turn-1"} for index in idxs], ) turn2_inputs = [] @@ -196,13 +231,14 @@ def generate_multiturn( if use_category_temperatures: completions_turn_2 = _infer_grouped_by_temperature( model_spec=model, - provider=provider, max_tokens=max_tokens, model_kwargs=model_kwargs, - base_model=chat_model, + base_model=base_model, inputs=turn2_inputs, + instruction_ids=[f"{index}:turn-2" for index in idxs], temperatures=temperatures, use_tqdm=use_tqdm, + inference_cache=inference_cache, ) else: completions_turn_2 = do_inference( @@ -210,6 +246,7 @@ def generate_multiturn( inputs=turn2_inputs, use_tqdm=use_tqdm, stage="generation", + cache_metadata=[{"instruction_id": f"{index}:turn-2"} for index in idxs], ) return pd.DataFrame( @@ -227,20 +264,26 @@ def generate_base( truncate_input_chars: int | None = 8192, max_tokens: int | None = 32768, use_tqdm: bool = False, + inference_cache: CompletionInferenceCache | None = None, **engine_kwargs, ) -> pd.DataFrame: - chat_model = make_model(model, max_tokens=max_tokens, **engine_kwargs) + chat_model = prepare_model( + model, + max_tokens=max_tokens, + cache=inference_cache, + **engine_kwargs, + ) inputs = [ truncate(instruction, max_len=truncate_input_chars) for instruction in instructions ] - completions = batch_inference_once( - chat_model, - inputs, - max_tokens=max_tokens, + completions = do_inference( + chat_model=chat_model, + inputs=inputs, stage="generation", + cache_metadata=[{"instruction_id": str(index)} for index in instructions.index], ) df_outputs = pd.DataFrame( diff --git a/judgearena/utils/__init__.py b/judgearena/utils/__init__.py index 797e1b49..833d13de 100644 --- a/judgearena/utils/__init__.py +++ b/judgearena/utils/__init__.py @@ -13,12 +13,9 @@ compute_pref_summary, ) from judgearena.utils.io import ( - Timeblock, - cache_function_dataframe, data_root, download_all, download_hf, - generation_cache_token, read_df, safe_parse_int, ) @@ -33,13 +30,10 @@ "BattleReport", "PrefSummary", "Report", - "Timeblock", - "cache_function_dataframe", "compute_pref_summary", "data_root", "download_all", "download_hf", - "generation_cache_token", "read_df", "safe_parse_int", "safe_text", diff --git a/judgearena/utils/io.py b/judgearena/utils/io.py index 4e49fde9..aeab573e 100644 --- a/judgearena/utils/io.py +++ b/judgearena/utils/io.py @@ -2,10 +2,7 @@ from __future__ import annotations -import hashlib import os -import time -from collections.abc import Callable from pathlib import Path import pandas as pd @@ -74,104 +71,5 @@ def download_all(): download_hf(name=task_id, local_path=local_path_tables) -class Timeblock: - """Timer context manager""" - - def __init__(self, name: str | None = None, verbose: bool = True): - self.name = name - self.verbose = verbose - - def __enter__(self): - """Start a new timer as a context manager""" - self.start = time.time() - return self - - def __exit__(self, *args): - """Stop the context manager timer""" - self.end = time.time() - self.duration = self.end - self.start - if self.verbose: - logger.info("%s", self) - - def __str__(self): - name = self.name if self.name else "block" - msg = f"{name} took {self.duration} seconds" - return msg - - -def generation_cache_token(kwargs: dict[str, object]) -> str: - """Short, deterministic token of generation kwargs for cache-key busting. - - Folds the resolved per-role generation kwargs (sampling params, max_tokens, - chat_template, ...) into a stable 16-char hash so that changing any of them - invalidates cached completions. Hashing keeps the cache name bounded even - when a long ``chat_template`` is present. - """ - serialized = "_".join(f"{k}={kwargs[k]!r}" for k in sorted(kwargs)) - return hashlib.sha256(serialized.encode()).hexdigest()[:16] - - -def cache_function_dataframe( - fun: Callable[[], pd.DataFrame], - cache_name: str, - ignore_cache: bool = False, - cache_path: Path | None = None, - parquet: bool = False, -) -> pd.DataFrame: - """ - :param fun: a function whose dataframe result obtained `fun()` will be cached - :param cache_name: the cache of the function result is written into `{cache_path}/{cache_name}.csv.zip` - :param ignore_cache: whether to recompute even if the cache is present - :param cache_path: folder where to write cache files, default to ~/cache-zeroshot/ - :param parquet: whether to store the data in parquet, if not specified use csv.zip - :return: result of fun() - """ - if cache_path is None: - cache_path = data_root / "cache" - - if parquet: - cache_file = cache_path / (cache_name + ".parquet") - else: - cache_file = cache_path / (cache_name + ".csv.zip") - cache_file.parent.mkdir(parents=True, exist_ok=True) - if cache_file.exists() and not ignore_cache: - logger.info("Loading cache %s", cache_file) - if parquet: - return pd.read_parquet(cache_file) - else: - return pd.read_csv(cache_file) - else: - logger.info( - "Cache %s not found or ignore_cache set to True, regenerating the file", - cache_file, - ) - with Timeblock("Evaluate function."): - df = fun() - assert isinstance(df, pd.DataFrame) - if parquet: - # object cols cannot be saved easily in parquet; numpy arrays must be - # deep-converted to plain Python so str() produces ast.literal_eval-safe - # repr (no "array([...])" syntax, which breaks literal_eval) - import numpy as np - - def _to_python(x): - """Recursively convert numpy arrays/scalars to Python lists/dicts.""" - if isinstance(x, np.ndarray): - return [_to_python(i) for i in x] - if isinstance(x, dict): - return {k: _to_python(v) for k, v in x.items()} - if isinstance(x, list): - return [_to_python(i) for i in x] - return x - - for col in df.select_dtypes(include="object").columns: - df[col] = df[col].apply(_to_python).astype(str) - df.to_parquet(cache_file, index=False) - return pd.read_parquet(cache_file) - else: - df.to_csv(cache_file, index=False) - return pd.read_csv(cache_file) - - if __name__ == "__main__": download_all() diff --git a/tests/test_config.py b/tests/test_config.py index e09d8b9e..ff0d2ffe 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -201,7 +201,7 @@ def test_judge_settings_include_inherited_engine_kwargs( monkeypatch.setattr( config_module, "get_packaged_task", lambda _task: _registered_task() ) - monkeypatch.setattr(execution_module, "make_model", lambda **kwargs: kwargs) + monkeypatch.setattr(execution_module, "prepare_model", lambda **kwargs: kwargs) data = _base_generate() data["model"]["engine_kwargs"] = { "temperature": 0.5, @@ -235,7 +235,7 @@ def test_judge_engine_inheritance_matches_the_runner( data["task"] = task data["model"]["engine_kwargs"] = {"max_tokens": 128} data["judge"]["model"] = judge_model - monkeypatch.setattr(execution_module, "make_model", lambda **kwargs: kwargs) + monkeypatch.setattr(execution_module, "prepare_model", lambda **kwargs: kwargs) cfg = RunConfig(**data) if task == "mt-bench": diff --git a/tests/test_elo_calibration.py b/tests/test_elo_calibration.py index 3a5e02a0..6fe4bc90 100644 --- a/tests/test_elo_calibration.py +++ b/tests/test_elo_calibration.py @@ -29,7 +29,7 @@ def test_non_pairscore_calibration_does_not_consume_rng_or_build_judge(monkeypat def fail_if_called(**_kwargs): raise AssertionError("calibration judge was built") - monkeypatch.setattr(elo_calibration, "make_model", fail_if_called) + monkeypatch.setattr(elo_calibration, "prepare_model", fail_if_called) result = calibrate_pairscore_temperature( pd.DataFrame(), pd.DataFrame(), @@ -43,6 +43,7 @@ def fail_if_called(**_kwargs): prompt=SimpleNamespace(parser=object()), truncate_input_chars=None, default_temperature=0.3, + arena="test-arena", ) assert result is None diff --git a/tests/test_estimate_elo_ratings.py b/tests/test_estimate_elo_ratings.py index ecd4ebf4..cbb75519 100644 --- a/tests/test_estimate_elo_ratings.py +++ b/tests/test_estimate_elo_ratings.py @@ -16,6 +16,7 @@ ) from judgearena.benchmarks.elo.runner import run_elo from judgearena.benchmarks.elo.scoring import BradleyTerryMetric +from judgearena.cache_sqlite import JudgementCache from judgearena.config import RunConfig, load_config from judgearena.evaluate import JudgeAnnotation, judge_and_parse_prefs from judgearena.models import DummyModel, make_model @@ -87,13 +88,6 @@ def mock_generate(instructions, model, **kwargs): monkeypatch.setattr(estimate_elo_ratings, "generate_instructions", mock_generate) - def _run_without_cache(fun, **_kwargs): - return fun() - - monkeypatch.setattr( - estimate_elo_ratings, "cache_function_dataframe", _run_without_cache - ) - def _default_args(*, result_folder: str, **kwargs) -> RunConfig: task = kwargs.pop("task", "elo-comparia") @@ -106,6 +100,7 @@ def _default_args(*, result_folder: str, **kwargs) -> RunConfig: strip_thinking_before_judging = kwargs.pop("strip_thinking_before_judging", False) calibrate_temperature = kwargs.pop("calibrate_temperature", False) battle_thinking_token_budget = kwargs.pop("battle_thinking_token_budget", None) + store_root = kwargs.pop("store_root", None) assert not kwargs, f"unexpected kwargs: {kwargs}" judge: dict[str, object] = { "model": judge_model, @@ -124,7 +119,7 @@ def _default_args(*, result_folder: str, **kwargs) -> RunConfig: "languages": languages, "calibrate_temperature": calibrate_temperature, }, - run={"result_folder": result_folder}, + run={"result_folder": result_folder, "store_root": store_root}, ) @@ -670,8 +665,13 @@ def fake_calibrate(delta_s, y): estimate_elo_ratings, "load_battles", lambda _task: _arena_df(900) ) + store_root = tmp_path / "cache" result = run_elo_with_task( - _default_args(result_folder=str(tmp_path), calibrate_temperature=True) + _default_args( + result_folder=str(tmp_path), + calibrate_temperature=True, + store_root=str(store_root), + ) ) assert captured["n_pairs"] >= 10 @@ -680,6 +680,9 @@ def fake_calibrate(delta_s, y): assert battles["pref"].tolist() == pytest.approx( [1 / (1 + math.exp(-0.42 * 10))] * len(battles) ) + judgement_db = next((store_root / "judgements").rglob("judgements.db")) + with JudgementCache(judgement_db) as cache: + assert cache.query()["instruction_id"].str.startswith("ComparIA:").all() def test_extract_turn_text_tolerates_moderated_turns(): diff --git a/tests/test_generate_and_evaluate.py b/tests/test_generate_and_evaluate.py index 39f79faf..b790fe67 100644 --- a/tests/test_generate_and_evaluate.py +++ b/tests/test_generate_and_evaluate.py @@ -82,13 +82,6 @@ def mock_external_data_and_cache(monkeypatch): ), ) - def _run_without_cache(fun, **_kwargs): - return fun() - - monkeypatch.setattr( - generate_and_evaluate, "cache_function_dataframe", _run_without_cache - ) - def _mock_alpaca_judge(monkeypatch, messages) -> dict[str, object]: from judgearena.benchmarks.pairwise.scoring import alpaca_eval @@ -103,7 +96,7 @@ def make_fake_judge(**kwargs): captured.update(kwargs) return FakeJudge() - monkeypatch.setattr(benchmark_execution, "make_model", make_fake_judge) + monkeypatch.setattr(benchmark_execution, "prepare_model", make_fake_judge) monkeypatch.setattr( alpaca_eval, "_length_controlled_metrics", @@ -378,7 +371,7 @@ def batch(self, inputs, **_kwargs): return FakeJudge() - monkeypatch.setattr(benchmark_execution, "make_model", fake_make_model) + monkeypatch.setattr(benchmark_execution, "prepare_model", fake_make_model) prefs = run_pairwise( _cfg( @@ -439,7 +432,7 @@ def test_pairwise_bootstraps_use_run_seed(monkeypatch, tmp_path): ) monkeypatch.setattr( benchmark_execution, - "make_model", + "prepare_model", lambda **_kwargs: FakeListLLM( responses=["score A: 10 score B: 0", "score A: 0 score B: 10"] ), @@ -597,7 +590,7 @@ def test_run_pairwise_weights_and_preserves_incomplete_alpaca_annotations( def test_all_missing_alpaca_judgments_save_empty_results(monkeypatch, tmp_path, capsys): monkeypatch.setattr( benchmark_execution, - "make_model", + "prepare_model", lambda **_kwargs: FakeListLLM(responses=["M"]), ) diff --git a/tests/test_meta_eval_annotations.py b/tests/test_meta_eval_annotations.py index dc007940..b39625be 100644 --- a/tests/test_meta_eval_annotations.py +++ b/tests/test_meta_eval_annotations.py @@ -21,6 +21,8 @@ def _sample(): [ { "battle_id": "arena:q1", + "model_a": "alpha", + "model_b": "beta", "conversation_a": [ {"role": "user", "content": "Same prompt"}, {"role": "assistant", "content": "Alpha answer"}, diff --git a/tests/test_mt_bench_downloads.py b/tests/test_mt_bench_downloads.py index 86dbf967..c42c80da 100644 --- a/tests/test_mt_bench_downloads.py +++ b/tests/test_mt_bench_downloads.py @@ -161,10 +161,6 @@ def test_generate_mt_bench_completions_uses_pregenerated_baseline(monkeypatch): ) generated_models = [] - monkeypatch.setattr( - mt_bench_runner, "cache_function_dataframe", lambda fun, **_kwargs: fun() - ) - def fake_generate_multiturn(**kwargs): generated_models.append(kwargs["model"]) return pd.DataFrame( @@ -323,7 +319,7 @@ def fake_make_model(**kwargs): captured["make_model"] = kwargs return object() - monkeypatch.setattr(mt_bench_runner, "make_model", fake_make_model) + monkeypatch.setattr(mt_bench_runner, "prepare_model", fake_make_model) def fake_run_mt_bench_fastchat(**kwargs): captured["fastchat"] = kwargs @@ -394,7 +390,7 @@ def fake_make_model(**kwargs): captured["make_model"] = kwargs return object() - monkeypatch.setattr(mt_bench_runner, "make_model", fake_make_model) + monkeypatch.setattr(mt_bench_runner, "prepare_model", fake_make_model) def fake_run_mt_bench_fastchat(**kwargs): captured["fastchat"] = kwargs @@ -464,7 +460,7 @@ def fake_run_mt_bench_preset(**kwargs): return pd.Series([0.0], dtype=float) captured = {} - monkeypatch.setattr(mt_bench_runner, "make_model", fake_make_model) + monkeypatch.setattr(mt_bench_runner, "prepare_model", fake_make_model) monkeypatch.setattr( mt_bench_runner, "_run_mt_bench_preset", @@ -499,9 +495,6 @@ def test_generate_mt_bench_completions_forwards_thinking_controls(monkeypatch): ) captured: dict[str, dict] = {} - monkeypatch.setattr( - mt_bench_runner, "cache_function_dataframe", lambda fun, **_kwargs: fun() - ) monkeypatch.setattr( mt_bench_runner, "load_mt_bench_model_answers", @@ -578,7 +571,7 @@ def test_run_mt_bench_forwards_strip_thinking_to_fastchat_judge(monkeypatch, tmp ), ), ) - monkeypatch.setattr(mt_bench_runner, "make_model", lambda **kwargs: object()) + monkeypatch.setattr(mt_bench_runner, "prepare_model", lambda **kwargs: object()) monkeypatch.setattr( mt_bench_runner, "_finalize_mt_bench_run", lambda **kwargs: kwargs["prefs"] ) diff --git a/tests/test_seed_plumbing.py b/tests/test_seed_plumbing.py index 47522d05..6ab1a535 100644 --- a/tests/test_seed_plumbing.py +++ b/tests/test_seed_plumbing.py @@ -4,7 +4,6 @@ from judgearena.config import RunConfig, build_run_config from judgearena.models import make_model -from judgearena.utils import generation_cache_token def test_make_model_dummy_captures_temperature_and_seed(): @@ -93,17 +92,6 @@ def test_baseline_sampling_params_inherit_from_model_when_unset(): } -def test_generation_cache_token_is_sensitive_to_sampling_params(): - base = {"max_tokens": 32768, "temperature": 0.0, "seed": 1} - same = {"seed": 1, "temperature": 0.0, "max_tokens": 32768} # order independent - changed_seed = {**base, "seed": 2} - changed_temp = {**base, "temperature": 1.0} - - assert generation_cache_token(base) == generation_cache_token(same) - assert generation_cache_token(base) != generation_cache_token(changed_seed) - assert generation_cache_token(base) != generation_cache_token(changed_temp) - - def test_model_args_per_role_kwargs_are_independent(): cfg = RunConfig( task="alpaca-eval-ja", diff --git a/tests/test_usage.py b/tests/test_usage.py index 1304ee07..51e7a0a7 100644 --- a/tests/test_usage.py +++ b/tests/test_usage.py @@ -106,7 +106,13 @@ def fake_runner(cfg, task): def test_generation_preserves_sync_async_routing(monkeypatch): model = FakeModel(["batch"]) - monkeypatch.setattr(generate_module, "make_model", lambda *args, **kwargs: model) + prepared_kwargs = {} + + def fake_prepare_model(*args, **kwargs): + prepared_kwargs.update(kwargs) + return model + + monkeypatch.setattr(generate_module, "prepare_model", fake_prepare_model) instructions = pd.Series(["one"]) instruction_outputs = generate_module.generate_instructions( instructions, "Dummy/model", use_tqdm=True @@ -116,4 +122,4 @@ def test_generation_preserves_sync_async_routing(monkeypatch): instructions, "Dummy/model", max_tokens=123, use_tqdm=True ) assert base_outputs["completion"].tolist() == ["batch"] - assert model.batch_kwargs["max_tokens"] == 123 + assert prepared_kwargs["max_tokens"] == 123 From 4cd84829db1ee6900f0b69a3eb172404c3a46bce Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:00:41 +0200 Subject: [PATCH 5/8] fix: support nullable model_b and transactional in-place cache merge --- judgearena/cache_sqlite.py | 48 ++++++++++++++------------- tests/test_cache_sqlite.py | 66 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 91 insertions(+), 23 deletions(-) diff --git a/judgearena/cache_sqlite.py b/judgearena/cache_sqlite.py index 0325ff39..f9a53f09 100644 --- a/judgearena/cache_sqlite.py +++ b/judgearena/cache_sqlite.py @@ -4,7 +4,6 @@ import hashlib import json -import os import sqlite3 import uuid from datetime import UTC, datetime @@ -127,29 +126,32 @@ def _delete(self, conditions: list[str], params: list[Any]) -> int: return cursor.rowcount def merge_from(self, other_db: Path) -> int: - """Atomically merge another cache database using pushed_at last-write-wins.""" - self.close() - frames = [] - for db_path in (other_db, self.db_path): - if db_path.exists(): - with sqlite3.connect(db_path) as conn: - frames.append(pd.read_sql(f"SELECT * FROM {self.table}", conn)) - rows = ( - pd.concat(frames, ignore_index=True) - .sort_values("pushed_at", kind="stable") - .drop_duplicates("input_hash", keep="last") + """Merge another cache database in place using pushed_at last-write-wins.""" + conn = self._connect() + if not other_db.exists(): + return conn.execute(f"SELECT COUNT(*) FROM {self.table}").fetchone()[0] + + columns = [row[1] for row in conn.execute(f"PRAGMA table_info({self.table})")] + column_list = ", ".join(columns) + updates = ", ".join( + f"{column} = excluded.{column}" + for column in columns + if column != "input_hash" ) - - self.db_path.parent.mkdir(parents=True, exist_ok=True) - temporary_db = self.db_path.with_name(f".{self.db_path.name}.{uuid.uuid4()}") + conn.execute("ATTACH DATABASE ? AS incoming", (str(other_db),)) try: - with sqlite3.connect(temporary_db) as conn: - conn.execute(self.schema) - rows.to_sql(self.table, conn, if_exists="append", index=False) - os.replace(temporary_db, self.db_path) + with conn: + conn.execute( + f""" + INSERT INTO {self.table} ({column_list}) + SELECT {column_list} FROM incoming.{self.table} WHERE true + ON CONFLICT(input_hash) DO UPDATE SET {updates} + WHERE excluded.pushed_at > {self.table}.pushed_at + """ + ) finally: - temporary_db.unlink(missing_ok=True) - return len(rows) + conn.execute("DETACH DATABASE incoming") + return conn.execute(f"SELECT COUNT(*) FROM {self.table}").fetchone()[0] def close(self) -> None: if self._conn is not None: @@ -257,7 +259,7 @@ class JudgementCache(_SQLiteCache): benchmark TEXT NOT NULL, instruction_id TEXT NOT NULL, model_a TEXT NOT NULL, - model_b TEXT NOT NULL, + model_b TEXT, judge TEXT NOT NULL, top_logprobs TEXT, -- direct/reversed relative to the source model order, when applicable @@ -285,7 +287,7 @@ def save( str(row["benchmark"]), str(row["instruction_id"]), str(row["model_a"]), - str(row["model_b"]), + None if pd.isna(row.get("model_b")) else str(row["model_b"]), str(row["judge"]), ( stable_json_dumps(row["top_logprobs"]) diff --git a/tests/test_cache_sqlite.py b/tests/test_cache_sqlite.py index 28dc693e..6f7018fe 100644 --- a/tests/test_cache_sqlite.py +++ b/tests/test_cache_sqlite.py @@ -62,6 +62,8 @@ def test_completion_cache_uses_content_key_and_last_write(tmp_path): cache.save(first, pushed_by="alice") cache.save(second, pushed_by="bob") result = cache.query([input_hash("rendered prompt")]) + assert cache.query([]).empty + assert len(cache.query(None)) == 1 assert result["completion"].tolist() == ["second"] assert result["pushed_by"].tolist() == ["bob"] @@ -115,3 +117,67 @@ def test_judgement_cache_filters_and_deletes_by_candidate_model(tmp_path): assert json.loads(result.iloc[0]["top_logprobs"]) == {"M": -2.0, "m": -0.1} assert cache.delete(model="candidate") == 1 assert cache.query()["instruction_id"].tolist() == ["1"] + + +def test_judgement_cache_preserves_null_model_b(tmp_path): + row = pd.DataFrame( + [ + { + "judge_input": "pointwise prompt", + "judge_completion": "Rating: [[8]]", + "benchmark": "mt-bench-101", + "instruction_id": "1", + "model_a": "candidate", + "model_b": None, + "judge": "VLLM/Qwen/Qwen3-8B", + "top_logprobs": None, + "orientation": "single", + } + ] + ) + + with JudgementCache(tmp_path / JUDGEMENT_DB_NAME) as cache: + cache.save(row, pushed_by="alice") + assert cache.query().iloc[0]["model_b"] is None + + +def test_merge_from_updates_live_database_in_place(tmp_path): + local_path = tmp_path / "local" / COMPLETION_DB_NAME + incoming_path = tmp_path / "incoming" / COMPLETION_DB_NAME + shared = pd.DataFrame( + [ + { + "input_text": "shared", + "completion": "local", + "benchmark": "arena-hard", + "instruction_id": "1", + "model": "VLLM/Qwen/Qwen3-8B", + } + ] + ) + with CompletionCache(incoming_path) as incoming: + incoming.save( + pd.concat( + [ + shared.assign(completion="incoming"), + shared.assign( + input_text="new", + completion="new", + instruction_id="2", + ), + ], + ignore_index=True, + ), + pushed_by="bob", + ) + incoming._connect().execute( + "UPDATE completions SET pushed_at = '2030-01-01T00:00:00+00:00'" + ) + incoming._connect().commit() + + with CompletionCache(local_path) as cache: + cache.save(shared, pushed_by="alice") + inode = local_path.stat().st_ino + assert cache.merge_from(incoming_path) == 2 + assert local_path.stat().st_ino == inode + assert cache.query()["completion"].tolist() == ["incoming", "new"] From 72987de7a5fef49b0ae36d467d2c6ac563196d57 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:00:46 +0200 Subject: [PATCH 6/8] refactor: consolidate InferenceResult and harden do_inference caching --- judgearena/inference.py | 48 ++++++++++++------ judgearena/models.py | 94 ++++++++++++++++++++++++----------- tests/test_inference_cache.py | 51 ++++++++++++++++--- 3 files changed, 143 insertions(+), 50 deletions(-) diff --git a/judgearena/inference.py b/judgearena/inference.py index 2a2e24d1..463b995e 100644 --- a/judgearena/inference.py +++ b/judgearena/inference.py @@ -7,7 +7,7 @@ from collections.abc import Callable from dataclasses import dataclass, field from pathlib import Path -from typing import Any, ClassVar +from typing import Any, ClassVar, NotRequired, TypedDict import pandas as pd @@ -21,6 +21,7 @@ stable_json_dumps, write_descriptor, ) +from judgearena.usage import RequestUsage _ROLE_MAP = {"human": "user", "ai": "assistant", "system": "system"} @@ -63,11 +64,26 @@ def build_model_descriptor( @dataclass(frozen=True) -class CachedInferenceResult: - """Provider output fields required by downstream parsing.""" +class InferenceResult: + """A text completion and optional provider response details.""" text: str first_token_top_logprobs: dict[str, float] | None = None + usage: RequestUsage | None = None + + +class CompletionCacheRowMetadata(TypedDict): + instruction_id: str + + +class JudgementCacheRowMetadata(TypedDict): + instruction_id: str + model_a: str + model_b: str | None + orientation: NotRequired[str | None] + + +CacheRowMetadata = CompletionCacheRowMetadata | JudgementCacheRowMetadata @dataclass @@ -77,7 +93,7 @@ class PreparedModel: model_spec: str descriptor: dict[str, Any] | None factory: Callable[[], Any] - cache: InferenceCache | None = None + cache: InferenceCache[Any] | None = None _model: Any = field(default=None, init=False, repr=False) def materialize(self) -> Any: @@ -87,7 +103,7 @@ def materialize(self) -> Any: @dataclass(frozen=True) -class InferenceCache(ABC): +class InferenceCache[CacheRowMetadataT: CacheRowMetadata](ABC): """Share cache lifecycle while subclasses define role-specific rows.""" store_root: Path @@ -116,7 +132,7 @@ def save_outputs( model: PreparedModel, input_texts: list[str], outputs: list[Any], - metadata: list[dict[str, Any]], + metadata: list[CacheRowMetadataT], indices: list[int], ) -> None: rows = [ @@ -137,16 +153,16 @@ def make_row( model: PreparedModel, input_text: str, output: Any, - metadata: dict[str, Any], + metadata: CacheRowMetadataT, ) -> dict[str, Any]: """Convert one inference output to its role-specific storage row.""" @abstractmethod - def cached_result(self, row: pd.Series) -> CachedInferenceResult: + def cached_result(self, row: pd.Series) -> InferenceResult: """Restore output fields from a stored row.""" -class CompletionInferenceCache(InferenceCache): +class CompletionInferenceCache(InferenceCache[CompletionCacheRowMetadata]): """Cache generated model completions.""" kind = "completions" @@ -159,7 +175,7 @@ def make_row( model: PreparedModel, input_text: str, output: Any, - metadata: dict[str, Any], + metadata: CompletionCacheRowMetadata, ) -> dict[str, Any]: return { "input_text": input_text, @@ -169,11 +185,11 @@ def make_row( "model": model.model_spec, } - def cached_result(self, row: pd.Series) -> CachedInferenceResult: - return CachedInferenceResult(text=str(row["completion"])) + def cached_result(self, row: pd.Series) -> InferenceResult: + return InferenceResult(text=str(row["completion"])) -class JudgementInferenceCache(InferenceCache): +class JudgementInferenceCache(InferenceCache[JudgementCacheRowMetadata]): """Cache raw judge completions.""" kind = "judgements" @@ -186,7 +202,7 @@ def make_row( model: PreparedModel, input_text: str, output: Any, - metadata: dict[str, Any], + metadata: JudgementCacheRowMetadata, ) -> dict[str, Any]: return { "judge_input": input_text, @@ -200,9 +216,9 @@ def make_row( "orientation": metadata.get("orientation"), } - def cached_result(self, row: pd.Series) -> CachedInferenceResult: + def cached_result(self, row: pd.Series) -> InferenceResult: top_logprobs = row["top_logprobs"] - return CachedInferenceResult( + return InferenceResult( text=str(row["judge_completion"]), first_token_top_logprobs=( json.loads(top_logprobs) if pd.notna(top_logprobs) else None diff --git a/judgearena/models.py b/judgearena/models.py index 8e3adc57..79b64de6 100644 --- a/judgearena/models.py +++ b/judgearena/models.py @@ -6,20 +6,24 @@ import json import math import os +import sqlite3 import time import warnings from collections.abc import Mapping -from dataclasses import dataclass, replace +from dataclasses import replace from langchain_community.llms import LlamaCpp from langchain_openai import ChatOpenAI +from pandas.errors import DatabaseError from tqdm.asyncio import tqdm from tqdm.contrib.logging import logging_redirect_tqdm from judgearena.cache_sqlite import input_hash from judgearena.constants import VLLM_REASONING_END_STR, VLLM_REASONING_START_STR from judgearena.inference import ( + CacheRowMetadata, InferenceCache, + InferenceResult, PreparedModel, build_model_descriptor, canonicalize_chat_input, @@ -30,6 +34,7 @@ logger = get_logger(__name__) +_CACHE_OPERATION_ERRORS = (OSError, sqlite3.Error, DatabaseError) DEFAULT_VLLM_JUDGE_THINKING_TOKEN_BUDGET = 512 _THINKING_MODEL_PARSER_BY_SUBSTRING = ( @@ -482,15 +487,6 @@ async def ainvoke(self, input_item, **invoke_kwargs): ) -@dataclass(frozen=True) -class InferenceResult: - """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: """Extract first-token top logprobs from a langchain AIMessage, if any.""" metadata = getattr(response, "response_metadata", None) or {} @@ -670,7 +666,7 @@ def batch_inference_once( return [result.text for result in results] -def _do_inference_uncached( +def _run_backend_inference( chat_model, inputs, use_tqdm: bool = False, @@ -786,12 +782,12 @@ def do_inference( return_top_logprobs: bool = False, *, stage: str = "unspecified", - cache_metadata: list[dict] | None = None, + cache_row_metadata: list[CacheRowMetadata] | None = None, ): """Reuse raw provider outputs and invoke the backend only for cache misses.""" inputs = list(inputs) if not isinstance(chat_model, PreparedModel): - return _do_inference_uncached( + return _run_backend_inference( chat_model, inputs, use_tqdm, @@ -801,23 +797,53 @@ def do_inference( cache = chat_model.cache if cache is None or chat_model.descriptor is None: - return _do_inference_uncached( + return _run_backend_inference( chat_model.materialize(), inputs, use_tqdm, return_top_logprobs, stage=stage, ) - if cache_metadata is None or len(cache_metadata) != len(inputs): - raise ValueError("cache_metadata must contain one row per inference input.") + if cache_row_metadata is None or len(cache_row_metadata) != len(inputs): + raise ValueError("cache_row_metadata must contain one row per inference input.") input_texts = [canonicalize_chat_input(item) for item in inputs] input_hashes = [input_hash(input_text) for input_text in input_texts] - with cache.open_store(chat_model) as store: - cached_rows = store.query(input_hashes).set_index("input_hash") + try: + store = cache.open_store(chat_model) + except _CACHE_OPERATION_ERRORS as exc: + logger.warning( + "Cache open failed at %s: %s. Continuing without caching.", + cache.store_root, + exc, + ) + return _run_backend_inference( + chat_model.materialize(), + inputs, + use_tqdm, + return_top_logprobs, + stage=stage, + ) + + try: + try: + cached_rows = store.query(input_hashes).set_index("input_hash") + except _CACHE_OPERATION_ERRORS as exc: + logger.warning( + "Cache read failed at %s: %s. Continuing without caching.", + store.db_path, + exc, + ) + return _run_backend_inference( + chat_model.materialize(), + inputs, + use_tqdm, + return_top_logprobs, + stage=stage, + ) results: list[InferenceResult | None] = [ ( - InferenceResult(**cache.cached_result(cached_rows.loc[key]).__dict__) + cache.cached_result(cached_rows.loc[key]) if key in cached_rows.index else None ) @@ -827,7 +853,7 @@ def do_inference( index for index, result in enumerate(results) if result is None ] if missing_indices: - generated = _do_inference_uncached( + generated = _run_backend_inference( chat_model.materialize(), [inputs[index] for index in missing_indices], use_tqdm, @@ -836,14 +862,26 @@ def do_inference( ) for index, result in zip(missing_indices, generated, strict=True): results[index] = result - cache.save_outputs( - store, - chat_model, - input_texts, - generated, - cache_metadata, - missing_indices, - ) + try: + cache.save_outputs( + store, + chat_model, + input_texts, + generated, + cache_row_metadata, + missing_indices, + ) + except _CACHE_OPERATION_ERRORS as exc: + logger.warning( + "Cache write failed at %s: %s. Preserving generated results.", + store.db_path, + exc, + ) + finally: + try: + store.close() + except _CACHE_OPERATION_ERRORS as exc: + logger.warning("Cache close failed at %s: %s.", store.db_path, exc) resolved_results = [result for result in results if result is not None] if return_top_logprobs: diff --git a/tests/test_inference_cache.py b/tests/test_inference_cache.py index 3343fb97..686e2628 100644 --- a/tests/test_inference_cache.py +++ b/tests/test_inference_cache.py @@ -1,3 +1,6 @@ +import sqlite3 + +import pytest from langchain_core.messages import AIMessage import judgearena.models as models @@ -22,7 +25,7 @@ def test_full_hit_does_not_materialize_model(tmp_path, monkeypatch): do_inference( prepare_model("Dummy/test-model", cache=cache), ["prompt"], - cache_metadata=metadata, + cache_row_metadata=metadata, ) def fail_if_materialized(*_args, **_kwargs): @@ -33,7 +36,7 @@ def fail_if_materialized(*_args, **_kwargs): outputs = do_inference( prepare_model("Dummy/test-model", cache=cache), ["prompt"], - cache_metadata=metadata, + cache_row_metadata=metadata, ) assert outputs == ["generated:prompt"] @@ -47,7 +50,7 @@ def test_mixed_hits_and_misses_preserve_order(tmp_path, monkeypatch): do_inference( prepare_model("Dummy/test-model", cache=cache), ["hit"], - cache_metadata=[{"instruction_id": "hit"}], + cache_row_metadata=[{"instruction_id": "hit"}], ) backend = EchoModel() @@ -55,7 +58,7 @@ def test_mixed_hits_and_misses_preserve_order(tmp_path, monkeypatch): outputs = do_inference( prepare_model("Dummy/test-model", cache=cache), ["miss-a", "hit", "miss-b"], - cache_metadata=[ + cache_row_metadata=[ {"instruction_id": "a"}, {"instruction_id": "hit"}, {"instruction_id": "b"}, @@ -92,15 +95,51 @@ def batch(self, inputs, **_kwargs): prepare_model("Dummy/judge", cache=cache), ["judge prompt"], return_top_logprobs=True, - cache_metadata=metadata, + cache_row_metadata=metadata, ) second = do_inference( prepare_model("Dummy/judge", cache=cache), ["judge prompt"], return_top_logprobs=True, - cache_metadata=metadata, + cache_row_metadata=metadata, ) assert second[0].text == first[0].text assert second[0].first_token_top_logprobs == first[0].first_token_top_logprobs assert second[0].usage is None + + +def test_cache_write_failure_preserves_generated_results(tmp_path, monkeypatch): + cache = CompletionInferenceCache(tmp_path, "arena-hard") + backend = EchoModel() + monkeypatch.setattr(models, "make_model", lambda *_args, **_kwargs: backend) + + def fail_save(*_args, **_kwargs): + raise sqlite3.OperationalError("read-only") + + monkeypatch.setattr(CompletionInferenceCache, "save_outputs", fail_save) + + outputs = do_inference( + prepare_model("Dummy/test-model", cache=cache), + ["prompt"], + cache_row_metadata=[{"instruction_id": "1"}], + ) + + assert outputs == ["generated:prompt"] + assert backend.calls == [["prompt"]] + + +def test_cache_descriptor_validation_still_fails_loudly(tmp_path, monkeypatch): + cache = CompletionInferenceCache(tmp_path, "arena-hard") + + def fail_open(*_args, **_kwargs): + raise ValueError("descriptor mismatch") + + monkeypatch.setattr(CompletionInferenceCache, "open_store", fail_open) + + with pytest.raises(ValueError, match="descriptor mismatch"): + do_inference( + prepare_model("Dummy/test-model", cache=cache), + ["prompt"], + cache_row_metadata=[{"instruction_id": "1"}], + ) From 22e9b732647837a42179c4b3447741275e4ff644 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:00:52 +0200 Subject: [PATCH 7/8] refactor: share VLLM sampling defaults and unify provider input modes --- judgearena/constants.py | 3 +++ judgearena/inference.py | 22 +++++++++++++--------- judgearena/models.py | 14 +++++++++++--- tests/test_chat_vllm.py | 2 ++ 4 files changed, 29 insertions(+), 12 deletions(-) diff --git a/judgearena/constants.py b/judgearena/constants.py index 2c1e9bd4..f4abf95d 100644 --- a/judgearena/constants.py +++ b/judgearena/constants.py @@ -8,3 +8,6 @@ VLLM_REASONING_END_STR = ( "I have to give the solution based on the thinking directly now." ) + +VLLM_DEFAULT_TEMPERATURE = 0.6 +VLLM_DEFAULT_TOP_P = 0.95 diff --git a/judgearena/inference.py b/judgearena/inference.py index 0005e285..41779630 100644 --- a/judgearena/inference.py +++ b/judgearena/inference.py @@ -22,10 +22,9 @@ stable_json_dumps, write_descriptor, ) +from judgearena.constants import VLLM_DEFAULT_TEMPERATURE, VLLM_DEFAULT_TOP_P _ROLE_MAP = {"human": "user", "ai": "assistant", "system": "system"} -_CHAT_PROVIDERS = {"ChatOpenAI", "Dummy", "OpenRouter", "VLLM"} -_TEXT_PROVIDERS = {"LlamaCpp", "OpenAI", "Together"} VLLM_EXECUTION_ONLY_KWARGS = { "enforce_eager", "gpu_memory_utilization", @@ -33,14 +32,19 @@ } InputMode = Literal["chat", "text"] +_PROVIDER_INPUT_MODES: dict[str, InputMode] = { + "ChatOpenAI": "chat", + "Dummy": "chat", + "LlamaCpp": "text", + "OpenAI": "text", + "OpenRouter": "chat", + "Together": "text", + "VLLM": "chat", +} def provider_input_mode(provider: str) -> InputMode | None: - if provider in _CHAT_PROVIDERS: - return "chat" - if provider in _TEXT_PROVIDERS: - return "text" - return None + return _PROVIDER_INPUT_MODES.get(provider) def _canonical_messages(input_item: Any) -> list[dict[str, Any]]: @@ -91,12 +95,12 @@ def build_model_descriptor( descriptor_kwargs = resolved_kwargs.copy() if provider == "VLLM": descriptor_kwargs["temperature"] = ( - 0.6 + VLLM_DEFAULT_TEMPERATURE if descriptor_kwargs.get("temperature") is None else descriptor_kwargs["temperature"] ) descriptor_kwargs["top_p"] = ( - 0.95 + VLLM_DEFAULT_TOP_P if descriptor_kwargs.get("top_p") is None else descriptor_kwargs["top_p"] ) diff --git a/judgearena/models.py b/judgearena/models.py index f370c003..278f5f63 100644 --- a/judgearena/models.py +++ b/judgearena/models.py @@ -17,7 +17,12 @@ from tqdm.contrib.logging import logging_redirect_tqdm from judgearena.cache_sqlite import input_hash -from judgearena.constants import VLLM_REASONING_END_STR, VLLM_REASONING_START_STR +from judgearena.constants import ( + VLLM_DEFAULT_TEMPERATURE, + VLLM_DEFAULT_TOP_P, + VLLM_REASONING_END_STR, + VLLM_REASONING_START_STR, +) from judgearena.inference import ( InferenceCache, PreparedModel, @@ -261,8 +266,10 @@ def __init__( self._sampling_params_kwargs = { "max_tokens": max_tokens, - "temperature": 0.6 if temperature is None else float(temperature), - "top_p": 0.95 if top_p is None else float(top_p), + "temperature": ( + VLLM_DEFAULT_TEMPERATURE if temperature is None else float(temperature) + ), + "top_p": VLLM_DEFAULT_TOP_P if top_p is None else float(top_p), } if top_k is not None: self._sampling_params_kwargs["top_k"] = int(top_k) @@ -894,6 +901,7 @@ def _route_sampling_params( return engine_kwargs +# Prevent local-engine settings from leaking into hosted provider constructors. _VLLM_ONLY_KWARGS = ( "max_model_len", "chat_template", diff --git a/tests/test_chat_vllm.py b/tests/test_chat_vllm.py index ddaad5fa..982e04bb 100644 --- a/tests/test_chat_vllm.py +++ b/tests/test_chat_vllm.py @@ -64,6 +64,8 @@ def test_chat_vllm_enables_reasoning_support_for_qwen_thinking_budget(monkeypatc gpu_memory_utilization=0.7, ) + assert captured["sampling_kwargs"]["temperature"] == models.VLLM_DEFAULT_TEMPERATURE + assert captured["sampling_kwargs"]["top_p"] == models.VLLM_DEFAULT_TOP_P assert captured["sampling_kwargs"]["thinking_token_budget"] == 64 assert "structured_outputs" not in captured["sampling_kwargs"] assert captured["reasoning_config_kwargs"] == { From 5c40e4d46bed2fb5fa552dc7e45d36e2cd5b9333 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:00:55 +0200 Subject: [PATCH 8/8] feat: default store_root to XDG cache with null opt-out --- judgearena/benchmarks/elo/calibration.py | 2 +- judgearena/benchmarks/elo/runner.py | 2 +- judgearena/benchmarks/meta_eval/annotate.py | 2 +- .../benchmarks/mt_bench/pairwise_judging.py | 2 +- judgearena/benchmarks/pairwise/runner.py | 2 +- judgearena/config.py | 10 +++- judgearena/evaluate.py | 22 +++++---- judgearena/generate.py | 18 ++++++-- tests/test_config.py | 46 ++++++++++++++++++- 9 files changed, 85 insertions(+), 21 deletions(-) diff --git a/judgearena/benchmarks/elo/calibration.py b/judgearena/benchmarks/elo/calibration.py index 10575f35..17293a63 100644 --- a/judgearena/benchmarks/elo/calibration.py +++ b/judgearena/benchmarks/elo/calibration.py @@ -123,7 +123,7 @@ def calibrate_pairscore_temperature( prompt_preset=prompt.preset_name, parse=prompt.parser, truncate_input_chars=truncate_input_chars, - cache_metadata=[ + cache_row_metadata=[ { "instruction_id": f"{arena}:{row.question_id}", "model_a": row.model_a, diff --git a/judgearena/benchmarks/elo/runner.py b/judgearena/benchmarks/elo/runner.py index 7815784c..f1566da5 100644 --- a/judgearena/benchmarks/elo/runner.py +++ b/judgearena/benchmarks/elo/runner.py @@ -198,7 +198,7 @@ def run_judge() -> pd.DataFrame: parse=resolved_prompt.parser, truncate_input_chars=cfg.generation.truncate_judge_input_chars, use_tqdm=use_tqdm, - cache_metadata=[ + cache_row_metadata=[ { "instruction_id": str(instructions.index[index]), "model_a": ( diff --git a/judgearena/benchmarks/meta_eval/annotate.py b/judgearena/benchmarks/meta_eval/annotate.py index 843ed6ff..ca5a067a 100644 --- a/judgearena/benchmarks/meta_eval/annotate.py +++ b/judgearena/benchmarks/meta_eval/annotate.py @@ -161,7 +161,7 @@ def annotate_sample( parse=parser, truncate_input_chars=cfg.generation.truncate_judge_input_chars, use_tqdm=cfg.run.use_tqdm, - cache_metadata=[ + cache_row_metadata=[ { "instruction_id": battle["battle_id"], "model_a": battle["model_a"], diff --git a/judgearena/benchmarks/mt_bench/pairwise_judging.py b/judgearena/benchmarks/mt_bench/pairwise_judging.py index f07669f3..cf519786 100644 --- a/judgearena/benchmarks/mt_bench/pairwise_judging.py +++ b/judgearena/benchmarks/mt_bench/pairwise_judging.py @@ -99,7 +99,7 @@ def infer_pairwise_judgments_by_prompt_groups( inputs=prompt_inputs, use_tqdm=use_tqdm, stage="judging", - cache_metadata=[ + cache_row_metadata=[ { "instruction_id": ( f"{items[item_index].question_id}:turn-{items[item_index].turn}" diff --git a/judgearena/benchmarks/pairwise/runner.py b/judgearena/benchmarks/pairwise/runner.py index 63980e88..681271bb 100644 --- a/judgearena/benchmarks/pairwise/runner.py +++ b/judgearena/benchmarks/pairwise/runner.py @@ -312,7 +312,7 @@ def _load_or_generate_completions( parse=group_prompt.parser, truncate_input_chars=cfg.generation.truncate_judge_input_chars, use_tqdm=cfg.run.use_tqdm, - cache_metadata=[ + cache_row_metadata=[ { "instruction_id": str(index), "model_a": judged_model_a.loc[index], diff --git a/judgearena/config.py b/judgearena/config.py index 8e2280e5..5e3165c9 100644 --- a/judgearena/config.py +++ b/judgearena/config.py @@ -3,6 +3,7 @@ from __future__ import annotations import argparse +import os from pathlib import Path from typing import Literal @@ -32,6 +33,11 @@ _ACTIVE_CLI_ARGS: list[str] | None = None +def _default_store_root() -> str: + cache_home = Path(os.getenv("XDG_CACHE_HOME") or Path.home() / ".cache") + return str(cache_home.expanduser() / "judgearena") + + def _drop_none(kwargs: dict[str, object]) -> dict[str, object]: return {k: v for k, v in kwargs.items() if v is not None} @@ -410,8 +416,8 @@ class RunArgs(BaseModel): """Directory where annotations, results, and the resolved ``config.yaml`` are written (under a per-run subfolder).""" - store_root: str | None = None - """Root directory for content-addressed completion and judgement caches.""" + store_root: str | None = Field(default_factory=_default_store_root) + """Root directory for content-addressed caches. Set to null to disable.""" use_tqdm: bool = False """Show a tqdm progress bar (not compatible with vLLM).""" diff --git a/judgearena/evaluate.py b/judgearena/evaluate.py index b5b4c6cd..3c3ad2e2 100644 --- a/judgearena/evaluate.py +++ b/judgearena/evaluate.py @@ -1,5 +1,8 @@ +from __future__ import annotations + import json from dataclasses import dataclass +from typing import TYPE_CHECKING import pandas as pd from langchain_core.prompts import ChatPromptTemplate @@ -16,6 +19,9 @@ ) from judgearena.utils import strip_thinking_tags, truncate +if TYPE_CHECKING: + from judgearena.inference import JudgementCacheRowMetadata + logger = get_logger(__name__) # Re-exported for existing callers; implementations live with the presets. @@ -111,7 +117,7 @@ def annotate_battles( prompt_preset: str = DEFAULT_JUDGE_PROMPT_PRESET, strip_thinking_before_judging: bool = False, collect_top_logprobs: bool = False, - cache_metadata: list[dict] | None = None, + cache_row_metadata: list[JudgementCacheRowMetadata] | None = None, ) -> list[JudgeAnnotation]: """ Directly evaluate from list of instructions and completions @@ -183,7 +189,7 @@ def annotate_battles( use_tqdm=use_tqdm, return_top_logprobs=collect_top_logprobs, stage="judging", - cache_metadata=cache_metadata, + cache_row_metadata=cache_row_metadata, ) if not collect_top_logprobs: judge_results = [InferenceResult(text=text) for text in judge_results] @@ -231,7 +237,7 @@ def judge_and_parse_prefs( truncate_input_chars: int = 8192, use_tqdm: bool = False, parse: JudgeParser | None = None, - cache_metadata: list[dict] | None = None, + cache_row_metadata: list[JudgementCacheRowMetadata] | None = None, ) -> tuple[list[JudgeAnnotation], list[JudgeAnnotation] | None, pd.Series]: """Run judge annotation and parse preferences, handling swap_mode='both'. @@ -269,12 +275,12 @@ def judge_and_parse_prefs( truncate_input_chars=truncate_input_chars, use_tqdm=use_tqdm, collect_top_logprobs=parse.requires_top_logprobs, - cache_metadata=cache_metadata, + cache_row_metadata=cache_row_metadata, ) annotations_reversed = None if swap_mode == "both": - reversed_cache_metadata = ( + reversed_cache_row_metadata = ( [ { **metadata, @@ -286,9 +292,9 @@ def judge_and_parse_prefs( else "direct" ), } - for metadata in cache_metadata + for metadata in cache_row_metadata ] - if cache_metadata is not None + if cache_row_metadata is not None else None ) annotations_reversed = annotate_battles( @@ -303,7 +309,7 @@ def judge_and_parse_prefs( truncate_input_chars=truncate_input_chars, use_tqdm=use_tqdm, collect_top_logprobs=parse.requires_top_logprobs, - cache_metadata=reversed_cache_metadata, + cache_row_metadata=reversed_cache_row_metadata, ) def _none_to_nan(x): diff --git a/judgearena/generate.py b/judgearena/generate.py index 4ca968f5..e528f92d 100644 --- a/judgearena/generate.py +++ b/judgearena/generate.py @@ -46,7 +46,9 @@ def generate_instructions( inputs=inputs, use_tqdm=use_tqdm, stage="generation", - cache_metadata=[{"instruction_id": str(index)} for index in instructions.index], + cache_row_metadata=[ + {"instruction_id": str(index)} for index in instructions.index + ], ) df_outputs = pd.DataFrame( data={ @@ -109,7 +111,7 @@ def _infer_grouped_by_temperature( inputs=group_inputs, use_tqdm=use_tqdm, stage="generation", - cache_metadata=[ + cache_row_metadata=[ {"instruction_id": instruction_ids[index]} for index in idxs ], ) @@ -189,7 +191,9 @@ def generate_multiturn( inputs=turn1_inputs, use_tqdm=use_tqdm, stage="generation", - cache_metadata=[{"instruction_id": f"{index}:turn-1"} for index in idxs], + cache_row_metadata=[ + {"instruction_id": f"{index}:turn-1"} for index in idxs + ], ) turn2_inputs = [] @@ -246,7 +250,9 @@ def generate_multiturn( inputs=turn2_inputs, use_tqdm=use_tqdm, stage="generation", - cache_metadata=[{"instruction_id": f"{index}:turn-2"} for index in idxs], + cache_row_metadata=[ + {"instruction_id": f"{index}:turn-2"} for index in idxs + ], ) return pd.DataFrame( @@ -283,7 +289,9 @@ def generate_base( chat_model=chat_model, inputs=inputs, stage="generation", - cache_metadata=[{"instruction_id": str(index)} for index in instructions.index], + cache_row_metadata=[ + {"instruction_id": str(index)} for index in instructions.index + ], ) df_outputs = pd.DataFrame( diff --git a/tests/test_config.py b/tests/test_config.py index ff0d2ffe..6637b318 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,3 +1,4 @@ +from pathlib import Path from types import SimpleNamespace import pytest @@ -6,7 +7,13 @@ import judgearena.config as config_module from judgearena import cli as cli_module from judgearena.benchmarks import execution as execution_module -from judgearena.config import EloArgs, RunConfig, dump_config, load_config +from judgearena.config import ( + EloArgs, + RunConfig, + build_run_config, + dump_config, + load_config, +) from judgearena.tasks.schema import EloScoringSpec, MetricSpec @@ -34,6 +41,43 @@ def test_generate_config_constructs(): assert cfg.elo is None +def test_store_root_defaults_to_xdg_cache_home(monkeypatch, tmp_path): + monkeypatch.setenv("XDG_CACHE_HOME", str(tmp_path)) + + cfg = RunConfig(**_base_generate()) + + assert cfg.run.store_root == str(tmp_path / "judgearena") + + +def test_store_root_defaults_to_user_cache(monkeypatch): + monkeypatch.delenv("XDG_CACHE_HOME", raising=False) + + cfg = RunConfig(**_base_generate()) + + assert cfg.run.store_root == str(Path.home() / ".cache" / "judgearena") + + +def test_store_root_null_disables_cache_from_cli(): + cfg = build_run_config( + [ + "--task", + "alpaca-eval", + "--model.name", + "Dummy/a", + "--model.baseline", + "Dummy/b", + "--judge.model", + "Dummy/j", + "--run.store_root", + "null", + ] + ) + + assert cfg.run.store_root is None + assert execution_module.build_completion_cache(cfg) is None + assert execution_module.build_judgement_cache(cfg) is None + + def test_load_config_ignores_unused_legacy_judge_fields(tmp_path): yaml_path = tmp_path / "run.yaml" yaml_path.write_text(