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/6] 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/6] 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/6] 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 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 4/6] 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 5/6] 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 6/6] 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"] == {