diff --git a/tests/entrypoints/pooling/scoring/test_io_processor_unit.py b/tests/entrypoints/pooling/scoring/test_io_processor_unit.py index 5a0ff9363e..b8511ea2c9 100644 --- a/tests/entrypoints/pooling/scoring/test_io_processor_unit.py +++ b/tests/entrypoints/pooling/scoring/test_io_processor_unit.py @@ -1,20 +1,25 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for ScoringIOProcessor post-tokenization helpers.""" +"""Unit tests for ScoringIOProcessor helpers.""" from dataclasses import dataclass from types import SimpleNamespace import pytest +import torch from tokenizers import Tokenizer, models, pre_tokenizers, processors from transformers import TokenizersBackend from vllm import TokensPrompt from vllm.entrypoints.pooling.scoring.io_processor import ( + BiEncoderIOProcessor, CrossEncoderIOProcessor, + JinaRankingIOProcessor, + LateInteractionIOProcessor, _apply_post_tokenization_to_token_type_ids, ) from vllm.entrypoints.pooling.scoring.utils import compress_token_type_ids +from vllm.outputs import PoolingOutput, PoolingRequestOutput, RequestError from vllm.renderers import TokenizeParams pytestmark = pytest.mark.skip_global_cleanup @@ -113,3 +118,59 @@ def test_token_type_ids_stay_aligned_with_a_truncated_padded_prompt(): assert first_doc == 10 assert prompt_token_ids[:first_doc] == list(range(10, num_query)) assert prompt_token_ids[first_doc:40] == list(range(num_query, 50)) + + +def _pooling_output( + request_id: str, + data: torch.Tensor, + error: RequestError | None = None, +) -> PoolingRequestOutput: + return PoolingRequestOutput( + request_id=request_id, + outputs=PoolingOutput(data), + prompt_token_ids=[1], + num_cached_tokens=0, + finished=True, + error=error, + ) + + +@pytest.mark.parametrize( + "processor_type", [BiEncoderIOProcessor, LateInteractionIOProcessor] +) +@pytest.mark.parametrize("failed_index", [0, 1]) +def test_paired_scoring_preserves_request_error(processor_type, failed_index): + processor = processor_type.__new__(processor_type) + processor.pad_token_id = None + error = RequestError( + code="multimodal_cache_miss", + message="Multi-modal processor cache miss.", + retryable=True, + ) + outputs = [ + _pooling_output("query", torch.tensor([1.0, 2.0])), + _pooling_output("document", torch.tensor([3.0, 4.0])), + ] + outputs[failed_index] = _pooling_output( + outputs[failed_index].request_id, + torch.empty(0), + error, + ) + + (result,) = processor._post_process(outputs, n_queries=1) + + assert result.request_id == "query_document" + assert result.error is error + assert result.outputs is outputs[failed_index].outputs + + +def test_jina_scoring_preserves_request_error_before_indexing_output(): + processor = JinaRankingIOProcessor.__new__(JinaRankingIOProcessor) + error = RequestError( + code="multimodal_cache_miss", + message="Multi-modal processor cache miss.", + retryable=True, + ) + failed = _pooling_output("request", torch.empty(0), error) + + assert processor._post_process([failed], n_queries=1) == [failed] diff --git a/tests/entrypoints/pooling/scoring/test_late_interaction_serving.py b/tests/entrypoints/pooling/scoring/test_late_interaction_serving.py index def1dc264f..f295c5fc04 100644 --- a/tests/entrypoints/pooling/scoring/test_late_interaction_serving.py +++ b/tests/entrypoints/pooling/scoring/test_late_interaction_serving.py @@ -49,10 +49,60 @@ def _query_key(context: PoolingServeContext) -> str: @pytest.mark.asyncio -async def test_colliding_request_ids_use_distinct_query_cache_keys( +@pytest.mark.skip_global_cleanup +@pytest.mark.parametrize("failed_stage", ["query", "doc"]) +async def test_flash_late_interaction_cleans_up_queries_on_failure( + failed_stage: str, monkeypatch: pytest.MonkeyPatch, ): serving = object.__new__(ServingScores) + serving.io_processor = Mock() + serving.engine_client = Mock(abort=AsyncMock(), collective_rpc=AsyncMock()) + + context = _make_context() + monkeypatch.setattr(serving, "_init_ctx", AsyncMock(return_value=context)) + monkeypatch.setattr(serving, "_preprocessing", AsyncMock()) + + async def encode_queries(ctx): + ctx.late_interaction_query_keys = ["query-key"] + if failed_stage == "query": + raise RuntimeError("query failed") + + async def encode_docs(ctx): + ctx.late_interaction_doc_keys = ["doc-key"] + if failed_stage == "doc": + raise RuntimeError("doc failed") + + monkeypatch.setattr( + serving, + "_flash_late_interaction_encode_queries", + AsyncMock(side_effect=encode_queries), + ) + monkeypatch.setattr( + serving, + "_flash_late_interaction_encode_docs", + AsyncMock(side_effect=encode_docs), + ) + + with pytest.raises(RuntimeError, match=f"{failed_stage} failed"): + await serving.flash_late_interaction() + + expected_doc_keys = [] if failed_stage == "query" else ["doc-key"] + serving.engine_client.abort.assert_awaited_once_with( + ["query-key", *expected_doc_keys] + ) + serving.engine_client.collective_rpc.assert_awaited_once_with( + "release_late_interaction_query_cache", + args=(["query-key"],), + ) + + +@pytest.mark.asyncio +async def test_colliding_request_ids_use_distinct_late_interaction_keys( + monkeypatch: pytest.MonkeyPatch, +): + serving = object.__new__(ServingScores) + serving.engine_client = Mock(abort=AsyncMock(), collective_rpc=AsyncMock()) prepare_generators = AsyncMock() monkeypatch.setattr(serving, "_prepare_generators", prepare_generators) monkeypatch.setattr(serving, "_collect_batch", AsyncMock()) @@ -68,9 +118,24 @@ async def test_colliding_request_ids_use_distinct_query_cache_keys( first_doc_key = _query_key(prepared_contexts[1]) second_query_key = _query_key(prepared_contexts[2]) second_doc_key = _query_key(prepared_contexts[3]) + first_doc_request_ids = prepared_contexts[1].prompt_request_ids + second_doc_request_ids = prepared_contexts[3].prompt_request_ids assert first_query_key == first_doc_key assert second_query_key == second_doc_key assert first_query_key != second_query_key assert "caller-controlled-id" not in first_query_key assert "caller-controlled-id" not in second_query_key + assert first_doc_request_ids is not None + assert second_doc_request_ids is not None + assert first.late_interaction_doc_keys == first_doc_request_ids + assert second.late_interaction_doc_keys == second_doc_request_ids + assert first_doc_request_ids != second_doc_request_ids + assert all("caller-controlled-id" not in key for key in first_doc_request_ids) + assert all("caller-controlled-id" not in key for key in second_doc_request_ids) + + assert first.late_interaction_query_keys is not None + await serving._cleanup_flash_late_interaction(first) + serving.engine_client.abort.assert_awaited_once_with( + [*first.late_interaction_query_keys, *first_doc_request_ids] + ) diff --git a/tests/test_outputs.py b/tests/test_outputs.py index 7b234884c5..9acc1f7b74 100644 --- a/tests/test_outputs.py +++ b/tests/test_outputs.py @@ -2,8 +2,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest +import torch -from vllm.outputs import RequestOutput +from vllm.outputs import ( + ClassificationRequestOutput, + EmbeddingRequestOutput, + PoolingOutput, + PoolingRequestOutput, + RequestError, + RequestOutput, + ScoringRequestOutput, +) pytestmark = pytest.mark.cpu_test @@ -19,3 +28,36 @@ def test_request_output_forward_compatible(): example_arg_added_in_new_version="some_value", ) assert output is not None + + +@pytest.mark.parametrize( + ("output_type", "empty_value"), + [ + (EmbeddingRequestOutput, []), + (ClassificationRequestOutput, []), + (ScoringRequestOutput, 0.0), + ], +) +@pytest.mark.skip_global_cleanup +def test_specialized_pooling_outputs_preserve_request_error(output_type, empty_value): + error = RequestError( + code="multimodal_cache_miss", + message="Multi-modal processor cache miss.", + retryable=True, + ) + base_output = PoolingRequestOutput( + request_id="request", + outputs=PoolingOutput(torch.empty(0)), + prompt_token_ids=[1, 2], + num_cached_tokens=0, + finished=True, + error=error, + ) + + output = output_type.from_base(base_output) + + assert output.error is error + value = getattr(output.outputs, "embedding", None) + value = getattr(output.outputs, "probs", value) + value = getattr(output.outputs, "score", value) + assert value == empty_value diff --git a/tests/v1/engine/test_llm_engine.py b/tests/v1/engine/test_llm_engine.py index c9832f7d8e..b09f9405d7 100644 --- a/tests/v1/engine/test_llm_engine.py +++ b/tests/v1/engine/test_llm_engine.py @@ -2,11 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import random from typing import TYPE_CHECKING +from unittest.mock import MagicMock, call import pytest from vllm import LLM from vllm.sampling_params import SamplingParams, StructuredOutputsParams +from vllm.v1.engine.llm_engine import LLMEngine from vllm.v1.metrics.reader import Counter, Gauge, Histogram, Metric, Vector if TYPE_CHECKING: @@ -18,6 +20,38 @@ MODEL = "facebook/opt-125m" DTYPE = "half" +@pytest.mark.skip_global_cleanup +def test_step_invalidates_mm_processor_cache_misses(): + engine = object.__new__(LLMEngine) + engine.should_execute_dummy_batch = False + engine.log_stats = False + engine.logger_manager = None + + mm_processor_cache = MagicMock() + engine.renderer = MagicMock(mm_processor_cache=mm_processor_cache) + + engine_core_output = MagicMock( + mm_cache_miss_hashes=["missing-hash-1", "missing-hash-2"] + ) + outputs = MagicMock( + outputs=[engine_core_output], + timestamp=0.0, + scheduler_stats=None, + ) + engine.engine_core = MagicMock() + engine.engine_core.get_output.return_value = outputs + + processed_outputs = MagicMock(reqs_to_abort=[], request_outputs=[]) + engine.output_processor = MagicMock() + engine.output_processor.process_outputs.return_value = processed_outputs + + assert engine.step() == [] + assert mm_processor_cache.invalidate.call_args_list == [ + call("missing-hash-1"), + call("missing-hash-2"), + ] + + def _vllm_model( apc: bool, vllm_runner: type[VllmRunner], diff --git a/tests/v1/engine/test_output_processor.py b/tests/v1/engine/test_output_processor.py index 40a6883aa9..44c3ee20b5 100644 --- a/tests/v1/engine/test_output_processor.py +++ b/tests/v1/engine/test_output_processor.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, Mock import numpy as np import pytest +import torch from tests.v1.engine.utils import ( NUM_PROMPT_LOGPROBS_UNDER_TEST, @@ -18,7 +19,12 @@ from tests.v1.engine.utils import ( from vllm import PoolingParams from vllm.logprobs import FlatLogprobs, Logprob, PromptLogprobs, SampleLogprobs from vllm.lora.request import LoRARequest -from vllm.outputs import CompletionOutput, RequestOutput, SamplingMask +from vllm.outputs import ( + CompletionOutput, + PoolingRequestOutput, + RequestOutput, + SamplingMask, +) from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.tokenizers import TokenizerLike from vllm.v1.engine import ( @@ -1630,3 +1636,92 @@ def test_request_output_add_merges_delta_sampling_masks(): assert merged.outputs[0].token_ids == [10, 20] assert merged.outputs[0].sampling_mask.token_ids == [[10, 11], [20]] + + +@pytest.mark.skip_global_cleanup +def test_pooling_cache_miss_is_a_request_error_and_batch_continues(): + output_processor = OutputProcessor(tokenizer=None, log_stats=False) + requests = [ + EngineCoreRequest( + request_id=f"request-{idx}-int", + external_req_id=f"request-{idx}", + prompt_token_ids=[1, 2], + mm_features=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + sampling_params=None, + pooling_params=PoolingParams(task="embed"), + ) + for idx in range(2) + ] + for request in requests: + output_processor.add_request(request, prompt=None) + + processed = output_processor.process_outputs( + [ + EngineCoreOutput( + request_id=requests[0].request_id, + new_token_ids=[], + finish_reason=FinishReason.ERROR, + mm_cache_miss_hashes=["missing-hash"], + ), + EngineCoreOutput( + request_id=requests[1].request_id, + new_token_ids=[], + pooling_output=torch.tensor([1.0, 2.0]), + finish_reason=FinishReason.STOP, + ), + ] + ) + + assert len(processed.request_outputs) == 2 + failed, succeeded = processed.request_outputs + assert isinstance(failed, PoolingRequestOutput) + assert failed.error is not None + assert failed.error.code == "multimodal_cache_miss" + assert failed.error.retryable + assert isinstance(succeeded, PoolingRequestOutput) + assert succeeded.error is None + assert torch.equal(succeeded.outputs.data, torch.tensor([1.0, 2.0])) + + +@pytest.mark.skip_global_cleanup +def test_async_pooling_cache_miss_does_not_fail_output_processor(): + output_processor = OutputProcessor(tokenizer=None, log_stats=False) + request = EngineCoreRequest( + request_id="request-int", + external_req_id="request", + prompt_token_ids=[1, 2], + mm_features=None, + arrival_time=0, + lora_request=None, + cache_salt=None, + data_parallel_rank=None, + sampling_params=None, + pooling_params=PoolingParams(task="embed"), + ) + queue = RequestOutputCollector( + output_kind=RequestOutputKind.FINAL_ONLY, + request_id=request.request_id, + ) + output_processor.add_request(request, prompt=None, queue=queue) + + processed = output_processor.process_outputs( + [ + EngineCoreOutput( + request_id=request.request_id, + new_token_ids=[], + finish_reason=FinishReason.ERROR, + mm_cache_miss_hashes=["missing-hash"], + ) + ] + ) + + assert processed.request_outputs == [] + output = queue.get_nowait() + assert isinstance(output, PoolingRequestOutput) + assert output.error is not None + assert output.error.retryable + assert not output_processor.has_unfinished_requests() diff --git a/tests/v1/worker/test_late_interaction_runner.py b/tests/v1/worker/test_late_interaction_runner.py index 4befbf72cc..35ce3670a7 100644 --- a/tests/v1/worker/test_late_interaction_runner.py +++ b/tests/v1/worker/test_late_interaction_runner.py @@ -137,6 +137,52 @@ def test_finished_request_releases_unscored_doc_use(): ) +@pytest.mark.skip_global_cleanup +def test_release_queries_is_targeted_and_idempotent(): + runner = LateInteractionRunner() + query_emb = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float32) + + for query_key in ("query-release", "query-keep"): + query_params = _make_pooling_params( + build_late_interaction_query_params(query_key=query_key, query_uses=1) + ) + runner.postprocess_pooler_output( + raw_pooler_output=[query_emb], + pooling_params=[query_params], + req_ids=[f"{query_key}-req"], + finished_mask=[True], + ) + + released_doc_params = _make_pooling_params( + build_late_interaction_doc_params(query_key="query-release") + ) + kept_doc_params = _make_pooling_params( + build_late_interaction_doc_params(query_key="query-keep") + ) + runner.register_request("released-doc", released_doc_params) + runner.register_request("kept-doc", kept_doc_params) + + runner.release_queries(["query-release"]) + runner.release_queries(["query-release"]) + + with pytest.raises(ValueError, match="query cache miss"): + runner.postprocess_pooler_output( + raw_pooler_output=[query_emb], + pooling_params=[released_doc_params], + req_ids=["released-doc"], + finished_mask=[True], + ) + + kept_output = runner.postprocess_pooler_output( + raw_pooler_output=[query_emb], + pooling_params=[kept_doc_params], + req_ids=["kept-doc"], + finished_mask=[True], + ) + assert isinstance(kept_output, list) + assert kept_output[0] is not None + + def test_invalid_query_uses_raises(): runner = LateInteractionRunner() bad_meta = LateInteractionParams( diff --git a/vllm/entrypoints/pooling/base/serving.py b/vllm/entrypoints/pooling/base/serving.py index c0c0c15e9c..92cdfaec25 100644 --- a/vllm/entrypoints/pooling/base/serving.py +++ b/vllm/entrypoints/pooling/base/serving.py @@ -21,6 +21,7 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.engine.serving import BaseServing from vllm.entrypoints.serve.engine.typing import AnyRequest from vllm.entrypoints.serve.utils.request_logger import RequestLogger +from vllm.exceptions import GenerationError, RetryableRequestError from vllm.lora.request import LoRARequest from vllm.renderers.base import BaseRenderer from vllm.tracing import ( @@ -206,6 +207,10 @@ class PoolingBaseServing(ABC, BaseServing): final_res_batch = [None] * num_inputs async for i, res in ctx.result_generator: + if res.error is not None: + if res.error.retryable: + raise RetryableRequestError(res.error.message) + raise GenerationError(res.error.message) final_res_batch[i] = res if None in final_res_batch: diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index 8ffa2e4ddd..91c5f762cf 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -191,6 +191,27 @@ class ScoringIOProcessor(PoolingIOProcessor): scoring_data = self.valid_inputs(data_1, data_2) return scoring_data + def _pair_error_output( + self, + output_1: PoolingRequestOutput, + output_2: PoolingRequestOutput, + ) -> PoolingRequestOutput | None: + failed = output_1 if output_1.error is not None else output_2 + if failed.error is None: + return None + + padding = [self.pad_token_id] if self.pad_token_id is not None else [] + return PoolingRequestOutput( + request_id=f"{output_1.request_id}_{output_2.request_id}", + outputs=failed.outputs, + prompt_token_ids=( + output_1.prompt_token_ids + padding + output_2.prompt_token_ids + ), + num_cached_tokens=(output_1.num_cached_tokens + output_2.num_cached_tokens), + finished=failed.finished, + error=failed.error, + ) + class BiEncoderIOProcessor(ScoringIOProcessor): name = "bi-encoder" @@ -310,6 +331,10 @@ class BiEncoderIOProcessor(ScoringIOProcessor): final_res_batch: list[PoolingRequestOutput] = [] for emb_1, emb_2 in zip(emb_data_1, emb_data_2): + if error_output := self._pair_error_output(emb_1, emb_2): + final_res_batch.append(error_output) + continue + pair_score = F.cosine_similarity( emb_1.outputs.data.float(), emb_2.outputs.data.float(), dim=0 ) @@ -352,6 +377,10 @@ class LateInteractionIOProcessor(BiEncoderIOProcessor): # Compute MaxSim scores for emb_1, emb_2 in zip(emb_data_1, emb_data_2): + if error_output := self._pair_error_output(emb_1, emb_2): + final_res_batch.append(error_output) + continue + # emb_1.outputs.data: [query_len, dim] # emb_2.outputs.data: [doc_len, dim] q_emb = emb_1.outputs.data @@ -907,6 +936,10 @@ class JinaRankingIOProcessor(LateInteractionIOProcessor, JinaRankingIOProcessorM final_res_batch: list[PoolingRequestOutput] = [] for i in range(len(outputs)): + if outputs[i].error is not None: + final_res_batch.append(outputs[i]) + continue + embeds = outputs[i].outputs.data.float() # The JinaForRanking model concatenates docs first, then query. diff --git a/vllm/entrypoints/pooling/scoring/serving.py b/vllm/entrypoints/pooling/scoring/serving.py index 674522b7eb..b84d3310d6 100644 --- a/vllm/entrypoints/pooling/scoring/serving.py +++ b/vllm/entrypoints/pooling/scoring/serving.py @@ -195,13 +195,36 @@ class ServingScores(PoolingServing): ctx = await self._init_ctx(self.io_processor, *args, **kwargs) await self._preprocessing(self.io_processor, ctx) - # stage 1: encode queries and cache token embeddings on workers. - await self._flash_late_interaction_encode_queries(ctx) - # stage 2: encode docs and return scalar scores from workers. - await self._flash_late_interaction_encode_docs(ctx) + try: + # stage 1: encode queries and cache token embeddings on workers. + await self._flash_late_interaction_encode_queries(ctx) + # stage 2: encode docs and return scalar scores from workers. + await self._flash_late_interaction_encode_docs(ctx) + except BaseException: + try: + await self._cleanup_flash_late_interaction(ctx) + except Exception: + logger.exception("Failed to clean up late-interaction query cache.") + raise return await self._postprocessing_async(self.io_processor, ctx) + async def _cleanup_flash_late_interaction(self, ctx: ScoringServeContext) -> None: + query_keys = ctx.late_interaction_query_keys + if not query_keys: + return + + doc_keys = ctx.late_interaction_doc_keys or [] + + try: + # Stop documents before removing the query tensors they reference. + await self.engine_client.abort([*query_keys, *doc_keys]) + finally: + await self.engine_client.collective_rpc( + "release_late_interaction_query_cache", + args=(query_keys,), + ) + async def _flash_late_interaction_encode_queries(self, ctx: ScoringServeContext): assert ctx.n_queries is not None assert ctx.engine_inputs is not None @@ -259,7 +282,9 @@ class ServingScores(PoolingServing): query_keys = ctx.late_interaction_query_keys if query_keys is None: raise RuntimeError("Late-interaction query keys were not initialized.") - doc_keys = [f"{ctx.request_id}-doc-{i}" for i in range(n_docs)] + doc_namespace = random_uuid() + doc_keys = [f"late-interaction-{doc_namespace}-doc-{i}" for i in range(n_docs)] + ctx.late_interaction_doc_keys = doc_keys for i in range(n_docs): query_idx = 0 if n_queries == 1 else i diff --git a/vllm/entrypoints/pooling/typing.py b/vllm/entrypoints/pooling/typing.py index 7ce26dd6cf..8e45ce33f4 100644 --- a/vllm/entrypoints/pooling/typing.py +++ b/vllm/entrypoints/pooling/typing.py @@ -115,6 +115,7 @@ class PoolingServeContext(Generic[PoolingRequestT]): ## for flash-late-interaction query_final_res_batch: list[PoolingRequestOutput] | None = None late_interaction_query_keys: list[str] | None = None + late_interaction_doc_keys: list[str] | None = None @dataclass diff --git a/vllm/entrypoints/serve/exception_handling/error_response.py b/vllm/entrypoints/serve/exception_handling/error_response.py index 505c5ac35c..46e2504663 100644 --- a/vllm/entrypoints/serve/exception_handling/error_response.py +++ b/vllm/entrypoints/serve/exception_handling/error_response.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from http import HTTPStatus -from vllm.exceptions import GenerationError +from vllm.exceptions import GenerationError, RetryableRequestError from vllm.logger import init_logger from ..engine.protocol import ErrorInfo, ErrorResponse @@ -56,7 +56,11 @@ def create_error_response( status_code = HTTPStatus.BAD_REQUEST param = None elif isinstance(exc, GenerationError): - err_type = "InternalServerError" + err_type = ( + "ServiceUnavailableError" + if isinstance(exc, RetryableRequestError) + else "InternalServerError" + ) status_code = exc.status_code param = None elif isinstance(exc, VLLMServerError): diff --git a/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py b/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py index c023413937..7b6954b8f7 100644 --- a/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py +++ b/vllm/entrypoints/serve/exception_handling/handlers/vllm_error.py @@ -70,8 +70,8 @@ async def generation_error_handler(req: Request, exc: GenerationError): """Handle GenerationError without logging stack traces. GenerationError is a known, expected error (e.g. KV cache load failure) - that should be returned to the client as a 500 response without polluting - server logs with stack traces. + that should be returned to the client without polluting server logs with + stack traces. """ err = create_error_response(exc) return JSONResponse(err.model_dump(), status_code=err.error.code) diff --git a/vllm/exceptions.py b/vllm/exceptions.py index c7900a0347..9bfd4634bc 100644 --- a/vllm/exceptions.py +++ b/vllm/exceptions.py @@ -130,6 +130,14 @@ class GenerationError(VLLMServerError): self.status_code = HTTPStatus.INTERNAL_SERVER_ERROR +class RetryableRequestError(GenerationError): + """A transient request-level failure that is safe for clients to retry.""" + + def __init__(self, message: str = "Request failed; please retry"): + super().__init__(message) + self.status_code = HTTPStatus.SERVICE_UNAVAILABLE + + class GracefulHTTPError(VLLMError): """Exception that should be translated into an HTTP error response. diff --git a/vllm/outputs.py b/vllm/outputs.py index 8bbbbdf40c..62247d93b2 100644 --- a/vllm/outputs.py +++ b/vllm/outputs.py @@ -18,6 +18,15 @@ from vllm.v1.metrics.stats import RequestSpecDecodeMetrics, RequestStateStats logger = init_logger(__name__) +@dataclass(frozen=True) +class RequestError: + """A request-level error that does not mark the engine as unhealthy.""" + + code: str + message: str + retryable: bool = False + + @dataclass class SamplingMask: """Per-token sampling support sets aligned with completion token IDs. @@ -259,6 +268,7 @@ class PoolingRequestOutput(Generic[_O]): prompt_token_ids (list[int]): A list of token IDs used in the prompt. num_cached_tokens: The number of tokens with prefix cache hit. finished (bool): A flag indicating whether the pooling is completed. + error: Structured request-level error information, if the request failed. """ @@ -269,20 +279,25 @@ class PoolingRequestOutput(Generic[_O]): prompt_token_ids: list[int], num_cached_tokens: int, finished: bool, + *, + error: RequestError | None = None, ): self.request_id = request_id self.prompt_token_ids = prompt_token_ids self.num_cached_tokens = num_cached_tokens self.finished = finished self.outputs = outputs + self.error = error def __repr__(self) -> str: + error = f", error={self.error!r}" if self.error is not None else "" return ( f"{type(self).__name__}(request_id={self.request_id!r}, " f"outputs={self.outputs!r}, " f"prompt_token_ids={self.prompt_token_ids}, " f"num_cached_tokens={self.num_cached_tokens}, " - f"finished={self.finished})" + f"finished={self.finished}" + f"{error})" ) @@ -319,12 +334,18 @@ class EmbeddingRequestOutput(PoolingRequestOutput[EmbeddingOutput]): def from_base( request_output: PoolingRequestOutput, ) -> "EmbeddingRequestOutput": + outputs = ( + EmbeddingOutput([]) + if request_output.error is not None + else EmbeddingOutput.from_base(request_output.outputs) + ) return EmbeddingRequestOutput( request_id=request_output.request_id, - outputs=EmbeddingOutput.from_base(request_output.outputs), + outputs=outputs, prompt_token_ids=request_output.prompt_token_ids, num_cached_tokens=request_output.num_cached_tokens, finished=request_output.finished, + error=request_output.error, ) @@ -362,12 +383,18 @@ class ClassificationRequestOutput(PoolingRequestOutput[ClassificationOutput]): def from_base( request_output: PoolingRequestOutput, ) -> "ClassificationRequestOutput": + outputs = ( + ClassificationOutput([]) + if request_output.error is not None + else ClassificationOutput.from_base(request_output.outputs) + ) return ClassificationRequestOutput( request_id=request_output.request_id, - outputs=ClassificationOutput.from_base(request_output.outputs), + outputs=outputs, prompt_token_ids=request_output.prompt_token_ids, num_cached_tokens=request_output.num_cached_tokens, finished=request_output.finished, + error=request_output.error, ) @@ -402,10 +429,16 @@ class ScoringRequestOutput(PoolingRequestOutput[ScoringOutput]): def from_base( request_output: PoolingRequestOutput, ) -> "ScoringRequestOutput": + outputs = ( + ScoringOutput(0.0) + if request_output.error is not None + else ScoringOutput.from_base(request_output.outputs) + ) return ScoringRequestOutput( request_id=request_output.request_id, - outputs=ScoringOutput.from_base(request_output.outputs), + outputs=outputs, prompt_token_ids=request_output.prompt_token_ids, num_cached_tokens=request_output.num_cached_tokens, finished=request_output.finished, + error=request_output.error, ) diff --git a/vllm/v1/engine/llm_engine.py b/vllm/v1/engine/llm_engine.py index d14b3b9755..c3fb3efe7e 100644 --- a/vllm/v1/engine/llm_engine.py +++ b/vllm/v1/engine/llm_engine.py @@ -323,6 +323,14 @@ class LLMEngine: engine_core_timestamp=outputs.timestamp, iteration_stats=iteration_stats, ) + + mm_processor_cache = self.renderer.mm_processor_cache + if mm_processor_cache is not None: + for engine_core_output in outputs.outputs: + if engine_core_output.mm_cache_miss_hashes: + for mm_hash in engine_core_output.mm_cache_miss_hashes: + mm_processor_cache.invalidate(mm_hash) + self.output_processor.update_scheduler_stats(outputs.scheduler_stats) # 3) Abort any reqs that finished due to stop strings. diff --git a/vllm/v1/engine/output_processor.py b/vllm/v1/engine/output_processor.py index 89b3a9cf09..c22c5474c7 100644 --- a/vllm/v1/engine/output_processor.py +++ b/vllm/v1/engine/output_processor.py @@ -16,6 +16,7 @@ from vllm.outputs import ( CompletionOutput, PoolingOutput, PoolingRequestOutput, + RequestError, RequestOutput, SamplingMask, ) @@ -304,6 +305,7 @@ class RequestState: stop_reason: int | str | None, kv_transfer_params: dict[str, Any] | None = None, ec_transfer_params: dict[str, Any] | None = None, + error: RequestError | None = None, ) -> RequestOutput | PoolingRequestOutput | None: finished = finish_reason is not None final_only = self.output_kind == RequestOutputKind.FINAL_ONLY @@ -342,6 +344,7 @@ class RequestState: external_req_id, [self._new_pooling_output(pooling_output)], finished, + error=error, ) output = self._new_completion_output(new_token_ids, finish_reason, stop_reason) @@ -369,6 +372,7 @@ class RequestState: finished: bool, kv_transfer_params: dict[str, Any] | None = None, ec_transfer_params: dict[str, Any] | None = None, + error: RequestError | None = None, ) -> RequestOutput | PoolingRequestOutput: # If prompt embeds were used, put placeholder prompt token ids prompt_token_ids = self.prompt_token_ids @@ -385,6 +389,7 @@ class RequestState: num_cached_tokens=self.num_cached_tokens, prompt_token_ids=prompt_token_ids, finished=finished, + error=error, ) assert self.logprobs_processor is not None if self.output_kind == RequestOutputKind.DELTA: @@ -715,6 +720,19 @@ class OutputProcessor: if engine_core_output.spec_decode_metrics is not None: req_state.spec_decode_metrics = engine_core_output.spec_decode_metrics + request_error = None + if ( + req_state.detokenizer is None + and engine_core_output.mm_cache_miss_hashes + and finish_reason == FinishReason.ERROR + ): + request_error = RequestError( + code="multimodal_cache_miss", + message="Multi-modal processor cache miss.", + retryable=True, + ) + pooling_output = EMPTY_CPU_TENSOR + if pooling_output is None: assert req_state.detokenizer is not None assert req_state.logprobs_processor is not None @@ -742,6 +760,7 @@ class OutputProcessor: stop_reason, kv_transfer_params, ec_transfer_params, + error=request_error, ): if req_state.streaming_input: request_output.finished = False diff --git a/vllm/v1/pool/late_interaction_runner.py b/vllm/v1/pool/late_interaction_runner.py index 7ea39f5754..6e8fc31861 100644 --- a/vllm/v1/pool/late_interaction_runner.py +++ b/vllm/v1/pool/late_interaction_runner.py @@ -42,6 +42,18 @@ class LateInteractionRunner: self._query_uses.clear() self._doc_query_keys.clear() + def release_queries(self, query_keys: Iterable[str]) -> None: + query_keys = set(query_keys) + for query_key in query_keys: + self._query_cache.pop(query_key, None) + self._query_uses.pop(query_key, None) + + self._doc_query_keys = { + req_id: query_key + for req_id, query_key in self._doc_query_keys.items() + if query_key not in query_keys + } + def register_request( self, req_id: str, pooling_params: PoolingParams | None ) -> None: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b40fb4eac1..1b9d0608e2 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1012,6 +1012,10 @@ class GPUModelRunner(LoRAModelRunnerMixin): if self.pooling_runner is not None: self.pooling_runner.clear() + def release_late_interaction_query_cache(self, query_keys: list[str]) -> None: + if self.pooling_runner is not None: + self.pooling_runner.release_late_interaction_query_cache(query_keys) + @torch.inference_mode() def profile_cudagraph_memory(self) -> int: """Estimate the GPU memory required to capture CUDA graphs.""" diff --git a/vllm/v1/worker/gpu/pool/pooling_runner.py b/vllm/v1/worker/gpu/pool/pooling_runner.py index 1a81046b17..be27019dd0 100644 --- a/vllm/v1/worker/gpu/pool/pooling_runner.py +++ b/vllm/v1/worker/gpu/pool/pooling_runner.py @@ -98,6 +98,9 @@ class PoolingRunner: def clear(self) -> None: self.late_interaction_runner.clear() + def release_late_interaction_query_cache(self, query_keys: list[str]) -> None: + self.late_interaction_runner.release_queries(query_keys) + def _get_pooling_metadata( self, input_batch: InputBatch, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 9ba7d89140..047fdc5ff0 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -993,6 +993,9 @@ class GPUModelRunner( self.encoder_cache.clear() self.late_interaction_runner.clear() + def release_late_interaction_query_cache(self, query_keys: list[str]) -> None: + self.late_interaction_runner.release_queries(query_keys) + def _get_positions(self, num_tokens: Any): if isinstance(num_tokens, int): if self.uses_mrope: diff --git a/vllm/v1/worker/worker_base.py b/vllm/v1/worker/worker_base.py index 017b1230cd..95a0c3475d 100644 --- a/vllm/v1/worker/worker_base.py +++ b/vllm/v1/worker/worker_base.py @@ -143,6 +143,13 @@ class WorkerBase: if callable(reset_fn): reset_fn() + def release_late_interaction_query_cache(self, query_keys: list[str]) -> None: + release_fn = getattr( + self.model_runner, "release_late_interaction_query_cache", None + ) + if callable(release_fn): + release_fn(query_keys) + def get_model(self) -> nn.Module: raise NotImplementedError