Judge tuning (1/5): Fix vLLM compatibility and context failures

judgearena/cache/inference.py

@@ -199,8 +199,10 @@ def save_outputs(
                 metadata=metadata[index],
             )
             for index, output in zip(indices, outputs, strict=True)
+            if output.error is None
         ]
-        store.save(pd.DataFrame(rows), pushed_by=self.pushed_by)
+        if rows:
+            store.save(pd.DataFrame(rows), pushed_by=self.pushed_by)
 
     @abstractmethod
     def make_row(

judgearena/inference.py

@@ -14,3 +14,4 @@ class InferenceResult:
     text: str
     first_token_top_logprobs: dict[str, float] | None = None
     usage: RequestUsage | None = None
+    error: str | None = None

judgearena/models.py

@@ -419,14 +419,43 @@ def _run_raw_batch(self, inputs: list):
             prompts = [self._to_raw_text(inp) for inp in inputs]
             outputs = self.llm.generate(prompts, self.sampling_params)
         else:
+            from vllm.exceptions import VLLMValidationError
+
             messages_batch = [self._to_messages(inp) for inp in inputs]
-            outputs = self.llm.chat(
-                messages_batch,
-                self.sampling_params,
-                add_generation_prompt=True,
-                chat_template=self.chat_template,
-                chat_template_kwargs=self._chat_template_kwargs,
-            )
+            valid_indices = []
+            for index, messages in enumerate(messages_batch):
+                try:
+                    # Private vLLM API (0.19-0.25): renders and length-checks one
+                    # chat request exactly as LLM.chat() does.
+                    self.llm._preprocess_chat_one(
+                        messages,
+                        add_generation_prompt=True,
+                        chat_template=self.chat_template,
+                        chat_template_kwargs=self._chat_template_kwargs,
+                        tokenization_kwargs={"max_output_tokens": 1},
+                    )
+                except VLLMValidationError as error:
+                    if error.parameter != "input_tokens":
+                        raise
+                else:
+                    valid_indices.append(index)
+            outputs = [None] * len(inputs)
+            if valid_indices:
+                generated = self.llm.chat(
+                    [messages_batch[index] for index in valid_indices],
+                    self.sampling_params,
+                    add_generation_prompt=True,
+                    chat_template=self.chat_template,
+                    chat_template_kwargs=self._chat_template_kwargs,
+                )
+                for index, output in zip(valid_indices, generated, strict=True):
+                    outputs[index] = output
+            if len(valid_indices) < len(inputs):
+                logger.warning(
+                    "Skipped %d requests exceeding the model context; "
+                    "recording failed annotations.",
+                    len(inputs) - len(valid_indices),
+                )
         return outputs
 
     def batch(self, inputs: list, **invoke_kwargs) -> list:
@@ -437,10 +466,17 @@ def batch(self, inputs: list, **invoke_kwargs) -> list:
         """
         usage_stage = invoke_kwargs.pop("usage_stage", None)
         outputs = self._run_raw_batch(inputs)
-        if self._top_logprobs is None and usage_stage is None:
+        if (
+            self._top_logprobs is None
+            and usage_stage is None
+            and all(out is not None for out in outputs)
+        ):
             return [out.outputs[0].text for out in outputs]
         results = []
         for out in outputs:
+            if out is None:
+                results.append(InferenceResult(text="", error="context_length"))
+                continue
             generation = out.outputs[0]
             top = None
             if getattr(generation, "logprobs", None):
@@ -611,7 +647,7 @@ def _to_inference_result(
     default_model: str | None,
 ) -> InferenceResult:
     if isinstance(response, InferenceResult):
-        if response.usage is not None:
+        if response.usage is not None or response.error is not None:
             return response
         return replace(
             response,
@@ -649,7 +685,7 @@ def _collect_inference_results(
         )
         for response in responses
     ]
-    request_usage = [result.usage for result in results]
+    request_usage = [result.usage for result in results if result.usage is not None]
     record_usage(request_usage)
 
     batch_usage = RunUsage(tuple(request_usage))

pyproject.toml

@@ -111,5 +111,6 @@ indent-style = "space"
 vllm = [
     "vllm>=0.19.1,<1.0.0",
     "transformers>=5.5.1,<6.0.0",
+    "openai>=2.30.0", # vLLM 0.25 imports responses.NamespaceTool.
 ]
 llamacpp = ["llama-cpp-python>=0.3.0"]

tests/test_chat_vllm.py

@@ -1,4 +1,5 @@
 import json
+import sqlite3
 import sys
 from types import SimpleNamespace
 
@@ -9,6 +10,9 @@
 def _install_fake_vllm(monkeypatch):
     captured = {}
 
+    class FakeValidationError(ValueError):
+        parameter = "input_tokens"
+
     class FakeSamplingParams:
         def __init__(self, **kwargs):
             captured["sampling_kwargs"] = kwargs
@@ -28,6 +32,10 @@ def __init__(self, *, model, trust_remote_code, **kwargs):
         def get_tokenizer(self):
             return SimpleNamespace(chat_template="{{ messages }}")
 
+        def _preprocess_chat_one(self, messages, **kwargs):
+            if messages[0]["content"] == "overlong":
+                raise FakeValidationError("too long")
+
         def chat(self, messages, sampling_params, **kwargs):
             captured["chat_call"] = {
                 "messages": messages,
@@ -46,6 +54,11 @@ def chat(self, messages, sampling_params, **kwargs):
         "vllm",
         SimpleNamespace(LLM=FakeLLM, SamplingParams=FakeSamplingParams),
     )
+    monkeypatch.setitem(
+        sys.modules,
+        "vllm.exceptions",
+        SimpleNamespace(VLLMValidationError=FakeValidationError),
+    )
     monkeypatch.setitem(
         sys.modules,
         "vllm.config.reasoning",
@@ -321,3 +334,38 @@ async def fail_if_called(*_args, **_kwargs):
     monkeypatch.setattr(chat_model, "ainvoke", fail_if_called)
 
     assert models.do_inference(chat_model, ["hello"], use_tqdm=True) == ["ok"]
+
+
+def test_chat_vllm_skips_overlong_requests_without_caching(tmp_path, monkeypatch):
+    from judgearena.cache.inference import JudgementInferenceCache
+
+    captured, _ = _install_fake_vllm(monkeypatch)
+    backend = models.ChatVLLM(model="test/model", max_tokens=128)
+    monkeypatch.setattr(models, "make_model", lambda *_args, **_kwargs: backend)
+    cache = JudgementInferenceCache(tmp_path, "meta-eval-test")
+    metadata = [
+        {
+            "instruction_id": str(i),
+            "model_a": "a",
+            "model_b": "b",
+            "orientation": "direct",
+        }
+        for i in range(2)
+    ]
+    prepared = models.prepare_model("Dummy/judge", cache=cache)
+    with track_usage() as tracker:
+        assert models.do_inference(
+            prepared, ["overlong", "valid"], cache_row_metadata=metadata
+        ) == ["", "ok"]
+    assert len(tracker.snapshot().requests) == 1
+    assert [r.error for r in backend.batch(["overlong", "valid"])] == [
+        "context_length",
+        None,
+    ]
+    assert captured["chat_call"]["messages"] == [[{"role": "user", "content": "valid"}]]
+    assert models.do_inference(
+        prepared, ["overlong", "valid"], cache_row_metadata=metadata
+    ) == ["", "ok"]
+    db = next(tmp_path.rglob("judgements.db"))
+    with sqlite3.connect(db) as connection:
+        assert connection.execute("SELECT COUNT(*) FROM judgements").fetchone()[0] == 1