diff --git a/tests/entrypoints/test_async_omni_pause_sleep_routing.py b/tests/entrypoints/test_async_omni_pause_sleep_routing.py index d3a60e5a6..887cf8885 100644 --- a/tests/entrypoints/test_async_omni_pause_sleep_routing.py +++ b/tests/entrypoints/test_async_omni_pause_sleep_routing.py @@ -33,9 +33,9 @@ def _make_omni(*, stage_types: list[str]) -> AsyncOmni: omni._final_output_handler = lambda: None omni.reset_mm_cache = AsyncMock() - stage_clients = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types] + stage_configs = [SimpleNamespace(stage_type=stage_type) for stage_type in stage_types] omni.engine = SimpleNamespace( - stage_clients=stage_clients, + stage_configs=stage_configs, stage_vllm_configs=[None] * len(stage_types), collective_rpc_async=AsyncMock(return_value=[True]), ) diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 969bbaf91..1b33850da 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -19,7 +19,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.config.omni_config import VllmOmniConfig from vllm_omni.config.pipeline_registry import OMNI_PIPELINES -from vllm_omni.config.stage_config import StageConfig +from vllm_omni.config.stage_config import StageConfig, StageType from vllm_omni.engine.async_omni_engine import StageRuntimeInfo from vllm_omni.engine.messages import ErrorMessage, OutputMessage from vllm_omni.entrypoints.async_omni import AsyncOmni @@ -107,7 +107,13 @@ class FakeAsyncOmniEngine: self.stage_metadata = stage_metadata or [THREE_STAGE_META[-1]] self.num_stages = len(self.stage_metadata) self.stage_configs = [ - StageConfig(stage_id=i, model_stage="dummy-model").to_omegaconf() for i in range(self.num_stages) + StageConfig( + stage_id=i, + model_stage="dummy-model", + stage_type=StageType(metadata.stage_type), + is_comprehension=getattr(metadata, "is_comprehension", False), + ).to_omegaconf() + for i, metadata in enumerate(self.stage_metadata) ] self.default_sampling_params_list = default_sampling_params_list or [ SamplingParams(max_tokens=8) for _ in range(self.num_stages) @@ -481,7 +487,7 @@ def test_model_config_and_vllm_config_forward_from_comprehension_stage(): io_processor = SimpleNamespace(name="io-processor") omni = object.__new__(AsyncOmni) omni.engine = SimpleNamespace( - stage_clients=[SimpleNamespace(is_comprehension=False), SimpleNamespace(is_comprehension=True)], + stage_configs=[SimpleNamespace(is_comprehension=False), SimpleNamespace(is_comprehension=True)], stage_vllm_configs=[None, vllm_config], ) omni.input_processor = input_processor @@ -502,7 +508,7 @@ def test_openai_serving_models_can_consume_async_omni_compat_attrs(): io_processor = SimpleNamespace(name="io-processor") omni = object.__new__(AsyncOmni) omni.engine = SimpleNamespace( - stage_clients=[SimpleNamespace(is_comprehension=True)], + stage_configs=[SimpleNamespace(is_comprehension=True)], stage_vllm_configs=[vllm_config], ) omni.input_processor = input_processor @@ -524,10 +530,11 @@ def test_get_diffusion_od_config_returns_diffusion_stage_config(): diffusion_od_config = object() omni = object.__new__(AsyncOmni) omni.engine = SimpleNamespace( + stage_configs=[SimpleNamespace(stage_type="llm"), SimpleNamespace(stage_type="diffusion")], stage_clients=[ SimpleNamespace(stage_type="llm"), SimpleNamespace(stage_type="diffusion", od_config=diffusion_od_config), - ] + ], ) assert omni.get_diffusion_od_config() is diffusion_od_config @@ -537,15 +544,33 @@ def test_get_diffusion_od_config_falls_back_to_inner_engine(): diffusion_od_config = object() omni = object.__new__(AsyncOmni) omni.engine = SimpleNamespace( + stage_configs=[SimpleNamespace(stage_type="llm"), SimpleNamespace(stage_type="diffusion")], stage_clients=[ SimpleNamespace(stage_type="llm"), SimpleNamespace(stage_type="diffusion", _engine=SimpleNamespace(od_config=diffusion_od_config)), - ] + ], ) assert omni.get_diffusion_od_config() is diffusion_od_config +@pytest.mark.asyncio +async def test_async_omni_rejects_diffusion_list_prompt_without_stage_clients( + monkeypatch: pytest.MonkeyPatch, +): + engine = FakeAsyncOmniEngine(stage_metadata=DIFFUSION_ONLY_META) + engine.stage_clients = [] + _patch_engine(monkeypatch, engine) + + app = AsyncOmni("dummy-model") + try: + with pytest.raises(ValueError, match="single prompt per request"): + async for _ in app.generate(prompt=["prompt-1", "prompt-2"], request_id="req-1"): + pass + finally: + app.shutdown() + + @pytest.mark.asyncio async def test_async_omni_yields_only_final_stage_outputs(monkeypatch: pytest.MonkeyPatch): engine = FakeAsyncOmniEngine( diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py index 0a89f4141..fbb30187f 100644 --- a/vllm_omni/engine/omni_engine_base.py +++ b/vllm_omni/engine/omni_engine_base.py @@ -422,7 +422,7 @@ class OmniEngineBase: for client in self.stage_clients ] supported_tasks: set[str] = set() - if any(getattr(client, "is_comprehension", False) for client in self.stage_clients): + if any(getattr(stage_config, "is_comprehension", False) for stage_config in self.stage_configs): supported_tasks.add("generate") if any(meta.final_output_type == "audio" for meta in self.stage_metadata): supported_tasks.add("speech") diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index ba4e75a11..7d1e702ef 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -186,7 +186,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): # Reject diffusion list-prompt early with a clear API error. if isinstance(prompt, list) and any( - getattr(client, "stage_type", "") == "diffusion" for client in getattr(self.engine, "stage_clients", []) + stage_config.stage_type == "diffusion" for stage_config in self.engine.stage_configs ): raise ValueError( "Diffusion stages accept only a single prompt per request. " @@ -626,7 +626,7 @@ class AsyncOmni(AsyncOmniBase, EngineClient): def _split_stage_ids_by_type(self, stage_ids: list[int] | None = None) -> tuple[list[int], list[int]]: """Split stage ids into AR/LLM (EngineCore) vs diffusion (worker RPC).""" - n_stages = len(self.engine.stage_clients) + n_stages = len(self.engine.stage_configs) if stage_ids is None: stage_ids = list(range(n_stages)) else: @@ -640,8 +640,8 @@ class AsyncOmni(AsyncOmniBase, EngineClient): ar_stage_ids: list[int] = [] diffusion_stage_ids: list[int] = [] for sid in stage_ids: - client = self.engine.stage_clients[sid] - if getattr(client, "stage_type", "llm") == "diffusion": + stage_config = self.engine.stage_configs[sid] + if stage_config.stage_type == "diffusion": diffusion_stage_ids.append(sid) else: ar_stage_ids.append(sid) diff --git a/vllm_omni/entrypoints/async_omni_base.py b/vllm_omni/entrypoints/async_omni_base.py index 6fa98f8cd..35572e303 100644 --- a/vllm_omni/entrypoints/async_omni_base.py +++ b/vllm_omni/entrypoints/async_omni_base.py @@ -155,13 +155,13 @@ class AsyncOmniBase(OmniBase): def _get_comprehension_stage_index(self) -> int | None: fallback_idx: int | None = None - for idx, stage_client in enumerate(self.engine.stage_clients): + for idx, stage_config in enumerate(self.engine.stage_configs): stage_vllm_config = self.engine.stage_vllm_configs[idx] if stage_vllm_config is None: continue if fallback_idx is None: fallback_idx = idx - if stage_client.is_comprehension: + if stage_config.is_comprehension: return idx return fallback_idx @@ -186,13 +186,11 @@ class AsyncOmniBase(OmniBase): def get_diffusion_od_config(self) -> Any | None: """Return the diffusion-stage config when the pipeline has one.""" - saw_diffusion_stage = False + saw_diffusion_stage = any(stage_config.stage_type == "diffusion" for stage_config in self.engine.stage_configs) for stage_client in self.engine.stage_clients: if getattr(stage_client, "stage_type", None) != "diffusion": continue - saw_diffusion_stage = True - od_config = getattr(stage_client, "od_config", None) if od_config is not None: return od_config