[Bugfix]Fix distributed stage0 multimodal cache routing (#3740)

Signed-off-by: bjf-frz <240110042+bjf-frz@users.noreply.github.com>
Co-authored-by: bjf-frz <240110042+bjf-frz@users.noreply.github.com>

tests/engine/test_async_omni_engine_input.py

@@ -3,6 +3,7 @@ from pytest_mock import MockerFixture
 from vllm.sampling_params import SamplingParams
 from vllm.v1.engine import EngineCoreRequest
 
+from vllm_omni.distributed.omni_coordinator import ReplicaInfo, ReplicaStatus
 from vllm_omni.engine import OmniEngineCoreRequest
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo
 from vllm_omni.engine.stage_pool import StagePool
@@ -95,6 +96,44 @@ class _FakeStageClient:
     stage_type = "llm"
     final_output = False
 
+    def __init__(self, input_address: str | None = None):
+        if input_address is not None:
+            self.client_addresses = {"input_address": input_address}
+
+
+class _FakeHub:
+    def __init__(self, replicas: list[ReplicaInfo]):
+        self._replicas = replicas
+
+    def get_replicas_for_stage(self, stage_id: int):
+        return type(
+            "ReplicaList",
+            (),
+            {"replicas": [rep for rep in self._replicas if rep.stage_id == stage_id]},
+        )()
+
+
+class _RoundRobinLB:
+    def __init__(self):
+        self._next = 0
+
+    def select(self, task, replicas):  # noqa: ARG002
+        idx = self._next % len(replicas)
+        self._next += 1
+        return idx
+
+
+def _replica(input_addr: str) -> ReplicaInfo:
+    return ReplicaInfo(
+        input_addr=input_addr,
+        output_addr=input_addr.replace("input", "output"),
+        stage_id=0,
+        status=ReplicaStatus.UP,
+        queue_length=0,
+        last_heartbeat=0.0,
+        registered_at=0.0,
+    )
+
 
 def test_build_add_request_message_scopes_mm_uuids_to_selected_stage0_replica(mocker: MockerFixture):
     engine = object.__new__(AsyncOmniEngine)
@@ -132,6 +171,93 @@ def test_build_add_request_message_scopes_mm_uuids_to_selected_stage0_replica(mo
     assert seen_uuids[0].removeprefix("stage0:rep0:") == seen_uuids[1].removeprefix("stage0:rep1:")
 
 
+@pytest.mark.asyncio
+async def test_build_add_request_message_scopes_mm_uuids_to_distributed_stage0_replica(mocker: MockerFixture):
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.model = "test-model"
+    engine.default_sampling_params_list = [params]
+    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
+    engine.supported_tasks = ("generate",)
+
+    addr0 = "tcp://host-a:1000/input"
+    addr1 = "tcp://host-b:1000/input"
+    stage_pool = StagePool(0, [_FakeStageClient(addr0), _FakeStageClient(addr1)])
+    stage_pool.attach_hub(_FakeHub([_replica(addr0), _replica(addr1)]))
+    stage_pool.attach_load_balancer(_RoundRobinLB())
+    engine.stage_pools = [stage_pool]
+
+    seen_uuids: list[str] = []
+
+    def process_inputs(**kwargs):
+        prompt = kwargs["prompt"]
+        seen_uuids.append(prompt["multi_modal_uuids"]["image"][0])
+        return _make_engine_core_request(kwargs["request_id"])
+
+    input_processor = mocker.Mock()
+    input_processor.process_inputs.side_effect = process_inputs
+    engine.input_processor = input_processor
+
+    for request_id in ("req-1", "req-2"):
+        engine._build_add_request_message(
+            request_id=request_id,
+            prompt={
+                "prompt": "describe",
+                "multi_modal_data": {"image": "same-image"},
+            },
+            sampling_params_list=[params],
+            final_stage_id=0,
+        )
+
+    assert seen_uuids[0].startswith("stage0:rep0:")
+    assert seen_uuids[1].startswith("stage0:rep1:")
+    assert stage_pool.get_bound_replica_id("req-1") == 0
+    assert stage_pool.get_bound_replica_id("req-2") == 1
+    assert await stage_pool.pick("req-1") == 0
+    assert await stage_pool.pick("req-2") == 1
+
+
+def test_build_add_request_message_skips_distributed_mm_scope_when_no_replica(mocker: MockerFixture):
+    engine = object.__new__(AsyncOmniEngine)
+    params = SamplingParams(max_tokens=8)
+    engine.model = "test-model"
+    engine.default_sampling_params_list = [params]
+    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
+    engine.supported_tasks = ("generate",)
+
+    addr0 = "tcp://host-a:1000/input"
+    addr1 = "tcp://host-b:1000/input"
+    stage_pool = StagePool(0, [_FakeStageClient(addr0), _FakeStageClient(addr1)])
+    stage_pool.attach_hub(_FakeHub([]))
+    stage_pool.attach_load_balancer(_RoundRobinLB())
+    engine.stage_pools = [stage_pool]
+
+    seen_prompt: dict | None = None
+
+    def process_inputs(**kwargs):
+        nonlocal seen_prompt
+        seen_prompt = kwargs["prompt"]
+        return _make_engine_core_request(kwargs["request_id"])
+
+    input_processor = mocker.Mock()
+    input_processor.process_inputs.side_effect = process_inputs
+    engine.input_processor = input_processor
+
+    engine._build_add_request_message(
+        request_id="req-no-replica",
+        prompt={
+            "prompt": "describe",
+            "multi_modal_data": {"image": "same-image"},
+        },
+        sampling_params_list=[params],
+        final_stage_id=0,
+    )
+
+    assert seen_prompt is not None
+    assert "multi_modal_uuids" not in seen_prompt
+    assert stage_pool.get_bound_replica_id("req-no-replica") is None
+
+
 def test_stage_pool_replica_count_falls_back_to_clients():
     class PoolWithoutLiveNumReplicas:
         clients = [object(), None, object()]

vllm_omni/engine/async_omni_engine.py

@@ -1557,21 +1557,30 @@ class AsyncOmniEngine:
         if stage0_pool.stage_type == "diffusion" or self._stage_pool_replica_count(stage0_pool) <= 1:
             return None
 
-        # This synchronous request path can safely pre-bind local in-process
-        # replicas. Distributed head mode still uses the async picker inside
-        # StagePool.submit_initial().
-        if self._stage_pool_is_distributed(stage0_pool):
-            logger.debug(
-                "[AsyncOmniEngine] Skipping stage-0 multimodal cache scoping for distributed routing req=%s",
-                request_id,
-            )
-            return None
-
         prompts = prompt if isinstance(prompt, list) else [prompt]
         if not any(isinstance(p, dict) and p.get("multi_modal_data") for p in prompts):
             return None
 
-        replica_id = stage0_pool.select_replica_id(request_id)
+        if self._stage_pool_is_distributed(stage0_pool):
+            preselect_replica_id = getattr(stage0_pool, "preselect_replica_id", None)
+            if not callable(preselect_replica_id):
+                logger.debug(
+                    "[AsyncOmniEngine] Skipping stage-0 multimodal cache scoping for distributed routing "
+                    "without preselect support req=%s",
+                    request_id,
+                )
+                return None
+            replica_id = preselect_replica_id(request_id)
+            if replica_id is None:
+                logger.debug(
+                    "[AsyncOmniEngine] Skipping stage-0 multimodal cache scoping for distributed routing "
+                    "because no serviceable replica is available yet req=%s",
+                    request_id,
+                )
+                return None
+        else:
+            replica_id = stage0_pool.select_replica_id(request_id)
+
         for p in prompts:
             self._ensure_stage_replica_mm_uuids(
                 p,

vllm_omni/engine/stage_pool.py

@@ -301,6 +301,51 @@ class StagePool:
                 raise RuntimeError(f"no UP replica for stage {self.stage_id} after {self.DISPATCH_WAIT_TIMEOUT_S:.1f}s")
             await asyncio.sleep(min(self.DISPATCH_RETRY_INTERVAL_S, deadline - now))
 
+    def preselect_replica_id(
+        self,
+        request_id: str,
+        task: Task | None = None,
+        *,
+        affinity_request_id: str | None = None,
+    ) -> int | None:
+        """Synchronously pick and bind a replica before request preprocessing.
+
+        The main-thread input preprocessing path cannot await :meth:`pick`, but
+        multimodal cache UUID scoping needs to know the same replica that
+        :meth:`submit_initial` will later use. In distributed mode this checks
+        the hub's cached replica snapshot once and records the selected input
+        address in ``_affinity`` so the async submit path reuses the route. If
+        no replica is currently serviceable, return ``None`` and let the async
+        submit-time router wait without blocking the caller.
+        """
+        if self._hub is None or self._lb is None:
+            return self.select_replica_id(request_id, affinity_request_id=affinity_request_id)
+
+        bound_addr = self._affinity.get(request_id)
+        if bound_addr is not None:
+            replica_id = self._serviceable_replica_id_for_addr(bound_addr)
+            if replica_id is not None:
+                return replica_id
+            self._affinity.pop(request_id, None)
+
+        if affinity_request_id is not None:
+            parent_addr = self._affinity.get(affinity_request_id)
+            if parent_addr is not None:
+                replica_id = self._serviceable_replica_id_for_addr(parent_addr)
+                if replica_id is not None:
+                    self._affinity[request_id] = parent_addr
+                    return replica_id
+
+        task = task or Task(request_id=request_id)
+        candidates = self._collect_serviceable_replicas()
+        if not candidates:
+            return None
+
+        lb_idx = self._lb.select(task, [rep for rep, _ in candidates])
+        replica_info, replica_id = candidates[lb_idx]
+        self._affinity[request_id] = replica_info.input_addr
+        return replica_id
+
     def _collect_serviceable_replicas(self) -> list[tuple[ReplicaInfo, int]]:
         """Return list of ``(ReplicaInfo, replica_id)`` for UP, attached replicas."""
         if self._hub is None: