diff --git a/judgearena/constants.py b/judgearena/constants.py index 2c1e9bd4..f4abf95d 100644 --- a/judgearena/constants.py +++ b/judgearena/constants.py @@ -8,3 +8,6 @@ VLLM_REASONING_END_STR = ( "I have to give the solution based on the thinking directly now." ) + +VLLM_DEFAULT_TEMPERATURE = 0.6 +VLLM_DEFAULT_TOP_P = 0.95 diff --git a/judgearena/inference.py b/judgearena/inference.py index 463b995e..a14df174 100644 --- a/judgearena/inference.py +++ b/judgearena/inference.py @@ -6,8 +6,9 @@ from abc import ABC, abstractmethod from collections.abc import Callable from dataclasses import dataclass, field +from importlib import metadata as importlib_metadata from pathlib import Path -from typing import Any, ClassVar, NotRequired, TypedDict +from typing import Any, ClassVar, Literal, NotRequired, TypedDict import pandas as pd @@ -21,28 +22,58 @@ stable_json_dumps, write_descriptor, ) +from judgearena.constants import VLLM_DEFAULT_TEMPERATURE, VLLM_DEFAULT_TOP_P from judgearena.usage import RequestUsage _ROLE_MAP = {"human": "user", "ai": "assistant", "system": "system"} +VLLM_EXECUTION_ONLY_KWARGS = { + "enforce_eager", + "gpu_memory_utilization", + "tensor_parallel_size", +} +InputMode = Literal["chat", "text"] +_PROVIDER_INPUT_MODES: dict[str, InputMode] = { + "ChatOpenAI": "chat", + "Dummy": "chat", + "LlamaCpp": "text", + "OpenAI": "text", + "OpenRouter": "chat", + "Together": "text", + "VLLM": "chat", +} -def canonicalize_chat_input(input_item: Any) -> str: - """Serialize a logical model input for content-addressed cache lookup.""" + +def provider_input_mode(provider: str) -> InputMode | None: + return _PROVIDER_INPUT_MODES.get(provider) + + +def _canonical_messages(input_item: Any) -> list[dict[str, Any]]: if isinstance(input_item, str): - payload = {"type": "text", "text": input_item} - elif hasattr(input_item, "to_messages"): - payload = { - "type": "messages", - "messages": [ - { - "role": _ROLE_MAP.get(message.type, message.type), - "content": message.content, - } - for message in input_item.to_messages() - ], - } + return [{"role": "user", "content": input_item}] + if hasattr(input_item, "to_messages"): + return [ + { + "role": _ROLE_MAP.get(message.type, message.type), + "content": message.content, + } + for message in input_item.to_messages() + ] + raise TypeError(f"Unsupported inference input: {type(input_item)!r}") + + +def canonicalize_model_input(input_item: Any, input_mode: InputMode) -> str: + """Serialize the chat messages or flattened text sent to a provider.""" + if input_mode == "text": + if isinstance(input_item, str): + text = input_item + elif hasattr(input_item, "to_string"): + text = input_item.to_string() + else: + raise TypeError(f"Unsupported inference input: {type(input_item)!r}") + payload = {"type": "text", "text": text} else: - raise TypeError(f"Unsupported inference input: {type(input_item)!r}") + payload = {"type": "messages", "messages": _canonical_messages(input_item)} return stable_json_dumps(payload) @@ -50,17 +81,51 @@ def build_model_descriptor( provider: str, model_name: str, resolved_kwargs: dict[str, Any], + endpoint: str | None = None, ) -> dict[str, Any] | None: - """Describe output-affecting settings without constructing the backend.""" - if provider != "Dummy": + """Describe output-affecting settings without constructing the backend. + + Omitted hosted settings retain provider defaults. Local-engine versions are + part of the key. VLLM tokenizer-template changes require a distinct model + revision or an explicit ``chat_template``. + """ + input_mode = provider_input_mode(provider) + if input_mode is None: return None - return { + + descriptor_kwargs = resolved_kwargs.copy() + if provider == "VLLM": + descriptor_kwargs["temperature"] = ( + VLLM_DEFAULT_TEMPERATURE + if descriptor_kwargs.get("temperature") is None + else descriptor_kwargs["temperature"] + ) + descriptor_kwargs["top_p"] = ( + VLLM_DEFAULT_TOP_P + if descriptor_kwargs.get("top_p") is None + else descriptor_kwargs["top_p"] + ) + descriptor_kwargs["chat_template"] = descriptor_kwargs.get("chat_template") + descriptor_kwargs = { + key: value + for key, value in descriptor_kwargs.items() + if key not in VLLM_EXECUTION_ONLY_KWARGS + } + + descriptor = { "schema_version": "judgearena-inference-cache/v1", "provider": provider, "model": model_name, - "input_mode": "chat", - "model_kwargs": resolved_kwargs, + "input_mode": input_mode, + "model_kwargs": descriptor_kwargs, } + if provider == "VLLM": + descriptor["backend_version"] = importlib_metadata.version("vllm") + elif provider == "LlamaCpp": + descriptor["backend_version"] = importlib_metadata.version("llama-cpp-python") + if endpoint is not None: + descriptor["endpoint"] = endpoint.rstrip("/") + return descriptor @dataclass(frozen=True) diff --git a/judgearena/models.py b/judgearena/models.py index 4b1a5d9b..adc41472 100644 --- a/judgearena/models.py +++ b/judgearena/models.py @@ -10,7 +10,6 @@ import time import warnings from collections.abc import Mapping -from contextlib import contextmanager from dataclasses import replace from langchain_community.llms import LlamaCpp @@ -20,14 +19,19 @@ from tqdm.contrib.logging import logging_redirect_tqdm from judgearena.cache_sqlite import input_hash -from judgearena.constants import VLLM_REASONING_END_STR, VLLM_REASONING_START_STR +from judgearena.constants import ( + VLLM_DEFAULT_TEMPERATURE, + VLLM_DEFAULT_TOP_P, + VLLM_REASONING_END_STR, + VLLM_REASONING_START_STR, +) from judgearena.inference import ( CacheRowMetadata, InferenceCache, InferenceResult, PreparedModel, build_model_descriptor, - canonicalize_chat_input, + canonicalize_model_input, ) from judgearena.log import get_logger from judgearena.usage import RequestUsage, RunUsage, record_usage @@ -267,8 +271,10 @@ def __init__( self._sampling_params_kwargs = { "max_tokens": max_tokens, - "temperature": 0.6 if temperature is None else float(temperature), - "top_p": 0.95 if top_p is None else float(top_p), + "temperature": ( + VLLM_DEFAULT_TEMPERATURE if temperature is None else float(temperature) + ), + "top_p": VLLM_DEFAULT_TOP_P if top_p is None else float(top_p), } if top_k is not None: self._sampling_params_kwargs["top_k"] = int(top_k) @@ -776,51 +782,6 @@ def batch_with_retry(batch_inputs, max_retries=5, base_delay=1.0): return [result.text for result in results] -@contextmanager -def _cache_store(cache, model): - try: - store = cache.open_store(model) - except _CACHE_OPERATION_ERRORS as exc: - logger.warning( - "Cache open failed at %s: %s. Continuing without caching.", - cache.store_root, - exc, - ) - yield None - return - - try: - yield store - finally: - try: - store.close() - except _CACHE_OPERATION_ERRORS as exc: - logger.warning("Cache close failed at %s: %s.", store.db_path, exc) - - -def _read_cached_rows(store, input_hashes): - try: - return store.query(input_hashes).set_index("input_hash") - except _CACHE_OPERATION_ERRORS as exc: - logger.warning( - "Cache read failed at %s: %s. Continuing without caching.", - store.db_path, - exc, - ) - return None - - -def _save_cache_outputs(cache, store, model, input_texts, outputs, metadata, indices): - try: - cache.save_outputs(store, model, input_texts, outputs, metadata, indices) - except _CACHE_OPERATION_ERRORS as exc: - logger.warning( - "Cache write failed at %s: %s. Preserving generated results.", - store.db_path, - exc, - ) - - def do_inference( chat_model, inputs, @@ -853,19 +814,34 @@ def do_inference( if cache_row_metadata is None or len(cache_row_metadata) != len(inputs): raise ValueError("cache_row_metadata must contain one row per inference input.") - input_texts = [canonicalize_chat_input(item) for item in inputs] + input_mode = chat_model.descriptor["input_mode"] + input_texts = [canonicalize_model_input(item, input_mode) for item in inputs] input_hashes = [input_hash(input_text) for input_text in input_texts] - with _cache_store(cache, chat_model) as store: - if store is None: - return _run_backend_inference( - chat_model.materialize(), - inputs, - use_tqdm, - return_top_logprobs, - stage=stage, + try: + store = cache.open_store(chat_model) + except _CACHE_OPERATION_ERRORS as exc: + logger.warning( + "Cache open failed at %s: %s. Continuing without caching.", + cache.store_root, + exc, + ) + return _run_backend_inference( + chat_model.materialize(), + inputs, + use_tqdm, + return_top_logprobs, + stage=stage, + ) + + try: + try: + cached_rows = store.query(input_hashes).set_index("input_hash") + except _CACHE_OPERATION_ERRORS as exc: + logger.warning( + "Cache read failed at %s: %s. Continuing without caching.", + store.db_path, + exc, ) - cached_rows = _read_cached_rows(store, input_hashes) - if cached_rows is None: return _run_backend_inference( chat_model.materialize(), inputs, @@ -894,15 +870,26 @@ def do_inference( ) for index, result in zip(missing_indices, generated, strict=True): results[index] = result - _save_cache_outputs( - cache, - store, - chat_model, - input_texts, - generated, - cache_row_metadata, - missing_indices, - ) + try: + cache.save_outputs( + store, + chat_model, + input_texts, + generated, + cache_row_metadata, + missing_indices, + ) + except _CACHE_OPERATION_ERRORS as exc: + logger.warning( + "Cache write failed at %s: %s. Preserving generated results.", + store.db_path, + exc, + ) + finally: + try: + store.close() + except _CACHE_OPERATION_ERRORS as exc: + logger.warning("Cache close failed at %s: %s.", store.db_path, exc) resolved_results = [result for result in results if result is not None] if return_top_logprobs: @@ -952,6 +939,62 @@ def _route_sampling_params( return engine_kwargs +# Prevent local-engine settings from leaking into hosted provider constructors. +_VLLM_ONLY_KWARGS = ( + "max_model_len", + "chat_template", + "language_model_only", + "gpu_memory_utilization", + "enforce_eager", + "tensor_parallel_size", + "quantization", + "kv_cache_dtype", + "reasoning_parser", + "reasoning_config", + "trust_remote_code", +) +_REMOTE_ENDPOINTS = { + "ChatOpenAI": "https://api.openai.com/v1", + "OpenAI": "https://api.openai.com/v1", + "OpenRouter": "https://openrouter.ai/api/v1", + "Together": "https://api.together.xyz/v1/completions", +} + + +def _resolve_model_config( + model: str, + max_tokens: int | None, + engine_kwargs: dict, +) -> tuple[str, str, dict, dict]: + """Resolve shared provider and sampling inputs once.""" + resolved_kwargs = engine_kwargs.copy() + resolved_kwargs["max_tokens"] = max_tokens or 8192 + sampling = { + key: resolved_kwargs.pop(key, None) + for key in ("temperature", "top_p", "top_k", "seed", "top_logprobs") + } + provider, model_name = _split_model_spec(model) + if provider != "VLLM": + for key in _VLLM_ONLY_KWARGS: + resolved_kwargs.pop(key, None) + return provider, model_name, resolved_kwargs, sampling + + +def _resolve_model_endpoint(provider: str, resolved_kwargs: dict) -> str | None: + if provider == "OpenRouter": + return _REMOTE_ENDPOINTS[provider] + for key in ("base_url", "openai_api_base", "together_api_base"): + if resolved_kwargs.get(key): + return str(resolved_kwargs[key]) + if provider in {"ChatOpenAI", "OpenAI"}: + return ( + os.getenv("OPENAI_BASE_URL") + or os.getenv("OPENAI_API_BASE") + or _REMOTE_ENDPOINTS[provider] + ) + return _REMOTE_ENDPOINTS.get(provider) + + def prepare_model( model: str, max_tokens: int | None = 8192, @@ -960,13 +1003,21 @@ def prepare_model( **engine_kwargs, ) -> PreparedModel: """Prepare cache identity without constructing the provider backend.""" - provider, model_name = _split_model_spec(model) - resolved_kwargs = {**engine_kwargs, "max_tokens": max_tokens or 8192} + provider, model_name, resolved_kwargs, sampling = _resolve_model_config( + model, max_tokens, engine_kwargs + ) descriptor = ( - build_model_descriptor(provider, model_name, resolved_kwargs) + build_model_descriptor( + provider, + model_name, + {**resolved_kwargs, **sampling}, + endpoint=_resolve_model_endpoint(provider, resolved_kwargs), + ) if cache is not None else None ) + if cache is not None and descriptor is None: + logger.warning("Caching is not supported for %s; running uncached.", model) factory_kwargs = engine_kwargs.copy() return PreparedModel( model_spec=model, @@ -992,40 +1043,14 @@ def make_model(model: str, max_tokens: int | None = 8192, **engine_kwargs): ``top_k``, ``seed``. vLLM-only keys (``max_model_len``, ``chat_template``) are stripped before reaching hosted providers. """ - # Avoid mutating the original engine_kwargs dictionary - # NOTE: this is a shallow copy since we are not modifying any - # mutable objects in the dictionary. - engine_kwargs = engine_kwargs.copy() - - # Dedicated arguments like max_tokens always win over engine_kwargs. - engine_kwargs["max_tokens"] = max_tokens or 8192 - - temperature = engine_kwargs.pop("temperature", None) - top_p = engine_kwargs.pop("top_p", None) - top_k = engine_kwargs.pop("top_k", None) - seed = engine_kwargs.pop("seed", None) - top_logprobs = engine_kwargs.pop("top_logprobs", None) - - model_provider, model_name = _split_model_spec(model) - - # vLLM-engine-only kwargs must not leak to remote-API providers - # (OpenRouter, OpenAI, Together): langchain-openai forwards unknown - # kwargs via model_kwargs into chat.completions.create, which rejects them. - if model_provider != "VLLM": - for key in ( - "max_model_len", - "chat_template", - "language_model_only", - "gpu_memory_utilization", - "enforce_eager", - "tensor_parallel_size", - "quantization", - "kv_cache_dtype", - "reasoning_parser", - "reasoning_config", - "trust_remote_code", - ): - engine_kwargs.pop(key, None) + model_provider, model_name, engine_kwargs, sampling = _resolve_model_config( + model, max_tokens, engine_kwargs + ) + temperature = sampling["temperature"] + top_p = sampling["top_p"] + top_k = sampling["top_k"] + seed = sampling["seed"] + top_logprobs = sampling["top_logprobs"] if model_provider == "Dummy": if top_logprobs is not None: diff --git a/tests/test_chat_vllm.py b/tests/test_chat_vllm.py index ddaad5fa..982e04bb 100644 --- a/tests/test_chat_vllm.py +++ b/tests/test_chat_vllm.py @@ -64,6 +64,8 @@ def test_chat_vllm_enables_reasoning_support_for_qwen_thinking_budget(monkeypatc gpu_memory_utilization=0.7, ) + assert captured["sampling_kwargs"]["temperature"] == models.VLLM_DEFAULT_TEMPERATURE + assert captured["sampling_kwargs"]["top_p"] == models.VLLM_DEFAULT_TOP_P assert captured["sampling_kwargs"]["thinking_token_budget"] == 64 assert "structured_outputs" not in captured["sampling_kwargs"] assert captured["reasoning_config_kwargs"] == { diff --git a/tests/test_inference_cache.py b/tests/test_inference_cache.py index 686e2628..d61c4d90 100644 --- a/tests/test_inference_cache.py +++ b/tests/test_inference_cache.py @@ -1,10 +1,18 @@ +import json import sqlite3 import pytest from langchain_core.messages import AIMessage +from langchain_core.prompts import ChatPromptTemplate +import judgearena.inference as inference import judgearena.models as models -from judgearena.inference import CompletionInferenceCache, JudgementInferenceCache +from judgearena.inference import ( + CompletionInferenceCache, + JudgementInferenceCache, + canonicalize_model_input, + provider_input_mode, +) from judgearena.models import InferenceResult, do_inference, prepare_model from judgearena.usage import track_usage @@ -109,6 +117,101 @@ def batch(self, inputs, **_kwargs): assert second[0].usage is None +def test_vllm_descriptor_contains_output_configuration(tmp_path, monkeypatch): + monkeypatch.setattr(inference.importlib_metadata, "version", lambda _name: "0.10.2") + cache = CompletionInferenceCache(tmp_path, "arena-hard") + + descriptor = prepare_model( + "VLLM/Qwen/Qwen3-8B", + max_tokens=32, + cache=cache, + temperature=0.2, + enforce_eager=True, + gpu_memory_utilization=0.9, + max_model_len=4096, + tensor_parallel_size=2, + ).descriptor + + assert descriptor["backend_version"] == "0.10.2" + assert descriptor["model_kwargs"]["temperature"] == 0.2 + assert descriptor["model_kwargs"]["top_p"] == 0.95 + assert descriptor["model_kwargs"]["max_model_len"] == 4096 + assert "tensor_parallel_size" not in descriptor["model_kwargs"] + + +@pytest.mark.parametrize( + ("model_spec", "expected_mode", "expected_endpoint"), + [ + ("Dummy/model", "chat", None), + ("VLLM/org/model", "chat", None), + ("OpenRouter/org/model", "chat", "https://openrouter.ai/api/v1"), + ("ChatOpenAI/model", "chat", "https://api.openai.com/v1"), + ("OpenAI/model", "text", "https://api.openai.com/v1"), + ("Together/org/model", "text", "https://api.together.xyz/v1/completions"), + ("LlamaCpp/./models/model.gguf", "text", None), + ], +) +def test_supported_provider_descriptors( + tmp_path, + monkeypatch, + model_spec, + expected_mode, + expected_endpoint, +): + versions = {"vllm": "0.10.2", "llama-cpp-python": "0.3.0"} + monkeypatch.setattr(inference.importlib_metadata, "version", versions.__getitem__) + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + cache = CompletionInferenceCache(tmp_path, "arena-hard") + + descriptor = prepare_model(model_spec, cache=cache).descriptor + + assert descriptor["input_mode"] == expected_mode + assert descriptor.get("endpoint") == expected_endpoint + + +def test_hosted_descriptor_hashes_routing_and_endpoint(tmp_path, monkeypatch): + cache = CompletionInferenceCache(tmp_path, "arena-hard") + unpinned = prepare_model("OpenRouter/org/model", cache=cache).descriptor + + pinned = prepare_model( + "OpenRouter/org/model", + cache=cache, + extra_body={"provider": {"order": ["Together"], "allow_fallbacks": False}}, + ).descriptor + assert unpinned != pinned + + monkeypatch.setenv("OPENAI_BASE_URL", "https://ambient.example/v1") + ambient = prepare_model("ChatOpenAI/model", cache=cache).descriptor + assert ambient["endpoint"] == "https://ambient.example/v1" + + +def test_unsupported_provider_runs_uncached(tmp_path, caplog): + model = prepare_model( + "Unsupported/model", + cache=CompletionInferenceCache(tmp_path, "arena-hard"), + ) + + assert model.descriptor is None + assert "Caching is not supported" in caplog.text + + +@pytest.mark.parametrize( + ("provider", "expected_type"), + [("OpenRouter", "messages"), ("Together", "text"), ("VLLM", "messages")], +) +def test_input_canonicalization_matches_provider_mode(provider, expected_type): + prompt = ChatPromptTemplate.from_messages( + [("system", "System"), ("user", "Question")] + ).invoke({}) + + payload = json.loads( + canonicalize_model_input(prompt, provider_input_mode(provider)) + ) + + assert payload["type"] == expected_type + + def test_cache_write_failure_preserves_generated_results(tmp_path, monkeypatch): cache = CompletionInferenceCache(tmp_path, "arena-hard") backend = EchoModel()