diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py index 9dd577c93..a1ec472f6 100644 --- a/tests/engine/test_async_omni_engine_input.py +++ b/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()] diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 420f808dc..ad8dbff6f 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/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, diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 68de78788..effe64705 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/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: