diff --git a/judgearena/cache/inference.py b/judgearena/cache/inference.py index f6aef317..63b79d07 100644 --- a/judgearena/cache/inference.py +++ b/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( diff --git a/judgearena/inference.py b/judgearena/inference.py index 939416ef..51e49b71 100644 --- a/judgearena/inference.py +++ b/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 diff --git a/judgearena/models.py b/judgearena/models.py index 464fa431..585db28c 100644 --- a/judgearena/models.py +++ b/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)) diff --git a/pyproject.toml b/pyproject.toml index 2f3586b3..2e214bdc 100644 --- a/pyproject.toml +++ b/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"] diff --git a/tests/test_chat_vllm.py b/tests/test_chat_vllm.py index 7d54c9dd..b8291bfb 100644 --- a/tests/test_chat_vllm.py +++ b/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