From b2e1ed5e1cea5848d82def3a567e385d01ee907e Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Sun, 19 Jul 2026 00:17:56 +0200 Subject: [PATCH 01/13] refactor: expose reusable arena turn text extraction --- judgearena/arenas_utils.py | 14 ++++++++++---- judgearena/estimate_elo_ratings.py | 8 ++++---- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/judgearena/arenas_utils.py b/judgearena/arenas_utils.py index e85750f..9608c62 100644 --- a/judgearena/arenas_utils.py +++ b/judgearena/arenas_utils.py @@ -10,8 +10,8 @@ logger = get_logger(__name__) -def _extract_instruction_text(turn: dict) -> str: - """Extract plain instruction text from a conversation first turn. +def extract_turn_text(turn: dict) -> str: + """Extract plain text from an arena conversation turn. Handles both the 100k schema (content is a plain string) and the 140k schema (content is an array of {type, text, ...} objects). @@ -19,7 +19,13 @@ def _extract_instruction_text(turn: dict) -> str: content = turn["content"] if isinstance(content, str): return content - return " ".join(block["text"] for block in content if block.get("type") == "text") + if isinstance(content, (list, tuple)): + return " ".join( + str(block.get("text", "")) + for block in content + if isinstance(block, dict) and block.get("type") == "text" + ) + return str(content) KNOWN_ARENAS = ["LMArena-100k", "LMArena-55k", "LMArena-140k", "ComparIA"] @@ -139,7 +145,7 @@ def get_winner( df["question_id"] = df["id"] df["lang"] = df["conversation_a"].apply( - lambda conv: detect_language(_extract_instruction_text(conv[0])).lower() + lambda conv: detect_language(extract_turn_text(conv[0])).lower() ) cols = [ diff --git a/judgearena/estimate_elo_ratings.py b/judgearena/estimate_elo_ratings.py index 51ba6e2..7c10c55 100644 --- a/judgearena/estimate_elo_ratings.py +++ b/judgearena/estimate_elo_ratings.py @@ -6,7 +6,7 @@ import pandas as pd from sklearn.linear_model import LogisticRegression -from judgearena.arenas_utils import _extract_instruction_text, load_arena_dataframe +from judgearena.arenas_utils import extract_turn_text, load_arena_dataframe from judgearena.cli_common import BaseCliArgs from judgearena.evaluate import judge_and_parse_prefs from judgearena.generate import generate_instructions @@ -178,7 +178,7 @@ def main(args: CliEloArgs) -> dict: # Extract user instructions (first turn of conversation_a) instructions = pd.Series( [ - _extract_instruction_text(row["conversation_a"][0]) + extract_turn_text(row["conversation_a"][0]) for _, row in df_battles.iterrows() ], name="instruction", @@ -246,9 +246,9 @@ def replace_slash(s: str) -> str: opponent_completions = [ ( - _extract_instruction_text(row["conversation_a"][1]) + extract_turn_text(row["conversation_a"][1]) if use_model_a_as_opponent[i] - else _extract_instruction_text(row["conversation_b"][1]) + else extract_turn_text(row["conversation_b"][1]) ) for i, (_, row) in enumerate(df_battles.iterrows()) ] From 34884899ac329bb662d2105beef8f771ffaef8d8 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Sun, 19 Jul 2026 00:19:31 +0200 Subject: [PATCH 02/13] fix: break utility import initialization cycle --- judgearena/utils.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/judgearena/utils.py b/judgearena/utils.py index 993ef01..8211a05 100644 --- a/judgearena/utils.py +++ b/judgearena/utils.py @@ -14,10 +14,6 @@ from tqdm.asyncio import tqdm from tqdm.contrib.logging import logging_redirect_tqdm -from judgearena.instruction_dataset.arena_hard import ( - download_arena_hard, - is_arena_hard_dataset, -) from judgearena.log import get_logger logger = get_logger(__name__) @@ -58,6 +54,13 @@ def read_df(filename: Path, **pandas_kwargs) -> pd.DataFrame: return pd.read_parquet(filename, **pandas_kwargs) +# The instruction_dataset package imports these data helpers during initialization. +from judgearena.instruction_dataset.arena_hard import ( # noqa: E402 + download_arena_hard, + is_arena_hard_dataset, +) + + def compute_pref_summary(prefs: pd.Series) -> dict[str, float | int]: """Compute win/loss/tie stats for preference series (0=A, 0.5=tie, 1=B).""" prefs = pd.Series(prefs, dtype="float64") From 9744f97385d2b352e24a9067f67abcde3a79f713 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Sun, 19 Jul 2026 00:26:27 +0200 Subject: [PATCH 03/13] feat: add meta-eval sampling, prompts, cache, and metrics Introduce the core judge meta-evaluation building blocks and package prompt resources so later CLI wiring can reuse them unchanged. --- judgearena/meta_eval/__init__.py | 1 + judgearena/meta_eval/cache.py | 165 +++++++ judgearena/meta_eval/cost.py | 57 +++ judgearena/meta_eval/metrics.py | 420 ++++++++++++++++++ judgearena/meta_eval/parsers.py | 147 ++++++ judgearena/meta_eval/prompts.py | 63 +++ .../prompts/alpaca_eval_pair_score_user.txt | 32 ++ .../meta_eval/prompts/alpaca_eval_system.txt | 1 + .../meta_eval/prompts/alpaca_eval_user.txt | 30 ++ .../meta_eval/prompts/arena_hard_system.txt | 19 + .../meta_eval/prompts/arena_hard_user.txt | 10 + judgearena/meta_eval/sampling.py | 112 +++++ pyproject.toml | 2 + 13 files changed, 1059 insertions(+) create mode 100644 judgearena/meta_eval/__init__.py create mode 100644 judgearena/meta_eval/cache.py create mode 100644 judgearena/meta_eval/cost.py create mode 100644 judgearena/meta_eval/metrics.py create mode 100644 judgearena/meta_eval/parsers.py create mode 100644 judgearena/meta_eval/prompts.py create mode 100644 judgearena/meta_eval/prompts/alpaca_eval_pair_score_user.txt create mode 100644 judgearena/meta_eval/prompts/alpaca_eval_system.txt create mode 100644 judgearena/meta_eval/prompts/alpaca_eval_user.txt create mode 100644 judgearena/meta_eval/prompts/arena_hard_system.txt create mode 100644 judgearena/meta_eval/prompts/arena_hard_user.txt create mode 100644 judgearena/meta_eval/sampling.py diff --git a/judgearena/meta_eval/__init__.py b/judgearena/meta_eval/__init__.py new file mode 100644 index 0000000..a0c6298 --- /dev/null +++ b/judgearena/meta_eval/__init__.py @@ -0,0 +1 @@ +"""Judge meta-evaluation against human-labeled arena battles.""" diff --git a/judgearena/meta_eval/cache.py b/judgearena/meta_eval/cache.py new file mode 100644 index 0000000..40acd7e --- /dev/null +++ b/judgearena/meta_eval/cache.py @@ -0,0 +1,165 @@ +"""SQLite-backed cache for meta-evaluation judge annotations.""" + +from __future__ import annotations + +import sqlite3 +from dataclasses import astuple, dataclass, field, fields +from datetime import UTC, datetime +from itertools import groupby +from pathlib import Path + +from judgearena.utils import data_root + +DEFAULT_DB_DIR = data_root / "cache" / "db" + + +@dataclass(frozen=True) +class AnnotationEntry: + benchmark: str + instruction_id: str + model_a: str + model_b: str + judge: str + judge_input: str + judge_completion: str + reasoning_content: str = "" + input_tokens: int = 0 + output_tokens: int = 0 + reasoning_tokens: int = 0 + date: str = field(default_factory=lambda: datetime.now(UTC).isoformat()) + + @classmethod + def field_names(cls) -> tuple[str, ...]: + return tuple(field_info.name for field_info in fields(cls)) + + @classmethod + def key_fields(cls) -> tuple[str, ...]: + return ("benchmark", "instruction_id", "model_a", "model_b", "judge") + + +@dataclass(frozen=True) +class AnnotationKey: + benchmark: str + instruction_id: str + model_a: str + model_b: str + judge: str + + @classmethod + def field_names(cls) -> tuple[str, ...]: + return tuple(field_info.name for field_info in fields(cls)) + + +def _db_path(db_dir: Path, benchmark: str, judge: str) -> Path: + return db_dir / benchmark / f"{judge.replace('/', '_')}.db" + + +class AnnotationCache: + """Persistent per-battle cache matching the original meta-eval pipeline.""" + + def __init__(self, db_dir: Path | str = DEFAULT_DB_DIR) -> None: + self._db_dir = Path(db_dir) + self._connections: dict[tuple[str, str], sqlite3.Connection] = {} + + def batch_get_annotations( + self, keys: list[AnnotationKey] + ) -> list[AnnotationEntry | None]: + column_names = ", ".join(AnnotationEntry.field_names()) + results = [] + for key in keys: + row = ( + self._connection(key.benchmark, key.judge) + .execute( + f"SELECT {column_names} FROM annotations " + f"WHERE {self._where_clause()}", + astuple(key), + ) + .fetchone() + ) + results.append(AnnotationEntry(*row) if row else None) + return results + + def batch_put(self, entries: list[AnnotationEntry]) -> None: + if not entries: + return + column_names = ", ".join(AnnotationEntry.field_names()) + placeholders = ", ".join("?" for _ in AnnotationEntry.field_names()) + sql = ( + f"INSERT OR REPLACE INTO annotations ({column_names}) " + f"VALUES ({placeholders})" + ) + + def cache_partition(entry: AnnotationEntry) -> tuple[str, str]: + return entry.benchmark, entry.judge + + for (benchmark, judge), group in groupby( + sorted(entries, key=cache_partition), + key=cache_partition, + ): + connection = self._connection(benchmark, judge) + connection.executemany(sql, [astuple(entry) for entry in group]) + connection.commit() + + def close(self) -> None: + for connection in self._connections.values(): + connection.close() + self._connections.clear() + + def _connection(self, benchmark: str, judge: str) -> sqlite3.Connection: + key = (benchmark, judge) + if key not in self._connections: + path = _db_path(self._db_dir, benchmark, judge) + path.parent.mkdir(parents=True, exist_ok=True) + connection = sqlite3.connect( + str(path), + check_same_thread=False, + timeout=30, + ) + connection.execute("PRAGMA journal_mode=WAL") + self._connections[key] = connection + self._create_table(connection) + return self._connections[key] + + @staticmethod + def _create_table(connection: sqlite3.Connection) -> None: + integer_fields = {"input_tokens", "output_tokens", "reasoning_tokens"} + default_text_fields = {"reasoning_content", "date"} + column_definitions = ", ".join( + ( + f"{name} INTEGER NOT NULL DEFAULT 0" + if name in integer_fields + else ( + f"{name} TEXT NOT NULL DEFAULT ''" + if name in default_text_fields + else f"{name} TEXT NOT NULL" + ) + ) + for name in AnnotationEntry.field_names() + ) + key_columns = ", ".join(AnnotationEntry.key_fields()) + connection.execute( + "CREATE TABLE IF NOT EXISTS annotations " + f"({column_definitions}, UNIQUE ({key_columns}))" + ) + connection.execute( + "CREATE INDEX IF NOT EXISTS idx_annotation_key " + f"ON annotations ({key_columns})" + ) + migrations = [ + "ALTER TABLE annotations ADD COLUMN " + "reasoning_content TEXT NOT NULL DEFAULT ''", + "ALTER TABLE annotations ADD COLUMN input_tokens INTEGER NOT NULL DEFAULT 0", + "ALTER TABLE annotations ADD COLUMN output_tokens INTEGER NOT NULL DEFAULT 0", + "ALTER TABLE annotations ADD COLUMN " + "reasoning_tokens INTEGER NOT NULL DEFAULT 0", + ] + for migration in migrations: + try: + connection.execute(migration) + except sqlite3.OperationalError: + pass + connection.commit() + + @staticmethod + def _where_clause() -> str: + return " AND ".join(f"{column} = ?" for column in AnnotationKey.field_names()) diff --git a/judgearena/meta_eval/cost.py b/judgearena/meta_eval/cost.py new file mode 100644 index 0000000..f2a6ffc --- /dev/null +++ b/judgearena/meta_eval/cost.py @@ -0,0 +1,57 @@ +"""Cost estimation helpers for meta-evaluation annotations.""" + +from __future__ import annotations + +import json + +from judgearena.utils import data_root + +_PRICING_CACHE_FILE = data_root / "cache" / "openrouter_pricing.json" +_openrouter_pricing_cache: dict[str, tuple[float, float]] = {} + + +def _openrouter_model_key(model_name: str) -> str: + if model_name.count("/") >= 2: + return "/".join(model_name.split("/")[-2:]) + return model_name + + +def load_openrouter_pricing() -> dict[str, tuple[float, float]]: + if _openrouter_pricing_cache: + return _openrouter_pricing_cache + if not _PRICING_CACHE_FILE.exists(): + return {} + with _PRICING_CACHE_FILE.open(encoding="utf-8") as handle: + raw = json.load(handle) + for model_id, prices in raw.items(): + prompt, completion = prices + _openrouter_pricing_cache[model_id] = (float(prompt), float(completion)) + return _openrouter_pricing_cache + + +def lookup_openrouter_pricing(model_name: str) -> tuple[float, float] | None: + pricing = load_openrouter_pricing() + key = _openrouter_model_key(model_name) + return pricing.get(key) + + +def estimate_token_count(text: str) -> int: + return len(text) // 4 if isinstance(text, str) else 0 + + +def estimate_annotation_cost_usd( + *, + judge_input: str, + judge_completion: str, + judge_model: str, +) -> tuple[float | None, str]: + """Estimate cost from text length and cached OpenRouter reference pricing.""" + pricing = lookup_openrouter_pricing(judge_model) + if pricing is None: + return None, "unavailable" + + input_price, output_price = pricing + prompt_tokens = estimate_token_count(judge_input) + completion_tokens = estimate_token_count(judge_completion) + cost = (prompt_tokens * input_price + completion_tokens * output_price) / 1e6 + return float(cost), "estimated" diff --git a/judgearena/meta_eval/metrics.py b/judgearena/meta_eval/metrics.py new file mode 100644 index 0000000..e47c0f5 --- /dev/null +++ b/judgearena/meta_eval/metrics.py @@ -0,0 +1,420 @@ +"""Metrics for judge meta-evaluation against human labels.""" + +from __future__ import annotations + +import math + +import numpy as np +import pandas as pd +from scipy.stats import spearmanr +from sklearn.linear_model import LogisticRegression +from sklearn.metrics import cohen_kappa_score + +from judgearena.estimate_elo_ratings import compute_bradley_terry + +WINNER_LABELS = ["model_a", "model_b", "tie"] + + +def _cohen_kappa(y_true: list[str], y_pred: list[str]) -> float: + if len(set(y_true) | set(y_pred)) < 2: + return float("nan") + return float(cohen_kappa_score(y_true, y_pred, labels=WINNER_LABELS)) + + +def _finite_std(values: list[float]) -> float: + finite = np.asarray([value for value in values if math.isfinite(value)]) + return float(np.std(finite)) if len(finite) else float("nan") + + +def bootstrap_std( + y_true: list[str], + y_pred: list[str], + *, + n_bootstraps: int, + seed: int, +) -> tuple[float, float]: + if not y_true: + return float("nan"), float("nan") + rng = np.random.default_rng(seed) + y_true_arr = np.array(y_true) + y_pred_arr = np.array(y_pred) + acc_samples: list[float] = [] + kappa_samples: list[float] = [] + for _ in range(n_bootstraps): + idx = rng.choice(len(y_true_arr), size=len(y_true_arr), replace=True) + acc_samples.append(float(np.mean(y_true_arr[idx] == y_pred_arr[idx]))) + kappa_samples.append( + _cohen_kappa( + y_true_arr[idx].tolist(), + y_pred_arr[idx].tolist(), + ) + ) + return float(np.std(acc_samples)), _finite_std(kappa_samples) + + +def compute_agreement_metrics( + winner_human: list[str], + winner_llm: list[str], + *, + n_bootstraps: int, + seed: int, +) -> dict[str, float | int]: + n_all = len(winner_human) + if n_all == 0: + nan = float("nan") + return { + "n": 0, + "accuracy": nan, + "acc_se": nan, + "kappa": nan, + "kappa_se": nan, + "n_nt": 0, + "accuracy_nt": nan, + "acc_se_nt": nan, + "kappa_nt": nan, + "kappa_se_nt": nan, + } + + acc_all = ( + sum(h == pred for h, pred in zip(winner_human, winner_llm, strict=True)) / n_all + ) + kappa_all = _cohen_kappa(winner_human, winner_llm) + acc_se, kappa_se = bootstrap_std( + winner_human, + winner_llm, + n_bootstraps=n_bootstraps, + seed=seed, + ) + + no_tie = [ + (h, pred) + for h, pred in zip(winner_human, winner_llm, strict=True) + if h != "tie" + ] + wh_nt, wl_nt = zip(*no_tie, strict=False) if no_tie else ([], []) + n_nt = len(wh_nt) + if n_nt: + acc_nt = sum(h == pred for h, pred in zip(wh_nt, wl_nt, strict=True)) / n_nt + kappa_nt = _cohen_kappa(list(wh_nt), list(wl_nt)) + acc_se_nt, kappa_se_nt = bootstrap_std( + list(wh_nt), + list(wl_nt), + n_bootstraps=n_bootstraps, + seed=seed + 1, + ) + else: + acc_nt = float("nan") + kappa_nt = float("nan") + acc_se_nt = float("nan") + kappa_se_nt = float("nan") + + return { + "n": n_all, + "accuracy": acc_all, + "acc_se": acc_se, + "kappa": kappa_all, + "kappa_se": kappa_se, + "n_nt": n_nt, + "accuracy_nt": acc_nt, + "acc_se_nt": acc_se_nt, + "kappa_nt": kappa_nt, + "kappa_se_nt": kappa_se_nt, + } + + +def compute_soft_bradley_terry( + df: pd.DataFrame, + pref_col: str = "pref_llm", + scale: float = 400, + base: float = 10, + init_rating: float = 1000, +) -> dict[str, float]: + df = df.dropna(subset=[pref_col]).copy() + if df.empty: + return {} + + all_models = sorted(set(df["model_a"].unique()) | set(df["model_b"].unique())) + models = pd.Series(np.arange(len(all_models)), index=all_models) + p = len(models) + n_battles = len(df) + x = np.zeros([2 * n_battles, p]) + y = np.zeros(2 * n_battles) + sample_weights = np.zeros(2 * n_battles) + + for idx, (_, row) in enumerate(df.iterrows()): + m_a = row["model_a"] + m_b = row["model_b"] + pref = row[pref_col] + x[2 * idx, models[m_a]] = +np.log(base) + x[2 * idx, models[m_b]] = -np.log(base) + y[2 * idx] = 1.0 + sample_weights[2 * idx] = 1.0 - pref + x[2 * idx + 1, models[m_a]] = +np.log(base) + x[2 * idx + 1, models[m_b]] = -np.log(base) + y[2 * idx + 1] = 0.0 + sample_weights[2 * idx + 1] = pref + + nonzero = sample_weights > 0 + x = x[nonzero] + y = y[nonzero] + sample_weights = sample_weights[nonzero] + if len(x) == 0: + return {} + + try: + lr = LogisticRegression(fit_intercept=False, C=1e10, tol=1e-6, max_iter=1000) + lr.fit(x, y, sample_weight=sample_weights) + except ValueError: + return {} + elo_scores = scale * lr.coef_[0] + init_rating + return dict(pd.Series(elo_scores, index=models.index)) + + +def format_metric(value: float | None, se: float | None, *, digits: int = 2) -> str: + if value is None or not math.isfinite(value): + return "n/a" + if se is None or not math.isfinite(se): + return f"{value:.{digits}f}" + return f"{value:.{digits}f} ± {se:.{digits}f}" + + +def _bt_ratings(df_sub: pd.DataFrame) -> tuple[dict[str, float], dict[str, float]]: + try: + human = compute_bradley_terry( + df_sub[["model_a", "model_b", "winner"]], + "winner", + ) + llm = compute_bradley_terry( + df_sub[["model_a", "model_b", "winner_llm"]].rename( + columns={"winner_llm": "winner"} + ), + "winner", + ) + except ValueError: + return {}, {} + return human, llm + + +def _bt_ratings_soft(df_sub: pd.DataFrame) -> tuple[dict[str, float], dict[str, float]]: + human = compute_bradley_terry(df_sub[["model_a", "model_b", "winner"]], "winner") + llm = compute_soft_bradley_terry(df_sub[["model_a", "model_b", "pref_llm"]]) + return human, llm + + +def _rating_vectors( + df_sub: pd.DataFrame, + *, + soft: bool, +) -> tuple[np.ndarray, np.ndarray]: + ratings_fn = _bt_ratings_soft if soft else _bt_ratings + human, llm = ratings_fn(df_sub) + common = sorted(set(human) & set(llm)) + return ( + np.array([human[model] for model in common]), + np.array([llm[model] for model in common]), + ) + + +def _bootstrap_rank_metric( + hv: np.ndarray, + lv: np.ndarray, + *, + metric: str, + n_bootstraps: int, + seed: int, +) -> tuple[float, float]: + if len(hv) == 0: + return float("nan"), float("nan") + rng = np.random.default_rng(seed) + samples: list[float] = [] + for _ in range(n_bootstraps): + idx = rng.choice(len(hv), size=len(hv), replace=True) + if metric == "spearman": + if len(np.unique(hv[idx])) < 2 or len(np.unique(lv[idx])) < 2: + continue + value = float(spearmanr(hv[idx], lv[idx])[0]) + else: + value = float(np.mean(np.abs(hv[idx] - lv[idx]))) + if math.isfinite(value): + samples.append(value) + if not samples: + return float("nan"), float("nan") + return float(np.mean(samples)), float(np.std(samples)) + + +def spearman_with_se( + df_sub: pd.DataFrame, + *, + n_bootstraps: int, + seed: int, + soft: bool = False, +) -> str: + human, llm = _rating_vectors(df_sub, soft=soft) + if len(human) == 0: + return "n/a" + if len(np.unique(human)) < 2 or len(np.unique(llm)) < 2: + return "n/a" + rho, _ = spearmanr(human, llm) + if rho is None or not math.isfinite(float(rho)): + return "n/a" + _, se = _bootstrap_rank_metric( + human, + llm, + metric="spearman", + n_bootstraps=n_bootstraps, + seed=seed, + ) + return format_metric(float(rho), se) + + +def mae_elo_with_se( + df_sub: pd.DataFrame, + *, + n_bootstraps: int, + seed: int, + soft: bool = False, +) -> str: + human, llm = _rating_vectors(df_sub, soft=soft) + if len(human) == 0: + return "n/a" + mae = float(np.mean(np.abs(human - llm))) + _, se = _bootstrap_rank_metric( + human, + llm, + metric="mae", + n_bootstraps=n_bootstraps, + seed=seed, + ) + return format_metric(mae, se, digits=1) + + +def summarize_language_splits( + df_ann: pd.DataFrame, + *, + exclude_human_ties: bool, + n_bootstraps: int, + seed: int, +) -> dict[str, dict[str, str | int]]: + rows: dict[str, dict[str, str | int]] = {} + for label, mask in [ + ("English", df_ann["lang"] == "en"), + ("Multilingual", df_ann["lang"] != "en"), + ]: + df_sub = df_ann[mask] + if exclude_human_ties: + df_sub = df_sub[df_sub["winner"] != "tie"] + metrics = compute_agreement_metrics( + df_sub["winner"].tolist(), + df_sub["winner_llm"].tolist(), + n_bootstraps=n_bootstraps, + seed=seed, + ) + entry: dict[str, str | int] = {"n": metrics["n"]} + if metrics["n"] == 0: + entry.update( + { + "kappa": "n/a", + "spearman": "n/a", + "spearman_soft": "n/a", + "mae_elo": "n/a", + "mae_soft_elo": "n/a", + } + ) + else: + entry["kappa"] = format_metric( + float(metrics["kappa"]), + float(metrics["kappa_se"]), + ) + entry["spearman"] = spearman_with_se( + df_sub, + n_bootstraps=n_bootstraps, + seed=seed + 2, + soft=False, + ) + entry["spearman_soft"] = spearman_with_se( + df_sub, + n_bootstraps=n_bootstraps, + seed=seed + 3, + soft=True, + ) + entry["mae_elo"] = mae_elo_with_se( + df_sub, + n_bootstraps=n_bootstraps, + seed=seed + 4, + soft=False, + ) + entry["mae_soft_elo"] = mae_elo_with_se( + df_sub, + n_bootstraps=n_bootstraps, + seed=seed + 5, + soft=True, + ) + rows[label] = entry + return rows + + +def compute_elo_gap_summary( + df_top: pd.DataFrame, + df_ann: pd.DataFrame, + top_models: list[str], + *, + n_battles_list: list[int], + n_seeds: int, + seed: int, + exclude_ties: bool, +) -> pd.DataFrame: + df_battles = df_top[["model_a", "model_b", "winner"]].copy() + human_ratings = compute_bradley_terry(df_battles, "winner") + rows: list[dict[str, float | int | str]] = [] + + for num_battles in n_battles_list: + for offset in range(n_seeds): + rng = np.random.default_rng(seed + offset) + gaps: list[float] = [] + for model in top_models: + model_mask_ann = (df_ann["model_a"] == model) | ( + df_ann["model_b"] == model + ) + if model_mask_ann.sum() < num_battles: + continue + model_mask_top = (df_battles["model_a"] == model) | ( + df_battles["model_b"] == model + ) + other_human = df_battles[~model_mask_top].copy() + sample = df_ann[model_mask_ann].sample( + n=num_battles, + replace=False, + random_state=int(rng.integers(0, 2**32 - 1)), + ) + if exclude_ties: + sample = sample[sample["winner_llm"] != "tie"] + if sample.empty: + continue + model_llm = sample[["model_a", "model_b", "winner_llm"]].rename( + columns={"winner_llm": "winner"} + ) + hybrid = pd.concat([other_human, model_llm], ignore_index=True) + hybrid_ratings = compute_bradley_terry(hybrid, "winner") + if model in hybrid_ratings and model in human_ratings: + gaps.append( + abs(hybrid_ratings[model] - human_ratings[model]), + ) + if gaps: + rows.append( + { + "num_battles": num_battles, + "seed": offset, + "mean_gap": float(np.mean(gaps)), + "exclude_ties": exclude_ties, + } + ) + + if not rows: + return pd.DataFrame(columns=["num_battles", "mean", "se", "exclude_ties"]) + + df_rows = pd.DataFrame(rows) + return ( + df_rows.groupby(["num_battles", "exclude_ties"])["mean_gap"] + .agg(mean="mean", se=lambda values: values.std() / np.sqrt(len(values))) + .reset_index() + ) diff --git a/judgearena/meta_eval/parsers.py b/judgearena/meta_eval/parsers.py new file mode 100644 index 0000000..828dcb2 --- /dev/null +++ b/judgearena/meta_eval/parsers.py @@ -0,0 +1,147 @@ +"""Winner and preference parsers for meta-evaluation prompt modes.""" + +from __future__ import annotations + +import json +import re + +import numpy as np +import pandas as pd + +from judgearena.evaluate import PairScore + +TIE_EPSILON = 0.01 +META_EVAL_PAIRSCORE_TEMPERATURE = 0.5 + +_ARENA_HARD_PATTERNS = [ + re.compile(r"\[\[([AB<>=]+)\]\]"), + re.compile(r"\[([AB<>=]+)\]"), +] +_ARENA_HARD_LIKERT_TO_WINNER = { + "A>>B": "model_a", + "A>B": "model_a", + "A=B": "tie", + "B>A": "model_b", + "B>>A": "model_b", + "B< PairScore: + parser = PairScore() + parser.temperature = temperature + return parser + + +def parse_pairscore_pref(judge_completion: str, *, temperature: float) -> float: + score = pair_score_parser(temperature).parse_model_raw(judge_completion) + if score is None or np.isnan(score): + return 0.5 + return float(score) + + +def parse_pairscore_winner( + judge_completion: str, + *, + temperature: float, + eps: float = TIE_EPSILON, +) -> str: + score = parse_pairscore_pref(judge_completion, temperature=temperature) + if abs(score - 0.5) < eps: + return "tie" + if score > 0.5 + eps: + return "model_b" + if score < 0.5 - eps: + return "model_a" + return "tie" + + +def parse_arena_hard_winner(judge_completion: str) -> str: + if not isinstance(judge_completion, str): + return "tie" + for pattern in _ARENA_HARD_PATTERNS: + matches = pattern.findall(judge_completion.upper()) + matches = [match for match in matches if match] + if matches: + return _ARENA_HARD_LIKERT_TO_WINNER.get(matches[-1].strip(), "tie") + return "tie" + + +def parse_alpaca_eval_winner(judge_completion: str) -> str: + if not isinstance(judge_completion, str): + return "tie" + + text = judge_completion + fenced = re.search(r"```json\s*(.*?)\s*```", text, re.DOTALL) + if fenced: + text = fenced.group(1) + else: + obj_match = re.search( + r'\{[^{}]*"ordered_models"[^{}]*\[[^\[\]]*\][^{}]*\}', + text, + re.DOTALL, + ) + if obj_match: + text = obj_match.group(0) + + try: + data = json.loads(text) + ordered = data.get("ordered_models", []) + m_entry = next((entry for entry in ordered if entry.get("model") == "m"), None) + if m_entry is None: + return "tie" + rank_m = m_entry["rank"] + if rank_m == 1: + return "model_a" + if rank_m == 2: + return "model_b" + except (json.JSONDecodeError, KeyError, TypeError): + pass + return "tie" + + +def winner_to_pref(winner: str) -> float: + return {"model_a": 0.0, "model_b": 1.0}.get(winner, 0.5) + + +def parse_winner(judge_completion: str, prompt_mode: str) -> str: + if prompt_mode == "arena-hard": + return parse_arena_hard_winner(judge_completion) + if prompt_mode == "alpaca-eval": + return parse_alpaca_eval_winner(judge_completion) + return parse_pairscore_winner( + judge_completion, + temperature=META_EVAL_PAIRSCORE_TEMPERATURE, + ) + + +def parse_pref(judge_completion: str, prompt_mode: str) -> float: + if prompt_mode in ("arena-hard", "alpaca-eval"): + return winner_to_pref(parse_winner(judge_completion, prompt_mode)) + return parse_pairscore_pref( + judge_completion, + temperature=META_EVAL_PAIRSCORE_TEMPERATURE, + ) + + +def add_parsed_columns(df: pd.DataFrame, prompt_mode: str) -> pd.DataFrame: + out = df.copy() + completions = out["judge_completion"].tolist() + out["winner_llm"] = [ + parse_winner(completion, prompt_mode) for completion in completions + ] + out["pref_llm"] = [ + parse_pref(completion, prompt_mode) for completion in completions + ] + return out + + +def invert_winner(winner: str) -> str: + if winner == "model_a": + return "model_b" + if winner == "model_b": + return "model_a" + return winner diff --git a/judgearena/meta_eval/prompts.py b/judgearena/meta_eval/prompts.py new file mode 100644 index 0000000..cd225e4 --- /dev/null +++ b/judgearena/meta_eval/prompts.py @@ -0,0 +1,63 @@ +"""Named prompt templates and mode registry for meta-evaluation.""" + +from __future__ import annotations + +from dataclasses import dataclass +from importlib.resources import files + +from judgearena.evaluate import load_judge_system_and_user_prompt +from judgearena.meta_eval.cli_args import PROMPT_MODES + + +@dataclass(frozen=True) +class PromptModeSpec: + name: str + system_prompt: str | None = None + user_prompt_template: str | None = None + + +def _read_prompt(filename: str) -> str: + return ( + files("judgearena.meta_eval") + .joinpath("prompts", filename) + .read_text(encoding="utf-8") + ) + + +def resolve_prompt_mode( + prompt_mode: str, + *, + provide_explanation: bool = False, +) -> PromptModeSpec: + if prompt_mode not in PROMPT_MODES: + raise ValueError(f"Unknown prompt mode {prompt_mode!r}.") + + if prompt_mode == "standard": + system_prompt, user_prompt_template = load_judge_system_and_user_prompt( + provide_explanation=provide_explanation, + ) + return PromptModeSpec( + name=prompt_mode, + system_prompt=system_prompt, + user_prompt_template=user_prompt_template, + ) + + if prompt_mode == "arena-hard": + return PromptModeSpec( + name=prompt_mode, + system_prompt=_read_prompt("arena_hard_system.txt"), + user_prompt_template=_read_prompt("arena_hard_user.txt"), + ) + + if prompt_mode == "alpaca-eval": + return PromptModeSpec( + name=prompt_mode, + system_prompt=_read_prompt("alpaca_eval_system.txt"), + user_prompt_template=_read_prompt("alpaca_eval_user.txt"), + ) + + return PromptModeSpec( + name=prompt_mode, + system_prompt=_read_prompt("alpaca_eval_system.txt"), + user_prompt_template=_read_prompt("alpaca_eval_pair_score_user.txt"), + ) diff --git a/judgearena/meta_eval/prompts/alpaca_eval_pair_score_user.txt b/judgearena/meta_eval/prompts/alpaca_eval_pair_score_user.txt new file mode 100644 index 0000000..9b676d6 --- /dev/null +++ b/judgearena/meta_eval/prompts/alpaca_eval_pair_score_user.txt @@ -0,0 +1,32 @@ +I require a leaderboard for various large language models. I'll provide you with prompts given to these models and their corresponding responses. Your task is to assess these responses, ranking the models in order of preference from a human perspective. + +## Prompt + +{{ + "instruction": "{user_prompt}", +}} + +## Model Outputs + +Here are the unordered outputs from the models. Each output is associated with a specific model, identified by a unique model identifier. + +{{ + {{ + "model": "model A", + "output": "{completion_A}" + }}, + {{ + "model": "model B", + "output": "{completion_B}" + }} +}} + +## Task + +Evaluate and rank the models based on the quality and relevance of their outputs. The ranking should be such that the model with the highest quality output is ranked first. + +## Your output, do not repeat the input above +``` +score_A: +score_B: +``` diff --git a/judgearena/meta_eval/prompts/alpaca_eval_system.txt b/judgearena/meta_eval/prompts/alpaca_eval_system.txt new file mode 100644 index 0000000..41b078f --- /dev/null +++ b/judgearena/meta_eval/prompts/alpaca_eval_system.txt @@ -0,0 +1 @@ +You are a highly efficient assistant, who evaluates and ranks large language models (LLMs) based on the quality of their responses to given prompts. This process will create a leaderboard reflecting the most accurate and human-preferred answers. diff --git a/judgearena/meta_eval/prompts/alpaca_eval_user.txt b/judgearena/meta_eval/prompts/alpaca_eval_user.txt new file mode 100644 index 0000000..c290708 --- /dev/null +++ b/judgearena/meta_eval/prompts/alpaca_eval_user.txt @@ -0,0 +1,30 @@ +I require a leaderboard for various large language models. I'll provide you with prompts given to these models and their corresponding responses. Your task is to assess these responses, ranking the models in order of preference from a human perspective. + +## Prompt + +{{ + "instruction": "{user_prompt}", +}} + +## Model Outputs + +Here are the unordered outputs from the models. Each output is associated with a specific model, identified by a unique model identifier. + +{{ + {{ + "model": "m", + "output": "{completion_A}" + }}, + {{ + "model": "M", + "output": "{completion_B}" + }} +}} + +## Task + +Evaluate and rank the models based on the quality and relevance of their outputs. The ranking should be such that the model with the highest quality output is ranked first. + +Output ONLY a JSON object in this exact format (no other text): +{{"ordered_models": [{{"model": "m", "rank": 1}}, {{"model": "M", "rank": 2}}]}} +where rank 1 = best. Replace ranks according to your evaluation. diff --git a/judgearena/meta_eval/prompts/arena_hard_system.txt b/judgearena/meta_eval/prompts/arena_hard_system.txt new file mode 100644 index 0000000..0c0d8d1 --- /dev/null +++ b/judgearena/meta_eval/prompts/arena_hard_system.txt @@ -0,0 +1,19 @@ +Please act as an impartial judge and evaluate the quality of the responses provided by two AI assistants to the user prompt displayed below. You will be given assistant A's answer and assistant B's answer. Your job is to evaluate which assistant's answer is better. + +Begin your evaluation by generating your own answer to the prompt. You must provide your answers before judging any answers. + +When evaluating the assistants' answers, compare both assistants' answers with your answer. You must identify and correct any mistakes or inaccurate information. + +Then consider if the assistant's answers are helpful, relevant, and concise. Helpful means the answer correctly responds to the prompt or follows the instructions. Note when user prompt has any ambiguity or more than one interpretation, it is more helpful and appropriate to ask for clarifications or more information from the user than providing an answer based on assumptions. Relevant means all parts of the response closely connect or are appropriate to what is being asked. Concise means the response is clear and not verbose or excessive. + +Then consider the creativity and novelty of the assistant's answers when needed. Finally, identify any missing important information in the assistants' answers that would be beneficial to include when responding to the user prompt. + +After providing your explanation, you must output only one of the following choices as your final verdict with a label: + +1. Assistant A is significantly better: [[A>>B]] +2. Assistant A is slightly better: [[A>B]] +3. Tie, relatively the same: [[A=B]] +4. Assistant B is slightly better: [[B>A]] +5. Assistant B is significantly better: [[B>>A]] + +Example output: "My final verdict is tie: [[A=B]]". diff --git a/judgearena/meta_eval/prompts/arena_hard_user.txt b/judgearena/meta_eval/prompts/arena_hard_user.txt new file mode 100644 index 0000000..8f2db9e --- /dev/null +++ b/judgearena/meta_eval/prompts/arena_hard_user.txt @@ -0,0 +1,10 @@ +<|User Prompt|> +{user_prompt} + +<|The Start of Assistant A's Answer|> +{completion_A} +<|The End of Assistant A's Answer|> + +<|The Start of Assistant B's Answer|> +{completion_B} +<|The End of Assistant B's Answer|> diff --git a/judgearena/meta_eval/sampling.py b/judgearena/meta_eval/sampling.py new file mode 100644 index 0000000..0311bd8 --- /dev/null +++ b/judgearena/meta_eval/sampling.py @@ -0,0 +1,112 @@ +"""Deterministic arena battle sampling for meta-evaluation.""" + +from __future__ import annotations + +import pandas as pd + +from judgearena.arenas_utils import KNOWN_ARENAS, load_arena_dataframe +from judgearena.log import get_logger + +logger = get_logger(__name__) + + +class MetaEvalSamplingError(ValueError): + """Raised when filtering or sampling yields an unusable subset.""" + + +def normalize_human_winner(winner: object) -> str: + text = str(winner) + if "tie" in text: + return "tie" + return text + + +def count_battles_per_model(df: pd.DataFrame) -> dict[str, int]: + return ( + pd.concat([df["model_a"], df["model_b"]], ignore_index=True) + .value_counts() + .to_dict() + ) + + +def load_reference_arena_battles( + reference_arena: str, + *, + languages: list[str] | None = None, +) -> pd.DataFrame: + if reference_arena not in KNOWN_ARENAS: + raise MetaEvalSamplingError( + f"Unsupported reference arena {reference_arena!r}; " + f"expected one of {KNOWN_ARENAS}." + ) + + df = load_arena_dataframe(arena=reference_arena).copy() + df["winner"] = df["winner"].map(normalize_human_winner) + + if languages: + df = df[df["lang"].isin(languages)].copy() + if df.empty: + langs = ", ".join(languages) + raise MetaEvalSamplingError( + f"No battles remain after filtering to languages: {langs}." + ) + + if df.empty: + raise MetaEvalSamplingError( + f"No battles found for reference arena {reference_arena!r}." + ) + return df + + +def select_top_models( + df: pd.DataFrame, + *, + top_models: int, +) -> tuple[list[str], pd.DataFrame]: + battle_counts = count_battles_per_model(df) + if not battle_counts: + raise MetaEvalSamplingError("Cannot select top models from an empty dataframe.") + + top = sorted(battle_counts, key=battle_counts.__getitem__, reverse=True)[ + :top_models + ] + top_set = set(top) + df_top = df[df["model_a"].isin(top_set) & df["model_b"].isin(top_set)].copy() + if df_top.empty: + raise MetaEvalSamplingError( + f"No battles remain among the top {top_models} models." + ) + return top, df_top + + +def sample_battles_per_model( + df_top: pd.DataFrame, + top_models: list[str], + *, + battles_per_model: int, + seed: int, +) -> pd.DataFrame: + per_model_samples: list[pd.DataFrame] = [] + for model_index, model in enumerate(top_models): + model_mask = (df_top["model_a"] == model) | (df_top["model_b"] == model) + df_model = df_top[model_mask] + if df_model.empty: + logger.warning("Model %s has no battles among top models; skipping.", model) + continue + sample_size = min(battles_per_model, len(df_model)) + sampled = df_model.sample( + n=sample_size, + replace=False, + random_state=seed + model_index, + ) + per_model_samples.append(sampled) + + if not per_model_samples: + raise MetaEvalSamplingError( + "Sampling produced no battles; reduce top_models or battles_per_model." + ) + + df_sample = pd.concat(per_model_samples, ignore_index=True) + if df_sample.empty: + raise MetaEvalSamplingError("Sampled battle set is empty.") + return df_sample diff --git a/pyproject.toml b/pyproject.toml index 2d421a9..5d0f8d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,6 +41,7 @@ dependencies = [ "pandas>=2.3.2", "pyyaml>=6.0.2", "scikit-learn>=1.8.0", + "scipy>=1.16.0", "seaborn>=0.13.2", "tqdm>=4.67.1", ] @@ -57,6 +58,7 @@ exclude = ["slurmpilot_scripts*"] [tool.setuptools.package-data] "judgearena.criteria" = ["data/*.yaml"] "judgearena.prompts" = ["*.txt"] +"judgearena.meta_eval" = ["prompts/*.txt"] [dependency-groups] dev = [ From 8498aa892c7d5d9d406c72c537f75b15dd0804fa Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Sun, 19 Jul 2026 00:26:32 +0200 Subject: [PATCH 04/13] feat: wire meta-eval annotation runner into CLI Add the annotation orchestration path and expose meta-eval as a flat --task route alongside generate+judge and ELO workflows. --- judgearena/cli.py | 30 +++- judgearena/cli_common.py | 15 +- judgearena/meta_eval/annotate.py | 273 +++++++++++++++++++++++++++++++ judgearena/meta_eval/cli_args.py | 142 ++++++++++++++++ judgearena/meta_eval/runner.py | 259 +++++++++++++++++++++++++++++ 5 files changed, 706 insertions(+), 13 deletions(-) create mode 100644 judgearena/meta_eval/annotate.py create mode 100644 judgearena/meta_eval/cli_args.py create mode 100644 judgearena/meta_eval/runner.py diff --git a/judgearena/cli.py b/judgearena/cli.py index eb94c83..2e91320 100644 --- a/judgearena/cli.py +++ b/judgearena/cli.py @@ -20,10 +20,16 @@ from judgearena.generate_and_evaluate import CliArgs from judgearena.generate_and_evaluate import main as main_generate_and_evaluate from judgearena.log import configure_logging, get_logger +from judgearena.meta_eval.cli_args import ( + add_meta_eval_arguments, + build_meta_eval_args, +) +from judgearena.meta_eval.runner import run_or_exit as main_meta_eval logger = get_logger(__name__) ELO_TASK_PREFIX = "elo-" +META_EVAL_TASK = "meta-eval" # Lowercase CLI task name -> canonical arena identifier used inside # ``judgearena.arenas_utils.KNOWN_ARENAS`` and the ``benchmark`` column of @@ -42,8 +48,9 @@ def _build_parser() -> argparse.ArgumentParser: prog="judgearena", description=( "Run a judge-based evaluation. Use `--task ` for generate+judge " - "benchmarks (e.g. alpaca-eval, arena-hard-v2.0, mt-bench) or " - "`--task elo-` for ELO rating (e.g. elo-lmarena-140k, elo-comparia)." + "benchmarks (e.g. alpaca-eval, arena-hard-v2.0, mt-bench), " + "`--task meta-eval` for judge meta-evaluation against human arena labels, " + "or `--task elo-` for ELO rating (e.g. elo-lmarena-140k, elo-comparia)." ), ) parser.add_argument( @@ -51,8 +58,8 @@ def _build_parser() -> argparse.ArgumentParser: help=( "Task to run. Generate+judge tasks: `alpaca-eval`, `arena-hard-v0.1`, " "`arena-hard-v2.0`, `m-arena-hard`, `m-arena-hard-{lang}`, `m-arena-hard-EU`, " - "`mt-bench`, `fluency-{lang}`. ELO tasks: `elo-lmarena-100k`, `elo-lmarena-140k`, " - "`elo-lmarena`, `elo-comparia`." + "`mt-bench`, `fluency-{lang}`, `meta-eval`. ELO tasks: `elo-lmarena-100k`, " + "`elo-lmarena-140k`, `elo-lmarena`, `elo-comparia`." ), ) parser.add_argument( @@ -87,7 +94,9 @@ def _build_parser() -> argparse.ArgumentParser: "--languages", nargs="+", default=None, - help="[elo] Language codes to evaluate, e.g. `en fr de`.", + help=( + "[elo/meta-eval] Language codes to evaluate, e.g. `en es fr` (ISO 639-1)." + ), ) parser.add_argument( "--n_instructions_per_language", @@ -99,13 +108,13 @@ def _build_parser() -> argparse.ArgumentParser: "--n_bootstraps", type=int, default=20, - help="[elo] Bootstrap samples for ELO confidence intervals.", + help="[elo/meta-eval] Bootstrap samples for uncertainty estimates.", ) parser.add_argument( "--seed", type=int, default=0, - help="[elo] Random seed for reproducibility.", + help="[elo/meta-eval] Random seed for reproducibility.", ) parser.add_argument( "--baseline_model", @@ -113,6 +122,7 @@ def _build_parser() -> argparse.ArgumentParser: default=None, help="[elo] Model anchored at 1000 ELO (ratings are reported relative to it).", ) + add_meta_eval_arguments(parser) add_common_arguments(parser) return parser @@ -243,7 +253,11 @@ def cli(argv: list[str] | None = None) -> None: configure_logging(resolve_verbosity(args), log_file=args.log_file) task = _resolve_task(args) model_a = _resolve_model_a(args) - if task.startswith(ELO_TASK_PREFIX): + if task == META_EVAL_TASK: + meta_args = build_meta_eval_args(args) + logger.debug("Running with CLI args: %s", meta_args.__dict__) + main_meta_eval(meta_args) + elif task.startswith(ELO_TASK_PREFIX): if task not in ELO_TASK_TO_ARENA: raise SystemExit( f"Unknown elo task {task!r}; expected one of {list(ELO_TASK_TO_ARENA)}." diff --git a/judgearena/cli_common.py b/judgearena/cli_common.py index 58ce78b..eaab581 100644 --- a/judgearena/cli_common.py +++ b/judgearena/cli_common.py @@ -11,6 +11,9 @@ import json from dataclasses import dataclass, field +DEFAULT_MAX_OUT_TOKENS_MODELS = 32768 +DEFAULT_MAX_OUT_TOKENS_JUDGE = 32768 + @dataclass class BaseCliArgs: @@ -23,8 +26,8 @@ class BaseCliArgs: swap_mode: str = "fixed" ignore_cache: bool = False truncate_all_input_chars: int = 8192 - max_out_tokens_models: int = 32768 - max_out_tokens_judge: int = 32768 + max_out_tokens_models: int = DEFAULT_MAX_OUT_TOKENS_MODELS + max_out_tokens_judge: int = DEFAULT_MAX_OUT_TOKENS_JUDGE max_model_len: int | None = None chat_template: str | None = None result_folder: str = "results" @@ -77,7 +80,9 @@ def add_common_arguments(parser: argparse.ArgumentParser) -> None: "Model comparison order mode. 'fixed': always use model order A-B. " "'both': correct for model order bias by evaluating each instruction " "twice, once as A-B and once as B-A, and concatenating the results. " - "This helps account for judge position bias. Default is 'fixed'." + "This helps account for judge position bias. For meta-eval, overall " + "agreement uses both passes while ranking and ELO-gap analyses retain " + "the forward pass. Default is 'fixed'." ), ) parser.add_argument( @@ -110,7 +115,7 @@ def add_common_arguments(parser: argparse.ArgumentParser) -> None: "--max_out_tokens_models", type=int, required=False, - default=32768, + default=DEFAULT_MAX_OUT_TOKENS_MODELS, help=( "Generation token budget for each model A/B response. For VLLM, " "keep this <= --max_model_len (if provided)." @@ -120,7 +125,7 @@ def add_common_arguments(parser: argparse.ArgumentParser) -> None: "--max_out_tokens_judge", type=int, required=False, - default=32768, + default=DEFAULT_MAX_OUT_TOKENS_JUDGE, help=( "Generation token budget for the judge response (reasoning + scores). " "For VLLM, keep this <= --max_model_len (if provided)." diff --git a/judgearena/meta_eval/annotate.py b/judgearena/meta_eval/annotate.py new file mode 100644 index 0000000..3536221 --- /dev/null +++ b/judgearena/meta_eval/annotate.py @@ -0,0 +1,273 @@ +"""Run LLM judge annotations for meta-evaluation battles.""" + +from __future__ import annotations + +import pandas as pd + +from judgearena.arenas_utils import extract_turn_text +from judgearena.evaluate import JudgeAnnotation, annotate_battles +from judgearena.meta_eval.cache import ( + AnnotationCache, + AnnotationEntry, + AnnotationKey, +) +from judgearena.meta_eval.cli_args import CliMetaEvalArgs +from judgearena.meta_eval.cost import ( + estimate_annotation_cost_usd, + estimate_token_count, +) +from judgearena.meta_eval.parsers import add_parsed_columns, invert_winner + + +def _battle_texts(df_batch: pd.DataFrame) -> tuple[list[str], list[str], list[str]]: + instructions = [extract_turn_text(conv[0]) for conv in df_batch["conversation_a"]] + completions_a = [ + extract_turn_text(conv[1]) if len(conv) > 1 else "" + for conv in df_batch["conversation_a"] + ] + completions_b = [ + extract_turn_text(conv[1]) if len(conv) > 1 else "" + for conv in df_batch["conversation_b"] + ] + return instructions, completions_a, completions_b + + +def _swap_batch(df_batch: pd.DataFrame) -> pd.DataFrame: + swapped = df_batch.copy() + swapped["conversation_a"] = df_batch["conversation_b"] + swapped["conversation_b"] = df_batch["conversation_a"] + swapped["model_a"] = df_batch["model_b"] + swapped["model_b"] = df_batch["model_a"] + return swapped + + +def _annotations_to_frame( + df_batch: pd.DataFrame, + annotations, + *, + prompt_mode: str, + judge_model: str, +) -> pd.DataFrame: + rows = [] + for ann, (_, battle) in zip(annotations, df_batch.iterrows(), strict=True): + judge_input = ann.judge_input or "" + cost_usd, cost_source = estimate_annotation_cost_usd( + judge_input=judge_input, + judge_completion=ann.judge_completion, + judge_model=judge_model, + ) + rows.append( + { + "question_id": battle["question_id"], + "model_a": battle["model_a"], + "model_b": battle["model_b"], + "winner": battle["winner"], + "lang": battle["lang"], + "benchmark": battle["benchmark"], + "instruction": ann.instruction, + "completion_a": ann.completion_A, + "completion_b": ann.completion_B, + "judge_input": ann.judge_input, + "judge_completion": ann.judge_completion, + "estimated_input_tokens": estimate_token_count(judge_input), + "estimated_output_tokens": estimate_token_count(ann.judge_completion), + "cost_usd": cost_usd, + "cost_source": cost_source, + } + ) + return add_parsed_columns(pd.DataFrame(rows), prompt_mode) + + +def _judge_cache_name(args: CliMetaEvalArgs) -> str: + if args.prompt_mode == "standard": + return args.judge_model + return f"{args.judge_model}::{args.prompt_mode}" + + +def _cache_keys( + df_batch: pd.DataFrame, + *, + judge: str, +) -> list[AnnotationKey]: + return [ + AnnotationKey( + benchmark=str(battle["benchmark"]), + instruction_id=str(battle["question_id"]), + model_a=str(battle["model_a"]), + model_b=str(battle["model_b"]), + judge=judge, + ) + for _, battle in df_batch.iterrows() + ] + + +def _annotation_from_entry( + entry: AnnotationEntry, + *, + instruction: str, + completion_a: str, + completion_b: str, +) -> JudgeAnnotation: + return JudgeAnnotation( + instruction=instruction, + completion_A=completion_a, + completion_B=completion_b, + judge_completion=entry.judge_completion, + judge_input=entry.judge_input, + ) + + +def _run_cached_batch( + df_batch: pd.DataFrame, + args: CliMetaEvalArgs, + *, + judge_chat_model, + annotation_cache: AnnotationCache, + prompt_spec, + swapped: bool, +) -> pd.DataFrame: + working = _swap_batch(df_batch) if swapped else df_batch + instructions, completions_a, completions_b = _battle_texts(working) + judge = _judge_cache_name(args) + keys = _cache_keys(working, judge=judge) + cached_entries = ( + [None] * len(keys) + if args.ignore_cache + else annotation_cache.batch_get_annotations(keys) + ) + missing_indices = [ + index for index, entry in enumerate(cached_entries) if entry is None + ] + + if missing_indices: + new_annotations = annotate_battles( + judge_chat_model=judge_chat_model, + instructions=[instructions[index] for index in missing_indices], + completions_A=[completions_a[index] for index in missing_indices], + completions_B=[completions_b[index] for index in missing_indices], + system_prompt=prompt_spec.system_prompt, + user_prompt_template=prompt_spec.user_prompt_template, + truncate_input_chars=args.truncate_all_input_chars, + provide_explanation=args.provide_explanation, + ) + new_entries = [ + AnnotationEntry( + **key.__dict__, + judge_input=annotation.judge_input or "", + judge_completion=annotation.judge_completion, + ) + for key, annotation in zip( + [keys[index] for index in missing_indices], + new_annotations, + strict=True, + ) + ] + annotation_cache.batch_put(new_entries) + for index, entry in zip(missing_indices, new_entries, strict=True): + cached_entries[index] = entry + + annotations = [ + _annotation_from_entry( + entry, + instruction=instruction, + completion_a=completion_a, + completion_b=completion_b, + ) + for entry, instruction, completion_a, completion_b in zip( + cached_entries, + instructions, + completions_a, + completions_b, + strict=True, + ) + if entry is not None + ] + return _annotations_to_frame( + working, + annotations, + prompt_mode=args.prompt_mode, + judge_model=args.judge_model, + ) + + +def _normalize_pass_frame( + pass_frame: pd.DataFrame, + original_batch: pd.DataFrame, + *, + orientation: str, +) -> pd.DataFrame: + normalized = pass_frame.copy() + normalized["orientation"] = orientation + normalized["presented_model_a"] = pass_frame["model_a"].tolist() + normalized["presented_model_b"] = pass_frame["model_b"].tolist() + normalized["presented_completion_a"] = pass_frame["completion_a"].tolist() + normalized["presented_completion_b"] = pass_frame["completion_b"].tolist() + normalized["model_a"] = original_batch["model_a"].tolist() + normalized["model_b"] = original_batch["model_b"].tolist() + normalized["winner"] = original_batch["winner"].tolist() + if orientation == "swapped": + normalized["completion_a"] = pass_frame["completion_b"].tolist() + normalized["completion_b"] = pass_frame["completion_a"].tolist() + normalized["winner_llm"] = [ + invert_winner(winner) for winner in pass_frame["winner_llm"] + ] + normalized["pref_llm"] = 1.0 - pass_frame["pref_llm"] + return normalized + + +def annotate_sample( + df_sample: pd.DataFrame, + args: CliMetaEvalArgs, + *, + judge_chat_model, + prompt_spec, + annotation_cache: AnnotationCache | None = None, +) -> pd.DataFrame: + n_total = len(df_sample) + n_batches = (n_total + args.batch_size - 1) // args.batch_size + parts: list[pd.DataFrame] = [] + owns_cache = annotation_cache is None + cache = annotation_cache or AnnotationCache() + + try: + for batch_idx in range(n_batches): + start = batch_idx * args.batch_size + end = min(start + args.batch_size, n_total) + df_batch = df_sample.iloc[start:end].copy() + batch_df = _run_cached_batch( + df_batch, + args, + judge_chat_model=judge_chat_model, + annotation_cache=cache, + prompt_spec=prompt_spec, + swapped=False, + ) + parts.append( + _normalize_pass_frame( + batch_df, + df_batch, + orientation="forward", + ) + ) + + if args.swap_mode == "both": + swapped_df = _run_cached_batch( + df_batch, + args, + judge_chat_model=judge_chat_model, + annotation_cache=cache, + prompt_spec=prompt_spec, + swapped=True, + ) + parts.append( + _normalize_pass_frame( + swapped_df, + df_batch, + orientation="swapped", + ) + ) + finally: + if owns_cache: + cache.close() + + return pd.concat(parts, ignore_index=True) diff --git a/judgearena/meta_eval/cli_args.py b/judgearena/meta_eval/cli_args.py new file mode 100644 index 0000000..7f8c53d --- /dev/null +++ b/judgearena/meta_eval/cli_args.py @@ -0,0 +1,142 @@ +"""CLI argument dataclass for meta-evaluation.""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass + +from judgearena.cli_common import ( + DEFAULT_MAX_OUT_TOKENS_MODELS, + BaseCliArgs, + parse_engine_kwargs, + resolve_verbosity, +) + +PROMPT_MODES = ( + "standard", + "arena-hard", + "alpaca-eval", + "alpaca-eval-pair-score", +) + + +@dataclass +class CliMetaEvalArgs(BaseCliArgs): + """CLI arguments for judge meta-evaluation.""" + + reference_arena: str = "LMArena-140k" + prompt_mode: str = "standard" + top_models: int = 20 + battles_per_model: int = 50 + batch_size: int = 50 + languages: list[str] | None = None + n_bootstraps: int = 20 + seed: int = 0 + elo_gap_battles: list[int] | None = None + elo_gap_seeds: int = 10 + exclude_human_ties: bool = True + + def __post_init__(self) -> None: + super().__post_init__() + if self.prompt_mode not in PROMPT_MODES: + raise ValueError( + f"Unsupported prompt_mode {self.prompt_mode!r}; " + f"expected one of {PROMPT_MODES}." + ) + if self.elo_gap_battles is None: + self.elo_gap_battles = [10, 20, 30, 40, 50] + + +def add_meta_eval_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--reference_arena", + default="LMArena-140k", + help="[meta-eval] Human-labeled reference arena to sample battles from.", + ) + parser.add_argument( + "--prompt_mode", + choices=list(PROMPT_MODES), + default="standard", + help="[meta-eval] Named judge prompt mode.", + ) + parser.add_argument( + "--top_models", + type=int, + default=20, + help="[meta-eval] Number of top models by battle count to include.", + ) + parser.add_argument( + "--battles_per_model", + type=int, + default=50, + help="[meta-eval] Battles sampled per top model.", + ) + parser.add_argument( + "--batch_size", + type=int, + default=50, + help="[meta-eval] Annotation batch size.", + ) + parser.add_argument( + "--elo_gap_battles", + nargs="+", + type=int, + default=[10, 20, 30, 40, 50], + help="[meta-eval] Battle counts for ELO-gap analysis.", + ) + parser.add_argument( + "--elo_gap_seeds", + type=int, + default=10, + help="[meta-eval] Random seeds for ELO-gap subsampling.", + ) + parser.add_argument( + "--include_human_ties", + action="store_true", + help="[meta-eval] Include human-labeled ties in agreement metrics.", + ) + + +def build_meta_eval_args(args: argparse.Namespace) -> CliMetaEvalArgs: + if args.model_A is not None or args.model_B is not None: + raise SystemExit( + "--model_A/--model_B are not used for meta-eval; only --judge_model is required." + ) + if args.n_instructions is not None: + raise SystemExit( + "--n_instructions is not used for meta-eval; use --top_models and " + "--battles_per_model to control the sample." + ) + if args.max_out_tokens_models != DEFAULT_MAX_OUT_TOKENS_MODELS: + raise SystemExit( + "--max_out_tokens_models is not used for meta-eval because no model " + "completions are generated." + ) + return CliMetaEvalArgs( + reference_arena=args.reference_arena, + prompt_mode=args.prompt_mode, + top_models=args.top_models, + battles_per_model=args.battles_per_model, + batch_size=args.batch_size, + languages=args.languages, + n_bootstraps=args.n_bootstraps, + seed=args.seed, + elo_gap_battles=args.elo_gap_battles, + elo_gap_seeds=args.elo_gap_seeds, + exclude_human_ties=not args.include_human_ties, + judge_model=args.judge_model, + n_instructions=args.n_instructions, + provide_explanation=args.provide_explanation, + swap_mode=args.swap_mode, + ignore_cache=args.ignore_cache, + truncate_all_input_chars=args.truncate_all_input_chars, + max_out_tokens_models=args.max_out_tokens_models, + max_out_tokens_judge=args.max_out_tokens_judge, + max_model_len=args.max_model_len, + chat_template=args.chat_template, + result_folder=args.result_folder, + engine_kwargs=parse_engine_kwargs(args.engine_kwargs), + verbosity=resolve_verbosity(args), + log_file=args.log_file, + no_log_file=args.no_log_file, + ) diff --git a/judgearena/meta_eval/runner.py b/judgearena/meta_eval/runner.py new file mode 100644 index 0000000..fb5b55c --- /dev/null +++ b/judgearena/meta_eval/runner.py @@ -0,0 +1,259 @@ +"""Main entrypoint for judge meta-evaluation.""" + +from __future__ import annotations + +import json +from dataclasses import asdict +from datetime import UTC, datetime +from pathlib import Path + +import pandas as pd + +from judgearena.log import attach_file_handler, get_logger, make_run_log_path +from judgearena.meta_eval.annotate import annotate_sample +from judgearena.meta_eval.cli_args import CliMetaEvalArgs +from judgearena.meta_eval.metrics import ( + compute_agreement_metrics, + compute_elo_gap_summary, + format_metric, + summarize_language_splits, +) +from judgearena.meta_eval.prompts import resolve_prompt_mode +from judgearena.meta_eval.sampling import ( + MetaEvalSamplingError, + load_reference_arena_battles, + sample_battles_per_model, + select_top_models, +) +from judgearena.repro import _to_jsonable, write_run_metadata +from judgearena.utils import make_model + +logger = get_logger(__name__) + + +def _result_folder_name(args: CliMetaEvalArgs, started_at: datetime) -> str: + name = ( + f"meta-eval-{args.reference_arena}-{args.prompt_mode}-" + f"{args.judge_model}-{args.swap_mode}" + ) + return name.replace("/", "_") + f"-{started_at.strftime('%Y%m%d_%H%M%S')}" + + +def _build_summary_csv( + language_summary: dict[str, dict[str, str | int]], +) -> pd.DataFrame: + rows = [] + for split, metrics in language_summary.items(): + rows.append({"split": split, **metrics}) + return pd.DataFrame(rows) + + +def _agreement_view( + metrics: dict[str, float | int], + *, + exclude_human_ties: bool, +) -> dict[str, float | int | str]: + suffix = "_nt" if exclude_human_ties else "" + n_key = "n_nt" if exclude_human_ties else "n" + accuracy = float(metrics[f"accuracy{suffix}"]) + accuracy_se = float(metrics[f"acc_se{suffix}"]) + kappa = float(metrics[f"kappa{suffix}"]) + kappa_se = float(metrics[f"kappa_se{suffix}"]) + return { + "n": int(metrics[n_key]), + "accuracy": accuracy, + "accuracy_se": accuracy_se, + "kappa": kappa, + "kappa_se": kappa_se, + "accuracy_formatted": format_metric(accuracy, accuracy_se, digits=3), + "kappa_formatted": format_metric(kappa, kappa_se, digits=3), + } + + +def _annotation_telemetry( + df_ann: pd.DataFrame, + *, + swap_mode: str, +) -> dict[str, object]: + costs = pd.to_numeric(df_ann["cost_usd"], errors="coerce").dropna() + sources = df_ann["cost_source"].dropna() + total_cost = float(costs.sum()) if not costs.empty else None + cost_per_1k = float(costs.mean() * 1000) if not costs.empty else None + if costs.empty: + logger.warning( + "OpenRouter reference pricing is unavailable; cost fields will be null." + ) + + return { + "judge_passes_per_battle": 2 if swap_mode == "both" else 1, + "judgement_count": len(df_ann), + "estimated_input_tokens": int( + pd.to_numeric(df_ann["estimated_input_tokens"], errors="coerce") + .fillna(0) + .sum() + ), + "estimated_output_tokens": int( + pd.to_numeric(df_ann["estimated_output_tokens"], errors="coerce") + .fillna(0) + .sum() + ), + "token_count_source": "estimated_chars_div_4", + "total_cost_usd": total_cost, + "cost_per_1k_judgements_usd": cost_per_1k, + "cost_source_counts": { + str(source): int(count) for source, count in sources.value_counts().items() + }, + } + + +def _compute_results( + *, + args: CliMetaEvalArgs, + top_models: list[str], + df_top: pd.DataFrame, + df_sample: pd.DataFrame, + df_ann: pd.DataFrame, +) -> dict: + agreement_metrics = compute_agreement_metrics( + df_ann["winner"].tolist(), + df_ann["winner_llm"].tolist(), + n_bootstraps=args.n_bootstraps, + seed=args.seed, + ) + primary_view = "no_human_ties" if args.exclude_human_ties else "all" + agreement = { + "primary_view": primary_view, + "all": _agreement_view( + agreement_metrics, + exclude_human_ties=False, + ), + "no_human_ties": _agreement_view( + agreement_metrics, + exclude_human_ties=True, + ), + } + ranking_annotations = df_ann[df_ann["orientation"] == "forward"].copy() + language_summary = summarize_language_splits( + ranking_annotations, + exclude_human_ties=args.exclude_human_ties, + n_bootstraps=args.n_bootstraps, + seed=args.seed, + ) + elo_gap_all = compute_elo_gap_summary( + df_top, + ranking_annotations, + top_models, + n_battles_list=args.elo_gap_battles or [], + n_seeds=args.elo_gap_seeds, + seed=args.seed, + exclude_ties=False, + ) + elo_gap_no_tie = compute_elo_gap_summary( + df_top, + ranking_annotations, + top_models, + n_battles_list=args.elo_gap_battles or [], + n_seeds=args.elo_gap_seeds, + seed=args.seed + 1000, + exclude_ties=True, + ) + return { + "task": "meta-eval", + "reference_arena": args.reference_arena, + "prompt_mode": args.prompt_mode, + "judge_model": args.judge_model, + "top_models": top_models, + "sample_size": len(df_sample), + "ranking_annotation_count": len(ranking_annotations), + "agreement": agreement, + "language_summary": language_summary, + "elo_gap_all": elo_gap_all.to_dict(orient="records"), + "elo_gap_exclude_ties": elo_gap_no_tie.to_dict(orient="records"), + **_annotation_telemetry(df_ann, swap_mode=args.swap_mode), + } + + +def main(args: CliMetaEvalArgs) -> dict: + started_at = datetime.now(UTC) + res_folder = Path(args.result_folder) / _result_folder_name(args, started_at) + res_folder.mkdir(parents=True, exist_ok=True) + + if not args.no_log_file: + attach_file_handler(make_run_log_path(res_folder)) + + with open(res_folder / "args.json", "w", encoding="utf-8") as handle: + json.dump(asdict(args), handle, indent=2) + + prompt_spec = resolve_prompt_mode( + args.prompt_mode, + provide_explanation=args.provide_explanation, + ) + judge_chat_model = make_model( + model=args.judge_model, + max_tokens=args.max_out_tokens_judge, + max_model_len=args.max_model_len, + chat_template=args.chat_template, + **args.engine_kwargs, + ) + + logger.info("Loading reference arena %s", args.reference_arena) + df = load_reference_arena_battles( + args.reference_arena, + languages=args.languages, + ) + top_models, df_top = select_top_models(df, top_models=args.top_models) + df_sample = sample_battles_per_model( + df_top, + top_models, + battles_per_model=args.battles_per_model, + seed=args.seed, + ) + + logger.info( + "Meta-eval sample: %d battles among top %d models", + len(df_sample), + len(top_models), + ) + + df_ann = annotate_sample( + df_sample, + args, + judge_chat_model=judge_chat_model, + prompt_spec=prompt_spec, + ) + df_ann.to_parquet(res_folder / "annotations.parquet", index=False) + + results = _compute_results( + args=args, + top_models=top_models, + df_top=df_top, + df_sample=df_sample, + df_ann=df_ann, + ) + + with open(res_folder / "results.json", "w", encoding="utf-8") as handle: + json.dump(_to_jsonable(results), handle, indent=2, allow_nan=False) + + summary_csv = _build_summary_csv(results["language_summary"]) + summary_csv.to_csv(res_folder / "summary.csv", index=False) + + write_run_metadata( + output_dir=res_folder, + entrypoint="judgearena.meta_eval.runner", + run=asdict(args), + results=results, + input_payloads={"question_id": df_sample["question_id"].astype(str).tolist()}, + judge_system_prompt=prompt_spec.system_prompt, + judge_user_prompt_template=prompt_spec.user_prompt_template, + started_at_utc=started_at, + ) + + logger.info("Meta-eval results saved to %s", res_folder) + return results + + +def run_or_exit(args: CliMetaEvalArgs) -> dict: + try: + return main(args) + except MetaEvalSamplingError as exc: + raise SystemExit(str(exc)) from exc From 742770096070c8b41fe18d7cdf5a4c422ce6a752 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Sun, 19 Jul 2026 00:26:36 +0200 Subject: [PATCH 05/13] test: cover meta-eval pipeline and packaged prompts Add unit/integration coverage for meta-eval behavior and ensure prompt resources remain loadable from wheel and sdist installs. --- tests/test_meta_eval.py | 649 ++++++++++++++++++++++++++++++++++++++++ tests/test_smoke.py | 13 + 2 files changed, 662 insertions(+) create mode 100644 tests/test_meta_eval.py diff --git a/tests/test_meta_eval.py b/tests/test_meta_eval.py new file mode 100644 index 0000000..d85ec77 --- /dev/null +++ b/tests/test_meta_eval.py @@ -0,0 +1,649 @@ +"""Tests for judge meta-evaluation.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pandas as pd +import pytest + +import judgearena.meta_eval.annotate as meta_annotate +import judgearena.meta_eval.cost as meta_cost +import judgearena.meta_eval.runner as meta_eval_runner +import judgearena.meta_eval.sampling as meta_sampling +from judgearena import cli as cli_module +from judgearena.arenas_utils import extract_turn_text +from judgearena.evaluate import JudgeAnnotation, PairScore +from judgearena.meta_eval.cache import AnnotationCache, AnnotationEntry, AnnotationKey +from judgearena.meta_eval.cli_args import CliMetaEvalArgs +from judgearena.meta_eval.metrics import ( + compute_agreement_metrics, + compute_elo_gap_summary, +) +from judgearena.meta_eval.parsers import ( + META_EVAL_PAIRSCORE_TEMPERATURE, + invert_winner, + parse_alpaca_eval_winner, + parse_arena_hard_winner, + parse_pairscore_winner, + parse_pref, + parse_winner, +) +from judgearena.meta_eval.prompts import PromptModeSpec, resolve_prompt_mode +from judgearena.meta_eval.sampling import ( + MetaEvalSamplingError, + load_reference_arena_battles, + sample_battles_per_model, + select_top_models, +) +from judgearena.repro import METADATA_FILENAME + + +def _conversation( + instruction: str, + answer_a: str, + answer_b: str, + *, + structured: bool = False, +) -> tuple[list[dict], list[dict]]: + if structured: + user_content = [{"type": "text", "text": instruction, "image": None}] + assistant_content = [{"type": "text", "text": answer_a}] + assistant_b_content = [{"type": "text", "text": answer_b}] + else: + user_content = instruction + assistant_content = answer_a + assistant_b_content = answer_b + conv_a = [ + {"role": "user", "content": user_content}, + {"role": "assistant", "content": assistant_content}, + ] + conv_b = [ + {"role": "user", "content": user_content}, + {"role": "assistant", "content": assistant_b_content}, + ] + return conv_a, conv_b + + +@pytest.fixture +def synthetic_arena_df() -> pd.DataFrame: + rows = [] + models = [f"model-{idx}" for idx in range(5)] + for idx in range(120): + model_a = models[idx % len(models)] + model_b = models[(idx + 1) % len(models)] + winner = ["model_a", "model_b", "tie"][idx % 3] + lang = "en" if idx % 2 == 0 else "es" + conv_a, conv_b = _conversation( + f"Question {idx}", + f"Answer A {idx}", + f"Answer B {idx}", + structured=idx % 4 == 0, + ) + rows.append( + { + "question_id": f"q-{idx}", + "tstamp": 1_700_000_000 + idx, + "model_a": model_a, + "model_b": model_b, + "winner": winner, + "conversation_a": conv_a, + "conversation_b": conv_b, + "benchmark": "LMArena-140k", + "lang": lang, + } + ) + return pd.DataFrame(rows) + + +@pytest.fixture +def meta_args(tmp_path: Path) -> CliMetaEvalArgs: + return CliMetaEvalArgs( + reference_arena="LMArena-140k", + prompt_mode="standard", + top_models=3, + battles_per_model=4, + batch_size=8, + languages=["en", "es"], + n_bootstraps=20, + seed=7, + judge_model="Dummy/judge", + result_folder=str(tmp_path / "results"), + no_log_file=True, + ) + + +def _judge_annotations(*, instructions, completions_A, completions_B, **_kwargs): + return [ + JudgeAnnotation( + instruction=instruction, + completion_A=completion_a, + completion_B=completion_b, + judge_completion="score_A: 9\nscore_B: 1", + judge_input="judge prompt", + ) + for instruction, completion_a, completion_b in zip( + instructions, + completions_A, + completions_B, + strict=True, + ) + ] + + +@pytest.fixture +def stub_meta_eval_runner(monkeypatch, synthetic_arena_df): + monkeypatch.setattr( + meta_eval_runner, + "load_reference_arena_battles", + lambda reference_arena, languages=None: synthetic_arena_df, + ) + monkeypatch.setattr(meta_eval_runner, "make_model", lambda **_kwargs: object()) + + +def test_cli_meta_eval_dispatch(monkeypatch): + captured: dict[str, object] = {} + + def fake_main(args: CliMetaEvalArgs) -> None: + captured["args"] = args + + monkeypatch.setattr(cli_module, "main_meta_eval", fake_main) + cli_module.cli( + [ + "--task", + "meta-eval", + "--judge", + "Dummy/J", + "--reference_arena", + "LMArena-140k", + "--prompt_mode", + "arena-hard", + "--languages", + "en", + "es", + "--top_models", + "5", + "--battles_per_model", + "10", + ] + ) + args: CliMetaEvalArgs = captured["args"] + assert args.reference_arena == "LMArena-140k" + assert args.prompt_mode == "arena-hard" + assert args.languages == ["en", "es"] + assert args.top_models == 5 + + +@pytest.mark.parametrize( + ("flag", "value", "message"), + [ + ("--model_A", "Dummy/A", "--model_A/--model_B are not used"), + ("--n_instructions", "10", "--n_instructions is not used"), + ( + "--max_out_tokens_models", + "1024", + "--max_out_tokens_models is not used", + ), + ], +) +def test_cli_meta_eval_rejects_irrelevant_flags( + monkeypatch, + flag, + value, + message, +): + monkeypatch.setattr(cli_module, "main_meta_eval", lambda _args: None) + with pytest.raises(SystemExit, match=message): + cli_module.cli( + [ + "--task", + "meta-eval", + "--judge", + "Dummy/J", + flag, + value, + ] + ) + + +def test_extract_text_structured_content(): + conv_a, _ = _conversation("hello", "A", "B", structured=True) + assert extract_turn_text(conv_a[0]) == "hello" + assert ( + extract_turn_text( + {"content": [{"type": "text", "text": "part one"}]}, + ) + == "part one" + ) + + +def test_language_filter_empty_raises(monkeypatch, synthetic_arena_df): + monkeypatch.setattr( + meta_sampling, + "load_arena_dataframe", + lambda arena: synthetic_arena_df, + ) + with pytest.raises(MetaEvalSamplingError, match="languages"): + load_reference_arena_battles("LMArena-140k", languages=["zz"]) + + +def test_sampling_is_deterministic(synthetic_arena_df): + df = synthetic_arena_df.copy() + top_models, df_top = select_top_models(df, top_models=3) + first = sample_battles_per_model( + df_top, + top_models, + battles_per_model=3, + seed=11, + ) + second = sample_battles_per_model( + df_top, + top_models, + battles_per_model=3, + seed=11, + ) + assert first["question_id"].tolist() == second["question_id"].tolist() + assert len(first) == 9 + + +@pytest.mark.parametrize( + ("completion", "expected"), + [ + ("My verdict is [[A>>B]]", "model_a"), + ("Final: [[B>A]]", "model_b"), + ("No signal", "tie"), + ], +) +def test_parse_arena_hard_winner(completion, expected): + assert parse_arena_hard_winner(completion) == expected + + +def test_parse_alpaca_eval_winner(): + completion = ( + '```json\n{"ordered_models": [{"model": "m", "rank": 2}, ' + '{"model": "M", "rank": 1}]}\n```' + ) + assert parse_alpaca_eval_winner(completion) == "model_b" + + +def test_pairscore_temperature_and_tie(): + completion = "score_A: 5\nscore_B: 5" + benchmark = PairScore() + benchmark.temperature = 0.3 + meta = PairScore() + meta.temperature = META_EVAL_PAIRSCORE_TEMPERATURE + assert benchmark.parse_model_raw(completion) == 0.5 + assert meta.parse_model_raw(completion) == 0.5 + assert parse_pairscore_winner(completion, temperature=0.5) == "tie" + assert ( + parse_pairscore_winner( + "score_A: 10\nscore_B: 0", + temperature=0.5, + ) + == "model_a" + ) + + +@pytest.mark.parametrize( + ("mode", "completion", "expected"), + [ + ("standard", "score_A: 9\nscore_B: 1", "model_a"), + ("arena-hard", "[[A=B]]", "tie"), + ( + "alpaca-eval", + '{"ordered_models": [{"model": "m", "rank": 1}]}', + "model_a", + ), + ("alpaca-eval-pair-score", "score_A: 9\nscore_B: 1", "model_a"), + ], +) +def test_parse_winner_modes(mode, completion, expected): + assert parse_winner(completion, mode) == expected + + +@pytest.mark.parametrize( + "mode", + ["arena-hard", "alpaca-eval", "alpaca-eval-pair-score"], +) +def test_file_backed_prompt_modes_load_packaged_resources(mode): + prompt = resolve_prompt_mode(mode) + assert prompt.system_prompt + assert prompt.user_prompt_template + + +def test_parse_pref_continuous_semantics(): + assert parse_pref("score_A: 10\nscore_B: 0", "standard") < 0.5 + assert parse_pref("score_A: 0\nscore_B: 10", "standard") > 0.5 + assert parse_pref("[[A=B]]", "arena-hard") == 0.5 + + +def test_swap_inversion(): + assert invert_winner("model_a") == "model_b" + assert invert_winner("tie") == "tie" + + +def test_agreement_metrics_on_fixture(): + human = ["model_a", "model_b", "tie", "model_a"] + llm = ["model_a", "model_a", "tie", "model_b"] + metrics = compute_agreement_metrics( + human, + llm, + n_bootstraps=10, + seed=0, + ) + assert metrics["n"] == 4 + assert metrics["accuracy"] == 0.5 + assert metrics["n_nt"] == 3 + + +def _cache_key(**overrides) -> AnnotationKey: + values = { + "benchmark": "LMArena-140k", + "instruction_id": "q-1", + "model_a": "model-a", + "model_b": "model-b", + "judge": "Dummy/judge", + } + values.update(overrides) + return AnnotationKey(**values) + + +def _cache_entry(completion: str = "score_A: 9\nscore_B: 1", **overrides): + key = _cache_key(**overrides) + return AnnotationEntry( + **key.__dict__, + judge_input="judge prompt", + judge_completion=completion, + ) + + +def test_annotation_cache_persists_and_preserves_batch_order(tmp_path): + db_dir = tmp_path / "db" + first = AnnotationCache(db_dir) + first.batch_put( + [ + _cache_entry(instruction_id="q-2", completion="second"), + _cache_entry(instruction_id="q-1", completion="first"), + ] + ) + first.close() + + second = AnnotationCache(db_dir) + entries = second.batch_get_annotations( + [ + _cache_key(instruction_id="q-1"), + _cache_key(instruction_id="q-2"), + _cache_key(instruction_id="missing"), + ] + ) + assert [entry.judge_completion if entry else None for entry in entries] == [ + "first", + "second", + None, + ] + second.close() + + +def test_annotation_cache_distinguishes_prompt_mode_and_model_order(tmp_path): + cache = AnnotationCache(tmp_path / "db") + cache.batch_put( + [ + _cache_entry(judge="Dummy/judge::arena-hard"), + _cache_entry(model_a="model-b", model_b="model-a"), + ] + ) + entries = cache.batch_get_annotations( + [ + _cache_key(judge="Dummy/judge"), + _cache_key(judge="Dummy/judge::arena-hard"), + _cache_key(model_a="model-b", model_b="model-a"), + ] + ) + assert entries[0] is None + assert all(entry is not None for entry in entries[1:]) + cache.close() + + +def test_annotate_sample_uses_cache_and_inverts_swapped_pass( + monkeypatch, + synthetic_arena_df, + meta_args, + tmp_path, +): + calls = {"count": 0} + + def fake_annotate_battles(**kwargs): + calls["count"] += 1 + return _judge_annotations(**kwargs) + + monkeypatch.setattr(meta_annotate, "annotate_battles", fake_annotate_battles) + meta_args.swap_mode = "both" + sample = synthetic_arena_df.iloc[:1] + cache = AnnotationCache(tmp_path / "db") + prompt_spec = PromptModeSpec( + name="standard", + system_prompt="system", + user_prompt_template="user", + ) + + annotations = meta_annotate.annotate_sample( + sample, + meta_args, + judge_chat_model=object(), + prompt_spec=prompt_spec, + annotation_cache=cache, + ) + assert len(annotations) == 2 + assert annotations["orientation"].tolist() == ["forward", "swapped"] + assert annotations["winner"].tolist() == [sample.iloc[0]["winner"]] * 2 + assert annotations["winner_llm"].tolist() == ["model_a", "model_b"] + assert annotations["model_a"].tolist() == [sample.iloc[0]["model_a"]] * 2 + assert annotations["presented_model_a"].tolist() == [ + sample.iloc[0]["model_a"], + sample.iloc[0]["model_b"], + ] + assert annotations["completion_a"].nunique() == 1 + assert ( + annotations.iloc[1]["presented_completion_a"] + == annotations.iloc[0]["completion_b"] + ) + assert calls["count"] == 2 + + meta_annotate.annotate_sample( + sample, + meta_args, + judge_chat_model=object(), + prompt_spec=prompt_spec, + annotation_cache=cache, + ) + assert calls["count"] == 2 + + meta_args.ignore_cache = True + meta_annotate.annotate_sample( + sample, + meta_args, + judge_chat_model=object(), + prompt_spec=prompt_spec, + annotation_cache=cache, + ) + assert calls["count"] == 4 + cache.close() + + +def test_cost_uses_offline_reference_pricing(monkeypatch, tmp_path): + pricing_file = tmp_path / "openrouter_pricing.json" + pricing_file.write_text( + json.dumps({"provider/model": [1.0, 2.0]}), + encoding="utf-8", + ) + monkeypatch.setattr(meta_cost, "_PRICING_CACHE_FILE", pricing_file) + meta_cost._openrouter_pricing_cache.clear() + + cost, source = meta_cost.estimate_annotation_cost_usd( + judge_input="a" * 40, + judge_completion="b" * 20, + judge_model="OpenRouter/provider/model", + ) + assert cost == pytest.approx((10 * 1.0 + 5 * 2.0) / 1e6) + assert source == "estimated" + meta_cost._openrouter_pricing_cache.clear() + + +def test_swapped_pass_telemetry_counts_both_judgements(): + annotations = pd.DataFrame( + { + "cost_usd": [0.1, 0.2], + "cost_source": ["estimated", "estimated"], + "estimated_input_tokens": [10, 11], + "estimated_output_tokens": [5, 6], + } + ) + telemetry = meta_eval_runner._annotation_telemetry( + annotations, + swap_mode="both", + ) + assert telemetry["judgement_count"] == 2 + assert telemetry["total_cost_usd"] == pytest.approx(0.3) + assert telemetry["estimated_input_tokens"] == 21 + assert telemetry["estimated_output_tokens"] == 11 + + +def test_degenerate_agreement_metrics_do_not_warn(recwarn): + metrics = compute_agreement_metrics( + ["model_a"] * 4, + ["model_a"] * 4, + n_bootstraps=10, + seed=0, + ) + assert metrics["accuracy"] == 1.0 + assert pd.isna(metrics["kappa"]) + assert not recwarn.list + + +def test_integration_meta_eval_artifacts( + monkeypatch, + meta_args: CliMetaEvalArgs, + stub_meta_eval_runner, +): + def fake_annotate(df_sample, args, **kwargs): + frame = df_sample[ + ["question_id", "model_a", "model_b", "winner", "lang", "benchmark"] + ].copy() + winners = frame["winner"].where(frame["winner"] == "model_a", "model_b") + return frame.assign( + orientation="forward", + instruction="instr", + completion_a="A", + completion_b="B", + judge_input="prompt", + judge_completion="score_A: 9\nscore_B: 1", + estimated_input_tokens=2, + estimated_output_tokens=5, + cost_usd=0.001, + cost_source="estimated", + winner_llm=winners, + pref_llm=winners.map({"model_a": 0.0, "model_b": 1.0}), + ) + + monkeypatch.setattr(meta_eval_runner, "annotate_sample", fake_annotate) + + results = meta_eval_runner.main(meta_args) + output_dirs = list(Path(meta_args.result_folder).glob("meta-eval-*")) + assert len(output_dirs) == 1 + out = output_dirs[0] + assert (out / "args.json").exists() + assert (out / "annotations.parquet").exists() + assert (out / "results.json").exists() + assert (out / "summary.csv").exists() + assert (out / METADATA_FILENAME).exists() + metadata = json.loads((out / METADATA_FILENAME).read_text(encoding="utf-8")) + assert metadata["entrypoint"] == "judgearena.meta_eval.runner" + assert results["agreement"]["primary_view"] == "no_human_ties" + assert results["agreement"]["all"]["n"] > results["agreement"]["no_human_ties"]["n"] + assert results["judgement_count"] == results["sample_size"] + assert results["total_cost_usd"] == pytest.approx(results["sample_size"] * 0.001) + assert "English" in results["language_summary"] + + +def test_swap_mode_both_artifact_reproduces_overall_agreement( + monkeypatch, + meta_args, + stub_meta_eval_runner, + tmp_path, +): + cache_class = AnnotationCache + monkeypatch.setattr( + meta_annotate, + "AnnotationCache", + lambda: cache_class(tmp_path / "cache"), + ) + monkeypatch.setattr(meta_annotate, "annotate_battles", _judge_annotations) + meta_args.swap_mode = "both" + meta_args.battles_per_model = 2 + meta_args.elo_gap_battles = [1] + meta_args.elo_gap_seeds = 1 + + results = meta_eval_runner.main(meta_args) + output_dir = next(Path(meta_args.result_folder).glob("meta-eval-*")) + annotations = pd.read_parquet(output_dir / "annotations.parquet") + forward = annotations[annotations["orientation"] == "forward"] + swapped = annotations[annotations["orientation"] == "swapped"] + + assert len(annotations) == 2 * results["sample_size"] + assert len(forward) == len(swapped) == results["sample_size"] + assert results["judgement_count"] == len(annotations) + assert results["ranking_annotation_count"] == len(forward) + assert (swapped["model_a"].to_numpy() == swapped["presented_model_b"]).all() + assert (swapped["model_b"].to_numpy() == swapped["presented_model_a"]).all() + assert (swapped["completion_a"].to_numpy() == forward["completion_a"]).all() + assert (swapped["completion_b"].to_numpy() == forward["completion_b"]).all() + + recomputed = compute_agreement_metrics( + annotations["winner"].tolist(), + annotations["winner_llm"].tolist(), + n_bootstraps=meta_args.n_bootstraps, + seed=meta_args.seed, + ) + assert results["agreement"]["all"]["accuracy"] == recomputed["accuracy"] + assert results["agreement"]["all"]["kappa"] == recomputed["kappa"] + + +def test_elo_gap_summary_runs(synthetic_arena_df): + top_models, df_top = select_top_models(synthetic_arena_df, top_models=3) + df_sample = sample_battles_per_model( + df_top, + top_models, + battles_per_model=3, + seed=1, + ) + df_ann = df_sample.copy() + df_ann["winner_llm"] = df_ann["winner"] + df_ann["pref_llm"] = 0.5 + summary = compute_elo_gap_summary( + df_top, + df_ann, + top_models, + n_battles_list=[2], + n_seeds=2, + seed=0, + exclude_ties=False, + ) + assert not summary.empty + + +def test_empty_language_subset_reports_na( + monkeypatch, + synthetic_arena_df, + meta_args: CliMetaEvalArgs, +): + meta_args.languages = ["en"] + + def _raise_empty(_reference_arena, languages=None): + raise MetaEvalSamplingError( + "No battles remain after filtering to languages: en." + ) + + monkeypatch.setattr(meta_eval_runner, "load_reference_arena_battles", _raise_empty) + with pytest.raises(SystemExit): + meta_eval_runner.run_or_exit(meta_args) diff --git a/tests/test_smoke.py b/tests/test_smoke.py index 4b81cfc..97ba31a 100644 --- a/tests/test_smoke.py +++ b/tests/test_smoke.py @@ -10,6 +10,14 @@ from judgearena.criteria.defaults import CRITERIA_BY_NAME +META_EVAL_PROMPT_RESOURCES = ( + "arena_hard_system.txt", + "arena_hard_user.txt", + "alpaca_eval_system.txt", + "alpaca_eval_user.txt", + "alpaca_eval_pair_score_user.txt", +) + def _assert_non_empty_text_resource(package: str, relative_path: str) -> None: content = files(package).joinpath(relative_path).read_text() @@ -30,6 +38,11 @@ def main() -> None: _assert_non_empty_text_resource("judgearena.prompts", "prompt.txt") _assert_non_empty_text_resource("judgearena.prompts", "system-prompt.txt") _assert_non_empty_text_resource("judgearena.criteria", "data/default.yaml") + for filename in META_EVAL_PROMPT_RESOURCES: + _assert_non_empty_text_resource( + "judgearena.meta_eval", + f"prompts/{filename}", + ) print("✅ All integrity checks passed: Imports, Criteria, and Resources are valid.") From b33243a2e39b4d3c628dad2f2f0de7d4d8d6fe63 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Sun, 19 Jul 2026 00:26:39 +0200 Subject: [PATCH 06/13] docs: document meta-eval task usage and artifacts Describe the new meta-eval CLI path, prompt modes, swap-mode semantics, and temporary SQLite cache limitations for operators. --- README.md | 54 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/README.md b/README.md index 306443d..4830b80 100644 --- a/README.md +++ b/README.md @@ -207,6 +207,7 @@ Task names follow [LMHarness](https://github.com/EleutherAI/lm-evaluation-harnes | `m-arena-hard-EU` | All EU languages combined | | `mt-bench` | Multi-turn benchmark with FastChat-compatible pairwise judging | | `fluency-{lang}` | Fluency evaluation for pretrained models (`finnish`, `french`, `german`, `spanish`, `swedish`) | +| `meta-eval` | Judge meta-evaluation against human-labeled arena battles (not generator ELO) | For Arena-Hard, JudgeArena resolves baseline metadata by task version: - `arena-hard-v0.1`: `gpt-4-0314` @@ -221,6 +222,59 @@ For Arena-Hard, JudgeArena resolves baseline metadata by task version: | `elo-lmarena` | Union of all `LMArena-*` variants | | `elo-comparia` | Battles sampled from the ComparIA arena | +## Judge meta-evaluation (`meta-eval`) + +`meta-eval` measures how well an LLM judge agrees with **human-labeled arena battles**. +It does **not** generate model completions and does **not** estimate a generator model's ELO rating. +Use `elo-*` tasks for generator ELO estimation instead. + +The task samples battles from a reference arena (default `LMArena-140k`), asks the judge to label fixed human battles, and reports accuracy, Cohen's kappa, Bradley-Terry ranking agreement, Spearman correlation, ELO MAE, bootstrap uncertainty, and ELO-gap analyses. + +```bash +judgearena \ + --task meta-eval \ + --judge_model OpenRouter/deepseek/deepseek-v3.2 \ + --reference_arena LMArena-140k \ + --prompt_mode standard \ + --languages en es \ + --top_models 20 \ + --battles_per_model 50 \ + --n_bootstraps 1000 \ + --seed 0 +``` + +### Prompt modes + +Only named prompt modes are supported (no custom prompt files): + +| `--prompt_mode` | Description | +|---------------------------|--------------------------------------------------| +| `standard` | Default PairScore judge prompt | +| `arena-hard` | Arena-Hard Likert verdict parsing | +| `alpaca-eval` | Alpaca-Eval JSON ordering prompt | +| `alpaca-eval-pair-score` | Alpaca-Eval prompt with PairScore output | + +PairScore meta-eval uses temperature `0.5` (paper setting). The standard generate+judge benchmark path keeps PairScore temperature `0.3`. + +Language filters use ISO 639-1 codes such as `en es fr`. + +With `--swap_mode both`, overall accuracy and kappa use both judge orderings +(the reversed verdict is inverted), while language, ranking, and ELO-gap analyses +retain one forward-order row per sampled battle to match the reference +meta-evaluation methodology. `annotations.parquet` stores one row per judge pass +with an `orientation` column; swapped winners and preferences are normalized +back to the original model ordering. Cost totals include both passes. + +Results are written under `[result_folder]/meta-eval-*` as `args.json`, +`annotations.parquet`, `results.json`, `summary.csv`, logs, and +`run-metadata.v1.json`. Judge annotations are cached per battle in +`$JUDGEARENA_DATA/cache/db/{benchmark}/{judge}.db`. This temporary SQLite WAL +cache should only be used by a single host; do not share the same database +concurrently across NFS-mounted compute nodes. Annotation artifacts include +character-based token estimates. Equivalent OpenRouter cost is reported only +when `[data_root]/cache/openrouter_pricing.json` already contains the judge +model; meta-eval never fetches pricing from compute nodes. + ## 📈 Estimating ELO Ratings JudgeArena can estimate the ELO rating of a model by running it against opponents sampled from a human preference arena (`LMArena-100k`, `LMArena-140k`, or `ComparIA`). From a79d6b922cc70faed0acef452b4ce2e90ca332ca Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Mon, 20 Jul 2026 14:22:23 +0200 Subject: [PATCH 07/13] fix: serialize judge prompts before caching Store rendered prompt text rather than LangChain prompt objects so SQLite-backed meta-eval runs complete after annotation. Includes-AI-Code: true --- judgearena/evaluate.py | 2 +- tests/test_meta_eval.py | 19 ++++++++++++++++++- 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/judgearena/evaluate.py b/judgearena/evaluate.py index 7994c5f..d2aca18 100644 --- a/judgearena/evaluate.py +++ b/judgearena/evaluate.py @@ -262,7 +262,7 @@ def annotate_battles( ): annotations.append( JudgeAnnotation( - judge_input=judge_input, + judge_input=judge_input.to_string(), judge_completion=judge_completion, instruction=instruction, completion_A=completion_A, diff --git a/tests/test_meta_eval.py b/tests/test_meta_eval.py index 37855f5..a287964 100644 --- a/tests/test_meta_eval.py +++ b/tests/test_meta_eval.py @@ -14,7 +14,7 @@ import judgearena.meta_eval.sampling as meta_sampling from judgearena import cli as cli_module from judgearena.arenas_utils import extract_turn_text -from judgearena.evaluate import JudgeAnnotation, PairScore +from judgearena.evaluate import JudgeAnnotation, PairScore, annotate_battles from judgearena.meta_eval.cache import AnnotationCache, AnnotationEntry, AnnotationKey from judgearena.meta_eval.cli_args import CliMetaEvalArgs from judgearena.meta_eval.metrics import ( @@ -221,6 +221,23 @@ def test_extract_text_structured_content(): ) +def test_annotate_battles_serializes_judge_input(): + class FakeJudge: + @staticmethod + def batch(*, inputs, **_kwargs): + return ["score_A: 9\nscore_B: 1"] * len(inputs) + + annotation = annotate_battles( + judge_chat_model=FakeJudge(), + instructions=["Question"], + completions_A=["Answer A"], + completions_B=["Answer B"], + )[0] + + assert isinstance(annotation.judge_input, str) + assert "Question" in annotation.judge_input + + def test_language_filter_empty_raises(monkeypatch, synthetic_arena_df): monkeypatch.setattr( meta_sampling, From 31229a29839e96f6422a78049f99e2353ae29fb6 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:52:26 +0200 Subject: [PATCH 08/13] feat(cache): add unified inference storage and sync Introduce content-addressed SQLite cells with atomic metadata associations and deterministic Hugging Face merge semantics. Includes-AI-Code: true --- judgearena/store_sqlite.py | 422 ++++++++++++++++++++++++ judgearena/store_sync.py | 646 +++++++++++++++++++++++++++++++++++++ tests/conftest.py | 125 +++++++ tests/test_store_sqlite.py | 258 +++++++++++++++ tests/test_store_sync.py | 564 ++++++++++++++++++++++++++++++++ 5 files changed, 2015 insertions(+) create mode 100644 judgearena/store_sqlite.py create mode 100644 judgearena/store_sync.py create mode 100644 tests/conftest.py create mode 100644 tests/test_store_sqlite.py create mode 100644 tests/test_store_sync.py diff --git a/judgearena/store_sqlite.py b/judgearena/store_sqlite.py new file mode 100644 index 0000000..56c4e61 --- /dev/null +++ b/judgearena/store_sqlite.py @@ -0,0 +1,422 @@ +"""SQLite-backed unified inference cache store.""" + +from __future__ import annotations + +import hashlib +import json +import sqlite3 +import uuid +from datetime import UTC, datetime +from pathlib import Path +from typing import Any +from urllib.parse import quote + +import pandas as pd + +from judgearena.log import get_logger + +logger = get_logger(__name__) + +IN_QUERY_CHUNK_SIZE = 500 +INFERENCE_DB_NAME = "inference.db" + +INFERENCE_COLUMNS = ( + "input_hash", + "input_text", + "output_text", + "producer_metadata_json", + "pushed_by", + "pushed_at", + "run_id", +) +METADATA_COLUMNS = ( + "input_hash", + "metadata_hash", + "metadata_json", + "observed_at", + "run_id", +) + + +def stable_json_dumps(value: Any) -> str: + """Return deterministic JSON for hashing and descriptor comparison.""" + return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + + +def descriptor_hash(value: Any, *, length: int | None = 16) -> str: + """Return a stable SHA-256 digest for a JSON-serializable descriptor.""" + digest = hashlib.sha256(stable_json_dumps(value).encode("utf-8")).hexdigest() + return digest if length is None else digest[:length] + + +def metadata_hash(metadata: Any) -> str: + """Return the full content hash for optional row-level metadata.""" + return descriptor_hash(metadata, length=None) + + +def sanitize_path_component(value: str) -> str: + """Return a single safe path segment without separators or traversal.""" + normalized = str(value).strip() + if not normalized or normalized in {".", ".."}: + raise ValueError(f"Invalid path component: {value!r}") + for segment in normalized.replace("\\", "/").split("/"): + if segment in {"", ".", ".."}: + raise ValueError(f"Invalid path component: {value!r}") + return quote(normalized, safe="-_.~") + + +def store_folder( + store_root: Path | str, + task: str, + model_spec: str, + config_hash: str, +) -> Path: + """Return the local folder for one configuration-scoped inference cell.""" + provider, model_path = model_spec.split("/", 1) + return ( + Path(store_root).expanduser() + / "inference" + / sanitize_path_component(task) + / sanitize_path_component(provider) + / sanitize_path_component(model_path) + / sanitize_path_component(config_hash) + ) + + +def _normalize_metadata_json(value: Any) -> tuple[str, Any]: + """Return canonical JSON text and parsed value for one metadata payload.""" + if isinstance(value, str): + parsed = json.loads(value) + else: + parsed = value + return stable_json_dumps(parsed), parsed + + +def _resolve_metadata_hash( + parsed_metadata: Any, + provided_hash: Any, +) -> str: + """Return a metadata hash, ignoring missing or NaN caller values.""" + if provided_hash is not None and not ( + isinstance(provided_hash, float) and pd.isna(provided_hash) + ): + provided = str(provided_hash).strip() + if provided: + return provided + return metadata_hash(parsed_metadata) + + +def write_store_metadata(folder: Path | str, config: dict) -> Path: + """Write or validate the descriptor that identifies a configuration cell.""" + folder = Path(folder) + expected_hash = descriptor_hash(config) + if folder.name != expected_hash: + raise ValueError( + f"Cell folder {folder.name!r} does not match descriptor hash " + f"{expected_hash!r}." + ) + folder.mkdir(parents=True, exist_ok=True) + path = folder / "metadata.json" + serialized = stable_json_dumps(config) + if path.exists(): + existing = stable_json_dumps(json.loads(path.read_text(encoding="utf-8"))) + if existing != serialized: + raise ValueError( + f"Existing metadata at {path} does not match the requested descriptor." + ) + return path + temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") + temporary.write_text( + json.dumps(config, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + temporary.replace(path) + return path + + +class SQLiteInferenceStore: + """SQLite-backed store for content-addressed inference outputs.""" + + def __init__(self, db_path: Path | str, *, readonly: bool = False) -> None: + self.db_path = Path(db_path) + self.readonly = readonly + self._conn: sqlite3.Connection | None = None + + def _connect(self) -> sqlite3.Connection: + if self._conn is None: + if self.readonly: + if not self.db_path.exists(): + raise FileNotFoundError( + f"Inference store not found: {self.db_path}" + ) + uri = f"{self.db_path.resolve().as_uri()}?mode=ro" + self._conn = sqlite3.connect(uri, uri=True) + return self._conn + + self.db_path.parent.mkdir(parents=True, exist_ok=True) + self._conn = sqlite3.connect(self.db_path) + self._conn.execute("PRAGMA journal_mode=WAL") + self._conn.execute(""" + CREATE TABLE IF NOT EXISTS inference ( + input_hash TEXT PRIMARY KEY, + input_text TEXT NOT NULL, + output_text TEXT NOT NULL, + producer_metadata_json TEXT NOT NULL, + pushed_by TEXT NOT NULL, + pushed_at TEXT NOT NULL, + run_id TEXT NOT NULL + ) + """) + self._conn.execute(""" + CREATE TABLE IF NOT EXISTS inference_metadata ( + input_hash TEXT NOT NULL, + metadata_hash TEXT NOT NULL, + metadata_json TEXT NOT NULL, + observed_at TEXT NOT NULL, + run_id TEXT NOT NULL, + PRIMARY KEY (input_hash, metadata_hash) + ) + """) + self._conn.commit() + return self._conn + + def _insert_outputs( + self, + df: pd.DataFrame, + *, + pushed_by: str, + run_id: str, + replace: bool = False, + ) -> int: + required = {"input_hash", "input_text", "output_text"} + missing_cols = required - set(df.columns) + if missing_cols: + raise ValueError(f"DataFrame missing columns: {missing_cols}") + + now = datetime.now(UTC).isoformat() + rows = [ + ( + row["input_hash"], + row["input_text"], + row["output_text"], + row.get("producer_metadata_json", "{}"), + pushed_by, + now, + run_id, + ) + for _, row in df.iterrows() + ] + verb = "REPLACE" if replace else "IGNORE" + conn = self._connect() + written = 0 + for row in rows: + cursor = conn.execute( + f"INSERT OR {verb} INTO inference " + "(input_hash, input_text, output_text, producer_metadata_json, " + "pushed_by, pushed_at, run_id) VALUES (?, ?, ?, ?, ?, ?, ?)", + row, + ) + written += cursor.rowcount + return written + + def _insert_metadata( + self, + df: pd.DataFrame, + *, + run_id: str, + ) -> int: + required = {"input_hash", "metadata_json"} + missing_cols = required - set(df.columns) + if missing_cols: + raise ValueError(f"DataFrame missing columns: {missing_cols}") + + now = datetime.now(UTC).isoformat() + rows: list[tuple[str, str, str, str, str]] = [] + for _, row in df.iterrows(): + metadata_json, parsed_metadata = _normalize_metadata_json( + row["metadata_json"] + ) + rows.append( + ( + row["input_hash"], + _resolve_metadata_hash( + parsed_metadata, + row["metadata_hash"] if "metadata_hash" in df.columns else None, + ), + metadata_json, + now, + run_id, + ) + ) + conn = self._connect() + written = 0 + for row in rows: + cursor = conn.execute( + "INSERT OR IGNORE INTO inference_metadata " + "(input_hash, metadata_hash, metadata_json, observed_at, run_id) " + "VALUES (?, ?, ?, ?, ?)", + row, + ) + written += cursor.rowcount + return written + + def save_outputs( + self, + df: pd.DataFrame, + *, + pushed_by: str, + run_id: str | None = None, + replace: bool = False, + ) -> int: + """Insert inference rows, optionally replacing existing keys.""" + with self._connect(): + written = self._insert_outputs( + df, + pushed_by=pushed_by, + run_id=run_id or str(uuid.uuid4()), + replace=replace, + ) + logger.info("Wrote %d inference rows to %s", written, self.db_path) + return written + + def save_metadata( + self, + df: pd.DataFrame, + *, + run_id: str | None = None, + ) -> int: + """Associate optional row metadata with cached inference rows.""" + with self._connect(): + written = self._insert_metadata( + df, + run_id=run_id or str(uuid.uuid4()), + ) + logger.info("Wrote %d metadata associations to %s", written, self.db_path) + return written + + def save_outputs_and_metadata( + self, + outputs: pd.DataFrame, + metadata: pd.DataFrame, + *, + pushed_by: str, + run_id: str | None = None, + replace: bool = False, + ) -> tuple[int, int]: + """Atomically save inference outputs and their metadata associations.""" + resolved_run_id = run_id or str(uuid.uuid4()) + with self._connect(): + outputs_written = self._insert_outputs( + outputs, + pushed_by=pushed_by, + run_id=resolved_run_id, + replace=replace, + ) + metadata_written = self._insert_metadata( + metadata, + run_id=resolved_run_id, + ) + logger.info( + "Wrote %d inference rows and %d metadata associations to %s", + outputs_written, + metadata_written, + self.db_path, + ) + return outputs_written, metadata_written + + def query(self, input_hashes: list[str] | None = None) -> pd.DataFrame: + """Return inference rows, optionally restricted to input hashes.""" + conn = self._connect() + if input_hashes is None: + return pd.read_sql( + "SELECT * FROM inference ORDER BY input_hash", + conn, + ) + if not input_hashes: + return pd.read_sql("SELECT * FROM inference WHERE 0", conn) + + frames: list[pd.DataFrame] = [] + for chunk_start in range(0, len(input_hashes), IN_QUERY_CHUNK_SIZE): + chunk = input_hashes[chunk_start : chunk_start + IN_QUERY_CHUNK_SIZE] + placeholders = ",".join("?" * len(chunk)) + frames.append( + pd.read_sql( + f"SELECT * FROM inference WHERE input_hash IN ({placeholders})" + " ORDER BY input_hash", + conn, + params=chunk, + ) + ) + if len(frames) == 1: + return frames[0] + combined = pd.concat(frames, ignore_index=True) + return combined.sort_values("input_hash", kind="stable").reset_index(drop=True) + + def query_metadata(self, input_hashes: list[str] | None = None) -> pd.DataFrame: + """Return metadata association rows.""" + conn = self._connect() + if input_hashes is None: + return pd.read_sql( + "SELECT * FROM inference_metadata ORDER BY input_hash, metadata_hash", + conn, + ) + if not input_hashes: + return pd.read_sql("SELECT * FROM inference_metadata WHERE 0", conn) + + frames: list[pd.DataFrame] = [] + for chunk_start in range(0, len(input_hashes), IN_QUERY_CHUNK_SIZE): + chunk = input_hashes[chunk_start : chunk_start + IN_QUERY_CHUNK_SIZE] + placeholders = ",".join("?" * len(chunk)) + frames.append( + pd.read_sql( + f"SELECT * FROM inference_metadata " + f"WHERE input_hash IN ({placeholders}) " + "ORDER BY input_hash, metadata_hash", + conn, + params=chunk, + ) + ) + if len(frames) == 1: + return frames[0] + combined = pd.concat(frames, ignore_index=True) + return combined.sort_values( + ["input_hash", "metadata_hash"], + kind="stable", + ).reset_index(drop=True) + + def missing(self, input_hashes: list[str]) -> list[str]: + """Return input hashes absent from the store, preserving caller order.""" + if not input_hashes: + return [] + present = self.outputs_by_hash(input_hashes) + return [value for value in input_hashes if value not in present] + + def outputs_by_hash(self, input_hashes: list[str]) -> dict[str, str]: + """Return stored output text keyed by input hash.""" + if not input_hashes: + return {} + conn = self._connect() + outputs: dict[str, str] = {} + for chunk_start in range(0, len(input_hashes), IN_QUERY_CHUNK_SIZE): + chunk = input_hashes[chunk_start : chunk_start + IN_QUERY_CHUNK_SIZE] + placeholders = ",".join("?" * len(chunk)) + outputs.update( + dict( + conn.execute( + f"SELECT input_hash, output_text FROM inference " + f"WHERE input_hash IN ({placeholders})", + chunk, + ).fetchall() + ) + ) + return outputs + + def close(self) -> None: + if self._conn is not None: + self._conn.close() + self._conn = None + + def __enter__(self) -> SQLiteInferenceStore: + return self + + def __exit__(self, *_: object) -> None: + self.close() diff --git a/judgearena/store_sync.py b/judgearena/store_sync.py new file mode 100644 index 0000000..fba153f --- /dev/null +++ b/judgearena/store_sync.py @@ -0,0 +1,646 @@ +"""Hugging Face Hub synchronization for unified inference cache cells.""" + +from __future__ import annotations + +import hashlib +import json +import shutil +import sqlite3 +import tempfile +from pathlib import Path, PurePosixPath + +import pandas as pd +from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download +from huggingface_hub.utils import ( + EntryNotFoundError, + HfHubHTTPError, + RepositoryNotFoundError, +) + +from judgearena.log import get_logger +from judgearena.store_sqlite import ( + INFERENCE_COLUMNS, + INFERENCE_DB_NAME, + METADATA_COLUMNS, + sanitize_path_component, + write_store_metadata, +) + +logger = get_logger(__name__) + +DEFAULT_CACHE_REPO = "judge-arena/judge-arena-cache" +DEFAULT_INFERENCE_PREFIX = "inference" +_DEFAULT_MAX_RETRIES = 5 +_INFERENCE_TABLE = "inference" +_METADATA_TABLE = "inference_metadata" + + +def _output_rank(output_text: str) -> str: + return hashlib.sha256(str(output_text).encode("utf-8")).hexdigest() + + +def _remove_sidecars(db_path: Path) -> None: + for suffix in ("-wal", "-shm"): + Path(f"{db_path}{suffix}").unlink(missing_ok=True) + + +def _copy_database(source: Path, destination: Path) -> None: + destination.parent.mkdir(parents=True, exist_ok=True) + destination.unlink(missing_ok=True) + _remove_sidecars(destination) + with ( + sqlite3.connect(source) as source_conn, + sqlite3.connect(destination) as destination_conn, + ): + source_conn.backup(destination_conn) + + +def _read_table(db_path: Path, table: str) -> pd.DataFrame: + with sqlite3.connect(db_path) as conn: + return pd.read_sql(f'SELECT * FROM "{table}"', conn) + + +def _merge_inference_frames(frames: list[pd.DataFrame]) -> pd.DataFrame: + if not frames: + return pd.DataFrame(columns=list(INFERENCE_COLUMNS)) + merged = pd.concat(frames, ignore_index=True) + if merged.empty: + return merged + merged = merged.copy() + merged["_output_rank"] = merged["output_text"].map(_output_rank) + merged = merged.sort_values( + ["pushed_at", "run_id", "_output_rank"], + kind="stable", + na_position="first", + ) + merged = merged.drop_duplicates(subset=["input_hash"], keep="last") + return merged.drop(columns=["_output_rank"]) + + +def _merge_metadata_frames(frames: list[pd.DataFrame]) -> pd.DataFrame: + if not frames: + return pd.DataFrame(columns=list(METADATA_COLUMNS)) + merged = pd.concat(frames, ignore_index=True) + if merged.empty: + return merged + merged = merged.sort_values( + ["observed_at", "run_id"], + kind="stable", + na_position="first", + ) + return merged.drop_duplicates(subset=["input_hash", "metadata_hash"], keep="last") + + +def _write_merged_db( + inference: pd.DataFrame, + metadata: pd.DataFrame, + destination: Path, + *, + template: Path, +) -> None: + _copy_database(template, destination) + with sqlite3.connect(destination) as conn: + conn.execute(f'DELETE FROM "{_INFERENCE_TABLE}"') + conn.execute(f'DELETE FROM "{_METADATA_TABLE}"') + if not inference.empty: + inference = inference.astype(object).where(pd.notna(inference), None) + conn.executemany( + f'INSERT INTO "{_INFERENCE_TABLE}" ' + f"({', '.join(INFERENCE_COLUMNS)}) " + f"VALUES ({', '.join('?' for _ in INFERENCE_COLUMNS)})", + inference.loc[:, list(INFERENCE_COLUMNS)].itertuples( + index=False, + name=None, + ), + ) + if not metadata.empty: + metadata = metadata.astype(object).where(pd.notna(metadata), None) + conn.executemany( + f'INSERT INTO "{_METADATA_TABLE}" ' + f"({', '.join(METADATA_COLUMNS)}) " + f"VALUES ({', '.join('?' for _ in METADATA_COLUMNS)})", + metadata.loc[:, list(METADATA_COLUMNS)].itertuples( + index=False, + name=None, + ), + ) + conn.commit() + conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") + + +def _merge_dbs(sources: list[Path], destination: Path) -> None: + """Merge inference rows and union metadata associations.""" + if not sources: + raise ValueError("sources must not be empty") + sources = [Path(source) for source in sources] + inference_frames = [_read_table(source, _INFERENCE_TABLE) for source in sources] + metadata_frames = [_read_table(source, _METADATA_TABLE) for source in sources] + merged_inference = _merge_inference_frames(inference_frames) + merged_metadata = _merge_metadata_frames(metadata_frames) + _write_merged_db( + merged_inference, + merged_metadata, + destination, + template=sources[0], + ) + + +def _replace_db(destination: Path, source: Path) -> None: + destination.parent.mkdir(parents=True, exist_ok=True) + temporary = destination.with_suffix(f"{destination.suffix}.tmp") + temporary.unlink(missing_ok=True) + _remove_sidecars(destination) + shutil.copy2(source, temporary) + temporary.replace(destination) + _remove_sidecars(destination) + + +def _head_oid( + api: HfApi, + repo_id: str, + *, + repo_type: str, + revision: str, +) -> str | None: + try: + info = api.repo_info( + repo_id, + repo_type=repo_type, + revision=revision, + ) + except RepositoryNotFoundError: + return None + return getattr(info, "sha", None) or getattr(info, "oid", None) + + +def _download_remote_file( + repo_id: str, + path_in_repo: str, + *, + repo_type: str, + revision: str, + token: str | None, +) -> Path | None: + try: + path = hf_hub_download( + repo_id=repo_id, + filename=path_in_repo, + repo_type=repo_type, + revision=revision, + token=token, + ) + except EntryNotFoundError: + return None + return Path(path) + + +def _cell_db_path(cell_dir: Path | str) -> Path: + return Path(cell_dir) / INFERENCE_DB_NAME + + +def _cell_metadata_path(cell_dir: Path | str) -> Path: + return Path(cell_dir) / "metadata.json" + + +def validate_path_filters( + *, + prefix: str | None = None, + task: str | None = None, + provider: str | None = None, + model: str | None = None, + config_hash: str | None = None, +) -> str | None: + """Validate hierarchical filters and return a normalized path prefix.""" + if prefix: + normalized = prefix.strip("/") + parts = PurePosixPath(normalized).parts + if ( + not parts + or parts[0] != DEFAULT_INFERENCE_PREFIX + or len(parts) > 5 + or any(part in {"", ".", ".."} for part in parts) + ): + raise ValueError( + f"--prefix must be a cache-cell directory under " + f"{DEFAULT_INFERENCE_PREFIX!r}, got {normalized!r}." + ) + return normalized + if config_hash and not (task and provider and model): + raise ValueError("--config_hash requires --task, --provider, and --model.") + if model and not (task and provider): + raise ValueError("--model requires --task and --provider.") + if provider and not task: + raise ValueError("--provider requires --task.") + if not task: + return None + parts = [ + DEFAULT_INFERENCE_PREFIX, + sanitize_path_component(task), + ] + if provider: + parts.append(sanitize_path_component(provider)) + if model: + parts.append(sanitize_path_component(model)) + if config_hash: + parts.append(sanitize_path_component(config_hash)) + return "/".join(parts) + + +def rel_path_in_repo(local_path: Path | str, store_root: Path | str) -> str: + """Return a cell path relative to the store root.""" + return ( + Path(local_path) + .expanduser() + .relative_to(Path(store_root).expanduser()) + .as_posix() + ) + + +def _local_path_for_remote_cell(store_root: Path, path_in_repo: str) -> Path: + remote_path = PurePosixPath(path_in_repo) + if ( + remote_path.is_absolute() + or len(remote_path.parts) != 6 + or remote_path.parts[0] != DEFAULT_INFERENCE_PREFIX + or remote_path.name != INFERENCE_DB_NAME + or any(part in {"", ".", ".."} for part in remote_path.parts) + ): + raise ValueError(f"Invalid remote cache cell path: {path_in_repo!r}") + local_path = store_root.joinpath(*remote_path.parts).resolve() + local_path.relative_to(store_root.resolve()) + return local_path + + +def _path_matches_prefix(rel_path: str, normalized_prefix: str) -> bool: + """Return True when *rel_path* equals or extends *normalized_prefix*.""" + return rel_path == normalized_prefix or rel_path.startswith(f"{normalized_prefix}/") + + +def iter_cell_dbs( + store_root: Path | str, + *, + path_prefix: str | None = None, +) -> list[Path]: + """Return local inference.db cells, optionally filtered by path prefix.""" + root = Path(store_root).expanduser() + if not root.exists(): + return [] + normalized_prefix = path_prefix.strip("/") if path_prefix else None + cells: list[Path] = [] + for candidate in sorted(root.rglob(INFERENCE_DB_NAME)): + relative = rel_path_in_repo(candidate, root) + try: + expected = _local_path_for_remote_cell(root, relative) + except ValueError: + logger.warning("Ignoring noncanonical local cache cell: %s", candidate) + continue + if candidate.resolve() == expected: + cells.append(candidate) + if normalized_prefix is None: + return cells + return [ + cell + for cell in cells + if _path_matches_prefix(rel_path_in_repo(cell, root), normalized_prefix) + ] + + +def discover_remote_cell_dbs( + repo_id: str, + *, + path_prefix: str | None = None, + repo_type: str = "dataset", + revision: str = "main", + token: str | None = None, +) -> list[str]: + """List remote inference.db paths under an optional prefix.""" + api = HfApi(token=token) + try: + files = api.list_repo_files( + repo_id, + repo_type=repo_type, + revision=revision, + ) + except RepositoryNotFoundError: + return [] + normalized_prefix = path_prefix.strip("/") if path_prefix else None + db_paths = [ + path + for path in files + if path.endswith(f"/{INFERENCE_DB_NAME}") or path == INFERENCE_DB_NAME + ] + if normalized_prefix is None: + return sorted(db_paths) + return sorted( + path for path in db_paths if _path_matches_prefix(path, normalized_prefix) + ) + + +def _metadata_path_in_repo(db_path_in_repo: str) -> str: + return str(Path(db_path_in_repo).parent / "metadata.json").replace("\\", "/") + + +def _materialize_remote_metadata( + remote_metadata: Path, + local_metadata_path: Path, +) -> None: + """Validate or write the remote descriptor before any DB merge.""" + remote_config = json.loads(remote_metadata.read_text(encoding="utf-8")) + write_store_metadata(local_metadata_path.parent, remote_config) + + +def fetch_cell( + repo_id: str, + path_in_repo: str, + local_db_path: Path | str, + *, + repo_type: str = "dataset", + revision: str = "main", + token: str | None = None, +) -> bool: + """Fetch a remote cell after validating sibling metadata.""" + local_db_path = Path(local_db_path) + local_metadata_path = _cell_metadata_path(local_db_path.parent) + metadata_in_repo = _metadata_path_in_repo(path_in_repo) + + remote_db = _download_remote_file( + repo_id, + path_in_repo, + repo_type=repo_type, + revision=revision, + token=token, + ) + if remote_db is None: + return False + + remote_metadata = _download_remote_file( + repo_id, + metadata_in_repo, + repo_type=repo_type, + revision=revision, + token=token, + ) + if remote_metadata is None: + raise ValueError( + f"Remote cell {path_in_repo} is missing required {metadata_in_repo}." + ) + + local_metadata_path.parent.mkdir(parents=True, exist_ok=True) + _materialize_remote_metadata(remote_metadata, local_metadata_path) + + if local_db_path.exists(): + with tempfile.TemporaryDirectory() as temporary: + merged = Path(temporary) / "merged.db" + _merge_dbs([local_db_path, remote_db], merged) + _replace_db(local_db_path, merged) + else: + _copy_database(remote_db, local_db_path) + + logger.info("Fetched %s from %s", path_in_repo, repo_id) + return True + + +def fetch_cell_metadata( + repo_id: str, + metadata_path_in_repo: str, + local_metadata_path: Path | str, + *, + repo_type: str = "dataset", + revision: str = "main", + token: str | None = None, +) -> bool: + """Download remote metadata.json when present.""" + local_metadata_path = Path(local_metadata_path) + remote = _download_remote_file( + repo_id, + metadata_path_in_repo, + repo_type=repo_type, + revision=revision, + token=token, + ) + if remote is None: + return False + _materialize_remote_metadata(remote, local_metadata_path) + logger.info("Fetched %s from %s", metadata_path_in_repo, repo_id) + return True + + +def fetch_cells( + repo_id: str, + store_root: Path | str, + db_paths: list[Path | str], + *, + repo_type: str = "dataset", + revision: str = "main", + token: str | None = None, + strict: bool = False, +) -> None: + """Fetch cells and sibling metadata, warning instead of failing by default.""" + for db_path_value in db_paths: + db_path = Path(db_path_value) + path_in_repo = rel_path_in_repo(db_path, store_root) + try: + fetch_cell( + repo_id, + path_in_repo, + db_path, + repo_type=repo_type, + revision=revision, + token=token, + ) + except Exception as exc: # noqa: BLE001 + if strict: + raise + logger.warning("Cache fetch skipped for %s: %s", path_in_repo, exc) + + +def fetch_remote_cells( + repo_id: str, + store_root: Path | str, + *, + path_prefix: str | None = None, + repo_type: str = "dataset", + revision: str = "main", + token: str | None = None, + strict: bool = False, +) -> list[Path]: + """Discover and fetch remote cells into a possibly empty local store.""" + store_root = Path(store_root).expanduser() + store_root.mkdir(parents=True, exist_ok=True) + remote_paths = discover_remote_cell_dbs( + repo_id, + path_prefix=path_prefix, + repo_type=repo_type, + revision=revision, + token=token, + ) + local_paths: list[Path] = [] + for path_in_repo in remote_paths: + try: + local_paths.append(_local_path_for_remote_cell(store_root, path_in_repo)) + except ValueError as exc: + if strict: + raise + logger.warning("Remote cache cell skipped: %s", exc) + fetch_cells( + repo_id, + store_root, + local_paths, + repo_type=repo_type, + revision=revision, + token=token, + strict=strict, + ) + return local_paths + + +def push_cell( + repo_id: str, + path_in_repo: str, + local_db_path: Path | str, + *, + pushed_by: str, + local_metadata_path: Path | str | None = None, + repo_type: str = "dataset", + revision: str = "main", + token: str | None = None, + create_pr: bool = False, + ensure_repo: bool = False, + private: bool = True, + max_retries: int = _DEFAULT_MAX_RETRIES, +) -> str: + """Merge and upload a cell DB plus metadata in one optimistic commit.""" + local_db_path = Path(local_db_path) + if not local_db_path.exists(): + raise FileNotFoundError(local_db_path) + + metadata_path = ( + Path(local_metadata_path) + if local_metadata_path is not None + else local_db_path.parent / "metadata.json" + ) + if not metadata_path.exists(): + raise FileNotFoundError( + f"Missing metadata.json for cache cell at {local_db_path}" + ) + metadata_in_repo = _metadata_path_in_repo(path_in_repo) + + api = HfApi(token=token) + if ensure_repo: + api.create_repo( + repo_id, + repo_type=repo_type, + private=private, + exist_ok=True, + ) + + last_error: HfHubHTTPError | None = None + for attempt in range(max_retries): + parent = _head_oid( + api, + repo_id, + repo_type=repo_type, + revision=revision, + ) + remote = _download_remote_file( + repo_id, + path_in_repo, + repo_type=repo_type, + revision=revision, + token=token, + ) + sources = [local_db_path] if remote is None else [remote, local_db_path] + with tempfile.TemporaryDirectory() as temporary: + merged = Path(temporary) / "merged.db" + _merge_dbs(sources, merged) + operations = [ + CommitOperationAdd( + path_in_repo=path_in_repo, + path_or_fileobj=str(merged), + ), + CommitOperationAdd( + path_in_repo=metadata_in_repo, + path_or_fileobj=str(metadata_path), + ), + ] + try: + info = api.create_commit( + repo_id=repo_id, + repo_type=repo_type, + revision=revision, + operations=operations, + parent_commit=parent, + create_pr=create_pr, + commit_message=( + f"cache: {Path(path_in_repo).parent.as_posix()} (by {pushed_by})" + ), + ) + except HfHubHTTPError as exc: + status = getattr(exc.response, "status_code", None) + if status == 412 and attempt < max_retries - 1: + last_error = exc + logger.info( + "Push of %s hit a concurrent commit (412); retrying (%d/%d).", + path_in_repo, + attempt + 1, + max_retries, + ) + continue + raise + _replace_db(local_db_path, merged) + + result = ( + getattr(info, "pr_url", None) + or getattr(info, "oid", None) + or getattr(info, "sha", None) + or "" + ) + logger.info("Pushed %s to %s (%s)", path_in_repo, repo_id, result) + return result + + raise RuntimeError( + f"push_cell exhausted {max_retries} retries for {path_in_repo}" + ) from last_error + + +def push_cells( + repo_id: str, + store_root: Path | str, + db_paths: list[Path | str], + *, + pushed_by: str, + repo_type: str = "dataset", + revision: str = "main", + token: str | None = None, + create_pr: bool = False, + ensure_repo: bool = False, + private: bool = True, + strict: bool = False, +) -> None: + """Push existing local cells and their descriptor metadata.""" + for db_path_value in db_paths: + db_path = Path(db_path_value) + if not db_path.exists(): + message = f"Cache cell does not exist: {db_path}" + if strict: + raise FileNotFoundError(message) + logger.warning(message) + continue + path_in_repo = rel_path_in_repo(db_path, store_root) + try: + push_cell( + repo_id, + path_in_repo, + db_path, + pushed_by=pushed_by, + repo_type=repo_type, + revision=revision, + token=token, + create_pr=create_pr, + ensure_repo=ensure_repo, + private=private, + ) + except Exception as exc: # noqa: BLE001 + if strict: + raise + logger.warning("Cache push skipped for %s: %s", path_in_repo, exc) diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..ac9177c --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,125 @@ +import uuid +from pathlib import Path + +import pytest +from huggingface_hub.utils import ( + EntryNotFoundError, + HfHubHTTPError, + RepositoryNotFoundError, +) + +from judgearena import store_sync + + +class _Response: + def __init__(self, status_code: int) -> None: + self.status_code = status_code + self.headers: dict = {} + self.text = "" + self.reason = "Precondition Failed" + self.request = None + + +class _Info: + def __init__(self, oid=None, sha=None, pr_url=None) -> None: + self.oid = oid + self.sha = sha + self.pr_url = pr_url + + +class _FakeRepo: + def __init__(self) -> None: + self.files: dict[str, bytes] = {} + self.head: str | None = None + self.pending_conflict = False + self.inject = None + self.commit_calls = 0 + + +class _FakeHfApi: + def __init__(self, repo: _FakeRepo) -> None: + self.repo = repo + + def create_repo(self, *args, **kwargs): + return None + + def repo_info(self, repo_id, repo_type="dataset", revision="main"): + if self.repo.head is None: + raise RepositoryNotFoundError( + f"Repository {repo_id} not found", + response=_Response(404), + ) + return _Info(sha=self.repo.head) + + def list_repo_files(self, repo_id, repo_type="dataset", revision="main"): + if self.repo.head is None: + raise RepositoryNotFoundError( + f"Repository {repo_id} not found", + response=_Response(404), + ) + return sorted(self.repo.files) + + def create_commit( + self, + *, + repo_id, + repo_type="dataset", + revision="main", + operations, + parent_commit=None, + create_pr=False, + commit_message=None, + ): + self.repo.commit_calls += 1 + if self.repo.pending_conflict: + self.repo.pending_conflict = False + if self.repo.inject is not None: + self.repo.inject() + raise HfHubHTTPError( + "412 Precondition Failed", + response=_Response(412), + ) + if ( + not create_pr + and parent_commit is not None + and parent_commit != self.repo.head + ): + raise HfHubHTTPError( + "412 Precondition Failed", + response=_Response(412), + ) + if create_pr: + return _Info(pr_url="https://hf.co/pr/1") + for operation in operations: + self.repo.files[operation.path_in_repo] = Path( + operation.path_or_fileobj + ).read_bytes() + self.repo.head = uuid.uuid4().hex + return _Info(oid=self.repo.head, sha=self.repo.head) + + +@pytest.fixture +def fake_hub(monkeypatch, tmp_path): + repo = _FakeRepo() + + def fake_download( + repo_id, + filename, + repo_type="dataset", + revision="main", + token=None, + ): + if filename not in repo.files: + raise EntryNotFoundError(f"{filename} not found") + destination = tmp_path / "downloads" / filename + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(repo.files[filename]) + return str(destination) + + monkeypatch.setattr( + store_sync, + "HfApi", + lambda token=None: _FakeHfApi(repo), + ) + monkeypatch.setattr(store_sync, "hf_hub_download", fake_download) + return repo diff --git a/tests/test_store_sqlite.py b/tests/test_store_sqlite.py new file mode 100644 index 0000000..1e1dbc4 --- /dev/null +++ b/tests/test_store_sqlite.py @@ -0,0 +1,258 @@ +import json + +import pandas as pd +import pytest + +from judgearena.store_sqlite import ( + IN_QUERY_CHUNK_SIZE, + SQLiteInferenceStore, + descriptor_hash, + metadata_hash, + sanitize_path_component, + stable_json_dumps, + store_folder, + write_store_metadata, +) + +CELL_CONFIG = {"task": "arena", "model_spec": "VLLM/Qwen/judge"} + + +def _outputs(hashes: list[str]) -> pd.DataFrame: + return pd.DataFrame( + { + "input_hash": hashes, + "input_text": [f"input-{value}" for value in hashes], + "output_text": [f"output-{value}" for value in hashes], + } + ) + + +def test_inference_roundtrip_and_missing_order(tmp_path): + hashes = ["h0", "h1", "h2"] + with SQLiteInferenceStore(tmp_path / "inference.db") as store: + store.save_outputs(_outputs(hashes), pushed_by="test") + assert store.missing(hashes) == [] + assert store.missing([*hashes, "missing"]) == ["missing"] + result = store.query(["h1", "absent", "h0"]) + assert store.query([]).empty + + assert result["input_hash"].tolist() == ["h0", "h1"] + + +def test_readonly_store_queries_without_schema_writes(tmp_path): + db_path = tmp_path / "inference.db" + with SQLiteInferenceStore(db_path) as store: + store.save_outputs(_outputs(["h1"]), pushed_by="test") + + before = db_path.stat().st_mtime_ns + with SQLiteInferenceStore(db_path, readonly=True) as store: + assert store.outputs_by_hash(["h1", "missing"]) == {"h1": "output-h1"} + assert store.missing(["h1", "missing"]) == ["missing"] + assert db_path.stat().st_mtime_ns == before + + +def test_save_outputs_insert_only_vs_replace(tmp_path): + db_path = tmp_path / "inference.db" + with SQLiteInferenceStore(db_path) as store: + assert store.save_outputs(_outputs(["same"]), pushed_by="alice") == 1 + assert ( + store.save_outputs( + pd.DataFrame( + { + "input_hash": ["same"], + "input_text": ["input-same"], + "output_text": ["ignored"], + } + ), + pushed_by="bob", + ) + == 0 + ) + assert store.query(["same"])["output_text"].iloc[0] == "output-same" + assert ( + store.save_outputs( + pd.DataFrame( + { + "input_hash": ["same"], + "input_text": ["input-same"], + "output_text": ["updated"], + } + ), + pushed_by="bob", + replace=True, + ) + == 1 + ) + result = store.query(["same"]) + + assert result["output_text"].iloc[0] == "updated" + assert result["pushed_by"].iloc[0] == "bob" + + +def test_many_to_one_metadata_is_idempotent(tmp_path): + meta = {"question_id": "q-1", "role": "judge"} + meta_json = stable_json_dumps(meta) + with SQLiteInferenceStore(tmp_path / "inference.db") as store: + store.save_outputs(_outputs(["h0"]), pushed_by="test") + assert ( + store.save_metadata( + pd.DataFrame({"input_hash": ["h0"], "metadata_json": [meta_json]}), + run_id="run-a", + ) + == 1 + ) + assert ( + store.save_metadata( + pd.DataFrame({"input_hash": ["h0"], "metadata_json": [meta_json]}), + run_id="run-b", + ) + == 0 + ) + assert ( + store.save_metadata( + pd.DataFrame( + { + "input_hash": ["h0"], + "metadata_hash": [metadata_hash(meta)], + "metadata_json": [meta_json], + } + ), + run_id="run-c", + ) + == 0 + ) + rows = store.query_metadata(["h0"]) + + assert len(rows) == 1 + assert json.loads(rows["metadata_json"].iloc[0]) == meta + + +def test_output_and_metadata_batch_rolls_back_atomically(tmp_path): + with SQLiteInferenceStore(tmp_path / "inference.db") as store: + with pytest.raises(ValueError): + store.save_outputs_and_metadata( + _outputs(["h0"]), + pd.DataFrame( + { + "input_hash": ["h0"], + "metadata_json": ["not-json"], + } + ), + pushed_by="test", + ) + + assert store.query().empty + assert store.query_metadata().empty + + +def test_save_metadata_normalizes_dict_and_string_equivalently(tmp_path): + meta = {"b": 2, "a": 1} + with SQLiteInferenceStore(tmp_path / "inference.db") as store: + store.save_outputs(_outputs(["h0"]), pushed_by="test") + store.save_metadata( + pd.DataFrame({"input_hash": ["h0"], "metadata_json": [meta]}), + ) + rows = store.query_metadata(["h0"]) + + assert rows["metadata_hash"].iloc[0] == metadata_hash(meta) + assert json.loads(rows["metadata_json"].iloc[0]) == meta + + with SQLiteInferenceStore(tmp_path / "other.db") as store: + store.save_outputs(_outputs(["h0"]), pushed_by="test") + store.save_metadata( + pd.DataFrame( + { + "input_hash": ["h0"], + "metadata_json": [json.dumps(meta, sort_keys=False)], + } + ), + ) + other = store.query_metadata(["h0"]) + + assert other["metadata_hash"].iloc[0] == rows["metadata_hash"].iloc[0] + + +def test_save_metadata_ignores_nan_metadata_hash(tmp_path): + meta = {"question_id": "q-1"} + with SQLiteInferenceStore(tmp_path / "inference.db") as store: + store.save_outputs(_outputs(["h0"]), pushed_by="test") + store.save_metadata( + pd.DataFrame( + { + "input_hash": ["h0"], + "metadata_hash": [float("nan")], + "metadata_json": [stable_json_dumps(meta)], + } + ), + ) + rows = store.query_metadata(["h0"]) + + assert rows["metadata_hash"].iloc[0] == metadata_hash(meta) + + +def test_chunked_query_preserves_input_independent_order(tmp_path, monkeypatch): + hashes = [f"h{index}" for index in range(IN_QUERY_CHUNK_SIZE + 3)] + with SQLiteInferenceStore(tmp_path / "inference.db") as store: + store.save_outputs(_outputs(hashes), pushed_by="test") + monkeypatch.setattr( + "judgearena.store_sqlite.IN_QUERY_CHUNK_SIZE", + 2, + ) + result = store.query(hashes) + + assert result["input_hash"].tolist() == sorted(hashes) + + +def test_write_store_metadata_rejects_misnamed_folder(tmp_path): + config = { + "descriptor_schema_version": "judgearena-cache/v1", + "task": "arena-hard-v2.0", + "model_spec": "VLLM/Qwen/Qwen3.5-9B", + } + wrong_folder = ( + tmp_path + / "inference" + / "arena" + / "VLLM" + / "Qwen%2FQwen3.5-9B" + / "deadbeefdeadbeef" + ) + with pytest.raises(ValueError, match="does not match descriptor hash"): + write_store_metadata(wrong_folder, config) + + +def test_store_folder_and_metadata_validation(tmp_path): + config = { + "descriptor_schema_version": "judgearena-cache/v1", + "task": "arena-hard-v2.0", + "model_spec": "VLLM/Qwen/Qwen3.5-9B", + } + folder = store_folder( + tmp_path, + "arena-hard-v2.0", + "VLLM/Qwen/Qwen3.5-9B", + descriptor_hash(config), + ) + metadata_path = write_store_metadata(folder, config) + assert folder.parts[-3:-1] == ("VLLM", "Qwen%2FQwen3.5-9B") + assert json.loads(metadata_path.read_text()) == config + write_store_metadata(folder, config) + with pytest.raises(ValueError, match="does not match"): + write_store_metadata(folder, {**config, "task": "other"}) + + +def test_sanitize_path_component_rejects_traversal(): + assert sanitize_path_component("Qwen/Qwen3.5-9B") == "Qwen%2FQwen3.5-9B" + assert sanitize_path_component("Qwen--Qwen3.5-9B") != sanitize_path_component( + "Qwen/Qwen3.5-9B" + ) + with pytest.raises(ValueError, match="Invalid path component"): + sanitize_path_component("..") + with pytest.raises(ValueError, match="Invalid path component"): + sanitize_path_component("foo/../bar") + + +def test_hash_helpers_are_stable(): + payload = {"b": 2, "a": 1} + assert descriptor_hash(payload) == descriptor_hash({"a": 1, "b": 2}) + assert metadata_hash(payload) == descriptor_hash(payload, length=None) diff --git a/tests/test_store_sync.py b/tests/test_store_sync.py new file mode 100644 index 0000000..35c86ac --- /dev/null +++ b/tests/test_store_sync.py @@ -0,0 +1,564 @@ +import json +from pathlib import Path + +import pytest + +from judgearena import store_sync +from judgearena.store_sqlite import ( + INFERENCE_DB_NAME, + SQLiteInferenceStore, + descriptor_hash, + store_folder, + write_store_metadata, +) + +REPO_ID = "org/cache" +CELL_CONFIG = {"task": "arena", "model_spec": "VLLM/Qwen/judge"} +CELL_CONFIG_HASH = descriptor_hash(CELL_CONFIG) +MODEL_SPEC = "VLLM/Qwen/judge" +PATH_IN_REPO = ( + f"inference/arena/VLLM/Qwen%2Fjudge/{CELL_CONFIG_HASH}/{INFERENCE_DB_NAME}" +) +METADATA_IN_REPO = f"inference/arena/VLLM/Qwen%2Fjudge/{CELL_CONFIG_HASH}/metadata.json" + + +def _local_cell_db(tmp_path, root: str = "store") -> Path: + cell_dir = store_folder( + Path(tmp_path) / root, "arena", MODEL_SPEC, CELL_CONFIG_HASH + ) + return cell_dir / INFERENCE_DB_NAME + + +def _write_inference(path: Path, rows: list[dict]) -> None: + with SQLiteInferenceStore(path) as store: + conn = store._connect() + conn.executemany( + "INSERT OR REPLACE INTO inference " + "(input_hash, input_text, output_text, producer_metadata_json, " + "pushed_by, pushed_at, run_id) VALUES (?, ?, ?, ?, ?, ?, ?)", + [ + ( + row["input_hash"], + row.get("input_text", f"input-{row['input_hash']}"), + row["output_text"], + row.get("producer_metadata_json", "{}"), + row.get("pushed_by", "test"), + row["pushed_at"], + row.get("run_id", "run"), + ) + for row in rows + ], + ) + conn.commit() + conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") + + +def _write_metadata(path: Path, rows: list[dict]) -> None: + with SQLiteInferenceStore(path) as store: + conn = store._connect() + conn.executemany( + "INSERT OR REPLACE INTO inference_metadata " + "(input_hash, metadata_hash, metadata_json, observed_at, run_id) " + "VALUES (?, ?, ?, ?, ?)", + [ + ( + row["input_hash"], + row["metadata_hash"], + row["metadata_json"], + row["observed_at"], + row.get("run_id", "run"), + ) + for row in rows + ], + ) + conn.commit() + conn.execute("PRAGMA wal_checkpoint(TRUNCATE)") + + +def _read_outputs(path: Path) -> dict[str, str]: + with SQLiteInferenceStore(path) as store: + frame = store.query() + return dict(zip(frame["input_hash"], frame["output_text"], strict=True)) + + +def _read_metadata(path: Path) -> set[tuple[str, str]]: + with SQLiteInferenceStore(path) as store: + frame = store.query_metadata() + return { + (row.input_hash, row.metadata_hash) for row in frame.itertuples(index=False) + } + + +def test_validate_path_filters_enforces_hierarchy(): + assert store_sync.validate_path_filters(task="arena") == "inference/arena" + assert ( + store_sync.validate_path_filters(task="arena", provider="VLLM") + == "inference/arena/VLLM" + ) + with pytest.raises(ValueError, match="requires --task"): + store_sync.validate_path_filters(provider="VLLM") + with pytest.raises(ValueError, match="requires --task and --provider"): + store_sync.validate_path_filters(task="arena", model="Qwen/judge") + with pytest.raises(ValueError, match="requires --task, --provider, and --model"): + store_sync.validate_path_filters(task="arena", config_hash="abc123") + with pytest.raises(ValueError, match="cache-cell directory"): + store_sync.validate_path_filters(prefix="inference/arena/../outside") + + +def test_merge_dbs_unions_rows_and_newest_tuple_wins(tmp_path): + older = tmp_path / "older.db" + newer = tmp_path / "newer.db" + _write_inference( + older, + [ + { + "input_hash": "shared", + "output_text": "old", + "pushed_at": "2026-01-01", + "run_id": "run-a", + }, + { + "input_hash": "only-old", + "output_text": "old-only", + "pushed_at": "2026-01-01", + }, + ], + ) + _write_inference( + newer, + [ + { + "input_hash": "shared", + "output_text": "new", + "pushed_at": "2026-02-01", + "run_id": "run-b", + }, + { + "input_hash": "only-new", + "output_text": "new-only", + "pushed_at": "2026-02-01", + }, + ], + ) + _write_metadata( + older, + [ + { + "input_hash": "shared", + "metadata_hash": "meta-a", + "metadata_json": '{"a": 1}', + "observed_at": "2026-01-01", + } + ], + ) + _write_metadata( + newer, + [ + { + "input_hash": "shared", + "metadata_hash": "meta-b", + "metadata_json": '{"b": 2}', + "observed_at": "2026-02-01", + } + ], + ) + + merged = tmp_path / "merged.db" + store_sync._merge_dbs([older, newer], merged) + assert _read_outputs(merged) == { + "shared": "new", + "only-old": "old-only", + "only-new": "new-only", + } + assert _read_metadata(merged) == {("shared", "meta-a"), ("shared", "meta-b")} + + +def test_merge_dbs_uses_output_hash_tiebreaker(tmp_path): + left = tmp_path / "left.db" + right = tmp_path / "right.db" + _write_inference( + left, + [ + { + "input_hash": "shared", + "output_text": "zzz", + "pushed_at": "2026-01-01", + "run_id": "run-a", + } + ], + ) + _write_inference( + right, + [ + { + "input_hash": "shared", + "output_text": "aaa", + "pushed_at": "2026-01-01", + "run_id": "run-a", + } + ], + ) + merged = tmp_path / "merged.db" + store_sync._merge_dbs([left, right], merged) + assert _read_outputs(merged)["shared"] == "aaa" + + +def test_fetch_cell_merges_remote_into_local(fake_hub, tmp_path): + remote = tmp_path / "remote.db" + _write_inference( + remote, + [ + { + "input_hash": "remote", + "output_text": "R", + "pushed_at": "2026-01-01", + } + ], + ) + fake_hub.files[PATH_IN_REPO] = remote.read_bytes() + fake_hub.files[METADATA_IN_REPO] = json.dumps(CELL_CONFIG).encode("utf-8") + + local = _local_cell_db(tmp_path) + _write_inference( + local, + [ + { + "input_hash": "local", + "output_text": "L", + "pushed_at": "2026-01-01", + } + ], + ) + write_store_metadata(local.parent, CELL_CONFIG) + assert store_sync.fetch_cell(REPO_ID, PATH_IN_REPO, local) + assert _read_outputs(local) == {"local": "L", "remote": "R"} + + +def test_fetch_cell_rejects_metadata_mismatch_before_db_merge(fake_hub, tmp_path): + remote = tmp_path / "remote.db" + _write_inference( + remote, + [ + { + "input_hash": "remote", + "output_text": "R", + "pushed_at": "2026-01-01", + } + ], + ) + fake_hub.files[PATH_IN_REPO] = remote.read_bytes() + fake_hub.files[METADATA_IN_REPO] = json.dumps( + {**CELL_CONFIG, "task": "other"} + ).encode("utf-8") + + local = _local_cell_db(tmp_path) + _write_inference( + local, + [ + { + "input_hash": "local", + "output_text": "L", + "pushed_at": "2026-01-01", + } + ], + ) + write_store_metadata(local.parent, CELL_CONFIG) + + with pytest.raises(ValueError, match="does not match"): + store_sync.fetch_cell(REPO_ID, PATH_IN_REPO, local) + assert _read_outputs(local) == {"local": "L"} + + +def test_fetch_cell_rejects_misnamed_remote_cell_folder(fake_hub, tmp_path): + wrong_hash = "0000000000000000" + wrong_path = f"inference/arena/VLLM/Qwen%2Fjudge/{wrong_hash}/{INFERENCE_DB_NAME}" + wrong_metadata = f"inference/arena/VLLM/Qwen%2Fjudge/{wrong_hash}/metadata.json" + remote = tmp_path / "remote.db" + _write_inference( + remote, + [ + { + "input_hash": "remote", + "output_text": "R", + "pushed_at": "2026-01-01", + } + ], + ) + fake_hub.files[wrong_path] = remote.read_bytes() + fake_hub.files[wrong_metadata] = json.dumps(CELL_CONFIG).encode("utf-8") + + local = store_folder(tmp_path, "arena", MODEL_SPEC, wrong_hash) / INFERENCE_DB_NAME + with pytest.raises(ValueError, match="does not match descriptor hash"): + store_sync.fetch_cell(REPO_ID, wrong_path, local) + + +def test_fetch_cell_requires_remote_metadata(fake_hub, tmp_path): + remote = tmp_path / "remote.db" + _write_inference( + remote, + [ + { + "input_hash": "remote", + "output_text": "R", + "pushed_at": "2026-01-01", + } + ], + ) + fake_hub.files[PATH_IN_REPO] = remote.read_bytes() + + local = tmp_path / "store" / INFERENCE_DB_NAME + with pytest.raises(ValueError, match="missing required"): + store_sync.fetch_cell(REPO_ID, PATH_IN_REPO, local) + + +def test_fetch_cell_missing_remote_is_noop(fake_hub, tmp_path): + local = tmp_path / INFERENCE_DB_NAME + _write_inference( + local, + [ + { + "input_hash": "local", + "output_text": "L", + "pushed_at": "2026-01-01", + } + ], + ) + assert not store_sync.fetch_cell(REPO_ID, PATH_IN_REPO, local) + assert _read_outputs(local) == {"local": "L"} + + +def test_push_cell_retries_and_preserves_concurrent_rows(fake_hub, tmp_path): + fake_hub.head = "initial" + cell_dir = store_folder(tmp_path, "arena", MODEL_SPEC, CELL_CONFIG_HASH) + local = cell_dir / INFERENCE_DB_NAME + write_store_metadata(cell_dir, CELL_CONFIG) + _write_inference( + local, + [ + { + "input_hash": "local", + "output_text": "L", + "pushed_at": "2026-01-01", + } + ], + ) + concurrent = tmp_path / "concurrent.db" + _write_inference( + concurrent, + [ + { + "input_hash": "remote", + "output_text": "R", + "pushed_at": "2026-02-01", + } + ], + ) + + def inject_concurrent_write(): + fake_hub.files[PATH_IN_REPO] = concurrent.read_bytes() + fake_hub.head = "concurrent-head" + + fake_hub.pending_conflict = True + fake_hub.inject = inject_concurrent_write + store_sync.push_cell( + REPO_ID, + PATH_IN_REPO, + local, + pushed_by="alice", + ) + + assert fake_hub.commit_calls == 2 + uploaded = tmp_path / "uploaded.db" + uploaded.write_bytes(fake_hub.files[PATH_IN_REPO]) + expected = {"local": "L", "remote": "R"} + assert _read_outputs(uploaded) == expected + assert _read_outputs(local) == expected + assert METADATA_IN_REPO in fake_hub.files + + +def test_push_cell_requires_metadata(fake_hub, tmp_path): + fake_hub.head = "initial" + local = tmp_path / "local.db" + _write_inference( + local, + [ + { + "input_hash": "local", + "output_text": "L", + "pushed_at": "2026-01-01", + } + ], + ) + with pytest.raises(FileNotFoundError, match="Missing metadata.json"): + store_sync.push_cell(REPO_ID, PATH_IN_REPO, local, pushed_by="alice") + + +def test_push_cells_strict_rejects_missing_local_cell(tmp_path): + with pytest.raises(FileNotFoundError, match="Cache cell does not exist"): + store_sync.push_cells( + REPO_ID, + tmp_path, + [tmp_path / "missing" / INFERENCE_DB_NAME], + pushed_by="alice", + strict=True, + ) + + +def test_iter_cell_dbs_uses_segment_boundary_prefix(tmp_path): + arena_cell = ( + tmp_path + / "inference" + / "arena" + / "VLLM" + / "Qwen%2Fjudge" + / "abc123" + / INFERENCE_DB_NAME + ) + arena_cell.parent.mkdir(parents=True) + arena_cell.write_bytes(b"") + arena_hard_cell = ( + tmp_path + / "inference" + / "arena-hard-v2.0" + / "VLLM" + / "Qwen%2Fjudge" + / "def456" + / INFERENCE_DB_NAME + ) + arena_hard_cell.parent.mkdir(parents=True) + arena_hard_cell.write_bytes(b"") + + matched = store_sync.iter_cell_dbs(tmp_path, path_prefix="inference/arena") + assert matched == [arena_cell] + + +def test_iter_cell_dbs_ignores_noncanonical_layout(tmp_path): + invalid = tmp_path / "inference" / "too-shallow" / INFERENCE_DB_NAME + invalid.parent.mkdir(parents=True) + invalid.write_bytes(b"") + + assert store_sync.iter_cell_dbs(tmp_path) == [] + + +def test_push_cell_create_pr_updates_local_db_without_uploading_branch( + fake_hub, + tmp_path, +): + fake_hub.head = "initial" + cell_dir = store_folder(tmp_path, "arena", MODEL_SPEC, CELL_CONFIG_HASH) + local = cell_dir / INFERENCE_DB_NAME + write_store_metadata(cell_dir, CELL_CONFIG) + _write_inference( + local, + [ + { + "input_hash": "local", + "output_text": "L", + "pushed_at": "2026-01-01", + } + ], + ) + remote = tmp_path / "remote.db" + _write_inference( + remote, + [ + { + "input_hash": "remote", + "output_text": "R", + "pushed_at": "2026-02-01", + } + ], + ) + fake_hub.files[PATH_IN_REPO] = remote.read_bytes() + fake_hub.files[METADATA_IN_REPO] = json.dumps(CELL_CONFIG).encode("utf-8") + + result = store_sync.push_cell( + REPO_ID, + PATH_IN_REPO, + local, + pushed_by="alice", + create_pr=True, + ) + assert result == "https://hf.co/pr/1" + assert fake_hub.commit_calls == 1 + assert fake_hub.files[PATH_IN_REPO] == remote.read_bytes() + assert _read_outputs(local) == {"local": "L", "remote": "R"} + + +def test_fetch_remote_cells_bootstraps_empty_store_with_prefix(fake_hub, tmp_path): + remote = tmp_path / "remote.db" + _write_inference( + remote, + [ + { + "input_hash": "remote", + "output_text": "R", + "pushed_at": "2026-01-01", + } + ], + ) + fake_hub.files[PATH_IN_REPO] = remote.read_bytes() + fake_hub.files[METADATA_IN_REPO] = json.dumps(CELL_CONFIG).encode("utf-8") + fake_hub.head = "initial" + + store_root = tmp_path / "empty-store" + fetched = store_sync.fetch_remote_cells( + REPO_ID, + store_root, + path_prefix="inference/arena", + strict=True, + ) + assert len(fetched) == 1 + assert _read_outputs(fetched[0]) == {"remote": "R"} + assert json.loads((fetched[0].parent / "metadata.json").read_text()) == CELL_CONFIG + assert ( + store_sync.discover_remote_cell_dbs( + REPO_ID, + path_prefix="inference/other", + ) + == [] + ) + + +def test_fetch_remote_cells_rejects_noncanonical_remote_path(fake_hub, tmp_path): + path_in_repo = "inference/arena/VLLM/../../outside/hash/inference.db" + fake_hub.files[path_in_repo] = b"not-a-database" + fake_hub.head = "initial" + + with pytest.raises(ValueError, match="Invalid remote cache cell path"): + store_sync.fetch_remote_cells( + REPO_ID, + tmp_path / "store", + path_prefix="inference/arena", + strict=True, + ) + + assert not (tmp_path / "outside").exists() + + +def test_iter_cell_dbs_respects_prefix(tmp_path): + cell = ( + tmp_path + / "inference" + / "arena" + / "VLLM" + / "Qwen%2Fjudge" + / "abc123" + / INFERENCE_DB_NAME + ) + cell.parent.mkdir(parents=True) + cell.write_bytes(b"") + other = ( + tmp_path + / "inference" + / "other-task" + / "VLLM" + / "Model" + / "def456" + / INFERENCE_DB_NAME + ) + other.parent.mkdir(parents=True) + other.write_bytes(b"") + + assert len(store_sync.iter_cell_dbs(tmp_path)) == 2 + assert len(store_sync.iter_cell_dbs(tmp_path, path_prefix="inference/arena")) == 1 From 92d76689b6f5477b2147f28780545b9c17ce5394 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:52:40 +0200 Subject: [PATCH 09/13] feat(cache): cache at the inference boundary Add lazy prepared-model adapters and route inference through versioned, provider-aware descriptors without initializing backends on full hits. Includes-AI-Code: true --- judgearena/inference_cache.py | 285 +++++++++++ judgearena/model_adapters.py | 910 +++++++++++++++++++++++++++++++++ judgearena/models.py | 923 +++++++++++++++++++++------------- tests/test_chat_vllm.py | 19 + tests/test_inference_cache.py | 465 +++++++++++++++++ tests/test_model_adapters.py | 355 +++++++++++++ 6 files changed, 2606 insertions(+), 351 deletions(-) create mode 100644 judgearena/inference_cache.py create mode 100644 judgearena/model_adapters.py create mode 100644 tests/test_inference_cache.py create mode 100644 tests/test_model_adapters.py diff --git a/judgearena/inference_cache.py b/judgearena/inference_cache.py new file mode 100644 index 0000000..f92b400 --- /dev/null +++ b/judgearena/inference_cache.py @@ -0,0 +1,285 @@ +"""Cache-aware inference boundary backed by configuration-scoped SQLite cells.""" + +from __future__ import annotations + +import getpass +import uuid +from collections.abc import Callable, Sequence +from pathlib import Path +from typing import Any, Literal + +import pandas as pd + +from judgearena.log import get_logger +from judgearena.store_sqlite import ( + INFERENCE_DB_NAME, + SQLiteInferenceStore, + descriptor_hash, + stable_json_dumps, + store_folder, + write_store_metadata, +) +from judgearena.store_sync import DEFAULT_CACHE_REPO, fetch_cells, push_cells + +logger = get_logger(__name__) + +CacheMode = Literal["use", "off", "refresh"] +_VALID_CACHE_MODES = frozenset({"use", "off", "refresh"}) + + +class InferenceCache: + """Context manager for one run-scoped inference cache session.""" + + def __init__( + self, + store_root: Path | str, + task: str, + *, + mode: CacheMode = "use", + fetch: bool = False, + push: bool = False, + create_pr: bool = False, + cache_hf_repo: str = DEFAULT_CACHE_REPO, + pushed_by: str | None = None, + repo_type: str = "dataset", + revision: str = "main", + ) -> None: + if mode not in _VALID_CACHE_MODES: + raise ValueError( + f"Invalid cache mode {mode!r}; expected one of {sorted(_VALID_CACHE_MODES)}" + ) + self.store_root = Path(store_root).expanduser() + self.task = task + self.mode = mode + self.fetch = fetch + self.push = push + self.create_pr = create_pr + self.cache_hf_repo = cache_hf_repo + self.pushed_by = pushed_by or getpass.getuser() + self.repo_type = repo_type + self.revision = revision + self.run_id = str(uuid.uuid4()) + self._stores: dict[tuple[str, str], SQLiteInferenceStore] = {} + self._cell_folders: dict[tuple[str, str], Path] = {} + self._fetched_cells: set[tuple[str, str]] = set() + self._dirty_cells: set[tuple[str, str]] = set() + self._closed = False + + def _cell_key(self, model_spec: str, descriptor: dict[str, Any]) -> tuple[str, str]: + return model_spec, descriptor_hash(descriptor) + + def _cell_folder(self, model_spec: str, descriptor: dict[str, Any]) -> Path: + key = self._cell_key(model_spec, descriptor) + if key not in self._cell_folders: + config_hash = key[1] + self._cell_folders[key] = store_folder( + self.store_root, + self.task, + model_spec, + config_hash, + ) + return self._cell_folders[key] + + def _open_store( + self, model_spec: str, descriptor: dict[str, Any] + ) -> SQLiteInferenceStore: + key = self._cell_key(model_spec, descriptor) + if key not in self._stores: + folder = self._cell_folder(model_spec, descriptor) + write_store_metadata(folder, descriptor) + db_path = folder / INFERENCE_DB_NAME + self._stores[key] = SQLiteInferenceStore(db_path) + if self.fetch and key not in self._fetched_cells: + fetch_cells( + self.cache_hf_repo, + self.store_root, + [db_path], + repo_type=self.repo_type, + revision=self.revision, + strict=False, + ) + self._fetched_cells.add(key) + return self._stores[key] + + def get_or_run( + self, + *, + model_spec: str, + descriptor: dict[str, Any], + canonical_inputs: Sequence[str], + original_inputs: Sequence[Any], + miss_runner: Callable[[list[Any]], list[str]], + row_metadata: Sequence[dict[str, Any] | None] | None = None, + producer_metadata: dict[str, Any] | None = None, + ) -> list[str]: + """Return outputs in caller order, deduplicating identical canonical inputs.""" + if len(canonical_inputs) != len(original_inputs): + raise ValueError( + "canonical_inputs and original_inputs must have equal length" + ) + if not canonical_inputs: + return [] + + if row_metadata is not None and len(row_metadata) != len(original_inputs): + raise ValueError("row_metadata length must match original_inputs") + + if self.mode == "off": + outputs = miss_runner(list(original_inputs)) + if len(outputs) != len(original_inputs): + raise ValueError("miss_runner returned unexpected number of outputs") + return outputs + + input_hashes = [ + descriptor_hash(canonical, length=None) for canonical in canonical_inputs + ] + unique_order: list[str] = [] + seen_hashes: set[str] = set() + for input_hash in input_hashes: + if input_hash not in seen_hashes: + unique_order.append(input_hash) + seen_hashes.add(input_hash) + + hash_to_canonical = { + input_hash: canonical + for input_hash, canonical in zip( + input_hashes, canonical_inputs, strict=True + ) + } + hash_to_original: dict[str, Any] = {} + for input_hash, original in zip(input_hashes, original_inputs, strict=True): + hash_to_original.setdefault(input_hash, original) + + store = self._open_store(model_spec, descriptor) + cell_key = self._cell_key(model_spec, descriptor) + + if self.mode == "refresh": + missing_hashes = unique_order + cached_by_hash: dict[str, str] = {} + else: + missing_hashes = store.missing(unique_order) + cached_rows = store.query( + [h for h in unique_order if h not in missing_hashes] + ) + cached_by_hash = { + row["input_hash"]: row["output_text"] + for _, row in cached_rows.iterrows() + } + + metadata_df = self._row_metadata_frame( + input_hashes=input_hashes, + row_metadata=row_metadata, + ) + metadata_written = 0 + if missing_hashes: + miss_inputs = [hash_to_original[h] for h in missing_hashes] + new_outputs = miss_runner(miss_inputs) + if len(new_outputs) != len(missing_hashes): + raise ValueError("miss_runner returned unexpected number of outputs") + + producer_json = stable_json_dumps(producer_metadata or {}) + outputs_df = pd.DataFrame( + { + "input_hash": missing_hashes, + "input_text": [hash_to_canonical[h] for h in missing_hashes], + "output_text": new_outputs, + "producer_metadata_json": [producer_json] * len(missing_hashes), + } + ) + if metadata_df is not None: + _, metadata_written = store.save_outputs_and_metadata( + outputs_df, + metadata_df, + pushed_by=self.pushed_by, + run_id=self.run_id, + replace=self.mode == "refresh", + ) + else: + store.save_outputs( + outputs_df, + pushed_by=self.pushed_by, + run_id=self.run_id, + replace=self.mode == "refresh", + ) + self._dirty_cells.add(cell_key) + if self.mode == "refresh": + cached_by_hash.update( + dict(zip(missing_hashes, new_outputs, strict=True)) + ) + else: + cached_by_hash.update(store.outputs_by_hash(missing_hashes)) + + elif metadata_df is not None: + metadata_written = store.save_metadata(metadata_df, run_id=self.run_id) + if metadata_written: + self._dirty_cells.add(cell_key) + return [cached_by_hash[h] for h in input_hashes] + + @staticmethod + def _row_metadata_frame( + *, + input_hashes: list[str], + row_metadata: Sequence[dict[str, Any] | None] | None, + ) -> pd.DataFrame | None: + if not row_metadata: + return None + rows = [] + for input_hash, metadata in zip(input_hashes, row_metadata, strict=True): + if metadata is None: + continue + rows.append( + { + "input_hash": input_hash, + "metadata_json": stable_json_dumps(metadata), + } + ) + if not rows: + return None + return pd.DataFrame(rows) + + def close(self) -> None: + if self._closed: + return + self._closed = True + for store in self._stores.values(): + store.close() + self._stores.clear() + + def __enter__(self) -> InferenceCache: + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc: BaseException | None, + tb: object, + ) -> None: + dirty_cells = set(self._dirty_cells) + cell_folders = dict(self._cell_folders) + try: + self.close() + finally: + if exc_type is None and self.push and dirty_cells: + resolved_paths = [ + cell_folders[key] / INFERENCE_DB_NAME + for key in dirty_cells + if key in cell_folders + and (cell_folders[key] / INFERENCE_DB_NAME).exists() + ] + if resolved_paths: + try: + push_cells( + self.cache_hf_repo, + self.store_root, + resolved_paths, + pushed_by=self.pushed_by, + repo_type=self.repo_type, + revision=self.revision, + create_pr=self.create_pr, + strict=False, + ) + except Exception as push_exc: # noqa: BLE001 + logger.warning( + "Cache push failed after inference run %s: %s", + self.run_id, + push_exc, + ) diff --git a/judgearena/model_adapters.py b/judgearena/model_adapters.py new file mode 100644 index 0000000..46a6d78 --- /dev/null +++ b/judgearena/model_adapters.py @@ -0,0 +1,910 @@ +"""Cache descriptors, provider canonicalization, and PreparedModel adapters.""" + +from __future__ import annotations + +import importlib.metadata +import warnings +from collections.abc import Callable +from typing import Any + +from langchain_community.llms import LlamaCpp +from langchain_openai import ChatOpenAI +from langchain_openai.llms import OpenAI +from langchain_together.llms import Together + +from judgearena.store_sqlite import stable_json_dumps + +DESCRIPTOR_SCHEMA_VERSION = "judgearena-inference-descriptor/v1" +HOSTED_ADAPTER_VERSION = "judgearena-hosted-adapter/v1" +LOCAL_LLAMACPP_ADAPTER_VERSION = "judgearena-local-llamacpp-adapter/v1" + +_UNCACHED_MODEL_WARNINGS: set[type] = set() + +_SECRET_EXACT_KEYS = frozenset( + { + "api_key", + "openai_api_key", + "together_api_key", + "token", + "password", + "default_headers", + "headers", + } +) + +_DESCRIPTOR_KEY_DENYLIST = frozenset( + { + "callback", + "callbacks", + "callback_manager", + "client", + "async_client", + "http_client", + "http_async_client", + "root_client", + "root_async_client", + "cache", + "verbose", + "rate_limiter", + "streaming", + "disable_streaming", + "tags", + "name", + "metadata", + } +) + +_TRANSIENT_MESSAGE_KEYS = frozenset({"id", "response_metadata", "usage_metadata"}) + + +def top_k_from_settings(settings: dict[str, Any]) -> int | None: + """Extract ``top_k`` without treating zero as missing.""" + if "top_k" in settings: + value = settings["top_k"] + return None if value is None else int(value) + model_kwargs = settings.get("model_kwargs") + if isinstance(model_kwargs, dict) and "top_k" in model_kwargs: + value = model_kwargs["top_k"] + return None if value is None else int(value) + return None + + +def _is_secret_key(key: str) -> bool: + lowered = key.lower().replace("-", "_") + if lowered in _SECRET_EXACT_KEYS: + return True + blocked_fragments = ( + "api_key", + "auth", + "authorization", + "bearer", + "credential", + "cookie", + "password", + "secret", + "header", + ) + if any(fragment in lowered for fragment in blocked_fragments): + return True + return lowered.endswith("_token") or lowered.endswith("_key") + + +def _provider_package_version(distribution: str) -> str | None: + try: + return importlib.metadata.version(distribution) + except importlib.metadata.PackageNotFoundError: + return None + + +def normalize_descriptor_value(value: Any) -> Any | None: + """Return a JSON-safe value, or None when normalization is unsafe.""" + if value is None or isinstance(value, (bool, int, float, str)): + return value + if isinstance(value, (list, tuple)): + normalized_items = [] + for item in value: + normalized = normalize_descriptor_value(item) + if normalized is None and item is not None: + return None + normalized_items.append(normalized) + return normalized_items + if isinstance(value, (set, frozenset)): + normalized_items = [] + for item in value: + normalized = normalize_descriptor_value(item) + if normalized is None and item is not None: + return None + normalized_items.append(normalized) + return sorted(normalized_items, key=stable_json_dumps) + if isinstance(value, dict): + normalized_dict: dict[str, Any] = {} + for key, item in sorted(value.items(), key=lambda pair: str(pair[0])): + key_str = str(key) + if _is_secret_key(key_str): + continue + normalized = normalize_descriptor_value(item) + if normalized is None and item is not None: + return None + normalized_dict[key_str] = normalized + return normalized_dict + return None + + +def normalize_reasoning_config(value: Any) -> Any | None: + if hasattr(value, "model_dump"): + return normalize_descriptor_value(value.model_dump()) + if hasattr(value, "reasoning_start_str"): + return normalize_descriptor_value( + { + "reasoning_start_str": getattr(value, "reasoning_start_str", None), + "reasoning_end_str": getattr(value, "reasoning_end_str", None), + } + ) + return None + + +def normalize_constructor_settings(settings: dict[str, Any]) -> dict[str, Any] | None: + """Normalize every JSON-safe non-secret constructor setting conservatively.""" + normalized: dict[str, Any] = {} + for key in sorted(settings): + key_str = str(key) + if key_str in _DESCRIPTOR_KEY_DENYLIST or _is_secret_key(key_str): + continue + value = settings[key] + if value is None: + continue + if key_str == "reasoning_config": + normalized_value = normalize_reasoning_config(value) + else: + normalized_value = normalize_descriptor_value(value) + if normalized_value is None: + return None + normalized[key_str] = normalized_value + return normalized + + +def _resolved_sampling_from_kwargs(resolved: Any) -> dict[str, Any] | None: + sampling = normalize_descriptor_value(dict(resolved.sampling_params_kwargs)) + if sampling is None: + return None + sampling["max_tokens"] = resolved.max_tokens + return sampling + + +def build_vllm_descriptor(model_spec: str, resolved: Any) -> dict[str, Any] | None: + """Build a descriptor from fully resolved vLLM settings.""" + vllm_version = _provider_package_version("vllm") + if vllm_version is None: + return None + + engine_settings = normalize_constructor_settings( + { + **dict(resolved.vllm_kwargs), + "chat_template": resolved.chat_template, + "chat_template_kwargs": resolved.chat_template_kwargs, + } + ) + if engine_settings is None: + return None + + sampling = _resolved_sampling_from_kwargs(resolved) + if sampling is None: + return None + + return { + "descriptor_schema_version": DESCRIPTOR_SCHEMA_VERSION, + "provider": "VLLM", + "model_spec": model_spec, + "model": resolved.model_path, + "input_mode": resolved.input_mode, + "chat_template": resolved.chat_template, + "chat_template_kwargs": resolved.chat_template_kwargs, + "sampling": sampling, + "engine_settings": engine_settings, + "vllm_version": vllm_version, + } + + +def build_hosted_descriptor( + *, + provider: str, + model_spec: str, + model_name: str, + max_tokens: int | None, + input_mode: str, + base_url: str | None, + constructor_settings: dict[str, Any], + resolved_sampling: dict[str, Any], +) -> dict[str, Any] | None: + settings = normalize_constructor_settings(constructor_settings) + if settings is None: + return None + sampling_settings = dict(resolved_sampling) + if max_tokens is not None: + sampling_settings["max_tokens"] = max_tokens + sampling = normalize_descriptor_value(sampling_settings) + if sampling is None: + return None + + descriptor: dict[str, Any] = { + "descriptor_schema_version": DESCRIPTOR_SCHEMA_VERSION, + "provider": provider, + "model_spec": model_spec, + "model": model_name, + "input_mode": input_mode, + "hosted_adapter_version": HOSTED_ADAPTER_VERSION, + "server_defaults": "unobserved", + "sampling": sampling, + "engine_settings": settings, + } + if base_url is not None: + descriptor["base_url"] = base_url + return descriptor + + +def build_llamacpp_descriptor( + *, + model_spec: str, + model_name: str, + max_tokens: int, + constructor_settings: dict[str, Any], + resolved_sampling: dict[str, Any], +) -> dict[str, Any] | None: + llama_cpp_version = _provider_package_version("llama-cpp-python") + if llama_cpp_version is None: + return None + settings = normalize_constructor_settings(constructor_settings) + if settings is None: + return None + sampling = normalize_descriptor_value( + {**resolved_sampling, "max_tokens": max_tokens} + ) + if sampling is None: + return None + return { + "descriptor_schema_version": DESCRIPTOR_SCHEMA_VERSION, + "provider": "LlamaCpp", + "model_spec": model_spec, + "model": model_name, + "input_mode": "raw", + "local_adapter_version": LOCAL_LLAMACPP_ADAPTER_VERSION, + "llama_cpp_python_version": llama_cpp_version, + "sampling": sampling, + "engine_settings": settings, + } + + +def build_dummy_descriptor( + *, + model_spec: str, + max_tokens: int, + input_mode: str, + constructor_settings: dict[str, Any], + resolved_sampling: dict[str, Any], +) -> dict[str, Any] | None: + settings = normalize_constructor_settings(constructor_settings) + if settings is None: + return None + sampling = normalize_descriptor_value( + {**resolved_sampling, "max_tokens": max_tokens} + ) + if sampling is None: + return None + return { + "descriptor_schema_version": DESCRIPTOR_SCHEMA_VERSION, + "provider": "Dummy", + "model_spec": model_spec, + "model": model_spec.partition("/")[2], + "input_mode": input_mode, + "sampling": sampling, + "engine_settings": settings, + } + + +def build_producer_metadata(*, provider: str) -> dict[str, Any]: + metadata: dict[str, Any] = { + "provider": provider, + "descriptor_schema_version": DESCRIPTOR_SCHEMA_VERSION, + } + if provider == "VLLM": + version = _provider_package_version("vllm") + if version is not None: + metadata["vllm_version"] = version + elif provider in {"OpenRouter", "OpenAI", "ChatOpenAI"}: + metadata["hosted_adapter_version"] = HOSTED_ADAPTER_VERSION + version = _provider_package_version("langchain-openai") + if version is not None: + metadata["langchain_openai_version"] = version + elif provider == "Together": + metadata["hosted_adapter_version"] = HOSTED_ADAPTER_VERSION + version = _provider_package_version("langchain-together") + if version is not None: + metadata["langchain_together_version"] = version + elif provider == "LlamaCpp": + metadata["local_adapter_version"] = LOCAL_LLAMACPP_ADAPTER_VERSION + llama_version = _provider_package_version("llama-cpp-python") + if llama_version is not None: + metadata["llama_cpp_python_version"] = llama_version + lc_version = _provider_package_version("langchain-community") + if lc_version is not None: + metadata["langchain_community_version"] = lc_version + elif provider == "Dummy": + pass + return metadata + + +def effective_sampling( + *, + temperature: float | None, + top_p: float | None, + top_k: int | None, + seed: int | None, +) -> dict[str, Any]: + sampling: dict[str, Any] = {} + if temperature is not None: + sampling["temperature"] = float(temperature) + if top_p is not None: + sampling["top_p"] = float(top_p) + if top_k is not None: + sampling["top_k"] = int(top_k) + if seed is not None: + sampling["seed"] = int(seed) + return sampling + + +def _vllm_role(message: Any) -> str: + role_map = {"human": "user", "ai": "assistant", "system": "system"} + if isinstance(message, dict): + role = str(message.get("role", "user")) + elif hasattr(message, "type"): + role = str(message.type) + elif isinstance(message, tuple) and message: + role = str(message[0]) + else: + role = "user" + return role_map.get(role, role) + + +def _vllm_content(message: Any) -> Any: + if isinstance(message, dict): + return message.get("content", "") + if isinstance(message, tuple) and len(message) > 1: + return message[1] + return getattr(message, "content", "") + + +def vllm_input_to_messages(input_item: Any) -> list[dict[str, Any]]: + role_map = {"human": "user", "ai": "assistant", "system": "system"} + if hasattr(input_item, "to_messages"): + lc_messages = input_item.to_messages() + return [ + {"role": role_map.get(msg.type, msg.type), "content": msg.content} + for msg in lc_messages + ] + if isinstance(input_item, list) and input_item and isinstance(input_item[0], tuple): + return [ + {"role": role_map.get(role, role), "content": content} + for role, content in input_item + ] + if isinstance(input_item, list) and input_item and isinstance(input_item[0], dict): + return [ + { + **message, + "role": role_map.get( + message.get("role") or "user", message.get("role") or "user" + ), + } + for message in input_item + ] + if isinstance(input_item, str): + return [{"role": "user", "content": input_item}] + raise ValueError(f"Unsupported input type: {type(input_item)}") + + +def vllm_input_to_raw_text(input_item: Any) -> str: + if isinstance(input_item, str): + return input_item + if hasattr(input_item, "to_string"): + return input_item.to_string() + if isinstance(input_item, list) and input_item and isinstance(input_item[0], dict): + return "\n".join(str(msg["content"]) for msg in input_item) + raise ValueError(f"Cannot extract raw text from: {type(input_item)}") + + +def canonicalize_vllm_input(item: Any, *, input_mode: str) -> str: + if input_mode == "raw": + payload = {"kind": "raw", "text": vllm_input_to_raw_text(item)} + else: + if isinstance(item, str): + messages = [{"role": "user", "content": item}] + else: + messages = vllm_input_to_messages(item) + payload = {"kind": "chat", "messages": messages} + return stable_json_dumps(payload) + + +def _canonicalize_message_field(value: Any) -> Any | None: + return normalize_descriptor_value(value) + + +def _canonicalize_hosted_message(message: Any) -> dict[str, Any] | None: + if isinstance(message, dict): + raw = { + key: value + for key, value in message.items() + if key not in _TRANSIENT_MESSAGE_KEYS + } + role = raw.pop("role", "user") + if role == "human": + role = "user" + entry: dict[str, Any] = {"role": role} + for key, value in raw.items(): + normalized = _canonicalize_message_field(value) + if normalized is None and value is not None: + return None + if normalized is not None: + entry[key] = normalized + return entry + + entry: dict[str, Any] = {"role": _vllm_role(message)} + for field in ("content", "name", "additional_kwargs", "tool_calls", "tool_call_id"): + if hasattr(message, field): + value = getattr(message, field) + if value in (None, {}, []): + continue + normalized = _canonicalize_message_field(value) + if normalized is None: + return None + entry[field] = normalized + return entry + + +def canonicalize_hosted_chat_input(item: Any) -> str: + if isinstance(item, str): + messages = [{"role": "user", "content": item}] + elif hasattr(item, "to_messages"): + messages = [] + for message in item.to_messages(): + canonical = _canonicalize_hosted_message(message) + if canonical is None: + raise ValueError(f"Unsupported hosted chat message: {type(message)!r}") + messages.append(canonical) + elif isinstance(item, list) and item: + messages = [] + for message in item: + canonical = _canonicalize_hosted_message(message) + if canonical is None: + raise ValueError(f"Unsupported hosted chat message: {type(message)!r}") + messages.append(canonical) + else: + raise ValueError(f"Unsupported hosted chat input type: {type(item)!r}") + return stable_json_dumps({"kind": "chat", "messages": messages}) + + +def canonicalize_raw_input(item: Any) -> str: + if isinstance(item, str): + text = item + elif hasattr(item, "to_string"): + text = item.to_string() + else: + text = vllm_input_to_raw_text(item) + return stable_json_dumps({"kind": "raw", "text": text}) + + +def canonicalize_dummy_input(item: Any, *, input_mode: str) -> str: + if input_mode == "raw": + return canonicalize_raw_input(item) + return canonicalize_hosted_chat_input(item) + + +class PreparedModel: + """Lazy model wrapper with cache descriptors and deferred backend init.""" + + def __init__( + self, + *, + provider: str, + model_spec: str, + model_name: str, + max_tokens: int | None, + engine_kwargs: dict[str, Any], + sampling: dict[str, Any], + input_mode: str, + descriptor: dict[str, Any] | None, + materialize: Callable[[PreparedModel], Any], + producer_metadata: dict[str, Any], + base_url: str | None = None, + vllm_resolved: Any | None = None, + ) -> None: + self.provider = provider + self.model_spec = model_spec + self.model_name = model_name + self.max_tokens = max_tokens + self.engine_kwargs = dict(engine_kwargs) + self.sampling = dict(sampling) + self.input_mode = input_mode + self._descriptor = descriptor + self._materialize_fn = materialize + self._producer_metadata = dict(producer_metadata) + self.base_url = base_url + self._vllm_resolved = vllm_resolved + self._backend: Any | None = None + + def cache_descriptor(self) -> dict[str, Any] | None: + return None if self._descriptor is None else dict(self._descriptor) + + def producer_metadata(self) -> dict[str, Any]: + return dict(self._producer_metadata) + + def canonicalize_input(self, item: Any) -> str: + if self.provider == "VLLM": + return canonicalize_vllm_input(item, input_mode=self.input_mode) + if self.provider == "Dummy": + return canonicalize_dummy_input(item, input_mode=self.input_mode) + if self.input_mode == "raw": + return canonicalize_raw_input(item) + return canonicalize_hosted_chat_input(item) + + def _sync_descriptor_sampling_field(self, field: str, value: Any) -> None: + if self._descriptor is None: + return + self._descriptor.setdefault("sampling", {})[field] = value + engine_settings = self._descriptor.get("engine_settings") + if self.provider != "VLLM" and isinstance(engine_settings, dict): + engine_settings[field] = value + + def set_temperature(self, temperature: float) -> None: + value = float(temperature) + self.sampling["temperature"] = value + self.engine_kwargs["temperature"] = value + self._sync_descriptor_sampling_field("temperature", value) + if self._vllm_resolved is not None: + self._vllm_resolved.sampling_params_kwargs["temperature"] = value + if self._backend is not None: + if hasattr(self._backend, "set_temperature"): + self._backend.set_temperature(value) + elif hasattr(self._backend, "temperature"): + self._backend.temperature = value + + def materialize(self) -> Any: + if self._backend is None: + self._backend = self._materialize_fn(self) + return self._backend + + def batch(self, inputs: list, **invoke_kwargs) -> list[str]: + return self.materialize().batch(inputs, **invoke_kwargs) + + def invoke(self, input_item, **invoke_kwargs) -> str: + return self.materialize().invoke(input_item, **invoke_kwargs) + + async def ainvoke(self, input_item, **invoke_kwargs): + return await self.materialize().ainvoke(input_item, **invoke_kwargs) + + def __getattr__(self, name: str) -> Any: + if name.startswith("_"): + raise AttributeError(name) + return getattr(self.materialize(), name) + + +class CachedModelAdapter(PreparedModel): + """Cache adapter wrapping an already-constructed backend.""" + + def __init__( + self, + *, + backend: Any, + provider: str, + model_spec: str, + model_name: str, + max_tokens: int | None, + engine_kwargs: dict[str, Any], + sampling: dict[str, Any], + input_mode: str, + descriptor: dict[str, Any] | None, + producer_metadata: dict[str, Any], + base_url: str | None = None, + vllm_resolved: Any | None = None, + ) -> None: + super().__init__( + provider=provider, + model_spec=model_spec, + model_name=model_name, + max_tokens=max_tokens, + engine_kwargs=engine_kwargs, + sampling=sampling, + input_mode=input_mode, + descriptor=descriptor, + producer_metadata=producer_metadata, + base_url=base_url, + vllm_resolved=vllm_resolved, + materialize=lambda prepared: backend, + ) + self._backend = backend + + +def _warn_uncached_model_type(chat_model: Any) -> None: + model_type = type(chat_model) + if model_type in _UNCACHED_MODEL_WARNINGS: + return + warnings.warn( + f"Inference cache skipped for unsupported model type {model_type.__name__}; " + "running uncached.", + stacklevel=3, + ) + _UNCACHED_MODEL_WARNINGS.add(model_type) + + +def resolve_hosted_base_url(settings: dict[str, Any]) -> str | None: + """Return the hosted endpoint URL from constructor settings.""" + for key in ("base_url", "openai_api_base"): + value = settings.get(key) + if value: + return str(value) + return None + + +def hosted_provider_for_endpoint( + base_url: str | None, + default_provider: str, +) -> str: + """Return OpenRouter when the endpoint targets OpenRouter.""" + if base_url and "openrouter.ai" in base_url: + return "OpenRouter" + return default_provider + + +def _extract_langchain_constructor_settings(model: Any) -> dict[str, Any]: + settings: dict[str, Any] = {} + model_fields = getattr(type(model), "model_fields", {}) + fields_set = getattr(model, "model_fields_set", None) + fields_to_capture = fields_set if isinstance(fields_set, set) else model_fields + canonical_names = { + "model_name": "model", + "openai_api_base": "base_url", + } + for field in fields_to_capture: + if field in _DESCRIPTOR_KEY_DENYLIST or _is_secret_key(field): + continue + value = getattr(model, field, None) + if value is None: + continue + if field == "model_kwargs" and not value: + continue + settings[canonical_names.get(field, field)] = value + model_kwargs = getattr(model, "model_kwargs", None) + if isinstance(model_kwargs, dict) and model_kwargs: + settings["model_kwargs"] = dict(model_kwargs) + return settings + + +def _optional_max_tokens(model: Any) -> int | None: + value = getattr(model, "max_tokens", None) + return None if value is None else int(value) + + +def adapt_dummy_backend( + model: Any, *, model_spec: str | None = None +) -> PreparedModel | None: + """Build a cache adapter for a constructed DummyModel backend.""" + name = getattr(model, "name", None) + init_kwargs = getattr(model, "init_kwargs", None) + if name is None or init_kwargs is None: + return None + provider = "Dummy" + spec = model_spec or name + _provider, _, model_name = spec.partition("/") + settings = dict(init_kwargs) + sampling = effective_sampling( + temperature=settings.get("temperature"), + top_p=settings.get("top_p"), + top_k=top_k_from_settings(settings), + seed=settings.get("seed"), + ) + descriptor = build_dummy_descriptor( + model_spec=spec, + max_tokens=int(settings.get("max_tokens", 8192)), + input_mode="chat", + constructor_settings=settings, + resolved_sampling=sampling, + ) + return CachedModelAdapter( + backend=model, + provider=provider, + model_spec=spec, + model_name=model_name or spec, + max_tokens=int(settings.get("max_tokens", 8192)), + engine_kwargs=settings, + sampling=sampling, + input_mode="chat", + descriptor=descriptor, + producer_metadata=build_producer_metadata(provider=provider), + ) + + +def adapt_vllm_backend( + model: Any, *, model_spec: str | None = None +) -> PreparedModel | None: + """Build a cache adapter for a constructed ChatVLLM backend.""" + model_path = getattr(model, "model_path", None) + resolved = getattr(model, "_resolved", None) + max_tokens = getattr(model, "max_tokens", None) + if model_path is None or resolved is None or max_tokens is None: + return None + spec = model_spec or f"VLLM/{model_path}" + descriptor = build_vllm_descriptor(spec, resolved) + if descriptor is None: + return None + return CachedModelAdapter( + backend=model, + provider="VLLM", + model_spec=spec, + model_name=model_path, + max_tokens=max_tokens, + engine_kwargs=dict(resolved.vllm_kwargs), + sampling=dict(resolved.sampling_params_kwargs), + input_mode=resolved.input_mode, + descriptor=descriptor, + producer_metadata=build_producer_metadata(provider="VLLM"), + vllm_resolved=resolved, + ) + + +def wrap_known_model( + model: Any, *, model_spec: str | None = None +) -> PreparedModel | None: + """Return a cache adapter for a known constructed backend, or None.""" + if isinstance(model, PreparedModel): + return model + + to_adapter = getattr(model, "to_prepared_cache_adapter", None) + if callable(to_adapter): + return to_adapter(model_spec=model_spec) + + if isinstance(model, ChatOpenAI): + settings = _extract_langchain_constructor_settings(model) + base_url = resolve_hosted_base_url(settings) + provider = hosted_provider_for_endpoint(base_url, "ChatOpenAI") + spec = model_spec or f"{provider}/{model.model_name}" + max_tokens = _optional_max_tokens(model) + sampling = effective_sampling( + temperature=settings.get("temperature"), + top_p=settings.get("top_p"), + top_k=top_k_from_settings(settings), + seed=settings.get("seed"), + ) + descriptor = build_hosted_descriptor( + provider=provider, + model_spec=spec, + model_name=model.model_name, + max_tokens=max_tokens, + input_mode="chat", + base_url=base_url, + constructor_settings=settings, + resolved_sampling=sampling, + ) + return CachedModelAdapter( + backend=model, + provider=provider, + model_spec=spec, + model_name=model.model_name, + max_tokens=max_tokens, + engine_kwargs=settings, + sampling=sampling, + input_mode="chat", + descriptor=descriptor, + producer_metadata=build_producer_metadata(provider=provider), + base_url=base_url, + ) + + if isinstance(model, OpenAI): + spec = model_spec or f"OpenAI/{model.model_name}" + settings = _extract_langchain_constructor_settings(model) + base_url = resolve_hosted_base_url(settings) + max_tokens = _optional_max_tokens(model) + sampling = effective_sampling( + temperature=settings.get("temperature"), + top_p=settings.get("top_p"), + top_k=None, + seed=settings.get("seed"), + ) + descriptor = build_hosted_descriptor( + provider="OpenAI", + model_spec=spec, + model_name=model.model_name, + max_tokens=max_tokens, + input_mode="raw", + base_url=base_url, + constructor_settings=settings, + resolved_sampling=sampling, + ) + return CachedModelAdapter( + backend=model, + provider="OpenAI", + model_spec=spec, + model_name=model.model_name, + max_tokens=max_tokens, + engine_kwargs=settings, + sampling=sampling, + input_mode="raw", + descriptor=descriptor, + producer_metadata=build_producer_metadata(provider="OpenAI"), + base_url=base_url, + ) + + if isinstance(model, Together): + spec = model_spec or f"Together/{model.model}" + settings = _extract_langchain_constructor_settings(model) + base_url = resolve_hosted_base_url(settings) + max_tokens = _optional_max_tokens(model) + sampling = effective_sampling( + temperature=settings.get("temperature"), + top_p=settings.get("top_p"), + top_k=top_k_from_settings(settings), + seed=None, + ) + descriptor = build_hosted_descriptor( + provider="Together", + model_spec=spec, + model_name=model.model, + max_tokens=max_tokens, + input_mode="raw", + base_url=base_url, + constructor_settings=settings, + resolved_sampling=sampling, + ) + return CachedModelAdapter( + backend=model, + provider="Together", + model_spec=spec, + model_name=model.model, + max_tokens=max_tokens, + engine_kwargs=settings, + sampling=sampling, + input_mode="raw", + descriptor=descriptor, + producer_metadata=build_producer_metadata(provider="Together"), + base_url=base_url, + ) + + if isinstance(model, LlamaCpp): + spec = model_spec or f"LlamaCpp/{getattr(model, 'model_path', 'unknown')}" + settings = _extract_langchain_constructor_settings(model) + sampling = effective_sampling( + temperature=settings.get("temperature"), + top_p=settings.get("top_p"), + top_k=settings.get("top_k"), + seed=settings.get("seed"), + ) + descriptor = build_llamacpp_descriptor( + model_spec=spec, + model_name=str(getattr(model, "model_path", "unknown")), + max_tokens=int(settings.get("max_tokens", 8192) or 8192), + constructor_settings=settings, + resolved_sampling=sampling, + ) + return CachedModelAdapter( + backend=model, + provider="LlamaCpp", + model_spec=spec, + model_name=str(getattr(model, "model_path", "unknown")), + max_tokens=int(settings.get("max_tokens", 8192) or 8192), + engine_kwargs=settings, + sampling=sampling, + input_mode="raw", + descriptor=descriptor, + producer_metadata=build_producer_metadata(provider="LlamaCpp"), + ) + + return None + + +def resolve_cacheable_model( + chat_model: Any, + *, + model_spec: str | None = None, +) -> PreparedModel | None: + wrapped = wrap_known_model(chat_model, model_spec=model_spec) + if wrapped is not None: + return wrapped + _warn_uncached_model_type(chat_model) + return None diff --git a/judgearena/models.py b/judgearena/models.py index c0772ae..e2cc23d 100644 --- a/judgearena/models.py +++ b/judgearena/models.py @@ -7,20 +7,44 @@ import os import time import warnings +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any from langchain_community.llms import LlamaCpp from langchain_openai import ChatOpenAI +from langchain_openai.llms import OpenAI +from langchain_together.llms import Together from tqdm.asyncio import tqdm from tqdm.contrib.logging import logging_redirect_tqdm from judgearena.constants import VLLM_REASONING_END_STR, VLLM_REASONING_START_STR from judgearena.log import get_logger +from judgearena.model_adapters import ( + DESCRIPTOR_SCHEMA_VERSION, + HOSTED_ADAPTER_VERSION, + PreparedModel, + adapt_dummy_backend, + adapt_vllm_backend, + build_dummy_descriptor, + build_hosted_descriptor, + build_llamacpp_descriptor, + build_producer_metadata, + build_vllm_descriptor, + effective_sampling, + hosted_provider_for_endpoint, + resolve_cacheable_model, + resolve_hosted_base_url, + top_k_from_settings, + vllm_input_to_messages, + vllm_input_to_raw_text, +) from judgearena.utils.io import safe_parse_int logger = get_logger(__name__) - DEFAULT_VLLM_JUDGE_THINKING_TOKEN_BUDGET = 512 +OPENROUTER_BASE_URL = "https://openrouter.ai/api/v1" _THINKING_MODEL_PARSER_BY_SUBSTRING = ( ("qwen3", "qwen3"), ("smollm3", "qwen3"), @@ -46,14 +70,7 @@ def _split_model_spec(model_spec: str) -> tuple[str, str]: def is_thinking_model(model_name: str) -> bool: - """Return True for reasoning models that emit `...` traces. - - Covers the Qwen3 family (e.g. `Qwen/Qwen3.5-9B`) and SmolLM3 (e.g. - `HuggingFaceTB/SmolLM3-3B`) plus `allenai/Olmo-3-7B-Think`; all emit - ``/`` traces so vLLM's budget enforcement and our - tag-stripping apply uniformly. Matching is case-insensitive to tolerate - mixed-case HF repo ids like `HuggingFaceTB/SmolLM3-3B`. - """ + """Return True for reasoning models that emit thinking traces.""" return _default_reasoning_parser_for_model(model_name) is not None @@ -83,9 +100,6 @@ def build_default_judge_model_kwargs( judge_model_kwargs["thinking_token_budget"] = ( DEFAULT_VLLM_JUDGE_THINKING_TOKEN_BUDGET ) - # FP8 weights leave little KV headroom on consumer-class GPUs; default - # to FP8 KV cache so judges like Skywork-70B-FP8 fit comfortably on - # 2x L40S at 32k context. Explicit caller overrides still win. if "kv_cache_dtype" not in judge_model_kwargs and "fp8" in model_name.lower(): judge_model_kwargs["kv_cache_dtype"] = "fp8" return judge_model_kwargs @@ -99,33 +113,29 @@ def _resolve_chat_template_kwargs( chat_template_kwargs = dict(explicit_chat_template_kwargs or {}) if disable_thinking and "enable_thinking" not in chat_template_kwargs: chat_template_kwargs["enable_thinking"] = False - return chat_template_kwargs or None +def _hf_from_pretrained_kwargs(engine_kwargs: dict[str, Any]) -> dict[str, Any]: + kwargs: dict[str, Any] = { + "trust_remote_code": bool(engine_kwargs.get("trust_remote_code", True)), + } + if engine_kwargs.get("revision") is not None: + kwargs["revision"] = engine_kwargs["revision"] + tokenizer_revision = engine_kwargs.get("tokenizer_revision") + if tokenizer_revision is not None: + kwargs["tokenizer_revision"] = tokenizer_revision + return kwargs + + def _is_retryable_error(e: Exception) -> bool: - """Return True if the exception is a transient server error that should be retried. - - Handles three formats: - - String representation contains the HTTP code (most providers) - - ValueError raised by langchain-openai with a dict arg: {'message': ..., 'code': 429} - - json.JSONDecodeError when a gateway returns a non-JSON (e.g. HTML) body - """ - # langchain-openai raises ValueError(response_dict.get("error")) where the - # error value is a dict like {'message': '...', 'code': 408} _RETRYABLE_CODES = {408, 429, 502, 503, 504} if isinstance(e, ValueError) and e.args: arg = e.args[0] if isinstance(arg, dict) and arg.get("code") in _RETRYABLE_CODES: return True - - # Gateways (e.g. OpenRouter) intermittently return non-JSON error bodies that - # surface as JSONDecodeError while parsing the response. These are transient - # and safe to retry. JSONDecodeError subclasses ValueError but carries a - # string arg, so it bypasses the dict-code check above. if isinstance(e, json.JSONDecodeError): return True - error_str = str(e) return ( any(str(code) in error_str for code in _RETRYABLE_CODES) @@ -176,21 +186,183 @@ def invoke(self, input, **invoke_kwargs) -> str: async def ainvoke(self, input, **invoke_kwargs): return self.message + def set_temperature(self, temperature: float) -> None: + self.init_kwargs["temperature"] = float(temperature) + + def to_prepared_cache_adapter( + self, + *, + model_spec: str | None = None, + ) -> PreparedModel | None: + """Return a cache adapter for this constructed Dummy backend.""" + return adapt_dummy_backend(self, model_spec=model_spec) + + +@dataclass +class VLLMResolvedSettings: + """Shared, lazy-resolved vLLM configuration for descriptors and ChatVLLM.""" + + model_path: str + max_tokens: int + chat_template: str | None + chat_template_kwargs: dict[str, object] | None + use_generate: bool + input_mode: str + sampling_params_kwargs: dict[str, Any] + vllm_kwargs: dict[str, Any] + explicit_reasoning_settings: bool = False + + +def resolve_vllm_settings( + model: str, + max_tokens: int = 8192, + chat_template: str | None = None, + temperature: float | None = None, + top_p: float | None = None, + top_k: int | None = None, + seed: int | None = None, + **vllm_kwargs: Any, +) -> VLLMResolvedSettings: + """Resolve vLLM chat/raw mode and sampling without loading ``vllm.LLM``.""" + from vllm.config.reasoning import ReasoningConfig + + engine_kwargs = dict(vllm_kwargs) + disable_thinking = bool(engine_kwargs.pop("disable_thinking", False)) + thinking_token_budget = engine_kwargs.pop("thinking_token_budget", None) + explicit_chat_template_kwargs = engine_kwargs.pop("chat_template_kwargs", None) + explicit_reasoning_settings = ( + "reasoning_parser" in engine_kwargs or "reasoning_config" in engine_kwargs + ) + chat_template_kwargs = _resolve_chat_template_kwargs( + explicit_chat_template_kwargs=explicit_chat_template_kwargs, + disable_thinking=disable_thinking, + ) + hf_kwargs = _hf_from_pretrained_kwargs(engine_kwargs) + config_hf_kwargs = { + key: hf_kwargs[key] + for key in ("trust_remote_code", "revision") + if key in hf_kwargs + } + tokenizer_hf_kwargs = dict(config_hf_kwargs) + if "tokenizer_revision" in hf_kwargs: + tokenizer_hf_kwargs["revision"] = hf_kwargs["tokenizer_revision"] + + max_model_len = engine_kwargs.get("max_model_len") + if max_model_len is not None: + try: + from transformers import AutoConfig + + config = AutoConfig.from_pretrained(model, **config_hf_kwargs) + model_max_pos = getattr(config, "max_position_embeddings", None) + if model_max_pos is not None and max_model_len > model_max_pos: + warnings.warn( + f"Capping max_model_len from {max_model_len} to " + f"{model_max_pos} (max_position_embeddings) for '{model}'.", + stacklevel=2, + ) + engine_kwargs["max_model_len"] = model_max_pos + except Exception as exc: + warnings.warn( + "Could not validate max_model_len against " + f"max_position_embeddings for '{model}': {exc}. " + "Proceeding without clamping; vLLM may raise if the value is too large.", + RuntimeWarning, + stacklevel=2, + ) + + if seed is not None: + engine_kwargs.setdefault("seed", int(seed)) + + sampling_params_kwargs: dict[str, Any] = { + "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), + } + if top_k is not None: + sampling_params_kwargs["top_k"] = int(top_k) + if seed is not None: + sampling_params_kwargs["seed"] = int(seed) + + if thinking_token_budget is not None: + if max_tokens is not None: + thinking_token_budget = min(int(thinking_token_budget), int(max_tokens)) + if explicit_reasoning_settings: + sampling_params_kwargs["thinking_token_budget"] = int(thinking_token_budget) + elif is_thinking_model(model): + reasoning_parser = _default_reasoning_parser_for_model(model) + assert reasoning_parser is not None + engine_kwargs.setdefault( + "reasoning_config", + ReasoningConfig( + reasoning_start_str=VLLM_REASONING_START_STR, + reasoning_end_str=VLLM_REASONING_END_STR, + ), + ) + engine_kwargs.setdefault("reasoning_parser", reasoning_parser) + sampling_params_kwargs["thinking_token_budget"] = int(thinking_token_budget) + else: + warnings.warn( + f"Model '{model}' is not in JudgeArena's built-in thinking-model " + "defaults (Qwen3/SmolLM3/Olmo-3-7B-Think). Ignoring " + "thinking_token_budget unless reasoning_parser or " + "reasoning_config is provided explicitly.", + stacklevel=2, + ) + + if chat_template: + use_generate = False + input_mode = "chat" + effective_chat_template: str | None = chat_template + else: + try: + from transformers import AutoTokenizer + + tokenizer_id = engine_kwargs.get("tokenizer", model) + tokenizer = AutoTokenizer.from_pretrained( + tokenizer_id, + **tokenizer_hf_kwargs, + ) + has_template = bool(getattr(tokenizer, "chat_template", None)) + except Exception as exc: + raise RuntimeError( + f"Could not resolve chat template metadata for '{model}': {exc}" + ) from exc + if not has_template: + warnings.warn( + f"Model '{model}' tokenizer does not define a chat template. " + "Falling back to llm.generate() (no chat formatting). " + "Override with --chat_template if this model needs one.", + stacklevel=2, + ) + use_generate = True + input_mode = "raw" + effective_chat_template = None + if disable_thinking: + warnings.warn( + f"Model '{model}' has no chat template, so disable_thinking " + "cannot be applied when falling back to llm.generate().", + stacklevel=2, + ) + else: + use_generate = False + input_mode = "chat" + effective_chat_template = None + + return VLLMResolvedSettings( + model_path=model, + max_tokens=max_tokens, + chat_template=effective_chat_template, + chat_template_kwargs=chat_template_kwargs, + use_generate=use_generate, + input_mode=input_mode, + sampling_params_kwargs=sampling_params_kwargs, + vllm_kwargs=engine_kwargs, + explicit_reasoning_settings=explicit_reasoning_settings, + ) + class ChatVLLM: - """VLLM wrapper that auto-detects whether to use chat() or generate(). - - Chat template handling: - - If ``chat_template`` is explicitly provided, always uses ``llm.chat()`` - with that template (useful for models whose tokenizer lacks a template - but you know the correct one). - - If the tokenizer defines a chat template, uses ``llm.chat()`` and lets - vLLM apply the tokenizer's template automatically. - - If no chat template is found (typical for base/pretrained models), - falls back to ``llm.generate()`` and emits a warning. This avoids the - ``ValueError`` raised by ``transformers >= v4.44`` which removed the - default chat template. - """ + """VLLM wrapper that auto-detects whether to use chat() or generate().""" def __init__( self, @@ -201,193 +373,73 @@ def __init__( top_p: float | None = None, top_k: int | None = None, seed: int | None = None, + *, + _resolved: VLLMResolvedSettings | None = None, **vllm_kwargs, ): from vllm import LLM, SamplingParams - from vllm.config.reasoning import ReasoningConfig - - self.model_path = model - self.max_tokens = max_tokens - disable_thinking = bool(vllm_kwargs.pop("disable_thinking", False)) - thinking_token_budget = vllm_kwargs.pop("thinking_token_budget", None) - explicit_chat_template_kwargs = vllm_kwargs.pop("chat_template_kwargs", None) - explicit_reasoning_settings = ( - "reasoning_parser" in vllm_kwargs or "reasoning_config" in vllm_kwargs - ) - self._chat_template_kwargs = _resolve_chat_template_kwargs( - explicit_chat_template_kwargs=explicit_chat_template_kwargs, - disable_thinking=disable_thinking, - ) - - # Cap max_model_len to the model's max_position_embeddings so that - # vLLM doesn't reject an overly large context window. - max_model_len = vllm_kwargs.get("max_model_len") - if max_model_len is not None: - try: - from transformers import AutoConfig - - config = AutoConfig.from_pretrained(model, trust_remote_code=True) - model_max_pos = getattr(config, "max_position_embeddings", None) - if model_max_pos is not None and max_model_len > model_max_pos: - warnings.warn( - f"Capping max_model_len from {max_model_len} to " - f"{model_max_pos} (max_position_embeddings) for '{model}'.", - stacklevel=2, - ) - vllm_kwargs["max_model_len"] = model_max_pos - except Exception as e: - warnings.warn( - "Could not validate max_model_len against " - f"max_position_embeddings for '{model}': {e}. " - "Proceeding without clamping; vLLM may raise if the value is too large.", - RuntimeWarning, - stacklevel=2, - ) - if seed is not None: - vllm_kwargs.setdefault("seed", int(seed)) + if _resolved is None: + _resolved = resolve_vllm_settings( + model, + max_tokens=max_tokens, + chat_template=chat_template, + temperature=temperature, + top_p=top_p, + top_k=top_k, + seed=seed, + **vllm_kwargs, + ) - 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), - } - if top_k is not None: - self._sampling_params_kwargs["top_k"] = int(top_k) - if seed is not None: - self._sampling_params_kwargs["seed"] = int(seed) - if thinking_token_budget is not None: - if max_tokens is not None: - thinking_token_budget = min(int(thinking_token_budget), int(max_tokens)) - if explicit_reasoning_settings: - self._sampling_params_kwargs["thinking_token_budget"] = int( - thinking_token_budget - ) - elif is_thinking_model(model): - reasoning_parser = _default_reasoning_parser_for_model(model) - assert reasoning_parser is not None # guarded by is_thinking_model() - vllm_kwargs.setdefault( - "reasoning_config", - ReasoningConfig( - reasoning_start_str=VLLM_REASONING_START_STR, - # Shared forced end marker so vLLM can enforce the - # thinking-token budget for offline `LLM.chat()`. The - # parser itself still varies by model family (e.g. OLMo - # uses `olmo3`) on vLLM's OpenAI server. - reasoning_end_str=VLLM_REASONING_END_STR, - ), - ) - vllm_kwargs.setdefault("reasoning_parser", reasoning_parser) - self._sampling_params_kwargs["thinking_token_budget"] = int( - thinking_token_budget - ) - else: - warnings.warn( - f"Model '{model}' is not in JudgeArena's built-in thinking-model " - "defaults (Qwen3/SmolLM3/Olmo-3-7B-Think). Ignoring " - "thinking_token_budget unless reasoning_parser or " - "reasoning_config is provided explicitly.", - stacklevel=2, - ) + self._resolved = _resolved + self.model_path = _resolved.model_path + self.max_tokens = _resolved.max_tokens + self._chat_template_kwargs = _resolved.chat_template_kwargs + self._sampling_params_kwargs = dict(_resolved.sampling_params_kwargs) self.sampling_params = SamplingParams(**self._sampling_params_kwargs) + self.chat_template = _resolved.chat_template + self._use_generate = _resolved.use_generate + llm_init_kwargs = dict(_resolved.vllm_kwargs) + trust_remote_code = bool(llm_init_kwargs.pop("trust_remote_code", True)) self.llm = _init_llm_with_retry( - LLM, model=model, trust_remote_code=True, **vllm_kwargs + LLM, + model=_resolved.model_path, + trust_remote_code=trust_remote_code, + **llm_init_kwargs, ) self.tokenizer = self.llm.get_tokenizer() - # Resolve chat template: - # 1. Explicit override always wins → use chat() with that template - # 2. If tokenizer has one, use it → use chat() (pass None to vLLM) - # 3. No template found → fall back to generate() for base models - if chat_template: - self.chat_template = chat_template - self._use_generate = False - logger.info("ChatVLLM: using explicit chat template for '%s'", model) + if self.chat_template: + logger.info( + "ChatVLLM: using explicit chat template for '%s'", + _resolved.model_path, + ) + elif self._use_generate: + logger.info( + "ChatVLLM: falling back to llm.generate() for '%s'", + _resolved.model_path, + ) else: - if not getattr(self.tokenizer, "chat_template", None): - warnings.warn( - f"Model '{model}' tokenizer does not define a chat template. " - f"Falling back to llm.generate() (no chat formatting). " - f"Override with --chat_template if this model needs one.", - stacklevel=2, - ) - self.chat_template = None - self._use_generate = True - if disable_thinking: - warnings.warn( - f"Model '{model}' has no chat template, so disable_thinking " - "cannot be applied when falling back to llm.generate().", - stacklevel=2, - ) - else: - self.chat_template = None # let vLLM use the tokenizer's own - self._use_generate = False - logger.info("ChatVLLM: using tokenizer's chat template for '%s'", model) + logger.info( + "ChatVLLM: using tokenizer's chat template for '%s'", + _resolved.model_path, + ) def set_temperature(self, temperature: float) -> None: from vllm import SamplingParams self._sampling_params_kwargs["temperature"] = float(temperature) self.sampling_params = SamplingParams(**self._sampling_params_kwargs) + self._resolved.sampling_params_kwargs["temperature"] = float(temperature) def _to_messages(self, input_item) -> list[dict]: - """Convert LangChain prompt input to OpenAI-style messages.""" - # Map LangChain message types to OpenAI roles - role_map = {"human": "user", "ai": "assistant", "system": "system"} - - # Handle ChatPromptValue from LangChain - if hasattr(input_item, "to_messages"): - lc_messages = input_item.to_messages() - return [ - {"role": role_map.get(msg.type, msg.type), "content": msg.content} - for msg in lc_messages - ] - # Handle list of tuples like [("system", "..."), ("user", "...")] - elif ( - isinstance(input_item, list) - and input_item - and isinstance(input_item[0], tuple) - ): - return [ - {"role": role if role != "human" else "user", "content": content} - for role, content in input_item - ] - # Handle already formatted messages - elif ( - isinstance(input_item, list) - and input_item - and isinstance(input_item[0], dict) - ): - return input_item - # Handle plain string (wrap as user message) - elif isinstance(input_item, str): - return [{"role": "user", "content": input_item}] - else: - raise ValueError(f"Unsupported input type: {type(input_item)}") + return vllm_input_to_messages(input_item) def _to_raw_text(self, input_item) -> str: - """Extract raw text from an input item for use with llm.generate().""" - if isinstance(input_item, str): - return input_item - # ChatPromptValue from LangChain - if hasattr(input_item, "to_string"): - return input_item.to_string() - # List of dicts (messages) - concatenate contents - if ( - isinstance(input_item, list) - and input_item - and isinstance(input_item[0], dict) - ): - return "\n".join(msg["content"] for msg in input_item) - raise ValueError(f"Cannot extract raw text from: {type(input_item)}") + return vllm_input_to_raw_text(input_item) def _run_raw_batch(self, inputs: list): - """Process a batch of inputs using vllm.LLM.chat() or llm.generate(). - - Uses ``llm.chat()`` when a chat template is available (instruct models), - and ``llm.generate()`` when no template is found (base models). - """ if self._use_generate: prompts = [self._to_raw_text(inp) for inp in inputs] outputs = self.llm.generate(prompts, self.sampling_params) @@ -403,62 +455,60 @@ def _run_raw_batch(self, inputs: list): return outputs def batch(self, inputs: list, **invoke_kwargs) -> list[str]: - """Return the text completion for each input in *inputs*.""" outputs = self._run_raw_batch(inputs) return [out.outputs[0].text for out in outputs] def invoke(self, input_item, **invoke_kwargs) -> str: - """Process a single input.""" - results = self.batch([input_item], **invoke_kwargs) - return results[0] + return self.batch([input_item], **invoke_kwargs)[0] async def ainvoke(self, input_item, **invoke_kwargs): - """Async version - runs sync version in executor for compatibility.""" - import asyncio - loop = asyncio.get_event_loop() return await loop.run_in_executor( None, lambda: self.invoke(input_item, **invoke_kwargs) ) + def to_prepared_cache_adapter( + self, + *, + model_spec: str | None = None, + ) -> PreparedModel | None: + """Return a cache adapter for this constructed ChatVLLM backend.""" + return adapt_vllm_backend(self, model_spec=model_spec) + -def do_inference(chat_model, inputs, use_tqdm: bool = False): - """Run inference over *inputs*, returning a list of text completions. +def _normalize_inference_outputs(results: list[Any]) -> list[str]: + return [x.content if hasattr(x, "content") else x for x in results] - Retries on rate-limit/server errors with exponential backoff. The async - path (``use_tqdm=True``) retries individual calls; the batch path splits - into ``4**attempt`` chunks on failure. - """ - invoke_kwargs = { - # "stop": ["```"], - # "max_tokens": 100, - } + +def _do_inference_uncached( + chat_model: Any, + inputs: Sequence[Any], + *, + use_tqdm: bool = False, +) -> list[str]: + invoke_kwargs: dict[str, Any] = {} if use_tqdm: - # perform inference asynchronously to be able to update tqdm, chat_model.batch does not work as it blocks until - # all requests are received - # JUDGEARENA_JUDGE_MAX_CONCURRENCY caps simultaneous in-flight ainvokes - # (e.g. against OpenRouter). Unset = unbounded, preserving prior behaviour. cap = safe_parse_int("JUDGEARENA_JUDGE_MAX_CONCURRENCY") cap = cap if cap and cap > 0 else None - async def process_with_real_progress(chat_model, inputs, pbar): + async def process_with_real_progress(model, batch_inputs, pbar): sem = asyncio.Semaphore(cap) if cap else None async def process_single(input_item, max_retries=5, base_delay=1.0): for attempt in range(max_retries): try: - result = await chat_model.ainvoke(input_item, **invoke_kwargs) + result = await model.ainvoke(input_item, **invoke_kwargs) pbar.update(1) return result - except Exception as e: - if attempt == max_retries - 1 or not _is_retryable_error(e): + except Exception as exc: + if attempt == max_retries - 1 or not _is_retryable_error(exc): raise delay = base_delay * (2**attempt) logger.warning( "Retry because of a server error, %d/%d: %s. Waiting %ss...", attempt + 1, max_retries, - e, + exc, delay, ) await asyncio.sleep(delay) @@ -469,14 +519,12 @@ async def gated(inp): async with sem: return await process_single(inp) - # asyncio.gather preserves order (unlike as_completed) - results = await asyncio.gather(*[gated(inp) for inp in inputs]) - return results + return await asyncio.gather(*[gated(inp) for inp in batch_inputs]) with logging_redirect_tqdm(), tqdm(total=len(inputs)) as pbar: res = asyncio.run( process_with_real_progress( - chat_model=chat_model, inputs=inputs, pbar=pbar + chat_model=chat_model, batch_inputs=inputs, pbar=pbar ) ) else: @@ -494,28 +542,74 @@ def batch_with_retry(batch_inputs, max_retries=5, base_delay=1.0): for chunk in chunks: results.extend(chat_model.batch(inputs=chunk, **invoke_kwargs)) return results - except Exception as e: - if attempt == max_retries - 1 or not _is_retryable_error(e): + except Exception as exc: + if attempt == max_retries - 1 or not _is_retryable_error(exc): raise delay = base_delay * (2**attempt) - next_chunks = 4 ** (attempt + 1) logger.warning( - "Retry because of a server error, %d/%d: %s. Waiting %ss, then splitting into %d chunks...", + "Retry because of a server error, %d/%d: %s. Waiting %ss...", attempt + 1, max_retries, - e, + exc, delay, - next_chunks, ) time.sleep(delay) + raise RuntimeError("batch_with_retry exhausted retries without returning") res = batch_with_retry(inputs) - # Not sure why the API of Langchain returns sometime a string and sometimes an AIMessage object - # is it because of using Chat and barebones models? - # when using OpenAI, the output is AIMessage not a string... - res = [x.content if hasattr(x, "content") else x for x in res] - return res + return _normalize_inference_outputs(res) + + +def do_inference( + chat_model, + inputs, + use_tqdm: bool = False, + *, + cache: Any | None = None, + cache_meta: dict[str, Any] | None = None, +): + """Run inference over *inputs*, optionally via the unified inference cache.""" + if not inputs: + return [] + + if cache is not None and getattr(cache, "mode", None) == "off": + return _do_inference_uncached(chat_model, inputs, use_tqdm=use_tqdm) + + if cache is None: + return _do_inference_uncached(chat_model, inputs, use_tqdm=use_tqdm) + + cacheable = resolve_cacheable_model(chat_model) + if cacheable is None: + return _do_inference_uncached(chat_model, inputs, use_tqdm=use_tqdm) + + descriptor = cacheable.cache_descriptor() + if descriptor is None: + warnings.warn( + f"Inference cache skipped for {cacheable.model_spec}; descriptor is unsafe.", + stacklevel=2, + ) + return _do_inference_uncached(chat_model, inputs, use_tqdm=use_tqdm) + + cache_meta = cache_meta or {} + row_metadata = cache_meta.get("metadata") + if row_metadata is not None and len(row_metadata) != len(inputs): + raise ValueError("cache_meta['metadata'] length must match inputs") + + canonical_inputs = [cacheable.canonicalize_input(item) for item in inputs] + + def miss_runner(miss_inputs: list[Any]) -> list[str]: + return _do_inference_uncached(cacheable, miss_inputs, use_tqdm=use_tqdm) + + return cache.get_or_run( + model_spec=cacheable.model_spec, + descriptor=descriptor, + canonical_inputs=canonical_inputs, + original_inputs=list(inputs), + miss_runner=miss_runner, + row_metadata=row_metadata, + producer_metadata=cacheable.producer_metadata(), + ) def _route_sampling_params( @@ -529,15 +623,6 @@ def _route_sampling_params( top_k_via_model_kwargs: bool = False, provider: str = "", ) -> dict: - """Route the cross-backend sampling params onto a provider's constructor kwargs. - - Only params that are explicitly set (not ``None``) are applied. - ``top_k_via_model_kwargs`` tunnels ``top_k`` through ``model_kwargs`` for - OpenAI-compatible backends that do not expose it directly. When - ``supported_fields`` is provided, a param the target class cannot accept - (and cannot be tunneled) is dropped with a warning rather than being passed - through and raising at construction time. - """ for key, value in ( ("temperature", temperature), ("top_p", top_p), @@ -560,25 +645,38 @@ def _route_sampling_params( return engine_kwargs -def make_model(model: str, max_tokens: int | None = 8192, **engine_kwargs): - """Instantiate a model wrapper from a provider/model-name string. - - Args: - model: Format ``{Provider}/{model_path}``, e.g. - ``VLLM/meta-llama/Llama-3.3-70B-Instruct``. - max_tokens: Maximum tokens the model may generate. - **engine_kwargs: Engine-specific options forwarded to the model wrapper. - Common keys honoured across backends: ``temperature``, ``top_p``, - ``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() +def _provider_model_class(model_provider: str): + model_classes = [LlamaCpp, ChatOpenAI, Together, OpenAI] + model_cls_dict = {model_cls.__name__: model_cls for model_cls in model_classes} + assert model_provider in model_cls_dict, ( + f"{model_provider} not available, choose among {list(model_cls_dict.keys())}" + ) + return model_cls_dict[model_provider] + + +def _vllm_constructor_kwargs(raw_engine_kwargs: dict[str, Any]) -> dict[str, Any]: + return { + key: value + for key, value in raw_engine_kwargs.items() + if key + not in { + "max_tokens", + "temperature", + "top_p", + "top_k", + "seed", + "chat_template", + } + } - # Dedicated arguments like max_tokens always win over engine_kwargs. - engine_kwargs["max_tokens"] = max_tokens or 8192 + +def make_model( + model: str, max_tokens: int | None = 8192, **engine_kwargs +) -> PreparedModel: + """Prepare a lazy model wrapper from a provider/model-name string.""" + engine_kwargs = engine_kwargs.copy() + resolved_max_tokens = max_tokens or 8192 + engine_kwargs["max_tokens"] = resolved_max_tokens temperature = engine_kwargs.pop("temperature", None) top_p = engine_kwargs.pop("top_p", None) @@ -586,25 +684,7 @@ def make_model(model: str, max_tokens: int | None = 8192, **engine_kwargs): seed = engine_kwargs.pop("seed", 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) + raw_engine_kwargs = dict(engine_kwargs) if model_provider == "Dummy": dummy_kwargs = {k: v for k, v in engine_kwargs.items() if v is not None} @@ -615,88 +695,229 @@ def make_model(model: str, max_tokens: int | None = 8192, **engine_kwargs): top_k=top_k, seed=seed, ) - return DummyModel(model, **dummy_kwargs) + sampling = effective_sampling( + temperature=dummy_kwargs.get("temperature"), + top_p=dummy_kwargs.get("top_p"), + top_k=dummy_kwargs.get("top_k"), + seed=dummy_kwargs.get("seed"), + ) + descriptor = build_dummy_descriptor( + model_spec=model, + max_tokens=resolved_max_tokens, + input_mode="chat", + constructor_settings=dummy_kwargs, + resolved_sampling=sampling, + ) + return PreparedModel( + provider=model_provider, + model_spec=model, + model_name=model_name, + max_tokens=resolved_max_tokens, + engine_kwargs=dummy_kwargs, + sampling=sampling, + input_mode="chat", + descriptor=descriptor, + producer_metadata=build_producer_metadata(provider=model_provider), + materialize=lambda prepared: DummyModel( + prepared.model_spec, + **prepared.engine_kwargs, + ), + ) - logger.info("Loading %s(model=%s)", model_provider, model_name) + logger.info("Preparing %s(model=%s)", model_provider, model_name) - # Use our custom ChatVLLM wrapper which properly applies chat templates if model_provider == "VLLM": - engine_kwargs = {k: v for k, v in engine_kwargs.items() if v is not None} - engine_kwargs["chat_template"] = engine_kwargs.get("chat_template", None) + vllm_kwargs = {k: v for k, v in raw_engine_kwargs.items() if v is not None} + vllm_kwargs["chat_template"] = vllm_kwargs.get("chat_template", None) _route_sampling_params( - engine_kwargs, + vllm_kwargs, + temperature=temperature, + top_p=top_p, + top_k=top_k, + seed=seed, + ) + resolved = resolve_vllm_settings( + model_name, + max_tokens=resolved_max_tokens, + chat_template=vllm_kwargs.get("chat_template"), temperature=temperature, top_p=top_p, top_k=top_k, seed=seed, + **_vllm_constructor_kwargs(vllm_kwargs), ) + descriptor = build_vllm_descriptor(model, resolved) + sampling = dict(resolved.sampling_params_kwargs) + + def materialize_vllm(prepared: PreparedModel) -> ChatVLLM: + assert prepared._vllm_resolved is not None + kwargs = dict(prepared.engine_kwargs) + return ChatVLLM( + model=model_name, + max_tokens=prepared.max_tokens, + chat_template=kwargs.get("chat_template"), + _resolved=prepared._vllm_resolved, + **_vllm_constructor_kwargs(kwargs), + ) - return ChatVLLM( - model=model_name, - **engine_kwargs, + return PreparedModel( + provider=model_provider, + model_spec=model, + model_name=model_name, + max_tokens=resolved_max_tokens, + engine_kwargs=vllm_kwargs, + sampling=sampling, + input_mode=resolved.input_mode, + descriptor=descriptor, + producer_metadata=build_producer_metadata(provider=model_provider), + materialize=materialize_vllm, + vllm_resolved=resolved, ) + hosted_kwargs = dict(engine_kwargs) + 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", + "disable_thinking", + "thinking_token_budget", + "chat_template_kwargs", + "revision", + "tokenizer_revision", + ): + hosted_kwargs.pop(key, None) + if model_provider == "OpenRouter": - # Special case we need to override API url and key - openai_kwargs = dict(engine_kwargs) - # OpenAI-compatible chat backends expose temperature/top_p/seed directly - # but not top_k, which has to be tunneled through model_kwargs. _route_sampling_params( - openai_kwargs, + hosted_kwargs, temperature=temperature, top_p=top_p, top_k=top_k, seed=seed, top_k_via_model_kwargs=True, ) - return ChatOpenAI( - api_key=os.getenv("OPENROUTER_API_KEY"), - base_url="https://openrouter.ai/api/v1", - model=model_name, - **openai_kwargs, + hosted_kwargs.setdefault("base_url", OPENROUTER_BASE_URL) + hosted_kwargs.setdefault("api_key", os.getenv("OPENROUTER_API_KEY")) + hosted_kwargs.setdefault("model", model_name) + base_url = str(hosted_kwargs["base_url"]) + sampling = effective_sampling( + temperature=hosted_kwargs.get("temperature"), + top_p=hosted_kwargs.get("top_p"), + top_k=top_k_from_settings(hosted_kwargs), + seed=hosted_kwargs.get("seed"), ) - else: - model_classes = [ - LlamaCpp, - ChatOpenAI, - ] - try: - from langchain_together.llms import Together - - model_classes.append(Together) - except ImportError as e: - logger.debug("Optional provider not available: %s", e) - try: - from langchain_openai.llms import OpenAI - - model_classes.append(OpenAI) - except ImportError as e: - logger.debug("Optional provider not available: %s", e) - model_cls_dict = {model_cls.__name__: model_cls for model_cls in model_classes} - assert model_provider in model_cls_dict, ( - f"{model_provider} not available, choose among {list(model_cls_dict.keys())}" + descriptor = build_hosted_descriptor( + provider=model_provider, + model_spec=model, + model_name=model_name, + max_tokens=resolved_max_tokens, + input_mode="chat", + base_url=base_url, + constructor_settings=hosted_kwargs, + resolved_sampling=sampling, + ) + return PreparedModel( + provider=model_provider, + model_spec=model, + model_name=model_name, + max_tokens=resolved_max_tokens, + engine_kwargs=hosted_kwargs, + sampling=sampling, + input_mode="chat", + descriptor=descriptor, + base_url=base_url, + producer_metadata=build_producer_metadata(provider=model_provider), + materialize=lambda prepared: ChatOpenAI(**prepared.engine_kwargs), ) - model_cls = model_cls_dict[model_provider] - if model_provider == "LlamaCpp": - engine_kwargs["model_path"] = model_name - else: - engine_kwargs["model"] = model_name - # Route sampling params against the target class's accepted fields: - # tunnel top_k via model_kwargs when the class lacks a top_k field, - # and drop any param the class cannot accept (e.g. Together has no - # ``seed``) instead of raising at construction time. - supported_fields = set(getattr(model_cls, "model_fields", {})) - _route_sampling_params( - engine_kwargs, - temperature=temperature, - top_p=top_p, - top_k=top_k, - seed=seed, - supported_fields=supported_fields, - top_k_via_model_kwargs=( - "top_k" not in supported_fields and "model_kwargs" in supported_fields - ), + model_cls = _provider_model_class(model_provider) + provider_kwargs = dict(hosted_kwargs) + if model_provider == "LlamaCpp": + provider_kwargs["model_path"] = model_name + input_mode = "raw" + else: + provider_kwargs["model"] = model_name + input_mode = "raw" if model_provider in {"OpenAI", "Together"} else "chat" + + supported_fields = set(getattr(model_cls, "model_fields", {})) + _route_sampling_params( + provider_kwargs, + temperature=temperature, + top_p=top_p, + top_k=top_k, + seed=seed, + supported_fields=supported_fields, + top_k_via_model_kwargs=( + "top_k" not in supported_fields and "model_kwargs" in supported_fields + ), + provider=model_provider, + ) + sampling = effective_sampling( + temperature=provider_kwargs.get("temperature"), + top_p=provider_kwargs.get("top_p"), + top_k=top_k_from_settings(provider_kwargs), + seed=provider_kwargs.get("seed"), + ) + base_url = resolve_hosted_base_url(provider_kwargs) + if model_cls is ChatOpenAI: + model_provider = hosted_provider_for_endpoint(base_url, model_provider) + if model_provider == "OpenRouter": + model = f"OpenRouter/{model_name}" + + if model_provider == "LlamaCpp": + descriptor = build_llamacpp_descriptor( + model_spec=model, + model_name=model_name, + max_tokens=resolved_max_tokens, + constructor_settings=provider_kwargs, + resolved_sampling=sampling, + ) + else: + descriptor = build_hosted_descriptor( provider=model_provider, + model_spec=model, + model_name=model_name, + max_tokens=resolved_max_tokens, + input_mode=input_mode, + base_url=base_url, + constructor_settings=provider_kwargs, + resolved_sampling=sampling, ) - return model_cls(**engine_kwargs) + + return PreparedModel( + provider=model_provider, + model_spec=model, + model_name=model_name, + max_tokens=resolved_max_tokens, + engine_kwargs=provider_kwargs, + sampling=sampling, + input_mode=input_mode, + descriptor=descriptor, + base_url=base_url, + producer_metadata=build_producer_metadata(provider=model_provider), + materialize=lambda prepared: model_cls(**prepared.engine_kwargs), + ) + + +__all__ = [ + "ChatVLLM", + "DESCRIPTOR_SCHEMA_VERSION", + "DummyModel", + "HOSTED_ADAPTER_VERSION", + "PreparedModel", + "VLLMResolvedSettings", + "build_default_judge_model_kwargs", + "do_inference", + "is_thinking_model", + "make_model", + "resolve_vllm_settings", +] diff --git a/tests/test_chat_vllm.py b/tests/test_chat_vllm.py index 0c1bac8..36b95b2 100644 --- a/tests/test_chat_vllm.py +++ b/tests/test_chat_vllm.py @@ -45,6 +45,25 @@ def chat(self, messages, sampling_params, **kwargs): "vllm.config.reasoning", SimpleNamespace(ReasoningConfig=FakeReasoningConfig), ) + + class FakeAutoTokenizer: + @staticmethod + def from_pretrained(model, trust_remote_code=True): + return SimpleNamespace(chat_template="{{ messages }}") + + class FakeAutoConfig: + @staticmethod + def from_pretrained(model, trust_remote_code=True): + return SimpleNamespace(max_position_embeddings=8192) + + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace( + AutoTokenizer=FakeAutoTokenizer, + AutoConfig=FakeAutoConfig, + ), + ) return captured, FakeReasoningConfig diff --git a/tests/test_inference_cache.py b/tests/test_inference_cache.py new file mode 100644 index 0000000..4a559fe --- /dev/null +++ b/tests/test_inference_cache.py @@ -0,0 +1,465 @@ +import json +import sys +from types import SimpleNamespace + +import pytest +from langchain_core.messages import HumanMessage, SystemMessage +from langchain_core.prompt_values import ChatPromptValue + +import judgearena.model_adapters as model_adapters +import judgearena.models as models +from judgearena.inference_cache import InferenceCache +from judgearena.models import ( + HOSTED_ADAPTER_VERSION, + PreparedModel, + do_inference, + make_model, +) +from judgearena.store_sqlite import ( + SQLiteInferenceStore, + descriptor_hash, + stable_json_dumps, + store_folder, +) + + +def _install_fake_vllm(monkeypatch): + captured = {"llm_init": False} + + class FakeSamplingParams: + def __init__(self, **kwargs): + captured["sampling_kwargs"] = kwargs + + class FakeReasoningConfig: + def __init__(self, **kwargs): + captured["reasoning_config_kwargs"] = kwargs + + class FakeLLM: + def __init__(self, *, model, trust_remote_code, **kwargs): + captured["llm_init"] = True + captured["llm_init_args"] = { + "model": model, + "trust_remote_code": trust_remote_code, + "kwargs": kwargs, + } + + def get_tokenizer(self): + return SimpleNamespace(chat_template="{{ messages }}") + + def chat(self, messages, sampling_params, **kwargs): + return [SimpleNamespace(outputs=[SimpleNamespace(text="generated")])] + + monkeypatch.setitem( + sys.modules, + "vllm", + SimpleNamespace(LLM=FakeLLM, SamplingParams=FakeSamplingParams), + ) + monkeypatch.setitem( + sys.modules, + "vllm.config.reasoning", + SimpleNamespace(ReasoningConfig=FakeReasoningConfig), + ) + + class FakeAutoTokenizer: + @staticmethod + def from_pretrained(model, trust_remote_code=True): + return SimpleNamespace(chat_template="{{ messages }}") + + class FakeAutoConfig: + @staticmethod + def from_pretrained(model, trust_remote_code=True): + return SimpleNamespace(max_position_embeddings=8192) + + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace( + AutoTokenizer=FakeAutoTokenizer, + AutoConfig=FakeAutoConfig, + ), + ) + return captured + + +def _seed_cache(tmp_path, task, model, inputs, outputs): + descriptor = model.cache_descriptor() + assert descriptor is not None + canonical = [model.canonicalize_input(item) for item in inputs] + with InferenceCache(tmp_path, task, mode="refresh") as cache: + cache.get_or_run( + model_spec=model.model_spec, + descriptor=descriptor, + canonical_inputs=canonical, + original_inputs=inputs, + miss_runner=lambda miss_inputs: outputs[: len(miss_inputs)], + producer_metadata=model.producer_metadata(), + ) + + +def test_vllm_full_cache_hit_skips_llm_init(monkeypatch, tmp_path): + captured = _install_fake_vllm(monkeypatch) + model = make_model("VLLM/Qwen/Qwen3.5-9B", max_tokens=16, temperature=0.0) + inputs = ["hello", "world"] + _seed_cache(tmp_path, "arena", model, inputs, ["cached-a", "cached-b"]) + + hit_model = make_model("VLLM/Qwen/Qwen3.5-9B", max_tokens=16, temperature=0.0) + with InferenceCache(tmp_path, "arena", mode="use") as cache: + results = do_inference(hit_model, inputs, cache=cache) + + assert results == ["cached-a", "cached-b"] + assert captured["llm_init"] is False + + +def test_within_batch_dedupe_runs_inference_once(tmp_path): + model = make_model("Dummy/cache-dedupe", max_tokens=8) + inputs = ["same", "same", "other"] + seen: list[str] = [] + + def miss_runner(miss_inputs): + seen.extend(miss_inputs) + return [f"out-{item}" for item in miss_inputs] + + descriptor = model.cache_descriptor() + canonical = [model.canonicalize_input(item) for item in inputs] + with InferenceCache(tmp_path, "arena", mode="refresh") as cache: + first = cache.get_or_run( + model_spec=model.model_spec, + descriptor=descriptor, + canonical_inputs=canonical, + original_inputs=inputs, + miss_runner=miss_runner, + producer_metadata=model.producer_metadata(), + ) + + assert first == ["out-same", "out-same", "out-other"] + assert seen == ["same", "other"] + + seen.clear() + with InferenceCache(tmp_path, "arena", mode="use") as cache: + second = do_inference(model, inputs, cache=cache) + + assert second == first + assert seen == [] + + +def test_use_mode_returns_row_won_by_concurrent_insert(monkeypatch, tmp_path): + model = make_model("Dummy/concurrent", max_tokens=8) + descriptor = model.cache_descriptor() + assert descriptor is not None + canonical = [model.canonicalize_input("input")] + original_save = SQLiteInferenceStore.save_outputs_and_metadata + + def racing_save(self, outputs, metadata, **kwargs): + competing = outputs.copy() + competing["output_text"] = "concurrent-winner" + self.save_outputs(competing, pushed_by="other", run_id="other") + return original_save(self, outputs, metadata, **kwargs) + + monkeypatch.setattr( + SQLiteInferenceStore, + "save_outputs_and_metadata", + racing_save, + ) + with InferenceCache(tmp_path, "arena", mode="use") as cache: + result = cache.get_or_run( + model_spec=model.model_spec, + descriptor=descriptor, + canonical_inputs=canonical, + original_inputs=["input"], + miss_runner=lambda _: ["fresh-loser"], + row_metadata=[{"question_id": "q1"}], + ) + + assert result == ["concurrent-winner"] + + +def test_changed_input_is_cache_miss(tmp_path): + model = make_model("Dummy/cache-change", max_tokens=8) + _seed_cache(tmp_path, "arena", model, ["alpha"], ["out-alpha"]) + + with InferenceCache(tmp_path, "arena", mode="use") as cache: + results = do_inference(model, ["beta"], cache=cache) + + assert results == ["cache-change"] + + +def test_seed_separates_descriptor_cells(tmp_path): + model_a = make_model("Dummy/cache-seed", max_tokens=8, seed=1) + model_b = make_model("Dummy/cache-seed", max_tokens=8, seed=2) + assert model_a.cache_descriptor() != model_b.cache_descriptor() + + _seed_cache(tmp_path, "arena", model_a, ["prompt"], ["seed-1"]) + with InferenceCache(tmp_path, "arena", mode="use") as cache: + hit_a = do_inference(model_a, ["prompt"], cache=cache) + miss_b = do_inference(model_b, ["prompt"], cache=cache) + + assert hit_a == ["seed-1"] + assert miss_b == ["cache-seed"] + + +def test_refresh_replaces_cached_output(tmp_path): + model = make_model("Dummy/cache-refresh", max_tokens=8) + _seed_cache(tmp_path, "arena", model, ["prompt"], ["old"]) + + with InferenceCache(tmp_path, "arena", mode="refresh") as cache: + results = do_inference(model, ["prompt"], cache=cache) + + assert results == ["cache-refresh"] + + +def test_metadata_associations_are_saved(tmp_path): + model = make_model("Dummy/cache-meta", max_tokens=8) + metadata = [{"question_id": "q-1"}, {"question_id": "q-2"}] + with InferenceCache(tmp_path, "arena", mode="refresh") as cache: + do_inference( + model, + ["one", "two"], + cache=cache, + cache_meta={"metadata": metadata}, + ) + + descriptor = model.cache_descriptor() + folder = store_folder( + tmp_path, + "arena", + model.model_spec, + descriptor_hash(descriptor), + ) + + with SQLiteInferenceStore(folder / "inference.db") as store: + rows = store.query_metadata() + + assert len(rows) == 2 + saved = {json.loads(value)["question_id"] for value in rows["metadata_json"]} + assert saved == {"q-1", "q-2"} + + +def test_secret_values_are_not_descriptorized(monkeypatch): + monkeypatch.setenv("OPENROUTER_API_KEY", "dummy") + model = make_model( + "OpenRouter/google/gemma-3-4b-it", + max_tokens=16, + api_key="super-secret", + default_headers={"Authorization": "Bearer secret"}, + ) + assert model.engine_kwargs["api_key"] == "super-secret" + assert model.engine_kwargs["default_headers"]["Authorization"] == "Bearer secret" + descriptor = model.cache_descriptor() + assert descriptor is not None + serialized = stable_json_dumps(descriptor) + assert "super-secret" not in serialized + assert "Authorization" not in serialized + assert "api_key" not in serialized + + +def test_provider_canonicalization_distinguishes_chat_and_raw(): + chat_model = make_model("Dummy/chat", max_tokens=8) + raw_model = make_model("Dummy/raw", max_tokens=8) + raw_model.input_mode = "raw" + + chat_payload = chat_model.canonicalize_input("hello") + raw_payload = raw_model.canonicalize_input("hello") + assert json.loads(chat_payload)["kind"] == "chat" + assert json.loads(raw_payload)["kind"] == "raw" + assert chat_payload != raw_payload + + prompt = ChatPromptValue( + messages=[SystemMessage(content="sys"), HumanMessage(content="hi", id="tmp")] + ) + chat_canonical = json.loads(chat_model.canonicalize_input(prompt)) + assert chat_canonical["messages"] == [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + ] + assert "id" not in stable_json_dumps(chat_canonical) + + +def test_set_temperature_direct_and_mutated_descriptors_match(tmp_path): + direct = make_model("Dummy/temp-hash", max_tokens=8, temperature=0.9) + mutated = make_model("Dummy/temp-hash", max_tokens=8) + mutated.set_temperature(0.9) + assert direct.cache_descriptor() == mutated.cache_descriptor() + assert descriptor_hash(direct.cache_descriptor()) == descriptor_hash( + mutated.cache_descriptor() + ) + + +def test_set_temperature_changes_cache_cell(tmp_path): + cold = make_model("Dummy/temp", max_tokens=8, temperature=0.2) + _seed_cache(tmp_path, "arena", cold, ["x"], ["cold-hit"]) + + model = make_model("Dummy/temp", max_tokens=8, temperature=0.2) + model.set_temperature(0.9) + assert model.cache_descriptor()["sampling"]["temperature"] == 0.9 + backend = model.materialize() + assert backend.init_kwargs["temperature"] == 0.9 + + with InferenceCache(tmp_path, "arena", mode="use") as cache: + results = do_inference(model, ["x"], cache=cache) + + assert results == ["temp"] + + +def test_hosted_adapter_version_is_part_of_descriptor(monkeypatch): + monkeypatch.setenv("OPENROUTER_API_KEY", "dummy") + model = make_model("OpenRouter/openai/gpt-4o-mini", max_tokens=8) + descriptor = model.cache_descriptor() + assert descriptor["hosted_adapter_version"] == HOSTED_ADAPTER_VERSION + assert descriptor["server_defaults"] == "unobserved" + metadata = model.producer_metadata() + assert metadata["hosted_adapter_version"] == HOSTED_ADAPTER_VERSION + assert "langchain_openai_version" in metadata + + +def test_direct_chat_vllm_can_be_cached_when_constructed(monkeypatch, tmp_path): + captured = _install_fake_vllm(monkeypatch) + chat_model = models.ChatVLLM( + model="Qwen/Qwen3.5-9B", + max_tokens=16, + gpu_memory_utilization=0.7, + ) + assert captured["llm_init"] is True + + wrapped = model_adapters.wrap_known_model(chat_model) + assert wrapped is not None + _seed_cache(tmp_path, "arena", wrapped, ["hello"], ["cached-vllm"]) + + with InferenceCache(tmp_path, "arena", mode="use") as cache: + results = do_inference(chat_model, ["hello"], cache=cache) + + assert results == ["cached-vllm"] + + +def test_off_mode_never_reads_cache(tmp_path): + model = make_model("Dummy/off", max_tokens=8) + _seed_cache(tmp_path, "arena", model, ["prompt"], ["cached"]) + + with InferenceCache(tmp_path, "arena", mode="off") as cache: + results = do_inference(model, ["prompt"], cache=cache) + + assert results == ["off"] + + +def test_do_inference_off_bypasses_cache_resolution(monkeypatch, tmp_path): + model = make_model("Dummy/off-bypass", max_tokens=8) + + def fail_resolution(*args, **kwargs): + raise AssertionError("resolve_cacheable_model must not run in off mode") + + monkeypatch.setattr(models, "resolve_cacheable_model", fail_resolution) + + with InferenceCache(tmp_path, "arena", mode="off") as cache: + results = do_inference(model, ["prompt"], cache=cache) + + assert results == ["off-bypass"] + + +def test_prepared_model_proxy_materializes_backend(): + model = make_model("Dummy/proxy", max_tokens=8) + assert isinstance(model, PreparedModel) + assert model.init_kwargs["max_tokens"] == 8 + + +def test_descriptor_schema_version_is_present(): + model = make_model("Dummy/schema", max_tokens=8) + descriptor = model.cache_descriptor() + assert descriptor["descriptor_schema_version"] == models.DESCRIPTOR_SCHEMA_VERSION + + +def test_hosted_adapter_version_must_be_bumped_when_request_shaping_changes(): + """Guardrail: request-shaping edits must bump HOSTED_ADAPTER_VERSION.""" + assert HOSTED_ADAPTER_VERSION == "judgearena-hosted-adapter/v1" + + +def test_off_mode_runs_every_input_without_dedupe(tmp_path): + model = make_model("Dummy/off-dedupe", max_tokens=8) + calls: list[str] = [] + + def counting_runner(inputs): + calls.extend(inputs) + return [f"out-{index}" for index, _ in enumerate(inputs)] + + with InferenceCache(tmp_path, "arena", mode="off") as cache: + results = cache.get_or_run( + model_spec=model.model_spec, + descriptor=model.cache_descriptor(), + canonical_inputs=[model.canonicalize_input("same")] * 2, + original_inputs=["same", "same"], + miss_runner=counting_runner, + producer_metadata=model.producer_metadata(), + ) + + assert results == ["out-0", "out-1"] + assert calls == ["same", "same"] + + +def test_metadata_only_association_marks_cell_dirty(tmp_path): + model = make_model("Dummy/meta-dirty", max_tokens=8) + _seed_cache(tmp_path, "arena", model, ["prompt"], ["cached-out"]) + + with InferenceCache(tmp_path, "arena", mode="use", push=False) as cache: + do_inference( + model, + ["prompt"], + cache=cache, + cache_meta={"metadata": [{"question_id": "q-hit"}]}, + ) + assert cache._dirty_cells + + descriptor = model.cache_descriptor() + folder = store_folder( + tmp_path, + "arena", + model.model_spec, + descriptor_hash(descriptor), + ) + with SQLiteInferenceStore(folder / "inference.db") as store: + rows = store.query_metadata() + assert len(rows) == 1 + + +def test_close_before_push_closes_sqlite_before_upload(monkeypatch, tmp_path): + model = make_model("Dummy/push-close", max_tokens=8) + states: list[bool] = [] + active_caches: list[InferenceCache] = [] + + import judgearena.inference_cache as inference_cache_mod + + original_enter = InferenceCache.__enter__ + + def track_enter(self): + active_caches.append(self) + return original_enter(self) + + monkeypatch.setattr(InferenceCache, "__enter__", track_enter) + + def spy_push(*args, **kwargs): + cache = active_caches[-1] + states.append(bool(cache._stores)) + return None + + monkeypatch.setattr(inference_cache_mod, "push_cells", spy_push) + + with InferenceCache(tmp_path, "arena", mode="refresh", push=True) as cache: + do_inference(model, ["x"], cache=cache) + + assert states == [False] + + +def test_invalid_cache_mode_rejected(tmp_path): + with pytest.raises(ValueError, match="Invalid cache mode"): + InferenceCache(tmp_path, "arena", mode="bogus") # type: ignore[arg-type] + + +def test_vllm_make_model_is_lazy_until_materialized(monkeypatch): + captured = _install_fake_vllm(monkeypatch) + model = make_model( + "VLLM/Qwen/Qwen3.5-9B", + max_tokens=16, + thinking_token_budget=64, + gpu_memory_utilization=0.7, + ) + assert isinstance(model, PreparedModel) + assert captured["llm_init"] is False + model.materialize() + assert captured["llm_init"] is True diff --git a/tests/test_model_adapters.py b/tests/test_model_adapters.py new file mode 100644 index 0000000..f97b7e0 --- /dev/null +++ b/tests/test_model_adapters.py @@ -0,0 +1,355 @@ +import json +import sys +from types import SimpleNamespace + +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage +from langchain_core.prompt_values import ChatPromptValue + +import judgearena.model_adapters as adapters +import judgearena.models as models +from judgearena.model_adapters import ( + HOSTED_ADAPTER_VERSION, + LOCAL_LLAMACPP_ADAPTER_VERSION, + build_producer_metadata, + build_vllm_descriptor, + effective_sampling, + normalize_constructor_settings, + top_k_from_settings, + wrap_known_model, +) +from judgearena.models import DummyModel, make_model, resolve_vllm_settings + + +def test_normalize_constructor_settings_redacts_secret_keys(): + settings = { + "temperature": 0.5, + "api_key": "secret-value", + "default_headers": {"Authorization": "Bearer x"}, + "model_kwargs": {"top_k": 40}, + } + normalized = normalize_constructor_settings(settings) + assert normalized == {"temperature": 0.5, "model_kwargs": {"top_k": 40}} + serialized = json.dumps(normalized) + assert "secret-value" not in serialized + assert "Authorization" not in serialized + + +def test_make_model_openrouter_endpoint_unifies_provider(monkeypatch): + monkeypatch.setenv("OPENROUTER_API_KEY", "dummy") + model = make_model( + "ChatOpenAI/google/gemma-3-4b-it", + max_tokens=16, + base_url="https://openrouter.ai/api/v1", + ) + descriptor = model.cache_descriptor() + assert descriptor is not None + assert model.provider == "OpenRouter" + assert model.model_spec == "OpenRouter/google/gemma-3-4b-it" + assert descriptor["provider"] == "OpenRouter" + assert descriptor["base_url"] == "https://openrouter.ai/api/v1" + + +def test_lazy_and_materialized_hosted_descriptors_match(monkeypatch): + monkeypatch.setenv("OPENROUTER_API_KEY", "dummy") + monkeypatch.setenv("OPENAI_API_KEY", "dummy") + monkeypatch.setenv("TOGETHER_API_KEY", "dummy") + specs = [ + "OpenRouter/openai/gpt-4o-mini", + "OpenAI/gpt-3.5-turbo-instruct", + "Together/meta-llama/Llama-3.3-70B-Instruct-Turbo", + ] + + for spec in specs: + lazy = make_model(spec, max_tokens=8, temperature=0.2) + wrapped = wrap_known_model(lazy.materialize(), model_spec=lazy.model_spec) + assert wrapped is not None + assert wrapped.cache_descriptor() == lazy.cache_descriptor() + + +def test_wrapped_hosted_model_preserves_unset_max_tokens(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "dummy") + backend = models.ChatOpenAI(model="gpt-4o-mini") + + wrapped = wrap_known_model(backend) + + assert wrapped is not None + assert wrapped.max_tokens is None + assert "max_tokens" not in wrapped.cache_descriptor()["sampling"] + + +def test_make_model_captures_base_url_for_generic_hosted(monkeypatch): + monkeypatch.setenv("OPENAI_API_KEY", "dummy") + model = make_model( + "OpenAI/text-davinci-003", + max_tokens=8, + openai_api_base="https://example.test/v1", + ) + descriptor = model.cache_descriptor() + assert descriptor is not None + assert descriptor["base_url"] == "https://example.test/v1" + + +def test_make_model_preserves_constructor_secrets(monkeypatch): + monkeypatch.setenv("OPENROUTER_API_KEY", "env-key") + model = make_model( + "OpenRouter/google/gemma-3-4b-it", + max_tokens=16, + api_key="runtime-secret", + default_headers={"Authorization": "Bearer secret"}, + ) + assert model.engine_kwargs["api_key"] == "runtime-secret" + assert model.engine_kwargs["default_headers"]["Authorization"] == "Bearer secret" + descriptor = model.cache_descriptor() + assert descriptor is not None + serialized = json.dumps(descriptor) + assert "runtime-secret" not in serialized + assert "Authorization" not in serialized + + +def test_top_k_from_settings_preserves_zero(): + assert top_k_from_settings({"top_k": 0}) == 0 + assert top_k_from_settings({"model_kwargs": {"top_k": 0}}) == 0 + assert ( + effective_sampling( + temperature=0.5, + top_p=0.9, + top_k=top_k_from_settings({"top_k": 0}), + seed=None, + )["top_k"] + == 0 + ) + + +def test_resolve_vllm_settings_uses_explicit_tokenizer(monkeypatch): + monkeypatch.setitem( + sys.modules, + "vllm.config.reasoning", + SimpleNamespace( + ReasoningConfig=lambda **kwargs: SimpleNamespace(**kwargs), + ), + ) + + seen: dict[str, object] = {} + + class FakeAutoTokenizer: + @staticmethod + def from_pretrained(model, **kwargs): + seen["tokenizer_id"] = model + seen["tokenizer_kwargs"] = kwargs + return SimpleNamespace(chat_template="{{ messages }}") + + class FakeAutoConfig: + @staticmethod + def from_pretrained(model, **kwargs): + return SimpleNamespace(max_position_embeddings=4096) + + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace( + AutoTokenizer=FakeAutoTokenizer, + AutoConfig=FakeAutoConfig, + ), + ) + + resolve_vllm_settings( + "org/base-model", + max_tokens=16, + tokenizer="org/custom-tokenizer", + revision="model-rev", + tokenizer_revision="tok-rev", + trust_remote_code=False, + ) + assert seen["tokenizer_id"] == "org/custom-tokenizer" + assert seen["tokenizer_kwargs"] == { + "trust_remote_code": False, + "revision": "tok-rev", + } + + +def test_vllm_descriptor_uses_fully_resolved_sampling(monkeypatch): + monkeypatch.setitem( + sys.modules, + "vllm.config.reasoning", + SimpleNamespace( + ReasoningConfig=lambda **kwargs: SimpleNamespace(**kwargs), + ), + ) + + class FakeAutoTokenizer: + @staticmethod + def from_pretrained(model, **kwargs): + return SimpleNamespace(chat_template="{{ messages }}") + + class FakeAutoConfig: + @staticmethod + def from_pretrained(model, **kwargs): + return SimpleNamespace(max_position_embeddings=4096) + + monkeypatch.setitem( + sys.modules, + "transformers", + SimpleNamespace( + AutoTokenizer=FakeAutoTokenizer, + AutoConfig=FakeAutoConfig, + ), + ) + + resolved = resolve_vllm_settings( + "Qwen/Qwen3.5-9B", + max_tokens=32, + max_model_len=8192, + thinking_token_budget=16, + gpu_memory_utilization=0.7, + ) + descriptor = build_vllm_descriptor("VLLM/Qwen/Qwen3.5-9B", resolved) + assert descriptor is not None + assert descriptor["sampling"]["temperature"] == 0.6 + assert descriptor["sampling"]["top_p"] == 0.95 + assert descriptor["sampling"]["max_tokens"] == 32 + assert descriptor["sampling"]["thinking_token_budget"] == 16 + assert descriptor["engine_settings"]["max_model_len"] == 4096 + assert descriptor["engine_settings"]["gpu_memory_utilization"] == 0.7 + assert descriptor["engine_settings"]["reasoning_parser"] == "qwen3" + assert "reasoning_config" in descriptor["engine_settings"] + assert "vllm_version" in descriptor + assert "langchain_openai_version" not in descriptor + + +def test_vllm_set_temperature_updates_sampling_only(): + descriptor = { + "sampling": {"temperature": 0.6, "max_tokens": 8}, + "engine_settings": {"max_model_len": 1024}, + } + resolved = SimpleNamespace(sampling_params_kwargs={"temperature": 0.6}) + model = adapters.PreparedModel( + provider="VLLM", + model_spec="VLLM/test/model", + model_name="test/model", + max_tokens=8, + engine_kwargs={"max_model_len": 1024}, + sampling={"temperature": 0.6, "max_tokens": 8}, + input_mode="chat", + descriptor=descriptor, + materialize=lambda _prepared: (_ for _ in ()).throw( + AssertionError("must stay lazy") + ), + producer_metadata={}, + vllm_resolved=resolved, + ) + + model.set_temperature(0.9) + + updated = model.cache_descriptor() + assert updated["sampling"]["temperature"] == 0.9 + assert "temperature" not in updated["engine_settings"] + assert resolved.sampling_params_kwargs["temperature"] == 0.9 + + +def test_llamacpp_descriptor_includes_local_engine_version(monkeypatch): + monkeypatch.setattr( + adapters, + "_provider_package_version", + lambda name: "0.42.0" if name == "llama-cpp-python" else "1.0.0", + ) + model = make_model("LlamaCpp/models/test.gguf", max_tokens=64) + descriptor = model.cache_descriptor() + assert descriptor is not None + assert descriptor["llama_cpp_python_version"] == "0.42.0" + assert descriptor["local_adapter_version"] == LOCAL_LLAMACPP_ADAPTER_VERSION + metadata = model.producer_metadata() + assert metadata["llama_cpp_python_version"] == "0.42.0" + assert metadata["descriptor_schema_version"] == adapters.DESCRIPTOR_SCHEMA_VERSION + + +def test_hosted_canonicalization_preserves_tool_fields(): + prompt = ChatPromptValue( + messages=[ + SystemMessage(content="sys"), + HumanMessage(content="call tool"), + AIMessage( + content="", + tool_calls=[ + { + "id": "call-1", + "name": "lookup", + "args": {"q": "paris"}, + } + ], + ), + ToolMessage(content='{"city":"Paris"}', tool_call_id="call-1"), + ] + ) + canonical = json.loads(adapters.canonicalize_hosted_chat_input(prompt)) + assert canonical["messages"][0] == {"role": "system", "content": "sys"} + tool_message = canonical["messages"][-1] + assert tool_message["role"] == "tool" + assert tool_message["tool_call_id"] == "call-1" + assert tool_message["content"] == '{"city":"Paris"}' + ai_message = canonical["messages"][2] + assert ai_message["tool_calls"][0]["name"] == "lookup" + assert ai_message["tool_calls"][0]["id"] == "call-1" + assert "id" not in ai_message + + +def test_vllm_dict_message_roles_match_runtime_normalization(): + messages = [ + {"role": "human", "content": "hello"}, + {"content": "missing role"}, + ] + + normalized = adapters.vllm_input_to_messages(messages) + canonical = json.loads( + adapters.canonicalize_vllm_input(messages, input_mode="chat") + ) + + assert normalized == [ + {"role": "user", "content": "hello"}, + {"role": "user", "content": "missing role"}, + ] + assert canonical["messages"] == normalized + + +def test_raw_canonicalization_accepts_dict_messages(): + canonical = json.loads( + adapters.canonicalize_raw_input( + [{"role": "user", "content": "hello"}, {"content": "world"}] + ) + ) + assert canonical == {"kind": "raw", "text": "hello\nworld"} + + +def test_set_temperature_updates_materialized_chatopenai_temperature(monkeypatch): + monkeypatch.setenv("OPENROUTER_API_KEY", "dummy") + + class FakeChat: + model_fields = { + "temperature": object(), + "max_tokens": object(), + "model": object(), + } + + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + monkeypatch.setattr(models, "ChatOpenAI", FakeChat) + model = make_model("OpenRouter/openai/gpt-4o-mini", max_tokens=8, temperature=0.2) + backend = model.materialize() + model.set_temperature(0.8) + assert backend.temperature == 0.8 + assert model.cache_descriptor()["sampling"]["temperature"] == 0.8 + + +def test_wrap_known_dummy_model(): + backend = DummyModel("Dummy/wrapped", max_tokens=8, temperature=0.1) + wrapped = wrap_known_model(backend) + assert wrapped is not None + assert wrapped.cache_descriptor() is not None + assert wrapped.canonicalize_input("hello").startswith('{"kind"') + + +def test_producer_metadata_includes_adapter_schema(): + metadata = build_producer_metadata(provider="OpenRouter") + assert metadata["hosted_adapter_version"] == HOSTED_ADAPTER_VERSION + assert metadata["descriptor_schema_version"] == adapters.DESCRIPTOR_SCHEMA_VERSION + assert "langchain_openai_version" in metadata From 42a50612e82fa61d9320500c0ecc3fb6abb2895b Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:53:10 +0200 Subject: [PATCH 10/13] feat(cache): unify cache lifecycle across pipelines Thread one inference cache through generation, judging, ELO, MT-Bench, and meta-eval while removing the legacy pass-level cache paths. Includes-AI-Code: true --- judgearena/config.py | 135 +++- judgearena/estimate_elo_ratings.py | 421 +++++++------ judgearena/evaluate.py | 135 +++- judgearena/generate.py | 86 ++- judgearena/generate_and_evaluate.py | 473 ++++++++------ judgearena/log.py | 21 +- judgearena/meta_eval/annotate.py | 185 ++---- judgearena/meta_eval/cache.py | 165 ----- judgearena/meta_eval/cli_args.py | 20 +- judgearena/meta_eval/runner.py | 65 +- judgearena/mt_bench/fastchat_compat.py | 8 +- judgearena/mt_bench/mt_bench_utils.py | 68 +- judgearena/mt_bench/pairwise_judging.py | 26 +- judgearena/mt_bench/preset_judging.py | 8 +- judgearena/pairwise_baselines.py | 34 + judgearena/utils/__init__.py | 4 - judgearena/utils/io.py | 82 +-- scripts/fluency/generate_fluency.py | 14 +- .../translate_arena_hard.py | 2 - slurmpilot_scripts/launch_evaluation.py | 1 - .../launch_generation_and_evaluation.py | 1 - tests/test_cli.py | 102 +++ tests/test_config.py | 106 ++++ tests/test_estimate_elo_cache_threading.py | 330 ++++++++++ tests/test_estimate_elo_ratings.py | 38 +- tests/test_evaluate_cache_threading.py | 107 ++++ tests/test_generate_and_evaluate.py | 7 - ...t_generate_and_evaluate_cache_threading.py | 215 +++++++ tests/test_generate_cache_threading.py | 174 ++++++ tests/test_logging.py | 24 + tests/test_meta_eval.py | 120 +--- tests/test_meta_eval_cache_threading.py | 581 ++++++++++++++++++ tests/test_mt_bench_downloads.py | 149 ++++- tests/test_mt_bench_fastchat_compat.py | 34 + .../test_mt_bench_pairwise_cache_threading.py | 159 +++++ tests/test_mt_bench_preset_judging.py | 46 ++ tests/test_no_legacy_runtime_cache.py | 77 +++ tests/test_seed_plumbing.py | 21 +- 38 files changed, 3197 insertions(+), 1047 deletions(-) delete mode 100644 judgearena/meta_eval/cache.py create mode 100644 judgearena/pairwise_baselines.py create mode 100644 tests/test_estimate_elo_cache_threading.py create mode 100644 tests/test_evaluate_cache_threading.py create mode 100644 tests/test_generate_and_evaluate_cache_threading.py create mode 100644 tests/test_generate_cache_threading.py create mode 100644 tests/test_meta_eval_cache_threading.py create mode 100644 tests/test_mt_bench_pairwise_cache_threading.py create mode 100644 tests/test_no_legacy_runtime_cache.py diff --git a/judgearena/config.py b/judgearena/config.py index d86267a..e00462a 100644 --- a/judgearena/config.py +++ b/judgearena/config.py @@ -3,6 +3,9 @@ from __future__ import annotations import argparse +import getpass +from collections.abc import Iterator +from contextlib import contextmanager from pathlib import Path from typing import Literal @@ -10,6 +13,7 @@ from pydantic import BaseModel, ConfigDict, Field, model_validator from pydantic_settings import ( BaseSettings, + CliImplicitFlag, CliSettingsSource, PydanticBaseSettingsSource, SettingsConfigDict, @@ -17,7 +21,21 @@ ) from judgearena.constants import ELO_TASK_PREFIX, ELO_TASK_TO_ARENA, META_EVAL_TASK -from judgearena.generate_and_evaluate import native_pairwise_baseline +from judgearena.inference_cache import InferenceCache +from judgearena.pairwise_baselines import native_pairwise_baseline +from judgearena.store_sync import DEFAULT_CACHE_REPO + +CacheMode = Literal["use", "off", "refresh"] + +_CACHE_CLI_SHORTCUTS = { + "cache.store_root": "store_root", + "cache.cache_mode": "cache_mode", + "cache.cache_hf_repo": "cache_hf_repo", + "cache.cache_fetch": "cache_fetch", + "cache.cache_push": "cache_push", + "cache.cache_create_pr": "cache_create_pr", + "cache.pushed_by": "pushed_by", +} # Set by build_run_config() for the duration of RunConfig() construction. _ACTIVE_CONFIG_PATH: str | None = None @@ -28,6 +46,10 @@ def _drop_none(kwargs: dict[str, object]) -> dict[str, object]: return {k: v for k, v in kwargs.items() if v is not None} +def default_pushed_by() -> str: + return getpass.getuser() + + class ModelArgs(BaseModel): """The model(s) under evaluation and their generation/engine settings.""" @@ -369,8 +391,63 @@ class MetaEvalArgs(BaseModel): """Include human-labeled ties in the primary agreement view.""" +class CacheArgs(BaseModel): + """Unified inference cache settings (SQLite cells with optional HF sync).""" + + model_config = ConfigDict(use_attribute_docstrings=True) + + store_root: str | None = None + """Local root for inference cache cells. When unset, caching is disabled.""" + + cache_mode: CacheMode = "use" + """``use``: read and insert rows. ``off``: always infer. ``refresh``: replace rows.""" + + cache_hf_repo: str = DEFAULT_CACHE_REPO + """Hugging Face dataset repo used when ``cache_fetch`` or ``cache_push`` is set.""" + + cache_fetch: CliImplicitFlag[bool] = False + """Explicit opt-in to fetch remote cache cells before inference.""" + + cache_push: CliImplicitFlag[bool] = False + """Explicit opt-in to push locally produced cache rows after a successful run.""" + + cache_create_pr: CliImplicitFlag[bool] = False + """Push cache updates through a Hugging Face pull request (requires ``cache_push``).""" + + pushed_by: str = Field(default_factory=lambda: default_pushed_by()) + """Provenance label recorded on locally produced cache rows.""" + + @model_validator(mode="after") + def _validate_cache_options(self) -> CacheArgs: + if self.store_root is not None and not self.store_root.strip(): + raise ValueError("cache.store_root must be non-empty when provided.") + if self.cache_fetch or self.cache_push or self.cache_create_pr: + if not self.store_root: + raise ValueError( + "cache.store_root is required when cache_fetch, cache_push, " + "or cache_create_pr is enabled." + ) + if self.cache_fetch or self.cache_push: + if not self.cache_hf_repo.strip(): + raise ValueError( + "cache.cache_hf_repo must be non-empty when cache_fetch or " + "cache_push is enabled." + ) + if self.cache_create_pr and not self.cache_push: + raise ValueError( + "cache.cache_push is required when cache_create_pr is enabled." + ) + if self.cache_mode == "off" and (self.cache_fetch or self.cache_push): + raise ValueError( + "cache_fetch and cache_push cannot be enabled when cache_mode is off." + ) + if self.cache_mode == "refresh" and not self.store_root: + raise ValueError("cache.store_root is required when cache_mode is refresh.") + return self + + class RunArgs(BaseModel): - """Run-level settings: seed, output location, caching, and logging.""" + """Run-level settings: seed, output location, and logging.""" model_config = ConfigDict(use_attribute_docstrings=True) @@ -381,9 +458,6 @@ 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.""" - use_tqdm: bool = False """Show a tqdm progress bar (not compatible with vLLM).""" @@ -403,6 +477,7 @@ class RunConfig(BaseSettings): protected_namespaces=(), nested_model_default_partial_update=True, cli_avoid_json=False, + cli_shortcuts=_CACHE_CLI_SHORTCUTS, use_attribute_docstrings=True, ) @@ -428,7 +503,10 @@ class RunConfig(BaseSettings): """Judge meta-evaluation settings (only for ``meta-eval``).""" run: RunArgs = Field(default_factory=RunArgs) - """Run-level settings (seed, output, caching, logging).""" + """Run-level settings (seed, output, logging).""" + + cache: CacheArgs = Field(default_factory=CacheArgs) + """Unified inference cache settings.""" @model_validator(mode="after") def _validate(self) -> RunConfig: @@ -545,3 +623,48 @@ def dump_config(cfg: RunConfig, path: str | Path) -> None: Path(path).write_text( yaml.safe_dump(cfg.model_dump(), sort_keys=False), encoding="utf-8" ) + + +def meta_eval_cache_task(reference_arena: str) -> str: + """Return the single-segment cache namespace for one meta-eval reference arena.""" + sanitized_arena = reference_arena.replace("/", "_").replace("\\", "_") + return f"{META_EVAL_TASK}-{sanitized_arena}" + + +def inference_cache_task(cfg: RunConfig) -> str: + """Return the cache namespace for a run configuration.""" + if cfg.task == META_EVAL_TASK: + if cfg.meta_eval is None: + raise ValueError("meta_eval config is required for the meta-eval task.") + return meta_eval_cache_task(cfg.meta_eval.reference_arena) + return cfg.task + + +@contextmanager +def open_inference_cache( + cache_args: CacheArgs, + task: str, +) -> Iterator[InferenceCache | None]: + """Open a run-scoped inference cache, or yield ``None`` when disabled.""" + if cache_args.store_root is None: + yield None + return + + with InferenceCache( + store_root=cache_args.store_root, + task=task, + mode=cache_args.cache_mode, + fetch=cache_args.cache_fetch, + push=cache_args.cache_push, + create_pr=cache_args.cache_create_pr, + cache_hf_repo=cache_args.cache_hf_repo, + pushed_by=cache_args.pushed_by, + ) as cache: + yield cache + + +@contextmanager +def inference_cache_session(cfg: RunConfig) -> Iterator[InferenceCache | None]: + """Open a run-scoped inference cache, or yield ``None`` when disabled.""" + with open_inference_cache(cfg.cache, inference_cache_task(cfg)) as cache: + yield cache diff --git a/judgearena/estimate_elo_ratings.py b/judgearena/estimate_elo_ratings.py index 06aaa6d..b9c2d5a 100644 --- a/judgearena/estimate_elo_ratings.py +++ b/judgearena/estimate_elo_ratings.py @@ -2,9 +2,8 @@ import json import re from datetime import UTC, datetime -from functools import partial from pathlib import Path -from typing import TYPE_CHECKING +from typing import Any import numpy as np import pandas as pd @@ -12,6 +11,7 @@ from judgearena.arenas_utils import extract_turn_text, load_arena_dataframe from judgearena.battles import Leaderboard, summarize_bootstrap, write_battles +from judgearena.config import RunConfig, inference_cache_session from judgearena.evaluate import ( PairScore, calibrate_temperature, @@ -21,15 +21,13 @@ ) from judgearena.generate import generate_instructions from judgearena.generate_and_evaluate import _build_generation_kwargs +from judgearena.inference_cache import InferenceCache from judgearena.log import get_logger from judgearena.models import build_default_judge_model_kwargs, make_model from judgearena.repro import write_run_metadata -from judgearena.utils import cache_function_dataframe, compute_pref_summary +from judgearena.utils import compute_pref_summary from judgearena.utils.eval import PrefSummary, Report -if TYPE_CHECKING: - from judgearena.config import RunConfig - logger = get_logger(__name__) @@ -161,23 +159,6 @@ def select_seeded_random_arena_battles( return sampled.reset_index(drop=True), metadata -def _sampling_cache_token( - sampling_metadata: dict[str, object], - *, - n_instructions: int | None, - n_instructions_per_language: int | None, -) -> str: - mode = sampling_metadata.get("sampling_mode") - if mode == "seeded_random": - return ( - "seeded-random_" - f"{sampling_metadata['requested_rows']}_" - f"seed-{sampling_metadata['random_seed']}_" - f"{str(sampling_metadata['sample_fingerprint'])[:12]}" - ) - return f"head_{n_instructions}_{n_instructions_per_language}" - - def _slugify(value: str) -> str: slug = re.sub(r"[^a-z0-9]+", "-", value.lower()).strip("-") return slug or "model" @@ -328,6 +309,130 @@ def _prefs_to_battle_results( return df +def _battle_identity_fallback( + *, + instruction: str, + focal_model: str, + opponent_model: str, + position: str, +) -> str: + payload = { + "instruction": instruction, + "focal_model": focal_model, + "opponent_model": opponent_model, + "position": position, + } + return hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()[:16] + + +def _elo_generation_row_metadata( + *, + arena: str, + df_battles: pd.DataFrame, + instructions: list[str], +) -> list[dict[str, Any]]: + metadata: list[dict[str, Any]] = [] + for index, instruction in enumerate(instructions): + row: dict[str, Any] = { + "arena": arena, + "source": "elo-generation", + } + question_id = ( + df_battles.iloc[index]["question_id"] + if "question_id" in df_battles.columns + else None + ) + if question_id is not None and pd.notna(question_id): + row["question_id"] = str(question_id) + else: + row["instruction_sha256"] = hashlib.sha256( + instruction.encode("utf-8") + ).hexdigest() + metadata.append(row) + return metadata + + +def _elo_judge_row_metadata( + *, + arena: str, + df_battles: pd.DataFrame, + instructions: list[str], + focal_model: str, + opponent_models: list[str], + our_model_is_position_a: np.ndarray, +) -> list[dict[str, Any]]: + metadata: list[dict[str, Any]] = [] + for index, opponent_model in enumerate(opponent_models): + position = "A" if our_model_is_position_a[index] else "B" + row: dict[str, Any] = { + "arena": arena, + "source": "elo-judge", + "focal_model": focal_model, + "opponent_model": opponent_model, + "position": position, + } + question_id = ( + df_battles.iloc[index]["question_id"] + if "question_id" in df_battles.columns + else None + ) + if question_id is not None and pd.notna(question_id): + row["question_id"] = str(question_id) + else: + row["battle_identity"] = _battle_identity_fallback( + instruction=instructions[index], + focal_model=focal_model, + opponent_model=opponent_model, + position=position, + ) + metadata.append(row) + return metadata + + +def _elo_calibration_row_metadata( + *, + arena: str, + cal_battles: pd.DataFrame, +) -> list[dict[str, Any]]: + metadata: list[dict[str, Any]] = [] + for row_index in cal_battles.index: + row: dict[str, Any] = { + "arena": arena, + "source": "elo-calibration", + "purpose": "temperature_calibration", + } + question_id = ( + cal_battles.loc[row_index, "question_id"] + if "question_id" in cal_battles.columns + else None + ) + if question_id is not None and pd.notna(question_id): + row["question_id"] = str(question_id) + else: + row["arena_row_index"] = ( + int(row_index) + if isinstance(row_index, int | np.integer) + else str(row_index) + ) + metadata.append(row) + return metadata + + +def _parse_prefs_from_judge_completions( + judge_completions: list[str], + *, + swap_mode: str, + score_parser: PairScore, +) -> list[float]: + parsed = pd.Series( + [score_parser.parse_model_raw(completion) for completion in judge_completions] + ).apply(lambda value: float("nan") if value is None else value) + if swap_mode == "both": + n_half = len(judge_completions) // 2 + return combine_swapped_prefs(parsed[:n_half], parsed[n_half:]).tolist() + return parsed.tolist() + + def arena_anchor_battles(df_arena_all: pd.DataFrame) -> pd.DataFrame: """Human anchor battles from a loaded arena frame. @@ -352,8 +457,14 @@ def arena_anchor_battles(df_arena_all: pd.DataFrame) -> pd.DataFrame: return df -def main(cfg: "RunConfig") -> dict: +def main(cfg: RunConfig) -> dict: assert cfg.elo is not None # main is dispatched only for elo tasks + with inference_cache_session(cfg) as cache: + return _run_elo(cfg, cache=cache) + + +def _run_elo(cfg: RunConfig, *, cache: InferenceCache | None) -> dict: + assert cfg.elo is not None run_started_at = datetime.now(UTC) rng = np.random.default_rng(cfg.run.seed) @@ -417,47 +528,20 @@ def main(cfg: "RunConfig") -> dict: # previously called evaluated_generation_kwargs() directly and silently # dropped battle_thinking_token_budget). extra_kwargs = _build_generation_kwargs(cfg, cfg.model.name, role="A") - use_tqdm = False - gen_fun = partial( - generate_instructions, + use_tqdm = cfg.run.use_tqdm + instruction_text = instructions.tolist() + completions_df = generate_instructions( + instructions=instructions, + model=cfg.model.name, truncate_input_chars=cfg.generation.truncate_all_input_chars, use_tqdm=use_tqdm, + cache=cache, + row_metadata=_elo_generation_row_metadata( + arena=cfg.elo.arena, + df_battles=df_battles, + instructions=instruction_text, + ), **extra_kwargs, - ) - - def replace_slash(s: str) -> str: - return s.replace("/", "_") - - languages_str = "-".join(sorted(cfg.elo.languages)) if cfg.elo.languages else "all" - extra_kwargs_str = ( - "_".join(f"{k}={v}" for k, v in sorted(extra_kwargs.items())) - if extra_kwargs - else "" - ) - sampling_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"{cfg.elo.arena}_{replace_slash(cfg.model.name)}_" - f"{sampling_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"] @@ -503,113 +587,56 @@ def replace_slash(s: str) -> str: fallback_chat_template=cfg.model.chat_template, ), ) - - def run_judge() -> pd.DataFrame: - judge_chat_model = make_model( - model=cfg.judge.model, - **judge_extra_kwargs, - ) - annotations, annotations_reversed, prefs = judge_and_parse_prefs( - judge_chat_model=judge_chat_model, - instructions=instructions.tolist(), - completions_A=completions_A, - completions_B=completions_B, - swap_mode=cfg.judge.swap_mode, - provide_explanation=cfg.judge.provide_explanation, - strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging, - system_prompt=resolved_prompt.system_prompt, - user_prompt_template=resolved_prompt.user_prompt_template, - prompt_preset=resolved_prompt.preset_name, - truncate_input_chars=cfg.generation.truncate_judge_input_chars, - use_tqdm=use_tqdm, - ) - if annotations_reversed is None: - row_annotations = list(annotations) - row_use_model_a = use_model_a_as_opponent - row_our_pos_a = our_model_is_position_a - row_opponents = list(opponent_models) - else: - # swap_mode="both": dataframe carries 2n rows (AB then BA). - # Position metadata is duplicated; prefs are already oriented - # consistently by judge_and_parse_prefs as [pref_AB, 1 - pref_BA]. - row_annotations = list(annotations) + list(annotations_reversed) - row_use_model_a = np.concatenate( - [use_model_a_as_opponent, use_model_a_as_opponent] - ) - row_our_pos_a = np.concatenate( - [our_model_is_position_a, our_model_is_position_a] - ) - row_opponents = list(opponent_models) + list(opponent_models) - return pd.DataFrame( - { - "judge_completion": [a.judge_completion for a in row_annotations], - "instruction": [a.instruction for a in row_annotations], - "completion_A": [a.completion_A for a in row_annotations], - "completion_B": [a.completion_B for a in row_annotations], - "pref": prefs, - "use_model_a_as_opponent": row_use_model_a, - "our_model_is_position_a": row_our_pos_a, - "opponent_model": row_opponents, - } - ) - - # 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}", + judge_chat_model = make_model( + model=cfg.judge.model, + **judge_extra_kwargs, ) + row_metadata = _elo_judge_row_metadata( + arena=cfg.elo.arena, + df_battles=df_battles, + instructions=instruction_text, + focal_model=cfg.model.name, + opponent_models=opponent_models, + our_model_is_position_a=our_model_is_position_a, + ) + annotations, annotations_reversed, _ = judge_and_parse_prefs( + judge_chat_model=judge_chat_model, + instructions=instruction_text, + completions_A=completions_A, + completions_B=completions_B, + swap_mode=cfg.judge.swap_mode, + provide_explanation=cfg.judge.provide_explanation, + strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging, + system_prompt=resolved_prompt.system_prompt, + user_prompt_template=resolved_prompt.user_prompt_template, + prompt_preset=resolved_prompt.preset_name, + parser_mode=resolved_prompt.parser_mode, + truncate_input_chars=cfg.generation.truncate_judge_input_chars, + use_tqdm=use_tqdm, + cache=cache, + row_metadata=row_metadata, + ) + if annotations_reversed is None: + row_annotations = list(annotations) + else: + row_annotations = list(annotations) + list(annotations_reversed) + judge_completions = [annotation.judge_completion for annotation in row_annotations] - # 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() - our_model_is_position_a = df_judge["our_model_is_position_a"].to_numpy() - opponent_models = df_judge["opponent_model"].tolist() - prefs = df_judge["pref"].tolist() - - # Instruction-index join key per judged battle, so the saved battles link - # back to the arena initial table / completion cache without copying text. - # df_judge repeats the n sampled battles once (AB) or twice (AB then BA for - # swap_mode="both"), so tile the ids to its actual length. if "question_id" in df_battles.columns and len(df_battles): qids = df_battles["question_id"].tolist() - n_rep = (len(df_judge) + len(qids) - 1) // len(qids) - question_ids = (qids * n_rep)[: len(df_judge)] + n_rep = (len(row_annotations) + len(qids) - 1) // len(qids) + question_ids = (qids * n_rep)[: len(row_annotations)] else: - question_ids = [None] * len(df_judge) + question_ids = [None] * len(row_annotations) - logger.debug("First judge output:\n%s", df_judge["judge_completion"].iloc[0][:500]) + logger.debug("First judge output:\n%s", judge_completions[0][:500]) - # Map preferences back to model-name-level battle results. model_name = cfg.model.name - df_llm_judge = _prefs_to_battle_results( - prefs, - our_model_is_position_a, - opponent_models, - model_name, - judge_model=cfg.judge.model, - question_ids=question_ids, - ) - - # Normalize prefs so pref < 0.5 always means our model wins, then summarise - prefs_normalized = pd.Series( - [ - p if (p is None or is_pos_a) else (1 - p) - for p, is_pos_a in zip(prefs, our_model_is_position_a, strict=True) - ] - ) - summary = compute_pref_summary(prefs_normalized) # Anchor the llm-judge battles against the human arena battles. These are # rebuilt from the (revision-pinned) arena, not persisted per run. df_arena = arena_anchor_battles(df_arena_all) - df_results = pd.concat([df_llm_judge, df_arena], ignore_index=True) - # Compute human-only BT ratings as ground-truth reference human_elo = fit_bradley_terry( df_arena, pref_col="pref_hard", baseline_model=cfg.elo.baseline_model @@ -653,19 +680,26 @@ def run_judge() -> pd.DataFrame: for i in cal_battles.index ] - judge_chat_model_cal = make_model( - model=cfg.judge.model, - max_tokens=cfg.judge.max_out_tokens, - **judge_extra_kwargs, + cal_row_metadata = _elo_calibration_row_metadata( + arena=cfg.elo.arena, + cal_battles=cal_battles, ) - cal_annotations, _, cal_prefs = judge_and_parse_prefs( - judge_chat_model=judge_chat_model_cal, + cal_annotations, _, _ = judge_and_parse_prefs( + judge_chat_model=judge_chat_model, instructions=cal_instructions, completions_A=cal_completions_a, completions_B=cal_completions_b, swap_mode=cfg.judge.swap_mode, provide_explanation=cfg.judge.provide_explanation, + strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging, + system_prompt=resolved_prompt.system_prompt, + user_prompt_template=resolved_prompt.user_prompt_template, + prompt_preset=resolved_prompt.preset_name, + parser_mode=resolved_prompt.parser_mode, truncate_input_chars=cfg.generation.truncate_judge_input_chars, + use_tqdm=use_tqdm, + cache=cache, + row_metadata=cal_row_metadata, ) # Build (delta_s, y) pairs from calibration battles. @@ -700,43 +734,44 @@ def run_judge() -> pd.DataFrame: cfg.elo.soft_elo_temperature, ) - # Build the score parser used for the main evaluation run. score_parser = PairScore( temperature=calibrated_temperature if calibrated_temperature is not None - else cfg.elo.soft_elo_temperature + else cfg.elo.soft_elo_temperature, + parser_mode=resolved_prompt.parser_mode, + ) + prefs = _parse_prefs_from_judge_completions( + judge_completions, + swap_mode=cfg.judge.swap_mode, + score_parser=score_parser, ) - # The prefs cached in df_judge were parsed at the default T=0.3, and the - # judge cache key ignores temperature, so they cannot reflect - # --soft-elo-temperature (or a calibrated T*). Re-parse from the stored - # judge completions with this run's score_parser so the soft-ELO bootstrap - # uses the requested temperature. - if cfg.elo.soft_elo: - new_prefs_ab = pd.Series( - [score_parser.parse_model_raw(c) for c in df_judge["judge_completion"]] - ).apply(lambda x: float("nan") if x is None else x) - - if cfg.judge.swap_mode == "both": - # df_judge stores AB then BA completions; re-orient the halves the - # same way run_judge() did. - n_half = len(df_judge) // 2 - prefs = combine_swapped_prefs( - new_prefs_ab[:n_half], new_prefs_ab[n_half:] - ).tolist() - else: - prefs = new_prefs_ab.tolist() - - # Rebuild battle results with the re-parsed prefs. - df_llm_judge = _prefs_to_battle_results( - prefs, - our_model_is_position_a, - opponent_models, - model_name, - judge_model=cfg.judge.model, - question_ids=question_ids, + if cfg.judge.swap_mode == "both": + battle_our_pos_a = np.concatenate( + [our_model_is_position_a, our_model_is_position_a] ) - df_results = pd.concat([df_llm_judge, df_arena], ignore_index=True) + battle_opponents = list(opponent_models) + list(opponent_models) + else: + battle_our_pos_a = our_model_is_position_a + battle_opponents = opponent_models + + df_llm_judge = _prefs_to_battle_results( + prefs, + battle_our_pos_a, + battle_opponents, + model_name, + judge_model=cfg.judge.model, + question_ids=question_ids, + ) + df_results = pd.concat([df_llm_judge, df_arena], ignore_index=True) + + prefs_normalized = pd.Series( + [ + p if (p is None or is_pos_a) else (1 - p) + for p, is_pos_a in zip(prefs, battle_our_pos_a, strict=True) + ] + ) + summary = compute_pref_summary(prefs_normalized) n_bootstraps = cfg.elo.n_bootstraps use_soft = cfg.elo.soft_elo diff --git a/judgearena/evaluate.py b/judgearena/evaluate.py index d2aca18..1415502 100644 --- a/judgearena/evaluate.py +++ b/judgearena/evaluate.py @@ -1,5 +1,8 @@ +from __future__ import annotations + import re from dataclasses import dataclass +from typing import TYPE_CHECKING, Any import numpy as np import pandas as pd @@ -20,6 +23,22 @@ logger = get_logger(__name__) +if TYPE_CHECKING: + from judgearena.inference_cache import InferenceCache + + +def _oriented_row_metadata( + base: list[dict[str, Any] | None] | None, + orientation: str, + n: int, +) -> list[dict[str, Any]]: + oriented: list[dict[str, Any]] = [] + for index in range(n): + row = dict(base[index]) if base and base[index] is not None else {} + row["orientation"] = orientation + oriented.append(row) + return oriented + class PairScore: def __init__(self, *, temperature: float = 0.3, parser_mode: str = "score"): @@ -172,6 +191,60 @@ class JudgeAnnotation: prompt_preset: str = DEFAULT_JUDGE_PROMPT_PRESET +def render_judge_inputs( + instructions: list[str], + completions_A: list[str], + completions_B: list[str], + *, + system_prompt: str | None = None, + user_prompt_template: str | None = None, + truncate_input_chars: int | None = 8192, + provide_explanation: bool = False, + prompt_preset: str = DEFAULT_JUDGE_PROMPT_PRESET, + strip_thinking_before_judging: bool = False, + task: str | None = None, + system_file: str | None = None, + user_file: str | None = None, + multi_turn: bool = False, +) -> list: + """Render judge chat prompt values without running inference.""" + assert len(instructions) == len(completions_A) == len(completions_B) + + resolved_prompt = resolve_judge_prompts( + provide_explanation=provide_explanation, + multi_turn=multi_turn, + prompt_preset=prompt_preset, + system_prompt=system_prompt, + user_prompt_template=user_prompt_template, + task=task, + system_file=system_file, + user_file=user_file, + ) + + message_templates: list[tuple[str, str]] = [] + if resolved_prompt.system_prompt is not None: + message_templates.append(("system", resolved_prompt.system_prompt)) + message_templates.append(("user", resolved_prompt.user_prompt_template)) + prompt_template = ChatPromptTemplate.from_messages(message_templates) + + if strip_thinking_before_judging: + completions_A = [strip_thinking_tags(c) for c in completions_A] + completions_B = [strip_thinking_tags(c) for c in completions_B] + + return prompt_template.batch( + [ + { + "user_prompt": user_prompt, + "completion_A": truncate(completion_A, max_len=truncate_input_chars), + "completion_B": truncate(completion_B, max_len=truncate_input_chars), + } + for user_prompt, completion_A, completion_B in zip( + instructions, completions_A, completions_B, strict=True + ) + ] + ) + + def annotate_battles( judge_chat_model, instructions: list[str], @@ -184,6 +257,12 @@ def annotate_battles( provide_explanation: bool = False, prompt_preset: str = DEFAULT_JUDGE_PROMPT_PRESET, strip_thinking_before_judging: bool = False, + cache: InferenceCache | None = None, + row_metadata: list[dict[str, Any] | None] | None = None, + task: str | None = None, + system_file: str | None = None, + user_file: str | None = None, + multi_turn: bool = False, ) -> list[JudgeAnnotation]: """ Directly evaluate from list of instructions and completions @@ -212,43 +291,41 @@ def annotate_battles( :param use_tqdm: :return: """ - # alternatively pass list of tuples - assert len(instructions) == len(completions_A) == len(completions_B) - resolved_prompt = resolve_judge_prompts( provide_explanation=provide_explanation, + multi_turn=multi_turn, prompt_preset=prompt_preset, system_prompt=system_prompt, user_prompt_template=user_prompt_template, + task=task, + system_file=system_file, + user_file=user_file, ) - message_templates: list[tuple[str, str]] = [] - if resolved_prompt.system_prompt is not None: - message_templates.append(("system", resolved_prompt.system_prompt)) - message_templates.append(("user", resolved_prompt.user_prompt_template)) - prompt_template = ChatPromptTemplate.from_messages(message_templates) - if strip_thinking_before_judging: - completions_A = [strip_thinking_tags(c) for c in completions_A] - completions_B = [strip_thinking_tags(c) for c in completions_B] - - inputs = prompt_template.batch( - [ - { - "user_prompt": user_prompt, - "completion_A": truncate(completion_A, max_len=truncate_input_chars), - "completion_B": truncate(completion_B, max_len=truncate_input_chars), - } - for user_prompt, completion_A, completion_B in zip( - instructions, completions_A, completions_B, strict=True - ) - ] + inputs = render_judge_inputs( + instructions, + completions_A, + completions_B, + system_prompt=resolved_prompt.system_prompt, + user_prompt_template=resolved_prompt.user_prompt_template, + truncate_input_chars=truncate_input_chars, + provide_explanation=provide_explanation, + prompt_preset=prompt_preset, + strip_thinking_before_judging=strip_thinking_before_judging, + task=task, + system_file=system_file, + user_file=user_file, + multi_turn=multi_turn, ) logger.info("Start LLM judge annotation (%d annotations).", len(inputs)) + cache_meta = {"metadata": row_metadata} if row_metadata is not None else None judge_completions = do_inference( chat_model=judge_chat_model, inputs=inputs, use_tqdm=use_tqdm, + cache=cache, + cache_meta=cache_meta, ) annotations = [] @@ -299,7 +376,9 @@ def judge_and_parse_prefs( parser_mode: str = "score", truncate_input_chars: int = 8192, use_tqdm: bool = False, - score_parser: "PairScore | None" = None, + score_parser: PairScore | None = None, + cache: InferenceCache | None = None, + row_metadata: list[dict[str, Any] | None] | None = None, ) -> tuple[list[JudgeAnnotation], list[JudgeAnnotation] | None, pd.Series]: """Run judge annotation and parse preferences, handling swap_mode='both'. @@ -318,6 +397,9 @@ def judge_and_parse_prefs( judge_chat_model, ) + n = len(instructions) + direct_metadata = _oriented_row_metadata(row_metadata, "direct", n) + annotations = annotate_battles( judge_chat_model=judge_chat_model, instructions=instructions, @@ -330,10 +412,13 @@ def judge_and_parse_prefs( prompt_preset=prompt_preset, truncate_input_chars=truncate_input_chars, use_tqdm=use_tqdm, + cache=cache, + row_metadata=direct_metadata, ) annotations_reversed = None if swap_mode == "both": + reversed_metadata = _oriented_row_metadata(row_metadata, "reversed", n) annotations_reversed = annotate_battles( judge_chat_model=judge_chat_model, instructions=instructions, @@ -346,6 +431,8 @@ def judge_and_parse_prefs( prompt_preset=prompt_preset, truncate_input_chars=truncate_input_chars, use_tqdm=use_tqdm, + cache=cache, + row_metadata=reversed_metadata, ) def _none_to_nan(x): diff --git a/judgearena/generate.py b/judgearena/generate.py index c38b2bc..a7d4191 100644 --- a/judgearena/generate.py +++ b/judgearena/generate.py @@ -1,9 +1,49 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + import pandas as pd from langchain_core.prompts import ChatPromptTemplate from judgearena.models import do_inference, make_model from judgearena.utils import strip_thinking_tags, truncate +if TYPE_CHECKING: + from judgearena.inference_cache import InferenceCache + + +def _instruction_index_metadata(indices: list[Any]) -> list[dict[str, Any]]: + return [{"instruction_index": str(index)} for index in indices] + + +def _mt_bench_generation_metadata( + questions: pd.DataFrame, + idxs: list[Any], + *, + turn: int, +) -> list[dict[str, Any]]: + metadata: list[dict[str, Any]] = [] + for idx in idxs: + row = questions.loc[idx] + category = row.get("category") + metadata.append( + { + "instruction_index": str(idx), + "turn": turn, + "category": None if category is None else str(category), + } + ) + return metadata + + +def _subset_metadata( + row_metadata: list[dict[str, Any] | None] | None, + idxs: list[int], +) -> list[dict[str, Any] | None] | None: + if row_metadata is None: + return None + return [row_metadata[i] for i in idxs] + def generate_instructions( instructions: pd.Series, @@ -12,6 +52,8 @@ def generate_instructions( max_tokens: int | None = 32768, use_tqdm: bool = True, system_prompt: str | None = None, + cache: InferenceCache | None = None, + row_metadata: list[dict[str, Any] | None] | None = None, **engine_kwargs, ) -> pd.DataFrame: chat_model = make_model(model, max_tokens=max_tokens, **engine_kwargs) @@ -25,6 +67,7 @@ def generate_instructions( [("system", system_prompt), ("user", "{user_prompt}")] ) + idxs = instructions.index.tolist() inputs = prompt_template.batch( [ { @@ -38,11 +81,17 @@ def generate_instructions( chat_model=chat_model, inputs=inputs, use_tqdm=use_tqdm, + cache=cache, + cache_meta={ + "metadata": row_metadata + if row_metadata is not None + else _instruction_index_metadata(idxs) + }, ) df_outputs = pd.DataFrame( data={ "completion": completions, - "instruction_index": instructions.index.tolist(), + "instruction_index": idxs, }, ) return df_outputs @@ -66,6 +115,8 @@ def _infer_grouped_by_temperature( inputs: list, temperatures: list[float], use_tqdm: bool, + cache: InferenceCache | None = None, + row_metadata: list[dict[str, Any] | None] | None = None, ) -> list[str]: outputs: list[str] = [""] * len(inputs) groups: dict[float, list[int]] = {} @@ -75,6 +126,7 @@ def _infer_grouped_by_temperature( for temp in sorted(groups.keys()): idxs = groups[temp] group_inputs = [inputs[i] for i in idxs] + group_metadata = _subset_metadata(row_metadata, idxs) if provider in {"VLLM", "LlamaCpp"}: _set_temperature_on_model(base_model, temp) @@ -88,6 +140,10 @@ def _infer_grouped_by_temperature( chat_model=group_model, inputs=group_inputs, use_tqdm=use_tqdm, + cache=cache, + cache_meta={"metadata": group_metadata} + if group_metadata is not None + else None, ) for i, out in zip(idxs, group_outs, strict=True): outputs[i] = out @@ -103,6 +159,7 @@ def generate_multiturn( use_tqdm: bool = True, temperature_config: dict[str, float] | None = None, strip_thinking_before_turn_2_prompt: bool = False, + cache: InferenceCache | None = None, **model_kwargs, ) -> pd.DataFrame: """Generate two-turn completions for MT-Bench style questions.""" @@ -135,6 +192,7 @@ def generate_multiturn( for _, row in questions.iterrows() ] ) + turn1_metadata = _mt_bench_generation_metadata(questions, idxs, turn=1) if use_category_temperatures: completions_turn_1 = _infer_grouped_by_temperature( @@ -146,12 +204,16 @@ def generate_multiturn( inputs=turn1_inputs, temperatures=temperatures, use_tqdm=use_tqdm, + cache=cache, + row_metadata=turn1_metadata, ) else: completions_turn_1 = do_inference( chat_model=chat_model, inputs=turn1_inputs, use_tqdm=use_tqdm, + cache=cache, + cache_meta={"metadata": turn1_metadata}, ) turn2_inputs = [] @@ -174,7 +236,7 @@ def generate_multiturn( # Strip ... from the turn-1 answer before the # character cap fires so a long cap lands on the visible answer # rather than deep inside a reasoning block (which would destroy - # the closer and push the thinking fragment into turn 2). + # the closer and push the thinking fragment into turn 2. t1_answer_str = str(t1_answer) if strip_thinking_before_turn_2_prompt: t1_answer_str = strip_thinking_tags(t1_answer_str) @@ -190,6 +252,8 @@ def generate_multiturn( ) ) + turn2_metadata = _mt_bench_generation_metadata(questions, idxs, turn=2) + if use_category_temperatures: completions_turn_2 = _infer_grouped_by_temperature( model_spec=model, @@ -200,12 +264,16 @@ def generate_multiturn( inputs=turn2_inputs, temperatures=temperatures, use_tqdm=use_tqdm, + cache=cache, + row_metadata=turn2_metadata, ) else: completions_turn_2 = do_inference( chat_model=chat_model, inputs=turn2_inputs, use_tqdm=use_tqdm, + cache=cache, + cache_meta={"metadata": turn2_metadata}, ) return pd.DataFrame( @@ -223,25 +291,29 @@ def generate_base( truncate_input_chars: int | None = 8192, max_tokens: int | None = 32768, use_tqdm: bool = False, + cache: InferenceCache | None = None, **engine_kwargs, ) -> pd.DataFrame: - model = make_model(model, max_tokens=max_tokens, **engine_kwargs) + chat_model = make_model(model, max_tokens=max_tokens, **engine_kwargs) + idxs = instructions.index.tolist() inputs = [ truncate(instruction, max_len=truncate_input_chars) for instruction in instructions ] - completions = model.batch( + completions = do_inference( + chat_model=chat_model, inputs=inputs, - max_tokens=max_tokens, + use_tqdm=use_tqdm, + cache=cache, + cache_meta={"metadata": _instruction_index_metadata(idxs)}, ) - completions = [x.content if hasattr(x, "content") else x for x in completions] df_outputs = pd.DataFrame( data={ "completion": completions, - "instruction_index": instructions.index.tolist(), + "instruction_index": idxs, }, ) diff --git a/judgearena/generate_and_evaluate.py b/judgearena/generate_and_evaluate.py index aff889a..1a76e6f 100644 --- a/judgearena/generate_and_evaluate.py +++ b/judgearena/generate_and_evaluate.py @@ -3,27 +3,25 @@ and then evaluates them using a judge model. """ -from collections.abc import Mapping +from __future__ import annotations + +from collections.abc import Callable, Mapping from dataclasses import dataclass from datetime import UTC, datetime from pathlib import Path -from typing import TYPE_CHECKING +from typing import Any import pandas as pd +from judgearena.config import RunConfig, dump_config, inference_cache_session from judgearena.evaluate import judge_and_parse_prefs, resolve_run_judge_prompt from judgearena.generate import generate_base, generate_instructions +from judgearena.inference_cache import InferenceCache from judgearena.instruction_dataset import load_instructions from judgearena.instruction_dataset.arena_hard import ( - ARENA_HARD_BASELINES, download_arena_hard, is_arena_hard_dataset, ) -from judgearena.instruction_dataset.m_arenahard import ( - M_ARENA_HARD_BASELINES, - split_m_arena_hard_dataset, -) -from judgearena.instruction_dataset.mt_bench import MT_BENCH_BASELINES from judgearena.log import ( attach_file_handler, get_logger, @@ -35,32 +33,25 @@ make_model, ) from judgearena.mt_bench.mt_bench_utils import run_mt_bench -from judgearena.repro import write_run_metadata -from judgearena.utils import ( - cache_function_dataframe, - compute_pref_summary, - data_root, - download_hf, - generation_cache_token, - read_df, +from judgearena.pairwise_baselines import ( + ALPACA_EVAL_BASELINES, + PAIRWISE_BASELINES, + native_pairwise_baseline, ) +from judgearena.repro import write_run_metadata +from judgearena.utils import compute_pref_summary, data_root, download_hf, read_df from judgearena.utils.eval import BattleReport -if TYPE_CHECKING: - from judgearena.config import RunConfig - logger = get_logger(__name__) -ALPACA_EVAL_BASELINES: dict[str, str] = { - "alpaca-eval": "gpt4_1106_preview", -} - -PAIRWISE_BASELINES: dict[str, str | Mapping[str, str]] = { - **ALPACA_EVAL_BASELINES, - **ARENA_HARD_BASELINES, - **M_ARENA_HARD_BASELINES, - **MT_BENCH_BASELINES, -} +__all__ = [ + "ALPACA_EVAL_BASELINES", + "PAIRWISE_BASELINES", + "BaselinePlan", + "main", + "native_pairwise_baseline", + "try_load_dataset_completions", +] def try_load_dataset_completions( @@ -112,13 +103,13 @@ class BaselinePlan: baseline_by_index: pd.Series @classmethod - def flat(cls, model: str, *, index: pd.Index) -> "BaselinePlan": + def flat(cls, model: str, *, index: pd.Index) -> BaselinePlan: return cls( baseline_by_index=pd.Series(model, index=index, name="model_B", dtype=str) ) @classmethod - def per_row(cls, series: pd.Series) -> "BaselinePlan": + def per_row(cls, series: pd.Series) -> BaselinePlan: return cls(baseline_by_index=series.astype(str).rename("model_B")) @property @@ -145,17 +136,6 @@ def aligned_to(self, index: pd.Index) -> pd.Series: return self.baseline_by_index.loc[index] -def native_pairwise_baseline(task: str) -> str | Mapping[str, str] | None: - """Return the dataset-native pairwise baseline, if the task defines one.""" - if task in PAIRWISE_BASELINES: - return PAIRWISE_BASELINES[task] - parsed_m_arena_hard = split_m_arena_hard_dataset(task) - if parsed_m_arena_hard is not None: - version_key, _lang_or_subset = parsed_m_arena_hard - return PAIRWISE_BASELINES[version_key] - return None - - def _resolve_baseline_plan( *, task: str, model_b: str | None, instructions_df: pd.DataFrame ) -> BaselinePlan: @@ -192,7 +172,7 @@ def _resolve_baseline_plan( def _build_generation_kwargs( - cfg: "RunConfig", model_spec: str, *, role: str + cfg: RunConfig, model_spec: str, *, role: str ) -> dict[str, object]: """Battle-model kwargs, adding a thinking-token sub-budget when requested.""" if role == "A": @@ -220,49 +200,35 @@ def load_contexts(dataset: str) -> pd.Series: return pd.read_csv(path).loc[:, "instruction"] -def main(cfg: "RunConfig"): - """ - 1) take as input: - * task (dataset), make sure instruct-completion works - * model to generate output from - * llm used for judge - * number of annotations - * path to save annotations - 2) create completions - 3) create annotations - """ +def _setup_result_folder( + cfg: RunConfig, result_name: str, run_started_at: datetime +) -> Path: + run_ts = run_started_at.strftime("%Y%m%d_%H%M%S") + res_folder = Path(cfg.run.result_folder) / f"{result_name}-{run_ts}" + res_folder.mkdir(parents=True, exist_ok=True) + if not cfg.run.no_log_file: + attach_file_handler(make_run_log_path(res_folder)) + return res_folder - 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 +def _pairwise_result_name(cfg: RunConfig, baseline_plan: BaselinePlan) -> str: + name = ( + f"{cfg.task}-{cfg.model.name}-{baseline_plan.display_name}-{cfg.judge.model}" + f"-{cfg.judge.swap_mode}" + ) + return name.replace("/", "_") + + +def _mt_bench_result_name(cfg: RunConfig, model_b: str) -> str: + name = ( + f"{cfg.task}-{cfg.model.name}-{model_b}-{cfg.judge.model}-{cfg.judge.swap_mode}" + ) + return name.replace("/", "_") - if cfg.task == "mt-bench": - model_b = cfg.model.baseline or native_pairwise_baseline(cfg.task) - if not isinstance(model_b, str): - raise ValueError("MT-Bench requires a flat native baseline.") - name = f"{cfg.task}-{cfg.model.name}-{model_b}-{cfg.judge.model}" - name += f"-{cfg.judge.swap_mode}" - name = name.replace("/", "_") - run_ts = run_started_at.strftime("%Y%m%d_%H%M%S") - res_folder = Path(cfg.run.result_folder) / f"{name}-{run_ts}" - res_folder.mkdir(parents=True, exist_ok=True) - if not cfg.run.no_log_file: - attach_file_handler(make_run_log_path(res_folder)) - return run_mt_bench( - cfg, - ignore_cache, - res_folder=res_folder, - result_name=name, - ) - # Currrently, we run context evaluation +def _load_task_instructions(cfg: RunConfig) -> tuple[pd.DataFrame, pd.Series, bool]: is_fluency_task = "fluency" in cfg.task if is_fluency_task: - # if cfg.task = "fluency-french", we map to "french-contexts.csv" - # to match files in https://huggingface.co/datasets/geoalgo/multilingual-contexts-to-be-completed lang = cfg.task.split("-")[-1] instructions = load_contexts(f"{lang}-contexts.csv") instructions_df = pd.DataFrame({"instruction": instructions.values}) @@ -281,84 +247,83 @@ def main(cfg: "RunConfig"): if cfg.generation.n_instructions is not None: instructions_df = instructions_df.head(n_instructions) instructions = instructions.head(n_instructions) - - baseline_plan = _resolve_baseline_plan( - task=cfg.task, model_b=cfg.model.baseline, instructions_df=instructions_df + return instructions_df, instructions, is_fluency_task + + +def _align_completion_series(df: pd.DataFrame, *, index: pd.Index) -> pd.Series: + return df.set_index("instruction_index").loc[index, "completion"] + + +def _load_or_generate_completions( + *, + cfg: RunConfig, + model_spec: str, + role: str, + instructions: pd.Series, + generation_function: Callable[..., pd.DataFrame], + cache: InferenceCache | None, + n_instructions: int, +) -> pd.Series: + preloaded = try_load_dataset_completions(cfg.task, model_spec, n_instructions) + if preloaded is not None: + return _align_completion_series(preloaded, index=instructions.index) + + generation_kwargs = _build_generation_kwargs(cfg, model_spec, role=role) + generated = generation_function( + instructions=instructions, + model=model_spec, + truncate_input_chars=cfg.generation.truncate_all_input_chars, + use_tqdm=cfg.run.use_tqdm, + cache=cache, + **generation_kwargs, ) - - name = f"{cfg.task}-{cfg.model.name}-{baseline_plan.display_name}-{cfg.judge.model}" - name += f"-{cfg.judge.swap_mode}" - name = name.replace("/", "_") - run_ts = run_started_at.strftime("%Y%m%d_%H%M%S") - res_folder = Path(cfg.run.result_folder) / f"{name}-{run_ts}" - res_folder.mkdir(parents=True, exist_ok=True) - if not cfg.run.no_log_file: - attach_file_handler(make_run_log_path(res_folder)) - - logger.info( - "Using task %s and evaluating %s against baseline %s.", - cfg.task, - cfg.model.name, - baseline_plan.display_name, + return _align_completion_series(generated, index=instructions.index) + + +def _generate_battle_completions( + *, + cfg: RunConfig, + instructions: pd.Series, + baseline_plan: BaselinePlan, + generation_function: Callable[..., pd.DataFrame], + cache: InferenceCache | None, + n_instructions: int, +) -> tuple[pd.Series, pd.Series, pd.Series]: + completions_a = _load_or_generate_completions( + cfg=cfg, + model_spec=cfg.model.name, + role="A", + instructions=instructions, + generation_function=generation_function, + cache=cache, + n_instructions=n_instructions, ) - logger.info( - "Generating completions for task %s with model %s and baseline %s " - "(or loading them directly if present)", - cfg.task, - cfg.model.name, - baseline_plan.display_name, - ) - - # TODO currently we just support base models for fluency, we could also support instruction-tuned models - generation_function = generate_base if is_fluency_task else generate_instructions - - def _run_generation( - model_spec: str, *, generation_kwargs: dict[str, object] - ) -> pd.DataFrame: - return generation_function( - instructions=instructions, - model=model_spec, - truncate_input_chars=cfg.generation.truncate_all_input_chars, - use_tqdm=cfg.run.use_tqdm, - **generation_kwargs, - ) - - def _align_completion_series(df: pd.DataFrame) -> pd.Series: - return df.set_index("instruction_index").loc[instructions.index, "completion"] - - def _load_or_generate_completions(model_spec: str, *, role: str) -> pd.Series: - preloaded = try_load_dataset_completions(cfg.task, model_spec, n_instructions) - if preloaded is not None: - return _align_completion_series(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}" - ), - ) - return _align_completion_series(generated) - - completions_A = _load_or_generate_completions(cfg.model.name, role="A") - baseline_per_index = baseline_plan.aligned_to(instructions.index) if baseline_plan.is_flat: - completions_B = _load_or_generate_completions( - baseline_plan.single_model, role="B" + completions_b = _load_or_generate_completions( + cfg=cfg, + model_spec=baseline_plan.single_model, + role="B", + instructions=instructions, + generation_function=generation_function, + cache=cache, + n_instructions=n_instructions, ) else: per_baseline_completions = { - model: _load_or_generate_completions(model, role="B") + model: _load_or_generate_completions( + cfg=cfg, + model_spec=model, + role="B", + instructions=instructions, + generation_function=generation_function, + cache=cache, + n_instructions=n_instructions, + ) for model in baseline_plan.unique_models } - completions_B = pd.Series( + completions_b = pd.Series( [ per_baseline_completions[model].loc[instruction_index] for instruction_index, model in baseline_per_index.items() @@ -367,16 +332,27 @@ def _load_or_generate_completions(model_spec: str, *, role: str) -> pd.Series: name="completion", ) - logger.debug("First instruction/context: %s", instructions.values[0]) - logger.debug("First completion of %s:\n%s", cfg.model.name, completions_A.values[0]) - logger.debug( - "First completion of %s:\n%s", - baseline_plan.display_name, - completions_B.values[0], - ) - logger.info("Evaluating completions with judge %s.", cfg.judge.model) + return completions_a, completions_b, baseline_per_index - judge_chat_model = make_model( + +def _judge_row_metadata( + *, + instruction_indices: list[Any], + model_a: str, + baseline_per_row: pd.Series, +) -> list[dict[str, Any]]: + return [ + { + "instruction_index": str(instruction_index), + "model_A": model_a, + "model_B": str(baseline_per_row.loc[instruction_index]), + } + for instruction_index in instruction_indices + ] + + +def _build_judge_model(cfg: RunConfig): + return make_model( model=cfg.judge.model, **build_default_judge_model_kwargs( cfg.judge.model, @@ -387,32 +363,27 @@ def _load_or_generate_completions(model_spec: str, *, role: str) -> pd.Series: ), ) - # save the resolved config for results analysis (round-trippable via --config_path) - from judgearena.config import dump_config - - dump_config(cfg, res_folder / "config.yaml") - - logger.info("Saving results to %s", res_folder) - resolved_prompt = resolve_run_judge_prompt(cfg.task, cfg.judge) - - annotations, annotations_reversed, prefs = judge_and_parse_prefs( - judge_chat_model=judge_chat_model, - instructions=instructions.head(n_instructions).tolist(), - completions_A=completions_A.head(n_instructions).tolist(), - completions_B=completions_B.head(n_instructions).tolist(), - swap_mode=cfg.judge.swap_mode, - provide_explanation=cfg.judge.provide_explanation, - strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging, - system_prompt=resolved_prompt.system_prompt, - user_prompt_template=resolved_prompt.user_prompt_template, - prompt_preset=resolved_prompt.preset_name, - parser_mode=resolved_prompt.parser_mode, - truncate_input_chars=cfg.generation.truncate_judge_input_chars, - use_tqdm=cfg.run.use_tqdm, - ) - eval_instruction_index = instructions.head(n_instructions).index.tolist() +def _persist_pairwise_results( + *, + cfg: RunConfig, + res_folder: Path, + name: str, + baseline_plan: BaselinePlan, + instructions: pd.Series, + completions_a: pd.Series, + completions_b: pd.Series, + baseline_per_index: pd.Series, + annotations: list, + annotations_reversed: list | None, + prefs: pd.Series, + resolved_prompt, + run_started_at: datetime, +) -> pd.Series: + n_instructions = len(instructions) + eval_instruction_index = instructions.index.tolist() baseline_per_eval = baseline_per_index.loc[eval_instruction_index] + df = pd.DataFrame(annotations) df["instruction_index"] = eval_instruction_index df["model_A"] = cfg.model.name @@ -429,9 +400,7 @@ def _load_or_generate_completions(model_spec: str, *, role: str) -> pd.Series: df.to_csv(res_folder / f"{name}-annotations.csv", index=False) - # compute and report statistics summary = compute_pref_summary(prefs) - report = BattleReport( task=cfg.task, model_a=cfg.model.name, @@ -459,10 +428,6 @@ def _load_or_generate_completions(model_spec: str, *, role: str) -> pd.Series: report.render() report.save(res_folder / f"results-{name}.json") - eval_instructions = instructions.head(n_instructions).tolist() - eval_completions_A = completions_A.head(n_instructions).tolist() - eval_completions_B = completions_B.head(n_instructions).tolist() - try: write_run_metadata( output_dir=res_folder, @@ -471,16 +436,140 @@ def _load_or_generate_completions(model_spec: str, *, role: str) -> pd.Series: results=results, input_payloads={ "instruction_index": eval_instruction_index, - "instructions": eval_instructions, - "completions_A": eval_completions_A, - "completions_B": eval_completions_B, + "instructions": instructions.head(n_instructions).tolist(), + "completions_A": completions_a.head(n_instructions).tolist(), + "completions_B": completions_b.head(n_instructions).tolist(), "baseline_model_B": baseline_per_eval.tolist(), }, judge_system_prompt=resolved_prompt.system_prompt, judge_user_prompt_template=resolved_prompt.user_prompt_template, started_at_utc=run_started_at, ) - except OSError as e: - logger.warning("Failed to write run metadata: %s", e) + except OSError as exc: + logger.warning("Failed to write run metadata: %s", exc) return prefs + + +def _run_pairwise_task(cfg: RunConfig, *, run_started_at: datetime) -> pd.Series: + instructions_df, instructions, is_fluency_task = _load_task_instructions(cfg) + n_instructions = len(instructions) + + baseline_plan = _resolve_baseline_plan( + task=cfg.task, model_b=cfg.model.baseline, instructions_df=instructions_df + ) + name = _pairwise_result_name(cfg, baseline_plan) + res_folder = _setup_result_folder(cfg, name, run_started_at) + + logger.info( + "Using task %s and evaluating %s against baseline %s.", + cfg.task, + cfg.model.name, + baseline_plan.display_name, + ) + logger.info( + "Generating completions for task %s with model %s and baseline %s " + "(or loading them directly if present)", + cfg.task, + cfg.model.name, + baseline_plan.display_name, + ) + + generation_function = generate_base if is_fluency_task else generate_instructions + + with inference_cache_session(cfg) as cache: + completions_a, completions_b, baseline_per_index = _generate_battle_completions( + cfg=cfg, + instructions=instructions, + baseline_plan=baseline_plan, + generation_function=generation_function, + cache=cache, + n_instructions=n_instructions, + ) + + logger.debug("First instruction/context: %s", instructions.values[0]) + logger.debug( + "First completion of %s:\n%s", cfg.model.name, completions_a.values[0] + ) + logger.debug( + "First completion of %s:\n%s", + baseline_plan.display_name, + completions_b.values[0], + ) + logger.info("Evaluating completions with judge %s.", cfg.judge.model) + + dump_config(cfg, res_folder / "config.yaml") + logger.info("Saving results to %s", res_folder) + + resolved_prompt = resolve_run_judge_prompt(cfg.task, cfg.judge) + judge_chat_model = _build_judge_model(cfg) + eval_instruction_index = instructions.index.tolist() + row_metadata = _judge_row_metadata( + instruction_indices=eval_instruction_index, + model_a=cfg.model.name, + baseline_per_row=baseline_per_index, + ) + + annotations, annotations_reversed, prefs = judge_and_parse_prefs( + judge_chat_model=judge_chat_model, + instructions=instructions.tolist(), + completions_A=completions_a.tolist(), + completions_B=completions_b.tolist(), + swap_mode=cfg.judge.swap_mode, + provide_explanation=cfg.judge.provide_explanation, + strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging, + system_prompt=resolved_prompt.system_prompt, + user_prompt_template=resolved_prompt.user_prompt_template, + prompt_preset=resolved_prompt.preset_name, + parser_mode=resolved_prompt.parser_mode, + truncate_input_chars=cfg.generation.truncate_judge_input_chars, + use_tqdm=cfg.run.use_tqdm, + cache=cache, + row_metadata=row_metadata, + ) + + return _persist_pairwise_results( + cfg=cfg, + res_folder=res_folder, + name=name, + baseline_plan=baseline_plan, + instructions=instructions, + completions_a=completions_a, + completions_b=completions_b, + baseline_per_index=baseline_per_index, + annotations=annotations, + annotations_reversed=annotations_reversed, + prefs=prefs, + resolved_prompt=resolved_prompt, + run_started_at=run_started_at, + ) + + +def main(cfg: RunConfig): + """ + 1) take as input: + * task (dataset), make sure instruct-completion works + * model to generate output from + * llm used for judge + * number of annotations + * path to save annotations + 2) create completions + 3) create annotations + """ + run_started_at = datetime.now(UTC) + + if cfg.task == "mt-bench": + model_b = cfg.model.baseline or native_pairwise_baseline(cfg.task) + if not isinstance(model_b, str): + raise ValueError("MT-Bench requires a flat native baseline.") + result_name = _mt_bench_result_name(cfg, model_b) + res_folder = _setup_result_folder(cfg, result_name, run_started_at) + with inference_cache_session(cfg) as cache: + return run_mt_bench( + cfg, + cache=cache, + res_folder=res_folder, + result_name=result_name, + ) + + return _run_pairwise_task(cfg, run_started_at=run_started_at) diff --git a/judgearena/log.py b/judgearena/log.py index a2a35c7..045717c 100644 --- a/judgearena/log.py +++ b/judgearena/log.py @@ -93,7 +93,13 @@ def configure_logging( # --- console handler --- # Avoid duplicate handlers when configure_logging is called more than once # (e.g. in tests). - if not root.handlers: + console_handlers = [ + handler + for handler in root.handlers + if isinstance(handler, logging.StreamHandler) + and not isinstance(handler, logging.FileHandler) + ] + if not console_handlers: handler = logging.StreamHandler(sys.stderr) handler.setLevel(level) handler.setFormatter( @@ -101,11 +107,14 @@ def configure_logging( ) root.addHandler(handler) else: - for h in root.handlers: - if isinstance(h, logging.StreamHandler) and not isinstance( - h, logging.FileHandler - ): - h.setLevel(level) + for handler in console_handlers: + handler.setLevel(level) + if handler.stream is not sys.stderr: + try: + handler.setStream(sys.stderr) + except ValueError: + # Test/output capture can close the previous stream. + handler.stream = sys.stderr # --- file handler (explicit --log-file) --- if log_file is not None: diff --git a/judgearena/meta_eval/annotate.py b/judgearena/meta_eval/annotate.py index 7883717..5a643b2 100644 --- a/judgearena/meta_eval/annotate.py +++ b/judgearena/meta_eval/annotate.py @@ -2,15 +2,13 @@ from __future__ import annotations +from typing import Any + import pandas as pd from judgearena.arenas_utils import extract_turn_text from judgearena.evaluate import JudgeAnnotation, annotate_battles -from judgearena.meta_eval.cache import ( - AnnotationCache, - AnnotationEntry, - AnnotationKey, -) +from judgearena.inference_cache import InferenceCache from judgearena.meta_eval.cli_args import CliMetaEvalArgs from judgearena.meta_eval.cost import ( estimate_annotation_cost_usd, @@ -43,7 +41,7 @@ def _swap_batch(df_batch: pd.DataFrame) -> pd.DataFrame: def _annotations_to_frame( df_batch: pd.DataFrame, - annotations, + annotations: list[JudgeAnnotation], *, prompt_mode: str, judge_model: str, @@ -78,110 +76,51 @@ def _annotations_to_frame( return add_parsed_columns(pd.DataFrame(rows), prompt_mode) -def _judge_cache_name(args: CliMetaEvalArgs) -> str: - if args.prompt_mode == "standard": - return args.judge_model - return f"{args.judge_model}::{args.prompt_mode}" - - -def _cache_keys( +def _row_metadata( df_batch: pd.DataFrame, + args: CliMetaEvalArgs, *, - judge: str, -) -> list[AnnotationKey]: + orientation: str, +) -> list[dict[str, Any]]: return [ - AnnotationKey( - benchmark=str(battle["benchmark"]), - instruction_id=str(battle["question_id"]), - model_a=str(battle["model_a"]), - model_b=str(battle["model_b"]), - judge=judge, - ) + { + "reference_arena": args.reference_arena, + "benchmark": str(battle["benchmark"]), + "question_id": str(battle["question_id"]), + "presented_model_a": str(battle["model_a"]), + "presented_model_b": str(battle["model_b"]), + "prompt_mode": args.prompt_mode, + "orientation": orientation, + } for _, battle in df_batch.iterrows() ] -def _annotation_from_entry( - entry: AnnotationEntry, - *, - instruction: str, - completion_a: str, - completion_b: str, -) -> JudgeAnnotation: - return JudgeAnnotation( - instruction=instruction, - completion_A=completion_a, - completion_B=completion_b, - judge_completion=entry.judge_completion, - judge_input=entry.judge_input, - ) - - -def _run_cached_batch( +def _run_batch( df_batch: pd.DataFrame, args: CliMetaEvalArgs, *, judge_chat_model, - annotation_cache: AnnotationCache, prompt_spec, + cache: InferenceCache | None, swapped: bool, ) -> pd.DataFrame: working = _swap_batch(df_batch) if swapped else df_batch instructions, completions_a, completions_b = _battle_texts(working) - judge = _judge_cache_name(args) - keys = _cache_keys(working, judge=judge) - cached_entries = ( - [None] * len(keys) - if args.ignore_cache - else annotation_cache.batch_get_annotations(keys) + orientation = "swapped" if swapped else "forward" + annotations = annotate_battles( + judge_chat_model=judge_chat_model, + instructions=instructions, + completions_A=completions_a, + completions_B=completions_b, + system_prompt=prompt_spec.system_prompt, + user_prompt_template=prompt_spec.user_prompt_template, + truncate_input_chars=args.truncate_judge_input_chars, + provide_explanation=args.provide_explanation, + strip_thinking_before_judging=args.strip_thinking_before_judging, + cache=cache, + row_metadata=_row_metadata(working, args, orientation=orientation), ) - missing_indices = [ - index for index, entry in enumerate(cached_entries) if entry is None - ] - - if missing_indices: - new_annotations = annotate_battles( - judge_chat_model=judge_chat_model, - instructions=[instructions[index] for index in missing_indices], - completions_A=[completions_a[index] for index in missing_indices], - completions_B=[completions_b[index] for index in missing_indices], - system_prompt=prompt_spec.system_prompt, - user_prompt_template=prompt_spec.user_prompt_template, - truncate_input_chars=args.truncate_judge_input_chars, - provide_explanation=args.provide_explanation, - ) - new_entries = [ - AnnotationEntry( - **key.__dict__, - judge_input=annotation.judge_input or "", - judge_completion=annotation.judge_completion, - ) - for key, annotation in zip( - [keys[index] for index in missing_indices], - new_annotations, - strict=True, - ) - ] - annotation_cache.batch_put(new_entries) - for index, entry in zip(missing_indices, new_entries, strict=True): - cached_entries[index] = entry - - annotations = [ - _annotation_from_entry( - entry, - instruction=instruction, - completion_a=completion_a, - completion_b=completion_b, - ) - for entry, instruction, completion_a, completion_b in zip( - cached_entries, - instructions, - completions_a, - completions_b, - strict=True, - ) - if entry is not None - ] return _annotations_to_frame( working, annotations, @@ -221,53 +160,47 @@ def annotate_sample( *, judge_chat_model, prompt_spec, - annotation_cache: AnnotationCache | None = None, + cache: InferenceCache | None = None, ) -> pd.DataFrame: n_total = len(df_sample) n_batches = (n_total + args.batch_size - 1) // args.batch_size parts: list[pd.DataFrame] = [] - owns_cache = annotation_cache is None - cache = annotation_cache or AnnotationCache() - try: - for batch_idx in range(n_batches): - start = batch_idx * args.batch_size - end = min(start + args.batch_size, n_total) - df_batch = df_sample.iloc[start:end].copy() - batch_df = _run_cached_batch( + for batch_idx in range(n_batches): + start = batch_idx * args.batch_size + end = min(start + args.batch_size, n_total) + df_batch = df_sample.iloc[start:end].copy() + batch_df = _run_batch( + df_batch, + args, + judge_chat_model=judge_chat_model, + prompt_spec=prompt_spec, + cache=cache, + swapped=False, + ) + parts.append( + _normalize_pass_frame( + batch_df, + df_batch, + orientation="forward", + ) + ) + + if args.swap_mode == "both": + swapped_df = _run_batch( df_batch, args, judge_chat_model=judge_chat_model, - annotation_cache=cache, prompt_spec=prompt_spec, - swapped=False, + cache=cache, + swapped=True, ) parts.append( _normalize_pass_frame( - batch_df, + swapped_df, df_batch, - orientation="forward", + orientation="swapped", ) ) - if args.swap_mode == "both": - swapped_df = _run_cached_batch( - df_batch, - args, - judge_chat_model=judge_chat_model, - annotation_cache=cache, - prompt_spec=prompt_spec, - swapped=True, - ) - parts.append( - _normalize_pass_frame( - swapped_df, - df_batch, - orientation="swapped", - ) - ) - finally: - if owns_cache: - cache.close() - return pd.concat(parts, ignore_index=True) diff --git a/judgearena/meta_eval/cache.py b/judgearena/meta_eval/cache.py deleted file mode 100644 index 40acd7e..0000000 --- a/judgearena/meta_eval/cache.py +++ /dev/null @@ -1,165 +0,0 @@ -"""SQLite-backed cache for meta-evaluation judge annotations.""" - -from __future__ import annotations - -import sqlite3 -from dataclasses import astuple, dataclass, field, fields -from datetime import UTC, datetime -from itertools import groupby -from pathlib import Path - -from judgearena.utils import data_root - -DEFAULT_DB_DIR = data_root / "cache" / "db" - - -@dataclass(frozen=True) -class AnnotationEntry: - benchmark: str - instruction_id: str - model_a: str - model_b: str - judge: str - judge_input: str - judge_completion: str - reasoning_content: str = "" - input_tokens: int = 0 - output_tokens: int = 0 - reasoning_tokens: int = 0 - date: str = field(default_factory=lambda: datetime.now(UTC).isoformat()) - - @classmethod - def field_names(cls) -> tuple[str, ...]: - return tuple(field_info.name for field_info in fields(cls)) - - @classmethod - def key_fields(cls) -> tuple[str, ...]: - return ("benchmark", "instruction_id", "model_a", "model_b", "judge") - - -@dataclass(frozen=True) -class AnnotationKey: - benchmark: str - instruction_id: str - model_a: str - model_b: str - judge: str - - @classmethod - def field_names(cls) -> tuple[str, ...]: - return tuple(field_info.name for field_info in fields(cls)) - - -def _db_path(db_dir: Path, benchmark: str, judge: str) -> Path: - return db_dir / benchmark / f"{judge.replace('/', '_')}.db" - - -class AnnotationCache: - """Persistent per-battle cache matching the original meta-eval pipeline.""" - - def __init__(self, db_dir: Path | str = DEFAULT_DB_DIR) -> None: - self._db_dir = Path(db_dir) - self._connections: dict[tuple[str, str], sqlite3.Connection] = {} - - def batch_get_annotations( - self, keys: list[AnnotationKey] - ) -> list[AnnotationEntry | None]: - column_names = ", ".join(AnnotationEntry.field_names()) - results = [] - for key in keys: - row = ( - self._connection(key.benchmark, key.judge) - .execute( - f"SELECT {column_names} FROM annotations " - f"WHERE {self._where_clause()}", - astuple(key), - ) - .fetchone() - ) - results.append(AnnotationEntry(*row) if row else None) - return results - - def batch_put(self, entries: list[AnnotationEntry]) -> None: - if not entries: - return - column_names = ", ".join(AnnotationEntry.field_names()) - placeholders = ", ".join("?" for _ in AnnotationEntry.field_names()) - sql = ( - f"INSERT OR REPLACE INTO annotations ({column_names}) " - f"VALUES ({placeholders})" - ) - - def cache_partition(entry: AnnotationEntry) -> tuple[str, str]: - return entry.benchmark, entry.judge - - for (benchmark, judge), group in groupby( - sorted(entries, key=cache_partition), - key=cache_partition, - ): - connection = self._connection(benchmark, judge) - connection.executemany(sql, [astuple(entry) for entry in group]) - connection.commit() - - def close(self) -> None: - for connection in self._connections.values(): - connection.close() - self._connections.clear() - - def _connection(self, benchmark: str, judge: str) -> sqlite3.Connection: - key = (benchmark, judge) - if key not in self._connections: - path = _db_path(self._db_dir, benchmark, judge) - path.parent.mkdir(parents=True, exist_ok=True) - connection = sqlite3.connect( - str(path), - check_same_thread=False, - timeout=30, - ) - connection.execute("PRAGMA journal_mode=WAL") - self._connections[key] = connection - self._create_table(connection) - return self._connections[key] - - @staticmethod - def _create_table(connection: sqlite3.Connection) -> None: - integer_fields = {"input_tokens", "output_tokens", "reasoning_tokens"} - default_text_fields = {"reasoning_content", "date"} - column_definitions = ", ".join( - ( - f"{name} INTEGER NOT NULL DEFAULT 0" - if name in integer_fields - else ( - f"{name} TEXT NOT NULL DEFAULT ''" - if name in default_text_fields - else f"{name} TEXT NOT NULL" - ) - ) - for name in AnnotationEntry.field_names() - ) - key_columns = ", ".join(AnnotationEntry.key_fields()) - connection.execute( - "CREATE TABLE IF NOT EXISTS annotations " - f"({column_definitions}, UNIQUE ({key_columns}))" - ) - connection.execute( - "CREATE INDEX IF NOT EXISTS idx_annotation_key " - f"ON annotations ({key_columns})" - ) - migrations = [ - "ALTER TABLE annotations ADD COLUMN " - "reasoning_content TEXT NOT NULL DEFAULT ''", - "ALTER TABLE annotations ADD COLUMN input_tokens INTEGER NOT NULL DEFAULT 0", - "ALTER TABLE annotations ADD COLUMN output_tokens INTEGER NOT NULL DEFAULT 0", - "ALTER TABLE annotations ADD COLUMN " - "reasoning_tokens INTEGER NOT NULL DEFAULT 0", - ] - for migration in migrations: - try: - connection.execute(migration) - except sqlite3.OperationalError: - pass - connection.commit() - - @staticmethod - def _where_clause() -> str: - return " AND ".join(f"{column} = ?" for column in AnnotationKey.field_names()) diff --git a/judgearena/meta_eval/cli_args.py b/judgearena/meta_eval/cli_args.py index f301dcd..febf3a6 100644 --- a/judgearena/meta_eval/cli_args.py +++ b/judgearena/meta_eval/cli_args.py @@ -2,9 +2,11 @@ from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import asdict, dataclass, field +from typing import Any -from judgearena.config import RunConfig +from judgearena.config import CacheArgs, RunConfig +from judgearena.model_adapters import normalize_constructor_settings PROMPT_MODES = ( "standard", @@ -32,7 +34,7 @@ class CliMetaEvalArgs: exclude_human_ties: bool = True provide_explanation: bool = False swap_mode: str = "fixed" - ignore_cache: bool = False + strip_thinking_before_judging: bool = False truncate_judge_input_chars: int | None = None max_out_tokens_judge: int = 32768 max_model_len: int | None = None @@ -40,6 +42,7 @@ class CliMetaEvalArgs: result_folder: str = "results" engine_kwargs: dict[str, object] = field(default_factory=dict) no_log_file: bool = False + cache: CacheArgs = field(default_factory=CacheArgs) def __post_init__(self) -> None: if self.swap_mode not in {"fixed", "both"}: @@ -50,6 +53,14 @@ def __post_init__(self) -> None: f"expected one of {PROMPT_MODES}." ) + def to_jsonable(self) -> dict[str, Any]: + """Serialize runtime args for ``args.json`` without secret values.""" + payload = asdict(self) + payload["cache"] = self.cache.model_dump() + sanitized_engine_kwargs = normalize_constructor_settings(self.engine_kwargs) + payload["engine_kwargs"] = sanitized_engine_kwargs or {} + return payload + def meta_eval_args_from_config(cfg: RunConfig) -> CliMetaEvalArgs: """Map the shared hierarchical run config to meta-eval runtime arguments.""" @@ -76,7 +87,7 @@ def meta_eval_args_from_config(cfg: RunConfig) -> CliMetaEvalArgs: exclude_human_ties=not cfg.meta_eval.include_human_ties, provide_explanation=cfg.judge.provide_explanation, swap_mode=cfg.judge.swap_mode, - ignore_cache=cfg.run.ignore_cache, + strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging, truncate_judge_input_chars=cfg.generation.truncate_judge_input_chars, max_out_tokens_judge=max_out_tokens_judge, max_model_len=max_model_len, @@ -84,4 +95,5 @@ def meta_eval_args_from_config(cfg: RunConfig) -> CliMetaEvalArgs: result_folder=cfg.run.result_folder, engine_kwargs=judge_kwargs, no_log_file=cfg.run.no_log_file, + cache=cfg.cache.model_copy(), ) diff --git a/judgearena/meta_eval/runner.py b/judgearena/meta_eval/runner.py index 783615b..039fc75 100644 --- a/judgearena/meta_eval/runner.py +++ b/judgearena/meta_eval/runner.py @@ -3,12 +3,12 @@ from __future__ import annotations import json -from dataclasses import asdict from datetime import UTC, datetime from pathlib import Path import pandas as pd +from judgearena.config import meta_eval_cache_task, open_inference_cache from judgearena.log import attach_file_handler, get_logger, make_run_log_path from judgearena.meta_eval.annotate import annotate_sample from judgearena.meta_eval.cli_args import CliMetaEvalArgs @@ -182,7 +182,7 @@ def main(args: CliMetaEvalArgs) -> dict: attach_file_handler(make_run_log_path(res_folder)) with open(res_folder / "args.json", "w", encoding="utf-8") as handle: - json.dump(asdict(args), handle, indent=2) + json.dump(args.to_jsonable(), handle, indent=2) prompt_spec = resolve_prompt_mode( args.prompt_mode, @@ -215,38 +215,43 @@ def main(args: CliMetaEvalArgs) -> dict: len(top_models), ) - df_ann = annotate_sample( - df_sample, - args, - judge_chat_model=judge_chat_model, - prompt_spec=prompt_spec, - ) - df_ann.to_parquet(res_folder / "annotations.parquet", index=False) + cache_task = meta_eval_cache_task(args.reference_arena) + with open_inference_cache(args.cache, cache_task) as cache: + df_ann = annotate_sample( + df_sample, + args, + judge_chat_model=judge_chat_model, + prompt_spec=prompt_spec, + cache=cache, + ) + df_ann.to_parquet(res_folder / "annotations.parquet", index=False) - results = _compute_results( - args=args, - top_models=top_models, - df_top=df_top, - df_sample=df_sample, - df_ann=df_ann, - ) + results = _compute_results( + args=args, + top_models=top_models, + df_top=df_top, + df_sample=df_sample, + df_ann=df_ann, + ) - with open(res_folder / "results.json", "w", encoding="utf-8") as handle: - json.dump(_to_jsonable(results), handle, indent=2, allow_nan=False) + with open(res_folder / "results.json", "w", encoding="utf-8") as handle: + json.dump(_to_jsonable(results), handle, indent=2, allow_nan=False) - summary_csv = _build_summary_csv(results["language_summary"]) - summary_csv.to_csv(res_folder / "summary.csv", index=False) + summary_csv = _build_summary_csv(results["language_summary"]) + summary_csv.to_csv(res_folder / "summary.csv", index=False) - write_run_metadata( - output_dir=res_folder, - entrypoint="judgearena.meta_eval.runner", - run=asdict(args), - results=results, - input_payloads={"question_id": df_sample["question_id"].astype(str).tolist()}, - judge_system_prompt=prompt_spec.system_prompt, - judge_user_prompt_template=prompt_spec.user_prompt_template, - started_at_utc=started_at, - ) + write_run_metadata( + output_dir=res_folder, + entrypoint="judgearena.meta_eval.runner", + run=args.to_jsonable(), + results=results, + input_payloads={ + "question_id": df_sample["question_id"].astype(str).tolist() + }, + judge_system_prompt=prompt_spec.system_prompt, + judge_user_prompt_template=prompt_spec.user_prompt_template, + started_at_utc=started_at, + ) logger.info("Meta-eval results saved to %s", res_folder) return results diff --git a/judgearena/mt_bench/fastchat_compat.py b/judgearena/mt_bench/fastchat_compat.py index eb60af7..9b950c2 100644 --- a/judgearena/mt_bench/fastchat_compat.py +++ b/judgearena/mt_bench/fastchat_compat.py @@ -4,7 +4,7 @@ import math from dataclasses import dataclass -from typing import Any, Literal +from typing import TYPE_CHECKING, Any, Literal import pandas as pd @@ -25,6 +25,9 @@ from judgearena.prompts.registry import DEFAULT_JUDGE_PROMPT_PRESET from judgearena.utils import strip_thinking_tags +if TYPE_CHECKING: + from judgearena.inference_cache import InferenceCache + FASTCHAT_TEMPERATURE_CONFIG: dict[str, float] = { "writing": 0.7, "roleplay": 0.7, @@ -347,6 +350,7 @@ def judge_mt_bench_pairwise_fastchat( use_tqdm: bool, prompt_preset: str = DEFAULT_JUDGE_PROMPT_PRESET, strip_thinking_before_judging: bool = False, + cache: InferenceCache | None = None, ) -> tuple[pd.Series, list[dict[str, Any]], list[dict[str, object]], int]: """Run FastChat-style MT-Bench pairwise judging with bracketed verdict outputs.""" assert swap_mode in ("fixed", "both") @@ -368,6 +372,7 @@ def judge_mt_bench_pairwise_fastchat( items=items, use_tqdm=use_tqdm, swap_answers=False, + cache=cache, ) g2_judgments: list[str] | None = None @@ -377,6 +382,7 @@ def judge_mt_bench_pairwise_fastchat( items=items, use_tqdm=use_tqdm, swap_answers=True, + cache=cache, ) annotations: list[dict[str, Any]] = [] diff --git a/judgearena/mt_bench/mt_bench_utils.py b/judgearena/mt_bench/mt_bench_utils.py index 318a155..1e62875 100644 --- a/judgearena/mt_bench/mt_bench_utils.py +++ b/judgearena/mt_bench/mt_bench_utils.py @@ -13,6 +13,7 @@ import pandas as pd +from judgearena.config import dump_config from judgearena.generate import generate_multiturn from judgearena.instruction_dataset import load_instructions from judgearena.instruction_dataset.mt_bench import ( @@ -32,17 +33,14 @@ resolve_run_judge_prompt, ) from judgearena.repro import write_run_metadata -from judgearena.utils import ( - cache_function_dataframe, - compute_pref_summary, - generation_cache_token, -) +from judgearena.utils import compute_pref_summary from judgearena.utils.eval import BattleReport, _compute_grouped_stats logger = get_logger(__name__) if TYPE_CHECKING: from judgearena.config import RunConfig + from judgearena.inference_cache import InferenceCache def _align_mt_bench_completions( @@ -87,29 +85,8 @@ def _build_mt_bench_generation_kwargs( def _generate_mt_bench_completions( cfg: RunConfig, questions_df: pd.DataFrame, - ignore_cache: bool, + cache: InferenceCache | None = None, ) -> tuple[pd.DataFrame, pd.DataFrame]: - cache_prefix = "mt-bench" - - def _run_generation( - model_name: str, *, generation_kwargs: dict[str, object] - ) -> pd.DataFrame: - # MT-Bench's category-aware temperatures only kick in when the user has - # not explicitly pinned a per-role temperature; otherwise the config - # override should win for reproducibility. - temperature_config = ( - None if "temperature" in generation_kwargs else FASTCHAT_TEMPERATURE_CONFIG - ) - return generate_multiturn( - questions=questions_df, - model=model_name, - truncate_input_chars=cfg.generation.truncate_all_input_chars, - use_tqdm=cfg.run.use_tqdm, - temperature_config=temperature_config, - strip_thinking_before_turn_2_prompt=cfg.judge.strip_thinking_before_judging, - **generation_kwargs, - ) - def _load_or_generate(model_name: str, *, role: str) -> pd.DataFrame: loaded_answers = load_mt_bench_model_answers( model_name, n_instructions=cfg.generation.n_instructions @@ -120,19 +97,24 @@ 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=ignore_cache, - cache_name=( - f"{cache_prefix}_{model_name}_{cfg.generation.n_instructions}_" - f"{sampling_token}" - ), + # MT-Bench's category-aware temperatures only kick in when the user has + # not explicitly pinned a per-role temperature; otherwise the config + # override should win for reproducibility. + temperature_config = ( + None if "temperature" in generation_kwargs else FASTCHAT_TEMPERATURE_CONFIG + ) + generated_answers = generate_multiturn( + questions=questions_df, + model=model_name, + truncate_input_chars=cfg.generation.truncate_all_input_chars, + use_tqdm=cfg.run.use_tqdm, + temperature_config=temperature_config, + strip_thinking_before_turn_2_prompt=cfg.judge.strip_thinking_before_judging, + cache=cache, + **generation_kwargs, ) return _align_mt_bench_completions( questions_df=questions_df, @@ -177,8 +159,6 @@ def _save_mt_bench_results( """Persist MT-Bench arguments, annotations, aggregate results, and metadata.""" res_folder.mkdir(parents=True, exist_ok=True) - from judgearena.config import dump_config - dump_config(cfg, res_folder / "config.yaml") annotations_df.to_csv(res_folder / f"{result_name}-annotations.csv", index=False) @@ -262,6 +242,7 @@ def _run_mt_bench_fastchat( resolved_prompt: ResolvedJudgePrompt, fastchat_prompt_preset: str, started_at_utc: datetime, + cache: InferenceCache | None = None, ) -> pd.Series: prefs, annotations, combined_metadata, num_inconsistent = ( judge_mt_bench_pairwise_fastchat( @@ -278,6 +259,7 @@ def _run_mt_bench_fastchat( use_tqdm=cfg.run.use_tqdm, prompt_preset=fastchat_prompt_preset, strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging, + cache=cache, ) ) return _finalize_mt_bench_run( @@ -307,6 +289,7 @@ def _run_mt_bench_preset( judge_chat_model, resolved_prompt: ResolvedJudgePrompt, started_at_utc: datetime, + cache: InferenceCache | None = None, ) -> pd.Series: prefs, annotations, combined_metadata = judge_mt_bench_with_preset( judge_chat_model=judge_chat_model, @@ -325,6 +308,7 @@ def _run_mt_bench_preset( system_file=cfg.judge.system_prompt_file, user_file=cfg.judge.user_prompt_file, strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging, + cache=cache, ) return _finalize_mt_bench_run( cfg=cfg, @@ -343,7 +327,7 @@ def _run_mt_bench_preset( def run_mt_bench( cfg: RunConfig, - ignore_cache: bool, + cache: InferenceCache | None = None, *, res_folder: Path, result_name: str, @@ -368,7 +352,7 @@ def run_mt_bench( completions_a, completions_b = _generate_mt_bench_completions( cfg=cfg, questions_df=questions_df, - ignore_cache=ignore_cache, + cache=cache, ) resolved_prompt = resolve_run_judge_prompt(cfg.task, cfg.judge, multi_turn=True) if resolved_prompt.delegated and not cfg.judge.provide_explanation: @@ -395,6 +379,7 @@ def run_mt_bench( resolved_prompt=resolved_prompt, fastchat_prompt_preset=DEFAULT_JUDGE_PROMPT_PRESET, started_at_utc=run_started_at, + cache=cache, ) return _run_mt_bench_preset( cfg=cfg, @@ -406,4 +391,5 @@ def run_mt_bench( judge_chat_model=judge_chat_model, resolved_prompt=resolved_prompt, started_at_utc=run_started_at, + cache=cache, ) diff --git a/judgearena/mt_bench/pairwise_judging.py b/judgearena/mt_bench/pairwise_judging.py index 1acc110..ef4afa9 100644 --- a/judgearena/mt_bench/pairwise_judging.py +++ b/judgearena/mt_bench/pairwise_judging.py @@ -2,7 +2,7 @@ from collections.abc import Callable from dataclasses import dataclass -from typing import Protocol +from typing import TYPE_CHECKING, Any, Protocol import pandas as pd from langchain_core.prompts import ChatPromptTemplate @@ -11,6 +11,9 @@ from judgearena.mt_bench.common import iter_mt_bench_pairwise_rows from judgearena.utils import strip_thinking_tags +if TYPE_CHECKING: + from judgearena.inference_cache import InferenceCache + class MTBenchPairwisePrompt(Protocol): name: str @@ -70,12 +73,27 @@ def build_pairwise_chat_prompt_template( return ChatPromptTemplate.from_messages(message_templates) +def _mt_bench_judge_metadata( + item: MTBenchJudgeItem, + *, + swap_answers: bool, +) -> dict[str, Any]: + return { + "question_id": str(item.question_id), + "category": item.category, + "turn": item.turn, + "prompt": item.prompt_name, + "orientation": "reversed" if swap_answers else "direct", + } + + def infer_pairwise_judgments_by_prompt_groups( *, judge_chat_model, items: list[MTBenchJudgeItem], use_tqdm: bool, swap_answers: bool, + cache: InferenceCache | None = None, ) -> tuple[list[str], list[dict[str, str]]]: judgments: list[str] = [""] * len(items) used_prompt_kwargs: list[dict[str, str]] = [{} for _ in items] @@ -83,6 +101,7 @@ def infer_pairwise_judgments_by_prompt_groups( prompt = items[idxs[0]].prompt prompt_template = build_pairwise_chat_prompt_template(prompt) batch_kwargs: list[dict[str, str]] = [] + batch_metadata: list[dict[str, Any]] = [] for item_index in idxs: prompt_kwargs = dict(items[item_index].prompt_kwargs) if swap_answers: @@ -91,11 +110,16 @@ def infer_pairwise_judgments_by_prompt_groups( multi_turn=prompt.multi_turn, ) batch_kwargs.append(prompt_kwargs) + batch_metadata.append( + _mt_bench_judge_metadata(items[item_index], swap_answers=swap_answers) + ) prompt_inputs = prompt_template.batch(batch_kwargs) outputs = do_inference( chat_model=judge_chat_model, inputs=prompt_inputs, use_tqdm=use_tqdm, + cache=cache, + cache_meta={"metadata": batch_metadata}, ) for item_index, output, prompt_kwargs in zip( idxs, outputs, batch_kwargs, strict=True diff --git a/judgearena/mt_bench/preset_judging.py b/judgearena/mt_bench/preset_judging.py index e7ec5d2..478f5e4 100644 --- a/judgearena/mt_bench/preset_judging.py +++ b/judgearena/mt_bench/preset_judging.py @@ -2,7 +2,7 @@ import math from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any import pandas as pd @@ -23,6 +23,9 @@ resolve_judge_prompt, ) +if TYPE_CHECKING: + from judgearena.inference_cache import InferenceCache + @dataclass(frozen=True) class MTBenchPresetPrompt: @@ -154,6 +157,7 @@ def judge_mt_bench_with_preset( system_file: str | None = None, user_file: str | None = None, strip_thinking_before_judging: bool = False, + cache: InferenceCache | None = None, ) -> tuple[pd.Series, list[dict[str, Any]], list[dict[str, object]]]: assert swap_mode in ("fixed", "both") eval_single, eval_multi = resolve_mt_bench_turn_flags(turns_mode) @@ -176,6 +180,7 @@ def judge_mt_bench_with_preset( items=items, use_tqdm=use_tqdm, swap_answers=False, + cache=cache, ) annotations: list[dict[str, Any]] = [] @@ -238,6 +243,7 @@ def _append_results( items=items, use_tqdm=use_tqdm, swap_answers=True, + cache=cache, ) ) _append_results(swapped_judgments, swapped_prompt_kwargs, swapped=True) diff --git a/judgearena/pairwise_baselines.py b/judgearena/pairwise_baselines.py new file mode 100644 index 0000000..2a192ea --- /dev/null +++ b/judgearena/pairwise_baselines.py @@ -0,0 +1,34 @@ +"""Dataset-native pairwise baseline registry for generate-and-evaluate tasks.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from judgearena.instruction_dataset.arena_hard import ARENA_HARD_BASELINES +from judgearena.instruction_dataset.m_arenahard import ( + M_ARENA_HARD_BASELINES, + split_m_arena_hard_dataset, +) +from judgearena.instruction_dataset.mt_bench import MT_BENCH_BASELINES + +ALPACA_EVAL_BASELINES: dict[str, str] = { + "alpaca-eval": "gpt4_1106_preview", +} + +PAIRWISE_BASELINES: dict[str, str | Mapping[str, str]] = { + **ALPACA_EVAL_BASELINES, + **ARENA_HARD_BASELINES, + **M_ARENA_HARD_BASELINES, + **MT_BENCH_BASELINES, +} + + +def native_pairwise_baseline(task: str) -> str | Mapping[str, str] | None: + """Return the dataset-native pairwise baseline, if the task defines one.""" + if task in PAIRWISE_BASELINES: + return PAIRWISE_BASELINES[task] + parsed_m_arena_hard = split_m_arena_hard_dataset(task) + if parsed_m_arena_hard is not None: + version_key, _lang_or_subset = parsed_m_arena_hard + return PAIRWISE_BASELINES[version_key] + return None diff --git a/judgearena/utils/__init__.py b/judgearena/utils/__init__.py index 797e1b4..9a8758f 100644 --- a/judgearena/utils/__init__.py +++ b/judgearena/utils/__init__.py @@ -14,11 +14,9 @@ ) from judgearena.utils.io import ( Timeblock, - cache_function_dataframe, data_root, download_all, download_hf, - generation_cache_token, read_df, safe_parse_int, ) @@ -34,12 +32,10 @@ "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 dd90cdc..b757210 100644 --- a/judgearena/utils/io.py +++ b/judgearena/utils/io.py @@ -2,10 +2,8 @@ from __future__ import annotations -import hashlib import os import time -from collections.abc import Callable from pathlib import Path import pandas as pd @@ -16,6 +14,8 @@ download_arena_hard, is_arena_hard_dataset, ) +from judgearena.instruction_dataset.m_arenahard import M_ARENA_HARD_BASELINES +from judgearena.instruction_dataset.mt_bench import download_mt_bench from judgearena.log import get_logger logger = get_logger(__name__) @@ -70,8 +70,6 @@ def safe_parse_int(env_var: str) -> int | None: def download_all(): - from judgearena.instruction_dataset.m_arenahard import M_ARENA_HARD_BASELINES - logger.info("Downloading all datasets in %s", data_root) local_path_tables = data_root / "tables" for dataset in ( @@ -95,8 +93,6 @@ def download_all(): force_download=False, ) - from judgearena.instruction_dataset.mt_bench import download_mt_bench - download_mt_bench() @@ -125,79 +121,5 @@ def __str__(self): 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/scripts/fluency/generate_fluency.py b/scripts/fluency/generate_fluency.py index 5564dfa..a13b5e4 100644 --- a/scripts/fluency/generate_fluency.py +++ b/scripts/fluency/generate_fluency.py @@ -2,18 +2,8 @@ import pandas as pd from datasets import Dataset -from langchain_community.cache import SQLiteCache -from langchain_core.globals import set_llm_cache from judgearena.models import do_inference, make_model -from judgearena.utils import data_root - - -def set_langchain_cache(): - set_llm_cache(SQLiteCache(database_path=str(data_root / ".langchain.db"))) - - -set_langchain_cache() dataset_name = "geoalgo/multilingual-fluency" model = "OpenRouter/openai/gpt-5-mini" @@ -90,14 +80,14 @@ def generate_contexts( model: str, languages: list[str], n_sentences_to_generate: int, - ignore_cache: bool = False, + overwrite_existing: bool = False, ): for target_language in languages: data_path = Path(__file__).parent / "data" / f"{target_language}-contexts.csv" data_path.parent.mkdir(parents=True, exist_ok=True) - if not data_path.exists() or ignore_cache: + if not data_path.exists() or overwrite_existing: judge_chat_model = make_model(model, max_tokens=65536) print( diff --git a/scripts/multilingual_arena_hard/translate_arena_hard.py b/scripts/multilingual_arena_hard/translate_arena_hard.py index d0849c4..87cd9ef 100644 --- a/scripts/multilingual_arena_hard/translate_arena_hard.py +++ b/scripts/multilingual_arena_hard/translate_arena_hard.py @@ -20,8 +20,6 @@ from judgearena.instruction_dataset import load_instructions from judgearena.utils import do_inference, make_model -# set_langchain_cache() - dataset_name = "openeurollm/ArenaHard-EU-v0-bis" """ diff --git a/slurmpilot_scripts/launch_evaluation.py b/slurmpilot_scripts/launch_evaluation.py index c067fe6..5b2f7bb 100644 --- a/slurmpilot_scripts/launch_evaluation.py +++ b/slurmpilot_scripts/launch_evaluation.py @@ -17,7 +17,6 @@ "method_B": "gpt4_1106_preview", "judge_model": "VLLM/meta-llama/Meta-Llama-3-8B-instruct", "n_instructions": 10, - # "ignore_cache": None, }, src_dir=str(Path(__file__).parent.parent / "judgearena/"), n_cpus=1, diff --git a/slurmpilot_scripts/launch_generation_and_evaluation.py b/slurmpilot_scripts/launch_generation_and_evaluation.py index 6668a33..98816d5 100644 --- a/slurmpilot_scripts/launch_generation_and_evaluation.py +++ b/slurmpilot_scripts/launch_generation_and_evaluation.py @@ -75,7 +75,6 @@ "model_B": model, "judge_model": "VLLM/Qwen/Qwen2.5-32B-Instruct-GPTQ-Int8", "n_instructions": 100, - # "ignore_cache": None, } for model in multisynt_models + qwen_models ], diff --git a/tests/test_cli.py b/tests/test_cli.py index e5c69b8..44a4d79 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -301,3 +301,105 @@ def test_judge_side_kwargs_parsed_separately(capture_mains): assert cfg.judge.max_model_len == 65536 assert cfg.model.engine_kwargs == {"tensor_parallel_size": 1} assert cfg.judge.engine_kwargs == {"tensor_parallel_size": 4} + + +def test_cache_nested_cli_flags(capture_mains, tmp_path, monkeypatch): + monkeypatch.setattr("getpass.getuser", lambda: "cli-user") + store = tmp_path / "store" + cli_module.cli( + [ + "--task", + "alpaca-eval", + "--model.name", + "Dummy/A", + "--model.baseline", + "Dummy/B", + "--judge.model", + "Dummy/J", + "--cache.store_root", + str(store), + "--cache.cache_mode", + "refresh", + "--cache.cache_fetch", + "--cache.cache_push", + "--cache.pushed_by", + "cli-user", + ] + ) + cfg = capture_mains["cfg"] + assert cfg.cache.store_root == str(store) + assert cfg.cache.cache_mode == "refresh" + assert cfg.cache.cache_fetch is True + assert cfg.cache.cache_push is True + assert cfg.cache.pushed_by == "cli-user" + + +def test_cache_flat_cli_aliases(capture_mains, tmp_path): + store = tmp_path / "store" + cli_module.cli( + [ + "--task", + "alpaca-eval", + "--model.name", + "Dummy/A", + "--model.baseline", + "Dummy/B", + "--judge.model", + "Dummy/J", + "--store_root", + str(store), + "--cache_mode", + "use", + "--cache_hf_repo", + "org/repo", + "--cache_fetch", + "--cache_create_pr", + "--cache_push", + ] + ) + cfg = capture_mains["cfg"] + assert cfg.cache.store_root == str(store) + assert cfg.cache.cache_mode == "use" + assert cfg.cache.cache_hf_repo == "org/repo" + assert cfg.cache.cache_fetch is True + assert cfg.cache.cache_push is True + assert cfg.cache.cache_create_pr is True + + +def test_cache_cli_overrides_yaml_fetch(tmp_path, capture_mains): + yaml_path = tmp_path / "run.yaml" + yaml_path.write_text( + "task: alpaca-eval\n" + "model: {name: Dummy/A, baseline: Dummy/B}\n" + "judge: {model: Dummy/J}\n" + "cache:\n" + " store_root: /yaml/store\n" + " cache_fetch: true\n" + ) + cli_module.cli( + [ + "--config_path", + str(yaml_path), + "--no-cache_fetch", + ] + ) + cfg = capture_mains["cfg"] + assert cfg.cache.store_root == "/yaml/store" + assert cfg.cache.cache_fetch is False + + +def test_cache_fetch_without_store_root_errors(capture_mains): + with pytest.raises(SystemExit): + cli_module.cli( + [ + "--task", + "alpaca-eval", + "--model.name", + "Dummy/A", + "--model.baseline", + "Dummy/B", + "--judge.model", + "Dummy/J", + "--cache_fetch", + ] + ) diff --git a/tests/test_config.py b/tests/test_config.py index 7a9bca5..c307354 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -254,3 +254,109 @@ def test_build_run_config_elo_arena_derived(): ] ) assert cfg.elo is not None and cfg.elo.arena == "ComparIA" + + +def test_cache_defaults(): + cfg = RunConfig(**_base_generate()) + assert cfg.cache.store_root is None + assert cfg.cache.cache_mode == "use" + assert cfg.cache.cache_hf_repo == "judge-arena/judge-arena-cache" + assert cfg.cache.cache_fetch is False + assert cfg.cache.cache_push is False + assert cfg.cache.cache_create_pr is False + + +def test_cache_pushed_by_defaults_to_getuser(monkeypatch): + monkeypatch.setattr("judgearena.config.default_pushed_by", lambda: "test-user") + cfg = RunConfig(**_base_generate()) + assert cfg.cache.pushed_by == "test-user" + + +def test_cache_yaml_load(tmp_path): + from judgearena.config import load_config + + yaml_path = tmp_path / "cache.yaml" + yaml_path.write_text( + "task: alpaca-eval\n" + "model: {name: Dummy/a, baseline: Dummy/b}\n" + "judge: {model: Dummy/j}\n" + "cache:\n" + " store_root: /data/cache\n" + " cache_mode: refresh\n" + " cache_fetch: true\n" + " cache_push: true\n" + " pushed_by: yaml-user\n" + ) + cfg = load_config(yaml_path) + assert cfg.cache.store_root == "/data/cache" + assert cfg.cache.cache_mode == "refresh" + assert cfg.cache.cache_fetch is True + assert cfg.cache.cache_push is True + assert cfg.cache.pushed_by == "yaml-user" + + +@pytest.mark.parametrize( + ("kwargs", "match"), + [ + ( + {"cache_fetch": True}, + "cache.store_root is required", + ), + ( + {"store_root": "/tmp", "cache_fetch": True, "cache_hf_repo": " "}, + "cache_hf_repo must be non-empty", + ), + ( + {"store_root": "/tmp", "cache_create_pr": True}, + "cache_push is required", + ), + ( + {"store_root": "/tmp", "cache_mode": "off", "cache_fetch": True}, + "cache_fetch and cache_push cannot be enabled", + ), + ( + {"cache_mode": "refresh"}, + "cache.store_root is required when cache_mode is refresh", + ), + ( + {"store_root": " "}, + "cache.store_root must be non-empty", + ), + ], +) +def test_cache_validation_rejects_invalid_combinations(kwargs, match): + data = _base_generate() + data["cache"] = kwargs + with pytest.raises(ValidationError, match=match): + RunConfig(**data) + + +def test_inference_cache_session_yields_none_without_store_root(): + from judgearena.config import inference_cache_session + + cfg = RunConfig(**_base_generate()) + with inference_cache_session(cfg) as cache: + assert cache is None + + +def test_inference_cache_session_opens_cache(tmp_path, monkeypatch): + from judgearena.config import inference_cache_session, inference_cache_task + + monkeypatch.setattr("getpass.getuser", lambda: "session-user") + data = _base_generate() + data["cache"] = { + "store_root": str(tmp_path / "store"), + "cache_mode": "refresh", + "cache_fetch": True, + "cache_push": True, + "pushed_by": "session-user", + } + cfg = RunConfig(**data) + with inference_cache_session(cfg) as cache: + assert cache is not None + assert cache.store_root == tmp_path / "store" + assert cache.task == inference_cache_task(cfg) + assert cache.mode == "refresh" + assert cache.fetch is True + assert cache.push is True + assert cache.pushed_by == "session-user" diff --git a/tests/test_estimate_elo_cache_threading.py b/tests/test_estimate_elo_cache_threading.py new file mode 100644 index 0000000..7b6aaed --- /dev/null +++ b/tests/test_estimate_elo_cache_threading.py @@ -0,0 +1,330 @@ +from __future__ import annotations + +import numpy as np +import pandas as pd +import pytest + +import judgearena.estimate_elo_ratings as estimate_elo_ratings +import judgearena.evaluate as evaluate_module +import judgearena.generate as generate_module +import judgearena.models as models_module +from judgearena.config import RunConfig +from judgearena.estimate_elo_ratings import main + + +def _make_conversation(content_user: str, content_assistant: str) -> list[dict]: + return [ + {"role": "user", "content": content_user}, + {"role": "assistant", "content": content_assistant}, + ] + + +@pytest.fixture +def synthetic_arena_df() -> pd.DataFrame: + rng = np.random.default_rng(42) + rows = [] + for i in range(30): + ma, mb = rng.choice( + ["arena_model_alpha", "arena_model_beta", "arena_model_gamma"], + size=2, + replace=False, + ) + rows.append( + { + "question_id": f"q{i}", + "tstamp": 1700000000 + i, + "model_a": ma, + "model_b": mb, + "winner": rng.choice(["model_a", "model_b", "tie"]), + "conversation_a": _make_conversation( + f"Instruction {i}", f"Response A {i}" + ), + "conversation_b": _make_conversation( + f"Instruction {i}", f"Response B {i}" + ), + "benchmark": "TestArena", + "lang": rng.choice(["en", "fr"]), + } + ) + return pd.DataFrame(rows) + + +@pytest.fixture(autouse=True) +def mock_elo_deps(monkeypatch, synthetic_arena_df): + monkeypatch.setattr( + estimate_elo_ratings, + "load_arena_dataframe", + lambda arena: synthetic_arena_df, + ) + + +def _cfg_with_cache(tmp_path, **overrides) -> RunConfig: + payload = { + "task": "elo-comparia", + "model": {"name": "Dummy/my model"}, + "judge": {"model": "Dummy/score A: 0 score B: 10", "swap_mode": "fixed"}, + "generation": {"n_instructions": 5}, + "elo": {"arena": "ComparIA", "n_bootstraps": 2}, + "run": {"result_folder": str(tmp_path / "results"), "no_log_file": True}, + "cache": {"store_root": str(tmp_path / "cache")}, + } + payload.update(overrides) + return RunConfig(**payload) + + +def test_elo_uses_one_shared_cache_handle(monkeypatch, tmp_path): + captured: list[tuple[str, object]] = [] + real_gen = generate_module.do_inference + real_eval = evaluate_module.do_inference + + def spy_gen(*args, **kwargs): + cache = kwargs.get("cache") + if cache is not None: + captured.append(("gen", cache)) + return real_gen(*args, **kwargs) + + def spy_eval(*args, **kwargs): + cache = kwargs.get("cache") + if cache is not None: + captured.append(("eval", cache)) + return real_eval(*args, **kwargs) + + monkeypatch.setattr(generate_module, "do_inference", spy_gen) + monkeypatch.setattr(evaluate_module, "do_inference", spy_eval) + + main(_cfg_with_cache(tmp_path)) + + assert captured + assert len({id(cache) for _, cache in captured}) == 1 + assert any(role == "gen" for role, _ in captured) + assert any(role == "eval" for role, _ in captured) + + +def test_elo_second_run_reuses_cached_rows(monkeypatch, tmp_path): + uncached_calls = {"count": 0} + real_uncached = models_module._do_inference_uncached + + def counting_uncached(*args, **kwargs): + uncached_calls["count"] += 1 + return real_uncached(*args, **kwargs) + + monkeypatch.setattr(models_module, "_do_inference_uncached", counting_uncached) + + cfg = _cfg_with_cache(tmp_path) + first = main(cfg) + assert uncached_calls["count"] > 0 + + uncached_calls["count"] = 0 + second = main(cfg) + assert uncached_calls["count"] == 0 + assert second["winrate"] == pytest.approx(first["winrate"]) + + +def test_elo_judge_row_metadata_includes_arena_and_battle_fields(monkeypatch, tmp_path): + captured_metadata: list[dict] = [] + real_eval = evaluate_module.do_inference + + def spy_eval(*args, **kwargs): + cache_meta = kwargs.get("cache_meta") + if cache_meta is not None: + captured_metadata.extend(cache_meta.get("metadata", [])) + return real_eval(*args, **kwargs) + + monkeypatch.setattr(evaluate_module, "do_inference", spy_eval) + + main( + _cfg_with_cache( + tmp_path, + model={"name": "Dummy/focal-model"}, + ) + ) + + assert captured_metadata + first = captured_metadata[0] + assert first["arena"] == "ComparIA" + assert first["source"] == "elo-judge" + assert first["focal_model"] == "Dummy/focal-model" + assert first["opponent_model"] + assert first["position"] in {"A", "B"} + assert first["question_id"] == "q0" + assert first["orientation"] == "direct" + + +def _configure_calibration_arena(monkeypatch, synthetic_arena_df): + frames = [] + for block in range(20): + chunk = synthetic_arena_df.copy() + chunk["question_id"] = [f"q{block * len(chunk) + j}" for j in range(len(chunk))] + chunk.index = chunk.index + block * len(chunk) + frames.append(chunk) + large_arena_df = pd.concat(frames) + anchor_battles = pd.DataFrame( + { + "model_a": ["arena_model_alpha"] * len(large_arena_df), + "model_b": ["arena_model_beta"] * len(large_arena_df), + "winner": ["model_a", "model_b"] * (len(large_arena_df) // 2), + "pref": [0.0, 1.0] * (len(large_arena_df) // 2), + "pref_hard": [0.0, 1.0] * (len(large_arena_df) // 2), + "source": ["human"] * len(large_arena_df), + "question_id": large_arena_df["question_id"].tolist(), + }, + index=large_arena_df.index, + ) + monkeypatch.setattr( + estimate_elo_ratings, + "arena_anchor_battles", + lambda _df: anchor_battles, + ) + monkeypatch.setattr( + estimate_elo_ratings, + "load_arena_dataframe", + lambda arena: large_arena_df, + ) + return large_arena_df + + +def test_elo_calibration_reuses_shared_cache(monkeypatch, tmp_path, synthetic_arena_df): + _configure_calibration_arena(monkeypatch, synthetic_arena_df) + + captured: list[tuple[str, object]] = [] + real_eval = evaluate_module.do_inference + + def spy_eval(*args, **kwargs): + cache = kwargs.get("cache") + cache_meta = kwargs.get("cache_meta") + if cache is not None: + captured.append(("eval", cache)) + if cache_meta is not None: + for row in cache_meta.get("metadata", []): + if row.get("purpose") == "temperature_calibration": + captured.append(("cal_meta", row)) + return real_eval(*args, **kwargs) + + monkeypatch.setattr(evaluate_module, "do_inference", spy_eval) + + main( + _cfg_with_cache( + tmp_path, + elo={ + "arena": "ComparIA", + "n_bootstraps": 2, + "calibrate_temperature": True, + "calibration_size": 3, + }, + ) + ) + + eval_caches = [cache for role, cache in captured if role == "eval"] + assert eval_caches + assert len({id(cache) for cache in eval_caches}) == 1 + cal_rows = [row for role, row in captured if role == "cal_meta"] + assert cal_rows + assert cal_rows[0]["source"] == "elo-calibration" + assert cal_rows[0]["question_id"] + + +def test_elo_cache_hit_reparses_scores_with_recomputed_calibration( + monkeypatch, tmp_path, synthetic_arena_df +): + _configure_calibration_arena(monkeypatch, synthetic_arena_df) + uncached_calls = {"count": 0} + real_uncached = models_module._do_inference_uncached + + def counting_uncached(*args, **kwargs): + uncached_calls["count"] += 1 + return real_uncached(*args, **kwargs) + + monkeypatch.setattr(models_module, "_do_inference_uncached", counting_uncached) + cfg = _cfg_with_cache( + tmp_path, + judge={"model": "Dummy/score A: 2 score B: 1", "swap_mode": "fixed"}, + elo={ + "arena": "ComparIA", + "n_bootstraps": 1, + "calibrate_temperature": True, + "calibration_size": 20, + }, + ) + + monkeypatch.setattr(estimate_elo_ratings, "calibrate_temperature", lambda *_: 0.5) + first = main(cfg) + assert uncached_calls["count"] > 0 + + uncached_calls["count"] = 0 + monkeypatch.setattr(estimate_elo_ratings, "calibrate_temperature", lambda *_: 5.0) + second = main(cfg) + + assert uncached_calls["count"] == 0 + assert second["elo_mean"] != pytest.approx(first["elo_mean"]) + + +def test_elo_judge_metadata_fallback_without_question_id(monkeypatch, tmp_path): + arena_no_qid = pd.DataFrame( + { + "tstamp": [1700000000], + "model_a": ["arena_model_alpha"], + "model_b": ["arena_model_beta"], + "winner": ["model_a"], + "conversation_a": [ + [ + {"role": "user", "content": "Instruction 0"}, + {"role": "assistant", "content": "Response A 0"}, + ] + ], + "conversation_b": [ + [ + {"role": "user", "content": "Instruction 0"}, + {"role": "assistant", "content": "Response B 0"}, + ] + ], + "benchmark": "TestArena", + "lang": ["en"], + } + ) + monkeypatch.setattr( + estimate_elo_ratings, + "load_arena_dataframe", + lambda arena: arena_no_qid, + ) + + captured_metadata: list[dict] = [] + real_eval = evaluate_module.do_inference + + def spy_eval(*args, **kwargs): + cache_meta = kwargs.get("cache_meta") + if cache_meta is not None: + captured_metadata.extend(cache_meta.get("metadata", [])) + return real_eval(*args, **kwargs) + + monkeypatch.setattr(evaluate_module, "do_inference", spy_eval) + + main( + _cfg_with_cache( + tmp_path, + generation={"n_instructions": 1}, + elo={"arena": "ComparIA", "n_bootstraps": 1}, + ) + ) + + judge_rows = [row for row in captured_metadata if row.get("source") == "elo-judge"] + assert judge_rows + assert "question_id" not in judge_rows[0] + assert judge_rows[0]["battle_identity"] + + +def test_elo_swap_mode_both_doubles_llm_judged_battles(tmp_path): + result = main( + RunConfig( + task="elo-comparia", + model={"name": "Dummy/my model"}, + judge={ + "model": "Dummy/score A: 0 score B: 10", + "swap_mode": "both", + }, + generation={"n_instructions": 4}, + elo={"arena": "ComparIA", "n_bootstraps": 1}, + run={"result_folder": str(tmp_path), "no_log_file": True}, + ) + ) + assert result["llm_judged_battles"] == 8 + assert result["num_battles"] == 4 diff --git a/tests/test_estimate_elo_ratings.py b/tests/test_estimate_elo_ratings.py index 42b3258..38cf8dc 100644 --- a/tests/test_estimate_elo_ratings.py +++ b/tests/test_estimate_elo_ratings.py @@ -74,13 +74,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: arena = kwargs.pop("arena", "ComparIA") @@ -256,9 +249,13 @@ def spy_judge( completions_A, completions_B, swap_mode="fixed", + cache=None, + row_metadata=None, **kwargs, ): captured["swap_mode"] = swap_mode + captured["cache"] = cache + captured["row_metadata"] = row_metadata n = len(instructions) dummy = JudgeAnnotation( judge_completion="score A: 0 score B: 10", @@ -266,6 +263,8 @@ def spy_judge( completion_A="", completion_B="", ) + if swap_mode == "both": + return [dummy] * n, [dummy] * n, pd.Series([1.0] * (2 * n)) return [dummy] * n, None, pd.Series([1.0] * n) monkeypatch.setattr(estimate_elo_ratings, "judge_and_parse_prefs", spy_judge) @@ -281,9 +280,13 @@ def spy_judge( completions_B, swap_mode="fixed", strip_thinking_before_judging=False, + cache=None, + row_metadata=None, **kwargs, ): captured["strip_thinking_before_judging"] = strip_thinking_before_judging + captured["cache"] = cache + captured["row_metadata"] = row_metadata n = len(instructions) dummy = JudgeAnnotation( judge_completion="score A: 0 score B: 10", @@ -320,8 +323,9 @@ def test_main_strip_thinking_defaults_off(monkeypatch, tmp_path): def _spy_generate_capturing(captured): - def spy_generate(instructions, model, **kwargs): + def spy_generate(instructions, model, cache=None, **kwargs): captured["gen_kwargs"] = kwargs + captured["cache"] = cache return pd.DataFrame( { "completion": [f"c{i}" for i in range(len(instructions))], @@ -332,6 +336,24 @@ def spy_generate(instructions, model, **kwargs): return spy_generate +def test_main_generation_cache_metadata_includes_arena_question_identity( + monkeypatch, tmp_path +): + captured = {} + monkeypatch.setattr( + estimate_elo_ratings, "generate_instructions", _spy_generate_capturing(captured) + ) + + main(_default_args(result_folder=str(tmp_path), arena="ComparIA")) + + row_metadata = captured["gen_kwargs"]["row_metadata"] + assert len(row_metadata) == 10 + assert all(row["arena"] == "ComparIA" for row in row_metadata) + assert [row["question_id"] for row in row_metadata] == [ + f"q{index}" for index in range(10) + ] + + def test_main_thinking_budget_injected_for_thinking_model(monkeypatch, tmp_path): """battle_thinking_token_budget must reach generation for VLLM thinking models.""" captured = {} diff --git a/tests/test_evaluate_cache_threading.py b/tests/test_evaluate_cache_threading.py new file mode 100644 index 0000000..e94ccf9 --- /dev/null +++ b/tests/test_evaluate_cache_threading.py @@ -0,0 +1,107 @@ +from __future__ import annotations + +import math + +import judgearena.evaluate as evaluate_module +from judgearena.evaluate import annotate_battles, judge_and_parse_prefs +from judgearena.inference_cache import InferenceCache +from judgearena.models import make_model + + +class FakeJudge: + def __init__(self, response: str = "score A: 0 score B: 10"): + self.response = response + + def batch(self, *, inputs, **_kwargs): + return [self.response] * len(inputs) + + +def test_annotate_battles_forwards_cache_and_metadata(monkeypatch): + captured: list[dict] = [] + + def spy_do_inference(*, cache, cache_meta, **kwargs): + captured.append({"cache": cache, "cache_meta": cache_meta}) + return ["score A: 0 score B: 10"] + + monkeypatch.setattr(evaluate_module, "do_inference", spy_do_inference) + + row_metadata = [{"battle_id": "b-1"}] + with InferenceCache("/tmp/unused", "judge", mode="off") as cache: + annotate_battles( + judge_chat_model=FakeJudge(), + instructions=["Question"], + completions_A=["A"], + completions_B=["B"], + cache=cache, + row_metadata=row_metadata, + ) + + assert captured[0]["cache"] is cache + assert captured[0]["cache_meta"] == {"metadata": row_metadata} + + +def test_judge_and_parse_prefs_adds_orientation_and_forwards_both(monkeypatch): + captured: list[dict] = [] + + def spy_do_inference(*, cache_meta, **kwargs): + captured.append(cache_meta) + return ["score A: 0 score B: 10"] * len(kwargs["inputs"]) + + monkeypatch.setattr(evaluate_module, "do_inference", spy_do_inference) + + base_metadata = [{"question_id": "q-42"}] + + _, annotations_reversed, prefs = judge_and_parse_prefs( + judge_chat_model=FakeJudge(), + instructions=["Q1", "Q2"], + completions_A=["A1", "A2"], + completions_B=["B1", "B2"], + swap_mode="both", + row_metadata=base_metadata * 2, + ) + + assert annotations_reversed is not None + assert len(captured) == 2 + assert captured[0]["metadata"] == [ + {"question_id": "q-42", "orientation": "direct"}, + {"question_id": "q-42", "orientation": "direct"}, + ] + assert captured[1]["metadata"] == [ + {"question_id": "q-42", "orientation": "reversed"}, + {"question_id": "q-42", "orientation": "reversed"}, + ] + assert len(prefs) == 4 + + +def test_judge_and_parse_prefs_default_without_cache_unchanged(): + judge = make_model("Dummy/score A: 0 score B: 10") + _, annotations_reversed, prefs = judge_and_parse_prefs( + judge_chat_model=judge, + instructions=["Q"], + completions_A=["A"], + completions_B=["B"], + swap_mode="fixed", + ) + + assert annotations_reversed is None + assert len(prefs) == 1 + assert not math.isnan(float(prefs.iloc[0])) + + +def test_annotate_battles_without_metadata_omits_cache_meta(monkeypatch): + captured: list[dict] = [] + + def spy_do_inference(*, cache_meta=None, **kwargs): + captured.append({"cache_meta": cache_meta}) + return ["score A: 0 score B: 10"] + + monkeypatch.setattr(evaluate_module, "do_inference", spy_do_inference) + + annotate_battles( + judge_chat_model=FakeJudge(), + instructions=["Question"], + completions_A=["A"], + completions_B=["B"], + ) + + assert captured[0]["cache_meta"] is None diff --git a/tests/test_generate_and_evaluate.py b/tests/test_generate_and_evaluate.py index e03d2c9..2aca333 100644 --- a/tests/test_generate_and_evaluate.py +++ b/tests/test_generate_and_evaluate.py @@ -78,13 +78,6 @@ def mock_external_data_and_cache(monkeypatch): lambda dataset, model, n_instructions: None, ) - def _run_without_cache(fun, **_kwargs): - return fun() - - monkeypatch.setattr( - generate_and_evaluate, "cache_function_dataframe", _run_without_cache - ) - def _instructions(ids: list[str], categories: list[str] | None = None) -> pd.DataFrame: data = {"instruction": list(ids)} diff --git a/tests/test_generate_and_evaluate_cache_threading.py b/tests/test_generate_and_evaluate_cache_threading.py new file mode 100644 index 0000000..95bda01 --- /dev/null +++ b/tests/test_generate_and_evaluate_cache_threading.py @@ -0,0 +1,215 @@ +from __future__ import annotations + +import pandas as pd +import pytest + +import judgearena.evaluate as evaluate_module +import judgearena.generate as generate_module +import judgearena.generate_and_evaluate as gae +import judgearena.models as models_module +import judgearena.mt_bench.mt_bench_utils as mt_bench_utils +import judgearena.mt_bench.pairwise_judging as mt_pairwise +from judgearena.config import RunConfig +from judgearena.generate_and_evaluate import main as main_generate_and_eval + + +def _synthetic_instructions(n: int = 20) -> pd.DataFrame: + return pd.DataFrame( + {"instruction": [f"Synthetic instruction {i}" for i in range(n)]}, + index=pd.Index(range(n), name="instruction_index"), + ) + + +def _cfg_with_cache(tmp_path, **overrides) -> RunConfig: + payload = { + "task": "alpaca-eval", + "model": {"name": "Dummy/gen-a", "baseline": "Dummy/gen-b"}, + "judge": {"model": "Dummy/score A: 0 score B: 10", "swap_mode": "fixed"}, + "generation": {"n_instructions": 2}, + "run": {"result_folder": str(tmp_path / "results"), "no_log_file": True}, + "cache": {"store_root": str(tmp_path / "cache")}, + } + payload.update(overrides) + return RunConfig(**payload) + + +@pytest.fixture +def mock_gae_inputs(monkeypatch): + instructions = _synthetic_instructions() + + monkeypatch.setattr( + gae, + "load_instructions", + lambda dataset, n_instructions=None: ( + instructions.head(n_instructions) + if n_instructions is not None + else instructions + ), + ) + monkeypatch.setattr( + gae, + "try_load_dataset_completions", + lambda dataset, model, n_instructions: None, + ) + + +def test_gae_uses_one_shared_cache_handle(mock_gae_inputs, monkeypatch, tmp_path): + captured: list[tuple[str, object]] = [] + real_gen = generate_module.do_inference + real_eval = evaluate_module.do_inference + + def spy_gen(*args, **kwargs): + cache = kwargs.get("cache") + if cache is not None: + captured.append(("gen", cache)) + return real_gen(*args, **kwargs) + + def spy_eval(*args, **kwargs): + cache = kwargs.get("cache") + if cache is not None: + captured.append(("eval", cache)) + return real_eval(*args, **kwargs) + + monkeypatch.setattr(generate_module, "do_inference", spy_gen) + monkeypatch.setattr(evaluate_module, "do_inference", spy_eval) + + main_generate_and_eval(_cfg_with_cache(tmp_path)) + + assert captured + assert len({id(cache) for _, cache in captured}) == 1 + assert any(role == "gen" for role, _ in captured) + assert any(role == "eval" for role, _ in captured) + + +def test_mt_bench_uses_one_shared_cache_handle(monkeypatch, tmp_path): + questions = pd.DataFrame( + { + "category": ["writing"], + "turn_1": ["Question 1"], + "turn_2": ["Question 2"], + }, + index=pd.Index([1], name="instruction_index"), + ) + monkeypatch.setattr( + mt_bench_utils, + "load_instructions", + lambda dataset, n_instructions=None: questions, + ) + monkeypatch.setattr( + mt_bench_utils, + "load_mt_bench_model_answers", + lambda model, n_instructions=None: None, + ) + captured: list[tuple[str, object]] = [] + real_generation = generate_module.do_inference + real_judging = mt_pairwise.do_inference + + def spy_generation(*args, **kwargs): + cache = kwargs.get("cache") + if cache is not None: + captured.append(("generation", cache)) + return real_generation(*args, **kwargs) + + def spy_judging(*args, **kwargs): + cache = kwargs.get("cache") + if cache is not None: + captured.append(("judging", cache)) + return real_judging(*args, **kwargs) + + monkeypatch.setattr(generate_module, "do_inference", spy_generation) + monkeypatch.setattr(mt_pairwise, "do_inference", spy_judging) + cfg = RunConfig( + task="mt-bench", + model={"name": "Dummy/gen-a", "baseline": "Dummy/gen-b"}, + judge={"model": "Dummy/[[A]]", "swap_mode": "fixed"}, + generation={"n_instructions": 1}, + run={"result_folder": str(tmp_path / "results"), "no_log_file": True}, + cache={"store_root": str(tmp_path / "cache")}, + ) + + main_generate_and_eval(cfg) + + assert {role for role, _ in captured} == {"generation", "judging"} + assert len({id(cache) for _, cache in captured}) == 1 + + +def test_gae_second_run_reuses_cached_rows(mock_gae_inputs, monkeypatch, tmp_path): + uncached_calls = {"count": 0} + real_uncached = models_module._do_inference_uncached + + def counting_uncached(*args, **kwargs): + uncached_calls["count"] += 1 + return real_uncached(*args, **kwargs) + + monkeypatch.setattr(models_module, "_do_inference_uncached", counting_uncached) + + cfg = _cfg_with_cache(tmp_path) + prefs_first = main_generate_and_eval(cfg) + assert uncached_calls["count"] > 0 + + uncached_calls["count"] = 0 + prefs_second = main_generate_and_eval(cfg) + assert uncached_calls["count"] == 0 + assert prefs_second.tolist() == prefs_first.tolist() + + +def test_gae_preloaded_completions_bypass_generation( + mock_gae_inputs, monkeypatch, tmp_path +): + preloaded = pd.DataFrame( + { + "completion": ["preloaded-a", "preloaded-b"], + "instruction_index": [0, 1], + } + ) + generation_calls: list[str] = [] + real_gen = models_module._do_inference_uncached + + def track_generation(chat_model, inputs, **kwargs): + model_spec = getattr(chat_model, "model_spec", None) or getattr( + chat_model, "name", "unknown" + ) + generation_calls.append(str(model_spec)) + return real_gen(chat_model, inputs, **kwargs) + + def load_preloaded(dataset, model, n_instructions): + if model == "Dummy/gen-a": + return preloaded + return None + + monkeypatch.setattr(gae, "try_load_dataset_completions", load_preloaded) + monkeypatch.setattr(models_module, "_do_inference_uncached", track_generation) + + main_generate_and_eval(_cfg_with_cache(tmp_path)) + + assert not any("gen-a" in call for call in generation_calls) + assert any("gen-b" in call for call in generation_calls) + + +def test_gae_judge_row_metadata_includes_models_and_instruction_index( + mock_gae_inputs, monkeypatch, tmp_path +): + captured_metadata: list[dict] = [] + real_eval = evaluate_module.do_inference + + def spy_eval(*args, **kwargs): + cache_meta = kwargs.get("cache_meta") + if cache_meta is not None: + captured_metadata.extend(cache_meta.get("metadata", [])) + return real_eval(*args, **kwargs) + + monkeypatch.setattr(evaluate_module, "do_inference", spy_eval) + + main_generate_and_eval( + _cfg_with_cache( + tmp_path, + model={"name": "Dummy/gen-a", "baseline": "Dummy/gen-b"}, + ) + ) + + assert captured_metadata + first = captured_metadata[0] + assert first["instruction_index"] == "0" + assert first["model_A"] == "Dummy/gen-a" + assert first["model_B"] == "Dummy/gen-b" + assert first["orientation"] == "direct" diff --git a/tests/test_generate_cache_threading.py b/tests/test_generate_cache_threading.py new file mode 100644 index 0000000..03bd897 --- /dev/null +++ b/tests/test_generate_cache_threading.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +import json + +import pandas as pd + +import judgearena.generate as generate_module +from judgearena.generate import generate_base, generate_multiturn +from judgearena.inference_cache import InferenceCache +from judgearena.models import make_model +from judgearena.store_sqlite import SQLiteInferenceStore, descriptor_hash, store_folder + + +def test_generate_base_routes_through_do_inference(monkeypatch): + calls: list[dict] = [] + real_do_inference = generate_module.do_inference + + def spy_do_inference(*args, **kwargs): + calls.append({"args": args, "kwargs": kwargs}) + return real_do_inference(*args, **kwargs) + + monkeypatch.setattr(generate_module, "do_inference", spy_do_inference) + + instructions = pd.Series(["hello", "world"], index=[10, 20]) + df = generate_base(instructions, "Dummy/generate-base-path", use_tqdm=False) + + assert len(calls) == 1 + assert calls[0]["kwargs"]["use_tqdm"] is False + assert calls[0]["kwargs"]["cache"] is None + assert df["completion"].tolist() == ["generate-base-path"] * 2 + assert df["instruction_index"].tolist() == [10, 20] + + +def test_generate_base_forwards_cache_and_metadata(monkeypatch): + captured: list[dict] = [] + + def spy_do_inference(*, cache, cache_meta, **kwargs): + captured.append({"cache": cache, "cache_meta": cache_meta}) + return ["out-a", "out-b"] + + monkeypatch.setattr(generate_module, "do_inference", spy_do_inference) + + instructions = pd.Series(["a", "b"], index=["q1", "q2"]) + with InferenceCache("/tmp/unused", "gen-task", mode="off") as cache: + df = generate_base( + instructions, + "Dummy/ignored", + cache=cache, + ) + + assert df["completion"].tolist() == ["out-a", "out-b"] + assert captured[0]["cache"] is cache + assert captured[0]["cache_meta"] == { + "metadata": [ + {"instruction_index": "q1"}, + {"instruction_index": "q2"}, + ] + } + + +def test_generate_base_cache_hit_skips_backend_batch(tmp_path): + model = make_model("Dummy/cache-generate-base", max_tokens=8) + inputs = ["alpha", "beta"] + descriptor = model.cache_descriptor() + assert descriptor is not None + canonical = [model.canonicalize_input(item) for item in inputs] + metadata = [{"instruction_index": "0"}, {"instruction_index": "1"}] + + with InferenceCache(tmp_path, "gen", mode="refresh") as cache: + cache.get_or_run( + model_spec=model.model_spec, + descriptor=descriptor, + canonical_inputs=canonical, + original_inputs=inputs, + miss_runner=lambda miss_inputs: [f"cached-{item}" for item in miss_inputs], + row_metadata=metadata, + producer_metadata=model.producer_metadata(), + ) + + instructions = pd.Series(inputs, index=[0, 1]) + with InferenceCache(tmp_path, "gen", mode="use") as cache: + df = generate_base( + instructions, + "Dummy/cache-generate-base", + max_tokens=8, + cache=cache, + ) + + assert df["completion"].tolist() == ["cached-alpha", "cached-beta"] + + +def test_generate_multiturn_metadata_and_temperature_groups(monkeypatch): + calls: list[dict] = [] + + def spy_do_inference(*, inputs, cache_meta=None, **kwargs): + calls.append({"inputs": inputs, "cache_meta": cache_meta}) + return [f"out-{index}" for index in range(len(inputs))] + + monkeypatch.setattr(generate_module, "do_inference", spy_do_inference) + + questions = pd.DataFrame( + { + "category": ["writing", "math", "writing"], + "turn_1": ["Q1", "Q2", "Q3"], + "turn_2": ["Q1b", "Q2b", "Q3b"], + }, + index=pd.Index([1, 2, 3], name="instruction_index"), + ) + temperature_config = {"writing": 0.5, "math": 0.9} + + df = generate_multiturn( + questions, + "Dummy/multiturn", + temperature_config=temperature_config, + use_tqdm=False, + ) + + assert len(df) == 3 + assert len(calls) == 4 + + turn1_calls = calls[:2] + turn2_calls = calls[2:] + + assert turn1_calls[0]["cache_meta"]["metadata"] == [ + {"instruction_index": "1", "turn": 1, "category": "writing"}, + {"instruction_index": "3", "turn": 1, "category": "writing"}, + ] + assert turn1_calls[1]["cache_meta"]["metadata"] == [ + {"instruction_index": "2", "turn": 1, "category": "math"}, + ] + assert len(turn1_calls[0]["inputs"]) == 2 + assert len(turn1_calls[1]["inputs"]) == 1 + + assert turn2_calls[0]["cache_meta"]["metadata"] == [ + {"instruction_index": "1", "turn": 2, "category": "writing"}, + {"instruction_index": "3", "turn": 2, "category": "writing"}, + ] + assert turn2_calls[1]["cache_meta"]["metadata"] == [ + {"instruction_index": "2", "turn": 2, "category": "math"}, + ] + + +def test_generate_multiturn_saves_metadata_in_cache(tmp_path): + questions = pd.DataFrame( + { + "category": ["writing"], + "turn_1": ["Q1"], + "turn_2": ["Q2"], + }, + index=pd.Index([7], name="instruction_index"), + ) + + with InferenceCache(tmp_path, "mt-bench", mode="refresh") as cache: + generate_multiturn( + questions, + "Dummy/mt-meta", + use_tqdm=False, + cache=cache, + ) + + model = make_model("Dummy/mt-meta", max_tokens=8192) + descriptor = model.cache_descriptor() + folder = store_folder( + tmp_path, + "mt-bench", + model.model_spec, + descriptor_hash(descriptor), + ) + with SQLiteInferenceStore(folder / "inference.db") as store: + rows = store.query_metadata() + + saved = [json.loads(value) for value in rows["metadata_json"]] + assert {"instruction_index": "7", "turn": 1, "category": "writing"} in saved + assert {"instruction_index": "7", "turn": 2, "category": "writing"} in saved diff --git a/tests/test_logging.py b/tests/test_logging.py index e4dfd20..867c058 100644 --- a/tests/test_logging.py +++ b/tests/test_logging.py @@ -2,7 +2,9 @@ from __future__ import annotations +import io import logging +import sys import pytest @@ -76,6 +78,28 @@ def test_configure_logging_no_duplicate_handlers(): assert len(console_handlers) == 1 +def test_configure_logging_rebinds_console_to_current_stderr(monkeypatch): + first_stream = io.StringIO() + second_stream = io.StringIO() + monkeypatch.setattr(sys, "stderr", first_stream) + configure_logging(0) + monkeypatch.setattr(sys, "stderr", second_stream) + configure_logging(0) + + get_logger("judgearena.test_rebind").error("new stream") + + assert "new stream" not in first_stream.getvalue() + assert "new stream" in second_stream.getvalue() + + +def test_configure_logging_adds_console_when_only_file_handler_exists(tmp_path): + attach_file_handler(tmp_path / "run.log") + + configure_logging(0) + + assert _console_handler_level() == logging.INFO + + def test_env_var_overrides_verbosity(monkeypatch): """JUDGEARENA_LOG_LEVEL env-var should override the CLI verbosity flag.""" monkeypatch.setenv("JUDGEARENA_LOG_LEVEL", "warning") diff --git a/tests/test_meta_eval.py b/tests/test_meta_eval.py index a287964..ebb4226 100644 --- a/tests/test_meta_eval.py +++ b/tests/test_meta_eval.py @@ -14,8 +14,9 @@ import judgearena.meta_eval.sampling as meta_sampling from judgearena import cli as cli_module from judgearena.arenas_utils import extract_turn_text +from judgearena.config import meta_eval_cache_task from judgearena.evaluate import JudgeAnnotation, PairScore, annotate_battles -from judgearena.meta_eval.cache import AnnotationCache, AnnotationEntry, AnnotationKey +from judgearena.inference_cache import InferenceCache from judgearena.meta_eval.cli_args import CliMetaEvalArgs from judgearena.meta_eval.metrics import ( compute_agreement_metrics, @@ -357,75 +358,7 @@ def test_agreement_metrics_on_fixture(): assert metrics["n_nt"] == 3 -def _cache_key(**overrides) -> AnnotationKey: - values = { - "benchmark": "LMArena-140k", - "instruction_id": "q-1", - "model_a": "model-a", - "model_b": "model-b", - "judge": "Dummy/judge", - } - values.update(overrides) - return AnnotationKey(**values) - - -def _cache_entry(completion: str = "score_A: 9\nscore_B: 1", **overrides): - key = _cache_key(**overrides) - return AnnotationEntry( - **key.__dict__, - judge_input="judge prompt", - judge_completion=completion, - ) - - -def test_annotation_cache_persists_and_preserves_batch_order(tmp_path): - db_dir = tmp_path / "db" - first = AnnotationCache(db_dir) - first.batch_put( - [ - _cache_entry(instruction_id="q-2", completion="second"), - _cache_entry(instruction_id="q-1", completion="first"), - ] - ) - first.close() - - second = AnnotationCache(db_dir) - entries = second.batch_get_annotations( - [ - _cache_key(instruction_id="q-1"), - _cache_key(instruction_id="q-2"), - _cache_key(instruction_id="missing"), - ] - ) - assert [entry.judge_completion if entry else None for entry in entries] == [ - "first", - "second", - None, - ] - second.close() - - -def test_annotation_cache_distinguishes_prompt_mode_and_model_order(tmp_path): - cache = AnnotationCache(tmp_path / "db") - cache.batch_put( - [ - _cache_entry(judge="Dummy/judge::arena-hard"), - _cache_entry(model_a="model-b", model_b="model-a"), - ] - ) - entries = cache.batch_get_annotations( - [ - _cache_key(judge="Dummy/judge"), - _cache_key(judge="Dummy/judge::arena-hard"), - _cache_key(model_a="model-b", model_b="model-a"), - ] - ) - assert entries[0] is None - assert all(entry is not None for entry in entries[1:]) - cache.close() - - -def test_annotate_sample_uses_cache_and_inverts_swapped_pass( +def test_annotate_sample_inverts_swapped_pass( monkeypatch, synthetic_arena_df, meta_args, @@ -440,20 +373,24 @@ def fake_annotate_battles(**kwargs): monkeypatch.setattr(meta_annotate, "annotate_battles", fake_annotate_battles) meta_args.swap_mode = "both" sample = synthetic_arena_df.iloc[:1] - cache = AnnotationCache(tmp_path / "db") prompt_spec = PromptModeSpec( name="standard", system_prompt="system", user_prompt_template="user", ) - annotations = meta_annotate.annotate_sample( - sample, - meta_args, - judge_chat_model=object(), - prompt_spec=prompt_spec, - annotation_cache=cache, - ) + with InferenceCache( + tmp_path / "cache", + meta_eval_cache_task(meta_args.reference_arena), + mode="off", + ) as cache: + annotations = meta_annotate.annotate_sample( + sample, + meta_args, + judge_chat_model=object(), + prompt_spec=prompt_spec, + cache=cache, + ) assert len(annotations) == 2 assert annotations["orientation"].tolist() == ["forward", "swapped"] assert annotations["winner"].tolist() == [sample.iloc[0]["winner"]] * 2 @@ -470,26 +407,6 @@ def fake_annotate_battles(**kwargs): ) assert calls["count"] == 2 - meta_annotate.annotate_sample( - sample, - meta_args, - judge_chat_model=object(), - prompt_spec=prompt_spec, - annotation_cache=cache, - ) - assert calls["count"] == 2 - - meta_args.ignore_cache = True - meta_annotate.annotate_sample( - sample, - meta_args, - judge_chat_model=object(), - prompt_spec=prompt_spec, - annotation_cache=cache, - ) - assert calls["count"] == 4 - cache.close() - def test_cost_uses_offline_reference_pricing(monkeypatch, tmp_path): pricing_file = tmp_path / "openrouter_pricing.json" @@ -592,11 +509,8 @@ def test_swap_mode_both_artifact_reproduces_overall_agreement( stub_meta_eval_runner, tmp_path, ): - cache_class = AnnotationCache - monkeypatch.setattr( - meta_annotate, - "AnnotationCache", - lambda: cache_class(tmp_path / "cache"), + meta_args.cache = meta_args.cache.model_copy( + update={"store_root": str(tmp_path / "cache")} ) monkeypatch.setattr(meta_annotate, "annotate_battles", _judge_annotations) meta_args.swap_mode = "both" diff --git a/tests/test_meta_eval_cache_threading.py b/tests/test_meta_eval_cache_threading.py new file mode 100644 index 0000000..42a00de --- /dev/null +++ b/tests/test_meta_eval_cache_threading.py @@ -0,0 +1,581 @@ +"""Unified inference cache integration tests for meta-evaluation.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pandas as pd +import pytest + +import judgearena.evaluate as evaluate_module +import judgearena.meta_eval.annotate as meta_annotate +import judgearena.meta_eval.runner as meta_eval_runner +import judgearena.models as models_module +from judgearena.config import ( + CacheArgs, + RunConfig, + inference_cache_task, + meta_eval_cache_task, +) +from judgearena.inference_cache import InferenceCache +from judgearena.meta_eval.cli_args import CliMetaEvalArgs, meta_eval_args_from_config +from judgearena.meta_eval.prompts import PromptModeSpec +from judgearena.meta_eval.runner import main as meta_eval_main +from judgearena.models import make_model +from judgearena.store_sqlite import SQLiteInferenceStore, descriptor_hash, store_folder + + +def _meta_args_with_cache(tmp_path: Path, **overrides) -> CliMetaEvalArgs: + values = { + "reference_arena": "LMArena-140k", + "prompt_mode": "standard", + "top_models": 3, + "battles_per_model": 1, + "batch_size": 8, + "languages": ["en"], + "judge_model": "Dummy/score_A: 9\nscore_B: 1", + "result_folder": str(tmp_path / "results"), + "no_log_file": True, + "cache": CacheArgs(store_root=str(tmp_path / "cache")), + } + values.update(overrides) + return CliMetaEvalArgs(**values) + + +def _prompt_spec() -> PromptModeSpec: + return PromptModeSpec( + name="standard", + system_prompt="system", + user_prompt_template=( + "Question: {user_prompt}\nA: {completion_A}\nB: {completion_B}" + ), + ) + + +def _single_battle_frame() -> pd.DataFrame: + conv_a = [ + {"role": "user", "content": "Question 0"}, + {"role": "assistant", "content": "Answer A 0"}, + ] + conv_b = [ + {"role": "user", "content": "Question 0"}, + {"role": "assistant", "content": "Answer B 0"}, + ] + return pd.DataFrame( + [ + { + "question_id": "q-0", + "model_a": "model-0", + "model_b": "model-1", + "winner": "model_a", + "lang": "en", + "benchmark": "LMArena-140k", + "conversation_a": conv_a, + "conversation_b": conv_b, + } + ] + ) + + +def _base_meta_eval_payload(tmp_path: Path, **overrides) -> dict: + payload = { + "task": "meta-eval", + "judge": {"model": "Dummy/j"}, + "run": {"result_folder": str(tmp_path / "results"), "no_log_file": True}, + } + payload.update(overrides) + return payload + + +def test_meta_eval_args_from_config_carries_cache(tmp_path): + cfg = RunConfig( + **_base_meta_eval_payload( + tmp_path, + cache={"store_root": str(tmp_path / "cache"), "cache_mode": "refresh"}, + ) + ) + args = meta_eval_args_from_config(cfg) + assert args.cache.store_root == str(tmp_path / "cache") + assert args.cache.cache_mode == "refresh" + + +def test_meta_eval_args_from_config_carries_strip_thinking(tmp_path): + cfg = RunConfig( + **_base_meta_eval_payload( + tmp_path, + judge={ + "model": "Dummy/j", + "strip_thinking_before_judging": True, + }, + ) + ) + + args = meta_eval_args_from_config(cfg) + + assert args.strip_thinking_before_judging is True + + +def test_inference_cache_task_includes_reference_arena(tmp_path): + cfg = RunConfig(**_base_meta_eval_payload(tmp_path)) + assert inference_cache_task(cfg) == "meta-eval-LMArena-140k" + assert meta_eval_cache_task("LMArena-140k") == "meta-eval-LMArena-140k" + + +def test_meta_eval_forwards_strip_thinking_to_annotation(monkeypatch, tmp_path): + captured: dict[str, object] = {} + + def spy_annotate_battles(*, instructions, **kwargs): + captured.update(kwargs) + return [ + evaluate_module.JudgeAnnotation( + judge_completion="score A: 9 score B: 1", + instruction=instruction, + completion_A="A", + completion_B="B", + judge_input="rendered", + ) + for instruction in instructions + ] + + monkeypatch.setattr(meta_annotate, "annotate_battles", spy_annotate_battles) + args = _meta_args_with_cache(tmp_path, strip_thinking_before_judging=True) + + meta_annotate.annotate_sample( + _single_battle_frame(), + args, + judge_chat_model=object(), + prompt_spec=_prompt_spec(), + ) + + assert captured["strip_thinking_before_judging"] is True + + +def test_meta_eval_second_run_reuses_cached_rows(monkeypatch, tmp_path): + uncached_calls = {"count": 0} + real_uncached = models_module._do_inference_uncached + + def counting_uncached(*args, **kwargs): + uncached_calls["count"] += 1 + return real_uncached(*args, **kwargs) + + monkeypatch.setattr(models_module, "_do_inference_uncached", counting_uncached) + + args = _meta_args_with_cache(tmp_path, swap_mode="fixed") + sample = _single_battle_frame() + prompt_spec = _prompt_spec() + judge = make_model(args.judge_model, max_tokens=32) + + with InferenceCache( + args.cache.store_root, + meta_eval_cache_task(args.reference_arena), + mode="use", + ) as cache: + meta_annotate.annotate_sample( + sample, + args, + judge_chat_model=judge, + prompt_spec=prompt_spec, + cache=cache, + ) + assert uncached_calls["count"] > 0 + + uncached_calls["count"] = 0 + with InferenceCache( + args.cache.store_root, + meta_eval_cache_task(args.reference_arena), + mode="use", + ) as cache: + meta_annotate.annotate_sample( + sample, + args, + judge_chat_model=judge, + prompt_spec=prompt_spec, + cache=cache, + ) + assert uncached_calls["count"] == 0 + + +def test_meta_eval_swapped_orientations_store_distinct_associations( + monkeypatch, tmp_path +): + captured_metadata: list[dict] = [] + real_do_inference = evaluate_module.do_inference + + def spy_do_inference(*args, **kwargs): + cache_meta = kwargs.get("cache_meta") + if cache_meta is not None: + captured_metadata.extend(cache_meta.get("metadata", [])) + return real_do_inference(*args, **kwargs) + + monkeypatch.setattr(evaluate_module, "do_inference", spy_do_inference) + + args = _meta_args_with_cache(tmp_path, swap_mode="both") + prompt_spec = _prompt_spec() + judge = make_model(args.judge_model, max_tokens=32) + + with InferenceCache( + args.cache.store_root, + meta_eval_cache_task(args.reference_arena), + mode="use", + ) as cache: + meta_annotate.annotate_sample( + _single_battle_frame(), + args, + judge_chat_model=judge, + prompt_spec=prompt_spec, + cache=cache, + ) + + orientations = {row["orientation"] for row in captured_metadata} + assert orientations == {"forward", "swapped"} + assert all(row["question_id"] == "q-0" for row in captured_metadata) + assert all(row["reference_arena"] == "LMArena-140k" for row in captured_metadata) + + model = make_model(args.judge_model, max_tokens=32) + descriptor = model.cache_descriptor() + folder = store_folder( + tmp_path / "cache", + meta_eval_cache_task(args.reference_arena), + model.model_spec, + descriptor_hash(descriptor), + ) + with SQLiteInferenceStore(folder / "inference.db") as store: + metadata_rows = store.query_metadata() + stored_orientations = { + json.loads(row["metadata_json"])["orientation"] + for _, row in metadata_rows.iterrows() + } + assert stored_orientations == {"forward", "swapped"} + + +def test_meta_eval_changed_rendered_input_invalidates_only_that_row( + monkeypatch, tmp_path +): + uncached_calls = {"count": 0} + real_uncached = models_module._do_inference_uncached + + def counting_uncached(*args, **kwargs): + uncached_calls["count"] += 1 + return real_uncached(*args, **kwargs) + + monkeypatch.setattr(models_module, "_do_inference_uncached", counting_uncached) + + args = _meta_args_with_cache(tmp_path) + prompt_spec = _prompt_spec() + judge = make_model(args.judge_model, max_tokens=32) + original = _single_battle_frame() + changed = original.copy() + changed.iloc[0, changed.columns.get_loc("conversation_a")] = [ + {"role": "user", "content": "Changed question"}, + {"role": "assistant", "content": "Answer A 0"}, + ] + + with InferenceCache( + args.cache.store_root, + meta_eval_cache_task(args.reference_arena), + mode="use", + ) as cache: + meta_annotate.annotate_sample( + original, + args, + judge_chat_model=judge, + prompt_spec=prompt_spec, + cache=cache, + ) + first_count = uncached_calls["count"] + assert first_count > 0 + + uncached_calls["count"] = 0 + with InferenceCache( + args.cache.store_root, + meta_eval_cache_task(args.reference_arena), + mode="use", + ) as cache: + meta_annotate.annotate_sample( + changed, + args, + judge_chat_model=judge, + prompt_spec=prompt_spec, + cache=cache, + ) + assert uncached_calls["count"] == first_count + + +def test_meta_eval_parsing_and_costs_recompute_from_cached_output( + monkeypatch, tmp_path +): + args = _meta_args_with_cache(tmp_path) + prompt_spec = PromptModeSpec( + name="standard", + system_prompt="system", + user_prompt_template="user", + ) + judge = make_model("Dummy/score_A: 1\nscore_B: 9", max_tokens=32) + sample = _single_battle_frame() + + with InferenceCache( + args.cache.store_root, + meta_eval_cache_task(args.reference_arena), + mode="use", + ) as cache: + first = meta_annotate.annotate_sample( + sample, + args, + judge_chat_model=judge, + prompt_spec=prompt_spec, + cache=cache, + ) + + cost_calls = {"count": 0} + original_cost = meta_annotate.estimate_annotation_cost_usd + + def spy_cost(*, judge_input, judge_completion, judge_model): + cost_calls["count"] += 1 + return original_cost( + judge_input=judge_input, + judge_completion=judge_completion, + judge_model=judge_model, + ) + + monkeypatch.setattr(meta_annotate, "estimate_annotation_cost_usd", spy_cost) + + with InferenceCache( + args.cache.store_root, + meta_eval_cache_task(args.reference_arena), + mode="use", + ) as cache: + second = meta_annotate.annotate_sample( + sample, + args, + judge_chat_model=judge, + prompt_spec=prompt_spec, + cache=cache, + ) + + assert first.iloc[0]["winner_llm"] == "model_b" + assert second.iloc[0]["winner_llm"] == "model_b" + assert cost_calls["count"] > 0 + + +def test_meta_eval_runner_uses_one_shared_cache_handle( + monkeypatch, tmp_path, synthetic_arena_df +): + captured: list[object] = [] + + def spy_annotate_sample(df_sample, args, *, cache=None, **kwargs): + captured.append(cache) + return pd.DataFrame( + { + "question_id": df_sample["question_id"], + "model_a": df_sample["model_a"], + "model_b": df_sample["model_b"], + "winner": df_sample["winner"], + "lang": df_sample["lang"], + "benchmark": df_sample["benchmark"], + "orientation": "forward", + "instruction": "instr", + "completion_a": "A", + "completion_b": "B", + "judge_input": "prompt", + "judge_completion": "score_A: 9\nscore_B: 1", + "estimated_input_tokens": 2, + "estimated_output_tokens": 5, + "cost_usd": 0.001, + "cost_source": "estimated", + "winner_llm": df_sample["winner"], + "pref_llm": 0.0, + } + ) + + monkeypatch.setattr( + meta_eval_runner, + "load_reference_arena_battles", + lambda reference_arena, languages=None: synthetic_arena_df, + ) + monkeypatch.setattr( + meta_eval_runner, + "select_top_models", + lambda df, top_models: (["model-0", "model-1"], df), + ) + monkeypatch.setattr( + meta_eval_runner, + "sample_battles_per_model", + lambda df_top, models, battles_per_model, seed: df_top.head(1), + ) + monkeypatch.setattr(meta_eval_runner, "make_model", lambda **_kwargs: object()) + monkeypatch.setattr(meta_eval_runner, "annotate_sample", spy_annotate_sample) + + args = _meta_args_with_cache(tmp_path, top_models=2) + meta_eval_main(args) + + assert len(captured) == 1 + assert captured[0] is not None + + +def test_meta_eval_args_json_includes_cache_config( + tmp_path, monkeypatch, synthetic_arena_df +): + monkeypatch.setattr( + meta_eval_runner, + "load_reference_arena_battles", + lambda reference_arena, languages=None: synthetic_arena_df, + ) + monkeypatch.setattr( + meta_eval_runner, + "select_top_models", + lambda df, top_models: (["model-0", "model-1"], df), + ) + monkeypatch.setattr( + meta_eval_runner, + "sample_battles_per_model", + lambda df_top, models, battles_per_model, seed: df_top.head(1), + ) + monkeypatch.setattr(meta_eval_runner, "make_model", lambda **_kwargs: object()) + monkeypatch.setattr( + meta_eval_runner, + "annotate_sample", + lambda df_sample, args, **kwargs: pd.DataFrame( + { + "question_id": df_sample["question_id"], + "model_a": df_sample["model_a"], + "model_b": df_sample["model_b"], + "winner": df_sample["winner"], + "lang": df_sample["lang"], + "benchmark": df_sample["benchmark"], + "orientation": "forward", + "instruction": "instr", + "completion_a": "A", + "completion_b": "B", + "judge_input": "prompt", + "judge_completion": "score_A: 9\nscore_B: 1", + "estimated_input_tokens": 2, + "estimated_output_tokens": 5, + "cost_usd": 0.001, + "cost_source": "estimated", + "winner_llm": df_sample["winner"], + "pref_llm": 0.0, + } + ), + ) + + args = _meta_args_with_cache(tmp_path, top_models=2) + meta_eval_main(args) + output_dir = next(Path(args.result_folder).glob("meta-eval-*")) + args_payload = json.loads((output_dir / "args.json").read_text(encoding="utf-8")) + assert args_payload["cache"]["store_root"] == str(tmp_path / "cache") + assert "ignore_cache" not in args_payload + + +def test_meta_eval_args_serialization_redacts_engine_secrets(tmp_path): + args = _meta_args_with_cache(tmp_path) + args.engine_kwargs = { + "temperature": 0.2, + "api_key": "must-not-leak", + "default_headers": {"Authorization": "secret"}, + } + + payload = args.to_jsonable() + + assert payload["engine_kwargs"] == {"temperature": 0.2} + assert "must-not-leak" not in json.dumps(payload) + assert "secret" not in json.dumps(payload) + + +def _stub_meta_eval_sampling(monkeypatch, synthetic_arena_df: pd.DataFrame) -> None: + monkeypatch.setattr( + meta_eval_runner, + "load_reference_arena_battles", + lambda reference_arena, languages=None: synthetic_arena_df, + ) + monkeypatch.setattr( + meta_eval_runner, + "select_top_models", + lambda df, top_models: (["model-0", "model-1"], df), + ) + monkeypatch.setattr( + meta_eval_runner, + "sample_battles_per_model", + lambda df_top, models, battles_per_model, seed: df_top.head(1), + ) + + +def test_meta_eval_runner_creates_cache_cell_under_single_component_task( + monkeypatch, tmp_path, synthetic_arena_df +): + _stub_meta_eval_sampling(monkeypatch, synthetic_arena_df) + args = _meta_args_with_cache(tmp_path, top_models=2, swap_mode="fixed") + + meta_eval_main(args) + + task_root = ( + Path(args.cache.store_root) + / "inference" + / meta_eval_cache_task(args.reference_arena) + ) + db_files = list(task_root.rglob("inference.db")) + assert db_files, f"expected cache cell under {task_root}" + assert "meta-eval-LMArena-140k" in str(db_files[0]) + + +def test_meta_eval_runner_skips_push_when_downstream_processing_fails( + monkeypatch, tmp_path, synthetic_arena_df +): + import judgearena.inference_cache as inference_cache_mod + + push_calls: list[tuple] = [] + monkeypatch.setattr( + inference_cache_mod, + "push_cells", + lambda *args, **kwargs: push_calls.append((args, kwargs)), + ) + _stub_meta_eval_sampling(monkeypatch, synthetic_arena_df) + monkeypatch.setattr( + meta_eval_runner, + "_compute_results", + lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("downstream failed")), + ) + + args = _meta_args_with_cache( + tmp_path, + top_models=2, + cache=CacheArgs(store_root=str(tmp_path / "cache"), cache_push=True), + ) + + with pytest.raises(RuntimeError, match="downstream failed"): + meta_eval_main(args) + + assert push_calls == [] + task_root = ( + Path(args.cache.store_root) + / "inference" + / meta_eval_cache_task(args.reference_arena) + ) + assert list(task_root.rglob("inference.db")) + + +@pytest.fixture +def synthetic_arena_df() -> pd.DataFrame: + conv_a = [ + {"role": "user", "content": "Question 0"}, + {"role": "assistant", "content": "Answer A 0"}, + ] + conv_b = [ + {"role": "user", "content": "Question 0"}, + {"role": "assistant", "content": "Answer B 0"}, + ] + return pd.DataFrame( + [ + { + "question_id": "q-0", + "tstamp": 1, + "model_a": "model-0", + "model_b": "model-1", + "winner": "model_a", + "conversation_a": conv_a, + "conversation_b": conv_b, + "benchmark": "LMArena-140k", + "lang": "en", + } + ] + ) diff --git a/tests/test_mt_bench_downloads.py b/tests/test_mt_bench_downloads.py index de9282e..d1d3457 100644 --- a/tests/test_mt_bench_downloads.py +++ b/tests/test_mt_bench_downloads.py @@ -4,9 +4,11 @@ import pytest import judgearena.instruction_dataset.mt_bench as mt_bench +import judgearena.models as models_module import judgearena.mt_bench.mt_bench_utils as mt_bench_utils import judgearena.utils.io as utils_io from judgearena.config import RunConfig +from judgearena.inference_cache import InferenceCache from judgearena.prompts.registry import FASTCHAT_PAIRWISE_PROMPT_PRESET @@ -64,7 +66,7 @@ def _contexts_snapshot_stub(**_kwargs): monkeypatch.setattr(utils_io, "snapshot_download", _contexts_snapshot_stub) monkeypatch.setattr( - mt_bench, + utils_io, "download_mt_bench", lambda: calls.__setitem__("mt_bench", calls["mt_bench"] + 1), ) @@ -115,13 +117,11 @@ def test_generate_mt_bench_completions_uses_pregenerated_baseline(monkeypatch): index=pd.Index([1, 2], name="instruction_index"), ) generated_models = [] - - monkeypatch.setattr( - mt_bench_utils, "cache_function_dataframe", lambda fun, **_kwargs: fun() - ) + cache_values: list[object | None] = [] def fake_generate_multiturn(**kwargs): generated_models.append(kwargs["model"]) + cache_values.append(kwargs.get("cache")) return pd.DataFrame( { "instruction_index": [1, 2], @@ -158,13 +158,15 @@ def fake_generate_multiturn(**kwargs): generation={"n_instructions": 2}, ) - completions_a, completions_b = mt_bench_utils._generate_mt_bench_completions( - cfg=cfg, - questions_df=questions_df, - ignore_cache=False, - ) + with InferenceCache("/tmp/unused", "mt-bench", mode="off") as cache: + completions_a, completions_b = mt_bench_utils._generate_mt_bench_completions( + cfg=cfg, + questions_df=questions_df, + cache=cache, + ) assert generated_models == ["VLLM/example/model-a"] + assert cache_values == [cache] assert completions_a.loc[1, "completion_turn_1"] == "Gen A1" assert completions_b.loc[1, "completion_turn_1"] == "Base A1" assert completions_b.loc[2, "completion_turn_2"] == "Base B2" @@ -199,7 +201,6 @@ def test_generate_mt_bench_completions_reports_missing_baseline_rows(monkeypatch mt_bench_utils._generate_mt_bench_completions( cfg=cfg, questions_df=questions_df, - ignore_cache=False, ) @@ -262,7 +263,7 @@ def test_run_mt_bench_resolves_native_baseline_and_judge_controls( monkeypatch.setattr( mt_bench_utils, "_generate_mt_bench_completions", - lambda cfg, questions_df, ignore_cache: ( + lambda cfg, questions_df, cache=None: ( pd.DataFrame( {"completion_turn_1": ["A1"], "completion_turn_2": ["A2"]}, index=questions_df.index, @@ -308,7 +309,6 @@ def fake_run_mt_bench_fastchat(**kwargs): mt_bench_utils.run_mt_bench( cfg, - ignore_cache=False, res_folder=tmp_path, result_name="mt-bench-test", ) @@ -338,7 +338,7 @@ def test_run_mt_bench_defaults_to_delegated_fastchat(monkeypatch, tmp_path): monkeypatch.setattr( mt_bench_utils, "_generate_mt_bench_completions", - lambda cfg, questions_df, ignore_cache: ( + lambda cfg, questions_df, cache=None: ( pd.DataFrame( {"completion_turn_1": ["A1"], "completion_turn_2": ["A2"]}, index=questions_df.index, @@ -381,7 +381,6 @@ def fake_run_mt_bench_fastchat(**kwargs): mt_bench_utils.run_mt_bench( cfg, - ignore_cache=False, res_folder=tmp_path, result_name="mt-bench-test", ) @@ -408,7 +407,7 @@ def test_run_mt_bench_concrete_prompt_preset_uses_preset_judging(monkeypatch, tm monkeypatch.setattr( mt_bench_utils, "_generate_mt_bench_completions", - lambda cfg, questions_df, ignore_cache: ( + lambda cfg, questions_df, cache=None: ( pd.DataFrame( {"completion_turn_1": ["A1"], "completion_turn_2": ["A2"]}, index=questions_df.index, @@ -451,7 +450,6 @@ def fake_run_mt_bench_preset(**kwargs): mt_bench_utils.run_mt_bench( cfg, - ignore_cache=False, res_folder=tmp_path, result_name="mt-bench-test", ) @@ -469,9 +467,6 @@ def test_generate_mt_bench_completions_forwards_thinking_controls(monkeypatch): ) captured: dict[str, dict] = {} - monkeypatch.setattr( - mt_bench_utils, "cache_function_dataframe", lambda fun, **_kwargs: fun() - ) monkeypatch.setattr( mt_bench_utils, "load_mt_bench_model_answers", @@ -508,7 +503,6 @@ def fake_generate_multiturn(**kwargs): mt_bench_utils._generate_mt_bench_completions( cfg=cfg, questions_df=questions_df, - ignore_cache=False, ) thinking_call = captured["VLLM/Qwen/Qwen3.5-9B"] @@ -535,7 +529,7 @@ def test_run_mt_bench_forwards_strip_thinking_to_fastchat_judge(monkeypatch, tmp monkeypatch.setattr( mt_bench_utils, "_generate_mt_bench_completions", - lambda cfg, questions_df, ignore_cache: ( + lambda cfg, questions_df, cache=None: ( pd.DataFrame( {"completion_turn_1": ["A1"], "completion_turn_2": ["A2"]}, index=questions_df.index, @@ -567,9 +561,118 @@ def fake_judge(**kwargs): mt_bench_utils.run_mt_bench( cfg, - ignore_cache=False, res_folder=tmp_path, result_name="mt-bench-test", ) assert captured["judge"]["strip_thinking_before_judging"] is True + + +def test_run_mt_bench_forwards_cache_to_generation_and_judging(monkeypatch, tmp_path): + questions_df = pd.DataFrame( + {"turn_1": ["Q1"], "turn_2": ["Q1b"]}, + index=pd.Index([1], name="instruction_index"), + ) + captured: dict[str, object | None] = {} + + monkeypatch.setattr( + mt_bench_utils, + "load_instructions", + lambda dataset, n_instructions=None: questions_df, + ) + + def fake_generate(**kwargs): + captured["generation_cache"] = kwargs.get("cache") + return pd.DataFrame( + { + "instruction_index": [1], + "completion_turn_1": ["A1"], + "completion_turn_2": ["A2"], + } + ) + + monkeypatch.setattr(mt_bench_utils, "generate_multiturn", fake_generate) + monkeypatch.setattr( + mt_bench_utils, + "load_mt_bench_model_answers", + lambda model, n_instructions=None: None, + ) + monkeypatch.setattr(mt_bench_utils, "make_model", lambda **kwargs: object()) + monkeypatch.setattr( + mt_bench_utils, "_finalize_mt_bench_run", lambda **kwargs: kwargs["prefs"] + ) + + def fake_judge(**kwargs): + captured["judge_cache"] = kwargs.get("cache") + return pd.Series([0.0], dtype=float), [], [], 0 + + monkeypatch.setattr(mt_bench_utils, "judge_mt_bench_pairwise_fastchat", fake_judge) + + cfg = RunConfig( + task="mt-bench", + model={"name": "VLLM/example/model-a"}, + judge={"model": "VLLM/Judge"}, + generation={"n_instructions": 1}, + run={"result_folder": str(tmp_path)}, + ) + + with InferenceCache("/tmp/unused", "mt-bench", mode="off") as cache: + mt_bench_utils.run_mt_bench( + cfg, + cache=cache, + res_folder=tmp_path, + result_name="mt-bench-test", + ) + + assert captured["generation_cache"] is cache + assert captured["judge_cache"] is cache + + +def test_generate_mt_bench_completions_reuses_inference_cache(tmp_path, monkeypatch): + questions_df = pd.DataFrame( + { + "category": ["writing"], + "turn_1": ["Q1"], + "turn_2": ["Q2"], + }, + index=pd.Index([1], name="instruction_index"), + ) + backend_inputs: list[int] = [] + real_uncached = models_module._do_inference_uncached + + def counting_uncached(chat_model, inputs, *, use_tqdm=False): + backend_inputs.append(len(inputs)) + return real_uncached(chat_model, inputs, use_tqdm=use_tqdm) + + monkeypatch.setattr(models_module, "_do_inference_uncached", counting_uncached) + monkeypatch.setattr( + mt_bench_utils, + "load_mt_bench_model_answers", + lambda model, n_instructions=None: None, + ) + + cfg = RunConfig( + task="mt-bench", + model={"name": "Dummy/mt-cache-a", "baseline": "Dummy/mt-cache-b"}, + judge={"model": "Dummy/J"}, + generation={"n_instructions": 1}, + ) + + with InferenceCache(tmp_path, "mt-bench", mode="refresh") as cache: + first_a, _first_b = mt_bench_utils._generate_mt_bench_completions( + cfg=cfg, + questions_df=questions_df, + cache=cache, + ) + + assert sum(backend_inputs) == 4 + + with InferenceCache(tmp_path, "mt-bench", mode="use") as cache: + second_a, _second_b = mt_bench_utils._generate_mt_bench_completions( + cfg=cfg, + questions_df=questions_df, + cache=cache, + ) + + assert sum(backend_inputs) == 4 + assert first_a.loc[1, "completion_turn_1"] == second_a.loc[1, "completion_turn_1"] diff --git a/tests/test_mt_bench_fastchat_compat.py b/tests/test_mt_bench_fastchat_compat.py index 555ae1f..920fda6 100644 --- a/tests/test_mt_bench_fastchat_compat.py +++ b/tests/test_mt_bench_fastchat_compat.py @@ -3,6 +3,8 @@ import pandas as pd import pytest +import judgearena.mt_bench.fastchat_compat as fastchat_module +from judgearena.inference_cache import InferenceCache from judgearena.mt_bench.fastchat_compat import ( _conservative_winner, _map_verdict_to_winner, @@ -118,3 +120,35 @@ def test_judge_mt_bench_pairwise_fastchat_swap_mode_both_is_conservative(): assert annotations[0]["final_winner"] == "model_A" assert "B1" in annotations[0]["g2_user_prompt"] assert metadata == [{"question_id": 1, "category": "writing", "turn": 1}] + + +def test_judge_mt_bench_pairwise_fastchat_forwards_cache(monkeypatch): + captured: list[object | None] = [] + + def fake_infer(*, cache, **kwargs): + captured.append(cache) + return (["[[A]]"] * len(kwargs["items"]), [{}] * len(kwargs["items"])) + + monkeypatch.setattr( + fastchat_module, + "infer_pairwise_judgments_by_prompt_groups", + fake_infer, + ) + + with InferenceCache("/tmp/unused", "mt-judge", mode="off") as cache: + judge_mt_bench_pairwise_fastchat( + judge_chat_model=object(), + judge_model="judge", + questions=_questions_df(category="writing"), + completions_a=_completions_df("A"), + completions_b=_completions_df("B"), + model_a="model-a", + model_b="model-b", + turns_mode="single", + swap_mode="both", + truncate_input_chars=None, + use_tqdm=False, + cache=cache, + ) + + assert captured == [cache, cache] diff --git a/tests/test_mt_bench_pairwise_cache_threading.py b/tests/test_mt_bench_pairwise_cache_threading.py new file mode 100644 index 0000000..60c1fd9 --- /dev/null +++ b/tests/test_mt_bench_pairwise_cache_threading.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import judgearena.mt_bench.pairwise_judging as pairwise_module +from judgearena.inference_cache import InferenceCache +from judgearena.mt_bench.pairwise_judging import ( + MTBenchJudgeItem, + infer_pairwise_judgments_by_prompt_groups, +) + + +@dataclass(frozen=True) +class _Prompt: + name: str + system_prompt: str | None + user_prompt_template: str + multi_turn: bool + ref_based: bool = False + + +def _item( + *, + question_id: object, + category: str | None, + turn: int, + prompt: _Prompt, + prompt_kwargs: dict[str, str], +) -> MTBenchJudgeItem: + return MTBenchJudgeItem( + question_id=question_id, + category=category, + turn=turn, + prompt=prompt, + prompt_kwargs=prompt_kwargs, + ) + + +def test_infer_pairwise_judgments_metadata_order_and_orientation(monkeypatch): + captured: list[dict] = [] + + def spy_do_inference(*, cache_meta, **kwargs): + captured.append( + { + "cache_meta": cache_meta, + "input_count": len(kwargs["inputs"]), + } + ) + return ["judgment"] * len(kwargs["inputs"]) + + monkeypatch.setattr(pairwise_module, "do_inference", spy_do_inference) + + single_prompt = _Prompt( + name="default-single", + system_prompt=None, + user_prompt_template="{question} {answer_a} {answer_b}", + multi_turn=False, + ) + multi_prompt = _Prompt( + name="default-multi", + system_prompt=None, + user_prompt_template="{question_1} {answer_a_1}", + multi_turn=True, + ) + items = [ + _item( + question_id=1, + category="writing", + turn=1, + prompt=single_prompt, + prompt_kwargs={ + "question": "Q1", + "answer_a": "A1", + "answer_b": "B1", + }, + ), + _item( + question_id=2, + category="math", + turn=2, + prompt=multi_prompt, + prompt_kwargs={ + "question_1": "Q2a", + "question_2": "Q2b", + "answer_a_1": "A2a", + "answer_a_2": "A2b", + "answer_b_1": "B2a", + "answer_b_2": "B2b", + }, + ), + ] + + judgments, used_kwargs = infer_pairwise_judgments_by_prompt_groups( + judge_chat_model=object(), + items=items, + use_tqdm=False, + swap_answers=True, + ) + + assert judgments == ["judgment", "judgment"] + assert used_kwargs[0]["answer_a"] == "B1" + assert used_kwargs[0]["answer_b"] == "A1" + assert len(captured) == 2 + assert captured[0]["input_count"] == 1 + assert captured[0]["cache_meta"]["metadata"] == [ + { + "question_id": "1", + "category": "writing", + "turn": 1, + "prompt": "default-single", + "orientation": "reversed", + } + ] + assert captured[1]["cache_meta"]["metadata"] == [ + { + "question_id": "2", + "category": "math", + "turn": 2, + "prompt": "default-multi", + "orientation": "reversed", + } + ] + + +def test_infer_pairwise_judgments_forwards_cache(monkeypatch): + captured: list[object] = [] + + def spy_do_inference(*, cache, **kwargs): + captured.append(cache) + return ["judgment"] * len(kwargs["inputs"]) + + monkeypatch.setattr(pairwise_module, "do_inference", spy_do_inference) + + prompt = _Prompt( + name="default-single", + system_prompt=None, + user_prompt_template="{question}", + multi_turn=False, + ) + items = [ + _item( + question_id=9, + category="coding", + turn=1, + prompt=prompt, + prompt_kwargs={"question": "Q", "answer_a": "A", "answer_b": "B"}, + ) + ] + + with InferenceCache("/tmp/unused", "mt-judge", mode="off") as cache: + infer_pairwise_judgments_by_prompt_groups( + judge_chat_model=object(), + items=items, + use_tqdm=False, + swap_answers=False, + cache=cache, + ) + + assert captured == [cache] diff --git a/tests/test_mt_bench_preset_judging.py b/tests/test_mt_bench_preset_judging.py index 7c6a9e4..9c6b57c 100644 --- a/tests/test_mt_bench_preset_judging.py +++ b/tests/test_mt_bench_preset_judging.py @@ -3,6 +3,8 @@ import pandas as pd import pytest +import judgearena.mt_bench.preset_judging as preset_module +from judgearena.inference_cache import InferenceCache from judgearena.mt_bench.preset_judging import ( _build_mt_bench_preset_items, _select_preset_prompt, @@ -151,3 +153,47 @@ def test_judge_mt_bench_with_preset_parses_and_inverts_swapped_scores(): {"question_id": 1, "category": "writing", "turn": 1}, {"question_id": 1, "category": "writing", "turn": 1}, ] + + +def test_judge_mt_bench_with_preset_forwards_cache(monkeypatch): + captured: list[object | None] = [] + + def fake_infer(*, cache, items, swap_answers=False, **kwargs): + captured.append(cache) + used_kwargs = [dict(item.prompt_kwargs) for item in items] + if swap_answers: + used_kwargs = [ + { + **kwargs_, + "answer_a": kwargs_.get("answer_b", ""), + "answer_b": kwargs_.get("answer_a", ""), + } + for kwargs_ in used_kwargs + ] + n = len(items) + return (["score_A: 10\nscore_B: 0"] * n, used_kwargs) + + monkeypatch.setattr( + preset_module, + "infer_pairwise_judgments_by_prompt_groups", + fake_infer, + ) + + with InferenceCache("/tmp/unused", "mt-judge", mode="off") as cache: + judge_mt_bench_with_preset( + judge_chat_model=object(), + judge_model="judge", + questions=_questions_df(category="writing"), + completions_a=_completions_df("A"), + completions_b=_completions_df("B"), + model_a="model-a", + model_b="model-b", + turns_mode="single", + swap_mode="both", + truncate_input_chars=None, + use_tqdm=False, + prompt_preset="default", + cache=cache, + ) + + assert captured == [cache, cache] diff --git a/tests/test_no_legacy_runtime_cache.py b/tests/test_no_legacy_runtime_cache.py new file mode 100644 index 0000000..2f600e1 --- /dev/null +++ b/tests/test_no_legacy_runtime_cache.py @@ -0,0 +1,77 @@ +"""Regression guard against reintroducing legacy runtime caches.""" + +from __future__ import annotations + +import re +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +JUDGEARENA_ROOT = REPO_ROOT / "judgearena" + +FORBIDDEN_PATTERNS: list[tuple[str, re.Pattern[str]]] = [ + ("AnnotationCache", re.compile(r"\bAnnotationCache\b")), + ("cache_function_dataframe", re.compile(r"\bcache_function_dataframe\b")), + ("generation_cache_token", re.compile(r"\bgeneration_cache_token\b")), + ("ignore_cache", re.compile(r"\bignore_cache\b")), + ("set_langchain_cache", re.compile(r"\bset_langchain_cache\b")), + ("legacy cache/db path", re.compile(r"cache/db")), +] + +FORBIDDEN_IMPORTS = ( + "from judgearena.utils import cache_function_dataframe", + "from judgearena.utils.io import cache_function_dataframe", + "from judgearena.meta_eval.cache import", +) + +LEGACY_MIGRATION_ALLOWLIST = frozenset( + {("ignore_cache", "judgearena/cache_backfill_config.py")} +) + +INFERENCE_BATCH_INVOKE_ALLOWLIST = frozenset( + { + "judgearena/models.py", + "judgearena/model_adapters.py", + "judgearena/generate.py", + "judgearena/evaluate.py", + "judgearena/mt_bench/pairwise_judging.py", + } +) + + +def _judgearena_py_files() -> list[Path]: + return sorted(JUDGEARENA_ROOT.rglob("*.py")) + + +def test_meta_eval_cache_module_is_absent() -> None: + assert not (JUDGEARENA_ROOT / "meta_eval" / "cache.py").exists() + + +def test_judgearena_has_no_legacy_runtime_cache_symbols() -> None: + violations: list[str] = [] + for path in _judgearena_py_files(): + rel = path.relative_to(REPO_ROOT).as_posix() + text = path.read_text(encoding="utf-8") + for label, pattern in FORBIDDEN_PATTERNS: + if pattern.search(text) and (label, rel) not in LEGACY_MIGRATION_ALLOWLIST: + violations.append(f"{rel}: {label}") + for imp in FORBIDDEN_IMPORTS: + if imp in text: + violations.append(f"{rel}: import {imp}") + assert not violations, "Legacy runtime cache references found:\n" + "\n".join( + violations + ) + + +def test_model_inference_batch_invoke_is_allowlisted() -> None: + pattern = re.compile(r"\.(batch|invoke)\(") + violations: list[str] = [] + for path in _judgearena_py_files(): + rel = path.relative_to(REPO_ROOT).as_posix() + if rel in INFERENCE_BATCH_INVOKE_ALLOWLIST: + continue + if pattern.search(path.read_text(encoding="utf-8")): + violations.append(rel) + assert not violations, ( + "Direct .batch/.invoke outside allowlist (route through do_inference):\n" + + "\n".join(violations) + ) diff --git a/tests/test_seed_plumbing.py b/tests/test_seed_plumbing.py index 0e6b2cc..62b1875 100644 --- a/tests/test_seed_plumbing.py +++ b/tests/test_seed_plumbing.py @@ -4,7 +4,7 @@ from judgearena.config import RunConfig, build_run_config from judgearena.models import make_model -from judgearena.utils import generation_cache_token +from judgearena.store_sqlite import descriptor_hash def test_make_model_dummy_captures_temperature_and_seed(): @@ -93,15 +93,18 @@ 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} +def test_inference_descriptor_is_sensitive_to_sampling_params(): + base_a = make_model("Dummy/foo", max_tokens=32768, temperature=0.0, seed=1) + base_b = make_model("Dummy/foo", max_tokens=32768, temperature=0.0, seed=1) + changed_seed = make_model("Dummy/foo", max_tokens=32768, temperature=0.0, seed=2) + changed_temp = make_model("Dummy/foo", max_tokens=32768, temperature=1.0, seed=1) - 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) + desc_a = base_a.cache_descriptor() + desc_b = base_b.cache_descriptor() + assert desc_a is not None and desc_b is not None + assert descriptor_hash(desc_a) == descriptor_hash(desc_b) + assert descriptor_hash(desc_a) != descriptor_hash(changed_seed.cache_descriptor()) + assert descriptor_hash(desc_a) != descriptor_hash(changed_temp.cache_descriptor()) def test_model_args_per_role_kwargs_are_independent(): From 4833fc15ebeecd5861954143001b0071ef1a1595 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:53:17 +0200 Subject: [PATCH 11/13] feat(cache): add safe sync and backfill tooling Expose strict cache synchronization commands and conservatively reconstruct hosted judge rows from verifiable saved-run artifacts. Includes-AI-Code: true --- judgearena/cache_backfill.py | 324 ++++++++++ judgearena/cache_backfill_common.py | 59 ++ judgearena/cache_backfill_config.py | 166 ++++++ judgearena/cache_backfill_discovery.py | 340 +++++++++++ judgearena/cache_backfill_sources.py | 368 ++++++++++++ judgearena/cache_sync.py | 176 ++++++ pyproject.toml | 1 + tests/test_cache_backfill.py | 792 +++++++++++++++++++++++++ tests/test_cache_sync.py | 160 +++++ 9 files changed, 2386 insertions(+) create mode 100644 judgearena/cache_backfill.py create mode 100644 judgearena/cache_backfill_common.py create mode 100644 judgearena/cache_backfill_config.py create mode 100644 judgearena/cache_backfill_discovery.py create mode 100644 judgearena/cache_backfill_sources.py create mode 100644 judgearena/cache_sync.py create mode 100644 tests/test_cache_backfill.py create mode 100644 tests/test_cache_sync.py diff --git a/judgearena/cache_backfill.py b/judgearena/cache_backfill.py new file mode 100644 index 0000000..5271e98 --- /dev/null +++ b/judgearena/cache_backfill.py @@ -0,0 +1,324 @@ +"""Backfill hosted judge inference rows from saved run folders into the unified cache.""" + +from __future__ import annotations + +import json +import sqlite3 +import uuid +from collections import defaultdict +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import pandas as pd + +from judgearena.cache_backfill_discovery import ( + SKIP_REASON_BY_KIND, + ArtifactKind, + ClassifiedSource, + discover_sources, +) +from judgearena.cache_backfill_sources import ( + BackfillRow, + SourceExtraction, + extract_gae_rows, + extract_meta_eval_rows, + extract_mt_bench_rows, +) +from judgearena.log import get_logger +from judgearena.store_sqlite import ( + INFERENCE_DB_NAME, + SQLiteInferenceStore, + descriptor_hash, + stable_json_dumps, + store_folder, + write_store_metadata, +) + +logger = get_logger(__name__) + +BACKFILL_PUSHED_BY = "backfill" + + +@dataclass +class BackfillReport: + written: int = 0 + existing: int = 0 + rows_planned: int = 0 + skipped: dict[str, int] = field(default_factory=dict) + sources: dict[str, dict[str, int]] = field(default_factory=dict) + runs_processed: int = 0 + dry_run: bool = False + + def to_jsonable(self) -> dict[str, Any]: + return { + "written": self.written, + "existing": self.existing, + "rows_planned": self.rows_planned, + "skipped": dict(sorted(self.skipped.items())), + "sources": self.sources, + "runs_processed": self.runs_processed, + "dry_run": self.dry_run, + } + + +def _merge_skip_counts(target: dict[str, int], delta: dict[str, int]) -> None: + for reason, count in delta.items(): + target[reason] = target.get(reason, 0) + count + + +def _extract_rows(classified: ClassifiedSource) -> SourceExtraction: + run_dir = classified.run_dir + assert run_dir is not None + if classified.kind == ArtifactKind.GAE_RUN: + return extract_gae_rows(run_dir) + if classified.kind == ArtifactKind.MT_BENCH_RUN: + return extract_mt_bench_rows(run_dir) + if classified.kind == ArtifactKind.META_EVAL_RUN: + return extract_meta_eval_rows(run_dir) + raise ValueError(f"Unsupported migratable kind: {classified.kind}") + + +def _row_cache_key(row: BackfillRow) -> tuple[str, str, str, str]: + input_hash = descriptor_hash(row.canonical_input, length=None) + return row.task, row.model_spec, descriptor_hash(row.descriptor), input_hash + + +def _dedupe_and_detect_conflicts( + rows: list[BackfillRow], +) -> tuple[list[BackfillRow], int]: + grouped: dict[tuple[str, str, str, str], list[BackfillRow]] = defaultdict(list) + for row in rows: + grouped[_row_cache_key(row)].append(row) + + final: list[BackfillRow] = [] + dropped_rows = 0 + for key_rows in grouped.values(): + if len({row.output_text for row in key_rows}) > 1: + dropped_rows += len(key_rows) + continue + # Preserve duplicates so each distinct metadata association is written. + final.extend(key_rows) + + return final, dropped_rows + + +def _classify_rows_for_cell( + cell_rows: list[BackfillRow], + stored_outputs: dict[str, str], +) -> tuple[list[BackfillRow], list[BackfillRow], list[BackfillRow], int]: + to_write: list[BackfillRow] = [] + metadata_rows: list[BackfillRow] = [] + conflicting: list[BackfillRow] = [] + existing_count = 0 + + rows_by_hash: dict[str, list[BackfillRow]] = defaultdict(list) + for row in cell_rows: + rows_by_hash[descriptor_hash(row.canonical_input, length=None)].append(row) + + for input_hash, hash_rows in rows_by_hash.items(): + stored_output = stored_outputs.get(input_hash) + if stored_output is None: + to_write.append(hash_rows[0]) + metadata_rows.extend(hash_rows) + continue + if stored_output == hash_rows[0].output_text: + existing_count += 1 + metadata_rows.extend(hash_rows) + continue + conflicting.extend(hash_rows) + + return to_write, metadata_rows, conflicting, existing_count + + +def _write_rows( + rows: list[BackfillRow], + store_root: Path, + *, + dry_run: bool, +) -> tuple[int, int, dict[str, int]]: + if not rows: + return 0, 0, {} + + by_cell: dict[tuple[str, str, str], list[BackfillRow]] = defaultdict(list) + for row in rows: + by_cell[(row.task, row.model_spec, descriptor_hash(row.descriptor))].append(row) + + written = 0 + existing = 0 + skipped: dict[str, int] = {} + run_id = str(uuid.uuid4()) + + for (task, model_spec, config_hash), cell_rows in by_cell.items(): + descriptor = cell_rows[0].descriptor + try: + folder = store_folder(store_root, task, model_spec, config_hash) + db_path = folder / INFERENCE_DB_NAME + input_hashes = [ + descriptor_hash(row.canonical_input, length=None) for row in cell_rows + ] + unique_hashes = list(dict.fromkeys(input_hashes)) + + if dry_run: + if db_path.exists(): + with SQLiteInferenceStore(db_path, readonly=True) as store: + stored_outputs = store.outputs_by_hash(unique_hashes) + to_write, _, conflicting, existing_count = ( + _classify_rows_for_cell(cell_rows, stored_outputs) + ) + existing += existing_count + written += len(to_write) + if conflicting: + _merge_skip_counts( + skipped, + {"conflicting_existing_output": len(conflicting)}, + ) + else: + written += len(unique_hashes) + continue + + write_store_metadata(folder, descriptor) + with SQLiteInferenceStore(db_path) as store: + stored_outputs = store.outputs_by_hash(unique_hashes) + to_write, metadata_rows, conflicting, existing_count = ( + _classify_rows_for_cell(cell_rows, stored_outputs) + ) + existing += existing_count + if conflicting: + _merge_skip_counts( + skipped, + {"conflicting_existing_output": len(conflicting)}, + ) + + outputs_payload = [] + for row in to_write: + input_hash = descriptor_hash(row.canonical_input, length=None) + outputs_payload.append( + { + "input_hash": input_hash, + "input_text": row.canonical_input, + "output_text": row.output_text, + "producer_metadata_json": stable_json_dumps( + row.producer_metadata + ), + } + ) + + metadata_payload = [] + for row in metadata_rows: + metadata_payload.append( + { + "input_hash": descriptor_hash( + row.canonical_input, length=None + ), + "metadata_json": stable_json_dumps(row.row_metadata), + } + ) + + if outputs_payload and metadata_payload: + outputs_written, _ = store.save_outputs_and_metadata( + pd.DataFrame(outputs_payload), + pd.DataFrame(metadata_payload), + pushed_by=BACKFILL_PUSHED_BY, + run_id=run_id, + replace=False, + ) + written += outputs_written + elif outputs_payload: + written += store.save_outputs( + pd.DataFrame(outputs_payload), + pushed_by=BACKFILL_PUSHED_BY, + run_id=run_id, + replace=False, + ) + elif metadata_payload: + store.save_metadata(pd.DataFrame(metadata_payload), run_id=run_id) + except (OSError, ValueError, sqlite3.Error) as exc: + logger.warning( + "Cell integrity error for task=%s model=%s config=%s: %s", + task, + model_spec, + config_hash, + exc, + ) + _merge_skip_counts(skipped, {"cell_integrity_error": len(cell_rows)}) + + return written, existing, skipped + + +def backfill_sources( + sources: list[Path | str], + store_root: Path | str, + *, + dry_run: bool = False, +) -> BackfillReport: + """Discover saved judge runs and insert reconstructable rows into the store.""" + report = BackfillReport(dry_run=dry_run) + resolved_sources = [Path(source) for source in sources] + discovery = discover_sources(resolved_sources) + + for skipped in discovery.skipped: + reason = SKIP_REASON_BY_KIND.get(skipped.kind, "unknown_source") + report.skipped[reason] = report.skipped.get(reason, 0) + 1 + + extracted_rows: list[BackfillRow] = [] + for classified in discovery.migratable_runs: + run_dir = classified.run_dir + try: + extraction = _extract_rows(classified) + except Exception as exc: + logger.warning( + "Source extraction failed for %s: %s", + run_dir.name if run_dir else classified.path, + exc, + ) + _merge_skip_counts(report.skipped, {"source_extraction_failed": 1}) + continue + + _merge_skip_counts(report.skipped, extraction.skipped) + source_stats = report.sources.setdefault( + extraction.source_kind, + {"runs": 0, "rows_extracted": 0}, + ) + source_stats["runs"] += 1 + source_stats["rows_extracted"] += len(extraction.rows) + extracted_rows.extend(extraction.rows) + report.runs_processed += 1 + + deduped_rows, conflict_count = _dedupe_and_detect_conflicts(extracted_rows) + report.rows_planned = len(deduped_rows) + if conflict_count: + _merge_skip_counts(report.skipped, {"conflicting_outputs": conflict_count}) + + written, existing, write_skipped = _write_rows( + deduped_rows, + Path(store_root).expanduser(), + dry_run=dry_run, + ) + _merge_skip_counts(report.skipped, write_skipped) + report.written = written + report.existing = existing + return report + + +def write_report(report: BackfillReport, path: Path | str) -> None: + """Persist a JSON-safe backfill report.""" + output = Path(path) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text( + json.dumps(report.to_jsonable(), indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def log_report_summary(report: BackfillReport) -> None: + logger.info( + "Backfill complete: written=%d existing=%d rows_planned=%d " + "runs=%d dry_run=%s skipped=%s", + report.written, + report.existing, + report.rows_planned, + report.runs_processed, + report.dry_run, + report.skipped, + ) diff --git a/judgearena/cache_backfill_common.py b/judgearena/cache_backfill_common.py new file mode 100644 index 0000000..9a0417f --- /dev/null +++ b/judgearena/cache_backfill_common.py @@ -0,0 +1,59 @@ +"""Shared helpers for cache backfill extraction.""" + +from __future__ import annotations + +from pathlib import Path + +import pandas as pd +from langchain_core.messages import HumanMessage, SystemMessage +from langchain_core.prompt_values import ChatPromptValue + +HOSTED_BACKFILL_PROVIDERS = frozenset( + {"OpenRouter", "ChatOpenAI", "OpenAI", "Together", "Dummy"} +) + + +def provider_from_model_spec(model_spec: str) -> str: + provider, _, _ = model_spec.partition("/") + return provider + + +def is_backfillable_provider(model_spec: str) -> bool: + return provider_from_model_spec(model_spec) in HOSTED_BACKFILL_PROVIDERS + + +def source_run_id(run_dir: Path) -> str: + return run_dir.name + + +def chat_prompt_value( + *, system_prompt: str | None, user_prompt: str +) -> ChatPromptValue: + messages = [] + if system_prompt: + messages.append(SystemMessage(content=system_prompt)) + messages.append(HumanMessage(content=user_prompt)) + return ChatPromptValue(messages=messages) + + +def prompt_text(value: object) -> str | None: + if value is None or (isinstance(value, float) and pd.isna(value)): + return None + text = str(value) + if not text.strip() or text.strip().lower() == "nan": + return None + return text + + +def mt_swapped(value: object) -> bool: + if value is None or (isinstance(value, float) and pd.isna(value)): + return False + if isinstance(value, bool): + return value + if isinstance(value, (int, float)): + return bool(value) + return str(value).strip().lower() in {"1", "true", "yes"} + + +def increment(skipped: dict[str, int], reason: str, count: int = 1) -> None: + skipped[reason] = skipped.get(reason, 0) + count diff --git a/judgearena/cache_backfill_config.py b/judgearena/cache_backfill_config.py new file mode 100644 index 0000000..0ad765b --- /dev/null +++ b/judgearena/cache_backfill_config.py @@ -0,0 +1,166 @@ +"""Reconstruct model configurations from historical JudgeArena runs.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from judgearena.config import CacheArgs, RunConfig, load_config +from judgearena.meta_eval.cli_args import CliMetaEvalArgs +from judgearena.model_adapters import PreparedModel +from judgearena.models import build_default_judge_model_kwargs, make_model +from judgearena.repro import METADATA_FILENAME + + +def _cache_relevant_config(cfg: RunConfig) -> dict[str, Any]: + payload = cfg.model_dump(mode="json") + payload.pop("cache", None) + payload.pop("run", None) + return payload + + +def _load_modern_run_config(run_dir: Path) -> RunConfig | None: + candidates: list[tuple[str, RunConfig]] = [] + metadata_path = run_dir / METADATA_FILENAME + if metadata_path.exists(): + payload = json.loads(metadata_path.read_text(encoding="utf-8")) + run = payload.get("run") + if not isinstance(run, dict): + raise ValueError( + f"{METADATA_FILENAME} has no run config in {run_dir.name}." + ) + candidates.append((METADATA_FILENAME, RunConfig(**run))) + config_path = run_dir / "config.yaml" + if config_path.exists(): + candidates.append(("config.yaml", load_config(config_path))) + if not candidates: + return None + + expected = _cache_relevant_config(candidates[0][1]) + conflicting = [ + name + for name, candidate in candidates[1:] + if _cache_relevant_config(candidate) != expected + ] + if conflicting: + sources = ", ".join([candidates[0][0], *conflicting]) + raise ValueError(f"Conflicting run configs in {run_dir.name}: {sources}") + return candidates[0][1] + + +def _legacy_args_to_run_config(args: dict[str, Any]) -> RunConfig: + judge_model = args.get("judge_model") or args.get("judge", {}).get("model") + if not isinstance(judge_model, str): + raise ValueError("Legacy args missing judge_model.") + + engine_kwargs = dict(args.get("engine_kwargs") or {}) + judge_engine_kwargs = dict(args.get("judge_engine_kwargs") or {}) + judge_engine_kwargs.update(engine_kwargs) + truncate_all = args.get("truncate_all_input_chars", 8192) + truncate_judge = args.get("truncate_judge_input_chars") + if truncate_judge is None: + truncate_judge = truncate_all + + return RunConfig( + task=str(args["task"]), + model={ + "name": args.get("model_A") or args.get("model", {}).get("name"), + "baseline": args.get("model_B") or args.get("model", {}).get("baseline"), + "max_out_tokens": args.get("max_out_tokens_models") + or args.get("model", {}).get("max_out_tokens", 32768), + "max_model_len": args.get("max_model_len"), + "chat_template": args.get("chat_template"), + "engine_kwargs": engine_kwargs, + }, + judge={ + "model": judge_model, + "max_out_tokens": args.get("max_out_tokens_judge") + or args.get("judge", {}).get("max_out_tokens", 32768), + "max_model_len": args.get("max_model_len_judge") + or args.get("max_model_len"), + "chat_template": args.get("chat_template_judge") + or args.get("chat_template"), + "engine_kwargs": judge_engine_kwargs, + "provide_explanation": bool(args.get("provide_explanation", False)), + "swap_mode": args.get("swap_mode", "fixed"), + "prompt_preset": args.get("prompt_preset"), + "system_prompt_file": args.get("judge_system_prompt_file"), + "user_prompt_file": args.get("judge_user_prompt_file"), + "strip_thinking_before_judging": bool( + args.get("strip_thinking_before_judging", False) + ), + }, + generation={ + "n_instructions": args.get("n_instructions"), + "truncate_all_input_chars": truncate_all, + "truncate_judge_input_chars": truncate_judge, + }, + run={ + "result_folder": str(args.get("result_folder", "results")), + "seed": args.get("seed", 0), + }, + ) + + +def load_gae_run_config(run_dir: Path) -> RunConfig: + modern_cfg = _load_modern_run_config(run_dir) + if modern_cfg is not None: + return modern_cfg + args_paths = sorted(run_dir.glob("args-*.json")) + if not args_paths: + raise ValueError(f"No reconstructable config found under {run_dir.name}.") + if len(args_paths) > 1: + names = ", ".join(path.name for path in args_paths) + raise ValueError( + f"Ambiguous legacy args files under {run_dir.name}; " + f"expected one args-*.json or config/metadata: {names}" + ) + args = json.loads(args_paths[0].read_text(encoding="utf-8")) + return _legacy_args_to_run_config(args) + + +def load_meta_args(run_dir: Path) -> CliMetaEvalArgs: + args_path = run_dir / "args.json" + if not args_path.exists(): + raise ValueError(f"Meta-eval run missing args.json: {run_dir.name}") + payload = json.loads(args_path.read_text(encoding="utf-8")) + cache_payload = payload.pop("cache", {}) + payload.pop("ignore_cache", None) + if isinstance(cache_payload, dict): + cache_payload.pop("ignore_cache", None) + payload["cache"] = CacheArgs(**cache_payload) if cache_payload else CacheArgs() + return CliMetaEvalArgs(**payload) + + +def build_gae_judge_model(cfg: RunConfig) -> PreparedModel: + return make_model( + model=cfg.judge.model, + **build_default_judge_model_kwargs( + cfg.judge.model, + cfg.model.engine_kwargs, + judge_engine_kwargs_override=cfg.judge.model_kwargs( + fallback_chat_template=cfg.model.chat_template, + ), + ), + ) + + +def build_mt_judge_model(cfg: RunConfig, *, delegated: bool) -> PreparedModel: + judge_model_kwargs = cfg.judge.model_kwargs( + base_engine_kwargs=cfg.model.engine_kwargs, + fallback_chat_template=cfg.model.chat_template, + ) + if delegated and cfg.judge.temperature is None: + judge_model_kwargs.setdefault("temperature", 0.0) + return make_model(model=cfg.judge.model, **judge_model_kwargs) + + +def build_meta_judge_model(args: CliMetaEvalArgs) -> PreparedModel: + return make_model( + model=args.judge_model, + max_tokens=args.max_out_tokens_judge, + max_model_len=args.max_model_len, + chat_template=args.chat_template, + **args.engine_kwargs, + ) diff --git a/judgearena/cache_backfill_discovery.py b/judgearena/cache_backfill_discovery.py new file mode 100644 index 0000000..0b8577a --- /dev/null +++ b/judgearena/cache_backfill_discovery.py @@ -0,0 +1,340 @@ +"""Discover and classify saved run folders and cache artifacts for backfill.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass, field +from enum import StrEnum +from pathlib import Path + +import yaml + +from judgearena.constants import ELO_TASK_PREFIX, META_EVAL_TASK +from judgearena.repro import METADATA_FILENAME + +LEGACY_CELL_DB_NAMES = frozenset({"judgements.db", "completions.db"}) +PASS_LEVEL_CACHE_SUFFIXES = (".csv.zip", ".parquet", ".csv") +GENERATION_ONLY_MARKERS = frozenset( + { + "completions.parquet", + "completions.csv", + "completions.csv.zip", + "model_outputs.parquet", + } +) +GAE_REQUIRED_COLUMNS = frozenset( + {"instruction", "completion_A", "completion_B", "judge_input"} +) +GAE_OUTPUT_COLUMNS = frozenset({"judge_completion", "judge_output"}) + + +class ArtifactKind(StrEnum): + GAE_RUN = "gae_run" + MT_BENCH_RUN = "mt_bench_run" + META_EVAL_RUN = "meta_eval_run" + ELO_RUN = "elo_run" + LEGACY_CACHE_CELL = "legacy_cache_cell" + META_EVAL_IDENTITY_DB = "meta_eval_identity_db" + PASS_LEVEL_CACHE = "pass_level_cache" + GENERATION_ARTIFACT = "generation_artifact" + UNKNOWN = "unknown" + + +SKIP_REASON_BY_KIND: dict[ArtifactKind, str] = { + ArtifactKind.ELO_RUN: "elo_run_missing_inference_outputs", + ArtifactKind.LEGACY_CACHE_CELL: "legacy_cache_cell_unmigratable", + ArtifactKind.META_EVAL_IDENTITY_DB: "meta_eval_identity_db", + ArtifactKind.PASS_LEVEL_CACHE: "pass_level_cache_untrusted", + ArtifactKind.GENERATION_ARTIFACT: "generation_provenance_unknown", + ArtifactKind.UNKNOWN: "unknown_judge_run", +} + + +@dataclass +class ClassifiedSource: + path: Path + kind: ArtifactKind + run_dir: Path | None = None + + +@dataclass +class DiscoveryReport: + migratable_runs: list[ClassifiedSource] = field(default_factory=list) + skipped: list[ClassifiedSource] = field(default_factory=list) + + +def _glob_has_matches(resolved: Path, pattern: str) -> bool: + return next(resolved.glob(pattern), None) is not None + + +def _is_elo_task_name(value: str | None) -> bool: + return bool(value and value.startswith(ELO_TASK_PREFIX)) + + +def _looks_like_meta_eval_dir(path: Path) -> bool: + return ( + path.name.startswith(f"{META_EVAL_TASK}-") + or (path / "annotations.parquet").exists() + ) + + +def _csv_columns(csv_path: Path) -> set[str]: + header = csv_path.read_text(encoding="utf-8").splitlines()[:1] + if not header: + return set() + return {part.strip() for part in header[0].split(",")} + + +def _looks_like_mt_annotations(path: Path) -> bool: + for csv_path in path.glob("*-annotations.csv"): + columns = _csv_columns(csv_path) + if not columns: + continue + mt_markers = {"question_id", "turn", "category"} + if mt_markers.issubset(columns): + return True + if "g1_user_prompt" in columns or "user_prompt" in columns: + return True + return False + + +def _looks_like_gae_annotations(path: Path) -> bool: + for csv_path in path.glob("*-annotations.csv"): + columns = _csv_columns(csv_path) + if not columns: + continue + if not GAE_REQUIRED_COLUMNS.issubset(columns): + continue + if not GAE_OUTPUT_COLUMNS.intersection(columns): + continue + return True + return False + + +def _task_from_run_dir(run_dir: Path) -> str | None: + metadata_path = run_dir / METADATA_FILENAME + if metadata_path.exists(): + payload = json.loads(metadata_path.read_text(encoding="utf-8")) + run = payload.get("run") + if isinstance(run, dict): + task = run.get("task") + if isinstance(task, str): + return task + config_path = run_dir / "config.yaml" + if config_path.exists(): + payload = yaml.safe_load(config_path.read_text(encoding="utf-8")) or {} + task = payload.get("task") + if isinstance(task, str): + return task + for args_path in run_dir.glob("args-*.json"): + payload = json.loads(args_path.read_text(encoding="utf-8")) + task = payload.get("task") + if isinstance(task, str): + return task + args_path = run_dir / "args.json" + if args_path.exists(): + payload = json.loads(args_path.read_text(encoding="utf-8")) + if isinstance(payload, dict): + task = payload.get("task") + if isinstance(task, str): + return task + return None + + +def _classify_path(path: Path) -> ClassifiedSource: + resolved = path.resolve() + + if resolved.is_file(): + suffix = "".join(resolved.suffixes) or resolved.suffix + name = resolved.name + if name in GENERATION_ONLY_MARKERS: + return ClassifiedSource(resolved, ArtifactKind.GENERATION_ARTIFACT) + if ( + name.endswith(PASS_LEVEL_CACHE_SUFFIXES) + or suffix in PASS_LEVEL_CACHE_SUFFIXES + ): + return ClassifiedSource(resolved, ArtifactKind.PASS_LEVEL_CACHE) + if resolved.suffix == ".db": + parts = {part.lower() for part in resolved.parts} + if "cache" in parts and "db" in parts: + return ClassifiedSource(resolved, ArtifactKind.META_EVAL_IDENTITY_DB) + if name in LEGACY_CELL_DB_NAMES: + return ClassifiedSource(resolved, ArtifactKind.LEGACY_CACHE_CELL) + return ClassifiedSource(resolved, ArtifactKind.UNKNOWN) + + if not resolved.is_dir(): + return ClassifiedSource(resolved, ArtifactKind.UNKNOWN) + + if resolved.name in LEGACY_CELL_DB_NAMES or any( + (resolved / db_name).exists() for db_name in LEGACY_CELL_DB_NAMES + ): + return ClassifiedSource(resolved, ArtifactKind.LEGACY_CACHE_CELL) + + parts = {part.lower() for part in resolved.parts} + if ( + "cache" in parts + and "db" in parts + and any(child.suffix == ".db" for child in resolved.glob("*.db")) + ): + return ClassifiedSource(resolved, ArtifactKind.META_EVAL_IDENTITY_DB) + + task = _task_from_run_dir(resolved) + if _is_elo_task_name(task) or resolved.name.startswith(ELO_TASK_PREFIX): + return ClassifiedSource(resolved, ArtifactKind.ELO_RUN, run_dir=resolved) + + if (resolved / "annotations.parquet").exists() and ( + (resolved / "args.json").exists() or (resolved / METADATA_FILENAME).exists() + ): + return ClassifiedSource( + resolved, + ArtifactKind.META_EVAL_RUN, + run_dir=resolved, + ) + + if list(resolved.glob("*-annotations.csv")): + if task == "mt-bench" or _looks_like_mt_annotations(resolved): + return ClassifiedSource( + resolved, + ArtifactKind.MT_BENCH_RUN, + run_dir=resolved, + ) + if _is_elo_task_name(task): + return ClassifiedSource(resolved, ArtifactKind.ELO_RUN, run_dir=resolved) + if _looks_like_gae_annotations(resolved): + return ClassifiedSource(resolved, ArtifactKind.GAE_RUN, run_dir=resolved) + return ClassifiedSource(resolved, ArtifactKind.UNKNOWN) + + if _looks_like_meta_eval_dir(resolved) and (resolved / "args.json").exists(): + return ClassifiedSource( + resolved, + ArtifactKind.META_EVAL_RUN, + run_dir=resolved, + ) + + if any(resolved.joinpath(name).exists() for name in GENERATION_ONLY_MARKERS): + return ClassifiedSource(resolved, ArtifactKind.GENERATION_ARTIFACT) + + if any( + _glob_has_matches(resolved, f"*{suffix}") + for suffix in PASS_LEVEL_CACHE_SUFFIXES + ): + if not list(resolved.glob("*-annotations.csv")): + return ClassifiedSource(resolved, ArtifactKind.PASS_LEVEL_CACHE) + + return ClassifiedSource(resolved, ArtifactKind.UNKNOWN) + + +def _discover_run_dirs(source: Path) -> list[Path]: + classified = _classify_path(source) + if classified.run_dir is not None: + return [classified.run_dir] + + if not source.is_dir(): + return [] + + run_dirs: list[Path] = [] + seen: set[Path] = set() + for annotation_csv in source.rglob("*-annotations.csv"): + run_dir = annotation_csv.parent.resolve() + if run_dir not in seen: + seen.add(run_dir) + run_dirs.append(run_dir) + for annotation_parquet in source.rglob("annotations.parquet"): + run_dir = annotation_parquet.parent.resolve() + if run_dir not in seen: + seen.add(run_dir) + run_dirs.append(run_dir) + return sorted(run_dirs) + + +def _discover_nested_skipped_artifacts(source: Path) -> list[ClassifiedSource]: + skipped: list[ClassifiedSource] = [] + seen: set[Path] = set() + for db_name in LEGACY_CELL_DB_NAMES: + for db_path in source.rglob(db_name): + parent = db_path.parent.resolve() + if parent in seen: + continue + seen.add(parent) + classified = _classify_path(parent) + if classified.kind in SKIP_REASON_BY_KIND: + skipped.append(classified) + return skipped + + +def _collect_skipped_artifacts( + source: Path, seen_skipped: set[Path] +) -> list[ClassifiedSource]: + skipped: list[ClassifiedSource] = [] + for child in source.rglob("*"): + if child in seen_skipped: + continue + classified = _classify_path(child) + if classified.kind not in SKIP_REASON_BY_KIND: + continue + seen_skipped.add(child) + skipped.append(classified) + return skipped + + +def discover_sources(sources: list[Path]) -> DiscoveryReport: + """Discover migratable judge run folders and classify skipped artifacts.""" + report = DiscoveryReport() + seen_runs: set[Path] = set() + seen_skipped: set[Path] = set() + + for source in sources: + source = source.resolve() + direct = _classify_path(source) + + if direct.run_dir is not None: + if direct.run_dir not in seen_runs: + seen_runs.add(direct.run_dir) + report.migratable_runs.append(direct) + continue + + if source.is_file(): + if source not in seen_skipped: + seen_skipped.add(source) + report.skipped.append(direct) + continue + + run_dirs = _discover_run_dirs(source) + for nested in _discover_nested_skipped_artifacts(source): + if nested.path not in seen_skipped: + seen_skipped.add(nested.path) + report.skipped.append(nested) + + if not run_dirs: + for classified_child in _collect_skipped_artifacts(source, seen_skipped): + report.skipped.append(classified_child) + if ( + direct.kind in SKIP_REASON_BY_KIND + and source not in seen_skipped + and (source.is_file() or direct.kind != ArtifactKind.UNKNOWN) + ): + seen_skipped.add(source) + report.skipped.append(direct) + continue + + for run_dir in run_dirs: + if run_dir in seen_runs: + continue + classified = _classify_path(run_dir) + if classified.kind in SKIP_REASON_BY_KIND: + if run_dir not in seen_skipped: + seen_skipped.add(run_dir) + report.skipped.append(classified) + continue + if classified.kind in { + ArtifactKind.GAE_RUN, + ArtifactKind.MT_BENCH_RUN, + ArtifactKind.META_EVAL_RUN, + }: + seen_runs.add(run_dir) + report.migratable_runs.append(classified) + elif run_dir not in seen_skipped: + seen_skipped.add(run_dir) + report.skipped.append(classified) + + return report diff --git a/judgearena/cache_backfill_sources.py b/judgearena/cache_backfill_sources.py new file mode 100644 index 0000000..51af8cd --- /dev/null +++ b/judgearena/cache_backfill_sources.py @@ -0,0 +1,368 @@ +"""Extract backfill rows from saved GAE, MT-Bench, and meta-eval run folders.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import pandas as pd + +from judgearena.cache_backfill_common import ( + chat_prompt_value, + increment, + is_backfillable_provider, + mt_swapped, + prompt_text, + source_run_id, +) +from judgearena.cache_backfill_config import ( + build_gae_judge_model, + build_meta_judge_model, + build_mt_judge_model, + load_gae_run_config, + load_meta_args, +) +from judgearena.config import RunConfig, meta_eval_cache_task +from judgearena.evaluate import render_judge_inputs, resolve_run_judge_prompt +from judgearena.meta_eval.prompts import resolve_prompt_mode +from judgearena.model_adapters import PreparedModel + + +@dataclass(frozen=True) +class BackfillRow: + task: str + model_spec: str + descriptor: dict[str, Any] + canonical_input: str + output_text: str + row_metadata: dict[str, Any] + producer_metadata: dict[str, Any] + + +@dataclass +class SourceExtraction: + rows: list[BackfillRow] + skipped: dict[str, int] + source_kind: str + + +def _annotations_path(run_dir: Path) -> Path: + paths = sorted(run_dir.glob("*-annotations.csv")) + if len(paths) != 1: + names = ", ".join(path.name for path in paths) or "none" + raise ValueError( + f"Expected exactly one *-annotations.csv in {run_dir.name}; found {names}." + ) + return paths[0] + + +def _maybe_descriptor(model: PreparedModel) -> dict[str, Any] | None: + return model.cache_descriptor() + + +def _infer_gae_orientation(row: pd.Series, *, cfg: RunConfig) -> str | None: + model_a = str(row.get("model_A", "")) + model_b = str(row.get("model_B", "")) + focal = cfg.model.name + a_is_focal = model_a == focal + b_is_focal = model_b == focal + if a_is_focal and not b_is_focal: + return "direct" + if b_is_focal and not a_is_focal: + return "reversed" + return None + + +def _meta_eval_verify_completions(row: pd.Series) -> tuple[str, str]: + presented_a = prompt_text(row.get("presented_completion_a")) + presented_b = prompt_text(row.get("presented_completion_b")) + if presented_a is not None and presented_b is not None: + return presented_a, presented_b + completion_a = str(row.get("completion_a", "")) + completion_b = str(row.get("completion_b", "")) + if str(row.get("orientation", "forward")) == "swapped": + return completion_b, completion_a + return completion_a, completion_b + + +def extract_gae_rows(run_dir: Path) -> SourceExtraction: + cfg = load_gae_run_config(run_dir) + annotations_path = _annotations_path(run_dir) + df = pd.read_csv(annotations_path, keep_default_na=False) + skipped: dict[str, int] = {} + rows: list[BackfillRow] = [] + run_id = source_run_id(run_dir) + + if not is_backfillable_provider(cfg.judge.model): + increment(skipped, "local_engine_unsupported", len(df)) + return SourceExtraction(rows=[], skipped=skipped, source_kind="gae") + + judge_model = build_gae_judge_model(cfg) + descriptor = _maybe_descriptor(judge_model) + if descriptor is None: + increment(skipped, "local_engine_unsupported", len(df)) + return SourceExtraction(rows=[], skipped=skipped, source_kind="gae") + + resolved_prompt = resolve_run_judge_prompt(cfg.task, cfg.judge) + output_column = ( + "judge_completion" if "judge_completion" in df.columns else "judge_output" + ) + if output_column not in df.columns: + increment(skipped, "judge_input_unverifiable", len(df)) + return SourceExtraction(rows=[], skipped=skipped, source_kind="gae") + + if "judge_input" not in df.columns: + increment(skipped, "judge_input_unverifiable", len(df)) + return SourceExtraction(rows=[], skipped=skipped, source_kind="gae") + + rendered_inputs = render_judge_inputs( + df["instruction"].astype(str).tolist(), + df["completion_A"].astype(str).tolist(), + df["completion_B"].astype(str).tolist(), + system_prompt=resolved_prompt.system_prompt, + user_prompt_template=resolved_prompt.user_prompt_template, + truncate_input_chars=cfg.generation.truncate_judge_input_chars, + provide_explanation=cfg.judge.provide_explanation, + prompt_preset=resolved_prompt.preset_name, + strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging, + task=cfg.task, + system_file=cfg.judge.system_prompt_file, + user_file=cfg.judge.user_prompt_file, + ) + + producer = judge_model.producer_metadata() + + for idx, (_, row) in enumerate(df.iterrows()): + stored_input = row.get("judge_input") + if pd.isna(stored_input) or stored_input is None: + increment(skipped, "judge_input_unverifiable") + continue + rendered = rendered_inputs[idx].to_string() + if str(stored_input) != rendered: + increment(skipped, "judge_input_mismatch") + continue + output_text = prompt_text(row.get(output_column)) + if output_text is None: + increment(skipped, "judge_output_missing") + continue + + orientation = _infer_gae_orientation(row, cfg=cfg) + if orientation is None: + increment(skipped, "battle_orientation_unverifiable") + continue + prompt_input = rendered_inputs[idx] + rows.append( + BackfillRow( + task=cfg.task, + model_spec=judge_model.model_spec, + descriptor=descriptor, + canonical_input=judge_model.canonicalize_input(prompt_input), + output_text=output_text, + row_metadata={ + "task": cfg.task, + "instruction_index": str(row.get("instruction_index", idx)), + "presented_model_a": str(row.get("model_A", cfg.model.name)), + "presented_model_b": str(row.get("model_B", "")), + "orientation": orientation, + "source_run_id": run_id, + }, + producer_metadata=producer, + ) + ) + + return SourceExtraction(rows=rows, skipped=skipped, source_kind="gae") + + +def _mt_is_fastchat(df: pd.DataFrame) -> bool: + return "g1_user_prompt" in df.columns + + +def extract_mt_bench_rows(run_dir: Path) -> SourceExtraction: + cfg = load_gae_run_config(run_dir) + annotations_path = _annotations_path(run_dir) + df = pd.read_csv(annotations_path) + skipped: dict[str, int] = {} + rows: list[BackfillRow] = [] + run_id = source_run_id(run_dir) + + if not is_backfillable_provider(cfg.judge.model): + increment(skipped, "local_engine_unsupported", len(df)) + return SourceExtraction(rows=[], skipped=skipped, source_kind="mt_bench") + + resolved_prompt = resolve_run_judge_prompt(cfg.task, cfg.judge, multi_turn=True) + judge_model = build_mt_judge_model(cfg, delegated=resolved_prompt.delegated) + descriptor = _maybe_descriptor(judge_model) + if descriptor is None: + increment(skipped, "local_engine_unsupported", len(df)) + return SourceExtraction(rows=[], skipped=skipped, source_kind="mt_bench") + + producer = judge_model.producer_metadata() + fastchat = _mt_is_fastchat(df) + + for _, row in df.iterrows(): + candidates: list[tuple[str, str | None, str, str]] = [] + if fastchat: + g1_output = prompt_text(row.get("g1_judgment")) + g1_prompt = prompt_text(row.get("g1_user_prompt")) + if g1_output is not None and g1_prompt is not None: + candidates.append( + ( + "direct", + prompt_text(row.get("system_prompt")), + g1_prompt, + g1_output, + ) + ) + g2_output = prompt_text(row.get("g2_judgment")) + g2_prompt = prompt_text(row.get("g2_user_prompt")) + if g2_output is not None and g2_prompt is not None: + candidates.append( + ( + "reversed", + prompt_text(row.get("system_prompt")), + g2_prompt, + g2_output, + ) + ) + else: + output = prompt_text(row.get("judge_completion")) + if output is None: + increment(skipped, "judge_input_unverifiable") + continue + orientation = "reversed" if mt_swapped(row.get("swapped")) else "direct" + user_prompt = prompt_text(row.get("user_prompt")) + if user_prompt is None: + increment(skipped, "judge_input_unverifiable") + continue + candidates.append( + ( + orientation, + prompt_text(row.get("system_prompt")), + user_prompt, + output, + ) + ) + + for orientation, system_prompt, user_prompt, output_text in candidates: + prompt_input = chat_prompt_value( + system_prompt=system_prompt, + user_prompt=user_prompt, + ) + turn_value = row.get("turn") + rows.append( + BackfillRow( + task=cfg.task, + model_spec=judge_model.model_spec, + descriptor=descriptor, + canonical_input=judge_model.canonicalize_input(prompt_input), + output_text=output_text, + row_metadata={ + "question_id": str(row.get("question_id", "")), + "category": row.get("category"), + "turn": int(turn_value) if pd.notna(turn_value) else None, + "orientation": orientation, + "prompt": row.get("prompt_name"), + "source_run_id": run_id, + }, + producer_metadata=producer, + ) + ) + + return SourceExtraction(rows=rows, skipped=skipped, source_kind="mt_bench") + + +def extract_meta_eval_rows(run_dir: Path) -> SourceExtraction: + args = load_meta_args(run_dir) + df = pd.read_parquet(run_dir / "annotations.parquet") + skipped: dict[str, int] = {} + rows: list[BackfillRow] = [] + run_id = source_run_id(run_dir) + + if not is_backfillable_provider(args.judge_model): + increment(skipped, "local_engine_unsupported", len(df)) + return SourceExtraction(rows=[], skipped=skipped, source_kind="meta_eval") + + prompt_spec = resolve_prompt_mode( + args.prompt_mode, + provide_explanation=args.provide_explanation, + ) + judge_model = build_meta_judge_model(args) + descriptor = _maybe_descriptor(judge_model) + if descriptor is None: + increment(skipped, "local_engine_unsupported", len(df)) + return SourceExtraction(rows=[], skipped=skipped, source_kind="meta_eval") + + if "judge_input" not in df.columns: + increment(skipped, "judge_input_unverifiable", len(df)) + return SourceExtraction(rows=[], skipped=skipped, source_kind="meta_eval") + + output_column = ( + "judge_completion" if "judge_completion" in df.columns else "judge_output" + ) + if output_column not in df.columns: + increment(skipped, "judge_input_unverifiable", len(df)) + return SourceExtraction(rows=[], skipped=skipped, source_kind="meta_eval") + + verify_a: list[str] = [] + verify_b: list[str] = [] + for _, row in df.iterrows(): + completion_a, completion_b = _meta_eval_verify_completions(row) + verify_a.append(completion_a) + verify_b.append(completion_b) + + rendered_inputs = render_judge_inputs( + df["instruction"].astype(str).tolist(), + verify_a, + verify_b, + system_prompt=prompt_spec.system_prompt, + user_prompt_template=prompt_spec.user_prompt_template, + truncate_input_chars=args.truncate_judge_input_chars, + provide_explanation=args.provide_explanation, + ) + + task = meta_eval_cache_task(args.reference_arena) + producer = judge_model.producer_metadata() + + for idx, (_, row) in enumerate(df.iterrows()): + stored_input = row.get("judge_input") + if pd.isna(stored_input) or stored_input is None: + increment(skipped, "judge_input_unverifiable") + continue + rendered = rendered_inputs[idx].to_string() + if str(stored_input) != rendered: + increment(skipped, "judge_input_mismatch") + continue + output_text = prompt_text(row.get(output_column)) + if output_text is None: + increment(skipped, "judge_output_missing") + continue + + orientation = str(row.get("orientation", "forward")) + prompt_input = rendered_inputs[idx] + rows.append( + BackfillRow( + task=task, + model_spec=judge_model.model_spec, + descriptor=descriptor, + canonical_input=judge_model.canonicalize_input(prompt_input), + output_text=output_text, + row_metadata={ + "reference_arena": args.reference_arena, + "benchmark": str(row.get("benchmark", "")), + "question_id": str(row.get("question_id", "")), + "presented_model_a": str( + row.get("presented_model_a", row.get("model_a", "")) + ), + "presented_model_b": str( + row.get("presented_model_b", row.get("model_b", "")) + ), + "prompt_mode": args.prompt_mode, + "orientation": orientation, + "source_run_id": run_id, + }, + producer_metadata=producer, + ) + ) + + return SourceExtraction(rows=rows, skipped=skipped, source_kind="meta_eval") diff --git a/judgearena/cache_sync.py b/judgearena/cache_sync.py new file mode 100644 index 0000000..04439ae --- /dev/null +++ b/judgearena/cache_sync.py @@ -0,0 +1,176 @@ +"""Standalone CLI for synchronizing inference cache cells with Hugging Face.""" + +from __future__ import annotations + +import argparse +import getpass +import sys +from pathlib import Path + +from judgearena.cache_backfill import backfill_sources, log_report_summary, write_report +from judgearena.log import configure_logging, get_logger +from judgearena.store_sync import ( + DEFAULT_CACHE_REPO, + fetch_remote_cells, + iter_cell_dbs, + push_cells, + validate_path_filters, +) + +logger = get_logger(__name__) + + +def _add_filter_args(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--prefix", + help="Exact repository path prefix (overrides task/provider/model/config).", + ) + parser.add_argument("--task", help="Filter cells by benchmark task name.") + parser.add_argument("--provider", help="Filter cells by provider, e.g. VLLM.") + parser.add_argument( + "--model", + help="Filter cells by model path (slashes become '--' in folders).", + ) + parser.add_argument("--config_hash", help="Filter cells by descriptor hash.") + + +def _add_common(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--store_root", required=True, help="Local store root.") + parser.add_argument("--cache_hf_repo", default=DEFAULT_CACHE_REPO) + parser.add_argument("--repo_type", default="dataset") + parser.add_argument("--revision", default="main") + parser.add_argument("-v", "--verbose", action="count", default=0) + + +def _resolve_prefix(args: argparse.Namespace) -> str | None: + return validate_path_filters( + prefix=args.prefix, + task=args.task, + provider=args.provider, + model=args.model, + config_hash=args.config_hash, + ) + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="judgearena-cache", + description="Fetch or push shared inference cache cells.", + ) + subparsers = parser.add_subparsers(dest="command", required=True) + + fetch = subparsers.add_parser( + "fetch", + help="Discover and merge remote cells into the local store.", + ) + _add_common(fetch) + _add_filter_args(fetch) + + push = subparsers.add_parser( + "push", + help="Merge and upload local cells.", + ) + _add_common(push) + _add_filter_args(push) + push.add_argument("--pushed_by", default=getpass.getuser()) + push.add_argument("--create_pr", action="store_true") + push.add_argument( + "--ensure_repo", + action="store_true", + help="Create the repository if it does not exist.", + ) + push.add_argument( + "--public", + action="store_true", + help="Create a public repository when used with --ensure_repo.", + ) + + backfill = subparsers.add_parser( + "backfill", + help="Backfill hosted judge outputs from saved run folders.", + ) + backfill.add_argument( + "sources", + nargs="+", + type=Path, + help="Run folders or parents containing saved judge annotations.", + ) + backfill.add_argument("--store_root", required=True, help="Local store root.") + backfill.add_argument( + "--dry_run", + action="store_true", + help="Plan and report without writing inference rows.", + ) + backfill.add_argument( + "--report", + type=Path, + help="Optional path to write a JSON backfill report.", + ) + backfill.add_argument("-v", "--verbose", action="count", default=0) + return parser + + +def main(argv: list[str] | None = None) -> None: + args = _build_parser().parse_args(argv) + configure_logging(getattr(args, "verbose", 0)) + + if args.command == "backfill": + report = backfill_sources( + args.sources, + args.store_root, + dry_run=args.dry_run, + ) + if args.report is not None: + write_report(report, args.report) + log_report_summary(report) + return + + try: + path_prefix = _resolve_prefix(args) + except ValueError as exc: + logger.error("%s", exc) + sys.exit(1) + + if args.command == "fetch": + if path_prefix is None: + logger.error( + "Fetch requires a path filter. Provide --prefix or at least --task." + ) + sys.exit(1) + fetched = fetch_remote_cells( + args.cache_hf_repo, + args.store_root, + path_prefix=path_prefix, + repo_type=args.repo_type, + revision=args.revision, + strict=True, + ) + if not fetched: + logger.warning( + "No remote cells matched prefix %r under %s", + path_prefix, + args.cache_hf_repo, + ) + return + + db_paths = iter_cell_dbs(args.store_root, path_prefix=path_prefix) + if not db_paths: + logger.warning("No inference.db cells found under %s", args.store_root) + return + + push_cells( + args.cache_hf_repo, + args.store_root, + db_paths, + pushed_by=args.pushed_by, + repo_type=args.repo_type, + revision=args.revision, + create_pr=args.create_pr, + ensure_repo=args.ensure_repo, + private=not args.public, + strict=True, + ) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index dbb37e9..f6b09ba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,6 +4,7 @@ build-backend = "setuptools.build_meta" [project.scripts] judgearena = "judgearena.cli:cli" +judgearena-cache = "judgearena.cache_sync:main" [project] name = "judgearena" diff --git a/tests/test_cache_backfill.py b/tests/test_cache_backfill.py new file mode 100644 index 0000000..1ce0a66 --- /dev/null +++ b/tests/test_cache_backfill.py @@ -0,0 +1,792 @@ +from __future__ import annotations + +import json +from pathlib import Path + +import pandas as pd +import pytest + +import judgearena.generate_and_evaluate as gae +from judgearena import cache_sync +from judgearena.cache_backfill import ( + BACKFILL_PUSHED_BY, + BackfillReport, + backfill_sources, + write_report, +) +from judgearena.cache_backfill_discovery import ArtifactKind, discover_sources +from judgearena.cache_backfill_sources import ( + _infer_gae_orientation, + extract_gae_rows, + extract_mt_bench_rows, +) +from judgearena.config import RunConfig, dump_config, meta_eval_cache_task +from judgearena.evaluate import render_judge_inputs, resolve_run_judge_prompt +from judgearena.meta_eval.cli_args import CliMetaEvalArgs +from judgearena.meta_eval.prompts import resolve_prompt_mode +from judgearena.repro import write_run_metadata +from judgearena.store_sqlite import INFERENCE_DB_NAME, SQLiteInferenceStore + + +def _synthetic_instructions(n: int = 2) -> pd.DataFrame: + return pd.DataFrame( + {"instruction": [f"instruction {i}" for i in range(n)]}, + index=pd.Index([f"idx-{i}" for i in range(n)], name="instruction_index"), + ) + + +def _cfg_with_cache(tmp_path, **overrides) -> RunConfig: + payload = { + "task": "alpaca-eval", + "model": {"name": "Dummy/gen-a", "baseline": "Dummy/gen-b"}, + "judge": {"model": "Dummy/score A: 0 score B: 10", "swap_mode": "fixed"}, + "generation": {"n_instructions": 2}, + "run": {"result_folder": str(tmp_path / "results"), "no_log_file": True}, + "cache": {"store_root": str(tmp_path / "live-cache")}, + } + payload.update(overrides) + return RunConfig(**payload) + + +def _write_gae_annotations( + run_dir: Path, + cfg: RunConfig, + *, + swap_both: bool = False, +) -> None: + instructions = ["instruction 0", "instruction 1"] + completions_a = ["completion-a-0", "completion-a-1"] + completions_b = ["completion-b-0", "completion-b-1"] + resolved = resolve_run_judge_prompt(cfg.task, cfg.judge) + rendered = render_judge_inputs( + instructions, + completions_a, + completions_b, + system_prompt=resolved.system_prompt, + user_prompt_template=resolved.user_prompt_template, + truncate_input_chars=cfg.generation.truncate_judge_input_chars, + provide_explanation=cfg.judge.provide_explanation, + prompt_preset=resolved.preset_name, + strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging, + task=cfg.task, + ) + rows = [] + for idx, (instruction, ca, cb, judge_input) in enumerate( + zip(instructions, completions_a, completions_b, rendered, strict=True) + ): + rows.append( + { + "instruction": instruction, + "completion_A": ca, + "completion_B": cb, + "judge_completion": "Score A: 8\nScore B: 6", + "judge_input": judge_input.to_string(), + "instruction_index": f"idx-{idx}", + "model_A": cfg.model.name, + "model_B": cfg.model.baseline, + "judge": cfg.judge.model, + } + ) + if swap_both: + reversed_rendered = render_judge_inputs( + instructions, + completions_b, + completions_a, + system_prompt=resolved.system_prompt, + user_prompt_template=resolved.user_prompt_template, + truncate_input_chars=cfg.generation.truncate_judge_input_chars, + provide_explanation=cfg.judge.provide_explanation, + prompt_preset=resolved.preset_name, + strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging, + task=cfg.task, + ) + for idx, (instruction, ca, cb, judge_input) in enumerate( + zip( + instructions, + completions_b, + completions_a, + reversed_rendered, + strict=True, + ) + ): + rows.append( + { + "instruction": instruction, + "completion_A": ca, + "completion_B": cb, + "judge_completion": "Score A: 6\nScore B: 8", + "judge_input": judge_input.to_string(), + "instruction_index": f"idx-{idx}", + "model_A": cfg.model.baseline, + "model_B": cfg.model.name, + "judge": cfg.judge.model, + } + ) + run_dir.mkdir(parents=True, exist_ok=True) + dump_config(cfg, run_dir / "config.yaml") + pd.DataFrame(rows).to_csv(run_dir / "pair-annotations.csv", index=False) + write_run_metadata( + output_dir=run_dir, + entrypoint="tests", + run=cfg.model_dump(), + results={"n": len(rows)}, + input_payloads={ + "instruction_index": [row["instruction_index"] for row in rows] + }, + judge_system_prompt=resolved.system_prompt, + judge_user_prompt_template=resolved.user_prompt_template, + ) + + +def _write_legacy_gae_run(run_dir: Path) -> None: + run_dir.mkdir(parents=True, exist_ok=True) + args = { + "task": "alpaca-eval", + "model_A": "Dummy/gen-a", + "model_B": "Dummy/gen-b", + "judge_model": "Dummy/score A: 0 score B: 10", + "swap_mode": "fixed", + "provide_explanation": False, + "truncate_all_input_chars": 8192, + "engine_kwargs": {}, + } + (run_dir / "args-alpaca.json").write_text(json.dumps(args)) + cfg = RunConfig( + task=args["task"], + model={"name": args["model_A"], "baseline": args["model_B"]}, + judge={"model": args["judge_model"], "swap_mode": args["swap_mode"]}, + generation={"truncate_judge_input_chars": args["truncate_all_input_chars"]}, + ) + _write_gae_annotations(run_dir, cfg) + + +def _count_judge_inference_rows(store_root: Path, judge_model: str) -> int: + total = 0 + for db_path in store_root.rglob(INFERENCE_DB_NAME): + metadata = json.loads( + (db_path.parent / "metadata.json").read_text(encoding="utf-8") + ) + if metadata.get("model_spec") != judge_model: + continue + with SQLiteInferenceStore(db_path) as store: + total += len(store.query()) + return total + + +def test_discover_skips_elo_and_legacy_artifacts(tmp_path): + elo_dir = tmp_path / "elo-lmarena-100k-run" + elo_dir.mkdir() + (elo_dir / "results.json").write_text("{}") + + legacy_db = tmp_path / "cache" / "db" / "arena" / "judge.db" + legacy_db.parent.mkdir(parents=True) + legacy_db.write_text("sqlite") + + pass_cache = tmp_path / "tables" / "model_outputs" / "alpaca-eval.csv.zip" + pass_cache.parent.mkdir(parents=True) + pass_cache.write_text("zip") + + report = discover_sources([tmp_path]) + skipped_kinds = {item.kind for item in report.skipped} + assert ArtifactKind.META_EVAL_IDENTITY_DB in skipped_kinds + assert ArtifactKind.PASS_LEVEL_CACHE in skipped_kinds + assert ArtifactKind.ELO_RUN in skipped_kinds + + +def test_discover_classifies_standalone_completions_as_generation(tmp_path): + artifact = tmp_path / "completions.parquet" + artifact.write_bytes(b"not imported") + + report = discover_sources([artifact]) + + assert [item.kind for item in report.skipped] == [ArtifactKind.GENERATION_ARTIFACT] + + +def test_gae_current_config_backfill_and_idempotency(tmp_path): + cfg = _cfg_with_cache(tmp_path) + run_dir = tmp_path / "results" / "gae-run" + _write_gae_annotations(run_dir, cfg) + + store_root = tmp_path / "backfill-store" + first = backfill_sources([run_dir], store_root) + assert first.written == 2 + assert first.existing == 0 + + second = backfill_sources([run_dir], store_root) + assert second.written == 0 + assert second.existing == 2 + + db_path = next(store_root.rglob(INFERENCE_DB_NAME)) + with SQLiteInferenceStore(db_path) as store: + rows = store.query() + assert all(row["pushed_by"] == BACKFILL_PUSHED_BY for _, row in rows.iterrows()) + + +def test_legacy_gae_args_backfill_parity(tmp_path): + run_dir = tmp_path / "legacy-run" + _write_legacy_gae_run(run_dir) + store_root = tmp_path / "store" + report = backfill_sources([run_dir], store_root) + assert report.written == 2 + assert report.sources["gae"]["runs"] == 1 + + +def test_gae_swap_rows_backfill(tmp_path): + cfg = _cfg_with_cache( + tmp_path, + judge={"model": "Dummy/score A: 0 score B: 10", "swap_mode": "both"}, + ) + run_dir = tmp_path / "swap-run" + _write_gae_annotations(run_dir, cfg, swap_both=True) + + report = backfill_sources([run_dir], tmp_path / "store") + assert report.written == 4 + + +def test_gae_judge_input_mismatch_skipped(tmp_path): + cfg = _cfg_with_cache(tmp_path) + run_dir = tmp_path / "mismatch-run" + _write_gae_annotations(run_dir, cfg) + csv_path = run_dir / "pair-annotations.csv" + df = pd.read_csv(csv_path) + df.loc[0, "judge_input"] = "tampered" + df.to_csv(csv_path, index=False) + + report = backfill_sources([run_dir], tmp_path / "store") + assert report.written == 1 + assert report.skipped["judge_input_mismatch"] == 1 + + +def test_gae_missing_judge_output_is_not_backfilled(tmp_path): + cfg = _cfg_with_cache(tmp_path) + run_dir = tmp_path / "missing-output" + _write_gae_annotations(run_dir, cfg) + csv_path = run_dir / "pair-annotations.csv" + df = pd.read_csv(csv_path, keep_default_na=False) + df.loc[0, "judge_completion"] = "" + df.to_csv(csv_path, index=False) + + report = backfill_sources([run_dir], tmp_path / "store") + + assert report.written == 1 + assert report.skipped["judge_output_missing"] == 1 + + +def test_multiple_annotations_files_fail_closed(tmp_path): + cfg = _cfg_with_cache(tmp_path) + run_dir = tmp_path / "multiple-annotations" + _write_gae_annotations(run_dir, cfg) + (run_dir / "other-annotations.csv").write_bytes( + (run_dir / "pair-annotations.csv").read_bytes() + ) + + report = backfill_sources([run_dir], tmp_path / "store") + + assert report.written == 0 + assert report.skipped["source_extraction_failed"] == 1 + + +def test_conflicting_modern_configs_fail_closed(tmp_path): + cfg = _cfg_with_cache(tmp_path) + run_dir = tmp_path / "conflicting-configs" + _write_gae_annotations(run_dir, cfg) + conflicting = cfg.model_copy(deep=True) + conflicting.judge.model = "Dummy/different-judge" + dump_config(conflicting, run_dir / "config.yaml") + + report = backfill_sources([run_dir], tmp_path / "store") + + assert report.written == 0 + assert report.skipped["source_extraction_failed"] == 1 + + +def test_gae_ambiguous_orientation_is_skipped(tmp_path): + cfg = _cfg_with_cache(tmp_path) + run_dir = tmp_path / "ambiguous-orientation" + _write_gae_annotations(run_dir, cfg) + csv_path = run_dir / "pair-annotations.csv" + df = pd.read_csv(csv_path) + df.loc[0, ["model_A", "model_B"]] = ["other-a", "other-b"] + df.to_csv(csv_path, index=False) + + report = backfill_sources([run_dir], tmp_path / "store") + + assert report.written == 1 + assert report.skipped["battle_orientation_unverifiable"] == 1 + + +def test_gae_missing_judge_input_fail_closed(tmp_path): + cfg = _cfg_with_cache(tmp_path) + run_dir = tmp_path / "missing-input-run" + _write_gae_annotations(run_dir, cfg) + csv_path = run_dir / "pair-annotations.csv" + df = pd.read_csv(csv_path) + df = df.drop(columns=["judge_input"]) + df.to_csv(csv_path, index=False) + + report = backfill_sources([run_dir], tmp_path / "store") + assert report.written == 0 + assert report.skipped.get("unknown_judge_run", 0) == 1 + + +def test_local_engine_skipped_without_init(tmp_path, monkeypatch): + cfg = _cfg_with_cache( + tmp_path, + judge={"model": "VLLM/Qwen/Qwen2.5-0.5B-Instruct", "swap_mode": "fixed"}, + ) + run_dir = tmp_path / "vllm-run" + _write_gae_annotations(run_dir, cfg) + + def fail_init(*args, **kwargs): + raise AssertionError("VLLM should not initialize during backfill") + + monkeypatch.setattr( + "judgearena.models.ChatVLLM.__init__", + fail_init, + ) + + report = backfill_sources([run_dir], tmp_path / "store") + assert report.written == 0 + assert report.skipped["local_engine_unsupported"] == 2 + + +def test_mt_preset_backfill(tmp_path): + cfg = _cfg_with_cache( + tmp_path, + task="mt-bench", + model={"name": "Dummy/a", "baseline": "Dummy/b"}, + judge={"model": "Dummy/judge-output", "swap_mode": "fixed"}, + ) + run_dir = tmp_path / "mt-preset" + run_dir.mkdir(parents=True) + dump_config(cfg, run_dir / "config.yaml") + pd.DataFrame( + [ + { + "question_id": 1, + "category": "writing", + "turn": 1, + "model_A": "Dummy/a", + "model_B": "Dummy/b", + "judge": "Dummy/judge-output", + "prompt_name": "default-single", + "system_prompt": "system", + "user_prompt": "user body", + "judge_completion": "Score A: 8\nScore B: 6", + "swapped": False, + } + ] + ).to_csv(run_dir / "mt-annotations.csv", index=False) + + report = backfill_sources([run_dir], tmp_path / "store") + assert report.written == 1 + assert report.sources["mt_bench"]["runs"] == 1 + + +def test_mt_fastchat_g1_g2_backfill(tmp_path): + cfg = _cfg_with_cache( + tmp_path, + task="mt-bench", + model={"name": "Dummy/a", "baseline": "Dummy/b"}, + judge={"model": "Dummy/judge-output", "swap_mode": "both"}, + ) + run_dir = tmp_path / "mt-fastchat" + run_dir.mkdir(parents=True) + dump_config(cfg, run_dir / "config.yaml") + pd.DataFrame( + [ + { + "question_id": 1, + "category": "writing", + "turn": 1, + "model_A": "Dummy/a", + "model_B": "Dummy/b", + "judge": "Dummy/judge-output", + "prompt_name": "pair-v2", + "system_prompt": "system", + "g1_user_prompt": " g1 user ", + "g1_judgment": " [[A]] ", + "g2_user_prompt": "g2 user", + "g2_judgment": "[[B]]", + } + ] + ).to_csv(run_dir / "mt-annotations.csv", index=False) + + report = backfill_sources([run_dir], tmp_path / "store") + assert report.written == 2 + extraction = extract_mt_bench_rows(run_dir) + assert " g1 user " in extraction.rows[0].canonical_input + assert extraction.rows[0].output_text == " [[A]] " + + +def test_meta_eval_forward_and_swapped(tmp_path): + args = CliMetaEvalArgs( + judge_model="Dummy/meta-judge", + reference_arena="LMArena-140k", + prompt_mode="standard", + swap_mode="both", + ) + prompt_spec = resolve_prompt_mode(args.prompt_mode, provide_explanation=False) + forward_rendered = render_judge_inputs( + ["instruction"], + ["completion a"], + ["completion b"], + system_prompt=prompt_spec.system_prompt, + user_prompt_template=prompt_spec.user_prompt_template, + truncate_input_chars=args.truncate_judge_input_chars, + provide_explanation=False, + )[0].to_string() + swapped_rendered = render_judge_inputs( + ["instruction"], + ["completion b"], + ["completion a"], + system_prompt=prompt_spec.system_prompt, + user_prompt_template=prompt_spec.user_prompt_template, + truncate_input_chars=args.truncate_judge_input_chars, + provide_explanation=False, + )[0].to_string() + + run_dir = tmp_path / "meta-run" + run_dir.mkdir(parents=True) + (run_dir / "args.json").write_text(json.dumps(args.to_jsonable())) + pd.DataFrame( + [ + { + "question_id": "q1", + "benchmark": "arena", + "model_a": "m-a", + "model_b": "m-b", + "instruction": "instruction", + "completion_a": "completion a", + "completion_b": "completion b", + "presented_completion_a": "completion a", + "presented_completion_b": "completion b", + "judge_input": forward_rendered, + "judge_completion": "Score A: 8\nScore B: 6", + "orientation": "forward", + "presented_model_a": "m-a", + "presented_model_b": "m-b", + }, + { + "question_id": "q1", + "benchmark": "arena", + "model_a": "m-a", + "model_b": "m-b", + "instruction": "instruction", + "completion_a": "completion a", + "completion_b": "completion b", + "presented_completion_a": "completion b", + "presented_completion_b": "completion a", + "judge_input": swapped_rendered, + "judge_completion": "Score A: 6\nScore B: 8", + "orientation": "swapped", + "presented_model_a": "m-b", + "presented_model_b": "m-a", + }, + ] + ).to_parquet(run_dir / "annotations.parquet") + + report = backfill_sources([run_dir], tmp_path / "store") + assert report.written == 2 + assert report.rows_planned == 2 + task = meta_eval_cache_task(args.reference_arena) + assert any(task in str(path) for path in (tmp_path / "store").rglob("*")) + db_path = next((tmp_path / "store").rglob(INFERENCE_DB_NAME)) + with SQLiteInferenceStore(db_path) as store: + meta_rows = store.query_metadata() + assert all( + "source_run_id" in json.loads(row["metadata_json"]) + for _, row in meta_rows.iterrows() + ) + assert all( + "source_run_folder" not in json.loads(row["metadata_json"]) + for _, row in meta_rows.iterrows() + ) + + +def test_conflicting_outputs_skipped(tmp_path): + cfg = _cfg_with_cache(tmp_path) + run_a = tmp_path / "run-a" + run_b = tmp_path / "run-b" + _write_gae_annotations(run_a, cfg) + _write_gae_annotations(run_b, cfg) + df = pd.read_csv(run_b / "pair-annotations.csv") + df.loc[0, "judge_completion"] = "Score A: 1\nScore B: 9" + df.to_csv(run_b / "pair-annotations.csv", index=False) + + report = backfill_sources([run_a, run_b], tmp_path / "store") + assert report.skipped.get("conflicting_outputs", 0) == 2 + assert _count_judge_inference_rows(tmp_path / "store", cfg.judge.model) <= 2 + + +def test_identical_outputs_preserve_all_run_metadata(tmp_path): + cfg = _cfg_with_cache(tmp_path) + run_a = tmp_path / "run-a" + run_b = tmp_path / "run-b" + _write_gae_annotations(run_a, cfg) + _write_gae_annotations(run_b, cfg) + + store_root = tmp_path / "store" + report = backfill_sources([run_a, run_b], store_root) + + assert report.written == 2 + db_path = next(store_root.rglob(INFERENCE_DB_NAME)) + with SQLiteInferenceStore(db_path) as store: + metadata = store.query_metadata() + source_ids = { + json.loads(value)["source_run_id"] for value in metadata["metadata_json"] + } + assert source_ids == {"run-a", "run-b"} + assert len(metadata) == 4 + + +def test_dry_run_writes_nothing(tmp_path): + cfg = _cfg_with_cache(tmp_path) + run_dir = tmp_path / "dry-run" + _write_gae_annotations(run_dir, cfg) + report = backfill_sources([run_dir], tmp_path / "store", dry_run=True) + assert report.written == 2 + assert _count_judge_inference_rows(tmp_path / "store", cfg.judge.model) == 0 + + +def test_backfill_cli(tmp_path): + cfg = _cfg_with_cache(tmp_path) + run_dir = tmp_path / "cli-run" + _write_gae_annotations(run_dir, cfg) + report_path = tmp_path / "report.json" + + cache_sync.main( + [ + "backfill", + str(run_dir), + "--store_root", + str(tmp_path / "store"), + "--report", + str(report_path), + ] + ) + + payload = json.loads(report_path.read_text(encoding="utf-8")) + assert payload["written"] == 2 + assert _count_judge_inference_rows(tmp_path / "store", cfg.judge.model) == 2 + + +@pytest.fixture +def mock_gae_inputs(monkeypatch): + monkeypatch.setattr( + gae, + "load_instructions", + lambda dataset, n_instructions=None: _synthetic_instructions( + n_instructions or 2 + ), + ) + monkeypatch.setattr(gae, "try_load_dataset_completions", lambda *args: None) + + +def test_gae_live_run_backfill_reuses_cells(mock_gae_inputs, tmp_path): + cfg = _cfg_with_cache(tmp_path) + gae.main(cfg) + run_dir = next((tmp_path / "results").iterdir()) + + live_count = _count_judge_inference_rows(tmp_path / "live-cache", cfg.judge.model) + backfill_store = tmp_path / "backfill-store" + report = backfill_sources([run_dir], backfill_store) + assert report.written == live_count + assert _count_judge_inference_rows(backfill_store, cfg.judge.model) == live_count + + +def test_write_report_roundtrip(tmp_path): + report = BackfillReport( + written=3, + existing=1, + rows_planned=4, + skipped={"judge_input_mismatch": 2}, + ) + path = tmp_path / "nested" / "report.json" + write_report(report, path) + payload = json.loads(path.read_text(encoding="utf-8")) + assert payload["written"] == 3 + assert payload["rows_planned"] == 4 + assert payload["skipped"]["judge_input_mismatch"] == 2 + + +def test_gae_orientation_inference_reversed_rows(): + cfg = RunConfig( + task="alpaca-eval", + model={"name": "Dummy/gen-a", "baseline": "Dummy/gen-b"}, + judge={"model": "Dummy/judge"}, + ) + direct = pd.Series({"model_A": "Dummy/gen-a", "model_B": "Dummy/gen-b"}) + reversed_row = pd.Series({"model_A": "Dummy/gen-b", "model_B": "Dummy/gen-a"}) + assert _infer_gae_orientation(direct, cfg=cfg) == "direct" + assert _infer_gae_orientation(reversed_row, cfg=cfg) == "reversed" + + +def test_source_extraction_failure_continues_other_runs(tmp_path): + cfg = _cfg_with_cache(tmp_path) + good_run = tmp_path / "good-run" + bad_run = tmp_path / "bad-run" + _write_gae_annotations(good_run, cfg) + bad_run.mkdir() + pd.DataFrame( + { + "instruction": ["instruction 0"], + "completion_A": ["a"], + "completion_B": ["b"], + "judge_input": ["prompt"], + "judge_completion": ["Score A: 1\nScore B: 0"], + } + ).to_csv(bad_run / "pair-annotations.csv", index=False) + + report = backfill_sources([bad_run, good_run], tmp_path / "store") + assert report.skipped["source_extraction_failed"] == 1 + assert report.written == 2 + + +def test_unknown_annotation_csv_skipped_not_migrated(tmp_path): + run_dir = tmp_path / "custom-run" + run_dir.mkdir() + pd.DataFrame([{"foo": 1, "bar": 2}]).to_csv( + run_dir / "custom-annotations.csv", index=False + ) + report = discover_sources([run_dir]) + assert report.migratable_runs == [] + assert any(item.kind == ArtifactKind.UNKNOWN for item in report.skipped) + + +def test_conflicting_existing_output_skips_metadata(tmp_path): + cfg = _cfg_with_cache(tmp_path) + run_dir = tmp_path / "gae-run" + _write_gae_annotations(run_dir, cfg) + + store_root = tmp_path / "store" + first = backfill_sources([run_dir], store_root) + assert first.written == 2 + + df = pd.read_csv(run_dir / "pair-annotations.csv") + df.loc[0, "judge_completion"] = "Score A: 1\nScore B: 9" + df.to_csv(run_dir / "pair-annotations.csv", index=False) + + second = backfill_sources([run_dir], store_root) + assert second.skipped.get("conflicting_existing_output", 0) == 1 + db_path = next(store_root.rglob(INFERENCE_DB_NAME)) + with SQLiteInferenceStore(db_path) as store: + rows = store.query() + assert rows.iloc[0]["output_text"] == "Score A: 8\nScore B: 6" + + +def test_legacy_multiple_args_files_fail_closed(tmp_path): + run_dir = tmp_path / "ambiguous-legacy" + run_dir.mkdir() + (run_dir / "args-a.json").write_text(json.dumps({"task": "alpaca-eval"})) + (run_dir / "args-b.json").write_text(json.dumps({"task": "alpaca-eval"})) + pd.DataFrame( + { + "instruction": ["x"], + "completion_A": ["a"], + "completion_B": ["b"], + "judge_input": ["prompt"], + "judge_completion": ["Score A: 1\nScore B: 0"], + } + ).to_csv(run_dir / "pair-annotations.csv", index=False) + + report = backfill_sources([run_dir], tmp_path / "store") + assert report.skipped["source_extraction_failed"] == 1 + assert report.written == 0 + + +def test_legacy_truncate_all_input_chars_used_for_judge(tmp_path): + run_dir = tmp_path / "legacy-truncate" + _write_legacy_gae_run(run_dir) + extraction = extract_gae_rows(run_dir) + assert extraction.rows + assert extraction.skipped.get("judge_input_mismatch", 0) == 0 + + +def test_nested_legacy_db_does_not_block_parent_run_discovery(tmp_path): + cfg = _cfg_with_cache(tmp_path) + run_dir = tmp_path / "results" / "nested-run" + _write_gae_annotations(run_dir, cfg) + legacy_dir = tmp_path / "results" / "cache" / "nested" + legacy_dir.mkdir(parents=True) + (legacy_dir / "judgements.db").write_text("sqlite") + + report = discover_sources([tmp_path / "results"]) + migratable = {item.run_dir.name for item in report.migratable_runs} + assert "nested-run" in migratable + skipped_kinds = {item.kind for item in report.skipped} + assert ArtifactKind.LEGACY_CACHE_CELL in skipped_kinds + + backfill_report = backfill_sources([tmp_path / "results"], tmp_path / "store") + assert backfill_report.written == 2 + + +def test_mt_swapped_string_and_nan_prompts(tmp_path): + cfg = _cfg_with_cache( + tmp_path, + task="mt-bench", + model={"name": "Dummy/a", "baseline": "Dummy/b"}, + judge={"model": "Dummy/judge-output", "swap_mode": "fixed"}, + ) + run_dir = tmp_path / "mt-swapped-string" + run_dir.mkdir(parents=True) + dump_config(cfg, run_dir / "config.yaml") + pd.DataFrame( + [ + { + "question_id": 1, + "category": "writing", + "turn": 1, + "model_A": "Dummy/a", + "model_B": "Dummy/b", + "judge": "Dummy/judge-output", + "prompt_name": "default-single", + "system_prompt": float("nan"), + "user_prompt": "user body", + "judge_completion": "Score A: 8\nScore B: 6", + "swapped": "true", + } + ] + ).to_csv(run_dir / "mt-annotations.csv", index=False) + + report = backfill_sources([run_dir], tmp_path / "store") + assert report.written == 1 + db_path = next((tmp_path / "store").rglob(INFERENCE_DB_NAME)) + with SQLiteInferenceStore(db_path) as store: + row = store.query().iloc[0] + assert "nan" not in row["input_text"].lower() + + +def test_dry_run_does_not_touch_existing_db(tmp_path): + cfg = _cfg_with_cache(tmp_path) + run_dir = tmp_path / "dry-run-existing" + _write_gae_annotations(run_dir, cfg) + store_root = tmp_path / "store" + backfill_sources([run_dir], store_root, dry_run=False) + db_path = next(store_root.rglob(INFERENCE_DB_NAME)) + before = db_path.stat().st_mtime_ns + wal_path = Path(f"{db_path}-wal") + wal_before_exists = wal_path.exists() + + report = backfill_sources([run_dir], store_root, dry_run=True) + assert report.written == 0 + assert report.existing == 2 + assert db_path.stat().st_mtime_ns == before + assert wal_path.exists() == wal_before_exists + + +def test_cell_integrity_error_is_reported(tmp_path, monkeypatch): + cfg = _cfg_with_cache(tmp_path) + run_dir = tmp_path / "integrity-run" + _write_gae_annotations(run_dir, cfg) + + def fail_metadata(*args, **kwargs): + raise ValueError("metadata mismatch") + + monkeypatch.setattr( + "judgearena.cache_backfill.write_store_metadata", + fail_metadata, + ) + report = backfill_sources([run_dir], tmp_path / "store") + assert report.skipped.get("cell_integrity_error", 0) == 2 + assert report.written == 0 diff --git a/tests/test_cache_sync.py b/tests/test_cache_sync.py new file mode 100644 index 0000000..601944c --- /dev/null +++ b/tests/test_cache_sync.py @@ -0,0 +1,160 @@ +import getpass +import json +from pathlib import Path + +import pandas as pd +import pytest + +from judgearena import cache_sync +from judgearena.store_sqlite import ( + INFERENCE_DB_NAME, + SQLiteInferenceStore, + descriptor_hash, + store_folder, + write_store_metadata, +) + +REPO_ID = "org/cache" +CELL_CONFIG = {"task": "arena", "model_spec": "VLLM/Qwen/judge"} +CELL_CONFIG_HASH = descriptor_hash(CELL_CONFIG) +MODEL_SPEC = "VLLM/Qwen/judge" +PATH_IN_REPO = ( + f"inference/arena/VLLM/Qwen%2Fjudge/{CELL_CONFIG_HASH}/{INFERENCE_DB_NAME}" +) +METADATA_IN_REPO = f"inference/arena/VLLM/Qwen%2Fjudge/{CELL_CONFIG_HASH}/metadata.json" + + +def _local_cell_db(tmp_path) -> Path: + cell_dir = store_folder(tmp_path, "arena", MODEL_SPEC, CELL_CONFIG_HASH) + return cell_dir / INFERENCE_DB_NAME + + +def _write_inference(path: Path) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + with SQLiteInferenceStore(path) as store: + store.save_outputs( + pd.DataFrame( + { + "input_hash": ["local"], + "input_text": ["input-local"], + "output_text": ["L"], + } + ), + pushed_by="test", + ) + + +def test_fetch_requires_filter(tmp_path, capsys): + with pytest.raises(SystemExit) as exc: + cache_sync.main(["fetch", "--store_root", str(tmp_path)]) + assert exc.value.code == 1 + assert "requires a path filter" in capsys.readouterr().err + + +def test_fetch_rejects_invalid_filter_gaps(tmp_path, capsys): + with pytest.raises(SystemExit) as exc: + cache_sync.main( + [ + "fetch", + "--store_root", + str(tmp_path), + "--provider", + "VLLM", + ] + ) + assert exc.value.code == 1 + assert "requires --task" in capsys.readouterr().err + + +def test_fetch_bootstraps_filtered_remote_cells(fake_hub, tmp_path): + remote = tmp_path / "remote.db" + _write_inference(remote) + fake_hub.files[PATH_IN_REPO] = remote.read_bytes() + fake_hub.files[METADATA_IN_REPO] = json.dumps(CELL_CONFIG).encode("utf-8") + fake_hub.head = "initial" + + store_root = tmp_path / "store" + cache_sync.main( + [ + "fetch", + "--store_root", + str(store_root), + "--cache_hf_repo", + REPO_ID, + "--task", + "arena", + ] + ) + local_db = _local_cell_db(store_root) + assert local_db.exists() + assert (local_db.parent / "metadata.json").exists() + + +def test_push_uploads_local_cells(fake_hub, tmp_path): + local_db = _local_cell_db(tmp_path) + _write_inference(local_db) + write_store_metadata(local_db.parent, CELL_CONFIG) + fake_hub.head = "initial" + + cache_sync.main( + [ + "push", + "--store_root", + str(tmp_path), + "--cache_hf_repo", + REPO_ID, + "--task", + "arena", + ] + ) + assert PATH_IN_REPO in fake_hub.files + assert METADATA_IN_REPO in fake_hub.files + assert fake_hub.commit_calls == 1 + + +def test_push_defaults_pushed_by_to_current_user(fake_hub, tmp_path, monkeypatch): + local_db = _local_cell_db(tmp_path) + _write_inference(local_db) + write_store_metadata(local_db.parent, CELL_CONFIG) + fake_hub.head = "initial" + observed: list[str] = [] + + def capture_push(*args, **kwargs): + observed.append(kwargs.get("pushed_by", args[3] if len(args) > 3 else None)) + + monkeypatch.setattr(cache_sync, "push_cells", capture_push) + monkeypatch.setattr(getpass, "getuser", lambda: "unit-test-user") + + cache_sync.main( + [ + "push", + "--store_root", + str(tmp_path), + "--cache_hf_repo", + REPO_ID, + "--task", + "arena", + ] + ) + assert observed == ["unit-test-user"] + + +def test_push_create_pr(fake_hub, tmp_path): + local_db = _local_cell_db(tmp_path) + _write_inference(local_db) + write_store_metadata(local_db.parent, CELL_CONFIG) + fake_hub.head = "initial" + + cache_sync.main( + [ + "push", + "--store_root", + str(tmp_path), + "--cache_hf_repo", + REPO_ID, + "--task", + "arena", + "--create_pr", + ] + ) + assert PATH_IN_REPO not in fake_hub.files From 75c769999ee5432db5b1726abef4bd130402c979 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Mon, 20 Jul 2026 18:53:23 +0200 Subject: [PATCH 12/13] docs(cache): document unified cache workflows Describe cache modes, hosted-provider limits, offline Hugging Face synchronization, SQLite constraints, and conservative migration examples. Includes-AI-Code: true --- README.md | 243 ++++++++++++++++++++++++++++++++++++-- configs/alpaca_eval.yaml | 8 ++ configs/elo_comparia.yaml | 8 ++ 3 files changed, 247 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index d651900..f7d116b 100644 --- a/README.md +++ b/README.md @@ -53,9 +53,9 @@ judgearena \ ``` **What happens here?** -- Use completions available for `gpt4_1106_preview` in Alpaca-Eval task -- Generates completions for `model_B` if not already cached on `vLLM` -- Compares two models using `deepseek-chat-v3.1` which the cheapest option available on `OpenRouter` +- Loads precomputed completions for `gpt4_1106_preview` from the Alpaca-Eval dataset +- Generates completions for `model_B` when they are not already present locally or in an optional inference cache (see [Unified inference cache](#unified-inference-cache)) +- Compares the two models with `deepseek-chat-v3.1`, a low-cost option on OpenRouter It will then display the results of the battles: @@ -106,6 +106,209 @@ model defaults. Each run also writes the fully-resolved config to `--config_path` to reproduce the run. An ELO example is in [`configs/elo_comparia.yaml`](configs/elo_comparia.yaml). +### Unified inference cache + +Caching is **disabled by default**. Set `cache.store_root` (CLI: `--store_root` or +`--cache.store_root`) to enable the shared SQLite store. Generation, judging, +meta-eval, MT-Bench, and ELO all route through the same per-inference boundary. + +Each **configuration cell** is scoped by benchmark task, model spec, and a +descriptor hash derived from resolved engine/sampling settings. Within a cell, +rows are keyed by a SHA-256 hash of the **canonical rendered input** (for judges, +the fully rendered prompt including instruction, truncated completions, and +template). Changing any of those fields is a cache miss rather than a stale hit. + +**Store layout** (local paths mirror the default Hugging Face dataset +`judge-arena/judge-arena-cache`): + +```text +{store_root}/inference/{task}/{provider}/{percent-encoded-model}/{descriptor_hash}/ + metadata.json # descriptor for this cell + inference.db # WAL-mode SQLite (`inference` + `inference_metadata` tables) +``` + +**Modes** (`--cache_mode` / `--cache.cache_mode`): + +- `use` (default): read hits, infer misses, insert new rows (`INSERT OR IGNORE`). +- `off`: always infer, even when `store_root` is set. +- `refresh`: replace existing rows for the same input hash (`INSERT OR REPLACE`). + +**Run flags** — nested and flat aliases are equivalent (`--cache.store_root` ≡ +`--store_root`, `--cache.cache_fetch` ≡ `--cache_fetch`, and so on). Hugging Face +sync is **opt-in**: `cache_fetch`, `cache_push`, and `cache_create_pr` default to +`false`. When enabled, fetch pulls remote cells lazily as each local cell opens; +push uploads only cells written during the run and **only after a successful exit** +(failures and interrupts leave local SQLite updates but skip HF push). Use +`--no-cache_fetch` on compute nodes that cannot reach the Hub. +Main-pipeline sync is best-effort and logs failures without discarding local +results; the standalone `judgearena-cache fetch` and `push` commands fail loudly. + +**Stochastic sampling.** Identical canonical inputs in one batch are deduplicated to +a single inference call. Across runs, the first stored output wins locally; Hub +merges keep the newest row by `pushed_at`. For stochastic judges (for example +meta-eval PairScore at temperature `0.5`), set an explicit `--judge.seed` when you +need reproducible resampling, or use `--cache_mode refresh` / a separate +`store_root`. `--run.seed` drives battle sampling and shuffles; it is not +automatically forwarded as a backend seed. + +**Descriptors and provenance.** + +- **Local engines (vLLM, LlamaCpp):** descriptors capture fully resolved engine + settings and package versions (`vllm_version`, `llama_cpp_python_version`, …). +- **Hosted APIs (OpenRouter, OpenAI, Together):** only **explicitly set** client + sampling fields are hashed. Provider-side defaults and model aliases are + inherently best-effort — the client library version is recorded in row + provenance (`langchain_openai_version`, `hosted_adapter_version`, …) but is **not** + part of the descriptor hash. +- **Secrets and unsafe settings:** API keys, headers, and other sensitive + constructor fields are redacted. Non-JSON-safe or un-normalizable + `engine_kwargs` yield no descriptor; inference still runs, uncached, with a + warning. + +**Single-host SQLite.** Open cells use WAL mode. Do not share one open +`inference.db` across NFS or multiple compute nodes. Synchronize **closed** cells +via Hugging Face (`judgearena-cache fetch` / `push`) or a local store copy. + +**Interrupted runs.** Partial rows are persisted locally as they are written. Re-run +to hit local cache, use `refresh` to regenerate, or `judgearena-cache push` after +you verify the store. + +**ELO and meta-eval.** Cached judge completions are stored as raw outputs; ELO +preferences, calibration, and leaderboard statistics are always recomputed from +those outputs (and current run settings) rather than replaying prior aggregates. +Meta-eval uses task namespace `meta-eval-{reference_arena}` under `inference/`. + +#### Examples + +Local cache only (no network): + +```bash +judgearena \ + --task alpaca-eval \ + --model.name gpt4_1106_preview \ + --model.baseline VLLM/utter-project/EuroLLM-9B \ + --judge.model OpenRouter/deepseek/deepseek-chat-v3.1 \ + --generation.n_instructions 10 \ + --store_root ~/judgearena-data/cache +``` + +Run with explicit Hugging Face fetch and push: + +```bash +judgearena \ + --config_path configs/alpaca_eval.yaml \ + --cache.store_root ~/judgearena-data/cache \ + --cache.cache_hf_repo judge-arena/judge-arena-cache \ + --cache.cache_fetch \ + --cache.cache_push \ + --cache.pushed_by "$USER" +``` + +Offline workflow (login node → compute → login node): + +```bash +# login node: prefetch cells for the benchmark +judgearena-cache fetch \ + --store_root ~/judgearena-data/cache \ + --cache_hf_repo judge-arena/judge-arena-cache \ + --task alpaca-eval + +# compute node: local store only +judgearena \ + --config_path configs/alpaca_eval.yaml \ + --store_root ~/judgearena-data/cache \ + --no-cache_fetch + +# login node: merge and upload +judgearena-cache push \ + --store_root ~/judgearena-data/cache \ + --cache_hf_repo judge-arena/judge-arena-cache \ + --task alpaca-eval \ + --pushed_by "$USER" +``` + +If that path is on NFS, stage the **closed** store to node-local storage before +opening it on a compute node, then stage it back after the job. Never open the +same cell concurrently from multiple hosts. + +Filtered sync (`fetch` requires `--prefix` or at least `--task`; `push` without +filters uploads every local cell): + +```bash +judgearena-cache fetch \ + --store_root ~/judgearena-data/cache \ + --prefix inference/alpaca-eval/OpenRouter + +judgearena-cache push \ + --store_root ~/judgearena-data/cache \ + --task alpaca-eval \ + --provider OpenRouter \ + --model deepseek/deepseek-chat-v3.1 \ + --create_pr \ + --pushed_by "$USER" +``` + +Open a pull request from the main CLI instead: + +```bash +judgearena \ + --config_path configs/alpaca_eval.yaml \ + --store_root ~/judgearena-data/cache \ + --cache_fetch \ + --cache_push \ + --cache_create_pr +``` + +#### `judgearena-cache backfill` + +`judgearena-cache backfill` imports **verified hosted judge rows** from saved +generate-and-evaluate, MT-Bench, and meta-eval run folders into the unified store. +For generate-and-evaluate and meta-eval, it reconstructs the rendered judge input +from saved configuration and checks it against stored `judge_input` text. +MT-Bench stores the exact system and rendered user messages used for each pass, +which are canonicalized directly. +**Insert-only, no network.** + +Imported: hosted judge outputs (`OpenRouter`, `ChatOpenAI`, `OpenAI`, `Together`, +and `Dummy` in tests). Model-generation outputs are never backfilled because old +artifacts cannot prove whether they came from inference or precomputed datasets. + +Skipped: generation-only artifacts, ELO runs, local engines (vLLM / LlamaCpp), +legacy pass-level caches, identity-keyed meta-eval DBs under `cache/db`, legacy +`judgements.db` / `completions.db` cells, rows whose recomputed input does not +match saved `judge_input`, and conflicting outputs for the same key. + +Dry-run first, then write: + +```bash +judgearena-cache backfill results/ \ + --store_root ~/judgearena-data/cache \ + --dry_run \ + --report ~/judgearena-data/cache/backfill-report.json + +judgearena-cache backfill results/alpaca-eval-2026-04-01/ \ + --store_root ~/judgearena-data/cache \ + --report ~/judgearena-data/cache/backfill-report.json +``` + +The JSON report lists `written`, `existing`, `rows_planned`, `runs_processed`, +and per-reason `skipped` counts. + +Optional YAML (commented so defaults stay network-free): + +```yaml +# cache: +# store_root: ~/judgearena-data/cache +# cache_mode: use +# cache_hf_repo: judge-arena/judge-arena-cache +# # cache_fetch: true +# # cache_push: true +# pushed_by: your-user +``` + +Standalone cache CLI: `judgearena-cache {fetch,push,backfill}` (see +`judgearena-cache --help`). + ### Length and Token Parameters The evaluation scripts expose four different length controls with different roles: @@ -330,13 +533,13 @@ back to the original model ordering. Cost totals include both passes. Results are written under `[run.result_folder]/meta-eval-*` as `args.json`, `annotations.parquet`, `results.json`, `summary.csv`, logs, and -`run-metadata.v1.json`. Judge annotations are cached per battle in -`$JUDGEARENA_DATA/cache/db/{benchmark}/{judge}.db`. This temporary SQLite WAL -cache should only be used by a single host; do not share the same database -concurrently across NFS-mounted compute nodes. Annotation artifacts include -character-based token estimates. Equivalent OpenRouter cost is reported only -when `[data_root]/cache/openrouter_pricing.json` already contains the judge -model; meta-eval never fetches pricing from compute nodes. +`run-metadata.v1.json`. With `--store_root`, judge calls use the unified cache +under `{store_root}/inference/meta-eval-{reference_arena}/…` (content-addressed +by rendered judge input; see [Unified inference cache](#unified-inference-cache)). +Annotation artifacts include character-based token estimates. Equivalent +OpenRouter cost is reported only when +`[data_root]/cache/openrouter_pricing.json` already contains the judge model; +meta-eval never fetches pricing from compute nodes. ## 📈 Estimating ELO Ratings @@ -355,6 +558,10 @@ judgearena \ --generation.n_instructions 200 ``` +With `--store_root`, generation and judge calls reuse cached rows; ELO win rates, +calibration, and Bradley-Terry scores are always recomputed from the cached judge +outputs and the current run configuration (see [Unified inference cache](#unified-inference-cache)). + ### Key options | Flag | Default | Description | @@ -410,12 +617,24 @@ Win rate: 60.25% ### Offline Setup (Slurm/Air-Gapped Environments) -Pre-download all datasets before running jobs: +Pre-download datasets and models on a login node before submitting jobs: ```bash -python -c "from judgearena.utils import download_all; download_all()" # Download all datasets (optional) +python -c "from judgearena.utils import download_all; download_all()" # optional: all datasets +hf download --local-dir ~/models/my-model # model weights ``` +Prefetch shared inference cells when you use `--store_root`: + +```bash +judgearena-cache fetch \ + --store_root ~/judgearena-data/cache \ + --task alpaca-eval +``` + +On compute nodes, pass `--store_root` with `--no-cache_fetch`, then push from the +login node after the job finishes (see [Unified inference cache](#unified-inference-cache)). + Datasets are stored in: - `$JUDGEARENA_DATA` if set; otherwise `$OPENJURY_DATA` if set (legacy) - `~/judgearena-data/` if neither variable is set diff --git a/configs/alpaca_eval.yaml b/configs/alpaca_eval.yaml index 149e765..f50bfad 100644 --- a/configs/alpaca_eval.yaml +++ b/configs/alpaca_eval.yaml @@ -10,3 +10,11 @@ judge: model: OpenRouter/google/gemma-4-31b-it generation: n_instructions: 10 +# Optional unified inference cache (disabled unless store_root is set): +# cache: +# store_root: ~/judgearena-data/cache +# cache_mode: use +# cache_hf_repo: judge-arena/judge-arena-cache +# # cache_fetch: true # prefetch on login node via judgearena-cache instead +# # cache_push: true # or judgearena-cache push after the run +# pushed_by: your-user diff --git a/configs/elo_comparia.yaml b/configs/elo_comparia.yaml index c21730b..37e541f 100644 --- a/configs/elo_comparia.yaml +++ b/configs/elo_comparia.yaml @@ -14,3 +14,11 @@ elo: # Restrict the arena to these language codes. Defaults to English here; add # more to evaluate multilingually, e.g. languages: ["en", "fr", "de"]. languages: ["en"] +# Optional unified inference cache (disabled unless store_root is set): +# cache: +# store_root: ~/judgearena-data/cache +# cache_mode: use +# cache_hf_repo: judge-arena/judge-arena-cache +# # cache_fetch: true +# # cache_push: true +# pushed_by: your-user From 8dcaf70ece1a9018338122e5a9ebf766e3024213 Mon Sep 17 00:00:00 2001 From: Erlis Lushtaku <59629249+ErlisLushtaku@users.noreply.github.com> Date: Fri, 7 Aug 2026 01:22:34 +0200 Subject: [PATCH 13/13] refactor(cache): trim unified cache scope Keep the operational fetch/push CLI while removing migration tooling, documentation, and redundant pipeline test plumbing to make the cache stack reviewable. --- README.md | 193 +---- configs/alpaca_eval.yaml | 8 - configs/elo_comparia.yaml | 8 - judgearena/cache_backfill.py | 324 ------- judgearena/cache_backfill_common.py | 59 -- judgearena/cache_backfill_config.py | 166 ---- judgearena/cache_backfill_discovery.py | 340 -------- judgearena/cache_backfill_sources.py | 368 -------- judgearena/cache_sync.py | 56 +- judgearena/config.py | 2 +- judgearena/generate_and_evaluate.py | 10 +- judgearena/log.py | 21 +- judgearena/pairwise_baselines.py | 34 - tests/test_cache_backfill.py | 792 ------------------ tests/test_cache_sync.py | 138 +-- tests/test_cli.py | 67 +- tests/test_config.py | 41 - tests/test_estimate_elo_cache_threading.py | 330 -------- tests/test_estimate_elo_ratings.py | 11 +- tests/test_evaluate_cache_threading.py | 107 --- ...t_generate_and_evaluate_cache_threading.py | 33 - tests/test_generate_cache_threading.py | 174 ---- tests/test_inference_cache.py | 182 +--- tests/test_logging.py | 24 - tests/test_meta_eval_cache_threading.py | 581 ------------- tests/test_model_adapters.py | 272 +----- tests/test_mt_bench_downloads.py | 124 +-- tests/test_mt_bench_fastchat_compat.py | 34 - .../test_mt_bench_pairwise_cache_threading.py | 159 ---- tests/test_mt_bench_preset_judging.py | 46 - tests/test_no_legacy_runtime_cache.py | 6 +- tests/test_store_sqlite.py | 45 - tests/test_store_sync.py | 229 ----- 33 files changed, 153 insertions(+), 4831 deletions(-) delete mode 100644 judgearena/cache_backfill.py delete mode 100644 judgearena/cache_backfill_common.py delete mode 100644 judgearena/cache_backfill_config.py delete mode 100644 judgearena/cache_backfill_discovery.py delete mode 100644 judgearena/cache_backfill_sources.py delete mode 100644 judgearena/pairwise_baselines.py delete mode 100644 tests/test_cache_backfill.py delete mode 100644 tests/test_estimate_elo_cache_threading.py delete mode 100644 tests/test_evaluate_cache_threading.py delete mode 100644 tests/test_generate_cache_threading.py delete mode 100644 tests/test_meta_eval_cache_threading.py delete mode 100644 tests/test_mt_bench_pairwise_cache_threading.py diff --git a/README.md b/README.md index cf5e895..cf0c8cc 100644 --- a/README.md +++ b/README.md @@ -108,206 +108,71 @@ model defaults. Each run also writes the fully-resolved config to ### Unified inference cache -Caching is **disabled by default**. Set `cache.store_root` (CLI: `--store_root` or -`--cache.store_root`) to enable the shared SQLite store. Generation, judging, -meta-eval, MT-Bench, and ELO all route through the same per-inference boundary. - -Each **configuration cell** is scoped by benchmark task, model spec, and a -descriptor hash derived from resolved engine/sampling settings. Within a cell, -rows are keyed by a SHA-256 hash of the **canonical rendered input** (for judges, -the fully rendered prompt including instruction, truncated completions, and -template). Changing any of those fields is a cache miss rather than a stale hit. - -**Store layout** (local paths mirror the default Hugging Face dataset -`judge-arena/judge-arena-cache`): +Caching is disabled by default. Set `--store_root` (or +`cache.store_root` in YAML) to cache generation and judge calls from all +benchmarks at the shared `do_inference` boundary. Rows are content-addressed by +the canonical rendered input and stored in model/configuration-specific cells: ```text {store_root}/inference/{task}/{provider}/{percent-encoded-model}/{descriptor_hash}/ - metadata.json # descriptor for this cell - inference.db # WAL-mode SQLite (`inference` + `inference_metadata` tables) + metadata.json + inference.db ``` -**Modes** (`--cache_mode` / `--cache.cache_mode`): +Cache modes: -- `use` (default): read hits, infer misses, insert new rows (`INSERT OR IGNORE`). +- `use` (default): reuse hits and insert misses. - `off`: always infer, even when `store_root` is set. -- `refresh`: replace existing rows for the same input hash (`INSERT OR REPLACE`). - -**Run flags** — nested and flat aliases are equivalent (`--cache.store_root` ≡ -`--store_root`, `--cache.cache_fetch` ≡ `--cache_fetch`, and so on). Hugging Face -sync is **opt-in**: `cache_fetch`, `cache_push`, and `cache_create_pr` default to -`false`. When enabled, fetch pulls remote cells lazily as each local cell opens; -push uploads only cells written during the run and **only after a successful exit** -(failures and interrupts leave local SQLite updates but skip HF push). Use -`--no-cache_fetch` on compute nodes that cannot reach the Hub. -Main-pipeline sync is best-effort and logs failures without discarding local -results; the standalone `judgearena-cache fetch` and `push` commands fail loudly. - -**Stochastic sampling.** Identical canonical inputs in one batch are deduplicated to -a single inference call. Across runs, the first stored output wins locally; Hub -merges keep the newest row by `pushed_at`. For stochastic judges (for example -meta-eval PairScore at temperature `0.5`), set an explicit `--judge.seed` when you -need reproducible resampling, or use `--cache_mode refresh` / a separate -`store_root`. `--run.seed` drives battle sampling and shuffles; it is not -automatically forwarded as a backend seed. - -**Descriptors and provenance.** - -- **Local engines (vLLM, LlamaCpp):** descriptors capture fully resolved engine - settings and package versions (`vllm_version`, `llama_cpp_python_version`, …). -- **Hosted APIs (OpenRouter, OpenAI, Together):** only **explicitly set** client - sampling fields are hashed. Provider-side defaults and model aliases are - inherently best-effort — the client library version is recorded in row - provenance (`langchain_openai_version`, `hosted_adapter_version`, …) but is **not** - part of the descriptor hash. -- **Secrets and unsafe settings:** API keys, headers, and other sensitive - constructor fields are redacted. Non-JSON-safe or un-normalizable - `engine_kwargs` yield no descriptor; inference still runs, uncached, with a - warning. - -**Single-host SQLite.** Open cells use WAL mode. Do not share one open -`inference.db` across NFS or multiple compute nodes. Synchronize **closed** cells -via Hugging Face (`judgearena-cache fetch` / `push`) or a local store copy. - -**Interrupted runs.** Partial rows are persisted locally as they are written. Re-run -to hit local cache, use `refresh` to regenerate, or `judgearena-cache push` after -you verify the store. - -**ELO and meta-eval.** Cached judge completions are stored as raw outputs; ELO -preferences, calibration, and leaderboard statistics are always recomputed from -those outputs (and current run settings) rather than replaying prior aggregates. -Meta-eval uses task namespace `meta-eval-{reference_arena}` under `inference/`. - -#### Examples - -Local cache only (no network): +- `refresh`: regenerate and replace matching rows. + +Sampling parameters and explicit engine settings are part of the cell +descriptor, while API keys and headers are redacted. Hosted provider defaults +that are not visible to the client cannot be hashed. ELO and meta-eval cache raw +judge outputs and recompute preferences, calibration, costs, and aggregate +metrics on each run. + +Local-only example: ```bash judgearena \ - --task alpaca-eval \ - --model.name gpt4_1106_preview \ - --model.baseline VLLM/utter-project/EuroLLM-9B \ - --judge.model OpenRouter/deepseek/deepseek-chat-v3.1 \ - --generation.n_instructions 10 \ + --config_path configs/alpaca_eval.yaml \ --store_root ~/judgearena-data/cache ``` -Run with explicit Hugging Face fetch and push: +Hugging Face synchronization is opt-in. `--cache_fetch` fetches cells lazily; +`--cache_push` uploads cells written during a successful run. Pipeline sync is +best-effort, while the standalone commands fail loudly. ```bash judgearena \ --config_path configs/alpaca_eval.yaml \ - --cache.store_root ~/judgearena-data/cache \ - --cache.cache_hf_repo judge-arena/judge-arena-cache \ - --cache.cache_fetch \ - --cache.cache_push \ - --cache.pushed_by "$USER" + --store_root ~/judgearena-data/cache \ + --cache_fetch \ + --cache_push ``` -Offline workflow (login node → compute → login node): +For offline compute nodes, fetch and push closed cells from the login node: ```bash -# login node: prefetch cells for the benchmark judgearena-cache fetch \ --store_root ~/judgearena-data/cache \ - --cache_hf_repo judge-arena/judge-arena-cache \ --task alpaca-eval -# compute node: local store only judgearena \ --config_path configs/alpaca_eval.yaml \ --store_root ~/judgearena-data/cache \ --no-cache_fetch -# login node: merge and upload -judgearena-cache push \ - --store_root ~/judgearena-data/cache \ - --cache_hf_repo judge-arena/judge-arena-cache \ - --task alpaca-eval \ - --pushed_by "$USER" -``` - -If that path is on NFS, stage the **closed** store to node-local storage before -opening it on a compute node, then stage it back after the job. Never open the -same cell concurrently from multiple hosts. - -Filtered sync (`fetch` requires `--prefix` or at least `--task`; `push` without -filters uploads every local cell): - -```bash -judgearena-cache fetch \ - --store_root ~/judgearena-data/cache \ - --prefix inference/alpaca-eval/OpenRouter - judgearena-cache push \ --store_root ~/judgearena-data/cache \ --task alpaca-eval \ - --provider OpenRouter \ - --model deepseek/deepseek-chat-v3.1 \ - --create_pr \ --pushed_by "$USER" ``` -Open a pull request from the main CLI instead: - -```bash -judgearena \ - --config_path configs/alpaca_eval.yaml \ - --store_root ~/judgearena-data/cache \ - --cache_fetch \ - --cache_push \ - --cache_create_pr -``` - -#### `judgearena-cache backfill` - -`judgearena-cache backfill` imports **verified hosted judge rows** from saved -generate-and-evaluate, MT-Bench, and meta-eval run folders into the unified store. -For generate-and-evaluate and meta-eval, it reconstructs the rendered judge input -from saved configuration and checks it against stored `judge_input` text. -MT-Bench stores the exact system and rendered user messages used for each pass, -which are canonicalized directly. -**Insert-only, no network.** - -Imported: hosted judge outputs (`OpenRouter`, `ChatOpenAI`, `OpenAI`, `Together`, -and `Dummy` in tests). Model-generation outputs are never backfilled because old -artifacts cannot prove whether they came from inference or precomputed datasets. - -Skipped: generation-only artifacts, ELO runs, local engines (vLLM / LlamaCpp), -legacy pass-level caches, identity-keyed meta-eval DBs under `cache/db`, legacy -`judgements.db` / `completions.db` cells, rows whose recomputed input does not -match saved `judge_input`, and conflicting outputs for the same key. - -Dry-run first, then write: - -```bash -judgearena-cache backfill results/ \ - --store_root ~/judgearena-data/cache \ - --dry_run \ - --report ~/judgearena-data/cache/backfill-report.json - -judgearena-cache backfill results/alpaca-eval-2026-04-01/ \ - --store_root ~/judgearena-data/cache \ - --report ~/judgearena-data/cache/backfill-report.json -``` - -The JSON report lists `written`, `existing`, `rows_planned`, `runs_processed`, -and per-reason `skipped` counts. - -Optional YAML (commented so defaults stay network-free): - -```yaml -# cache: -# store_root: ~/judgearena-data/cache -# cache_mode: use -# cache_hf_repo: judge-arena/judge-arena-cache -# # cache_fetch: true -# # cache_push: true -# pushed_by: your-user -``` - -Standalone cache CLI: `judgearena-cache {fetch,push,backfill}` (see -`judgearena-cache --help`). +SQLite cells must not be open concurrently from multiple hosts or shared over +NFS. Stage closed cells to node-local storage or synchronize them through +Hugging Face. `judgearena-cache fetch` requires `--prefix` or at least `--task`; +unfiltered `push` uploads every local cell. ### Length and Token Parameters diff --git a/configs/alpaca_eval.yaml b/configs/alpaca_eval.yaml index f50bfad..149e765 100644 --- a/configs/alpaca_eval.yaml +++ b/configs/alpaca_eval.yaml @@ -10,11 +10,3 @@ judge: model: OpenRouter/google/gemma-4-31b-it generation: n_instructions: 10 -# Optional unified inference cache (disabled unless store_root is set): -# cache: -# store_root: ~/judgearena-data/cache -# cache_mode: use -# cache_hf_repo: judge-arena/judge-arena-cache -# # cache_fetch: true # prefetch on login node via judgearena-cache instead -# # cache_push: true # or judgearena-cache push after the run -# pushed_by: your-user diff --git a/configs/elo_comparia.yaml b/configs/elo_comparia.yaml index 37e541f..c21730b 100644 --- a/configs/elo_comparia.yaml +++ b/configs/elo_comparia.yaml @@ -14,11 +14,3 @@ elo: # Restrict the arena to these language codes. Defaults to English here; add # more to evaluate multilingually, e.g. languages: ["en", "fr", "de"]. languages: ["en"] -# Optional unified inference cache (disabled unless store_root is set): -# cache: -# store_root: ~/judgearena-data/cache -# cache_mode: use -# cache_hf_repo: judge-arena/judge-arena-cache -# # cache_fetch: true -# # cache_push: true -# pushed_by: your-user diff --git a/judgearena/cache_backfill.py b/judgearena/cache_backfill.py deleted file mode 100644 index 5271e98..0000000 --- a/judgearena/cache_backfill.py +++ /dev/null @@ -1,324 +0,0 @@ -"""Backfill hosted judge inference rows from saved run folders into the unified cache.""" - -from __future__ import annotations - -import json -import sqlite3 -import uuid -from collections import defaultdict -from dataclasses import dataclass, field -from pathlib import Path -from typing import Any - -import pandas as pd - -from judgearena.cache_backfill_discovery import ( - SKIP_REASON_BY_KIND, - ArtifactKind, - ClassifiedSource, - discover_sources, -) -from judgearena.cache_backfill_sources import ( - BackfillRow, - SourceExtraction, - extract_gae_rows, - extract_meta_eval_rows, - extract_mt_bench_rows, -) -from judgearena.log import get_logger -from judgearena.store_sqlite import ( - INFERENCE_DB_NAME, - SQLiteInferenceStore, - descriptor_hash, - stable_json_dumps, - store_folder, - write_store_metadata, -) - -logger = get_logger(__name__) - -BACKFILL_PUSHED_BY = "backfill" - - -@dataclass -class BackfillReport: - written: int = 0 - existing: int = 0 - rows_planned: int = 0 - skipped: dict[str, int] = field(default_factory=dict) - sources: dict[str, dict[str, int]] = field(default_factory=dict) - runs_processed: int = 0 - dry_run: bool = False - - def to_jsonable(self) -> dict[str, Any]: - return { - "written": self.written, - "existing": self.existing, - "rows_planned": self.rows_planned, - "skipped": dict(sorted(self.skipped.items())), - "sources": self.sources, - "runs_processed": self.runs_processed, - "dry_run": self.dry_run, - } - - -def _merge_skip_counts(target: dict[str, int], delta: dict[str, int]) -> None: - for reason, count in delta.items(): - target[reason] = target.get(reason, 0) + count - - -def _extract_rows(classified: ClassifiedSource) -> SourceExtraction: - run_dir = classified.run_dir - assert run_dir is not None - if classified.kind == ArtifactKind.GAE_RUN: - return extract_gae_rows(run_dir) - if classified.kind == ArtifactKind.MT_BENCH_RUN: - return extract_mt_bench_rows(run_dir) - if classified.kind == ArtifactKind.META_EVAL_RUN: - return extract_meta_eval_rows(run_dir) - raise ValueError(f"Unsupported migratable kind: {classified.kind}") - - -def _row_cache_key(row: BackfillRow) -> tuple[str, str, str, str]: - input_hash = descriptor_hash(row.canonical_input, length=None) - return row.task, row.model_spec, descriptor_hash(row.descriptor), input_hash - - -def _dedupe_and_detect_conflicts( - rows: list[BackfillRow], -) -> tuple[list[BackfillRow], int]: - grouped: dict[tuple[str, str, str, str], list[BackfillRow]] = defaultdict(list) - for row in rows: - grouped[_row_cache_key(row)].append(row) - - final: list[BackfillRow] = [] - dropped_rows = 0 - for key_rows in grouped.values(): - if len({row.output_text for row in key_rows}) > 1: - dropped_rows += len(key_rows) - continue - # Preserve duplicates so each distinct metadata association is written. - final.extend(key_rows) - - return final, dropped_rows - - -def _classify_rows_for_cell( - cell_rows: list[BackfillRow], - stored_outputs: dict[str, str], -) -> tuple[list[BackfillRow], list[BackfillRow], list[BackfillRow], int]: - to_write: list[BackfillRow] = [] - metadata_rows: list[BackfillRow] = [] - conflicting: list[BackfillRow] = [] - existing_count = 0 - - rows_by_hash: dict[str, list[BackfillRow]] = defaultdict(list) - for row in cell_rows: - rows_by_hash[descriptor_hash(row.canonical_input, length=None)].append(row) - - for input_hash, hash_rows in rows_by_hash.items(): - stored_output = stored_outputs.get(input_hash) - if stored_output is None: - to_write.append(hash_rows[0]) - metadata_rows.extend(hash_rows) - continue - if stored_output == hash_rows[0].output_text: - existing_count += 1 - metadata_rows.extend(hash_rows) - continue - conflicting.extend(hash_rows) - - return to_write, metadata_rows, conflicting, existing_count - - -def _write_rows( - rows: list[BackfillRow], - store_root: Path, - *, - dry_run: bool, -) -> tuple[int, int, dict[str, int]]: - if not rows: - return 0, 0, {} - - by_cell: dict[tuple[str, str, str], list[BackfillRow]] = defaultdict(list) - for row in rows: - by_cell[(row.task, row.model_spec, descriptor_hash(row.descriptor))].append(row) - - written = 0 - existing = 0 - skipped: dict[str, int] = {} - run_id = str(uuid.uuid4()) - - for (task, model_spec, config_hash), cell_rows in by_cell.items(): - descriptor = cell_rows[0].descriptor - try: - folder = store_folder(store_root, task, model_spec, config_hash) - db_path = folder / INFERENCE_DB_NAME - input_hashes = [ - descriptor_hash(row.canonical_input, length=None) for row in cell_rows - ] - unique_hashes = list(dict.fromkeys(input_hashes)) - - if dry_run: - if db_path.exists(): - with SQLiteInferenceStore(db_path, readonly=True) as store: - stored_outputs = store.outputs_by_hash(unique_hashes) - to_write, _, conflicting, existing_count = ( - _classify_rows_for_cell(cell_rows, stored_outputs) - ) - existing += existing_count - written += len(to_write) - if conflicting: - _merge_skip_counts( - skipped, - {"conflicting_existing_output": len(conflicting)}, - ) - else: - written += len(unique_hashes) - continue - - write_store_metadata(folder, descriptor) - with SQLiteInferenceStore(db_path) as store: - stored_outputs = store.outputs_by_hash(unique_hashes) - to_write, metadata_rows, conflicting, existing_count = ( - _classify_rows_for_cell(cell_rows, stored_outputs) - ) - existing += existing_count - if conflicting: - _merge_skip_counts( - skipped, - {"conflicting_existing_output": len(conflicting)}, - ) - - outputs_payload = [] - for row in to_write: - input_hash = descriptor_hash(row.canonical_input, length=None) - outputs_payload.append( - { - "input_hash": input_hash, - "input_text": row.canonical_input, - "output_text": row.output_text, - "producer_metadata_json": stable_json_dumps( - row.producer_metadata - ), - } - ) - - metadata_payload = [] - for row in metadata_rows: - metadata_payload.append( - { - "input_hash": descriptor_hash( - row.canonical_input, length=None - ), - "metadata_json": stable_json_dumps(row.row_metadata), - } - ) - - if outputs_payload and metadata_payload: - outputs_written, _ = store.save_outputs_and_metadata( - pd.DataFrame(outputs_payload), - pd.DataFrame(metadata_payload), - pushed_by=BACKFILL_PUSHED_BY, - run_id=run_id, - replace=False, - ) - written += outputs_written - elif outputs_payload: - written += store.save_outputs( - pd.DataFrame(outputs_payload), - pushed_by=BACKFILL_PUSHED_BY, - run_id=run_id, - replace=False, - ) - elif metadata_payload: - store.save_metadata(pd.DataFrame(metadata_payload), run_id=run_id) - except (OSError, ValueError, sqlite3.Error) as exc: - logger.warning( - "Cell integrity error for task=%s model=%s config=%s: %s", - task, - model_spec, - config_hash, - exc, - ) - _merge_skip_counts(skipped, {"cell_integrity_error": len(cell_rows)}) - - return written, existing, skipped - - -def backfill_sources( - sources: list[Path | str], - store_root: Path | str, - *, - dry_run: bool = False, -) -> BackfillReport: - """Discover saved judge runs and insert reconstructable rows into the store.""" - report = BackfillReport(dry_run=dry_run) - resolved_sources = [Path(source) for source in sources] - discovery = discover_sources(resolved_sources) - - for skipped in discovery.skipped: - reason = SKIP_REASON_BY_KIND.get(skipped.kind, "unknown_source") - report.skipped[reason] = report.skipped.get(reason, 0) + 1 - - extracted_rows: list[BackfillRow] = [] - for classified in discovery.migratable_runs: - run_dir = classified.run_dir - try: - extraction = _extract_rows(classified) - except Exception as exc: - logger.warning( - "Source extraction failed for %s: %s", - run_dir.name if run_dir else classified.path, - exc, - ) - _merge_skip_counts(report.skipped, {"source_extraction_failed": 1}) - continue - - _merge_skip_counts(report.skipped, extraction.skipped) - source_stats = report.sources.setdefault( - extraction.source_kind, - {"runs": 0, "rows_extracted": 0}, - ) - source_stats["runs"] += 1 - source_stats["rows_extracted"] += len(extraction.rows) - extracted_rows.extend(extraction.rows) - report.runs_processed += 1 - - deduped_rows, conflict_count = _dedupe_and_detect_conflicts(extracted_rows) - report.rows_planned = len(deduped_rows) - if conflict_count: - _merge_skip_counts(report.skipped, {"conflicting_outputs": conflict_count}) - - written, existing, write_skipped = _write_rows( - deduped_rows, - Path(store_root).expanduser(), - dry_run=dry_run, - ) - _merge_skip_counts(report.skipped, write_skipped) - report.written = written - report.existing = existing - return report - - -def write_report(report: BackfillReport, path: Path | str) -> None: - """Persist a JSON-safe backfill report.""" - output = Path(path) - output.parent.mkdir(parents=True, exist_ok=True) - output.write_text( - json.dumps(report.to_jsonable(), indent=2, sort_keys=True) + "\n", - encoding="utf-8", - ) - - -def log_report_summary(report: BackfillReport) -> None: - logger.info( - "Backfill complete: written=%d existing=%d rows_planned=%d " - "runs=%d dry_run=%s skipped=%s", - report.written, - report.existing, - report.rows_planned, - report.runs_processed, - report.dry_run, - report.skipped, - ) diff --git a/judgearena/cache_backfill_common.py b/judgearena/cache_backfill_common.py deleted file mode 100644 index 9a0417f..0000000 --- a/judgearena/cache_backfill_common.py +++ /dev/null @@ -1,59 +0,0 @@ -"""Shared helpers for cache backfill extraction.""" - -from __future__ import annotations - -from pathlib import Path - -import pandas as pd -from langchain_core.messages import HumanMessage, SystemMessage -from langchain_core.prompt_values import ChatPromptValue - -HOSTED_BACKFILL_PROVIDERS = frozenset( - {"OpenRouter", "ChatOpenAI", "OpenAI", "Together", "Dummy"} -) - - -def provider_from_model_spec(model_spec: str) -> str: - provider, _, _ = model_spec.partition("/") - return provider - - -def is_backfillable_provider(model_spec: str) -> bool: - return provider_from_model_spec(model_spec) in HOSTED_BACKFILL_PROVIDERS - - -def source_run_id(run_dir: Path) -> str: - return run_dir.name - - -def chat_prompt_value( - *, system_prompt: str | None, user_prompt: str -) -> ChatPromptValue: - messages = [] - if system_prompt: - messages.append(SystemMessage(content=system_prompt)) - messages.append(HumanMessage(content=user_prompt)) - return ChatPromptValue(messages=messages) - - -def prompt_text(value: object) -> str | None: - if value is None or (isinstance(value, float) and pd.isna(value)): - return None - text = str(value) - if not text.strip() or text.strip().lower() == "nan": - return None - return text - - -def mt_swapped(value: object) -> bool: - if value is None or (isinstance(value, float) and pd.isna(value)): - return False - if isinstance(value, bool): - return value - if isinstance(value, (int, float)): - return bool(value) - return str(value).strip().lower() in {"1", "true", "yes"} - - -def increment(skipped: dict[str, int], reason: str, count: int = 1) -> None: - skipped[reason] = skipped.get(reason, 0) + count diff --git a/judgearena/cache_backfill_config.py b/judgearena/cache_backfill_config.py deleted file mode 100644 index 0ad765b..0000000 --- a/judgearena/cache_backfill_config.py +++ /dev/null @@ -1,166 +0,0 @@ -"""Reconstruct model configurations from historical JudgeArena runs.""" - -from __future__ import annotations - -import json -from pathlib import Path -from typing import Any - -from judgearena.config import CacheArgs, RunConfig, load_config -from judgearena.meta_eval.cli_args import CliMetaEvalArgs -from judgearena.model_adapters import PreparedModel -from judgearena.models import build_default_judge_model_kwargs, make_model -from judgearena.repro import METADATA_FILENAME - - -def _cache_relevant_config(cfg: RunConfig) -> dict[str, Any]: - payload = cfg.model_dump(mode="json") - payload.pop("cache", None) - payload.pop("run", None) - return payload - - -def _load_modern_run_config(run_dir: Path) -> RunConfig | None: - candidates: list[tuple[str, RunConfig]] = [] - metadata_path = run_dir / METADATA_FILENAME - if metadata_path.exists(): - payload = json.loads(metadata_path.read_text(encoding="utf-8")) - run = payload.get("run") - if not isinstance(run, dict): - raise ValueError( - f"{METADATA_FILENAME} has no run config in {run_dir.name}." - ) - candidates.append((METADATA_FILENAME, RunConfig(**run))) - config_path = run_dir / "config.yaml" - if config_path.exists(): - candidates.append(("config.yaml", load_config(config_path))) - if not candidates: - return None - - expected = _cache_relevant_config(candidates[0][1]) - conflicting = [ - name - for name, candidate in candidates[1:] - if _cache_relevant_config(candidate) != expected - ] - if conflicting: - sources = ", ".join([candidates[0][0], *conflicting]) - raise ValueError(f"Conflicting run configs in {run_dir.name}: {sources}") - return candidates[0][1] - - -def _legacy_args_to_run_config(args: dict[str, Any]) -> RunConfig: - judge_model = args.get("judge_model") or args.get("judge", {}).get("model") - if not isinstance(judge_model, str): - raise ValueError("Legacy args missing judge_model.") - - engine_kwargs = dict(args.get("engine_kwargs") or {}) - judge_engine_kwargs = dict(args.get("judge_engine_kwargs") or {}) - judge_engine_kwargs.update(engine_kwargs) - truncate_all = args.get("truncate_all_input_chars", 8192) - truncate_judge = args.get("truncate_judge_input_chars") - if truncate_judge is None: - truncate_judge = truncate_all - - return RunConfig( - task=str(args["task"]), - model={ - "name": args.get("model_A") or args.get("model", {}).get("name"), - "baseline": args.get("model_B") or args.get("model", {}).get("baseline"), - "max_out_tokens": args.get("max_out_tokens_models") - or args.get("model", {}).get("max_out_tokens", 32768), - "max_model_len": args.get("max_model_len"), - "chat_template": args.get("chat_template"), - "engine_kwargs": engine_kwargs, - }, - judge={ - "model": judge_model, - "max_out_tokens": args.get("max_out_tokens_judge") - or args.get("judge", {}).get("max_out_tokens", 32768), - "max_model_len": args.get("max_model_len_judge") - or args.get("max_model_len"), - "chat_template": args.get("chat_template_judge") - or args.get("chat_template"), - "engine_kwargs": judge_engine_kwargs, - "provide_explanation": bool(args.get("provide_explanation", False)), - "swap_mode": args.get("swap_mode", "fixed"), - "prompt_preset": args.get("prompt_preset"), - "system_prompt_file": args.get("judge_system_prompt_file"), - "user_prompt_file": args.get("judge_user_prompt_file"), - "strip_thinking_before_judging": bool( - args.get("strip_thinking_before_judging", False) - ), - }, - generation={ - "n_instructions": args.get("n_instructions"), - "truncate_all_input_chars": truncate_all, - "truncate_judge_input_chars": truncate_judge, - }, - run={ - "result_folder": str(args.get("result_folder", "results")), - "seed": args.get("seed", 0), - }, - ) - - -def load_gae_run_config(run_dir: Path) -> RunConfig: - modern_cfg = _load_modern_run_config(run_dir) - if modern_cfg is not None: - return modern_cfg - args_paths = sorted(run_dir.glob("args-*.json")) - if not args_paths: - raise ValueError(f"No reconstructable config found under {run_dir.name}.") - if len(args_paths) > 1: - names = ", ".join(path.name for path in args_paths) - raise ValueError( - f"Ambiguous legacy args files under {run_dir.name}; " - f"expected one args-*.json or config/metadata: {names}" - ) - args = json.loads(args_paths[0].read_text(encoding="utf-8")) - return _legacy_args_to_run_config(args) - - -def load_meta_args(run_dir: Path) -> CliMetaEvalArgs: - args_path = run_dir / "args.json" - if not args_path.exists(): - raise ValueError(f"Meta-eval run missing args.json: {run_dir.name}") - payload = json.loads(args_path.read_text(encoding="utf-8")) - cache_payload = payload.pop("cache", {}) - payload.pop("ignore_cache", None) - if isinstance(cache_payload, dict): - cache_payload.pop("ignore_cache", None) - payload["cache"] = CacheArgs(**cache_payload) if cache_payload else CacheArgs() - return CliMetaEvalArgs(**payload) - - -def build_gae_judge_model(cfg: RunConfig) -> PreparedModel: - return make_model( - model=cfg.judge.model, - **build_default_judge_model_kwargs( - cfg.judge.model, - cfg.model.engine_kwargs, - judge_engine_kwargs_override=cfg.judge.model_kwargs( - fallback_chat_template=cfg.model.chat_template, - ), - ), - ) - - -def build_mt_judge_model(cfg: RunConfig, *, delegated: bool) -> PreparedModel: - judge_model_kwargs = cfg.judge.model_kwargs( - base_engine_kwargs=cfg.model.engine_kwargs, - fallback_chat_template=cfg.model.chat_template, - ) - if delegated and cfg.judge.temperature is None: - judge_model_kwargs.setdefault("temperature", 0.0) - return make_model(model=cfg.judge.model, **judge_model_kwargs) - - -def build_meta_judge_model(args: CliMetaEvalArgs) -> PreparedModel: - return make_model( - model=args.judge_model, - max_tokens=args.max_out_tokens_judge, - max_model_len=args.max_model_len, - chat_template=args.chat_template, - **args.engine_kwargs, - ) diff --git a/judgearena/cache_backfill_discovery.py b/judgearena/cache_backfill_discovery.py deleted file mode 100644 index 0b8577a..0000000 --- a/judgearena/cache_backfill_discovery.py +++ /dev/null @@ -1,340 +0,0 @@ -"""Discover and classify saved run folders and cache artifacts for backfill.""" - -from __future__ import annotations - -import json -from dataclasses import dataclass, field -from enum import StrEnum -from pathlib import Path - -import yaml - -from judgearena.constants import ELO_TASK_PREFIX, META_EVAL_TASK -from judgearena.repro import METADATA_FILENAME - -LEGACY_CELL_DB_NAMES = frozenset({"judgements.db", "completions.db"}) -PASS_LEVEL_CACHE_SUFFIXES = (".csv.zip", ".parquet", ".csv") -GENERATION_ONLY_MARKERS = frozenset( - { - "completions.parquet", - "completions.csv", - "completions.csv.zip", - "model_outputs.parquet", - } -) -GAE_REQUIRED_COLUMNS = frozenset( - {"instruction", "completion_A", "completion_B", "judge_input"} -) -GAE_OUTPUT_COLUMNS = frozenset({"judge_completion", "judge_output"}) - - -class ArtifactKind(StrEnum): - GAE_RUN = "gae_run" - MT_BENCH_RUN = "mt_bench_run" - META_EVAL_RUN = "meta_eval_run" - ELO_RUN = "elo_run" - LEGACY_CACHE_CELL = "legacy_cache_cell" - META_EVAL_IDENTITY_DB = "meta_eval_identity_db" - PASS_LEVEL_CACHE = "pass_level_cache" - GENERATION_ARTIFACT = "generation_artifact" - UNKNOWN = "unknown" - - -SKIP_REASON_BY_KIND: dict[ArtifactKind, str] = { - ArtifactKind.ELO_RUN: "elo_run_missing_inference_outputs", - ArtifactKind.LEGACY_CACHE_CELL: "legacy_cache_cell_unmigratable", - ArtifactKind.META_EVAL_IDENTITY_DB: "meta_eval_identity_db", - ArtifactKind.PASS_LEVEL_CACHE: "pass_level_cache_untrusted", - ArtifactKind.GENERATION_ARTIFACT: "generation_provenance_unknown", - ArtifactKind.UNKNOWN: "unknown_judge_run", -} - - -@dataclass -class ClassifiedSource: - path: Path - kind: ArtifactKind - run_dir: Path | None = None - - -@dataclass -class DiscoveryReport: - migratable_runs: list[ClassifiedSource] = field(default_factory=list) - skipped: list[ClassifiedSource] = field(default_factory=list) - - -def _glob_has_matches(resolved: Path, pattern: str) -> bool: - return next(resolved.glob(pattern), None) is not None - - -def _is_elo_task_name(value: str | None) -> bool: - return bool(value and value.startswith(ELO_TASK_PREFIX)) - - -def _looks_like_meta_eval_dir(path: Path) -> bool: - return ( - path.name.startswith(f"{META_EVAL_TASK}-") - or (path / "annotations.parquet").exists() - ) - - -def _csv_columns(csv_path: Path) -> set[str]: - header = csv_path.read_text(encoding="utf-8").splitlines()[:1] - if not header: - return set() - return {part.strip() for part in header[0].split(",")} - - -def _looks_like_mt_annotations(path: Path) -> bool: - for csv_path in path.glob("*-annotations.csv"): - columns = _csv_columns(csv_path) - if not columns: - continue - mt_markers = {"question_id", "turn", "category"} - if mt_markers.issubset(columns): - return True - if "g1_user_prompt" in columns or "user_prompt" in columns: - return True - return False - - -def _looks_like_gae_annotations(path: Path) -> bool: - for csv_path in path.glob("*-annotations.csv"): - columns = _csv_columns(csv_path) - if not columns: - continue - if not GAE_REQUIRED_COLUMNS.issubset(columns): - continue - if not GAE_OUTPUT_COLUMNS.intersection(columns): - continue - return True - return False - - -def _task_from_run_dir(run_dir: Path) -> str | None: - metadata_path = run_dir / METADATA_FILENAME - if metadata_path.exists(): - payload = json.loads(metadata_path.read_text(encoding="utf-8")) - run = payload.get("run") - if isinstance(run, dict): - task = run.get("task") - if isinstance(task, str): - return task - config_path = run_dir / "config.yaml" - if config_path.exists(): - payload = yaml.safe_load(config_path.read_text(encoding="utf-8")) or {} - task = payload.get("task") - if isinstance(task, str): - return task - for args_path in run_dir.glob("args-*.json"): - payload = json.loads(args_path.read_text(encoding="utf-8")) - task = payload.get("task") - if isinstance(task, str): - return task - args_path = run_dir / "args.json" - if args_path.exists(): - payload = json.loads(args_path.read_text(encoding="utf-8")) - if isinstance(payload, dict): - task = payload.get("task") - if isinstance(task, str): - return task - return None - - -def _classify_path(path: Path) -> ClassifiedSource: - resolved = path.resolve() - - if resolved.is_file(): - suffix = "".join(resolved.suffixes) or resolved.suffix - name = resolved.name - if name in GENERATION_ONLY_MARKERS: - return ClassifiedSource(resolved, ArtifactKind.GENERATION_ARTIFACT) - if ( - name.endswith(PASS_LEVEL_CACHE_SUFFIXES) - or suffix in PASS_LEVEL_CACHE_SUFFIXES - ): - return ClassifiedSource(resolved, ArtifactKind.PASS_LEVEL_CACHE) - if resolved.suffix == ".db": - parts = {part.lower() for part in resolved.parts} - if "cache" in parts and "db" in parts: - return ClassifiedSource(resolved, ArtifactKind.META_EVAL_IDENTITY_DB) - if name in LEGACY_CELL_DB_NAMES: - return ClassifiedSource(resolved, ArtifactKind.LEGACY_CACHE_CELL) - return ClassifiedSource(resolved, ArtifactKind.UNKNOWN) - - if not resolved.is_dir(): - return ClassifiedSource(resolved, ArtifactKind.UNKNOWN) - - if resolved.name in LEGACY_CELL_DB_NAMES or any( - (resolved / db_name).exists() for db_name in LEGACY_CELL_DB_NAMES - ): - return ClassifiedSource(resolved, ArtifactKind.LEGACY_CACHE_CELL) - - parts = {part.lower() for part in resolved.parts} - if ( - "cache" in parts - and "db" in parts - and any(child.suffix == ".db" for child in resolved.glob("*.db")) - ): - return ClassifiedSource(resolved, ArtifactKind.META_EVAL_IDENTITY_DB) - - task = _task_from_run_dir(resolved) - if _is_elo_task_name(task) or resolved.name.startswith(ELO_TASK_PREFIX): - return ClassifiedSource(resolved, ArtifactKind.ELO_RUN, run_dir=resolved) - - if (resolved / "annotations.parquet").exists() and ( - (resolved / "args.json").exists() or (resolved / METADATA_FILENAME).exists() - ): - return ClassifiedSource( - resolved, - ArtifactKind.META_EVAL_RUN, - run_dir=resolved, - ) - - if list(resolved.glob("*-annotations.csv")): - if task == "mt-bench" or _looks_like_mt_annotations(resolved): - return ClassifiedSource( - resolved, - ArtifactKind.MT_BENCH_RUN, - run_dir=resolved, - ) - if _is_elo_task_name(task): - return ClassifiedSource(resolved, ArtifactKind.ELO_RUN, run_dir=resolved) - if _looks_like_gae_annotations(resolved): - return ClassifiedSource(resolved, ArtifactKind.GAE_RUN, run_dir=resolved) - return ClassifiedSource(resolved, ArtifactKind.UNKNOWN) - - if _looks_like_meta_eval_dir(resolved) and (resolved / "args.json").exists(): - return ClassifiedSource( - resolved, - ArtifactKind.META_EVAL_RUN, - run_dir=resolved, - ) - - if any(resolved.joinpath(name).exists() for name in GENERATION_ONLY_MARKERS): - return ClassifiedSource(resolved, ArtifactKind.GENERATION_ARTIFACT) - - if any( - _glob_has_matches(resolved, f"*{suffix}") - for suffix in PASS_LEVEL_CACHE_SUFFIXES - ): - if not list(resolved.glob("*-annotations.csv")): - return ClassifiedSource(resolved, ArtifactKind.PASS_LEVEL_CACHE) - - return ClassifiedSource(resolved, ArtifactKind.UNKNOWN) - - -def _discover_run_dirs(source: Path) -> list[Path]: - classified = _classify_path(source) - if classified.run_dir is not None: - return [classified.run_dir] - - if not source.is_dir(): - return [] - - run_dirs: list[Path] = [] - seen: set[Path] = set() - for annotation_csv in source.rglob("*-annotations.csv"): - run_dir = annotation_csv.parent.resolve() - if run_dir not in seen: - seen.add(run_dir) - run_dirs.append(run_dir) - for annotation_parquet in source.rglob("annotations.parquet"): - run_dir = annotation_parquet.parent.resolve() - if run_dir not in seen: - seen.add(run_dir) - run_dirs.append(run_dir) - return sorted(run_dirs) - - -def _discover_nested_skipped_artifacts(source: Path) -> list[ClassifiedSource]: - skipped: list[ClassifiedSource] = [] - seen: set[Path] = set() - for db_name in LEGACY_CELL_DB_NAMES: - for db_path in source.rglob(db_name): - parent = db_path.parent.resolve() - if parent in seen: - continue - seen.add(parent) - classified = _classify_path(parent) - if classified.kind in SKIP_REASON_BY_KIND: - skipped.append(classified) - return skipped - - -def _collect_skipped_artifacts( - source: Path, seen_skipped: set[Path] -) -> list[ClassifiedSource]: - skipped: list[ClassifiedSource] = [] - for child in source.rglob("*"): - if child in seen_skipped: - continue - classified = _classify_path(child) - if classified.kind not in SKIP_REASON_BY_KIND: - continue - seen_skipped.add(child) - skipped.append(classified) - return skipped - - -def discover_sources(sources: list[Path]) -> DiscoveryReport: - """Discover migratable judge run folders and classify skipped artifacts.""" - report = DiscoveryReport() - seen_runs: set[Path] = set() - seen_skipped: set[Path] = set() - - for source in sources: - source = source.resolve() - direct = _classify_path(source) - - if direct.run_dir is not None: - if direct.run_dir not in seen_runs: - seen_runs.add(direct.run_dir) - report.migratable_runs.append(direct) - continue - - if source.is_file(): - if source not in seen_skipped: - seen_skipped.add(source) - report.skipped.append(direct) - continue - - run_dirs = _discover_run_dirs(source) - for nested in _discover_nested_skipped_artifacts(source): - if nested.path not in seen_skipped: - seen_skipped.add(nested.path) - report.skipped.append(nested) - - if not run_dirs: - for classified_child in _collect_skipped_artifacts(source, seen_skipped): - report.skipped.append(classified_child) - if ( - direct.kind in SKIP_REASON_BY_KIND - and source not in seen_skipped - and (source.is_file() or direct.kind != ArtifactKind.UNKNOWN) - ): - seen_skipped.add(source) - report.skipped.append(direct) - continue - - for run_dir in run_dirs: - if run_dir in seen_runs: - continue - classified = _classify_path(run_dir) - if classified.kind in SKIP_REASON_BY_KIND: - if run_dir not in seen_skipped: - seen_skipped.add(run_dir) - report.skipped.append(classified) - continue - if classified.kind in { - ArtifactKind.GAE_RUN, - ArtifactKind.MT_BENCH_RUN, - ArtifactKind.META_EVAL_RUN, - }: - seen_runs.add(run_dir) - report.migratable_runs.append(classified) - elif run_dir not in seen_skipped: - seen_skipped.add(run_dir) - report.skipped.append(classified) - - return report diff --git a/judgearena/cache_backfill_sources.py b/judgearena/cache_backfill_sources.py deleted file mode 100644 index 51af8cd..0000000 --- a/judgearena/cache_backfill_sources.py +++ /dev/null @@ -1,368 +0,0 @@ -"""Extract backfill rows from saved GAE, MT-Bench, and meta-eval run folders.""" - -from __future__ import annotations - -from dataclasses import dataclass -from pathlib import Path -from typing import Any - -import pandas as pd - -from judgearena.cache_backfill_common import ( - chat_prompt_value, - increment, - is_backfillable_provider, - mt_swapped, - prompt_text, - source_run_id, -) -from judgearena.cache_backfill_config import ( - build_gae_judge_model, - build_meta_judge_model, - build_mt_judge_model, - load_gae_run_config, - load_meta_args, -) -from judgearena.config import RunConfig, meta_eval_cache_task -from judgearena.evaluate import render_judge_inputs, resolve_run_judge_prompt -from judgearena.meta_eval.prompts import resolve_prompt_mode -from judgearena.model_adapters import PreparedModel - - -@dataclass(frozen=True) -class BackfillRow: - task: str - model_spec: str - descriptor: dict[str, Any] - canonical_input: str - output_text: str - row_metadata: dict[str, Any] - producer_metadata: dict[str, Any] - - -@dataclass -class SourceExtraction: - rows: list[BackfillRow] - skipped: dict[str, int] - source_kind: str - - -def _annotations_path(run_dir: Path) -> Path: - paths = sorted(run_dir.glob("*-annotations.csv")) - if len(paths) != 1: - names = ", ".join(path.name for path in paths) or "none" - raise ValueError( - f"Expected exactly one *-annotations.csv in {run_dir.name}; found {names}." - ) - return paths[0] - - -def _maybe_descriptor(model: PreparedModel) -> dict[str, Any] | None: - return model.cache_descriptor() - - -def _infer_gae_orientation(row: pd.Series, *, cfg: RunConfig) -> str | None: - model_a = str(row.get("model_A", "")) - model_b = str(row.get("model_B", "")) - focal = cfg.model.name - a_is_focal = model_a == focal - b_is_focal = model_b == focal - if a_is_focal and not b_is_focal: - return "direct" - if b_is_focal and not a_is_focal: - return "reversed" - return None - - -def _meta_eval_verify_completions(row: pd.Series) -> tuple[str, str]: - presented_a = prompt_text(row.get("presented_completion_a")) - presented_b = prompt_text(row.get("presented_completion_b")) - if presented_a is not None and presented_b is not None: - return presented_a, presented_b - completion_a = str(row.get("completion_a", "")) - completion_b = str(row.get("completion_b", "")) - if str(row.get("orientation", "forward")) == "swapped": - return completion_b, completion_a - return completion_a, completion_b - - -def extract_gae_rows(run_dir: Path) -> SourceExtraction: - cfg = load_gae_run_config(run_dir) - annotations_path = _annotations_path(run_dir) - df = pd.read_csv(annotations_path, keep_default_na=False) - skipped: dict[str, int] = {} - rows: list[BackfillRow] = [] - run_id = source_run_id(run_dir) - - if not is_backfillable_provider(cfg.judge.model): - increment(skipped, "local_engine_unsupported", len(df)) - return SourceExtraction(rows=[], skipped=skipped, source_kind="gae") - - judge_model = build_gae_judge_model(cfg) - descriptor = _maybe_descriptor(judge_model) - if descriptor is None: - increment(skipped, "local_engine_unsupported", len(df)) - return SourceExtraction(rows=[], skipped=skipped, source_kind="gae") - - resolved_prompt = resolve_run_judge_prompt(cfg.task, cfg.judge) - output_column = ( - "judge_completion" if "judge_completion" in df.columns else "judge_output" - ) - if output_column not in df.columns: - increment(skipped, "judge_input_unverifiable", len(df)) - return SourceExtraction(rows=[], skipped=skipped, source_kind="gae") - - if "judge_input" not in df.columns: - increment(skipped, "judge_input_unverifiable", len(df)) - return SourceExtraction(rows=[], skipped=skipped, source_kind="gae") - - rendered_inputs = render_judge_inputs( - df["instruction"].astype(str).tolist(), - df["completion_A"].astype(str).tolist(), - df["completion_B"].astype(str).tolist(), - system_prompt=resolved_prompt.system_prompt, - user_prompt_template=resolved_prompt.user_prompt_template, - truncate_input_chars=cfg.generation.truncate_judge_input_chars, - provide_explanation=cfg.judge.provide_explanation, - prompt_preset=resolved_prompt.preset_name, - strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging, - task=cfg.task, - system_file=cfg.judge.system_prompt_file, - user_file=cfg.judge.user_prompt_file, - ) - - producer = judge_model.producer_metadata() - - for idx, (_, row) in enumerate(df.iterrows()): - stored_input = row.get("judge_input") - if pd.isna(stored_input) or stored_input is None: - increment(skipped, "judge_input_unverifiable") - continue - rendered = rendered_inputs[idx].to_string() - if str(stored_input) != rendered: - increment(skipped, "judge_input_mismatch") - continue - output_text = prompt_text(row.get(output_column)) - if output_text is None: - increment(skipped, "judge_output_missing") - continue - - orientation = _infer_gae_orientation(row, cfg=cfg) - if orientation is None: - increment(skipped, "battle_orientation_unverifiable") - continue - prompt_input = rendered_inputs[idx] - rows.append( - BackfillRow( - task=cfg.task, - model_spec=judge_model.model_spec, - descriptor=descriptor, - canonical_input=judge_model.canonicalize_input(prompt_input), - output_text=output_text, - row_metadata={ - "task": cfg.task, - "instruction_index": str(row.get("instruction_index", idx)), - "presented_model_a": str(row.get("model_A", cfg.model.name)), - "presented_model_b": str(row.get("model_B", "")), - "orientation": orientation, - "source_run_id": run_id, - }, - producer_metadata=producer, - ) - ) - - return SourceExtraction(rows=rows, skipped=skipped, source_kind="gae") - - -def _mt_is_fastchat(df: pd.DataFrame) -> bool: - return "g1_user_prompt" in df.columns - - -def extract_mt_bench_rows(run_dir: Path) -> SourceExtraction: - cfg = load_gae_run_config(run_dir) - annotations_path = _annotations_path(run_dir) - df = pd.read_csv(annotations_path) - skipped: dict[str, int] = {} - rows: list[BackfillRow] = [] - run_id = source_run_id(run_dir) - - if not is_backfillable_provider(cfg.judge.model): - increment(skipped, "local_engine_unsupported", len(df)) - return SourceExtraction(rows=[], skipped=skipped, source_kind="mt_bench") - - resolved_prompt = resolve_run_judge_prompt(cfg.task, cfg.judge, multi_turn=True) - judge_model = build_mt_judge_model(cfg, delegated=resolved_prompt.delegated) - descriptor = _maybe_descriptor(judge_model) - if descriptor is None: - increment(skipped, "local_engine_unsupported", len(df)) - return SourceExtraction(rows=[], skipped=skipped, source_kind="mt_bench") - - producer = judge_model.producer_metadata() - fastchat = _mt_is_fastchat(df) - - for _, row in df.iterrows(): - candidates: list[tuple[str, str | None, str, str]] = [] - if fastchat: - g1_output = prompt_text(row.get("g1_judgment")) - g1_prompt = prompt_text(row.get("g1_user_prompt")) - if g1_output is not None and g1_prompt is not None: - candidates.append( - ( - "direct", - prompt_text(row.get("system_prompt")), - g1_prompt, - g1_output, - ) - ) - g2_output = prompt_text(row.get("g2_judgment")) - g2_prompt = prompt_text(row.get("g2_user_prompt")) - if g2_output is not None and g2_prompt is not None: - candidates.append( - ( - "reversed", - prompt_text(row.get("system_prompt")), - g2_prompt, - g2_output, - ) - ) - else: - output = prompt_text(row.get("judge_completion")) - if output is None: - increment(skipped, "judge_input_unverifiable") - continue - orientation = "reversed" if mt_swapped(row.get("swapped")) else "direct" - user_prompt = prompt_text(row.get("user_prompt")) - if user_prompt is None: - increment(skipped, "judge_input_unverifiable") - continue - candidates.append( - ( - orientation, - prompt_text(row.get("system_prompt")), - user_prompt, - output, - ) - ) - - for orientation, system_prompt, user_prompt, output_text in candidates: - prompt_input = chat_prompt_value( - system_prompt=system_prompt, - user_prompt=user_prompt, - ) - turn_value = row.get("turn") - rows.append( - BackfillRow( - task=cfg.task, - model_spec=judge_model.model_spec, - descriptor=descriptor, - canonical_input=judge_model.canonicalize_input(prompt_input), - output_text=output_text, - row_metadata={ - "question_id": str(row.get("question_id", "")), - "category": row.get("category"), - "turn": int(turn_value) if pd.notna(turn_value) else None, - "orientation": orientation, - "prompt": row.get("prompt_name"), - "source_run_id": run_id, - }, - producer_metadata=producer, - ) - ) - - return SourceExtraction(rows=rows, skipped=skipped, source_kind="mt_bench") - - -def extract_meta_eval_rows(run_dir: Path) -> SourceExtraction: - args = load_meta_args(run_dir) - df = pd.read_parquet(run_dir / "annotations.parquet") - skipped: dict[str, int] = {} - rows: list[BackfillRow] = [] - run_id = source_run_id(run_dir) - - if not is_backfillable_provider(args.judge_model): - increment(skipped, "local_engine_unsupported", len(df)) - return SourceExtraction(rows=[], skipped=skipped, source_kind="meta_eval") - - prompt_spec = resolve_prompt_mode( - args.prompt_mode, - provide_explanation=args.provide_explanation, - ) - judge_model = build_meta_judge_model(args) - descriptor = _maybe_descriptor(judge_model) - if descriptor is None: - increment(skipped, "local_engine_unsupported", len(df)) - return SourceExtraction(rows=[], skipped=skipped, source_kind="meta_eval") - - if "judge_input" not in df.columns: - increment(skipped, "judge_input_unverifiable", len(df)) - return SourceExtraction(rows=[], skipped=skipped, source_kind="meta_eval") - - output_column = ( - "judge_completion" if "judge_completion" in df.columns else "judge_output" - ) - if output_column not in df.columns: - increment(skipped, "judge_input_unverifiable", len(df)) - return SourceExtraction(rows=[], skipped=skipped, source_kind="meta_eval") - - verify_a: list[str] = [] - verify_b: list[str] = [] - for _, row in df.iterrows(): - completion_a, completion_b = _meta_eval_verify_completions(row) - verify_a.append(completion_a) - verify_b.append(completion_b) - - rendered_inputs = render_judge_inputs( - df["instruction"].astype(str).tolist(), - verify_a, - verify_b, - system_prompt=prompt_spec.system_prompt, - user_prompt_template=prompt_spec.user_prompt_template, - truncate_input_chars=args.truncate_judge_input_chars, - provide_explanation=args.provide_explanation, - ) - - task = meta_eval_cache_task(args.reference_arena) - producer = judge_model.producer_metadata() - - for idx, (_, row) in enumerate(df.iterrows()): - stored_input = row.get("judge_input") - if pd.isna(stored_input) or stored_input is None: - increment(skipped, "judge_input_unverifiable") - continue - rendered = rendered_inputs[idx].to_string() - if str(stored_input) != rendered: - increment(skipped, "judge_input_mismatch") - continue - output_text = prompt_text(row.get(output_column)) - if output_text is None: - increment(skipped, "judge_output_missing") - continue - - orientation = str(row.get("orientation", "forward")) - prompt_input = rendered_inputs[idx] - rows.append( - BackfillRow( - task=task, - model_spec=judge_model.model_spec, - descriptor=descriptor, - canonical_input=judge_model.canonicalize_input(prompt_input), - output_text=output_text, - row_metadata={ - "reference_arena": args.reference_arena, - "benchmark": str(row.get("benchmark", "")), - "question_id": str(row.get("question_id", "")), - "presented_model_a": str( - row.get("presented_model_a", row.get("model_a", "")) - ), - "presented_model_b": str( - row.get("presented_model_b", row.get("model_b", "")) - ), - "prompt_mode": args.prompt_mode, - "orientation": orientation, - "source_run_id": run_id, - }, - producer_metadata=producer, - ) - ) - - return SourceExtraction(rows=rows, skipped=skipped, source_kind="meta_eval") diff --git a/judgearena/cache_sync.py b/judgearena/cache_sync.py index 04439ae..f8bcb26 100644 --- a/judgearena/cache_sync.py +++ b/judgearena/cache_sync.py @@ -1,13 +1,11 @@ -"""Standalone CLI for synchronizing inference cache cells with Hugging Face.""" +"""CLI for synchronizing inference cache cells with Hugging Face.""" from __future__ import annotations import argparse import getpass import sys -from pathlib import Path -from judgearena.cache_backfill import backfill_sources, log_report_summary, write_report from judgearena.log import configure_logging, get_logger from judgearena.store_sync import ( DEFAULT_CACHE_REPO, @@ -27,14 +25,11 @@ def _add_filter_args(parser: argparse.ArgumentParser) -> None: ) parser.add_argument("--task", help="Filter cells by benchmark task name.") parser.add_argument("--provider", help="Filter cells by provider, e.g. VLLM.") - parser.add_argument( - "--model", - help="Filter cells by model path (slashes become '--' in folders).", - ) + parser.add_argument("--model", help="Filter cells by model path.") parser.add_argument("--config_hash", help="Filter cells by descriptor hash.") -def _add_common(parser: argparse.ArgumentParser) -> None: +def _add_common_args(parser: argparse.ArgumentParser) -> None: parser.add_argument("--store_root", required=True, help="Local store root.") parser.add_argument("--cache_hf_repo", default=DEFAULT_CACHE_REPO) parser.add_argument("--repo_type", default="dataset") @@ -63,14 +58,11 @@ def _build_parser() -> argparse.ArgumentParser: "fetch", help="Discover and merge remote cells into the local store.", ) - _add_common(fetch) + _add_common_args(fetch) _add_filter_args(fetch) - push = subparsers.add_parser( - "push", - help="Merge and upload local cells.", - ) - _add_common(push) + push = subparsers.add_parser("push", help="Merge and upload local cells.") + _add_common_args(push) _add_filter_args(push) push.add_argument("--pushed_by", default=getpass.getuser()) push.add_argument("--create_pr", action="store_true") @@ -84,46 +76,12 @@ def _build_parser() -> argparse.ArgumentParser: action="store_true", help="Create a public repository when used with --ensure_repo.", ) - - backfill = subparsers.add_parser( - "backfill", - help="Backfill hosted judge outputs from saved run folders.", - ) - backfill.add_argument( - "sources", - nargs="+", - type=Path, - help="Run folders or parents containing saved judge annotations.", - ) - backfill.add_argument("--store_root", required=True, help="Local store root.") - backfill.add_argument( - "--dry_run", - action="store_true", - help="Plan and report without writing inference rows.", - ) - backfill.add_argument( - "--report", - type=Path, - help="Optional path to write a JSON backfill report.", - ) - backfill.add_argument("-v", "--verbose", action="count", default=0) return parser def main(argv: list[str] | None = None) -> None: args = _build_parser().parse_args(argv) - configure_logging(getattr(args, "verbose", 0)) - - if args.command == "backfill": - report = backfill_sources( - args.sources, - args.store_root, - dry_run=args.dry_run, - ) - if args.report is not None: - write_report(report, args.report) - log_report_summary(report) - return + configure_logging(args.verbose) try: path_prefix = _resolve_prefix(args) diff --git a/judgearena/config.py b/judgearena/config.py index 9442a79..02af6bc 100644 --- a/judgearena/config.py +++ b/judgearena/config.py @@ -20,9 +20,9 @@ YamlConfigSettingsSource, ) +from judgearena.baselines import native_pairwise_baseline from judgearena.constants import ELO_TASK_PREFIX, ELO_TASK_TO_ARENA, META_EVAL_TASK from judgearena.inference_cache import InferenceCache -from judgearena.pairwise_baselines import native_pairwise_baseline from judgearena.store_sync import DEFAULT_CACHE_REPO CacheMode = Literal["use", "off", "refresh"] diff --git a/judgearena/generate_and_evaluate.py b/judgearena/generate_and_evaluate.py index 788eda3..4e99dad 100644 --- a/judgearena/generate_and_evaluate.py +++ b/judgearena/generate_and_evaluate.py @@ -13,6 +13,11 @@ import pandas as pd +from judgearena.baselines import ( + ALPACA_EVAL_BASELINES, + PAIRWISE_BASELINES, + native_pairwise_baseline, +) from judgearena.benchmark import ( BenchmarkAdapter, build_generation_kwargs, @@ -36,11 +41,6 @@ make_run_log_path, ) from judgearena.mt_bench.mt_bench_utils import run_mt_bench -from judgearena.pairwise_baselines import ( - ALPACA_EVAL_BASELINES, - PAIRWISE_BASELINES, - native_pairwise_baseline, -) from judgearena.repro import write_run_metadata from judgearena.utils import compute_pref_summary, data_root, download_hf, read_df from judgearena.utils.eval import BattleReport diff --git a/judgearena/log.py b/judgearena/log.py index 045717c..a2a35c7 100644 --- a/judgearena/log.py +++ b/judgearena/log.py @@ -93,13 +93,7 @@ def configure_logging( # --- console handler --- # Avoid duplicate handlers when configure_logging is called more than once # (e.g. in tests). - console_handlers = [ - handler - for handler in root.handlers - if isinstance(handler, logging.StreamHandler) - and not isinstance(handler, logging.FileHandler) - ] - if not console_handlers: + if not root.handlers: handler = logging.StreamHandler(sys.stderr) handler.setLevel(level) handler.setFormatter( @@ -107,14 +101,11 @@ def configure_logging( ) root.addHandler(handler) else: - for handler in console_handlers: - handler.setLevel(level) - if handler.stream is not sys.stderr: - try: - handler.setStream(sys.stderr) - except ValueError: - # Test/output capture can close the previous stream. - handler.stream = sys.stderr + for h in root.handlers: + if isinstance(h, logging.StreamHandler) and not isinstance( + h, logging.FileHandler + ): + h.setLevel(level) # --- file handler (explicit --log-file) --- if log_file is not None: diff --git a/judgearena/pairwise_baselines.py b/judgearena/pairwise_baselines.py deleted file mode 100644 index 2a192ea..0000000 --- a/judgearena/pairwise_baselines.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Dataset-native pairwise baseline registry for generate-and-evaluate tasks.""" - -from __future__ import annotations - -from collections.abc import Mapping - -from judgearena.instruction_dataset.arena_hard import ARENA_HARD_BASELINES -from judgearena.instruction_dataset.m_arenahard import ( - M_ARENA_HARD_BASELINES, - split_m_arena_hard_dataset, -) -from judgearena.instruction_dataset.mt_bench import MT_BENCH_BASELINES - -ALPACA_EVAL_BASELINES: dict[str, str] = { - "alpaca-eval": "gpt4_1106_preview", -} - -PAIRWISE_BASELINES: dict[str, str | Mapping[str, str]] = { - **ALPACA_EVAL_BASELINES, - **ARENA_HARD_BASELINES, - **M_ARENA_HARD_BASELINES, - **MT_BENCH_BASELINES, -} - - -def native_pairwise_baseline(task: str) -> str | Mapping[str, str] | None: - """Return the dataset-native pairwise baseline, if the task defines one.""" - if task in PAIRWISE_BASELINES: - return PAIRWISE_BASELINES[task] - parsed_m_arena_hard = split_m_arena_hard_dataset(task) - if parsed_m_arena_hard is not None: - version_key, _lang_or_subset = parsed_m_arena_hard - return PAIRWISE_BASELINES[version_key] - return None diff --git a/tests/test_cache_backfill.py b/tests/test_cache_backfill.py deleted file mode 100644 index 1ce0a66..0000000 --- a/tests/test_cache_backfill.py +++ /dev/null @@ -1,792 +0,0 @@ -from __future__ import annotations - -import json -from pathlib import Path - -import pandas as pd -import pytest - -import judgearena.generate_and_evaluate as gae -from judgearena import cache_sync -from judgearena.cache_backfill import ( - BACKFILL_PUSHED_BY, - BackfillReport, - backfill_sources, - write_report, -) -from judgearena.cache_backfill_discovery import ArtifactKind, discover_sources -from judgearena.cache_backfill_sources import ( - _infer_gae_orientation, - extract_gae_rows, - extract_mt_bench_rows, -) -from judgearena.config import RunConfig, dump_config, meta_eval_cache_task -from judgearena.evaluate import render_judge_inputs, resolve_run_judge_prompt -from judgearena.meta_eval.cli_args import CliMetaEvalArgs -from judgearena.meta_eval.prompts import resolve_prompt_mode -from judgearena.repro import write_run_metadata -from judgearena.store_sqlite import INFERENCE_DB_NAME, SQLiteInferenceStore - - -def _synthetic_instructions(n: int = 2) -> pd.DataFrame: - return pd.DataFrame( - {"instruction": [f"instruction {i}" for i in range(n)]}, - index=pd.Index([f"idx-{i}" for i in range(n)], name="instruction_index"), - ) - - -def _cfg_with_cache(tmp_path, **overrides) -> RunConfig: - payload = { - "task": "alpaca-eval", - "model": {"name": "Dummy/gen-a", "baseline": "Dummy/gen-b"}, - "judge": {"model": "Dummy/score A: 0 score B: 10", "swap_mode": "fixed"}, - "generation": {"n_instructions": 2}, - "run": {"result_folder": str(tmp_path / "results"), "no_log_file": True}, - "cache": {"store_root": str(tmp_path / "live-cache")}, - } - payload.update(overrides) - return RunConfig(**payload) - - -def _write_gae_annotations( - run_dir: Path, - cfg: RunConfig, - *, - swap_both: bool = False, -) -> None: - instructions = ["instruction 0", "instruction 1"] - completions_a = ["completion-a-0", "completion-a-1"] - completions_b = ["completion-b-0", "completion-b-1"] - resolved = resolve_run_judge_prompt(cfg.task, cfg.judge) - rendered = render_judge_inputs( - instructions, - completions_a, - completions_b, - system_prompt=resolved.system_prompt, - user_prompt_template=resolved.user_prompt_template, - truncate_input_chars=cfg.generation.truncate_judge_input_chars, - provide_explanation=cfg.judge.provide_explanation, - prompt_preset=resolved.preset_name, - strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging, - task=cfg.task, - ) - rows = [] - for idx, (instruction, ca, cb, judge_input) in enumerate( - zip(instructions, completions_a, completions_b, rendered, strict=True) - ): - rows.append( - { - "instruction": instruction, - "completion_A": ca, - "completion_B": cb, - "judge_completion": "Score A: 8\nScore B: 6", - "judge_input": judge_input.to_string(), - "instruction_index": f"idx-{idx}", - "model_A": cfg.model.name, - "model_B": cfg.model.baseline, - "judge": cfg.judge.model, - } - ) - if swap_both: - reversed_rendered = render_judge_inputs( - instructions, - completions_b, - completions_a, - system_prompt=resolved.system_prompt, - user_prompt_template=resolved.user_prompt_template, - truncate_input_chars=cfg.generation.truncate_judge_input_chars, - provide_explanation=cfg.judge.provide_explanation, - prompt_preset=resolved.preset_name, - strip_thinking_before_judging=cfg.judge.strip_thinking_before_judging, - task=cfg.task, - ) - for idx, (instruction, ca, cb, judge_input) in enumerate( - zip( - instructions, - completions_b, - completions_a, - reversed_rendered, - strict=True, - ) - ): - rows.append( - { - "instruction": instruction, - "completion_A": ca, - "completion_B": cb, - "judge_completion": "Score A: 6\nScore B: 8", - "judge_input": judge_input.to_string(), - "instruction_index": f"idx-{idx}", - "model_A": cfg.model.baseline, - "model_B": cfg.model.name, - "judge": cfg.judge.model, - } - ) - run_dir.mkdir(parents=True, exist_ok=True) - dump_config(cfg, run_dir / "config.yaml") - pd.DataFrame(rows).to_csv(run_dir / "pair-annotations.csv", index=False) - write_run_metadata( - output_dir=run_dir, - entrypoint="tests", - run=cfg.model_dump(), - results={"n": len(rows)}, - input_payloads={ - "instruction_index": [row["instruction_index"] for row in rows] - }, - judge_system_prompt=resolved.system_prompt, - judge_user_prompt_template=resolved.user_prompt_template, - ) - - -def _write_legacy_gae_run(run_dir: Path) -> None: - run_dir.mkdir(parents=True, exist_ok=True) - args = { - "task": "alpaca-eval", - "model_A": "Dummy/gen-a", - "model_B": "Dummy/gen-b", - "judge_model": "Dummy/score A: 0 score B: 10", - "swap_mode": "fixed", - "provide_explanation": False, - "truncate_all_input_chars": 8192, - "engine_kwargs": {}, - } - (run_dir / "args-alpaca.json").write_text(json.dumps(args)) - cfg = RunConfig( - task=args["task"], - model={"name": args["model_A"], "baseline": args["model_B"]}, - judge={"model": args["judge_model"], "swap_mode": args["swap_mode"]}, - generation={"truncate_judge_input_chars": args["truncate_all_input_chars"]}, - ) - _write_gae_annotations(run_dir, cfg) - - -def _count_judge_inference_rows(store_root: Path, judge_model: str) -> int: - total = 0 - for db_path in store_root.rglob(INFERENCE_DB_NAME): - metadata = json.loads( - (db_path.parent / "metadata.json").read_text(encoding="utf-8") - ) - if metadata.get("model_spec") != judge_model: - continue - with SQLiteInferenceStore(db_path) as store: - total += len(store.query()) - return total - - -def test_discover_skips_elo_and_legacy_artifacts(tmp_path): - elo_dir = tmp_path / "elo-lmarena-100k-run" - elo_dir.mkdir() - (elo_dir / "results.json").write_text("{}") - - legacy_db = tmp_path / "cache" / "db" / "arena" / "judge.db" - legacy_db.parent.mkdir(parents=True) - legacy_db.write_text("sqlite") - - pass_cache = tmp_path / "tables" / "model_outputs" / "alpaca-eval.csv.zip" - pass_cache.parent.mkdir(parents=True) - pass_cache.write_text("zip") - - report = discover_sources([tmp_path]) - skipped_kinds = {item.kind for item in report.skipped} - assert ArtifactKind.META_EVAL_IDENTITY_DB in skipped_kinds - assert ArtifactKind.PASS_LEVEL_CACHE in skipped_kinds - assert ArtifactKind.ELO_RUN in skipped_kinds - - -def test_discover_classifies_standalone_completions_as_generation(tmp_path): - artifact = tmp_path / "completions.parquet" - artifact.write_bytes(b"not imported") - - report = discover_sources([artifact]) - - assert [item.kind for item in report.skipped] == [ArtifactKind.GENERATION_ARTIFACT] - - -def test_gae_current_config_backfill_and_idempotency(tmp_path): - cfg = _cfg_with_cache(tmp_path) - run_dir = tmp_path / "results" / "gae-run" - _write_gae_annotations(run_dir, cfg) - - store_root = tmp_path / "backfill-store" - first = backfill_sources([run_dir], store_root) - assert first.written == 2 - assert first.existing == 0 - - second = backfill_sources([run_dir], store_root) - assert second.written == 0 - assert second.existing == 2 - - db_path = next(store_root.rglob(INFERENCE_DB_NAME)) - with SQLiteInferenceStore(db_path) as store: - rows = store.query() - assert all(row["pushed_by"] == BACKFILL_PUSHED_BY for _, row in rows.iterrows()) - - -def test_legacy_gae_args_backfill_parity(tmp_path): - run_dir = tmp_path / "legacy-run" - _write_legacy_gae_run(run_dir) - store_root = tmp_path / "store" - report = backfill_sources([run_dir], store_root) - assert report.written == 2 - assert report.sources["gae"]["runs"] == 1 - - -def test_gae_swap_rows_backfill(tmp_path): - cfg = _cfg_with_cache( - tmp_path, - judge={"model": "Dummy/score A: 0 score B: 10", "swap_mode": "both"}, - ) - run_dir = tmp_path / "swap-run" - _write_gae_annotations(run_dir, cfg, swap_both=True) - - report = backfill_sources([run_dir], tmp_path / "store") - assert report.written == 4 - - -def test_gae_judge_input_mismatch_skipped(tmp_path): - cfg = _cfg_with_cache(tmp_path) - run_dir = tmp_path / "mismatch-run" - _write_gae_annotations(run_dir, cfg) - csv_path = run_dir / "pair-annotations.csv" - df = pd.read_csv(csv_path) - df.loc[0, "judge_input"] = "tampered" - df.to_csv(csv_path, index=False) - - report = backfill_sources([run_dir], tmp_path / "store") - assert report.written == 1 - assert report.skipped["judge_input_mismatch"] == 1 - - -def test_gae_missing_judge_output_is_not_backfilled(tmp_path): - cfg = _cfg_with_cache(tmp_path) - run_dir = tmp_path / "missing-output" - _write_gae_annotations(run_dir, cfg) - csv_path = run_dir / "pair-annotations.csv" - df = pd.read_csv(csv_path, keep_default_na=False) - df.loc[0, "judge_completion"] = "" - df.to_csv(csv_path, index=False) - - report = backfill_sources([run_dir], tmp_path / "store") - - assert report.written == 1 - assert report.skipped["judge_output_missing"] == 1 - - -def test_multiple_annotations_files_fail_closed(tmp_path): - cfg = _cfg_with_cache(tmp_path) - run_dir = tmp_path / "multiple-annotations" - _write_gae_annotations(run_dir, cfg) - (run_dir / "other-annotations.csv").write_bytes( - (run_dir / "pair-annotations.csv").read_bytes() - ) - - report = backfill_sources([run_dir], tmp_path / "store") - - assert report.written == 0 - assert report.skipped["source_extraction_failed"] == 1 - - -def test_conflicting_modern_configs_fail_closed(tmp_path): - cfg = _cfg_with_cache(tmp_path) - run_dir = tmp_path / "conflicting-configs" - _write_gae_annotations(run_dir, cfg) - conflicting = cfg.model_copy(deep=True) - conflicting.judge.model = "Dummy/different-judge" - dump_config(conflicting, run_dir / "config.yaml") - - report = backfill_sources([run_dir], tmp_path / "store") - - assert report.written == 0 - assert report.skipped["source_extraction_failed"] == 1 - - -def test_gae_ambiguous_orientation_is_skipped(tmp_path): - cfg = _cfg_with_cache(tmp_path) - run_dir = tmp_path / "ambiguous-orientation" - _write_gae_annotations(run_dir, cfg) - csv_path = run_dir / "pair-annotations.csv" - df = pd.read_csv(csv_path) - df.loc[0, ["model_A", "model_B"]] = ["other-a", "other-b"] - df.to_csv(csv_path, index=False) - - report = backfill_sources([run_dir], tmp_path / "store") - - assert report.written == 1 - assert report.skipped["battle_orientation_unverifiable"] == 1 - - -def test_gae_missing_judge_input_fail_closed(tmp_path): - cfg = _cfg_with_cache(tmp_path) - run_dir = tmp_path / "missing-input-run" - _write_gae_annotations(run_dir, cfg) - csv_path = run_dir / "pair-annotations.csv" - df = pd.read_csv(csv_path) - df = df.drop(columns=["judge_input"]) - df.to_csv(csv_path, index=False) - - report = backfill_sources([run_dir], tmp_path / "store") - assert report.written == 0 - assert report.skipped.get("unknown_judge_run", 0) == 1 - - -def test_local_engine_skipped_without_init(tmp_path, monkeypatch): - cfg = _cfg_with_cache( - tmp_path, - judge={"model": "VLLM/Qwen/Qwen2.5-0.5B-Instruct", "swap_mode": "fixed"}, - ) - run_dir = tmp_path / "vllm-run" - _write_gae_annotations(run_dir, cfg) - - def fail_init(*args, **kwargs): - raise AssertionError("VLLM should not initialize during backfill") - - monkeypatch.setattr( - "judgearena.models.ChatVLLM.__init__", - fail_init, - ) - - report = backfill_sources([run_dir], tmp_path / "store") - assert report.written == 0 - assert report.skipped["local_engine_unsupported"] == 2 - - -def test_mt_preset_backfill(tmp_path): - cfg = _cfg_with_cache( - tmp_path, - task="mt-bench", - model={"name": "Dummy/a", "baseline": "Dummy/b"}, - judge={"model": "Dummy/judge-output", "swap_mode": "fixed"}, - ) - run_dir = tmp_path / "mt-preset" - run_dir.mkdir(parents=True) - dump_config(cfg, run_dir / "config.yaml") - pd.DataFrame( - [ - { - "question_id": 1, - "category": "writing", - "turn": 1, - "model_A": "Dummy/a", - "model_B": "Dummy/b", - "judge": "Dummy/judge-output", - "prompt_name": "default-single", - "system_prompt": "system", - "user_prompt": "user body", - "judge_completion": "Score A: 8\nScore B: 6", - "swapped": False, - } - ] - ).to_csv(run_dir / "mt-annotations.csv", index=False) - - report = backfill_sources([run_dir], tmp_path / "store") - assert report.written == 1 - assert report.sources["mt_bench"]["runs"] == 1 - - -def test_mt_fastchat_g1_g2_backfill(tmp_path): - cfg = _cfg_with_cache( - tmp_path, - task="mt-bench", - model={"name": "Dummy/a", "baseline": "Dummy/b"}, - judge={"model": "Dummy/judge-output", "swap_mode": "both"}, - ) - run_dir = tmp_path / "mt-fastchat" - run_dir.mkdir(parents=True) - dump_config(cfg, run_dir / "config.yaml") - pd.DataFrame( - [ - { - "question_id": 1, - "category": "writing", - "turn": 1, - "model_A": "Dummy/a", - "model_B": "Dummy/b", - "judge": "Dummy/judge-output", - "prompt_name": "pair-v2", - "system_prompt": "system", - "g1_user_prompt": " g1 user ", - "g1_judgment": " [[A]] ", - "g2_user_prompt": "g2 user", - "g2_judgment": "[[B]]", - } - ] - ).to_csv(run_dir / "mt-annotations.csv", index=False) - - report = backfill_sources([run_dir], tmp_path / "store") - assert report.written == 2 - extraction = extract_mt_bench_rows(run_dir) - assert " g1 user " in extraction.rows[0].canonical_input - assert extraction.rows[0].output_text == " [[A]] " - - -def test_meta_eval_forward_and_swapped(tmp_path): - args = CliMetaEvalArgs( - judge_model="Dummy/meta-judge", - reference_arena="LMArena-140k", - prompt_mode="standard", - swap_mode="both", - ) - prompt_spec = resolve_prompt_mode(args.prompt_mode, provide_explanation=False) - forward_rendered = render_judge_inputs( - ["instruction"], - ["completion a"], - ["completion b"], - system_prompt=prompt_spec.system_prompt, - user_prompt_template=prompt_spec.user_prompt_template, - truncate_input_chars=args.truncate_judge_input_chars, - provide_explanation=False, - )[0].to_string() - swapped_rendered = render_judge_inputs( - ["instruction"], - ["completion b"], - ["completion a"], - system_prompt=prompt_spec.system_prompt, - user_prompt_template=prompt_spec.user_prompt_template, - truncate_input_chars=args.truncate_judge_input_chars, - provide_explanation=False, - )[0].to_string() - - run_dir = tmp_path / "meta-run" - run_dir.mkdir(parents=True) - (run_dir / "args.json").write_text(json.dumps(args.to_jsonable())) - pd.DataFrame( - [ - { - "question_id": "q1", - "benchmark": "arena", - "model_a": "m-a", - "model_b": "m-b", - "instruction": "instruction", - "completion_a": "completion a", - "completion_b": "completion b", - "presented_completion_a": "completion a", - "presented_completion_b": "completion b", - "judge_input": forward_rendered, - "judge_completion": "Score A: 8\nScore B: 6", - "orientation": "forward", - "presented_model_a": "m-a", - "presented_model_b": "m-b", - }, - { - "question_id": "q1", - "benchmark": "arena", - "model_a": "m-a", - "model_b": "m-b", - "instruction": "instruction", - "completion_a": "completion a", - "completion_b": "completion b", - "presented_completion_a": "completion b", - "presented_completion_b": "completion a", - "judge_input": swapped_rendered, - "judge_completion": "Score A: 6\nScore B: 8", - "orientation": "swapped", - "presented_model_a": "m-b", - "presented_model_b": "m-a", - }, - ] - ).to_parquet(run_dir / "annotations.parquet") - - report = backfill_sources([run_dir], tmp_path / "store") - assert report.written == 2 - assert report.rows_planned == 2 - task = meta_eval_cache_task(args.reference_arena) - assert any(task in str(path) for path in (tmp_path / "store").rglob("*")) - db_path = next((tmp_path / "store").rglob(INFERENCE_DB_NAME)) - with SQLiteInferenceStore(db_path) as store: - meta_rows = store.query_metadata() - assert all( - "source_run_id" in json.loads(row["metadata_json"]) - for _, row in meta_rows.iterrows() - ) - assert all( - "source_run_folder" not in json.loads(row["metadata_json"]) - for _, row in meta_rows.iterrows() - ) - - -def test_conflicting_outputs_skipped(tmp_path): - cfg = _cfg_with_cache(tmp_path) - run_a = tmp_path / "run-a" - run_b = tmp_path / "run-b" - _write_gae_annotations(run_a, cfg) - _write_gae_annotations(run_b, cfg) - df = pd.read_csv(run_b / "pair-annotations.csv") - df.loc[0, "judge_completion"] = "Score A: 1\nScore B: 9" - df.to_csv(run_b / "pair-annotations.csv", index=False) - - report = backfill_sources([run_a, run_b], tmp_path / "store") - assert report.skipped.get("conflicting_outputs", 0) == 2 - assert _count_judge_inference_rows(tmp_path / "store", cfg.judge.model) <= 2 - - -def test_identical_outputs_preserve_all_run_metadata(tmp_path): - cfg = _cfg_with_cache(tmp_path) - run_a = tmp_path / "run-a" - run_b = tmp_path / "run-b" - _write_gae_annotations(run_a, cfg) - _write_gae_annotations(run_b, cfg) - - store_root = tmp_path / "store" - report = backfill_sources([run_a, run_b], store_root) - - assert report.written == 2 - db_path = next(store_root.rglob(INFERENCE_DB_NAME)) - with SQLiteInferenceStore(db_path) as store: - metadata = store.query_metadata() - source_ids = { - json.loads(value)["source_run_id"] for value in metadata["metadata_json"] - } - assert source_ids == {"run-a", "run-b"} - assert len(metadata) == 4 - - -def test_dry_run_writes_nothing(tmp_path): - cfg = _cfg_with_cache(tmp_path) - run_dir = tmp_path / "dry-run" - _write_gae_annotations(run_dir, cfg) - report = backfill_sources([run_dir], tmp_path / "store", dry_run=True) - assert report.written == 2 - assert _count_judge_inference_rows(tmp_path / "store", cfg.judge.model) == 0 - - -def test_backfill_cli(tmp_path): - cfg = _cfg_with_cache(tmp_path) - run_dir = tmp_path / "cli-run" - _write_gae_annotations(run_dir, cfg) - report_path = tmp_path / "report.json" - - cache_sync.main( - [ - "backfill", - str(run_dir), - "--store_root", - str(tmp_path / "store"), - "--report", - str(report_path), - ] - ) - - payload = json.loads(report_path.read_text(encoding="utf-8")) - assert payload["written"] == 2 - assert _count_judge_inference_rows(tmp_path / "store", cfg.judge.model) == 2 - - -@pytest.fixture -def mock_gae_inputs(monkeypatch): - monkeypatch.setattr( - gae, - "load_instructions", - lambda dataset, n_instructions=None: _synthetic_instructions( - n_instructions or 2 - ), - ) - monkeypatch.setattr(gae, "try_load_dataset_completions", lambda *args: None) - - -def test_gae_live_run_backfill_reuses_cells(mock_gae_inputs, tmp_path): - cfg = _cfg_with_cache(tmp_path) - gae.main(cfg) - run_dir = next((tmp_path / "results").iterdir()) - - live_count = _count_judge_inference_rows(tmp_path / "live-cache", cfg.judge.model) - backfill_store = tmp_path / "backfill-store" - report = backfill_sources([run_dir], backfill_store) - assert report.written == live_count - assert _count_judge_inference_rows(backfill_store, cfg.judge.model) == live_count - - -def test_write_report_roundtrip(tmp_path): - report = BackfillReport( - written=3, - existing=1, - rows_planned=4, - skipped={"judge_input_mismatch": 2}, - ) - path = tmp_path / "nested" / "report.json" - write_report(report, path) - payload = json.loads(path.read_text(encoding="utf-8")) - assert payload["written"] == 3 - assert payload["rows_planned"] == 4 - assert payload["skipped"]["judge_input_mismatch"] == 2 - - -def test_gae_orientation_inference_reversed_rows(): - cfg = RunConfig( - task="alpaca-eval", - model={"name": "Dummy/gen-a", "baseline": "Dummy/gen-b"}, - judge={"model": "Dummy/judge"}, - ) - direct = pd.Series({"model_A": "Dummy/gen-a", "model_B": "Dummy/gen-b"}) - reversed_row = pd.Series({"model_A": "Dummy/gen-b", "model_B": "Dummy/gen-a"}) - assert _infer_gae_orientation(direct, cfg=cfg) == "direct" - assert _infer_gae_orientation(reversed_row, cfg=cfg) == "reversed" - - -def test_source_extraction_failure_continues_other_runs(tmp_path): - cfg = _cfg_with_cache(tmp_path) - good_run = tmp_path / "good-run" - bad_run = tmp_path / "bad-run" - _write_gae_annotations(good_run, cfg) - bad_run.mkdir() - pd.DataFrame( - { - "instruction": ["instruction 0"], - "completion_A": ["a"], - "completion_B": ["b"], - "judge_input": ["prompt"], - "judge_completion": ["Score A: 1\nScore B: 0"], - } - ).to_csv(bad_run / "pair-annotations.csv", index=False) - - report = backfill_sources([bad_run, good_run], tmp_path / "store") - assert report.skipped["source_extraction_failed"] == 1 - assert report.written == 2 - - -def test_unknown_annotation_csv_skipped_not_migrated(tmp_path): - run_dir = tmp_path / "custom-run" - run_dir.mkdir() - pd.DataFrame([{"foo": 1, "bar": 2}]).to_csv( - run_dir / "custom-annotations.csv", index=False - ) - report = discover_sources([run_dir]) - assert report.migratable_runs == [] - assert any(item.kind == ArtifactKind.UNKNOWN for item in report.skipped) - - -def test_conflicting_existing_output_skips_metadata(tmp_path): - cfg = _cfg_with_cache(tmp_path) - run_dir = tmp_path / "gae-run" - _write_gae_annotations(run_dir, cfg) - - store_root = tmp_path / "store" - first = backfill_sources([run_dir], store_root) - assert first.written == 2 - - df = pd.read_csv(run_dir / "pair-annotations.csv") - df.loc[0, "judge_completion"] = "Score A: 1\nScore B: 9" - df.to_csv(run_dir / "pair-annotations.csv", index=False) - - second = backfill_sources([run_dir], store_root) - assert second.skipped.get("conflicting_existing_output", 0) == 1 - db_path = next(store_root.rglob(INFERENCE_DB_NAME)) - with SQLiteInferenceStore(db_path) as store: - rows = store.query() - assert rows.iloc[0]["output_text"] == "Score A: 8\nScore B: 6" - - -def test_legacy_multiple_args_files_fail_closed(tmp_path): - run_dir = tmp_path / "ambiguous-legacy" - run_dir.mkdir() - (run_dir / "args-a.json").write_text(json.dumps({"task": "alpaca-eval"})) - (run_dir / "args-b.json").write_text(json.dumps({"task": "alpaca-eval"})) - pd.DataFrame( - { - "instruction": ["x"], - "completion_A": ["a"], - "completion_B": ["b"], - "judge_input": ["prompt"], - "judge_completion": ["Score A: 1\nScore B: 0"], - } - ).to_csv(run_dir / "pair-annotations.csv", index=False) - - report = backfill_sources([run_dir], tmp_path / "store") - assert report.skipped["source_extraction_failed"] == 1 - assert report.written == 0 - - -def test_legacy_truncate_all_input_chars_used_for_judge(tmp_path): - run_dir = tmp_path / "legacy-truncate" - _write_legacy_gae_run(run_dir) - extraction = extract_gae_rows(run_dir) - assert extraction.rows - assert extraction.skipped.get("judge_input_mismatch", 0) == 0 - - -def test_nested_legacy_db_does_not_block_parent_run_discovery(tmp_path): - cfg = _cfg_with_cache(tmp_path) - run_dir = tmp_path / "results" / "nested-run" - _write_gae_annotations(run_dir, cfg) - legacy_dir = tmp_path / "results" / "cache" / "nested" - legacy_dir.mkdir(parents=True) - (legacy_dir / "judgements.db").write_text("sqlite") - - report = discover_sources([tmp_path / "results"]) - migratable = {item.run_dir.name for item in report.migratable_runs} - assert "nested-run" in migratable - skipped_kinds = {item.kind for item in report.skipped} - assert ArtifactKind.LEGACY_CACHE_CELL in skipped_kinds - - backfill_report = backfill_sources([tmp_path / "results"], tmp_path / "store") - assert backfill_report.written == 2 - - -def test_mt_swapped_string_and_nan_prompts(tmp_path): - cfg = _cfg_with_cache( - tmp_path, - task="mt-bench", - model={"name": "Dummy/a", "baseline": "Dummy/b"}, - judge={"model": "Dummy/judge-output", "swap_mode": "fixed"}, - ) - run_dir = tmp_path / "mt-swapped-string" - run_dir.mkdir(parents=True) - dump_config(cfg, run_dir / "config.yaml") - pd.DataFrame( - [ - { - "question_id": 1, - "category": "writing", - "turn": 1, - "model_A": "Dummy/a", - "model_B": "Dummy/b", - "judge": "Dummy/judge-output", - "prompt_name": "default-single", - "system_prompt": float("nan"), - "user_prompt": "user body", - "judge_completion": "Score A: 8\nScore B: 6", - "swapped": "true", - } - ] - ).to_csv(run_dir / "mt-annotations.csv", index=False) - - report = backfill_sources([run_dir], tmp_path / "store") - assert report.written == 1 - db_path = next((tmp_path / "store").rglob(INFERENCE_DB_NAME)) - with SQLiteInferenceStore(db_path) as store: - row = store.query().iloc[0] - assert "nan" not in row["input_text"].lower() - - -def test_dry_run_does_not_touch_existing_db(tmp_path): - cfg = _cfg_with_cache(tmp_path) - run_dir = tmp_path / "dry-run-existing" - _write_gae_annotations(run_dir, cfg) - store_root = tmp_path / "store" - backfill_sources([run_dir], store_root, dry_run=False) - db_path = next(store_root.rglob(INFERENCE_DB_NAME)) - before = db_path.stat().st_mtime_ns - wal_path = Path(f"{db_path}-wal") - wal_before_exists = wal_path.exists() - - report = backfill_sources([run_dir], store_root, dry_run=True) - assert report.written == 0 - assert report.existing == 2 - assert db_path.stat().st_mtime_ns == before - assert wal_path.exists() == wal_before_exists - - -def test_cell_integrity_error_is_reported(tmp_path, monkeypatch): - cfg = _cfg_with_cache(tmp_path) - run_dir = tmp_path / "integrity-run" - _write_gae_annotations(run_dir, cfg) - - def fail_metadata(*args, **kwargs): - raise ValueError("metadata mismatch") - - monkeypatch.setattr( - "judgearena.cache_backfill.write_store_metadata", - fail_metadata, - ) - report = backfill_sources([run_dir], tmp_path / "store") - assert report.skipped.get("cell_integrity_error", 0) == 2 - assert report.written == 0 diff --git a/tests/test_cache_sync.py b/tests/test_cache_sync.py index 601944c..dbb8911 100644 --- a/tests/test_cache_sync.py +++ b/tests/test_cache_sync.py @@ -1,11 +1,9 @@ -import getpass import json -from pathlib import Path import pandas as pd import pytest -from judgearena import cache_sync +from judgearena.cache_sync import main from judgearena.store_sqlite import ( INFERENCE_DB_NAME, SQLiteInferenceStore, @@ -16,116 +14,65 @@ REPO_ID = "org/cache" CELL_CONFIG = {"task": "arena", "model_spec": "VLLM/Qwen/judge"} -CELL_CONFIG_HASH = descriptor_hash(CELL_CONFIG) -MODEL_SPEC = "VLLM/Qwen/judge" -PATH_IN_REPO = ( - f"inference/arena/VLLM/Qwen%2Fjudge/{CELL_CONFIG_HASH}/{INFERENCE_DB_NAME}" -) -METADATA_IN_REPO = f"inference/arena/VLLM/Qwen%2Fjudge/{CELL_CONFIG_HASH}/metadata.json" - - -def _local_cell_db(tmp_path) -> Path: - cell_dir = store_folder(tmp_path, "arena", MODEL_SPEC, CELL_CONFIG_HASH) - return cell_dir / INFERENCE_DB_NAME +CELL_HASH = descriptor_hash(CELL_CONFIG) +CELL_PREFIX = f"inference/arena/VLLM/Qwen%2Fjudge/{CELL_HASH}" +DB_IN_REPO = f"{CELL_PREFIX}/{INFERENCE_DB_NAME}" +METADATA_IN_REPO = f"{CELL_PREFIX}/metadata.json" -def _write_inference(path: Path) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - with SQLiteInferenceStore(path) as store: +def _write_cell(store_root): + folder = store_folder(store_root, "arena", "VLLM/Qwen/judge", CELL_HASH) + write_store_metadata(folder, CELL_CONFIG) + with SQLiteInferenceStore(folder / INFERENCE_DB_NAME) as store: store.save_outputs( pd.DataFrame( { - "input_hash": ["local"], - "input_text": ["input-local"], - "output_text": ["L"], + "input_hash": ["h1"], + "input_text": ["input"], + "output_text": ["output"], } ), pushed_by="test", ) + return folder / INFERENCE_DB_NAME -def test_fetch_requires_filter(tmp_path, capsys): - with pytest.raises(SystemExit) as exc: - cache_sync.main(["fetch", "--store_root", str(tmp_path)]) - assert exc.value.code == 1 - assert "requires a path filter" in capsys.readouterr().err - - -def test_fetch_rejects_invalid_filter_gaps(tmp_path, capsys): - with pytest.raises(SystemExit) as exc: - cache_sync.main( - [ - "fetch", - "--store_root", - str(tmp_path), - "--provider", - "VLLM", - ] - ) - assert exc.value.code == 1 - assert "requires --task" in capsys.readouterr().err +def test_fetch_requires_filter(tmp_path): + with pytest.raises(SystemExit, match="1"): + main(["fetch", "--store_root", str(tmp_path)]) -def test_fetch_bootstraps_filtered_remote_cells(fake_hub, tmp_path): - remote = tmp_path / "remote.db" - _write_inference(remote) - fake_hub.files[PATH_IN_REPO] = remote.read_bytes() - fake_hub.files[METADATA_IN_REPO] = json.dumps(CELL_CONFIG).encode("utf-8") +def test_fetch_bootstraps_filtered_remote_cell(fake_hub, tmp_path): + remote_db = _write_cell(tmp_path / "remote") + fake_hub.files[DB_IN_REPO] = remote_db.read_bytes() + fake_hub.files[METADATA_IN_REPO] = json.dumps(CELL_CONFIG).encode() fake_hub.head = "initial" + local_root = tmp_path / "local" - store_root = tmp_path / "store" - cache_sync.main( + main( [ "fetch", "--store_root", - str(store_root), + str(local_root), "--cache_hf_repo", REPO_ID, "--task", "arena", ] ) - local_db = _local_cell_db(store_root) - assert local_db.exists() - assert (local_db.parent / "metadata.json").exists() - - -def test_push_uploads_local_cells(fake_hub, tmp_path): - local_db = _local_cell_db(tmp_path) - _write_inference(local_db) - write_store_metadata(local_db.parent, CELL_CONFIG) - fake_hub.head = "initial" - cache_sync.main( - [ - "push", - "--store_root", - str(tmp_path), - "--cache_hf_repo", - REPO_ID, - "--task", - "arena", - ] - ) - assert PATH_IN_REPO in fake_hub.files - assert METADATA_IN_REPO in fake_hub.files - assert fake_hub.commit_calls == 1 + with SQLiteInferenceStore( + store_folder(local_root, "arena", "VLLM/Qwen/judge", CELL_HASH) + / INFERENCE_DB_NAME + ) as store: + assert store.query()["output_text"].tolist() == ["output"] -def test_push_defaults_pushed_by_to_current_user(fake_hub, tmp_path, monkeypatch): - local_db = _local_cell_db(tmp_path) - _write_inference(local_db) - write_store_metadata(local_db.parent, CELL_CONFIG) +def test_push_uploads_local_cells(fake_hub, tmp_path): fake_hub.head = "initial" - observed: list[str] = [] - - def capture_push(*args, **kwargs): - observed.append(kwargs.get("pushed_by", args[3] if len(args) > 3 else None)) - - monkeypatch.setattr(cache_sync, "push_cells", capture_push) - monkeypatch.setattr(getpass, "getuser", lambda: "unit-test-user") + _write_cell(tmp_path) - cache_sync.main( + main( [ "push", "--store_root", @@ -136,25 +83,6 @@ def capture_push(*args, **kwargs): "arena", ] ) - assert observed == ["unit-test-user"] - -def test_push_create_pr(fake_hub, tmp_path): - local_db = _local_cell_db(tmp_path) - _write_inference(local_db) - write_store_metadata(local_db.parent, CELL_CONFIG) - fake_hub.head = "initial" - - cache_sync.main( - [ - "push", - "--store_root", - str(tmp_path), - "--cache_hf_repo", - REPO_ID, - "--task", - "arena", - "--create_pr", - ] - ) - assert PATH_IN_REPO not in fake_hub.files + assert DB_IN_REPO in fake_hub.files + assert METADATA_IN_REPO in fake_hub.files diff --git a/tests/test_cli.py b/tests/test_cli.py index 44a4d79..8f539b8 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -303,38 +303,7 @@ def test_judge_side_kwargs_parsed_separately(capture_mains): assert cfg.judge.engine_kwargs == {"tensor_parallel_size": 4} -def test_cache_nested_cli_flags(capture_mains, tmp_path, monkeypatch): - monkeypatch.setattr("getpass.getuser", lambda: "cli-user") - store = tmp_path / "store" - cli_module.cli( - [ - "--task", - "alpaca-eval", - "--model.name", - "Dummy/A", - "--model.baseline", - "Dummy/B", - "--judge.model", - "Dummy/J", - "--cache.store_root", - str(store), - "--cache.cache_mode", - "refresh", - "--cache.cache_fetch", - "--cache.cache_push", - "--cache.pushed_by", - "cli-user", - ] - ) - cfg = capture_mains["cfg"] - assert cfg.cache.store_root == str(store) - assert cfg.cache.cache_mode == "refresh" - assert cfg.cache.cache_fetch is True - assert cfg.cache.cache_push is True - assert cfg.cache.pushed_by == "cli-user" - - -def test_cache_flat_cli_aliases(capture_mains, tmp_path): +def test_cache_cli_flags(capture_mains, tmp_path): store = tmp_path / "store" cli_module.cli( [ @@ -349,43 +318,19 @@ def test_cache_flat_cli_aliases(capture_mains, tmp_path): "--store_root", str(store), "--cache_mode", - "use", - "--cache_hf_repo", - "org/repo", + "refresh", "--cache_fetch", - "--cache_create_pr", "--cache_push", + "--pushed_by", + "cli-user", ] ) cfg = capture_mains["cfg"] assert cfg.cache.store_root == str(store) - assert cfg.cache.cache_mode == "use" - assert cfg.cache.cache_hf_repo == "org/repo" + assert cfg.cache.cache_mode == "refresh" assert cfg.cache.cache_fetch is True assert cfg.cache.cache_push is True - assert cfg.cache.cache_create_pr is True - - -def test_cache_cli_overrides_yaml_fetch(tmp_path, capture_mains): - yaml_path = tmp_path / "run.yaml" - yaml_path.write_text( - "task: alpaca-eval\n" - "model: {name: Dummy/A, baseline: Dummy/B}\n" - "judge: {model: Dummy/J}\n" - "cache:\n" - " store_root: /yaml/store\n" - " cache_fetch: true\n" - ) - cli_module.cli( - [ - "--config_path", - str(yaml_path), - "--no-cache_fetch", - ] - ) - cfg = capture_mains["cfg"] - assert cfg.cache.store_root == "/yaml/store" - assert cfg.cache.cache_fetch is False + assert cfg.cache.pushed_by == "cli-user" def test_cache_fetch_without_store_root_errors(capture_mains): diff --git a/tests/test_config.py b/tests/test_config.py index c307354..aefb8f1 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -266,35 +266,6 @@ def test_cache_defaults(): assert cfg.cache.cache_create_pr is False -def test_cache_pushed_by_defaults_to_getuser(monkeypatch): - monkeypatch.setattr("judgearena.config.default_pushed_by", lambda: "test-user") - cfg = RunConfig(**_base_generate()) - assert cfg.cache.pushed_by == "test-user" - - -def test_cache_yaml_load(tmp_path): - from judgearena.config import load_config - - yaml_path = tmp_path / "cache.yaml" - yaml_path.write_text( - "task: alpaca-eval\n" - "model: {name: Dummy/a, baseline: Dummy/b}\n" - "judge: {model: Dummy/j}\n" - "cache:\n" - " store_root: /data/cache\n" - " cache_mode: refresh\n" - " cache_fetch: true\n" - " cache_push: true\n" - " pushed_by: yaml-user\n" - ) - cfg = load_config(yaml_path) - assert cfg.cache.store_root == "/data/cache" - assert cfg.cache.cache_mode == "refresh" - assert cfg.cache.cache_fetch is True - assert cfg.cache.cache_push is True - assert cfg.cache.pushed_by == "yaml-user" - - @pytest.mark.parametrize( ("kwargs", "match"), [ @@ -302,10 +273,6 @@ def test_cache_yaml_load(tmp_path): {"cache_fetch": True}, "cache.store_root is required", ), - ( - {"store_root": "/tmp", "cache_fetch": True, "cache_hf_repo": " "}, - "cache_hf_repo must be non-empty", - ), ( {"store_root": "/tmp", "cache_create_pr": True}, "cache_push is required", @@ -314,14 +281,6 @@ def test_cache_yaml_load(tmp_path): {"store_root": "/tmp", "cache_mode": "off", "cache_fetch": True}, "cache_fetch and cache_push cannot be enabled", ), - ( - {"cache_mode": "refresh"}, - "cache.store_root is required when cache_mode is refresh", - ), - ( - {"store_root": " "}, - "cache.store_root must be non-empty", - ), ], ) def test_cache_validation_rejects_invalid_combinations(kwargs, match): diff --git a/tests/test_estimate_elo_cache_threading.py b/tests/test_estimate_elo_cache_threading.py deleted file mode 100644 index 7b6aaed..0000000 --- a/tests/test_estimate_elo_cache_threading.py +++ /dev/null @@ -1,330 +0,0 @@ -from __future__ import annotations - -import numpy as np -import pandas as pd -import pytest - -import judgearena.estimate_elo_ratings as estimate_elo_ratings -import judgearena.evaluate as evaluate_module -import judgearena.generate as generate_module -import judgearena.models as models_module -from judgearena.config import RunConfig -from judgearena.estimate_elo_ratings import main - - -def _make_conversation(content_user: str, content_assistant: str) -> list[dict]: - return [ - {"role": "user", "content": content_user}, - {"role": "assistant", "content": content_assistant}, - ] - - -@pytest.fixture -def synthetic_arena_df() -> pd.DataFrame: - rng = np.random.default_rng(42) - rows = [] - for i in range(30): - ma, mb = rng.choice( - ["arena_model_alpha", "arena_model_beta", "arena_model_gamma"], - size=2, - replace=False, - ) - rows.append( - { - "question_id": f"q{i}", - "tstamp": 1700000000 + i, - "model_a": ma, - "model_b": mb, - "winner": rng.choice(["model_a", "model_b", "tie"]), - "conversation_a": _make_conversation( - f"Instruction {i}", f"Response A {i}" - ), - "conversation_b": _make_conversation( - f"Instruction {i}", f"Response B {i}" - ), - "benchmark": "TestArena", - "lang": rng.choice(["en", "fr"]), - } - ) - return pd.DataFrame(rows) - - -@pytest.fixture(autouse=True) -def mock_elo_deps(monkeypatch, synthetic_arena_df): - monkeypatch.setattr( - estimate_elo_ratings, - "load_arena_dataframe", - lambda arena: synthetic_arena_df, - ) - - -def _cfg_with_cache(tmp_path, **overrides) -> RunConfig: - payload = { - "task": "elo-comparia", - "model": {"name": "Dummy/my model"}, - "judge": {"model": "Dummy/score A: 0 score B: 10", "swap_mode": "fixed"}, - "generation": {"n_instructions": 5}, - "elo": {"arena": "ComparIA", "n_bootstraps": 2}, - "run": {"result_folder": str(tmp_path / "results"), "no_log_file": True}, - "cache": {"store_root": str(tmp_path / "cache")}, - } - payload.update(overrides) - return RunConfig(**payload) - - -def test_elo_uses_one_shared_cache_handle(monkeypatch, tmp_path): - captured: list[tuple[str, object]] = [] - real_gen = generate_module.do_inference - real_eval = evaluate_module.do_inference - - def spy_gen(*args, **kwargs): - cache = kwargs.get("cache") - if cache is not None: - captured.append(("gen", cache)) - return real_gen(*args, **kwargs) - - def spy_eval(*args, **kwargs): - cache = kwargs.get("cache") - if cache is not None: - captured.append(("eval", cache)) - return real_eval(*args, **kwargs) - - monkeypatch.setattr(generate_module, "do_inference", spy_gen) - monkeypatch.setattr(evaluate_module, "do_inference", spy_eval) - - main(_cfg_with_cache(tmp_path)) - - assert captured - assert len({id(cache) for _, cache in captured}) == 1 - assert any(role == "gen" for role, _ in captured) - assert any(role == "eval" for role, _ in captured) - - -def test_elo_second_run_reuses_cached_rows(monkeypatch, tmp_path): - uncached_calls = {"count": 0} - real_uncached = models_module._do_inference_uncached - - def counting_uncached(*args, **kwargs): - uncached_calls["count"] += 1 - return real_uncached(*args, **kwargs) - - monkeypatch.setattr(models_module, "_do_inference_uncached", counting_uncached) - - cfg = _cfg_with_cache(tmp_path) - first = main(cfg) - assert uncached_calls["count"] > 0 - - uncached_calls["count"] = 0 - second = main(cfg) - assert uncached_calls["count"] == 0 - assert second["winrate"] == pytest.approx(first["winrate"]) - - -def test_elo_judge_row_metadata_includes_arena_and_battle_fields(monkeypatch, tmp_path): - captured_metadata: list[dict] = [] - real_eval = evaluate_module.do_inference - - def spy_eval(*args, **kwargs): - cache_meta = kwargs.get("cache_meta") - if cache_meta is not None: - captured_metadata.extend(cache_meta.get("metadata", [])) - return real_eval(*args, **kwargs) - - monkeypatch.setattr(evaluate_module, "do_inference", spy_eval) - - main( - _cfg_with_cache( - tmp_path, - model={"name": "Dummy/focal-model"}, - ) - ) - - assert captured_metadata - first = captured_metadata[0] - assert first["arena"] == "ComparIA" - assert first["source"] == "elo-judge" - assert first["focal_model"] == "Dummy/focal-model" - assert first["opponent_model"] - assert first["position"] in {"A", "B"} - assert first["question_id"] == "q0" - assert first["orientation"] == "direct" - - -def _configure_calibration_arena(monkeypatch, synthetic_arena_df): - frames = [] - for block in range(20): - chunk = synthetic_arena_df.copy() - chunk["question_id"] = [f"q{block * len(chunk) + j}" for j in range(len(chunk))] - chunk.index = chunk.index + block * len(chunk) - frames.append(chunk) - large_arena_df = pd.concat(frames) - anchor_battles = pd.DataFrame( - { - "model_a": ["arena_model_alpha"] * len(large_arena_df), - "model_b": ["arena_model_beta"] * len(large_arena_df), - "winner": ["model_a", "model_b"] * (len(large_arena_df) // 2), - "pref": [0.0, 1.0] * (len(large_arena_df) // 2), - "pref_hard": [0.0, 1.0] * (len(large_arena_df) // 2), - "source": ["human"] * len(large_arena_df), - "question_id": large_arena_df["question_id"].tolist(), - }, - index=large_arena_df.index, - ) - monkeypatch.setattr( - estimate_elo_ratings, - "arena_anchor_battles", - lambda _df: anchor_battles, - ) - monkeypatch.setattr( - estimate_elo_ratings, - "load_arena_dataframe", - lambda arena: large_arena_df, - ) - return large_arena_df - - -def test_elo_calibration_reuses_shared_cache(monkeypatch, tmp_path, synthetic_arena_df): - _configure_calibration_arena(monkeypatch, synthetic_arena_df) - - captured: list[tuple[str, object]] = [] - real_eval = evaluate_module.do_inference - - def spy_eval(*args, **kwargs): - cache = kwargs.get("cache") - cache_meta = kwargs.get("cache_meta") - if cache is not None: - captured.append(("eval", cache)) - if cache_meta is not None: - for row in cache_meta.get("metadata", []): - if row.get("purpose") == "temperature_calibration": - captured.append(("cal_meta", row)) - return real_eval(*args, **kwargs) - - monkeypatch.setattr(evaluate_module, "do_inference", spy_eval) - - main( - _cfg_with_cache( - tmp_path, - elo={ - "arena": "ComparIA", - "n_bootstraps": 2, - "calibrate_temperature": True, - "calibration_size": 3, - }, - ) - ) - - eval_caches = [cache for role, cache in captured if role == "eval"] - assert eval_caches - assert len({id(cache) for cache in eval_caches}) == 1 - cal_rows = [row for role, row in captured if role == "cal_meta"] - assert cal_rows - assert cal_rows[0]["source"] == "elo-calibration" - assert cal_rows[0]["question_id"] - - -def test_elo_cache_hit_reparses_scores_with_recomputed_calibration( - monkeypatch, tmp_path, synthetic_arena_df -): - _configure_calibration_arena(monkeypatch, synthetic_arena_df) - uncached_calls = {"count": 0} - real_uncached = models_module._do_inference_uncached - - def counting_uncached(*args, **kwargs): - uncached_calls["count"] += 1 - return real_uncached(*args, **kwargs) - - monkeypatch.setattr(models_module, "_do_inference_uncached", counting_uncached) - cfg = _cfg_with_cache( - tmp_path, - judge={"model": "Dummy/score A: 2 score B: 1", "swap_mode": "fixed"}, - elo={ - "arena": "ComparIA", - "n_bootstraps": 1, - "calibrate_temperature": True, - "calibration_size": 20, - }, - ) - - monkeypatch.setattr(estimate_elo_ratings, "calibrate_temperature", lambda *_: 0.5) - first = main(cfg) - assert uncached_calls["count"] > 0 - - uncached_calls["count"] = 0 - monkeypatch.setattr(estimate_elo_ratings, "calibrate_temperature", lambda *_: 5.0) - second = main(cfg) - - assert uncached_calls["count"] == 0 - assert second["elo_mean"] != pytest.approx(first["elo_mean"]) - - -def test_elo_judge_metadata_fallback_without_question_id(monkeypatch, tmp_path): - arena_no_qid = pd.DataFrame( - { - "tstamp": [1700000000], - "model_a": ["arena_model_alpha"], - "model_b": ["arena_model_beta"], - "winner": ["model_a"], - "conversation_a": [ - [ - {"role": "user", "content": "Instruction 0"}, - {"role": "assistant", "content": "Response A 0"}, - ] - ], - "conversation_b": [ - [ - {"role": "user", "content": "Instruction 0"}, - {"role": "assistant", "content": "Response B 0"}, - ] - ], - "benchmark": "TestArena", - "lang": ["en"], - } - ) - monkeypatch.setattr( - estimate_elo_ratings, - "load_arena_dataframe", - lambda arena: arena_no_qid, - ) - - captured_metadata: list[dict] = [] - real_eval = evaluate_module.do_inference - - def spy_eval(*args, **kwargs): - cache_meta = kwargs.get("cache_meta") - if cache_meta is not None: - captured_metadata.extend(cache_meta.get("metadata", [])) - return real_eval(*args, **kwargs) - - monkeypatch.setattr(evaluate_module, "do_inference", spy_eval) - - main( - _cfg_with_cache( - tmp_path, - generation={"n_instructions": 1}, - elo={"arena": "ComparIA", "n_bootstraps": 1}, - ) - ) - - judge_rows = [row for row in captured_metadata if row.get("source") == "elo-judge"] - assert judge_rows - assert "question_id" not in judge_rows[0] - assert judge_rows[0]["battle_identity"] - - -def test_elo_swap_mode_both_doubles_llm_judged_battles(tmp_path): - result = main( - RunConfig( - task="elo-comparia", - model={"name": "Dummy/my model"}, - judge={ - "model": "Dummy/score A: 0 score B: 10", - "swap_mode": "both", - }, - generation={"n_instructions": 4}, - elo={"arena": "ComparIA", "n_bootstraps": 1}, - run={"result_folder": str(tmp_path), "no_log_file": True}, - ) - ) - assert result["llm_judged_battles"] == 8 - assert result["num_battles"] == 4 diff --git a/tests/test_estimate_elo_ratings.py b/tests/test_estimate_elo_ratings.py index 38cf8dc..6407ada 100644 --- a/tests/test_estimate_elo_ratings.py +++ b/tests/test_estimate_elo_ratings.py @@ -85,6 +85,7 @@ def _default_args(*, result_folder: str, **kwargs) -> RunConfig: swap_mode = kwargs.pop("swap_mode", "fixed") strip_thinking_before_judging = kwargs.pop("strip_thinking_before_judging", 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, @@ -100,6 +101,7 @@ def _default_args(*, result_folder: str, **kwargs) -> RunConfig: generation={"n_instructions": n_instructions}, elo={"arena": arena, "n_bootstraps": n_bootstraps, "languages": languages}, run={"result_folder": result_folder}, + cache={"store_root": store_root} if store_root is not None else {}, ) @@ -344,9 +346,16 @@ def test_main_generation_cache_metadata_includes_arena_question_identity( estimate_elo_ratings, "generate_instructions", _spy_generate_capturing(captured) ) - main(_default_args(result_folder=str(tmp_path), arena="ComparIA")) + main( + _default_args( + result_folder=str(tmp_path), + arena="ComparIA", + store_root=str(tmp_path / "cache"), + ) + ) row_metadata = captured["gen_kwargs"]["row_metadata"] + assert captured["cache"] is not None assert len(row_metadata) == 10 assert all(row["arena"] == "ComparIA" for row in row_metadata) assert [row["question_id"] for row in row_metadata] == [ diff --git a/tests/test_evaluate_cache_threading.py b/tests/test_evaluate_cache_threading.py deleted file mode 100644 index e94ccf9..0000000 --- a/tests/test_evaluate_cache_threading.py +++ /dev/null @@ -1,107 +0,0 @@ -from __future__ import annotations - -import math - -import judgearena.evaluate as evaluate_module -from judgearena.evaluate import annotate_battles, judge_and_parse_prefs -from judgearena.inference_cache import InferenceCache -from judgearena.models import make_model - - -class FakeJudge: - def __init__(self, response: str = "score A: 0 score B: 10"): - self.response = response - - def batch(self, *, inputs, **_kwargs): - return [self.response] * len(inputs) - - -def test_annotate_battles_forwards_cache_and_metadata(monkeypatch): - captured: list[dict] = [] - - def spy_do_inference(*, cache, cache_meta, **kwargs): - captured.append({"cache": cache, "cache_meta": cache_meta}) - return ["score A: 0 score B: 10"] - - monkeypatch.setattr(evaluate_module, "do_inference", spy_do_inference) - - row_metadata = [{"battle_id": "b-1"}] - with InferenceCache("/tmp/unused", "judge", mode="off") as cache: - annotate_battles( - judge_chat_model=FakeJudge(), - instructions=["Question"], - completions_A=["A"], - completions_B=["B"], - cache=cache, - row_metadata=row_metadata, - ) - - assert captured[0]["cache"] is cache - assert captured[0]["cache_meta"] == {"metadata": row_metadata} - - -def test_judge_and_parse_prefs_adds_orientation_and_forwards_both(monkeypatch): - captured: list[dict] = [] - - def spy_do_inference(*, cache_meta, **kwargs): - captured.append(cache_meta) - return ["score A: 0 score B: 10"] * len(kwargs["inputs"]) - - monkeypatch.setattr(evaluate_module, "do_inference", spy_do_inference) - - base_metadata = [{"question_id": "q-42"}] - - _, annotations_reversed, prefs = judge_and_parse_prefs( - judge_chat_model=FakeJudge(), - instructions=["Q1", "Q2"], - completions_A=["A1", "A2"], - completions_B=["B1", "B2"], - swap_mode="both", - row_metadata=base_metadata * 2, - ) - - assert annotations_reversed is not None - assert len(captured) == 2 - assert captured[0]["metadata"] == [ - {"question_id": "q-42", "orientation": "direct"}, - {"question_id": "q-42", "orientation": "direct"}, - ] - assert captured[1]["metadata"] == [ - {"question_id": "q-42", "orientation": "reversed"}, - {"question_id": "q-42", "orientation": "reversed"}, - ] - assert len(prefs) == 4 - - -def test_judge_and_parse_prefs_default_without_cache_unchanged(): - judge = make_model("Dummy/score A: 0 score B: 10") - _, annotations_reversed, prefs = judge_and_parse_prefs( - judge_chat_model=judge, - instructions=["Q"], - completions_A=["A"], - completions_B=["B"], - swap_mode="fixed", - ) - - assert annotations_reversed is None - assert len(prefs) == 1 - assert not math.isnan(float(prefs.iloc[0])) - - -def test_annotate_battles_without_metadata_omits_cache_meta(monkeypatch): - captured: list[dict] = [] - - def spy_do_inference(*, cache_meta=None, **kwargs): - captured.append({"cache_meta": cache_meta}) - return ["score A: 0 score B: 10"] - - monkeypatch.setattr(evaluate_module, "do_inference", spy_do_inference) - - annotate_battles( - judge_chat_model=FakeJudge(), - instructions=["Question"], - completions_A=["A"], - completions_B=["B"], - ) - - assert captured[0]["cache_meta"] is None diff --git a/tests/test_generate_and_evaluate_cache_threading.py b/tests/test_generate_and_evaluate_cache_threading.py index 95bda01..7994a6d 100644 --- a/tests/test_generate_and_evaluate_cache_threading.py +++ b/tests/test_generate_and_evaluate_cache_threading.py @@ -153,39 +153,6 @@ def counting_uncached(*args, **kwargs): assert prefs_second.tolist() == prefs_first.tolist() -def test_gae_preloaded_completions_bypass_generation( - mock_gae_inputs, monkeypatch, tmp_path -): - preloaded = pd.DataFrame( - { - "completion": ["preloaded-a", "preloaded-b"], - "instruction_index": [0, 1], - } - ) - generation_calls: list[str] = [] - real_gen = models_module._do_inference_uncached - - def track_generation(chat_model, inputs, **kwargs): - model_spec = getattr(chat_model, "model_spec", None) or getattr( - chat_model, "name", "unknown" - ) - generation_calls.append(str(model_spec)) - return real_gen(chat_model, inputs, **kwargs) - - def load_preloaded(dataset, model, n_instructions): - if model == "Dummy/gen-a": - return preloaded - return None - - monkeypatch.setattr(gae, "try_load_dataset_completions", load_preloaded) - monkeypatch.setattr(models_module, "_do_inference_uncached", track_generation) - - main_generate_and_eval(_cfg_with_cache(tmp_path)) - - assert not any("gen-a" in call for call in generation_calls) - assert any("gen-b" in call for call in generation_calls) - - def test_gae_judge_row_metadata_includes_models_and_instruction_index( mock_gae_inputs, monkeypatch, tmp_path ): diff --git a/tests/test_generate_cache_threading.py b/tests/test_generate_cache_threading.py deleted file mode 100644 index 03bd897..0000000 --- a/tests/test_generate_cache_threading.py +++ /dev/null @@ -1,174 +0,0 @@ -from __future__ import annotations - -import json - -import pandas as pd - -import judgearena.generate as generate_module -from judgearena.generate import generate_base, generate_multiturn -from judgearena.inference_cache import InferenceCache -from judgearena.models import make_model -from judgearena.store_sqlite import SQLiteInferenceStore, descriptor_hash, store_folder - - -def test_generate_base_routes_through_do_inference(monkeypatch): - calls: list[dict] = [] - real_do_inference = generate_module.do_inference - - def spy_do_inference(*args, **kwargs): - calls.append({"args": args, "kwargs": kwargs}) - return real_do_inference(*args, **kwargs) - - monkeypatch.setattr(generate_module, "do_inference", spy_do_inference) - - instructions = pd.Series(["hello", "world"], index=[10, 20]) - df = generate_base(instructions, "Dummy/generate-base-path", use_tqdm=False) - - assert len(calls) == 1 - assert calls[0]["kwargs"]["use_tqdm"] is False - assert calls[0]["kwargs"]["cache"] is None - assert df["completion"].tolist() == ["generate-base-path"] * 2 - assert df["instruction_index"].tolist() == [10, 20] - - -def test_generate_base_forwards_cache_and_metadata(monkeypatch): - captured: list[dict] = [] - - def spy_do_inference(*, cache, cache_meta, **kwargs): - captured.append({"cache": cache, "cache_meta": cache_meta}) - return ["out-a", "out-b"] - - monkeypatch.setattr(generate_module, "do_inference", spy_do_inference) - - instructions = pd.Series(["a", "b"], index=["q1", "q2"]) - with InferenceCache("/tmp/unused", "gen-task", mode="off") as cache: - df = generate_base( - instructions, - "Dummy/ignored", - cache=cache, - ) - - assert df["completion"].tolist() == ["out-a", "out-b"] - assert captured[0]["cache"] is cache - assert captured[0]["cache_meta"] == { - "metadata": [ - {"instruction_index": "q1"}, - {"instruction_index": "q2"}, - ] - } - - -def test_generate_base_cache_hit_skips_backend_batch(tmp_path): - model = make_model("Dummy/cache-generate-base", max_tokens=8) - inputs = ["alpha", "beta"] - descriptor = model.cache_descriptor() - assert descriptor is not None - canonical = [model.canonicalize_input(item) for item in inputs] - metadata = [{"instruction_index": "0"}, {"instruction_index": "1"}] - - with InferenceCache(tmp_path, "gen", mode="refresh") as cache: - cache.get_or_run( - model_spec=model.model_spec, - descriptor=descriptor, - canonical_inputs=canonical, - original_inputs=inputs, - miss_runner=lambda miss_inputs: [f"cached-{item}" for item in miss_inputs], - row_metadata=metadata, - producer_metadata=model.producer_metadata(), - ) - - instructions = pd.Series(inputs, index=[0, 1]) - with InferenceCache(tmp_path, "gen", mode="use") as cache: - df = generate_base( - instructions, - "Dummy/cache-generate-base", - max_tokens=8, - cache=cache, - ) - - assert df["completion"].tolist() == ["cached-alpha", "cached-beta"] - - -def test_generate_multiturn_metadata_and_temperature_groups(monkeypatch): - calls: list[dict] = [] - - def spy_do_inference(*, inputs, cache_meta=None, **kwargs): - calls.append({"inputs": inputs, "cache_meta": cache_meta}) - return [f"out-{index}" for index in range(len(inputs))] - - monkeypatch.setattr(generate_module, "do_inference", spy_do_inference) - - questions = pd.DataFrame( - { - "category": ["writing", "math", "writing"], - "turn_1": ["Q1", "Q2", "Q3"], - "turn_2": ["Q1b", "Q2b", "Q3b"], - }, - index=pd.Index([1, 2, 3], name="instruction_index"), - ) - temperature_config = {"writing": 0.5, "math": 0.9} - - df = generate_multiturn( - questions, - "Dummy/multiturn", - temperature_config=temperature_config, - use_tqdm=False, - ) - - assert len(df) == 3 - assert len(calls) == 4 - - turn1_calls = calls[:2] - turn2_calls = calls[2:] - - assert turn1_calls[0]["cache_meta"]["metadata"] == [ - {"instruction_index": "1", "turn": 1, "category": "writing"}, - {"instruction_index": "3", "turn": 1, "category": "writing"}, - ] - assert turn1_calls[1]["cache_meta"]["metadata"] == [ - {"instruction_index": "2", "turn": 1, "category": "math"}, - ] - assert len(turn1_calls[0]["inputs"]) == 2 - assert len(turn1_calls[1]["inputs"]) == 1 - - assert turn2_calls[0]["cache_meta"]["metadata"] == [ - {"instruction_index": "1", "turn": 2, "category": "writing"}, - {"instruction_index": "3", "turn": 2, "category": "writing"}, - ] - assert turn2_calls[1]["cache_meta"]["metadata"] == [ - {"instruction_index": "2", "turn": 2, "category": "math"}, - ] - - -def test_generate_multiturn_saves_metadata_in_cache(tmp_path): - questions = pd.DataFrame( - { - "category": ["writing"], - "turn_1": ["Q1"], - "turn_2": ["Q2"], - }, - index=pd.Index([7], name="instruction_index"), - ) - - with InferenceCache(tmp_path, "mt-bench", mode="refresh") as cache: - generate_multiturn( - questions, - "Dummy/mt-meta", - use_tqdm=False, - cache=cache, - ) - - model = make_model("Dummy/mt-meta", max_tokens=8192) - descriptor = model.cache_descriptor() - folder = store_folder( - tmp_path, - "mt-bench", - model.model_spec, - descriptor_hash(descriptor), - ) - with SQLiteInferenceStore(folder / "inference.db") as store: - rows = store.query_metadata() - - saved = [json.loads(value) for value in rows["metadata_json"]] - assert {"instruction_index": "7", "turn": 1, "category": "writing"} in saved - assert {"instruction_index": "7", "turn": 2, "category": "writing"} in saved diff --git a/tests/test_inference_cache.py b/tests/test_inference_cache.py index 4a559fe..756e78a 100644 --- a/tests/test_inference_cache.py +++ b/tests/test_inference_cache.py @@ -3,28 +3,22 @@ from types import SimpleNamespace import pytest -from langchain_core.messages import HumanMessage, SystemMessage -from langchain_core.prompt_values import ChatPromptValue -import judgearena.model_adapters as model_adapters -import judgearena.models as models from judgearena.inference_cache import InferenceCache -from judgearena.models import ( - HOSTED_ADAPTER_VERSION, - PreparedModel, - do_inference, - make_model, -) +from judgearena.models import do_inference, make_model from judgearena.store_sqlite import ( SQLiteInferenceStore, descriptor_hash, - stable_json_dumps, store_folder, ) def _install_fake_vllm(monkeypatch): captured = {"llm_init": False} + monkeypatch.setattr( + "judgearena.model_adapters._provider_package_version", + lambda name: "test-version", + ) class FakeSamplingParams: def __init__(self, **kwargs): @@ -234,102 +228,6 @@ def test_metadata_associations_are_saved(tmp_path): assert saved == {"q-1", "q-2"} -def test_secret_values_are_not_descriptorized(monkeypatch): - monkeypatch.setenv("OPENROUTER_API_KEY", "dummy") - model = make_model( - "OpenRouter/google/gemma-3-4b-it", - max_tokens=16, - api_key="super-secret", - default_headers={"Authorization": "Bearer secret"}, - ) - assert model.engine_kwargs["api_key"] == "super-secret" - assert model.engine_kwargs["default_headers"]["Authorization"] == "Bearer secret" - descriptor = model.cache_descriptor() - assert descriptor is not None - serialized = stable_json_dumps(descriptor) - assert "super-secret" not in serialized - assert "Authorization" not in serialized - assert "api_key" not in serialized - - -def test_provider_canonicalization_distinguishes_chat_and_raw(): - chat_model = make_model("Dummy/chat", max_tokens=8) - raw_model = make_model("Dummy/raw", max_tokens=8) - raw_model.input_mode = "raw" - - chat_payload = chat_model.canonicalize_input("hello") - raw_payload = raw_model.canonicalize_input("hello") - assert json.loads(chat_payload)["kind"] == "chat" - assert json.loads(raw_payload)["kind"] == "raw" - assert chat_payload != raw_payload - - prompt = ChatPromptValue( - messages=[SystemMessage(content="sys"), HumanMessage(content="hi", id="tmp")] - ) - chat_canonical = json.loads(chat_model.canonicalize_input(prompt)) - assert chat_canonical["messages"] == [ - {"role": "system", "content": "sys"}, - {"role": "user", "content": "hi"}, - ] - assert "id" not in stable_json_dumps(chat_canonical) - - -def test_set_temperature_direct_and_mutated_descriptors_match(tmp_path): - direct = make_model("Dummy/temp-hash", max_tokens=8, temperature=0.9) - mutated = make_model("Dummy/temp-hash", max_tokens=8) - mutated.set_temperature(0.9) - assert direct.cache_descriptor() == mutated.cache_descriptor() - assert descriptor_hash(direct.cache_descriptor()) == descriptor_hash( - mutated.cache_descriptor() - ) - - -def test_set_temperature_changes_cache_cell(tmp_path): - cold = make_model("Dummy/temp", max_tokens=8, temperature=0.2) - _seed_cache(tmp_path, "arena", cold, ["x"], ["cold-hit"]) - - model = make_model("Dummy/temp", max_tokens=8, temperature=0.2) - model.set_temperature(0.9) - assert model.cache_descriptor()["sampling"]["temperature"] == 0.9 - backend = model.materialize() - assert backend.init_kwargs["temperature"] == 0.9 - - with InferenceCache(tmp_path, "arena", mode="use") as cache: - results = do_inference(model, ["x"], cache=cache) - - assert results == ["temp"] - - -def test_hosted_adapter_version_is_part_of_descriptor(monkeypatch): - monkeypatch.setenv("OPENROUTER_API_KEY", "dummy") - model = make_model("OpenRouter/openai/gpt-4o-mini", max_tokens=8) - descriptor = model.cache_descriptor() - assert descriptor["hosted_adapter_version"] == HOSTED_ADAPTER_VERSION - assert descriptor["server_defaults"] == "unobserved" - metadata = model.producer_metadata() - assert metadata["hosted_adapter_version"] == HOSTED_ADAPTER_VERSION - assert "langchain_openai_version" in metadata - - -def test_direct_chat_vllm_can_be_cached_when_constructed(monkeypatch, tmp_path): - captured = _install_fake_vllm(monkeypatch) - chat_model = models.ChatVLLM( - model="Qwen/Qwen3.5-9B", - max_tokens=16, - gpu_memory_utilization=0.7, - ) - assert captured["llm_init"] is True - - wrapped = model_adapters.wrap_known_model(chat_model) - assert wrapped is not None - _seed_cache(tmp_path, "arena", wrapped, ["hello"], ["cached-vllm"]) - - with InferenceCache(tmp_path, "arena", mode="use") as cache: - results = do_inference(chat_model, ["hello"], cache=cache) - - assert results == ["cached-vllm"] - - def test_off_mode_never_reads_cache(tmp_path): model = make_model("Dummy/off", max_tokens=8) _seed_cache(tmp_path, "arena", model, ["prompt"], ["cached"]) @@ -340,37 +238,6 @@ def test_off_mode_never_reads_cache(tmp_path): assert results == ["off"] -def test_do_inference_off_bypasses_cache_resolution(monkeypatch, tmp_path): - model = make_model("Dummy/off-bypass", max_tokens=8) - - def fail_resolution(*args, **kwargs): - raise AssertionError("resolve_cacheable_model must not run in off mode") - - monkeypatch.setattr(models, "resolve_cacheable_model", fail_resolution) - - with InferenceCache(tmp_path, "arena", mode="off") as cache: - results = do_inference(model, ["prompt"], cache=cache) - - assert results == ["off-bypass"] - - -def test_prepared_model_proxy_materializes_backend(): - model = make_model("Dummy/proxy", max_tokens=8) - assert isinstance(model, PreparedModel) - assert model.init_kwargs["max_tokens"] == 8 - - -def test_descriptor_schema_version_is_present(): - model = make_model("Dummy/schema", max_tokens=8) - descriptor = model.cache_descriptor() - assert descriptor["descriptor_schema_version"] == models.DESCRIPTOR_SCHEMA_VERSION - - -def test_hosted_adapter_version_must_be_bumped_when_request_shaping_changes(): - """Guardrail: request-shaping edits must bump HOSTED_ADAPTER_VERSION.""" - assert HOSTED_ADAPTER_VERSION == "judgearena-hosted-adapter/v1" - - def test_off_mode_runs_every_input_without_dedupe(tmp_path): model = make_model("Dummy/off-dedupe", max_tokens=8) calls: list[str] = [] @@ -393,31 +260,6 @@ def counting_runner(inputs): assert calls == ["same", "same"] -def test_metadata_only_association_marks_cell_dirty(tmp_path): - model = make_model("Dummy/meta-dirty", max_tokens=8) - _seed_cache(tmp_path, "arena", model, ["prompt"], ["cached-out"]) - - with InferenceCache(tmp_path, "arena", mode="use", push=False) as cache: - do_inference( - model, - ["prompt"], - cache=cache, - cache_meta={"metadata": [{"question_id": "q-hit"}]}, - ) - assert cache._dirty_cells - - descriptor = model.cache_descriptor() - folder = store_folder( - tmp_path, - "arena", - model.model_spec, - descriptor_hash(descriptor), - ) - with SQLiteInferenceStore(folder / "inference.db") as store: - rows = store.query_metadata() - assert len(rows) == 1 - - def test_close_before_push_closes_sqlite_before_upload(monkeypatch, tmp_path): model = make_model("Dummy/push-close", max_tokens=8) states: list[bool] = [] @@ -449,17 +291,3 @@ def spy_push(*args, **kwargs): def test_invalid_cache_mode_rejected(tmp_path): with pytest.raises(ValueError, match="Invalid cache mode"): InferenceCache(tmp_path, "arena", mode="bogus") # type: ignore[arg-type] - - -def test_vllm_make_model_is_lazy_until_materialized(monkeypatch): - captured = _install_fake_vllm(monkeypatch) - model = make_model( - "VLLM/Qwen/Qwen3.5-9B", - max_tokens=16, - thinking_token_budget=64, - gpu_memory_utilization=0.7, - ) - assert isinstance(model, PreparedModel) - assert captured["llm_init"] is False - model.materialize() - assert captured["llm_init"] is True diff --git a/tests/test_logging.py b/tests/test_logging.py index 867c058..e4dfd20 100644 --- a/tests/test_logging.py +++ b/tests/test_logging.py @@ -2,9 +2,7 @@ from __future__ import annotations -import io import logging -import sys import pytest @@ -78,28 +76,6 @@ def test_configure_logging_no_duplicate_handlers(): assert len(console_handlers) == 1 -def test_configure_logging_rebinds_console_to_current_stderr(monkeypatch): - first_stream = io.StringIO() - second_stream = io.StringIO() - monkeypatch.setattr(sys, "stderr", first_stream) - configure_logging(0) - monkeypatch.setattr(sys, "stderr", second_stream) - configure_logging(0) - - get_logger("judgearena.test_rebind").error("new stream") - - assert "new stream" not in first_stream.getvalue() - assert "new stream" in second_stream.getvalue() - - -def test_configure_logging_adds_console_when_only_file_handler_exists(tmp_path): - attach_file_handler(tmp_path / "run.log") - - configure_logging(0) - - assert _console_handler_level() == logging.INFO - - def test_env_var_overrides_verbosity(monkeypatch): """JUDGEARENA_LOG_LEVEL env-var should override the CLI verbosity flag.""" monkeypatch.setenv("JUDGEARENA_LOG_LEVEL", "warning") diff --git a/tests/test_meta_eval_cache_threading.py b/tests/test_meta_eval_cache_threading.py deleted file mode 100644 index 42a00de..0000000 --- a/tests/test_meta_eval_cache_threading.py +++ /dev/null @@ -1,581 +0,0 @@ -"""Unified inference cache integration tests for meta-evaluation.""" - -from __future__ import annotations - -import json -from pathlib import Path - -import pandas as pd -import pytest - -import judgearena.evaluate as evaluate_module -import judgearena.meta_eval.annotate as meta_annotate -import judgearena.meta_eval.runner as meta_eval_runner -import judgearena.models as models_module -from judgearena.config import ( - CacheArgs, - RunConfig, - inference_cache_task, - meta_eval_cache_task, -) -from judgearena.inference_cache import InferenceCache -from judgearena.meta_eval.cli_args import CliMetaEvalArgs, meta_eval_args_from_config -from judgearena.meta_eval.prompts import PromptModeSpec -from judgearena.meta_eval.runner import main as meta_eval_main -from judgearena.models import make_model -from judgearena.store_sqlite import SQLiteInferenceStore, descriptor_hash, store_folder - - -def _meta_args_with_cache(tmp_path: Path, **overrides) -> CliMetaEvalArgs: - values = { - "reference_arena": "LMArena-140k", - "prompt_mode": "standard", - "top_models": 3, - "battles_per_model": 1, - "batch_size": 8, - "languages": ["en"], - "judge_model": "Dummy/score_A: 9\nscore_B: 1", - "result_folder": str(tmp_path / "results"), - "no_log_file": True, - "cache": CacheArgs(store_root=str(tmp_path / "cache")), - } - values.update(overrides) - return CliMetaEvalArgs(**values) - - -def _prompt_spec() -> PromptModeSpec: - return PromptModeSpec( - name="standard", - system_prompt="system", - user_prompt_template=( - "Question: {user_prompt}\nA: {completion_A}\nB: {completion_B}" - ), - ) - - -def _single_battle_frame() -> pd.DataFrame: - conv_a = [ - {"role": "user", "content": "Question 0"}, - {"role": "assistant", "content": "Answer A 0"}, - ] - conv_b = [ - {"role": "user", "content": "Question 0"}, - {"role": "assistant", "content": "Answer B 0"}, - ] - return pd.DataFrame( - [ - { - "question_id": "q-0", - "model_a": "model-0", - "model_b": "model-1", - "winner": "model_a", - "lang": "en", - "benchmark": "LMArena-140k", - "conversation_a": conv_a, - "conversation_b": conv_b, - } - ] - ) - - -def _base_meta_eval_payload(tmp_path: Path, **overrides) -> dict: - payload = { - "task": "meta-eval", - "judge": {"model": "Dummy/j"}, - "run": {"result_folder": str(tmp_path / "results"), "no_log_file": True}, - } - payload.update(overrides) - return payload - - -def test_meta_eval_args_from_config_carries_cache(tmp_path): - cfg = RunConfig( - **_base_meta_eval_payload( - tmp_path, - cache={"store_root": str(tmp_path / "cache"), "cache_mode": "refresh"}, - ) - ) - args = meta_eval_args_from_config(cfg) - assert args.cache.store_root == str(tmp_path / "cache") - assert args.cache.cache_mode == "refresh" - - -def test_meta_eval_args_from_config_carries_strip_thinking(tmp_path): - cfg = RunConfig( - **_base_meta_eval_payload( - tmp_path, - judge={ - "model": "Dummy/j", - "strip_thinking_before_judging": True, - }, - ) - ) - - args = meta_eval_args_from_config(cfg) - - assert args.strip_thinking_before_judging is True - - -def test_inference_cache_task_includes_reference_arena(tmp_path): - cfg = RunConfig(**_base_meta_eval_payload(tmp_path)) - assert inference_cache_task(cfg) == "meta-eval-LMArena-140k" - assert meta_eval_cache_task("LMArena-140k") == "meta-eval-LMArena-140k" - - -def test_meta_eval_forwards_strip_thinking_to_annotation(monkeypatch, tmp_path): - captured: dict[str, object] = {} - - def spy_annotate_battles(*, instructions, **kwargs): - captured.update(kwargs) - return [ - evaluate_module.JudgeAnnotation( - judge_completion="score A: 9 score B: 1", - instruction=instruction, - completion_A="A", - completion_B="B", - judge_input="rendered", - ) - for instruction in instructions - ] - - monkeypatch.setattr(meta_annotate, "annotate_battles", spy_annotate_battles) - args = _meta_args_with_cache(tmp_path, strip_thinking_before_judging=True) - - meta_annotate.annotate_sample( - _single_battle_frame(), - args, - judge_chat_model=object(), - prompt_spec=_prompt_spec(), - ) - - assert captured["strip_thinking_before_judging"] is True - - -def test_meta_eval_second_run_reuses_cached_rows(monkeypatch, tmp_path): - uncached_calls = {"count": 0} - real_uncached = models_module._do_inference_uncached - - def counting_uncached(*args, **kwargs): - uncached_calls["count"] += 1 - return real_uncached(*args, **kwargs) - - monkeypatch.setattr(models_module, "_do_inference_uncached", counting_uncached) - - args = _meta_args_with_cache(tmp_path, swap_mode="fixed") - sample = _single_battle_frame() - prompt_spec = _prompt_spec() - judge = make_model(args.judge_model, max_tokens=32) - - with InferenceCache( - args.cache.store_root, - meta_eval_cache_task(args.reference_arena), - mode="use", - ) as cache: - meta_annotate.annotate_sample( - sample, - args, - judge_chat_model=judge, - prompt_spec=prompt_spec, - cache=cache, - ) - assert uncached_calls["count"] > 0 - - uncached_calls["count"] = 0 - with InferenceCache( - args.cache.store_root, - meta_eval_cache_task(args.reference_arena), - mode="use", - ) as cache: - meta_annotate.annotate_sample( - sample, - args, - judge_chat_model=judge, - prompt_spec=prompt_spec, - cache=cache, - ) - assert uncached_calls["count"] == 0 - - -def test_meta_eval_swapped_orientations_store_distinct_associations( - monkeypatch, tmp_path -): - captured_metadata: list[dict] = [] - real_do_inference = evaluate_module.do_inference - - def spy_do_inference(*args, **kwargs): - cache_meta = kwargs.get("cache_meta") - if cache_meta is not None: - captured_metadata.extend(cache_meta.get("metadata", [])) - return real_do_inference(*args, **kwargs) - - monkeypatch.setattr(evaluate_module, "do_inference", spy_do_inference) - - args = _meta_args_with_cache(tmp_path, swap_mode="both") - prompt_spec = _prompt_spec() - judge = make_model(args.judge_model, max_tokens=32) - - with InferenceCache( - args.cache.store_root, - meta_eval_cache_task(args.reference_arena), - mode="use", - ) as cache: - meta_annotate.annotate_sample( - _single_battle_frame(), - args, - judge_chat_model=judge, - prompt_spec=prompt_spec, - cache=cache, - ) - - orientations = {row["orientation"] for row in captured_metadata} - assert orientations == {"forward", "swapped"} - assert all(row["question_id"] == "q-0" for row in captured_metadata) - assert all(row["reference_arena"] == "LMArena-140k" for row in captured_metadata) - - model = make_model(args.judge_model, max_tokens=32) - descriptor = model.cache_descriptor() - folder = store_folder( - tmp_path / "cache", - meta_eval_cache_task(args.reference_arena), - model.model_spec, - descriptor_hash(descriptor), - ) - with SQLiteInferenceStore(folder / "inference.db") as store: - metadata_rows = store.query_metadata() - stored_orientations = { - json.loads(row["metadata_json"])["orientation"] - for _, row in metadata_rows.iterrows() - } - assert stored_orientations == {"forward", "swapped"} - - -def test_meta_eval_changed_rendered_input_invalidates_only_that_row( - monkeypatch, tmp_path -): - uncached_calls = {"count": 0} - real_uncached = models_module._do_inference_uncached - - def counting_uncached(*args, **kwargs): - uncached_calls["count"] += 1 - return real_uncached(*args, **kwargs) - - monkeypatch.setattr(models_module, "_do_inference_uncached", counting_uncached) - - args = _meta_args_with_cache(tmp_path) - prompt_spec = _prompt_spec() - judge = make_model(args.judge_model, max_tokens=32) - original = _single_battle_frame() - changed = original.copy() - changed.iloc[0, changed.columns.get_loc("conversation_a")] = [ - {"role": "user", "content": "Changed question"}, - {"role": "assistant", "content": "Answer A 0"}, - ] - - with InferenceCache( - args.cache.store_root, - meta_eval_cache_task(args.reference_arena), - mode="use", - ) as cache: - meta_annotate.annotate_sample( - original, - args, - judge_chat_model=judge, - prompt_spec=prompt_spec, - cache=cache, - ) - first_count = uncached_calls["count"] - assert first_count > 0 - - uncached_calls["count"] = 0 - with InferenceCache( - args.cache.store_root, - meta_eval_cache_task(args.reference_arena), - mode="use", - ) as cache: - meta_annotate.annotate_sample( - changed, - args, - judge_chat_model=judge, - prompt_spec=prompt_spec, - cache=cache, - ) - assert uncached_calls["count"] == first_count - - -def test_meta_eval_parsing_and_costs_recompute_from_cached_output( - monkeypatch, tmp_path -): - args = _meta_args_with_cache(tmp_path) - prompt_spec = PromptModeSpec( - name="standard", - system_prompt="system", - user_prompt_template="user", - ) - judge = make_model("Dummy/score_A: 1\nscore_B: 9", max_tokens=32) - sample = _single_battle_frame() - - with InferenceCache( - args.cache.store_root, - meta_eval_cache_task(args.reference_arena), - mode="use", - ) as cache: - first = meta_annotate.annotate_sample( - sample, - args, - judge_chat_model=judge, - prompt_spec=prompt_spec, - cache=cache, - ) - - cost_calls = {"count": 0} - original_cost = meta_annotate.estimate_annotation_cost_usd - - def spy_cost(*, judge_input, judge_completion, judge_model): - cost_calls["count"] += 1 - return original_cost( - judge_input=judge_input, - judge_completion=judge_completion, - judge_model=judge_model, - ) - - monkeypatch.setattr(meta_annotate, "estimate_annotation_cost_usd", spy_cost) - - with InferenceCache( - args.cache.store_root, - meta_eval_cache_task(args.reference_arena), - mode="use", - ) as cache: - second = meta_annotate.annotate_sample( - sample, - args, - judge_chat_model=judge, - prompt_spec=prompt_spec, - cache=cache, - ) - - assert first.iloc[0]["winner_llm"] == "model_b" - assert second.iloc[0]["winner_llm"] == "model_b" - assert cost_calls["count"] > 0 - - -def test_meta_eval_runner_uses_one_shared_cache_handle( - monkeypatch, tmp_path, synthetic_arena_df -): - captured: list[object] = [] - - def spy_annotate_sample(df_sample, args, *, cache=None, **kwargs): - captured.append(cache) - return pd.DataFrame( - { - "question_id": df_sample["question_id"], - "model_a": df_sample["model_a"], - "model_b": df_sample["model_b"], - "winner": df_sample["winner"], - "lang": df_sample["lang"], - "benchmark": df_sample["benchmark"], - "orientation": "forward", - "instruction": "instr", - "completion_a": "A", - "completion_b": "B", - "judge_input": "prompt", - "judge_completion": "score_A: 9\nscore_B: 1", - "estimated_input_tokens": 2, - "estimated_output_tokens": 5, - "cost_usd": 0.001, - "cost_source": "estimated", - "winner_llm": df_sample["winner"], - "pref_llm": 0.0, - } - ) - - monkeypatch.setattr( - meta_eval_runner, - "load_reference_arena_battles", - lambda reference_arena, languages=None: synthetic_arena_df, - ) - monkeypatch.setattr( - meta_eval_runner, - "select_top_models", - lambda df, top_models: (["model-0", "model-1"], df), - ) - monkeypatch.setattr( - meta_eval_runner, - "sample_battles_per_model", - lambda df_top, models, battles_per_model, seed: df_top.head(1), - ) - monkeypatch.setattr(meta_eval_runner, "make_model", lambda **_kwargs: object()) - monkeypatch.setattr(meta_eval_runner, "annotate_sample", spy_annotate_sample) - - args = _meta_args_with_cache(tmp_path, top_models=2) - meta_eval_main(args) - - assert len(captured) == 1 - assert captured[0] is not None - - -def test_meta_eval_args_json_includes_cache_config( - tmp_path, monkeypatch, synthetic_arena_df -): - monkeypatch.setattr( - meta_eval_runner, - "load_reference_arena_battles", - lambda reference_arena, languages=None: synthetic_arena_df, - ) - monkeypatch.setattr( - meta_eval_runner, - "select_top_models", - lambda df, top_models: (["model-0", "model-1"], df), - ) - monkeypatch.setattr( - meta_eval_runner, - "sample_battles_per_model", - lambda df_top, models, battles_per_model, seed: df_top.head(1), - ) - monkeypatch.setattr(meta_eval_runner, "make_model", lambda **_kwargs: object()) - monkeypatch.setattr( - meta_eval_runner, - "annotate_sample", - lambda df_sample, args, **kwargs: pd.DataFrame( - { - "question_id": df_sample["question_id"], - "model_a": df_sample["model_a"], - "model_b": df_sample["model_b"], - "winner": df_sample["winner"], - "lang": df_sample["lang"], - "benchmark": df_sample["benchmark"], - "orientation": "forward", - "instruction": "instr", - "completion_a": "A", - "completion_b": "B", - "judge_input": "prompt", - "judge_completion": "score_A: 9\nscore_B: 1", - "estimated_input_tokens": 2, - "estimated_output_tokens": 5, - "cost_usd": 0.001, - "cost_source": "estimated", - "winner_llm": df_sample["winner"], - "pref_llm": 0.0, - } - ), - ) - - args = _meta_args_with_cache(tmp_path, top_models=2) - meta_eval_main(args) - output_dir = next(Path(args.result_folder).glob("meta-eval-*")) - args_payload = json.loads((output_dir / "args.json").read_text(encoding="utf-8")) - assert args_payload["cache"]["store_root"] == str(tmp_path / "cache") - assert "ignore_cache" not in args_payload - - -def test_meta_eval_args_serialization_redacts_engine_secrets(tmp_path): - args = _meta_args_with_cache(tmp_path) - args.engine_kwargs = { - "temperature": 0.2, - "api_key": "must-not-leak", - "default_headers": {"Authorization": "secret"}, - } - - payload = args.to_jsonable() - - assert payload["engine_kwargs"] == {"temperature": 0.2} - assert "must-not-leak" not in json.dumps(payload) - assert "secret" not in json.dumps(payload) - - -def _stub_meta_eval_sampling(monkeypatch, synthetic_arena_df: pd.DataFrame) -> None: - monkeypatch.setattr( - meta_eval_runner, - "load_reference_arena_battles", - lambda reference_arena, languages=None: synthetic_arena_df, - ) - monkeypatch.setattr( - meta_eval_runner, - "select_top_models", - lambda df, top_models: (["model-0", "model-1"], df), - ) - monkeypatch.setattr( - meta_eval_runner, - "sample_battles_per_model", - lambda df_top, models, battles_per_model, seed: df_top.head(1), - ) - - -def test_meta_eval_runner_creates_cache_cell_under_single_component_task( - monkeypatch, tmp_path, synthetic_arena_df -): - _stub_meta_eval_sampling(monkeypatch, synthetic_arena_df) - args = _meta_args_with_cache(tmp_path, top_models=2, swap_mode="fixed") - - meta_eval_main(args) - - task_root = ( - Path(args.cache.store_root) - / "inference" - / meta_eval_cache_task(args.reference_arena) - ) - db_files = list(task_root.rglob("inference.db")) - assert db_files, f"expected cache cell under {task_root}" - assert "meta-eval-LMArena-140k" in str(db_files[0]) - - -def test_meta_eval_runner_skips_push_when_downstream_processing_fails( - monkeypatch, tmp_path, synthetic_arena_df -): - import judgearena.inference_cache as inference_cache_mod - - push_calls: list[tuple] = [] - monkeypatch.setattr( - inference_cache_mod, - "push_cells", - lambda *args, **kwargs: push_calls.append((args, kwargs)), - ) - _stub_meta_eval_sampling(monkeypatch, synthetic_arena_df) - monkeypatch.setattr( - meta_eval_runner, - "_compute_results", - lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("downstream failed")), - ) - - args = _meta_args_with_cache( - tmp_path, - top_models=2, - cache=CacheArgs(store_root=str(tmp_path / "cache"), cache_push=True), - ) - - with pytest.raises(RuntimeError, match="downstream failed"): - meta_eval_main(args) - - assert push_calls == [] - task_root = ( - Path(args.cache.store_root) - / "inference" - / meta_eval_cache_task(args.reference_arena) - ) - assert list(task_root.rglob("inference.db")) - - -@pytest.fixture -def synthetic_arena_df() -> pd.DataFrame: - conv_a = [ - {"role": "user", "content": "Question 0"}, - {"role": "assistant", "content": "Answer A 0"}, - ] - conv_b = [ - {"role": "user", "content": "Question 0"}, - {"role": "assistant", "content": "Answer B 0"}, - ] - return pd.DataFrame( - [ - { - "question_id": "q-0", - "tstamp": 1, - "model_a": "model-0", - "model_b": "model-1", - "winner": "model_a", - "conversation_a": conv_a, - "conversation_b": conv_b, - "benchmark": "LMArena-140k", - "lang": "en", - } - ] - ) diff --git a/tests/test_model_adapters.py b/tests/test_model_adapters.py index f97b7e0..187f112 100644 --- a/tests/test_model_adapters.py +++ b/tests/test_model_adapters.py @@ -6,44 +6,41 @@ from langchain_core.prompt_values import ChatPromptValue import judgearena.model_adapters as adapters -import judgearena.models as models from judgearena.model_adapters import ( HOSTED_ADAPTER_VERSION, - LOCAL_LLAMACPP_ADAPTER_VERSION, build_producer_metadata, build_vllm_descriptor, - effective_sampling, normalize_constructor_settings, - top_k_from_settings, wrap_known_model, ) -from judgearena.models import DummyModel, make_model, resolve_vllm_settings +from judgearena.models import make_model, resolve_vllm_settings def test_normalize_constructor_settings_redacts_secret_keys(): - settings = { - "temperature": 0.5, - "api_key": "secret-value", - "default_headers": {"Authorization": "Bearer x"}, - "model_kwargs": {"top_k": 40}, - } - normalized = normalize_constructor_settings(settings) + normalized = normalize_constructor_settings( + { + "temperature": 0.5, + "api_key": "secret-value", + "default_headers": {"Authorization": "Bearer x"}, + "model_kwargs": {"top_k": 40}, + } + ) + assert normalized == {"temperature": 0.5, "model_kwargs": {"top_k": 40}} serialized = json.dumps(normalized) assert "secret-value" not in serialized assert "Authorization" not in serialized -def test_make_model_openrouter_endpoint_unifies_provider(monkeypatch): +def test_openrouter_endpoint_unifies_provider(monkeypatch): monkeypatch.setenv("OPENROUTER_API_KEY", "dummy") model = make_model( "ChatOpenAI/google/gemma-3-4b-it", max_tokens=16, base_url="https://openrouter.ai/api/v1", ) + descriptor = model.cache_descriptor() - assert descriptor is not None - assert model.provider == "OpenRouter" assert model.model_spec == "OpenRouter/google/gemma-3-4b-it" assert descriptor["provider"] == "OpenRouter" assert descriptor["base_url"] == "https://openrouter.ai/api/v1" @@ -53,121 +50,24 @@ def test_lazy_and_materialized_hosted_descriptors_match(monkeypatch): monkeypatch.setenv("OPENROUTER_API_KEY", "dummy") monkeypatch.setenv("OPENAI_API_KEY", "dummy") monkeypatch.setenv("TOGETHER_API_KEY", "dummy") - specs = [ + + for spec in ( "OpenRouter/openai/gpt-4o-mini", "OpenAI/gpt-3.5-turbo-instruct", "Together/meta-llama/Llama-3.3-70B-Instruct-Turbo", - ] - - for spec in specs: + ): lazy = make_model(spec, max_tokens=8, temperature=0.2) wrapped = wrap_known_model(lazy.materialize(), model_spec=lazy.model_spec) assert wrapped is not None assert wrapped.cache_descriptor() == lazy.cache_descriptor() -def test_wrapped_hosted_model_preserves_unset_max_tokens(monkeypatch): - monkeypatch.setenv("OPENAI_API_KEY", "dummy") - backend = models.ChatOpenAI(model="gpt-4o-mini") - - wrapped = wrap_known_model(backend) - - assert wrapped is not None - assert wrapped.max_tokens is None - assert "max_tokens" not in wrapped.cache_descriptor()["sampling"] - - -def test_make_model_captures_base_url_for_generic_hosted(monkeypatch): - monkeypatch.setenv("OPENAI_API_KEY", "dummy") - model = make_model( - "OpenAI/text-davinci-003", - max_tokens=8, - openai_api_base="https://example.test/v1", - ) - descriptor = model.cache_descriptor() - assert descriptor is not None - assert descriptor["base_url"] == "https://example.test/v1" - - -def test_make_model_preserves_constructor_secrets(monkeypatch): - monkeypatch.setenv("OPENROUTER_API_KEY", "env-key") - model = make_model( - "OpenRouter/google/gemma-3-4b-it", - max_tokens=16, - api_key="runtime-secret", - default_headers={"Authorization": "Bearer secret"}, - ) - assert model.engine_kwargs["api_key"] == "runtime-secret" - assert model.engine_kwargs["default_headers"]["Authorization"] == "Bearer secret" - descriptor = model.cache_descriptor() - assert descriptor is not None - serialized = json.dumps(descriptor) - assert "runtime-secret" not in serialized - assert "Authorization" not in serialized - - -def test_top_k_from_settings_preserves_zero(): - assert top_k_from_settings({"top_k": 0}) == 0 - assert top_k_from_settings({"model_kwargs": {"top_k": 0}}) == 0 - assert ( - effective_sampling( - temperature=0.5, - top_p=0.9, - top_k=top_k_from_settings({"top_k": 0}), - seed=None, - )["top_k"] - == 0 - ) - - -def test_resolve_vllm_settings_uses_explicit_tokenizer(monkeypatch): - monkeypatch.setitem( - sys.modules, - "vllm.config.reasoning", - SimpleNamespace( - ReasoningConfig=lambda **kwargs: SimpleNamespace(**kwargs), - ), - ) - - seen: dict[str, object] = {} - - class FakeAutoTokenizer: - @staticmethod - def from_pretrained(model, **kwargs): - seen["tokenizer_id"] = model - seen["tokenizer_kwargs"] = kwargs - return SimpleNamespace(chat_template="{{ messages }}") - - class FakeAutoConfig: - @staticmethod - def from_pretrained(model, **kwargs): - return SimpleNamespace(max_position_embeddings=4096) - - monkeypatch.setitem( - sys.modules, - "transformers", - SimpleNamespace( - AutoTokenizer=FakeAutoTokenizer, - AutoConfig=FakeAutoConfig, - ), - ) - - resolve_vllm_settings( - "org/base-model", - max_tokens=16, - tokenizer="org/custom-tokenizer", - revision="model-rev", - tokenizer_revision="tok-rev", - trust_remote_code=False, - ) - assert seen["tokenizer_id"] == "org/custom-tokenizer" - assert seen["tokenizer_kwargs"] == { - "trust_remote_code": False, - "revision": "tok-rev", - } - - def test_vllm_descriptor_uses_fully_resolved_sampling(monkeypatch): + monkeypatch.setattr( + adapters, + "_provider_package_version", + lambda name: "test-version", + ) monkeypatch.setitem( sys.modules, "vllm.config.reasoning", @@ -203,63 +103,17 @@ def from_pretrained(model, **kwargs): gpu_memory_utilization=0.7, ) descriptor = build_vllm_descriptor("VLLM/Qwen/Qwen3.5-9B", resolved) - assert descriptor is not None - assert descriptor["sampling"]["temperature"] == 0.6 - assert descriptor["sampling"]["top_p"] == 0.95 - assert descriptor["sampling"]["max_tokens"] == 32 - assert descriptor["sampling"]["thinking_token_budget"] == 16 + + assert descriptor["sampling"] == { + "max_tokens": 32, + "temperature": 0.6, + "thinking_token_budget": 16, + "top_p": 0.95, + } assert descriptor["engine_settings"]["max_model_len"] == 4096 - assert descriptor["engine_settings"]["gpu_memory_utilization"] == 0.7 assert descriptor["engine_settings"]["reasoning_parser"] == "qwen3" assert "reasoning_config" in descriptor["engine_settings"] assert "vllm_version" in descriptor - assert "langchain_openai_version" not in descriptor - - -def test_vllm_set_temperature_updates_sampling_only(): - descriptor = { - "sampling": {"temperature": 0.6, "max_tokens": 8}, - "engine_settings": {"max_model_len": 1024}, - } - resolved = SimpleNamespace(sampling_params_kwargs={"temperature": 0.6}) - model = adapters.PreparedModel( - provider="VLLM", - model_spec="VLLM/test/model", - model_name="test/model", - max_tokens=8, - engine_kwargs={"max_model_len": 1024}, - sampling={"temperature": 0.6, "max_tokens": 8}, - input_mode="chat", - descriptor=descriptor, - materialize=lambda _prepared: (_ for _ in ()).throw( - AssertionError("must stay lazy") - ), - producer_metadata={}, - vllm_resolved=resolved, - ) - - model.set_temperature(0.9) - - updated = model.cache_descriptor() - assert updated["sampling"]["temperature"] == 0.9 - assert "temperature" not in updated["engine_settings"] - assert resolved.sampling_params_kwargs["temperature"] == 0.9 - - -def test_llamacpp_descriptor_includes_local_engine_version(monkeypatch): - monkeypatch.setattr( - adapters, - "_provider_package_version", - lambda name: "0.42.0" if name == "llama-cpp-python" else "1.0.0", - ) - model = make_model("LlamaCpp/models/test.gguf", max_tokens=64) - descriptor = model.cache_descriptor() - assert descriptor is not None - assert descriptor["llama_cpp_python_version"] == "0.42.0" - assert descriptor["local_adapter_version"] == LOCAL_LLAMACPP_ADAPTER_VERSION - metadata = model.producer_metadata() - assert metadata["llama_cpp_python_version"] == "0.42.0" - assert metadata["descriptor_schema_version"] == adapters.DESCRIPTOR_SCHEMA_VERSION def test_hosted_canonicalization_preserves_tool_fields(): @@ -280,76 +134,38 @@ def test_hosted_canonicalization_preserves_tool_fields(): ToolMessage(content='{"city":"Paris"}', tool_call_id="call-1"), ] ) - canonical = json.loads(adapters.canonicalize_hosted_chat_input(prompt)) - assert canonical["messages"][0] == {"role": "system", "content": "sys"} - tool_message = canonical["messages"][-1] - assert tool_message["role"] == "tool" - assert tool_message["tool_call_id"] == "call-1" - assert tool_message["content"] == '{"city":"Paris"}' - ai_message = canonical["messages"][2] - assert ai_message["tool_calls"][0]["name"] == "lookup" - assert ai_message["tool_calls"][0]["id"] == "call-1" - assert "id" not in ai_message + + messages = json.loads(adapters.canonicalize_hosted_chat_input(prompt))["messages"] + assert messages[0] == {"role": "system", "content": "sys"} + assert messages[2]["tool_calls"][0]["name"] == "lookup" + assert messages[3] == { + "role": "tool", + "content": '{"city":"Paris"}', + "tool_call_id": "call-1", + } -def test_vllm_dict_message_roles_match_runtime_normalization(): +def test_vllm_canonicalization_matches_runtime_role_normalization(): messages = [ {"role": "human", "content": "hello"}, {"content": "missing role"}, ] - normalized = adapters.vllm_input_to_messages(messages) canonical = json.loads( adapters.canonicalize_vllm_input(messages, input_mode="chat") ) - assert normalized == [ - {"role": "user", "content": "hello"}, - {"role": "user", "content": "missing role"}, - ] - assert canonical["messages"] == normalized - - -def test_raw_canonicalization_accepts_dict_messages(): - canonical = json.loads( - adapters.canonicalize_raw_input( - [{"role": "user", "content": "hello"}, {"content": "world"}] - ) + assert ( + canonical["messages"] + == normalized + == [ + {"role": "user", "content": "hello"}, + {"role": "user", "content": "missing role"}, + ] ) - assert canonical == {"kind": "raw", "text": "hello\nworld"} - - -def test_set_temperature_updates_materialized_chatopenai_temperature(monkeypatch): - monkeypatch.setenv("OPENROUTER_API_KEY", "dummy") - - class FakeChat: - model_fields = { - "temperature": object(), - "max_tokens": object(), - "model": object(), - } - - def __init__(self, **kwargs): - self.__dict__.update(kwargs) - - monkeypatch.setattr(models, "ChatOpenAI", FakeChat) - model = make_model("OpenRouter/openai/gpt-4o-mini", max_tokens=8, temperature=0.2) - backend = model.materialize() - model.set_temperature(0.8) - assert backend.temperature == 0.8 - assert model.cache_descriptor()["sampling"]["temperature"] == 0.8 - - -def test_wrap_known_dummy_model(): - backend = DummyModel("Dummy/wrapped", max_tokens=8, temperature=0.1) - wrapped = wrap_known_model(backend) - assert wrapped is not None - assert wrapped.cache_descriptor() is not None - assert wrapped.canonicalize_input("hello").startswith('{"kind"') def test_producer_metadata_includes_adapter_schema(): metadata = build_producer_metadata(provider="OpenRouter") assert metadata["hosted_adapter_version"] == HOSTED_ADAPTER_VERSION assert metadata["descriptor_schema_version"] == adapters.DESCRIPTOR_SCHEMA_VERSION - assert "langchain_openai_version" in metadata diff --git a/tests/test_mt_bench_downloads.py b/tests/test_mt_bench_downloads.py index e87aca4..8ae0692 100644 --- a/tests/test_mt_bench_downloads.py +++ b/tests/test_mt_bench_downloads.py @@ -5,11 +5,9 @@ import judgearena.instruction_dataset.fluency as fluency_mod import judgearena.instruction_dataset.mt_bench as mt_bench -import judgearena.models as models_module import judgearena.mt_bench.mt_bench_utils as mt_bench_utils import judgearena.utils.io as utils_io from judgearena.config import RunConfig -from judgearena.inference_cache import InferenceCache from judgearena.prompts.registry import FASTCHAT_PAIRWISE_PROMPT_PRESET @@ -159,12 +157,12 @@ def fake_generate_multiturn(**kwargs): generation={"n_instructions": 2}, ) - with InferenceCache("/tmp/unused", "mt-bench", mode="off") as cache: - completions_a, completions_b = mt_bench_utils._generate_mt_bench_completions( - cfg=cfg, - questions_df=questions_df, - cache=cache, - ) + cache = object() + completions_a, completions_b = mt_bench_utils._generate_mt_bench_completions( + cfg=cfg, + questions_df=questions_df, + cache=cache, + ) assert generated_models == ["VLLM/example/model-a"] assert cache_values == [cache] @@ -567,113 +565,3 @@ def fake_judge(**kwargs): ) assert captured["judge"]["strip_thinking_before_judging"] is True - - -def test_run_mt_bench_forwards_cache_to_generation_and_judging(monkeypatch, tmp_path): - questions_df = pd.DataFrame( - {"turn_1": ["Q1"], "turn_2": ["Q1b"]}, - index=pd.Index([1], name="instruction_index"), - ) - captured: dict[str, object | None] = {} - - monkeypatch.setattr( - mt_bench_utils, - "load_instructions", - lambda dataset, n_instructions=None: questions_df, - ) - - def fake_generate(**kwargs): - captured["generation_cache"] = kwargs.get("cache") - return pd.DataFrame( - { - "instruction_index": [1], - "completion_turn_1": ["A1"], - "completion_turn_2": ["A2"], - } - ) - - monkeypatch.setattr(mt_bench_utils, "generate_multiturn", fake_generate) - monkeypatch.setattr( - mt_bench_utils, - "load_mt_bench_model_answers", - lambda model, n_instructions=None: None, - ) - monkeypatch.setattr(mt_bench_utils, "make_model", lambda **kwargs: object()) - monkeypatch.setattr( - mt_bench_utils, "_finalize_mt_bench_run", lambda **kwargs: kwargs["prefs"] - ) - - def fake_judge(**kwargs): - captured["judge_cache"] = kwargs.get("cache") - return pd.Series([0.0], dtype=float), [], [], 0 - - monkeypatch.setattr(mt_bench_utils, "judge_mt_bench_pairwise_fastchat", fake_judge) - - cfg = RunConfig( - task="mt-bench", - model={"name": "VLLM/example/model-a"}, - judge={"model": "VLLM/Judge"}, - generation={"n_instructions": 1}, - run={"result_folder": str(tmp_path)}, - ) - - with InferenceCache("/tmp/unused", "mt-bench", mode="off") as cache: - mt_bench_utils.run_mt_bench( - cfg, - cache=cache, - res_folder=tmp_path, - result_name="mt-bench-test", - ) - - assert captured["generation_cache"] is cache - assert captured["judge_cache"] is cache - - -def test_generate_mt_bench_completions_reuses_inference_cache(tmp_path, monkeypatch): - questions_df = pd.DataFrame( - { - "category": ["writing"], - "turn_1": ["Q1"], - "turn_2": ["Q2"], - }, - index=pd.Index([1], name="instruction_index"), - ) - backend_inputs: list[int] = [] - real_uncached = models_module._do_inference_uncached - - def counting_uncached(chat_model, inputs, *, use_tqdm=False): - backend_inputs.append(len(inputs)) - return real_uncached(chat_model, inputs, use_tqdm=use_tqdm) - - monkeypatch.setattr(models_module, "_do_inference_uncached", counting_uncached) - monkeypatch.setattr( - mt_bench_utils, - "load_mt_bench_model_answers", - lambda model, n_instructions=None: None, - ) - - cfg = RunConfig( - task="mt-bench", - model={"name": "Dummy/mt-cache-a", "baseline": "Dummy/mt-cache-b"}, - judge={"model": "Dummy/J"}, - generation={"n_instructions": 1}, - ) - - with InferenceCache(tmp_path, "mt-bench", mode="refresh") as cache: - first_a, _first_b = mt_bench_utils._generate_mt_bench_completions( - cfg=cfg, - questions_df=questions_df, - cache=cache, - ) - - assert sum(backend_inputs) == 4 - - with InferenceCache(tmp_path, "mt-bench", mode="use") as cache: - second_a, _second_b = mt_bench_utils._generate_mt_bench_completions( - cfg=cfg, - questions_df=questions_df, - cache=cache, - ) - - assert sum(backend_inputs) == 4 - assert first_a.loc[1, "completion_turn_1"] == second_a.loc[1, "completion_turn_1"] diff --git a/tests/test_mt_bench_fastchat_compat.py b/tests/test_mt_bench_fastchat_compat.py index 920fda6..555ae1f 100644 --- a/tests/test_mt_bench_fastchat_compat.py +++ b/tests/test_mt_bench_fastchat_compat.py @@ -3,8 +3,6 @@ import pandas as pd import pytest -import judgearena.mt_bench.fastchat_compat as fastchat_module -from judgearena.inference_cache import InferenceCache from judgearena.mt_bench.fastchat_compat import ( _conservative_winner, _map_verdict_to_winner, @@ -120,35 +118,3 @@ def test_judge_mt_bench_pairwise_fastchat_swap_mode_both_is_conservative(): assert annotations[0]["final_winner"] == "model_A" assert "B1" in annotations[0]["g2_user_prompt"] assert metadata == [{"question_id": 1, "category": "writing", "turn": 1}] - - -def test_judge_mt_bench_pairwise_fastchat_forwards_cache(monkeypatch): - captured: list[object | None] = [] - - def fake_infer(*, cache, **kwargs): - captured.append(cache) - return (["[[A]]"] * len(kwargs["items"]), [{}] * len(kwargs["items"])) - - monkeypatch.setattr( - fastchat_module, - "infer_pairwise_judgments_by_prompt_groups", - fake_infer, - ) - - with InferenceCache("/tmp/unused", "mt-judge", mode="off") as cache: - judge_mt_bench_pairwise_fastchat( - judge_chat_model=object(), - judge_model="judge", - questions=_questions_df(category="writing"), - completions_a=_completions_df("A"), - completions_b=_completions_df("B"), - model_a="model-a", - model_b="model-b", - turns_mode="single", - swap_mode="both", - truncate_input_chars=None, - use_tqdm=False, - cache=cache, - ) - - assert captured == [cache, cache] diff --git a/tests/test_mt_bench_pairwise_cache_threading.py b/tests/test_mt_bench_pairwise_cache_threading.py deleted file mode 100644 index 60c1fd9..0000000 --- a/tests/test_mt_bench_pairwise_cache_threading.py +++ /dev/null @@ -1,159 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass - -import judgearena.mt_bench.pairwise_judging as pairwise_module -from judgearena.inference_cache import InferenceCache -from judgearena.mt_bench.pairwise_judging import ( - MTBenchJudgeItem, - infer_pairwise_judgments_by_prompt_groups, -) - - -@dataclass(frozen=True) -class _Prompt: - name: str - system_prompt: str | None - user_prompt_template: str - multi_turn: bool - ref_based: bool = False - - -def _item( - *, - question_id: object, - category: str | None, - turn: int, - prompt: _Prompt, - prompt_kwargs: dict[str, str], -) -> MTBenchJudgeItem: - return MTBenchJudgeItem( - question_id=question_id, - category=category, - turn=turn, - prompt=prompt, - prompt_kwargs=prompt_kwargs, - ) - - -def test_infer_pairwise_judgments_metadata_order_and_orientation(monkeypatch): - captured: list[dict] = [] - - def spy_do_inference(*, cache_meta, **kwargs): - captured.append( - { - "cache_meta": cache_meta, - "input_count": len(kwargs["inputs"]), - } - ) - return ["judgment"] * len(kwargs["inputs"]) - - monkeypatch.setattr(pairwise_module, "do_inference", spy_do_inference) - - single_prompt = _Prompt( - name="default-single", - system_prompt=None, - user_prompt_template="{question} {answer_a} {answer_b}", - multi_turn=False, - ) - multi_prompt = _Prompt( - name="default-multi", - system_prompt=None, - user_prompt_template="{question_1} {answer_a_1}", - multi_turn=True, - ) - items = [ - _item( - question_id=1, - category="writing", - turn=1, - prompt=single_prompt, - prompt_kwargs={ - "question": "Q1", - "answer_a": "A1", - "answer_b": "B1", - }, - ), - _item( - question_id=2, - category="math", - turn=2, - prompt=multi_prompt, - prompt_kwargs={ - "question_1": "Q2a", - "question_2": "Q2b", - "answer_a_1": "A2a", - "answer_a_2": "A2b", - "answer_b_1": "B2a", - "answer_b_2": "B2b", - }, - ), - ] - - judgments, used_kwargs = infer_pairwise_judgments_by_prompt_groups( - judge_chat_model=object(), - items=items, - use_tqdm=False, - swap_answers=True, - ) - - assert judgments == ["judgment", "judgment"] - assert used_kwargs[0]["answer_a"] == "B1" - assert used_kwargs[0]["answer_b"] == "A1" - assert len(captured) == 2 - assert captured[0]["input_count"] == 1 - assert captured[0]["cache_meta"]["metadata"] == [ - { - "question_id": "1", - "category": "writing", - "turn": 1, - "prompt": "default-single", - "orientation": "reversed", - } - ] - assert captured[1]["cache_meta"]["metadata"] == [ - { - "question_id": "2", - "category": "math", - "turn": 2, - "prompt": "default-multi", - "orientation": "reversed", - } - ] - - -def test_infer_pairwise_judgments_forwards_cache(monkeypatch): - captured: list[object] = [] - - def spy_do_inference(*, cache, **kwargs): - captured.append(cache) - return ["judgment"] * len(kwargs["inputs"]) - - monkeypatch.setattr(pairwise_module, "do_inference", spy_do_inference) - - prompt = _Prompt( - name="default-single", - system_prompt=None, - user_prompt_template="{question}", - multi_turn=False, - ) - items = [ - _item( - question_id=9, - category="coding", - turn=1, - prompt=prompt, - prompt_kwargs={"question": "Q", "answer_a": "A", "answer_b": "B"}, - ) - ] - - with InferenceCache("/tmp/unused", "mt-judge", mode="off") as cache: - infer_pairwise_judgments_by_prompt_groups( - judge_chat_model=object(), - items=items, - use_tqdm=False, - swap_answers=False, - cache=cache, - ) - - assert captured == [cache] diff --git a/tests/test_mt_bench_preset_judging.py b/tests/test_mt_bench_preset_judging.py index 9c6b57c..7c6a9e4 100644 --- a/tests/test_mt_bench_preset_judging.py +++ b/tests/test_mt_bench_preset_judging.py @@ -3,8 +3,6 @@ import pandas as pd import pytest -import judgearena.mt_bench.preset_judging as preset_module -from judgearena.inference_cache import InferenceCache from judgearena.mt_bench.preset_judging import ( _build_mt_bench_preset_items, _select_preset_prompt, @@ -153,47 +151,3 @@ def test_judge_mt_bench_with_preset_parses_and_inverts_swapped_scores(): {"question_id": 1, "category": "writing", "turn": 1}, {"question_id": 1, "category": "writing", "turn": 1}, ] - - -def test_judge_mt_bench_with_preset_forwards_cache(monkeypatch): - captured: list[object | None] = [] - - def fake_infer(*, cache, items, swap_answers=False, **kwargs): - captured.append(cache) - used_kwargs = [dict(item.prompt_kwargs) for item in items] - if swap_answers: - used_kwargs = [ - { - **kwargs_, - "answer_a": kwargs_.get("answer_b", ""), - "answer_b": kwargs_.get("answer_a", ""), - } - for kwargs_ in used_kwargs - ] - n = len(items) - return (["score_A: 10\nscore_B: 0"] * n, used_kwargs) - - monkeypatch.setattr( - preset_module, - "infer_pairwise_judgments_by_prompt_groups", - fake_infer, - ) - - with InferenceCache("/tmp/unused", "mt-judge", mode="off") as cache: - judge_mt_bench_with_preset( - judge_chat_model=object(), - judge_model="judge", - questions=_questions_df(category="writing"), - completions_a=_completions_df("A"), - completions_b=_completions_df("B"), - model_a="model-a", - model_b="model-b", - turns_mode="single", - swap_mode="both", - truncate_input_chars=None, - use_tqdm=False, - prompt_preset="default", - cache=cache, - ) - - assert captured == [cache, cache] diff --git a/tests/test_no_legacy_runtime_cache.py b/tests/test_no_legacy_runtime_cache.py index 2f600e1..89a2a96 100644 --- a/tests/test_no_legacy_runtime_cache.py +++ b/tests/test_no_legacy_runtime_cache.py @@ -23,10 +23,6 @@ "from judgearena.meta_eval.cache import", ) -LEGACY_MIGRATION_ALLOWLIST = frozenset( - {("ignore_cache", "judgearena/cache_backfill_config.py")} -) - INFERENCE_BATCH_INVOKE_ALLOWLIST = frozenset( { "judgearena/models.py", @@ -52,7 +48,7 @@ def test_judgearena_has_no_legacy_runtime_cache_symbols() -> None: rel = path.relative_to(REPO_ROOT).as_posix() text = path.read_text(encoding="utf-8") for label, pattern in FORBIDDEN_PATTERNS: - if pattern.search(text) and (label, rel) not in LEGACY_MIGRATION_ALLOWLIST: + if pattern.search(text): violations.append(f"{rel}: {label}") for imp in FORBIDDEN_IMPORTS: if imp in text: diff --git a/tests/test_store_sqlite.py b/tests/test_store_sqlite.py index 1e1dbc4..53e9ae0 100644 --- a/tests/test_store_sqlite.py +++ b/tests/test_store_sqlite.py @@ -145,51 +145,6 @@ def test_output_and_metadata_batch_rolls_back_atomically(tmp_path): assert store.query_metadata().empty -def test_save_metadata_normalizes_dict_and_string_equivalently(tmp_path): - meta = {"b": 2, "a": 1} - with SQLiteInferenceStore(tmp_path / "inference.db") as store: - store.save_outputs(_outputs(["h0"]), pushed_by="test") - store.save_metadata( - pd.DataFrame({"input_hash": ["h0"], "metadata_json": [meta]}), - ) - rows = store.query_metadata(["h0"]) - - assert rows["metadata_hash"].iloc[0] == metadata_hash(meta) - assert json.loads(rows["metadata_json"].iloc[0]) == meta - - with SQLiteInferenceStore(tmp_path / "other.db") as store: - store.save_outputs(_outputs(["h0"]), pushed_by="test") - store.save_metadata( - pd.DataFrame( - { - "input_hash": ["h0"], - "metadata_json": [json.dumps(meta, sort_keys=False)], - } - ), - ) - other = store.query_metadata(["h0"]) - - assert other["metadata_hash"].iloc[0] == rows["metadata_hash"].iloc[0] - - -def test_save_metadata_ignores_nan_metadata_hash(tmp_path): - meta = {"question_id": "q-1"} - with SQLiteInferenceStore(tmp_path / "inference.db") as store: - store.save_outputs(_outputs(["h0"]), pushed_by="test") - store.save_metadata( - pd.DataFrame( - { - "input_hash": ["h0"], - "metadata_hash": [float("nan")], - "metadata_json": [stable_json_dumps(meta)], - } - ), - ) - rows = store.query_metadata(["h0"]) - - assert rows["metadata_hash"].iloc[0] == metadata_hash(meta) - - def test_chunked_query_preserves_input_independent_order(tmp_path, monkeypatch): hashes = [f"h{index}" for index in range(IN_QUERY_CHUNK_SIZE + 3)] with SQLiteInferenceStore(tmp_path / "inference.db") as store: diff --git a/tests/test_store_sync.py b/tests/test_store_sync.py index 35c86ac..3283d4b 100644 --- a/tests/test_store_sync.py +++ b/tests/test_store_sync.py @@ -269,64 +269,6 @@ def test_fetch_cell_rejects_metadata_mismatch_before_db_merge(fake_hub, tmp_path assert _read_outputs(local) == {"local": "L"} -def test_fetch_cell_rejects_misnamed_remote_cell_folder(fake_hub, tmp_path): - wrong_hash = "0000000000000000" - wrong_path = f"inference/arena/VLLM/Qwen%2Fjudge/{wrong_hash}/{INFERENCE_DB_NAME}" - wrong_metadata = f"inference/arena/VLLM/Qwen%2Fjudge/{wrong_hash}/metadata.json" - remote = tmp_path / "remote.db" - _write_inference( - remote, - [ - { - "input_hash": "remote", - "output_text": "R", - "pushed_at": "2026-01-01", - } - ], - ) - fake_hub.files[wrong_path] = remote.read_bytes() - fake_hub.files[wrong_metadata] = json.dumps(CELL_CONFIG).encode("utf-8") - - local = store_folder(tmp_path, "arena", MODEL_SPEC, wrong_hash) / INFERENCE_DB_NAME - with pytest.raises(ValueError, match="does not match descriptor hash"): - store_sync.fetch_cell(REPO_ID, wrong_path, local) - - -def test_fetch_cell_requires_remote_metadata(fake_hub, tmp_path): - remote = tmp_path / "remote.db" - _write_inference( - remote, - [ - { - "input_hash": "remote", - "output_text": "R", - "pushed_at": "2026-01-01", - } - ], - ) - fake_hub.files[PATH_IN_REPO] = remote.read_bytes() - - local = tmp_path / "store" / INFERENCE_DB_NAME - with pytest.raises(ValueError, match="missing required"): - store_sync.fetch_cell(REPO_ID, PATH_IN_REPO, local) - - -def test_fetch_cell_missing_remote_is_noop(fake_hub, tmp_path): - local = tmp_path / INFERENCE_DB_NAME - _write_inference( - local, - [ - { - "input_hash": "local", - "output_text": "L", - "pushed_at": "2026-01-01", - } - ], - ) - assert not store_sync.fetch_cell(REPO_ID, PATH_IN_REPO, local) - assert _read_outputs(local) == {"local": "L"} - - def test_push_cell_retries_and_preserves_concurrent_rows(fake_hub, tmp_path): fake_hub.head = "initial" cell_dir = store_folder(tmp_path, "arena", MODEL_SPEC, CELL_CONFIG_HASH) @@ -391,174 +333,3 @@ def test_push_cell_requires_metadata(fake_hub, tmp_path): ) with pytest.raises(FileNotFoundError, match="Missing metadata.json"): store_sync.push_cell(REPO_ID, PATH_IN_REPO, local, pushed_by="alice") - - -def test_push_cells_strict_rejects_missing_local_cell(tmp_path): - with pytest.raises(FileNotFoundError, match="Cache cell does not exist"): - store_sync.push_cells( - REPO_ID, - tmp_path, - [tmp_path / "missing" / INFERENCE_DB_NAME], - pushed_by="alice", - strict=True, - ) - - -def test_iter_cell_dbs_uses_segment_boundary_prefix(tmp_path): - arena_cell = ( - tmp_path - / "inference" - / "arena" - / "VLLM" - / "Qwen%2Fjudge" - / "abc123" - / INFERENCE_DB_NAME - ) - arena_cell.parent.mkdir(parents=True) - arena_cell.write_bytes(b"") - arena_hard_cell = ( - tmp_path - / "inference" - / "arena-hard-v2.0" - / "VLLM" - / "Qwen%2Fjudge" - / "def456" - / INFERENCE_DB_NAME - ) - arena_hard_cell.parent.mkdir(parents=True) - arena_hard_cell.write_bytes(b"") - - matched = store_sync.iter_cell_dbs(tmp_path, path_prefix="inference/arena") - assert matched == [arena_cell] - - -def test_iter_cell_dbs_ignores_noncanonical_layout(tmp_path): - invalid = tmp_path / "inference" / "too-shallow" / INFERENCE_DB_NAME - invalid.parent.mkdir(parents=True) - invalid.write_bytes(b"") - - assert store_sync.iter_cell_dbs(tmp_path) == [] - - -def test_push_cell_create_pr_updates_local_db_without_uploading_branch( - fake_hub, - tmp_path, -): - fake_hub.head = "initial" - cell_dir = store_folder(tmp_path, "arena", MODEL_SPEC, CELL_CONFIG_HASH) - local = cell_dir / INFERENCE_DB_NAME - write_store_metadata(cell_dir, CELL_CONFIG) - _write_inference( - local, - [ - { - "input_hash": "local", - "output_text": "L", - "pushed_at": "2026-01-01", - } - ], - ) - remote = tmp_path / "remote.db" - _write_inference( - remote, - [ - { - "input_hash": "remote", - "output_text": "R", - "pushed_at": "2026-02-01", - } - ], - ) - fake_hub.files[PATH_IN_REPO] = remote.read_bytes() - fake_hub.files[METADATA_IN_REPO] = json.dumps(CELL_CONFIG).encode("utf-8") - - result = store_sync.push_cell( - REPO_ID, - PATH_IN_REPO, - local, - pushed_by="alice", - create_pr=True, - ) - assert result == "https://hf.co/pr/1" - assert fake_hub.commit_calls == 1 - assert fake_hub.files[PATH_IN_REPO] == remote.read_bytes() - assert _read_outputs(local) == {"local": "L", "remote": "R"} - - -def test_fetch_remote_cells_bootstraps_empty_store_with_prefix(fake_hub, tmp_path): - remote = tmp_path / "remote.db" - _write_inference( - remote, - [ - { - "input_hash": "remote", - "output_text": "R", - "pushed_at": "2026-01-01", - } - ], - ) - fake_hub.files[PATH_IN_REPO] = remote.read_bytes() - fake_hub.files[METADATA_IN_REPO] = json.dumps(CELL_CONFIG).encode("utf-8") - fake_hub.head = "initial" - - store_root = tmp_path / "empty-store" - fetched = store_sync.fetch_remote_cells( - REPO_ID, - store_root, - path_prefix="inference/arena", - strict=True, - ) - assert len(fetched) == 1 - assert _read_outputs(fetched[0]) == {"remote": "R"} - assert json.loads((fetched[0].parent / "metadata.json").read_text()) == CELL_CONFIG - assert ( - store_sync.discover_remote_cell_dbs( - REPO_ID, - path_prefix="inference/other", - ) - == [] - ) - - -def test_fetch_remote_cells_rejects_noncanonical_remote_path(fake_hub, tmp_path): - path_in_repo = "inference/arena/VLLM/../../outside/hash/inference.db" - fake_hub.files[path_in_repo] = b"not-a-database" - fake_hub.head = "initial" - - with pytest.raises(ValueError, match="Invalid remote cache cell path"): - store_sync.fetch_remote_cells( - REPO_ID, - tmp_path / "store", - path_prefix="inference/arena", - strict=True, - ) - - assert not (tmp_path / "outside").exists() - - -def test_iter_cell_dbs_respects_prefix(tmp_path): - cell = ( - tmp_path - / "inference" - / "arena" - / "VLLM" - / "Qwen%2Fjudge" - / "abc123" - / INFERENCE_DB_NAME - ) - cell.parent.mkdir(parents=True) - cell.write_bytes(b"") - other = ( - tmp_path - / "inference" - / "other-task" - / "VLLM" - / "Model" - / "def456" - / INFERENCE_DB_NAME - ) - other.parent.mkdir(parents=True) - other.write_bytes(b"") - - assert len(store_sync.iter_cell_dbs(tmp_path)) == 2 - assert len(store_sync.iter_cell_dbs(tmp_path, path_prefix="inference/arena")) == 1