[Bugfix] Fix multimodal sender-cache invalidation race with queued renders

tests/multimodal/test_cache.py

@@ -1,14 +1,19 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import asyncio
 import multiprocessing as mp
+import threading
+from contextlib import suppress
 from types import SimpleNamespace
 from typing import cast
 from unittest.mock import Mock
 
 import numpy as np
 import pytest
 import torch
+from PIL import Image
 
+from vllm import PoolingParams
 from vllm.config import ModelConfig, ParallelConfig, VllmConfig
 from vllm.config.multimodal import MultiModalConfig
 from vllm.multimodal.cache import (
@@ -39,11 +44,20 @@
     PlaceholderRange,
 )
 from vllm.multimodal.processing import PromptInsertion
+from vllm.outputs import PoolingRequestOutput
 from vllm.renderers import renderer_from_config
+from vllm.sampling_params import RequestOutputKind
 from vllm.utils.async_utils import make_async
 from vllm.utils.mem_constants import GiB_bytes, MiB_bytes
+from vllm.v1.engine import (
+    EngineCoreOutput,
+    EngineCoreOutputs,
+    EngineCoreRequest,
+    FinishReason,
+)
 from vllm.v1.engine.async_llm import AsyncLLM
 from vllm.v1.engine.llm_engine import LLMEngine
+from vllm.v1.engine.output_processor import OutputProcessor, RequestOutputCollector
 from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
 
 from ..models.utils import build_model_context
@@ -1099,6 +1113,123 @@ async def call_release():
     assert receiver.get_and_update_item(payload, mm_hash) is item
 
 
+@pytest.mark.asyncio
+async def test_cache_drift_invalidation_waits_for_render_without_delaying_outputs(
+    monkeypatch,
+):
+    """Apply drift invalidation between renders, including already queued work."""
+    model_config = ModelConfig(
+        model="hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration",
+        max_model_len=2048,
+        mm_processor_cache_gb=1,
+    )
+    renderer = renderer_from_config(VllmConfig(model_config=model_config))
+    prompt = {
+        "prompt": "<|vision_start|><|image_pad|><|vision_end|>Describe.",
+        "multi_modal_data": {"image": Image.new("RGB", (56, 56))},
+    }
+    primed = (await renderer.render_cmpl_async([prompt]))[0]
+    mm_hash = primed["mm_hashes"]["image"][0]
+    processor = renderer.get_mm_processor()
+    get_missing = processor._get_cache_missing_items
+    entered = asyncio.Event()
+    release = threading.Event()
+    loop = asyncio.get_running_loop()
+
+    def hold_cache_hit(*args, **kwargs):
+        result = get_missing(*args, **kwargs)
+        assert result[0] == {"image": [True]}
+        loop.call_soon_threadsafe(entered.set)
+        assert release.wait(timeout=10)
+        return result
+
+    monkeypatch.setattr(processor, "_get_cache_missing_items", hold_cache_hit)
+    output_processor = OutputProcessor(tokenizer=None, log_stats=False)
+    collectors = []
+    for idx in range(2):
+        request = EngineCoreRequest(
+            request_id=str(idx),
+            external_req_id=str(idx),
+            prompt_token_ids=[1],
+            mm_features=None,
+            sampling_params=None,
+            pooling_params=PoolingParams(task="embed"),
+            arrival_time=0,
+            lora_request=None,
+            cache_salt=None,
+            data_parallel_rank=None,
+        )
+        collector = RequestOutputCollector(RequestOutputKind.FINAL_ONLY, str(idx))
+        output_processor.add_request(request, prompt=None, queue=collector)
+        collectors.append(collector)
+    outputs: asyncio.Queue[EngineCoreOutputs] = asyncio.Queue()
+    engine = SimpleNamespace(
+        output_handler=None,
+        engine_core=SimpleNamespace(get_output_async=outputs.get),
+        output_processor=output_processor,
+        log_stats=False,
+        logger_manager=None,
+        renderer=renderer,
+    )
+    AsyncLLM._run_output_handler(engine)
+    render = asyncio.create_task(renderer.render_cmpl_async([prompt]))
+    queued_render = None
+    try:
+        await asyncio.wait_for(entered.wait(), timeout=10)
+        # This render is queued before the output handler learns about the miss.
+        queued_render = renderer._process_multimodal_async(
+            renderer.get_tokenizer().encode(prompt["prompt"]),
+            prompt["multi_modal_data"],
+            mm_uuids=None,
+            mm_processor_kwargs=None,
+        )
+        outputs.put_nowait(
+            EngineCoreOutputs(
+                outputs=[
+                    EngineCoreOutput(
+                        request_id="0",
+                        new_token_ids=[],
+                        finish_reason=FinishReason.ERROR,
+                        mm_cache_miss_hashes=[mm_hash],
+                    )
+                ]
+            )
+        )
+        failed = await asyncio.wait_for(collectors[0].get(), timeout=5)
+        assert isinstance(failed, PoolingRequestOutput)
+        assert failed.error is not None
+        assert failed.error.code == "multimodal_cache_miss"
+        outputs.put_nowait(
+            EngineCoreOutputs(
+                outputs=[
+                    EngineCoreOutput(
+                        request_id="1",
+                        new_token_ids=[],
+                        finish_reason=FinishReason.STOP,
+                        pooling_output=torch.tensor([1.0]),
+                    )
+                ]
+            )
+        )
+        succeeded = await asyncio.wait_for(collectors[1].get(), timeout=5)
+        assert isinstance(succeeded, PoolingRequestOutput)
+        assert succeeded.error is None
+        monkeypatch.setattr(processor, "_get_cache_missing_items", get_missing)
+        release.set()
+        assert (await render)[0]["mm_kwargs"]["image"][0] is None
+        resent = await queued_render
+        assert resent["mm_kwargs"]["image"][0] is not None
+    finally:
+        release.set()
+        await asyncio.gather(render, return_exceptions=True)
+        if queued_render is not None:
+            await asyncio.gather(queued_render, return_exceptions=True)
+        engine.output_handler.cancel()
+        with suppress(asyncio.CancelledError):
+            await engine.output_handler
+        renderer.shutdown()
+
+
 _SLEEP_VISION_PROMPT = (
     "<|im_start|>system\nYou are a helpful assistant.<|im_end|>"
     "\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"

vllm/renderers/base.py

@@ -9,6 +9,7 @@
 from contextlib import ExitStack
 from dataclasses import replace
 from functools import cached_property
+from queue import Empty, SimpleQueue
 from typing import TYPE_CHECKING, Any, Generic, overload
 
 from typing_extensions import TypeVar
@@ -170,6 +171,7 @@ def __init__(self, config: "VllmConfig", tokenizer: _T | None) -> None:
         )
 
         self._mm_processor_cache: BaseMultiModalProcessorCache | None = None
+        self._pending_mm_cache_invalidations: SimpleQueue[str] = SimpleQueue()
 
         # A second processor with its own processor-only cache.
         # Used by the tokenize endpoint so that tokenize-only
@@ -391,6 +393,26 @@ async def clear_mm_cache_async(self) -> None:
         races with concurrent process_inputs on the mm_processor_cache."""
         await self._clear_mm_cache_async()
 
+    def queue_mm_cache_invalidations(self, mm_hashes: Sequence[str]) -> None:
+        """Mark stale sender entries without interrupting ongoing processing."""
+        if self.mm_processor_cache is not None:
+            for mm_hash in mm_hashes:
+                self._pending_mm_cache_invalidations.put(mm_hash)
+
+    def _drain_mm_cache_invalidations(self) -> None:
+        cache = self.mm_processor_cache
+        while True:
+            try:
+                mm_hash = self._pending_mm_cache_invalidations.get_nowait()
+            except Empty:
+                return
+            if cache is not None:
+                try:
+                    cache.invalidate(mm_hash)
+                except Exception:
+                    self._pending_mm_cache_invalidations.put(mm_hash)
+                    raise
+
     def shutdown(self) -> None:
         # Wait for any background MM warmup to finish first: it runs on the
         # _mm_executor and touches the mm_processor_cache, so it must be
@@ -901,6 +923,7 @@ def _process_multimodal(
         *,
         skip_mm_cache: bool = False,
     ) -> "MultiModalInput":
+        self._drain_mm_cache_invalidations()
         mm_processor = self.get_mm_processor()
 
         mm_req_id = f"renderer{self.api_process_rank}-mm-{self._mm_req_counter.inc(1)}"

vllm/v1/engine/async_llm.py

@@ -879,6 +879,17 @@ async def output_handler():
                     for start in range(0, num_outputs, chunk_size):
                         end = start + chunk_size
                         outputs_slice = engine_core_outputs[start:end]
+                        if mm_processor_cache is not None:
+                            missing_hashes = [
+                                mm_hash
+                                for eco in outputs_slice
+                                for mm_hash in eco.mm_cache_miss_hashes or ()
+                            ]
+                            if missing_hashes:
+                                # Mark before publishing errors so even queued
+                                # renders invalidate before their next cache lookup.
+                                renderer.queue_mm_cache_invalidations(missing_hashes)
+
                         # 2) Process EngineCoreOutputs.
                         processed_outputs = output_processor.process_outputs(
                             outputs_slice,
@@ -889,16 +900,6 @@ async def output_handler():
                         # NOTE: RequestOutputs are pushed to their queues.
                         assert not processed_outputs.request_outputs
 
-                        # 2b) Recover from P0/P1 cache drift: the engine flags hashes
-                        # it couldn't find (mm_cache_miss_hashes); drop them from the
-                        # P0 shadow so the client's retry resends the data and
-                        # repopulates P1. Hot-path no-op (field is None otherwise).
-                        if mm_processor_cache is not None:
-                            for eco in outputs_slice:
-                                if eco.mm_cache_miss_hashes:
-                                    for mm_hash in eco.mm_cache_miss_hashes:
-                                        mm_processor_cache.invalidate(mm_hash)
-
                         # Allow other asyncio tasks to run between chunks
                         if end < num_outputs:
                             await asyncio.sleep(0)