diff --git a/docs/configuration/README.md b/docs/configuration/README.md index b8640ecad84..fc450542289 100644 --- a/docs/configuration/README.md +++ b/docs/configuration/README.md @@ -14,6 +14,8 @@ For a specific example, see the [Qwen2.5-Omni deploy config](gh-file:vllm_omni/d For an introduction, see [Pipeline and deploy configurations](./stage_configs.md). +For developer-owned control flow, see [Static stage transitions](./stage_routing.md). + ## Memory Configuration - **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance diff --git a/docs/configuration/stage_routing.md b/docs/configuration/stage_routing.md new file mode 100644 index 00000000000..25c9c38c0fe --- /dev/null +++ b/docs/configuration/stage_routing.md @@ -0,0 +1,97 @@ +# Static stage transitions + +Model developers can set `PipelineConfig.stage_transitions` to an explicit tuple +of `(source_stage_id, target_stage_id)` pairs. This is part of the registered, +developer-owned pipeline topology. Deployment YAML continues to select hardware, +replicas and connector backends. + +With `stage_transitions=None`, requests follow the existing numeric sequence +`0 -> 1 -> ...`. Explicit transitions replace that sequence completely; omitted +edges do not gain a sequential fallback. + +For example, a pipeline with three registered stages can run `0 -> 2`: + +```python +PipelineConfig( + model_type="example", + stages=( + StagePipelineConfig(stage_id=0, model_stage="encoder"), + StagePipelineConfig(stage_id=1, model_stage="unused"), + StagePipelineConfig( + stage_id=2, + model_stage="decoder", + input_sources=(0,), + final_output=True, + final_output_type="audio", + ), + ), + stage_transitions=((0, 2),), +) +``` + +The explicit route starts at stage 0 and has at most one successor and one +predecessor per stage. Stage IDs must remain contiguous array indices, even when +the route skips stages or orders them differently. Duplicate edges, cycles, +self edges, unknown IDs and disconnected edges are rejected. The last stage of +the route must declare `final_output=True`. An empty tuple explicitly selects +stage 0 alone, so stage 0 must declare a final output in that case. + +`input_sources` still describes data dependencies. Every input source of an +active stage must precede that stage on the route. Changing a transition cannot +make a processor accept an incompatible payload or omit a required dependency. + +Configuration extensions use the actual route tail. For example, +`--forced-aligner` appends its pooling stage after that tail, uses it as the +audio input source and extends explicit transitions. A default `None` route +remains `None`. Typed and compatibility launch configurations use the same +extended pipeline. + +Requests execute a prefix of the configured route ending at `final_stage_id`. +If the model bypasses its initial stages for a request, execution starts at that +request's admitted entry stage within the prefix. +Requested output stages must occur in that prefix and declare final outputs; +they must include the endpoint. Streaming updates retain the admitted endpoint +and output stages. Public sampling arrays retain one slot per registered stage, +including inactive stages; IPC admission checks coverage through the highest +stage ID in the prefix. +Output modalities select the last matching stage in route order. For example, +under `0 -> 2 -> 1`, stage 1 can finish a request after stage 2. Stage sampling +parameters remain indexed by stage ID. +Worker connector endpoints remain fixed for the full route. A request ending +early can leave unread outgoing payloads, which follow the backend cleanup +contract below; this feature does not prune all data-plane sends per request. +An explicitly requested registered output modality is rejected if all stages +that produce it are inactive. Omitting modalities selects from active outputs. + +Async-chunk prewarming follows the same prefix and binds each upstream sender +through its StagePool before advertising its address to a receiver. The later +bridge submission reuses that binding. A lost bound replica fails its requests +instead of moving their senders after the receivers have been configured. +Connector payload send/receive endpoints use the projected transition, +including on headless workers. Backend projection respects `default_connector`; +request keys and chunk counters retain their existing per-request semantics. +Stages fed by an orchestrator bridge receive their real payload before dispatch. +Control edges and backend edges serve different purposes: a bridge-fed stage can +have an outgoing connector without an incoming one. A middle stage with both +connector directions requires the same backend on both edges, because one +worker owns one connector. Absent payload backends retain the existing local +IPC behavior; choosing a backend suitable for the deployment remains necessary. + +Omitted stages receive no request submissions or prewarming. They are still +initialized, so this feature does not reduce startup memory or process counts. +Normal and empty terminal outputs share the same completion condition: all +selected outputs have finished, and non-streaming async requests also wait for +every submitted stage. The endpoint's terminal output is retained until then, +so delayed upstream usage and stop reasons arrive before completion. +Cleanup retains the existing per-replica batched, best-effort resource release. +Abort acknowledgement does not certify physical release of transport buffers. +Changing a route does not strengthen a backend's release guarantees. Legacy V1 +full-payload and native KV transfer retain their existing cleanup paths; NIXL +ownership changes are reviewed separately. + +This first implementation supports static routes for turn-based requests. +Branches, joins, loops and token-dependent decisions need a separate execution +contract. Non-sequential duplex deployments are rejected before plugin loading; +their model-specific session contracts need a separate extension. Existing +native AR-to-DiT KV deployment restrictions still apply. +Streaming-input requests with a native MRv2 downstream receiver remain unsupported. diff --git a/docs/design/module/engine_orchestration.md b/docs/design/module/engine_orchestration.md index f9538f921a9..28f5a3ee52d 100644 --- a/docs/design/module/engine_orchestration.md +++ b/docs/design/module/engine_orchestration.md @@ -90,6 +90,11 @@ This document owns `AsyncOmniEngine`, `Orchestrator`, request-state creation, cross-stage routing, output ordering, companion tracking, control/RPC correlation, cancellation propagation, and terminal-state convergence. +After configuration resolution, the engine retains the effective `PipelineConfig`, +including injected stages such as ForcedAligner. Stage startup and orchestration +read transitions from that same configuration; the registered pipeline is not +modified. + It does not own stage placement, replica selection, stage-process startup, payload schemas, public protocol rendering, connector transport, or semantic error classification. Those responsibilities belong to the stage runtime, @@ -132,8 +137,14 @@ child must drop `parent_requests` once no children remain. EngineCore control RPCs (`sleep` / `wake_up` / `pause_scheduler` / `resume_scheduler`) must propagate worker exceptions rather than returning `{"supported": False}`. When a request has multiple final output stages, only the last abort message -is request-terminal. Output-processor abort state is committed only after the -physical EngineCore abort succeeds. +is request-terminal. Output-processor abort state is committed after the stage +client's abort call succeeds or a natural terminal confirms stage completion. +For AR clients, a successful call confirms message submission, not EngineCore +completion or physical release of transport buffers. +An abort failure still attempts the other stages and replicas. The request +remains closed to admission while stages with failed abort calls drain; their natural +terminals permit cleanup without publishing a normal completion. Correlated +abort failure also reports a request-scoped error to unblock its output stream. Changes that cross into stage runtime or public error behavior require review from that contract's owners. diff --git a/tests/config/test_forced_aligner_injection.py b/tests/config/test_forced_aligner_injection.py index 2037b22f1dd..f649dcb0a41 100644 --- a/tests/config/test_forced_aligner_injection.py +++ b/tests/config/test_forced_aligner_injection.py @@ -9,8 +9,11 @@ from vllm_omni.config.pipeline_registry import OMNI_PIPELINES as _PIPELINE_REGISTRY from vllm_omni.config.stage_config import ( DeployConfig, + PipelineConfig, StageExecutionType, + StagePipelineConfig, ) +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.engine.stage_init_utils import ( build_engine_args_dict_from_omni_stage_config, extract_stage_metadata_from_omni_stage_config, @@ -80,6 +83,51 @@ def test_inject_noop_without_forced_aligner(): assert len(ext_deploy.stages) == 0 +@pytest.mark.parametrize( + ("transitions", "order"), + [ + (None, (0, 1, 2)), + (((0, 1), (1, 2)), (0, 1, 2)), + (((0, 2),), (0, 2)), + (((2, 1), (0, 2)), (0, 2, 1)), + ((), (0,)), + ], +) +def test_injected_aligner_extends_actual_route_tail(transitions, order): + from vllm_omni.config.config_factory import StageConfigFactory + from vllm_omni.engine.omni_engine_base import OmniEngineBase + + pipeline = PipelineConfig( + model_type="test", + stages=tuple(StagePipelineConfig(i, "audio", final_output=True, final_output_type="audio") for i in range(3)), + stage_transitions=transitions, + ) + deploy = DeployConfig() + overrides = {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"} + extended, extended_deploy = inject_forced_aligner_stage(pipeline, deploy, overrides) + + assert extended.stages[-1].input_sources == (order[-1],) + route = StageRouting.from_transitions(len(extended.stages), extended.stage_transitions) + assert route.stage_order == (*order, 3) + assert pipeline.stage_transitions == transitions + assert len(pipeline.stages) == 3 and deploy.stages == [] + assert (extended.stage_transitions is None) == (transitions is None) + + typed = VllmOmniConfig.from_pipeline_config(extended, user_deploy_config=extended_deploy) + legacy = StageConfigFactory._resolve_legacy_from_registry(pipeline, overrides, user_deploy_config=deploy) + assert typed.pipeline_config == legacy.pipeline_config == extended + assert tuple(typed.stage_configs[-1].input_sources) == (order[-1],) + assert tuple(legacy.stage_configs[-1].input_sources) == (order[-1],) + + engine = object.__new__(OmniEngineBase) + engine.pipeline_config = pipeline + engine._set_pipeline_runtime_config(legacy.pipeline_config, None) + assert engine.pipeline_config is legacy.pipeline_config + assert StageRouting.from_transitions( + len(legacy.stage_configs), engine.pipeline_config.stage_transitions + ).stage_order == (*order, 3) + + @pytest.mark.parametrize("async_chunk", [True, False]) @pytest.mark.parametrize("cli_async_chunk", [None, True, False]) def test_aligner_uses_completed_audio_without_disabling_upstream_chunks(async_chunk, cli_async_chunk): diff --git a/tests/config/test_stage_routing.py b/tests/config/test_stage_routing.py new file mode 100644 index 00000000000..238047cbde5 --- /dev/null +++ b/tests/config/test_stage_routing.py @@ -0,0 +1,181 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest + +from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig +from vllm_omni.config.stage_routing import StageRouting +from vllm_omni.distributed.omni_connectors.utils.config import ( + ConnectorSpec, + OmniTransferConfig, + get_stage_connector_peer, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("num_stages", [1, 2, 5]) +def test_default_route_preserves_sequential_prefixes(num_stages): + route = StageRouting.from_transitions(num_stages) + assert route.stage_order == tuple(range(num_stages)) + for final in range(num_stages): + assert route.path_to(final) == tuple(range(final + 1)) + assert route.next_stage(final, final) is None + + +@pytest.mark.parametrize( + ("transitions", "order"), + [(((0, 2),), (0, 2)), (((2, 1), (0, 2)), (0, 2, 1)), ((), (0,))], +) +def test_explicit_route_has_no_implicit_edges(transitions, order): + route = StageRouting.from_transitions(3, transitions) + assert route.stage_order == order + assert route.next_stage(order[-1], order[-1]) is None + + +@pytest.mark.parametrize( + ("transitions", "error"), + [ + (((0, 3),), "invalid stage"), + (((-1, 1),), "invalid stage"), + (((0, True),), "invalid stage"), + (((0, 0),), "Self transition"), + (((0, 1), (1, 0)), "cycle"), + (((0, 1), (0, 2)), "branching"), + (((0, 1), (0, 1)), "duplicate"), + (((0, 2), (1, 2)), "multiple incoming"), + (((1, 2),), "unreachable"), + ], +) +def test_invalid_topology_is_rejected(transitions, error): + with pytest.raises(ValueError, match=error): + StageRouting.from_transitions(3, transitions) + + +def test_unreachable_request_endpoint_is_rejected(): + route = StageRouting.from_transitions(3, ((0, 2),)) + with pytest.raises(ValueError, match="not reachable"): + route.path_to(1) + assert route.next_stage(1, 2) is None + + +@pytest.mark.parametrize("transitions", [[(0, 2)], ([0, 2],), ((0,),), ((0, 1, 2),)]) +def test_mutable_or_malformed_transition_configuration_is_rejected(transitions): + with pytest.raises(ValueError, match="immutable tuple"): + StageRouting.from_transitions(3, transitions) + + +def test_pipeline_accepts_skipped_stage_with_satisfied_dependencies(): + pipeline = PipelineConfig( + model_type="test", + stages=( + StagePipelineConfig(0, "entry"), + StagePipelineConfig(1, "inactive", input_sources=(0,)), + StagePipelineConfig(2, "output", input_sources=(0,), final_output=True), + ), + stage_transitions=((0, 2),), + ) + assert pipeline.get_validation_errors() == [] + + +@pytest.mark.parametrize("sources", [(1,), (2,)]) +def test_pipeline_rejects_unsatisfied_route_dependency(sources): + with pytest.raises(ValueError, match="must precede"): + PipelineConfig( + model_type="test", + stages=( + StagePipelineConfig(0, "entry"), + StagePipelineConfig(1, "inactive"), + StagePipelineConfig(2, "output", input_sources=sources, final_output=True), + ), + stage_transitions=((0, 2),), + ) + + +def test_pipeline_rejects_route_without_terminal_output(): + with pytest.raises(ValueError, match="last stage"): + PipelineConfig( + model_type="test", + stages=(StagePipelineConfig(0, "entry"), StagePipelineConfig(1, "output", final_output=True)), + stage_transitions=(), + ) + + +@pytest.mark.parametrize("stage_id", [3, 1.0, True]) +def test_pipeline_rejects_ids_that_cannot_index_stage_arrays(stage_id): + with pytest.raises(ValueError, match="contiguous IDs"): + PipelineConfig( + model_type="test", + stages=(StagePipelineConfig(0, "entry"), StagePipelineConfig(stage_id, "output", final_output=True)), + stage_transitions=((0, 1),), + ) + + +@pytest.mark.parametrize(("stage_id", "source", "target"), [(0, None, 2), (2, 0, 1), (1, 2, None)]) +def test_connector_projection_preserves_both_route_directions(stage_id, source, target): + from vllm_omni.engine.stage_init_utils import get_stage_connector_spec + + route = StageRouting.from_transitions(3, ((0, 2), (2, 1))) + spec = get_stage_connector_spec(None, stage_id, True, route) + config = SimpleNamespace(stage_connector_config=spec) + assert get_stage_connector_peer(config, stage_id, "from_stage") == source + assert get_stage_connector_peer(config, stage_id, "to_stage") == target + + +def test_legacy_connector_endpoints_keep_numeric_defaults(): + config = SimpleNamespace(stage_connector_config={}) + assert get_stage_connector_peer(config, 1, "from_stage") == 0 + assert get_stage_connector_peer(config, 1, "to_stage") == 2 + + +@pytest.mark.parametrize("value", ["", " ", 1.5]) +def test_legacy_unset_connector_target_keeps_sequential_default(value): + config = SimpleNamespace(stage_connector_config={"to_stage": value}) + assert get_stage_connector_peer(config, 1, "to_stage") == 2 + + +def test_connector_projection_honors_default_without_mutating_deployment(): + from vllm_omni.engine.stage_init_utils import get_stage_connector_spec + + backend = ConnectorSpec("NixlConnector", {"zmq_port": 50051}) + config = OmniTransferConfig({("0", "1"): ConnectorSpec("SharedMemoryConnector")}, backend) + original = dict(config.connectors) + route = StageRouting.from_transitions(3, ((0, 2), (2, 1))) + spec = get_stage_connector_spec(config, 2, True, route) + assert spec["name"] == "NixlConnector" + assert (spec["from_stage"], spec["to_stage"]) == (0, 1) + assert spec["extra"]["outgoing"]["from_stage"] == 2 + assert spec["extra"]["outgoing"]["to_stage"] == 1 + assert config.connectors == original and config.default_connector is backend + assert backend.extra == {"zmq_port": 50051} + + +def test_connector_projection_preserves_bridge_fed_stage(): + from vllm_omni.engine.stage_init_utils import get_stage_connector_spec + + config = OmniTransferConfig({("2", "1"): ConnectorSpec("SharedMemoryConnector")}) + route = StageRouting.from_transitions(3, ((0, 2), (2, 1))) + spec = get_stage_connector_spec(config, 2, True, route) + assert spec["extra"]["role"] == "sender" + assert (spec["from_stage"], spec["to_stage"]) == (0, 1) + + +@pytest.mark.parametrize( + "backends", [("SharedMemoryConnector", "NixlConnector"), ("NixlConnector", "SharedMemoryConnector")] +) +def test_middle_stage_rejects_incompatible_connector_backends(backends): + from vllm_omni.engine.stage_init_utils import get_stage_connector_spec + + config = OmniTransferConfig({("0", "2"): ConnectorSpec(backends[0]), ("2", "1"): ConnectorSpec(backends[1])}) + route = StageRouting.from_transitions(3, ((0, 2), (2, 1))) + with pytest.raises(ValueError, match="same connector backend"): + get_stage_connector_spec(config, 2, True, route) + + +def test_inactive_stage_has_no_payload_peers(): + from vllm_omni.engine.stage_init_utils import get_stage_connector_spec + + spec = get_stage_connector_spec(None, 1, True, StageRouting.from_transitions(3, ((0, 2),))) + assert spec == {"from_stage": None, "to_stage": None, "extra": {"role": "sender"}} diff --git a/tests/core/sched/test_omni_scheduling_coordinator.py b/tests/core/sched/test_omni_scheduling_coordinator.py index 34f208b6f2c..3df1bb021c6 100644 --- a/tests/core/sched/test_omni_scheduling_coordinator.py +++ b/tests/core/sched/test_omni_scheduling_coordinator.py @@ -116,12 +116,14 @@ def test_chunk_registration_ready_and_terminal_lifecycle(): coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True) req = _make_request("internal", status=RequestStatus.WAITING) req.external_req_id = "external" + req.payload_sender_info = {"host": "10.0.0.1", "zmq_port": 51424} waiting = MockQueue([req]) coord.process_pending_chunks(waiting, [], set(), set()) assert req.status == RequestStatus.WAITING_FOR_CHUNK [handle] = coord.pending_chunk_registrations assert isinstance(handle, OmniChunkRecvHandle) assert (handle.request_id, handle.external_req_id) == ("internal", "external") + assert handle.payload_sender_info == req.payload_sender_info coord.restore_queues(waiting, []) coord.process_pending_chunks(waiting, [], {"internal"}, set()) assert req.status == RequestStatus.WAITING diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 7ad76fd7108..1b4c7b2d627 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -3981,3 +3981,39 @@ def test_save_async_boundary_holds_generation_without_request_counter(build_adap assert len(adapter._pending_save_reqs) == queued_before + 1 assert adapter._pending_save_reqs[-1]["request"] is follow_up + + +@pytest.mark.parametrize(("source", "target"), [(0, 2), (2, 1)]) +def test_configured_edge_transfers_ordered_chunks_through_real_shm(build_adapter, source, target): + """Exercise actual SHM put/get with scheduler threads stepped by the test.""" + from vllm_omni.config.stage_routing import StageRouting + from vllm_omni.engine.stage_init_utils import get_stage_connector_spec + + route = StageRouting.from_transitions(3, ((0, 2), (2, 1))) + sender, _ = build_adapter(stage_id=source) + receiver, _ = build_adapter(stage_id=target) + sender.config.stage_connector_config = get_stage_connector_spec(None, source, True, route) + receiver.config.stage_connector_config = get_stage_connector_spec(None, target, True, route) + send_connector = SharedMemoryConnector({"stage_id": source}) + recv_connector = SharedMemoryConnector({"stage_id": target}) + sender.connector = send_connector + receiver.connector = recv_connector + external_id = f"configured-route-{uuid.uuid4().hex}" + send_request = _req("sender", RequestStatus.RUNNING, external_req_id=external_id) + recv_request = _req("receiver", RequestStatus.WAITING, external_req_id=external_id) + try: + for index in range(3): + sender.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct( + codes=CodesStruct(audio=torch.tensor([index], dtype=torch.long)) + ) + sender.save_async(None, send_request) + sender._send_single_request(sender._pending_save_reqs.popleft()) + entry = _dequeue_load_entry(receiver, recv_request) + assert receiver._poll_single_request(entry) + assert receiver.get_req_chunk["receiver"] == index + 1 + payload = receiver._pending_ar_prompt_updates["receiver"][1] + assert payload["codes"]["audio"].tolist() == [index] + assert sender.put_req_chunk[external_id] == 3 + finally: + send_connector.close() + recv_connector.close() diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index bfa7be0e42d..17c6b540737 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -19,7 +19,7 @@ import pytest from omegaconf import OmegaConf from pydantic import ValidationError -from transformers import PretrainedConfig, Qwen3OmniMoeConfig +from transformers import PretrainedConfig, Qwen3Config, Qwen3OmniMoeConfig from vllm.config import VllmConfig from vllm.engine.arg_utils import EngineArgs from vllm.exceptions import VLLMValidationError @@ -36,6 +36,19 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@pytest.fixture(autouse=True) +def local_default_model(tmp_path_factory, monkeypatch): + """Exercise real config conversion without downloading the default model.""" + from vllm.engine.arg_utils import get_model_path + + model_path = tmp_path_factory.mktemp("default-model-config") + Qwen3Config(architectures=["Qwen3ForCausalLM"]).save_pretrained(model_path) + monkeypatch.setattr( + "vllm.engine.arg_utils.get_model_path", + lambda model, revision: str(model_path) if model == EngineArgs.model else get_model_path(model, revision), + ) + + @pytest.mark.parametrize("use_v2", [False, True]) def test_stage_runner_selects_matching_input_processor_validator(monkeypatch, use_v2): from vllm.sampling_params import SamplingParams @@ -89,6 +102,15 @@ def test_default_stage_id_is_concrete_int(): assert cfg.stage_id == 0 +@pytest.mark.parametrize(("stage_id", "source", "target"), [(2, 0, 1), (1, 2, None)]) +def test_projected_connector_endpoints_reach_worker_model_config(stage_id, source, target): + spec = {"name": "SharedMemoryConnector", "extra": {}, "from_stage": source, "to_stage": target} + config = OmniEngineArgs(stage_id=stage_id, stage_connector_spec=spec).create_model_config() + assert config.stage_connector_config["from_stage"] == source + assert config.stage_connector_config["to_stage"] == target + assert spec["extra"] == {} + + def test_full_payload_capability_reaches_omni_model_config(monkeypatch): """Keep the pipeline capability intact through the engine adapter.""" captured: dict[str, object] = {} @@ -186,6 +208,7 @@ class WorkerWithoutConnector: def test_multimodal_kwarg_overrides(monkeypatch): """Ensure that overrides in the multimodal config are preserved.""" + monkeypatch.setattr("vllm.engine.arg_utils.get_model_path", lambda model, revision: model) sig = inspect.signature(OmniEngineArgs) default_mm_cache = sig.parameters["mm_processor_cache_gb"].default override_val = default_mm_cache + 1 @@ -326,6 +349,7 @@ def fake_create_model_config(self): def test_remote_tokenizer_subfolder_download_does_not_report_failure(tmp_path, monkeypatch, mocker): """A successful remote tokenizer download must not enter the error path.""" + monkeypatch.setattr("vllm.engine.arg_utils.get_model_path", lambda model, revision: model) tokenizer_dir = tmp_path / "CosyVoice-BlankEN" tokenizer_dir.mkdir() baseline_config = Mock() @@ -463,6 +487,7 @@ def test_patch_valid_hf_config_without_model_type_preserves_keys(tmp_path): def test_remote_hf_config_error_reaches_parent_loader(monkeypatch): """Remote resolution failures must stay on vLLM's normal error path.""" + monkeypatch.setattr("vllm.engine.arg_utils.get_model_path", lambda model, revision: model) loader_error = OSError("remote config resolution failed") parent_inputs = {} diff --git a/tests/engine/test_async_omni_engine_input.py b/tests/engine/test_async_omni_engine_input.py index d5ed551449d..09ca0e4f702 100644 --- a/tests/engine/test_async_omni_engine_input.py +++ b/tests/engine/test_async_omni_engine_input.py @@ -13,6 +13,7 @@ from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.async_engine_utils import scope_stage_replica_mm_uuids from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo +from vllm_omni.engine.orchestrator import OrchestratorRequestState from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.engine.stage_pool import StagePool, StageUnavailableError from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt @@ -590,13 +591,17 @@ def process_inputs(**kwargs): assert await stage_pool.pick("req-1") == 0 assert await stage_pool.pick("req-2") == 1 - # Replica 0 leaves the hub after req-1's keys were scoped to it: a submit pinned to - # those keys fails instead of sending them to replica 1, while routing moves on. + # Keep the failed route until cleanup; the receiver may know this sender. stage_pool.attach_hub(_FakeHub([_replica(addr1)])) with pytest.raises(StageUnavailableError): await stage_pool.submit_initial( - "req-1", mocker.Mock(sampling_params_list=[params]), _make_engine_core_request(), expected_replica_id=0 + "req-1", + OrchestratorRequestState(request_id="req-1", sampling_params_list=[params]), + _make_engine_core_request(), + expected_replica_id=0, ) + assert stage_pool.get_bound_replica_id("req-1") == 0 + stage_pool.release("req-1") assert stage_pool.get_bound_replica_id("req-1") is None assert await stage_pool.pick("req-1") == 1 diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 3d5ff88b7ac..cb4f42394df 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import asyncio import concurrent.futures import contextlib import importlib @@ -16,8 +17,11 @@ from omegaconf import OmegaConf from vllm.v1.engine.utils import EngineZmqAddresses +from tests.engine.test_orchestrator import FakeStageClient, _build_stage_pools from tests.helpers.mock import patch_hf_snapshot_download -from vllm_omni.config.omni_config import OmniStageRuntimeConfig +from vllm_omni.config.omni_config import BaseVllmOmniStageConfig, OmniStageRuntimeConfig +from vllm_omni.config.stage_config import DeployConfig, PipelineConfig, StagePipelineConfig +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.diffusion.data import AttentionConfig from vllm_omni.engine import omni_engine_base as async_omni_engine_module from vllm_omni.engine.async_omni_engine import AsyncOmniEngine @@ -30,7 +34,8 @@ split_devices_for_replicas, stage_runtime_env, ) -from vllm_omni.engine.stage_runtime import StageRuntime +from vllm_omni.engine.stage_runtime import DistStageRuntime, StageRuntime +from vllm_omni.utils.forced_aligner import inject_forced_aligner_stage pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -313,19 +318,36 @@ def test_stage_engine_core_client_module_reload_keeps_forward_refs_deferred(): client_mod.__dict__.update(saved) -def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch): +@pytest.mark.parametrize("transitions", [None, ((0, 2),)]) +@pytest.mark.parametrize("with_aligner", [False, True]) +@pytest.mark.parametrize("distributed", [False, True]) +def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime( + monkeypatch, transitions, with_aligner, distributed +): import vllm_omni.engine.omni_engine_base as engine_mod engine = object.__new__(AsyncOmniEngine) - engine.stage_configs = [types.SimpleNamespace()] + engine.pipeline_config = PipelineConfig( + model_type="test", + stages=tuple(StagePipelineConfig(stage_id=sid, model_stage="test", final_output=sid == 2) for sid in range(3)), + stage_transitions=transitions, + ) + if with_aligner: + pipeline, _ = inject_forced_aligner_stage( + engine.pipeline_config, DeployConfig(), {"forced_aligner": "/models/Qwen3-ForcedAligner-0.6B"} + ) + engine._set_pipeline_runtime_config(pipeline, None) + engine.stage_configs = [ + BaseVllmOmniStageConfig(stage_pipeline_config=stage) for stage in engine.pipeline_config.stages + ] engine.model = "dummy-model" engine.config_path = "dummy-config" - engine.single_stage_mode = False + engine.single_stage_mode = distributed engine.async_chunk = False engine.tokenizer = None engine._single_stage_id_filter = None - engine._omni_master_address = None - engine._omni_master_port = None + engine._omni_master_address = "127.0.0.1" if distributed else None + engine._omni_master_port = 30001 if distributed else None engine._omni_dp_size_local = 1 engine._omni_heartbeat_timeout = 30.0 engine._omni_lb_policy = "random" @@ -335,25 +357,39 @@ def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_ engine._parallel_stage_init = False captured: dict[str, object] = {} - runtime = types.SimpleNamespace(stage_pools=[], initialize=lambda: None) + create_runtime = engine_mod.create_stage_runtime def _capture_create_stage_runtime(**kwargs): captured.update(kwargs) - return runtime + return create_runtime(**kwargs) monkeypatch.setattr(engine_mod, "create_stage_runtime", _capture_create_stage_runtime) + monkeypatch.setattr(StageRuntime, "initialize", lambda self: None) engine._initialize_stages(stage_init_timeout=7) assert captured["stage_init_timeout"] == 7 assert captured["log_stats"] is True assert captured["client_config"] is engine._client_config + assert captured["stage_transitions"] == engine.pipeline_config.stage_transitions + assert isinstance(engine._runtime, DistStageRuntime if distributed else StageRuntime) + expected_order = StageRouting.from_transitions(len(engine.stage_configs), engine.pipeline_config.stage_transitions) + assert engine._runtime._stage_routing == (expected_order if captured["stage_transitions"] is not None else None) + pools = _build_stage_pools([[FakeStageClient()] for _ in engine.stage_configs]) + orchestrator = engine._create_orchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=asyncio.Queue(), + rpc_async_queue=asyncio.Queue(), + stage_pools=pools, + ) + assert orchestrator._stage_routing == expected_order def test_async_omni_engine_initialize_stages_retains_stage0_prompt_transform(monkeypatch): import vllm_omni.engine.omni_engine_base as engine_mod engine = object.__new__(AsyncOmniEngine) + engine.pipeline_config = None engine.stage_configs = [types.SimpleNamespace()] engine.model = "dummy-model" engine.config_path = "dummy-config" diff --git a/tests/engine/test_cfg_companion_lifecycle.py b/tests/engine/test_cfg_companion_lifecycle.py index f696460f0c2..810cc2c6b80 100644 --- a/tests/engine/test_cfg_companion_lifecycle.py +++ b/tests/engine/test_cfg_companion_lifecycle.py @@ -18,10 +18,12 @@ from types import SimpleNamespace import pytest +from vllm.sampling_params import SamplingParams +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.messages import AbortRequestMessage, ErrorMessage -from vllm_omni.engine.orchestrator import Orchestrator +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -103,6 +105,7 @@ def _make_orchestrator(num_stages: int = 2) -> Orchestrator: orch._cfg_tracker = CfgCompanionTracker() orch.request_states = {} orch.stage_pools = [_FakePool("llm"), _FakePool("diffusion")][:num_stages] + orch._stage_routing = StageRouting.from_transitions(len(orch.stage_pools)) orch.output_async_queue = asyncio.Queue() orch.async_chunk = False orch.duplex_control_plane = None @@ -113,15 +116,12 @@ def _make_orchestrator(num_stages: int = 2) -> Orchestrator: return orch -def _req_state(final_stage_id: int = 1) -> SimpleNamespace: - return SimpleNamespace( +def _req_state(final_stage_id: int = 1) -> OrchestratorRequestState: + return OrchestratorRequestState( + request_id="p", final_stage_id=final_stage_id, - stage_submit_ts={}, - sampling_params_list=[SimpleNamespace(output_kind=None)] * (final_stage_id + 1), - streaming=SimpleNamespace(enabled=False, segment_finished=False), - running_counter_registered=False, + sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(final_stage_id + 1)], prompt={"prompt": "x"}, - pipeline_timings={}, ) diff --git a/tests/engine/test_configured_stage_routing.py b/tests/engine/test_configured_stage_routing.py new file mode 100644 index 00000000000..f9e33c79e06 --- /dev/null +++ b/tests/engine/test_configured_stage_routing.py @@ -0,0 +1,943 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Real orchestrator and StagePool logic with lightweight engine clients (L1).""" + +import asyncio +import threading +from dataclasses import replace +from types import SimpleNamespace + +import pytest +from vllm.sampling_params import SamplingParams +from vllm.v1.engine import EngineCoreOutputs, FinishReason + +from tests.engine.test_async_omni_engine_input import _FakeHub, _replica, _RoundRobinLB +from tests.engine.test_orchestrator import ( + FakeOutputProcessor, + FakePromptRequest, + FakeRunningCounter, + FakeStageClient, + _build_request_output, + _build_stage_pools, + _engine_core_outputs, + _enqueue_add_request, + _get_output_message, + _get_rpc_message, + _processed_request, + _shutdown_orchestrator, + _wait_for, + orchestrator_factory, +) +from tests.engine.test_orchestrator_stage_input_bridge import FakeInputProcessor +from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig +from vllm_omni.config.stage_routing import StageRouting +from vllm_omni.engine import OmniEngineCoreOutput +from vllm_omni.engine.messages import AbortRequestMessage, AbortResultMessage, ErrorMessage, StageSubmissionMessage +from vllm_omni.engine.omni_engine_base import OmniEngineBase +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState, StreamingSegmentState +from vllm_omni.engine.stage_pool import RequestClosedError +from vllm_omni.entrypoints.async_omni_base import AsyncOmniBase +from vllm_omni.entrypoints.client_request_state import ClientRequestState +from vllm_omni.entrypoints.omni_base import OmniBase +from vllm_omni.entrypoints.utils import get_final_stage_id_for_e2e +from vllm_omni.metrics.stats import OrchestratorAggregator +from vllm_omni.outputs.output_processor import MultimodalOutputProcessor + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +configured_orchestrator_factory = orchestrator_factory + + +def _direct_orchestrator(transitions=((0, 2),), *, async_chunk=False): + final_stage_id = StageRouting.from_transitions(3, transitions).stage_order[-1] + clients = [FakeStageClient(final_output=(sid == final_stage_id)) for sid in range(3)] + pools = _build_stage_pools([[client] for client in clients]) + orch = Orchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=asyncio.Queue(), + rpc_async_queue=asyncio.Queue(), + stage_pools=pools, + stage_transitions=transitions, + async_chunk=async_chunk, + ) + state = OrchestratorRequestState( + request_id="r", + sampling_params_list=[SamplingParams(max_tokens=1) for _ in pools], + final_stage_id=final_stage_id, + final_output_stage_ids={final_stage_id}, + prompt={"prompt_token_ids": [1, 2]}, + ) + orch.request_states["r"] = state + return orch, state, clients + + +def _submission(**overrides): + defaults = dict( + type="add_request", + request_id="r", + prompt=FakePromptRequest(request_id="r", prompt_token_ids=[1], resumable=False), + original_prompt={}, + output_prompt_text=None, + request_timestamp=1.0, + preprocess_ms=0.0, + enqueue_ts=0.0, + sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(3)], + final_stage_id=2, + final_output_stage_ids=[2], + ) + return StageSubmissionMessage(**(defaults | overrides)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transitions", [None, ((0, 1), (1, 2)), ((0, 2), (2, 1))]) +async def test_bridge_sender_is_bound_before_receiver_prewarm(mocker, transitions): + order = StageRouting.from_transitions(3, transitions).stage_order + bridge, receiver = order[1:] + clients = [[FakeStageClient(next_inputs=[{"prompt_token_ids": [11]}])] for _ in range(3)] + clients[bridge].append(FakeStageClient(next_inputs=[{"prompt_token_ids": [11]}])) + clients[receiver][0].final_output = True + for rid, client in enumerate(clients[bridge]): + mocker.patch.object( + client, "get_payload_sender_info", create=True, return_value={"host": "127.0.0.1", "zmq_port": 50051 + rid} + ) + configs: list[object] = [ + SimpleNamespace( + model_config=SimpleNamespace( + max_model_len=64, + stage_connector_config={"extra": {"role": "receiver" if sid == receiver else "sender"}}, + ) + ) + for sid in range(3) + ] + pools = _build_stage_pools(clients, stage_vllm_configs=configs) + pools[bridge].mark_replica_unavailable(0) + orch = Orchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=asyncio.Queue(), + rpc_async_queue=asyncio.Queue(), + stage_pools=pools, + stage_transitions=transitions, + async_chunk=True, + ) + state = OrchestratorRequestState( + request_id="r", + sampling_params_list=[SamplingParams(max_tokens=1) for _ in pools], + final_stage_id=receiver, + final_output_stage_ids={receiver}, + prompt={"prompt_token_ids": [1, 2]}, + ) + orch.request_states["r"] = state + pools[0].select_replica_id("r") + assert await orch._prewarm_async_chunk_stages( + "r", FakePromptRequest(request_id="r", prompt_token_ids=[1, 2], resumable=False), state + ) + assert pools[bridge].get_bound_replica_id("r") == 1 + assert bridge not in state.stage_submit_ts + assert not any(client.add_request_calls for client in clients[bridge]) + advertised = clients[receiver][0].add_request_calls[0][0].payload_sender_info + await orch._forward_to_next_stage("r", 0, _build_request_output("r"), state) + assert len(clients[bridge][1].add_request_calls) == 1 + assert not clients[bridge][0].add_request_calls + assert advertised == pools[bridge].get_bound_client("r").get_payload_sender_info() + await orch._cleanup_request_ids(["r"], abort=True) + assert all(pool.get_bound_replica_id("r") is None for pool in pools) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transitions", [((0, 2),), ((0, 2), (2, 1))]) +@pytest.mark.parametrize("raw_terminal", [False, True]) +async def test_endpoint_completion_waits_for_other_selected_output(transitions, raw_terminal): + orch, state, clients = _direct_orchestrator(transitions, async_chunk=True) + endpoint = state.final_stage_id + other = 0 + clients[other].final_output = True + state.final_output_stage_ids = {other, endpoint} + state.stage_submit_ts = {sid: 1.0 for sid in orch._request_stage_path(state)} + state.finished_stage_ids.update(set(state.stage_submit_ts) - {other, endpoint}) + if raw_terminal: + raw = OmniEngineCoreOutput(request_id="r", new_token_ids=[], finish_reason=FinishReason.STOP) + await orch._apply_raw_terminal_stage_finish(endpoint, raw, state) + await orch._finish_raw_terminal_requests(endpoint, 0, {"r"}) + else: + await orch._route_output(endpoint, 0, _build_request_output("r"), state, None) + assert orch.output_async_queue.empty() + assert state.pending_final_output.stage_id == endpoint + await orch._route_output(other, 0, _build_request_output("r"), state, None) + intermediate = orch.output_async_queue.get_nowait() + terminal = orch.output_async_queue.get_nowait() + assert intermediate.stage_id == other and not intermediate.finished + assert terminal.stage_id == endpoint and terminal.finished + assert orch.output_async_queue.empty() and "r" not in orch.request_states + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transitions", [None, ((0, 2), (2, 1))]) +async def test_abort_outputs_follow_route_and_finish_each_request_once(transitions): + orch, state, clients = _direct_orchestrator(transitions) + for client in clients: + client.final_output = True + state.final_output_stage_ids = {1, 2} + other = OrchestratorRequestState( + request_id="other", final_stage_id=state.final_stage_id, final_output_stage_ids={1, 2} + ) + orch.request_states["other"] = other + for pool in orch.stage_pools: + for rid in ("r", "other"): + pool.select_replica_id(rid) + outputs = await orch._cleanup_request_ids(["r", "other"], abort=True) + for rid in ("r", "other"): + selected = [msg for msg in outputs if msg.request_id == rid] + assert [msg.stage_id for msg in selected] == [sid for sid in orch._stage_routing.stage_order if sid in {1, 2}] + assert [msg.finished for msg in selected] == [False, True] + assert selected[-1].stage_id == state.final_stage_id + assert not orch.request_states + + +@pytest.mark.asyncio +async def test_empty_bridge_aborts_prewarmed_receiver_and_preserves_submission_metrics(): + orch, state, clients = _direct_orchestrator(((0, 2), (2, 1)), async_chunk=True) + orch.stage_pools[2].stage_vllm_config.model_config.stage_connector_config = {"extra": {"role": "sender"}} + endpoint_pool = orch.stage_pools[1] + receiver = FakeStageClient(final_output=True) + endpoint_pool.add_client("receiver-1", receiver, replica_id=1) + endpoint_pool.mark_replica_unavailable(0) + orch.stage_pools[0].select_replica_id("r") + state.stage_submit_ts[0] = 1.0 + assert await orch._prewarm_async_chunk_stages( + "r", FakePromptRequest(request_id="r", prompt_token_ids=[1], resumable=False), state + ) + submit_ts = state.stage_submit_ts[1] + await orch._forward_to_next_stage("r", 0, _build_request_output("r"), state) + terminal = orch.output_async_queue.get_nowait() + assert terminal.finished and terminal.stage_id == 1 and terminal.replica_id == 1 + assert terminal.stage_submit_ts == submit_ts + assert set(state.stage_submit_ts) == {0, 1} + assert receiver.abort_calls == [["r"]] + assert not clients[2].add_request_calls + assert "r" not in orch.request_states + assert all(pool.get_bound_replica_id("r") is None for pool in orch.stage_pools) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transitions", [None, ((0, 2),), ((0, 2), (2, 1))]) +@pytest.mark.parametrize("empty_inputs", [None, []]) +async def test_empty_diffusion_inputs_fail_only_request(mocker, transitions, empty_inputs): + orch, state, clients = _direct_orchestrator(transitions) + target = state.final_stage_id + source = orch._request_stage_path(state)[-2] + clients[target].stage_type = "diffusion" + clients[target].custom_process_input_func = mocker.Mock(return_value=empty_inputs) + orch.stage_pools[source].select_replica_id("r") + + await orch._forward_to_next_stage("r", source, _build_request_output("r"), state) + + terminal = orch.output_async_queue.get_nowait() + assert terminal.finished and terminal.stage_id == target + assert terminal.engine_outputs.error is not None + assert "r" not in orch.request_states + assert not clients[target].add_request_calls + assert all(pool.get_bound_replica_id("r") is None for pool in orch.stage_pools) + assert orch.rpc_async_queue.empty() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transitions", [((0, 2),), ((0, 2), (2, 1))]) +async def test_configured_route_runs_through_real_loop(configured_orchestrator_factory, transitions): + order = (0, 2) if len(transitions) == 1 else (0, 2, 1) + final = order[-1] + clients = [ + FakeStageClient(final_output=(sid == final), next_inputs=[{"prompt_token_ids": [sid + 10]}]) for sid in range(3) + ] + processors = [FakeOutputProcessor(request_outputs=[_build_request_output("r")]) for _ in clients] + harness = configured_orchestrator_factory(clients, output_processors=processors, stage_transitions=transitions) + try: + await _enqueue_add_request( + harness, + request_id="r", + prompt=FakePromptRequest(request_id="r", prompt_token_ids=[1, 2], resumable=False), + original_prompt={"prompt_token_ids": [1, 2]}, + sampling_params_list=[SamplingParams(max_tokens=1) for _ in clients], + final_stage_id=final, + ) + for sid in order: + await _wait_for(lambda: len(clients[sid].add_request_calls) == 1) + clients[sid].push_engine_core_outputs(_engine_core_outputs(f"raw-{sid}", float(sid + 1))) + output = await _get_output_message(harness) + assert output.stage_id == final and output.finished + await _wait_for(lambda: "r" not in harness.orchestrator.request_states) + for sid in set(range(3)) - set(order): + assert not clients[sid].add_request_calls + assert all(pool.get_bound_replica_id("r") is None for pool in harness.orchestrator.stage_pools) + finally: + await _shutdown_orchestrator(harness) + + +@pytest.mark.asyncio +async def test_async_prewarm_uses_actual_edge_and_sender(mocker): + orch, state, clients = _direct_orchestrator(async_chunk=True) + sender = mocker.patch.object(orch, "_build_payload_sender_info", return_value={"source": 0}) + emit = mocker.patch.object(orch, "_emit_tx_edge") + assert await orch._prewarm_async_chunk_stages( + "r", FakePromptRequest(request_id="r", prompt_token_ids=[1, 2], resumable=False), state + ) + assert not clients[1].add_request_calls + request = clients[2].add_request_calls[0][0] + assert request.payload_sender_info == {"source": 0} + sender.assert_called_once_with(0, request_id="r", req_state=state) + assert emit.call_args.kwargs["from_stage"] == 0 + assert emit.call_args.kwargs["to_stage"] == 2 + assert set(state.stage_submit_ts) == {2} + + +@pytest.mark.asyncio +async def test_async_prewarm_traverses_route_order(mocker): + orch, state, _ = _direct_orchestrator(((0, 2), (2, 1)), async_chunk=True) + state.final_stage_id = 1 + state.final_output_stage_ids = {1} + sender = mocker.patch.object(orch, "_build_payload_sender_info", return_value=None) + assert await orch._prewarm_async_chunk_stages( + "r", FakePromptRequest(request_id="r", prompt_token_ids=[1, 2], resumable=False), state + ) + assert list(state.stage_submit_ts) == [2, 1] + assert [call.args[0] for call in sender.call_args_list] == [0, 2] + + +@pytest.mark.asyncio +async def test_async_terminal_waits_for_submitted_route_only(): + orch, state, _ = _direct_orchestrator(async_chunk=True) + state.stage_submit_ts = {0: 1.0, 2: 2.0} + await orch._route_output(2, 0, _build_request_output("r"), state, None) + assert orch.output_async_queue.empty() + assert state.pending_final_output is not None + await orch._route_output(0, 0, _build_request_output("r"), state, None) + assert orch.output_async_queue.get_nowait().finished + assert orch.output_async_queue.empty() + assert "r" not in orch.request_states + + +@pytest.mark.asyncio +async def test_cancellation_during_prewarm_stops_later_route_edges(mocker): + orch, state, clients = _direct_orchestrator(((0, 2), (2, 1)), async_chunk=True) + state.final_stage_id = 1 + on_submit = mocker.patch.object(orch, "_on_stage_submitted") + + async def cancel_after_admission(request, **kwargs): + await orch._cleanup_request_ids(["r"], abort=True) + + mocker.patch.object(clients[2], "add_request_async", side_effect=cancel_after_admission) + assert not await orch._prewarm_async_chunk_stages( + "r", FakePromptRequest(request_id="r", prompt_token_ids=[1, 2], resumable=False), state + ) + assert not clients[1].add_request_calls + assert all(pool.get_bound_replica_id("r") is None for pool in orch.stage_pools) + on_submit.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("already_submitted", [False, True]) +async def test_cfg_companion_release_uses_actual_successor(mocker, already_submitted): + orch, state, _ = _direct_orchestrator() + orch._cfg_tracker.register_companion("r", "unconditional", "r-cfg") + orch._cfg_tracker.defer_parent("r", _build_request_output("r"), 0) + if already_submitted: + state.stage_submit_ts[2] = 1.0 + forward = mocker.patch.object(orch, "_forward_to_next_stage", new_callable=mocker.AsyncMock) + await orch._handle_cfg_companion_ready("r-cfg") + assert forward.await_count == (0 if already_submitted else 1) + + +@pytest.mark.asyncio +async def test_kv_ready_routes_once_to_configured_target(mocker): + orch, state, _ = _direct_orchestrator() + forward = mocker.patch.object(orch, "_forward_to_next_stage", new_callable=mocker.AsyncMock) + raw = EngineCoreOutputs( + outputs=[OmniEngineCoreOutput(request_id="r", new_token_ids=[], kv_transfer_params={"kv_ready": True})] + ) + await orch._handle_kv_ready_raw_outputs(0, raw) + forward.assert_awaited_once() + state.stage_submit_ts[2] = 1.0 + await orch._handle_kv_ready_raw_outputs(0, raw) + assert forward.await_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("final,outputs", [(1, [1]), (2, [1]), (2, [0])]) +async def test_invalid_request_route_fails_before_submission(final, outputs): + orch, _, clients = _direct_orchestrator() + orch.request_states.clear() + await orch._handle_add_request(_submission(final_stage_id=final, final_output_stage_ids=outputs)) + assert isinstance(orch.output_async_queue.get_nowait(), ErrorMessage) + assert not orch.request_states + assert all(not client.add_request_calls for client in clients) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("outputs", [[2.0], [True], [[2]]]) +async def test_bad_output_ids_fail_one_request_without_killing_dispatch(outputs): + orch, _, clients = _direct_orchestrator() + orch.request_states.clear() + await orch._handle_add_request(_submission(final_output_stage_ids=outputs)) + error = orch.output_async_queue.get_nowait() + assert isinstance(error, ErrorMessage) and not error.fatal + assert error.status_code == 400 + assert all(not client.add_request_calls for client in clients) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("count", [0, 1, 2]) +async def test_incomplete_sampling_array_fails_before_any_submission(count): + orch, _, clients = _direct_orchestrator() + orch.request_states.clear() + await orch._handle_add_request(_submission(sampling_params_list=[SamplingParams() for _ in range(count)])) + assert isinstance(orch.output_async_queue.get_nowait(), ErrorMessage) + assert all(not client.add_request_calls for client in clients) + + +@pytest.mark.asyncio +async def test_endpoint_must_be_selected_to_prevent_premature_completion(): + orch, _, clients = _direct_orchestrator() + clients[0].final_output = True + orch.request_states.clear() + await orch._handle_add_request(_submission(final_output_stage_ids=[0])) + error = orch.output_async_queue.get_nowait() + assert "final stage" in error.error + assert all(not client.add_request_calls for client in clients) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "overrides", + [ + {"final_stage_id": 0}, + {"final_output_stage_ids": [0, 2]}, + {"final_output_stage_ids": [2.0]}, + {"final_output_stage_ids": [True]}, + {"final_output_stage_ids": [[2]]}, + {"sampling_params_list": [SamplingParams()]}, + ], +) +async def test_streaming_update_preserves_admitted_route(overrides): + orch, _, clients = _direct_orchestrator() + await orch._handle_streaming_update(_submission(type="streaming_update", **overrides)) + assert isinstance(orch.output_async_queue.get_nowait(), ErrorMessage) + assert "r" not in orch.request_states + assert all(not client.add_request_calls for client in clients) + + +@pytest.mark.asyncio +async def test_streaming_update_can_reuse_sampling_params(): + orch, state, clients = _direct_orchestrator() + params = state.sampling_params_list + await orch._handle_streaming_update(_submission(type="streaming_update", sampling_params_list=[])) + assert state.sampling_params_list is params + assert state.final_stage_id == 2 and state.final_output_stage_ids == {2} + assert len(clients[0].add_request_calls) == 1 + assert orch.output_async_queue.empty() + + +def test_frontend_selects_final_outputs_in_route_order(mocker): + stages = [ + StagePipelineConfig(stage_id=0, model_stage="test", final_output=False), + StagePipelineConfig(stage_id=1, model_stage="test", final_output=True, final_output_type="audio"), + StagePipelineConfig(stage_id=2, model_stage="test", final_output=True, final_output_type="text"), + ] + base = object.__new__(OmniBase) + base._stage_meta_list = stages + base.output_modalities = ["text", "audio"] + pipeline = PipelineConfig( + model_type="test", + stages=tuple(stages), + stage_transitions=((0, 2), (2, 1)), + ) + base.engine = mocker.Mock(spec=OmniEngineBase, pipeline_config=pipeline) + assert base._compute_final_stage_id(["text", "audio"]) == 1 + assert base._compute_final_stage_id(["text"]) == 2 + assert base._compute_final_output_stage_ids(["text"]) == [2] + assert set(base._compute_final_output_stage_ids(["text", "audio"])) == {1, 2} + assert get_final_stage_id_for_e2e(["text", "audio"], ["text", "audio"], stages) == 2 + + +@pytest.mark.parametrize("modalities", [["audio"], ["text", "audio"]]) +@pytest.mark.parametrize("default_modalities", [["text", "audio"], ["text"]]) +def test_frontend_rejects_outputs_only_available_on_inactive_stages(modalities, default_modalities): + stages = [ + FakeStageClient(final_output=False), + FakeStageClient(final_output=True, final_output_type="audio"), + FakeStageClient(final_output=True, final_output_type="text"), + ] + with pytest.raises(ValueError, match="output modalities are unreachable"): + get_final_stage_id_for_e2e(modalities, default_modalities, stages, stage_order=(0, 2)) + assert get_final_stage_id_for_e2e(None, ["text", "audio"], stages, stage_order=(0, 2)) == 2 + assert get_final_stage_id_for_e2e(modalities, ["text", "audio"], stages) == (1 if modalities == ["audio"] else 2) + + +@pytest.mark.asyncio +async def test_cancel_while_waiting_for_replica_stops_old_request_admission(): + orch, state, clients = _direct_orchestrator(async_chunk=True) + pool = orch.stage_pools[2] + pool.add_client("tcp://127.0.0.1:5555", clients[2], replica_id=0) + polled = asyncio.Event() + + def no_replicas(_stage_id): + polled.set() + return SimpleNamespace(replicas=[]) + + pool.attach_hub(SimpleNamespace(get_replicas_for_stage=no_replicas)) + pool.attach_load_balancer(SimpleNamespace(select=lambda task, replicas: 0)) + task = asyncio.create_task( + orch._prewarm_async_chunk_stages( + "r", FakePromptRequest(request_id="r", prompt_token_ids=[1], resumable=False), state + ) + ) + await asyncio.wait_for(polled.wait(), timeout=1) + await orch._cleanup_request_ids(["r"], abort=True) + new_state = OrchestratorRequestState(request_id="r") + orch.request_states["r"] = new_state + pool.bind("r", "tcp://127.0.0.1:5555") + assert not await asyncio.wait_for(task, timeout=1) + assert orch.request_states["r"] is new_state + assert pool.get_bound_replica_id("r") == 0 + assert not clients[2].add_request_calls + await orch._cleanup_request_ids(["r"]) + + +@pytest.mark.asyncio +async def test_lost_preselected_replica_cannot_be_reselected(): + orch, state, clients = _direct_orchestrator(async_chunk=True) + pool = orch.stage_pools[2] + pool.add_client("replica-1", FakeStageClient(final_output=True), replica_id=1) + assert await pool.pick("r", request_state=state) == 0 + pool.evict_replica(0) + assert not await orch._dispatch_or_fail_request( + lambda: pool.submit_initial("r", state, SimpleNamespace()), + req_id="r", + stage_id=2, + operation="bridge submission", + ) + assert isinstance(orch.output_async_queue.get_nowait(), ErrorMessage) + assert "r" not in orch.request_states + assert not pool.clients[1].add_request_calls + + +@pytest.mark.asyncio +@pytest.mark.parametrize("async_chunk", [False, True]) +async def test_reordered_route_starts_at_admitted_nonzero_entry(mocker, async_chunk): + orch, _, clients = _direct_orchestrator(((0, 2), (2, 1)), async_chunk=async_chunk) + orch.request_states.clear() + processed = _processed_request("r", SamplingParams(max_tokens=1)) + build_entry = mocker.patch.object(orch, "_build_entry_stage_request", return_value=processed) + await orch._handle_add_request( + _submission( + entry_stage_id=2, + final_stage_id=1, + final_output_stage_ids=[1], + prompt={"prompt_token_ids": [1, 2]}, + ) + ) + state = orch.request_states["r"] + assert orch._request_stage_path(state) == (2, 1) + build_entry.assert_awaited_once() + assert clients[2].add_request_calls[0][0] is processed + assert not clients[0].add_request_calls + assert bool(clients[1].add_request_calls) is async_chunk + await orch._cleanup_request_ids(["r"], abort=True) + + +@pytest.mark.asyncio +async def test_closing_request_does_not_consume_its_abort_snapshot(mocker): + orch, state, _ = _direct_orchestrator() + state.closing = "aborting" + orch.request_states["healthy"] = OrchestratorRequestState(request_id="healthy", final_stage_id=2) + raw = EngineCoreOutputs( + outputs=[OmniEngineCoreOutput(request_id=key, new_token_ids=[1]) for key in ("r", "healthy")] + ) + process = mocker.patch.object(orch.stage_pools[0], "process_llm_raw_outputs", return_value=[]) + await orch._process_llm_stage_outputs(0, 0, raw, set()) + filtered = process.call_args.args[1] + assert filtered is not raw + assert [output.request_id for output in filtered.outputs] == ["healthy"] + assert [output.request_id for output in raw.outputs] == ["r", "healthy"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transitions", [None, ((0, 2), (2, 1))]) +async def test_prewarm_reuses_replica_selected_by_waiting_bridge(monkeypatch, transitions): + orch, state, clients = _direct_orchestrator(transitions, async_chunk=True) + _, bridge, receiver = orch._request_stage_path(state) + pool = orch.stage_pools[bridge] + pool.stage_vllm_config.model_config.async_chunk = False + replicas = [clients[bridge], FakeStageClient()] + addresses = ["tcp://127.0.0.1:50101", "tcp://127.0.0.1:50102"] + for replica_id, client in enumerate(replicas): + pool.add_client(addresses[replica_id], client, replica_id=replica_id) + monkeypatch.setattr( + client, + "get_payload_sender_info", + lambda port=50101 + replica_id: {"host": "127.0.0.1", "zmq_port": port}, + raising=False, + ) + orch.stage_pools[0].select_replica_id("r") + hub = _FakeHub([]) + pool.attach_hub(hub) + pool.attach_load_balancer(_RoundRobinLB()) + entered = {name: asyncio.Event() for name in ("bridge-submit", "prewarm")} + release = {name: asyncio.Event() for name in entered} + sleep = asyncio.sleep + + async def wait_for_capacity(delay): + task = asyncio.current_task() + assert task is not None + name = task.get_name() + if name in entered: + entered[name].set() + await release[name].wait() + else: + await sleep(delay) + + monkeypatch.setattr("vllm_omni.engine.stage_pool.asyncio.sleep", wait_for_capacity) + state.stage_submit_ts[bridge] = 1.0 + submit = asyncio.create_task( + pool.submit_initial("r", state, _processed_request("r", state.sampling_params_list[bridge])), + name="bridge-submit", + ) + prewarm = None + try: + await asyncio.wait_for(entered["bridge-submit"].wait(), 1) + prewarm = asyncio.create_task( + orch._prewarm_async_chunk_stages("r", FakePromptRequest("r", [1, 2], False), state), name="prewarm" + ) + await asyncio.wait_for(entered["prewarm"].wait(), 1) + hub._replicas = [replace(_replica(address), stage_id=bridge) for address in addresses] + release["bridge-submit"].set() + selected = await asyncio.wait_for(submit, 1) + release["prewarm"].set() + assert await asyncio.wait_for(prewarm, 1) + assert pool.get_bound_replica_id("r") == selected + advertised = clients[receiver].add_request_calls[0][0].payload_sender_info + assert advertised == replicas[selected].get_payload_sender_info() + await orch._cleanup_request_ids(["r"], abort=True) + assert replicas[selected].abort_calls == [["r"]] + assert not replicas[1 - selected].abort_calls + finally: + for task in (submit, prewarm): + if task is not None: + task.cancel() + await asyncio.gather(*(task for task in (submit, prewarm) if task is not None), return_exceptions=True) + + +@pytest.mark.asyncio +async def test_prewarm_does_not_bind_bridge_only_stages(): + orch, state, _ = _direct_orchestrator(None, async_chunk=True) + orch.stage_pools[0].select_replica_id("r") + for pool in orch.stage_pools[1:]: + pool.stage_vllm_config.model_config.async_chunk = False + assert await orch._prewarm_async_chunk_stages("r", FakePromptRequest("r", [1, 2], False), state) + assert all(pool.get_bound_replica_id("r") is None for pool in orch.stage_pools[1:]) + assert not state.stage_submit_ts + + +@pytest.mark.asyncio +@pytest.mark.parametrize("transitions", [None, ((0, 2),), ((0, 2), (2, 1))]) +@pytest.mark.parametrize("fail_first", [False, True]) +@pytest.mark.parametrize("finish_during_abort", [False, True]) +async def test_abort_failure_drains_without_resuming_or_leaking_request( + monkeypatch, mocker, transitions, fail_first, finish_during_abort +): + orch, state, clients = _direct_orchestrator(transitions, async_chunk=True) + orch._running_counter = FakeRunningCounter() + orch._register_running_request(state) + release_resources = mocker.spy(orch, "_release_stage_transfer_resources") + path = orch._request_stage_path(state) + failed_stage = path[0] if fail_first else path[-1] + for sid in path: + clients[sid].final_output = True + processor = MultimodalOutputProcessor(tokenizer=None, log_stats=False) + orch.stage_pools[sid]._output_processor = processor + request = _processed_request(f"internal-r-{sid}", state.sampling_params_list[sid]) + request.external_req_id = "r" + await orch.stage_pools[sid].submit_initial("r", state, request) + state.stage_submit_ts[sid] = 1.0 + failed_pool = orch.stage_pools[failed_stage] + processor = failed_pool.output_processor + healthy = OrchestratorRequestState( + request_id="healthy", + entry_stage_id=failed_stage, + final_stage_id=failed_stage, + final_output_stage_ids={failed_stage}, + sampling_params_list=state.sampling_params_list, + stage_submit_ts={failed_stage: 1.0}, + ) + orch.request_states["healthy"] = healthy + orch._register_running_request(healthy) + request = _processed_request("internal-healthy", state.sampling_params_list[failed_stage]) + request.external_req_id = "healthy" + await failed_pool.submit_initial("healthy", healthy, request) + + async def terminal(request_id): + terminals: set[str] = set() + processed = await orch._process_llm_stage_outputs( + failed_stage, + 0, + EngineCoreOutputs( + outputs=[ + OmniEngineCoreOutput( + request_id=request_id, + new_token_ids=[], + finish_reason=FinishReason.STOP, + ) + ] + ), + terminals, + ) + await orch._handle_processed_outputs(failed_stage, 0, processed) + await orch._finish_raw_terminal_requests(failed_stage, 0, terminals) + + async def fail_abort(request_ids): + if finish_during_abort: + await terminal(f"internal-r-{failed_stage}") + assert "r" in orch.request_states + raise RuntimeError("abort transport failed") + + monkeypatch.setattr(clients[failed_stage], "abort_requests_async", fail_abort) + await orch._handle_abort(AbortRequestMessage(request_ids=["r"], rpc_id="abort-r")) + assert not orch.rpc_async_queue.get_nowait().success + for sid in set(path) - {failed_stage}: + assert clients[sid].abort_calls == [[f"internal-r-{sid}"]] + assert orch.stage_pools[sid].get_bound_replica_id("r") is None + if not finish_during_abort: + assert state.closing + assert orch._running_counter.value == 2 + release_resources.assert_not_called() + with pytest.raises(RequestClosedError, match="r"): + await failed_pool.submit_initial("r", state, request) + await terminal(f"internal-r-{failed_stage}") + assert "r" not in orch.request_states + assert orch._running_counter.value == 1 + release_resources.assert_called_once_with(["r"]) + assert all(pool.get_bound_replica_id("r") is None for pool in orch.stage_pools) + assert f"internal-r-{failed_stage}" not in processor.request_states + assert "internal-healthy" in processor.request_states + assert failed_pool.get_bound_replica_id("healthy") == 0 + error = orch.output_async_queue.get_nowait() + assert isinstance(error, ErrorMessage) and error.request_id == "r" and not error.fatal + assert orch.output_async_queue.empty() + await terminal("internal-healthy") + assert orch.output_async_queue.get_nowait().request_id == "healthy" + assert not orch.request_states + assert not processor.request_states + assert orch._running_counter.value == 0 + + +@pytest.mark.asyncio +async def test_failed_abort_waits_for_each_remaining_stage(monkeypatch): + orch, state, clients = _direct_orchestrator() + state.stage_submit_ts = {0: 1.0, 2: 1.0} + for sid in (0, 2): + processor = MultimodalOutputProcessor(tokenizer=None, log_stats=False) + orch.stage_pools[sid]._output_processor = processor + request = _processed_request("r", state.sampling_params_list[sid]) + request.external_req_id = "r" + await orch.stage_pools[sid].submit_initial("r", state, request) + + async def fail_abort(request_ids): + raise RuntimeError("abort failed") + + for sid in (0, 2): + monkeypatch.setattr(clients[sid], "abort_requests_async", fail_abort) + await orch._handle_abort(AbortRequestMessage(request_ids=["r"], rpc_id="abort-r")) + assert not orch.rpc_async_queue.get_nowait().success + for sid in (2, 0): + terminals: set[str] = set() + await orch._process_llm_stage_outputs( + sid, + 0, + EngineCoreOutputs( + outputs=[ + OmniEngineCoreOutput( + request_id="r", + new_token_ids=[], + finish_reason=FinishReason.STOP, + ) + ] + ), + terminals, + ) + await orch._finish_raw_terminal_requests(sid, 0, terminals) + assert ("r" in orch.request_states) == (sid == 2) + assert all(not orch.stage_pools[sid].output_processor.request_states for sid in (0, 2)) + assert all(pool.get_bound_replica_id("r") is None for pool in orch.stage_pools) + + +@pytest.mark.asyncio +async def test_abort_failure_on_one_replica_does_not_skip_other_replicas(monkeypatch): + orch, state, clients = _direct_orchestrator() + pool = orch.stage_pools[2] + pool._output_processor = MultimodalOutputProcessor(tokenizer=None, log_stats=False) + pool.add_client("other", FakeStageClient(final_output=True), replica_id=1) + for replica_id, request_id in enumerate(("r", "other")): + request = _processed_request(request_id, state.sampling_params_list[2]) + request.external_req_id = request_id + await pool.submit_initial(request_id, state, request, pinned_replica_id=replica_id) + + async def fail_abort(request_ids): + raise RuntimeError("abort failed") + + monkeypatch.setattr(clients[2], "abort_requests_async", fail_abort) + with pytest.raises(RuntimeError, match="abort failed"): + await pool.abort_requests(["r", "other"]) + assert pool.get_bound_replica_id("r") == 0 + assert pool.get_bound_replica_id("other") is None + assert set(pool.output_processor.request_states) == {"r"} + assert pool.clients[1].abort_calls == [["other"]] + + +@pytest.mark.asyncio +async def test_draining_streaming_request_ignores_processed_segment_stop(monkeypatch): + orch, state, clients = _direct_orchestrator(async_chunk=True) + state.streaming.enabled = True + state.stage_submit_ts[2] = 1.0 + orch.stage_pools[2].select_replica_id("r") + + async def fail_abort(request_ids): + raise RuntimeError("abort failed") + + monkeypatch.setattr(clients[2], "abort_requests_async", fail_abort) + await orch._handle_abort(AbortRequestMessage(request_ids=["r"], rpc_id="abort-r")) + state.streaming.segments[2] = StreamingSegmentState(finished=True) + await orch._route_output(2, 0, _build_request_output("r"), state, None) + assert orch.request_states["r"] is state + assert 2 not in state.finished_stage_ids + state.streaming.segments[2] = StreamingSegmentState(finished=False) + await orch._route_output(2, 0, _build_request_output("r"), state, None) + assert "r" not in orch.request_states + + +@pytest.mark.asyncio +async def test_abort_failure_unblocks_frontend_output_stream(monkeypatch): + orch, state, clients = _direct_orchestrator() + orch.stage_pools[2].select_replica_id("r") + state.stage_submit_ts[2] = 1.0 + frontend = object.__new__(AsyncOmniBase) + frontend.request_states = {"r": ClientRequestState("r", queue=orch.output_async_queue, final_stage_id=2)} + stream = frontend._process_orchestrator_results("r", OrchestratorAggregator(3, False, 1.0, 2), 2, {}, 1.0) + waiting = asyncio.create_task(anext(stream)) + + async def fail_abort(request_ids): + raise RuntimeError("abort transport failed") + + abort = clients[2].abort_requests_async + monkeypatch.setattr(clients[2], "abort_requests_async", fail_abort) + try: + await orch._handle_abort(AbortRequestMessage(request_ids=["r"], rpc_id="abort-r")) + with pytest.raises(RuntimeError, match="abort transport failed"): + await asyncio.wait_for(waiting, 1) + finally: + waiting.cancel() + await asyncio.gather(waiting, return_exceptions=True) + await stream.aclose() + monkeypatch.setattr(clients[2], "abort_requests_async", abort) + await orch._cleanup_request_ids(["r"], abort=True) + + +@pytest.mark.asyncio +async def test_failed_abort_drains_and_reuses_id_through_real_loops(configured_orchestrator_factory, monkeypatch): + clients = [FakeStageClient(final_output=(sid == 2)) for sid in range(3)] + pools = _build_stage_pools([[client] for client in clients]) + for pool in pools: + pool._output_processor = MultimodalOutputProcessor(tokenizer=None, log_stats=False) + abort = clients[2].abort_requests_async + + async def fail_abort(request_ids): + raise RuntimeError("abort transport failed") + + monkeypatch.setattr(clients[2], "abort_requests_async", fail_abort) + harness = configured_orchestrator_factory(clients, stage_pools=pools, stage_transitions=((0, 2),), async_chunk=True) + orch = harness.orchestrator + params = [SamplingParams(max_tokens=1) for _ in clients] + + async def submit(): + request = _processed_request("r", params[0]) + request.external_req_id = "r" + await _enqueue_add_request( + harness, + request_id="r", + prompt=request, + original_prompt={"prompt_token_ids": [1, 2]}, + sampling_params_list=params, + final_stage_id=2, + ) + + try: + await submit() + await _wait_for(lambda: len(clients[2].add_request_calls) == 1) + old_state = orch.request_states["r"] + harness.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["r"], rpc_id="abort-r")) + result = await _get_rpc_message(harness) + assert isinstance(result, AbortResultMessage) + assert result.rpc_id == "abort-r" and not result.success + await _wait_for(lambda: not harness.output_sync_q.empty()) + error = harness.output_sync_q.get_nowait() + assert isinstance(error, ErrorMessage) and not error.fatal and error.request_id == "r" + assert orch.request_states["r"] is old_state and old_state.closing + assert pools[0].get_bound_replica_id("r") is None + assert pools[2].get_bound_replica_id("r") == 0 + clients[2].push_engine_core_outputs( + EngineCoreOutputs( + outputs=[OmniEngineCoreOutput(request_id="r", new_token_ids=[], finish_reason=FinishReason.STOP)] + ) + ) + await _wait_for(lambda: "r" not in orch.request_states) + assert all(pool.get_bound_replica_id("r") is None for pool in pools) + assert all(not pool.output_processor.request_states for pool in pools) + assert harness.output_sync_q.empty() + + monkeypatch.setattr(clients[2], "abort_requests_async", abort) + await submit() + await _wait_for(lambda: len(clients[2].add_request_calls) == 2) + assert orch.request_states["r"] is not old_state + assert not orch.request_states["r"].closing + harness.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["r"], rpc_id="abort-reused-r")) + result = await _get_rpc_message(harness) + assert isinstance(result, AbortResultMessage) + assert result.rpc_id == "abort-reused-r" and result.success + assert "r" not in orch.request_states + assert not clients[1].add_request_calls + finally: + await _shutdown_orchestrator(harness) + + +@pytest.mark.asyncio +async def test_cancelled_bridge_validation_cannot_fail_reused_request(monkeypatch): + orch, old_state, clients = _direct_orchestrator() + clients[2].next_inputs = [{"prompt": "late input"}] + processor = FakeInputProcessor() + orch._stage_input_processors[2] = processor + entered = asyncio.Event() + release = threading.Event() + loop = asyncio.get_running_loop() + + def reject_input(**kwargs): + loop.call_soon_threadsafe(entered.set) + assert release.wait(5), "Input validation was not released" + raise ValueError("late validation failure") + + monkeypatch.setattr(processor, "process_inputs", reject_input) + await orch.stage_pools[0].submit_initial("r", old_state, _processed_request("r", old_state.sampling_params_list[0])) + forward = asyncio.create_task(orch._forward_to_next_stage("r", 0, _build_request_output("r"), old_state)) + try: + await asyncio.wait_for(entered.wait(), 1) + await orch._handle_abort(AbortRequestMessage(request_ids=["r"], rpc_id="abort-old-r")) + assert orch.rpc_async_queue.get_nowait().success + await orch._handle_add_request(_submission()) + new_state = orch.request_states["r"] + assert new_state is not old_state + release.set() + await asyncio.wait_for(forward, 2) + assert orch.request_states["r"] is new_state + assert not new_state.closing + assert orch.stage_pools[0].get_bound_replica_id("r") == 0 + assert clients[0].abort_calls == [["r"]] + assert orch.output_async_queue.empty() + finally: + release.set() + await asyncio.gather(forward, return_exceptions=True) + await orch._cleanup_request_ids(["r"], abort=True) diff --git a/tests/engine/test_duplex_omni_engine.py b/tests/engine/test_duplex_omni_engine.py index 18211f94b8c..7489067db67 100644 --- a/tests/engine/test_duplex_omni_engine.py +++ b/tests/engine/test_duplex_omni_engine.py @@ -12,7 +12,7 @@ import pytest -from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig +from vllm_omni.config.stage_config import DuplexSessionRuntimeConfig, PipelineConfig, StagePipelineConfig from vllm_omni.engine.duplex.commands import Heartbeat from vllm_omni.engine.duplex.config import DuplexCapabilities, DuplexSessionConfig from vllm_omni.engine.duplex.delivery import DuplexOutputBuffer @@ -351,3 +351,16 @@ def execute(*args, **kwargs): engine.rpc_client.execute = execute await engine.touch_session_async("sid", activity="heartbeat") assert seen == [False] + + +def test_nonsequential_duplex_route_is_rejected_before_loading_plugin(): + engine = object.__new__(DuplexOmniEngine) + engine.model = "test" + engine.pipeline_config = PipelineConfig( + model_type="test", + stages=tuple(StagePipelineConfig(stage_id=sid, model_stage="test", final_output=sid == 2) for sid in range(3)), + duplex_plugin="test.plugin", + stage_transitions=((0, 2),), + ) + with pytest.raises(ValueError, match="route-aware duplex plugin"): + engine._validate_deployment() diff --git a/tests/engine/test_duplex_orchestrator.py b/tests/engine/test_duplex_orchestrator.py index 7f9ef886a5c..e055c11cf48 100644 --- a/tests/engine/test_duplex_orchestrator.py +++ b/tests/engine/test_duplex_orchestrator.py @@ -1013,7 +1013,8 @@ async def test_payload_sender_info_binds_from_explicit_session_state(monkeypatch assert pool.get_bound_replica_id(request_id) is None with monkeypatch.context() as patch: patch.delattr(orchestrator, "request_states") - assert orchestrator._build_payload_sender_info(0, request_id=request_id) == {"replica": 0} + assert orchestrator._build_payload_sender_info(0) == {"replica": 0} + assert orchestrator._build_payload_sender_info(0, request_id=request_id) is None assert pool.get_bound_replica_id(request_id) is None sender_info = orchestrator._build_payload_sender_info(0, request_id=request_id, req_state=state) assert sender_info == {"replica": 1} diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 44f6177090d..39b8a564e47 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -23,6 +23,7 @@ from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.engine import OmniEngineCoreOutput from vllm_omni.engine.errors import NativeKVHandoffError from vllm_omni.engine.messages import ( @@ -237,6 +238,7 @@ async def collective_rpc_async( class FakeOutputProcessor: def __init__(self, *, request_outputs: list[Any] | None = None) -> None: + self.request_states: dict[str, Any] = {} self.request_outputs = list(request_outputs or []) self.add_request_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] self.abort_calls: list[list[str]] = [] @@ -315,8 +317,8 @@ def _sampling_params(max_tokens: int = 4) -> SamplingParams: return SamplingParams(max_tokens=max_tokens) -def _engine_core_outputs(tag: str, timestamp: float) -> SimpleNamespace: - return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None, finished_requests=None) +def _engine_core_outputs(tag: str, timestamp: float) -> EngineCoreOutputs: + return EngineCoreOutputs(outputs=[EngineCoreOutput(request_id=tag, new_token_ids=[])], timestamp=timestamp) def _terminal_engine_core_outputs( @@ -428,6 +430,7 @@ def _build_harness( async_chunk: bool = False, log_stats: bool = False, stage_pools: list[StagePool] | None = None, + stage_transitions: tuple[tuple[int, int], ...] | None = None, ) -> OrchestratorFixture: """Build an Orchestrator test harness. @@ -461,6 +464,7 @@ async def _run() -> None: output_async_queue=output_queue.async_q, rpc_async_queue=rpc_queue.async_q, stage_pools=stage_pools, + stage_transitions=stage_transitions, async_chunk=async_chunk, log_stats=log_stats, ) @@ -785,6 +789,7 @@ async def test_replicated_stage1_keys_mm_cache_per_receiving_replica() -> None: stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(multimodal_config=mm_config)), ), ] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) seen_uuids: dict[str, str] = {} def process_inputs(*, request_id, prompt, params, **kwargs): @@ -1764,7 +1769,7 @@ async def test_stage_pool_pinned_submit_never_falls_back(operation: str, failure async def fail_selection(*args, **kwargs): raise AssertionError("pinned dispatch must not select another replica") - monkeypatch.setattr(pool, "_pick_or_select", fail_selection) + monkeypatch.setattr(pool, "pick", fail_selection) with pytest.raises(StageUnavailableError, match="pinned replica"): await getattr(pool, operation)("pinned", state, request, pinned_replica_id=pinned_replica_id) @@ -1784,7 +1789,7 @@ async def test_stage_pool_pinned_submit_rejects_mismatched_cache_replica(monkeyp async def fail_selection(*args, **kwargs): raise AssertionError("pinned dispatch must not select another replica") - monkeypatch.setattr(pool, "_pick_or_select", fail_selection) + monkeypatch.setattr(pool, "pick", fail_selection) with pytest.raises(StageUnavailableError, match="replica 0 keyed by pinned"): await pool.submit_initial("pinned", state, request, pinned_replica_id=1, expected_replica_id=0) @@ -1913,6 +1918,7 @@ async def submit_update(self, request_id, req_state, request, *, prompt_text=Non orchestrator.async_chunk = False orchestrator.request_states = {} orchestrator.stage_pools = [pool] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) add_request_calls: list[str] = [] async def record_add_request(msg) -> None: @@ -1995,6 +2001,7 @@ async def check_dispatch(request_id, req_state, request, **kwargs): def test_native_handoff_uses_bound_replica_and_reports_missing_binding(mocker) -> None: orchestrator = object.__new__(Orchestrator) orchestrator.stage_pools = [mocker.Mock()] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) orchestrator.stage_pools[0].get_bound_client.return_value = None req_state = SimpleNamespace(native_kv_transfer_id="xfer-req") output = SimpleNamespace(kv_transfer_params={"num_transfer_tokens": 4}) @@ -2059,6 +2066,7 @@ async def submit_update(self, request_id, req_state, request, *, prompt_text=Non ) } orchestrator.stage_pools = [pool] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) await orchestrator._handle_streaming_update( StageSubmissionMessage( @@ -2103,6 +2111,7 @@ def build_stage_metrics(self, outputs, **_kwargs): req_state.stage_submit_ts[0] = time.time() orchestrator.request_states = {"req-stream": req_state} orchestrator.stage_pools = [pool] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) routed: list[Any] = [] async def record_route(_stage_id, _replica_id, _output, _req_state, stage_metrics): @@ -2150,6 +2159,7 @@ def build_stage_metrics(self, outputs, **_kwargs): req_state.stage_submit_ts[0] = time.time() orchestrator.request_states = {"req-no-metrics": req_state} orchestrator.stage_pools = [pool] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) routed: list[Any] = [] async def record_route(_stage_id, _replica_id, _output, _req_state, stage_metrics): @@ -2558,7 +2568,7 @@ async def abort_requests_async(self, request_ids: list[str]) -> None: await super().abort_requests_async(request_ids) raise EngineDeadError("engine is still dead during abort") - dead, survivor = DeadClient(), FakeStageClient() + dead, survivor = DeadClient(final_output=True), FakeStageClient(final_output=True) processor = MultimodalOutputProcessor(tokenizer=None, log_stats=False) pool = StagePool( 0, @@ -2713,8 +2723,8 @@ async def test_collective_rpc_ignores_invalid_stage_ids(orchestrator_factory, ca @pytest.mark.asyncio async def test_multi_replica_cfg_companion_inherits_parent_affinity(orchestrator_factory) -> None: """CFG companions should be routed to the same stage-0 replica as their parent.""" - stage0_r0 = FakeStageClient(stage_type="llm", final_output=False) - stage0_r1 = FakeStageClient(stage_type="llm", final_output=False) + stage0_r0 = FakeStageClient(stage_type="llm", final_output=True) + stage0_r1 = FakeStageClient(stage_type="llm", final_output=True) default_vllm_cfg = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64)) stage_pools = _build_stage_pools( [[stage0_r0, stage0_r1]], @@ -2956,6 +2966,7 @@ def release_bindings(self, request_ids: list[str]) -> None: second = _Pool(fail_once=True) orchestrator = object.__new__(Orchestrator) orchestrator.stage_pools = [first, second] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) with pytest.raises(RuntimeError, match="stage abort failed"): await orchestrator._abort_request_ids(["req-a"]) @@ -2995,6 +3006,7 @@ async def test_raw_stage_error_reaches_caller_before_normal_terminal_routing(moc stage_pools=[], ) orchestrator.stage_pools = pools + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) orchestrator.request_states["stuck"] = OrchestratorRequestState(request_id="stuck", final_stage_id=1) mocker.patch.object(orchestrator, "_handle_kv_ready_raw_outputs", new_callable=mocker.AsyncMock) cleanup = mocker.patch.object(orchestrator, "_cleanup_request_ids", new_callable=mocker.AsyncMock) diff --git a/tests/engine/test_orchestrator_async_chunk_prewarm.py b/tests/engine/test_orchestrator_async_chunk_prewarm.py index 1c8a1383b25..38f55a5dffc 100644 --- a/tests/engine/test_orchestrator_async_chunk_prewarm.py +++ b/tests/engine/test_orchestrator_async_chunk_prewarm.py @@ -18,6 +18,7 @@ import torch from vllm.sampling_params import SamplingParams +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.data_entry_keys import ASYNC_CHUNK_PREWARM_NS from vllm_omni.engine.duplex.contracts import DuplexFence, DuplexStageRequestContext, DuplexStageSubmission from vllm_omni.engine.duplex_orchestrator import DuplexOrchestrator, DuplexOrchestratorRequestState @@ -50,6 +51,7 @@ class _FakePool: """Single-replica LLM pool that records the requests it is sent.""" stage_type = "llm" + final_output = True def __init__(self, role: str, stage_client: Any = None) -> None: self.stage_client = stage_client if stage_client is not None else SimpleNamespace() @@ -61,6 +63,9 @@ def __init__(self, role: str, stage_client: Any = None) -> None: def live_replica_ids(self) -> list[int]: return [0] + async def pick(self, request_id, *, request_state): + return 0 + async def submit_initial(self, request_id, _req_state, request, prompt_text=None, expected_replica_id=None): self.submitted.append(request) return 0 @@ -80,6 +85,7 @@ def _orchestrator(payload_func: Any = None, cls: type[Orchestrator] = Orchestrat code2wav_client = None if payload_func is None else _client(payload_func) orchestrator = object.__new__(cls) orchestrator.stage_pools = [_FakePool("sender"), _FakePool("receiver"), _FakePool("receiver", code2wav_client)] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) orchestrator.async_chunk = True orchestrator.request_states = {} orchestrator._emit_tx_edge = lambda **_kwargs: None @@ -196,6 +202,7 @@ async def test_prewarm_submits_plain_placeholder(case: dict[str, Any], expected_ ) streaming = case.get("streaming", False) req_state.streaming.enabled = streaming + orchestrator.request_states[req_state.request_id] = req_state stage0_kwargs = case.get("stage0", {"resumable": False}) prewarmed = await orchestrator._prewarm_async_chunk_stages( diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index 56314662b10..1da4d0b926e 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -580,7 +580,7 @@ def _build_bare_orchestrator(stage_pools) -> tuple[Orchestrator, tuple[janus.Que def _register_request(orchestrator: Orchestrator, request_id: str, *, submitted_stage: int | None = 0) -> None: - state = OrchestratorRequestState(request_id=request_id) + state = OrchestratorRequestState(request_id=request_id, final_stage_id=len(orchestrator.stage_pools) - 1) if submitted_stage is not None: state.stage_submit_ts[submitted_stage] = time.time() orchestrator.request_states[request_id] = state diff --git a/tests/engine/test_orchestrator_first_audio.py b/tests/engine/test_orchestrator_first_audio.py index dc720c44aab..8baeec7d407 100644 --- a/tests/engine/test_orchestrator_first_audio.py +++ b/tests/engine/test_orchestrator_first_audio.py @@ -10,6 +10,7 @@ from vllm.sampling_params import RequestOutputKind from vllm.v1.engine import FinishReason +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.messages import ErrorMessage, OutputMessage @@ -57,6 +58,9 @@ def register(self): def get_bound_replica_id(self, request_id): return self.bound_replica + async def pick(self, request_id, *, request_state): + return self.replica_id + def schedule_release_request_resources(self, request_ids): pass @@ -71,6 +75,7 @@ def _orchestrator(output_kind=RequestOutputKind.DELTA, *, registered=True, repli SimpleNamespace(final_output=False, schedule_release_request_resources=lambda request_ids: None), _CodecPool(output_kind, registered, replica_id), ] + obj._stage_routing = StageRouting.from_transitions(len(obj.stage_pools)) obj.request_states = {"r": OrchestratorRequestState(request_id="r", final_stage_id=final_stage_id)} obj.output_async_queue = asyncio.Queue() obj._cfg_tracker = CfgCompanionTracker() @@ -146,6 +151,7 @@ def _audio(output): async def test_single_stage_first_audio_obeys_output_kind_and_terminal_order(kind, suffix_overtook): obj = _orchestrator(kind, final_stage_id=0) obj.stage_pools = [obj.stage_pools[1]] + obj._stage_routing = StageRouting.from_transitions(len(obj.stage_pools)) suffix = _raw(required=True, samples=(20, 21)) terminal = _raw(required=True, terminal=True, samples=(30, 31)) if suffix_overtook: @@ -172,6 +178,7 @@ async def test_single_stage_first_audio_obeys_output_kind_and_terminal_order(kin async def test_single_stage_cancellation_discards_overtaking_terminal_and_late_first_audio(): obj = _orchestrator(final_stage_id=0) obj.stage_pools = [obj.stage_pools[1]] + obj._stage_routing = StageRouting.from_transitions(len(obj.stage_pools)) await _codec_output(obj, _raw(required=True, terminal=True), stage=0) assert obj.request_states["r"].pending_first_audio_outputs await obj._cleanup_request_ids(["r"], abort=True) @@ -183,6 +190,7 @@ async def test_single_stage_cancellation_discards_overtaking_terminal_and_late_f async def test_single_stage_first_audio_processing_failure_releases_waiting_terminal(): obj = _orchestrator(final_stage_id=0) obj.stage_pools = [obj.stage_pools[1]] + obj._stage_routing = StageRouting.from_transitions(len(obj.stage_pools)) await _codec_output(obj, _raw(required=True, terminal=True), stage=0) obj.stage_pools[0].process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed")) await obj._route_upstream_first_audio(0, 2, _raw(first=True)) @@ -264,7 +272,9 @@ async def test_first_audio_flush_uses_set_with_real_raw_terminal_handlers(termin await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11))) obj._finish_raw_terminal_requests.assert_awaited_once() - stage, replica, terminals = obj._finish_raw_terminal_requests.await_args.args + call = obj._finish_raw_terminal_requests.await_args + assert call is not None + stage, replica, terminals = call.args assert (stage, replica) == (1, 2) assert isinstance(terminals, set) assert terminals == ({"r"} if terminal else set()) @@ -280,11 +290,18 @@ async def test_prewarm_flushes_first_audio_after_successful_submission(mocker, p obj = _orchestrator(registered=False, final_stage_id=2 if processing_failed else 1) req_state = obj.request_states["r"] req_state.sampling_params_list = [None] * (req_state.final_stage_id + 1) + if processing_failed: + later_pool = SimpleNamespace( + submit_initial=mocker.AsyncMock(), schedule_release_request_resources=lambda request_ids: None + ) + obj.stage_pools.append(later_pool) + obj._stage_routing = StageRouting.from_transitions(len(obj.stage_pools)) await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11))) pool = obj.stage_pools[1] pool.stage_type = "llm" pool.stage_vllm_config = SimpleNamespace(model_config=SimpleNamespace(hf_config=None)) obj.stage_pools[0].get_bound_replica_id = lambda request_id: 0 + obj.stage_pools[0].pick = mocker.AsyncMock(return_value=0) obj._stage_receives_async_chunks = lambda stage: True obj._build_payload_sender_info = Mock(return_value={}) obj._on_stage_submitted = Mock() @@ -300,10 +317,6 @@ async def submit(*args, **kwargs): pool.submit_initial = submit if processing_failed: pool.process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed")) - later_pool = SimpleNamespace( - submit_initial=AsyncMock(), schedule_release_request_resources=lambda request_ids: None - ) - obj.stage_pools.append(later_pool) submitted = await obj._prewarm_async_chunk_stages("r", SimpleNamespace(prompt_token_ids=[0]), req_state) assert submitted is not processing_failed messages = _messages(obj) @@ -356,6 +369,8 @@ async def test_later_aligner_receives_full_codec_waveform(): from vllm_omni.model_executor.stage_input_processors.forced_aligner import code2wav2aligner obj = _orchestrator(RequestOutputKind.CUMULATIVE, final_stage_id=2) + obj.stage_pools.append(SimpleNamespace(final_output=True)) + obj._stage_routing = StageRouting.from_transitions(len(obj.stage_pools)) await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11))) await _codec_output(obj, _raw(required=True, terminal=True, samples=(20, 21))) messages = _messages(obj) @@ -436,3 +451,24 @@ async def test_invalid_first_audio_reports_error_and_releases_request(audio, cod obj._release_request_bindings.assert_called_once_with(["r"]) await obj._route_upstream_first_audio(0, 0, _raw(first=True)) assert obj.output_async_queue.empty() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("suffix_overtook", [False, True]) +async def test_first_audio_uses_configured_codec_and_keeps_chunk_order(suffix_overtook): + obj = _orchestrator(final_stage_id=2) + codec = obj.stage_pools[1] + obj.stage_pools = [obj.stage_pools[0], SimpleNamespace(final_output=False), codec] + obj._stage_routing = StageRouting.from_transitions(3, ((0, 2),)) + if suffix_overtook: + await _codec_output(obj, _raw(required=True, samples=(20, 21)), stage=2) + await _codec_output(obj, _raw(required=True, terminal=True, samples=(30, 31)), stage=2) + assert obj.output_async_queue.empty() + await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11))) + if not suffix_overtook: + await _codec_output(obj, _raw(required=True, samples=(20, 21)), stage=2) + await _codec_output(obj, _raw(required=True, terminal=True, samples=(30, 31)), stage=2) + messages = _messages(obj) + assert [_audio(message.engine_outputs).tolist() for message in messages] == [[10, 11], [20, 21], [30, 31]] + assert all(message.stage_id == 2 for message in messages) + assert messages[-1].finished diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index 5f3d1147145..d43b5053cc3 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -10,6 +10,7 @@ from vllm import SamplingParams from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.messages import OutputMessage from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState @@ -355,6 +356,7 @@ def test_forward_to_diffusion_attaches_kv_sender_info(): orchestrator.num_stages = 2 orchestrator.stage_pools = [sender_pool, diffusion_pool] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) orchestrator._cfg_tracker = CfgCompanionTracker() params = OmniDiffusionSamplingParams() @@ -364,6 +366,9 @@ def test_forward_to_diffusion_attaches_kv_sender_info(): sampling_params_list=[SamplingParams(max_tokens=4), params], final_stage_id=1, ) + for pool in orchestrator.stage_pools: + if pool.stage_type == "llm": + pool.select_replica_id(req_state.request_id) output = SimpleNamespace(request_id="req-1", finished=True) asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-1", sender_pool.stage_id, output, req_state)) @@ -384,6 +389,7 @@ def test_diffusion_resubmission_preserves_sender_endpoints(): diffusion_pool = StagePool(1, diffusion_stage) orchestrator.num_stages = 2 orchestrator.stage_pools = [sender_pool, diffusion_pool] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) orchestrator._cfg_tracker = CfgCompanionTracker() state = OrchestratorRequestState( request_id="streaming", @@ -391,6 +397,9 @@ def test_diffusion_resubmission_preserves_sender_endpoints(): sampling_params_list=[SamplingParams(max_tokens=4), OmniDiffusionSamplingParams()], final_stage_id=1, ) + for pool in orchestrator.stage_pools: + if pool.stage_type == "llm": + pool.select_replica_id(state.request_id) output = SimpleNamespace(request_id="streaming", finished=True) async def forward_twice(): @@ -414,6 +423,7 @@ def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info(): orchestrator.num_stages = 3 orchestrator.stage_pools = [source_pool, previous_pool, diffusion_pool] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) orchestrator._cfg_tracker = CfgCompanionTracker() params = OmniDiffusionSamplingParams() @@ -423,6 +433,9 @@ def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info(): sampling_params_list=[SamplingParams(max_tokens=4), SamplingParams(max_tokens=4), params], final_stage_id=2, ) + for pool in orchestrator.stage_pools: + if pool.stage_type == "llm": + pool.select_replica_id(req_state.request_id) output = SimpleNamespace(request_id="req-3", finished=True) asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-3", previous_pool.stage_id, output, req_state)) @@ -448,6 +461,7 @@ async def put(self, item): orchestrator.num_stages = 2 orchestrator.stage_pools = [sender_pool, diffusion_pool] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) orchestrator._cfg_tracker = CfgCompanionTracker() orchestrator.output_async_queue = _AsyncQueue() orchestrator.request_states = {} @@ -484,6 +498,7 @@ def test_prewarm_diffusion_attaches_kv_sender_info(): diffusion_pool = StagePool(1, diffusion_stage) orchestrator.stage_pools = [sender_pool, diffusion_pool] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) orchestrator.num_stages = 2 req_state = OrchestratorRequestState( @@ -492,6 +507,7 @@ def test_prewarm_diffusion_attaches_kv_sender_info(): sampling_params_list=[SamplingParams(max_tokens=4), OmniDiffusionSamplingParams()], final_stage_id=1, ) + orchestrator.request_states = {req_state.request_id: req_state} stage0_request = SimpleNamespace(prompt_token_ids=[1, 2, 3]) asyncio.run(Orchestrator._prewarm_async_chunk_stages(orchestrator, "req-2", stage0_request, req_state)) @@ -503,7 +519,7 @@ def test_prewarm_diffusion_attaches_kv_sender_info(): assert req_state.stage_submit_ts[1] > 0 -def test_prewarm_submits_bound_payload_endpoint_for_concurrent_replicas(): +def test_prewarm_submits_bound_payload_endpoint_for_concurrent_replicas(mocker): orchestrator = object.__new__(Orchestrator) endpoints = {"a": {"host": "10.0.0.2", "zmq_port": 52099}, "b": {"host": "10.0.0.3", "zmq_port": 54147}} source = SimpleNamespace( @@ -523,7 +539,10 @@ async def submit(key, state, request, **kwargs): submit_initial=submit, get_bound_replica_id=lambda key: {"a": 7, "b": 1}[key], ) + source.pick = mocker.AsyncMock(side_effect=lambda key, **kwargs: source.get_bound_replica_id(key)) + target.pick = mocker.AsyncMock(side_effect=lambda key, **kwargs: target.get_bound_replica_id(key)) orchestrator.stage_pools = [source, target] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) orchestrator._stage_receives_async_chunks = lambda stage: True orchestrator._record_duplex_stage_submission = lambda *args: None orchestrator._emit_tx_edge = lambda **kwargs: None @@ -535,17 +554,22 @@ async def dispatch(callback, **kwargs): orchestrator._dispatch_or_fail_request = dispatch async def run(): + states = { + key: OrchestratorRequestState( + request_id=key, + prompt={}, + sampling_params_list=[SamplingParams(), SamplingParams()], + final_stage_id=1, + ) + for key in endpoints + } + orchestrator.request_states = states await asyncio.gather( *[ orchestrator._prewarm_async_chunk_stages( key, SimpleNamespace(prompt_token_ids=[1]), - OrchestratorRequestState( - request_id=key, - prompt={}, - sampling_params_list=[SamplingParams(), SamplingParams()], - final_stage_id=1, - ), + states[key], ) for key in endpoints ] @@ -587,6 +611,7 @@ async def test_async_pipeline_defers_sync_stage_until_audio_finishes(mocker): ), ) orchestrator.stage_pools = [audio_pool, aligner_pool] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) state = OrchestratorRequestState( request_id="aligned", prompt={"additional_information": {"text": ["Hello world"]}}, @@ -594,6 +619,7 @@ async def test_async_pipeline_defers_sync_stage_until_audio_finishes(mocker): final_stage_id=1, final_output_stage_ids={0, 1}, ) + orchestrator.request_states = {state.request_id: state} await orchestrator._prewarm_async_chunk_stages( "aligned", mocker.Mock(spec=OmniEngineCoreRequest, prompt_token_ids=[1, 2]), state diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 70dc7ca9f88..cf9a6d43d04 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -14,6 +14,7 @@ from vllm.outputs import CompletionOutput, RequestOutput from vllm.sampling_params import SamplingParams +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.engine.duplex_orchestrator import DuplexOrchestrator from vllm_omni.engine.orchestrator import ( Orchestrator, @@ -100,7 +101,7 @@ def process_inputs(self, **kwargs): class FakePrewarmPool: - """Single-replica pool double that binds requests only on submission.""" + """Single-replica pool double that binds before advertising its sender.""" stage_type = "llm" @@ -115,6 +116,10 @@ def __init__(self, role: str) -> None: ) self.submitted: list[Any] = [] + async def pick(self, request_id, *, request_state): + self._bound_request_ids.add(request_id) + return 0 + async def submit_initial(self, request_id, _req_state, request, prompt_text=None): self.submitted.append(request) self._bound_request_ids.add(request_id) @@ -221,6 +226,7 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N if payload_sender_info is not None: stage1.stage_client.get_payload_sender_info = MagicMock(return_value=payload_sender_info) orchestrator.stage_pools = [stage0, stage1, stage2] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) orchestrator._emit_tx_edge = lambda **_kwargs: None orchestrator._on_stage_submitted = MagicMock() max_frames = [256] @@ -234,6 +240,7 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(3)], final_stage_id=2, ) + orchestrator.request_states = {req_state.request_id: req_state} prewarmed = await orchestrator._prewarm_async_chunk_stages( "req-prewarm", @@ -249,7 +256,7 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N assert prewarmed is True assert stage1.submitted == [] - assert stage1.get_bound_client("req-prewarm") is None + assert stage1.get_bound_client("req-prewarm") is stage1.stage_client assert len(stage2.submitted) == 1 assert stage2.get_bound_client("req-prewarm") is stage2.stage_client assert stage2.submitted[0].payload_sender_info == payload_sender_info @@ -281,6 +288,7 @@ async def test_async_prewarm_skips_stage_with_custom_process_input_func() -> Non stage1.stage_client = SimpleNamespace(custom_process_input_func=lambda *a, **k: None) stage2 = FakePrewarmPool("receiver") orchestrator.stage_pools = [stage0, stage1, stage2] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) orchestrator._emit_tx_edge = lambda **_kwargs: None orchestrator._on_stage_submitted = MagicMock() req_state = OrchestratorRequestState( @@ -289,6 +297,7 @@ async def test_async_prewarm_skips_stage_with_custom_process_input_func() -> Non sampling_params_list=[SamplingParams(max_tokens=1) for _ in range(3)], final_stage_id=2, ) + orchestrator.request_states = {req_state.request_id: req_state} prewarmed = await orchestrator._prewarm_async_chunk_stages( "req-prewarm-custom", @@ -315,6 +324,7 @@ async def test_async_route_forwards_to_outgoing_only_stage() -> None: stage0 = SimpleNamespace(final_output=False) stage1 = FakePrewarmPool("sender") orchestrator.stage_pools = [stage0, stage1] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) orchestrator._forward_to_next_stage = AsyncMock() req_state = OrchestratorRequestState( request_id="req-route", @@ -340,6 +350,7 @@ async def test_streaming_segment_does_not_complete_final_output_stage() -> None: cleanup_parent=lambda _request_id: [], ) orchestrator.stage_pools = [SimpleNamespace(final_output=True)] + orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) orchestrator.output_async_queue = asyncio.Queue() orchestrator._cleanup_request_ids = AsyncMock() @@ -358,7 +369,7 @@ async def test_streaming_segment_does_not_complete_final_output_stage() -> None: await orchestrator._route_output(0, 0, output, req_state, None) - assert req_state.finished_final_output_stage_ids == set() + assert req_state.finished_stage_ids == set() orchestrator._cleanup_request_ids.assert_not_awaited() routed = orchestrator.output_async_queue.get_nowait() assert routed.finished is False diff --git a/tests/engine/test_terminal_stage_metrics_order.py b/tests/engine/test_terminal_stage_metrics_order.py index a64fe8cc3d8..5851a312bb8 100644 --- a/tests/engine/test_terminal_stage_metrics_order.py +++ b/tests/engine/test_terminal_stage_metrics_order.py @@ -8,6 +8,7 @@ import pytest +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -23,6 +24,7 @@ def build(): final_output=True, build_stage_metrics=lambda *a, **k: SimpleNamespace(stage_id=1, num_tokens_out=0) ), ] + o._stage_routing = StageRouting.from_transitions(len(o.stage_pools)) state = OrchestratorRequestState( request_id="r", final_stage_id=1, diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index f704d9e784c..32b0af987cc 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -106,6 +106,7 @@ def __init__( ) -> None: self.model = model self.config_path = None + self.pipeline_config = None self.stage_metadata = stage_metadata or [THREE_STAGE_META[-1]] self.num_stages = len(self.stage_metadata) self.stage_configs = [ diff --git a/tests/entrypoints/test_serve.py b/tests/entrypoints/test_serve.py index 3cf8feb877e..7f1e9246384 100644 --- a/tests/entrypoints/test_serve.py +++ b/tests/entrypoints/test_serve.py @@ -20,6 +20,7 @@ normalize_and_validate_diffusion_engine_ingress_kwargs, ) from vllm_omni.config.resolver import OmniConfigResolution +from vllm_omni.config.stage_config import PipelineConfig, StagePipelineConfig from vllm_omni.engine.arg_utils import FRONTEND_ONLY_ID_FILTER from vllm_omni.engine.stage_engine_startup import StageReplicaResources from vllm_omni.engine.stage_runtime import StageEngineLaunch @@ -406,8 +407,9 @@ def test_serve_validate_rejects_sleep_mode_with_multiple_api_servers(mocker: Moc @pytest.mark.parametrize("typed", [False, True]) @pytest.mark.parametrize("downstream_async_chunk", [False, True]) +@pytest.mark.parametrize("transitions", [None, ((0, 2),)]) def test_build_multi_api_stage_runtime_matches_current_constructor( - mocker: MockerFixture, typed: bool, downstream_async_chunk: bool + mocker: MockerFixture, typed: bool, downstream_async_chunk: bool, transitions: tuple[tuple[int, int], ...] | None ) -> None: from vllm_omni.entrypoints.cli import serve as serve_module @@ -421,14 +423,14 @@ def test_build_multi_api_stage_runtime_matches_current_constructor( ), frozenset({"model"}), ) + num_stages = 3 if transitions is not None else 2 + pipeline_stages = tuple( + StagePipelineConfig(stage_id=i, model_stage="ar", final_output=i == num_stages - 1) for i in range(num_stages) + ) if typed: from vllm_omni.config.omni_config import VllmOmniARStageConfig - from vllm_omni.config.stage_config import StagePipelineConfig - stages = [ - VllmOmniARStageConfig(stage_pipeline_config=StagePipelineConfig(stage_id=i, model_stage="ar")) - for i in range(2) - ] + stages = [VllmOmniARStageConfig(stage_pipeline_config=stage) for stage in pipeline_stages] stages[1].connector_config.async_chunk = downstream_async_chunk else: stages = [ @@ -438,7 +440,7 @@ def test_build_multi_api_stage_runtime_matches_current_constructor( async_chunk=bool(i and downstream_async_chunk), enable_sleep_mode=False ), ) - for i in range(2) + for i in range(num_stages) ] mocker.patch("vllm_omni.entrypoints.omni_base.omni_snapshot_download", return_value="dummy-model") mocker.patch( @@ -448,7 +450,11 @@ def test_build_multi_api_stage_runtime_matches_current_constructor( mocker.patch("vllm_omni.entrypoints.utils.parse_stage_overrides", return_value={}) mocker.patch( "vllm_omni.config.resolver.resolve_omni_config", - return_value=OmniConfigResolution(config_path="dummy-config", stage_configs=tuple(stages)), + return_value=OmniConfigResolution( + config_path="dummy-config", + stage_configs=tuple(stages), + pipeline_config=PipelineConfig(model_type="test", stages=pipeline_stages, stage_transitions=transitions), + ), ) runtime = serve_module._build_multi_api_stage_runtime(args, 2) @@ -456,6 +462,9 @@ def test_build_multi_api_stage_runtime_matches_current_constructor( assert runtime._client_count == 1 assert runtime._client_index == 0 assert runtime._async_chunk is downstream_async_chunk + assert (runtime._stage_routing.stage_order if runtime._stage_routing is not None else None) == ( + (0, 2) if transitions is not None else None + ) @pytest.mark.parametrize("typed", [False, True]) @@ -931,7 +940,10 @@ def _make_stage_cfg(stage_id: int, stage_type: str) -> SimpleNamespace: ) -def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: MockerFixture) -> None: +@pytest.mark.parametrize("transitions", [None, ((0, 2),)]) +def test_run_headless_llm_registers_with_auto_assigned_replica_id( + mocker: MockerFixture, transitions: tuple[tuple[int, int], ...] | None +) -> None: """LLM headless: each loop iteration registers with auto-assigned replica_id (master picks a free slot) and spawns one ``StageEngineCoreProcManager`` per local replica.""" @@ -950,7 +962,15 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker mocker.patch( "vllm_omni.config.resolver.resolve_omni_config", - return_value=_resolved(stage_cfg), + return_value=OmniConfigResolution( + config_path="/fake/stages.yaml", + stage_configs=(stage_cfg, _make_stage_cfg(1, "llm"), _make_stage_cfg(2, "llm")), + pipeline_config=PipelineConfig( + model_type="test", + stages=tuple(StagePipelineConfig(stage_id=i, model_stage="ar", final_output=i == 2) for i in range(3)), + stage_transitions=transitions, + ), + ), ) mocker.patch("vllm_omni.engine.stage_init_utils.prepare_engine_environment") mocker.patch("vllm_omni.engine.stage_init_utils.load_omni_transfer_config_for_model", return_value=None) @@ -998,6 +1018,10 @@ def test_run_headless_llm_registers_with_auto_assigned_replica_id(mocker: Mocker assert kwargs["replica_id"] is None assert "socket_ownership" not in kwargs assert mock_connector_spec.call_args.kwargs["async_chunk"] is True + stage_routing = mock_connector_spec.call_args.kwargs["stage_routing"] + assert (stage_routing.stage_order if stage_routing is not None else None) == ( + (0, 2) if transitions is not None else None + ) assert mock_manager_cls.call_count == 1 mgr_kwargs = mock_manager_cls.call_args.kwargs diff --git a/tests/worker/test_configured_connector_routing.py b/tests/worker/test_configured_connector_routing.py new file mode 100644 index 00000000000..de2f9b91719 --- /dev/null +++ b/tests/worker/test_configured_connector_routing.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Configured MRv2 payload edges over production SHM, with manually stepped I/O.""" + +import uuid +from types import SimpleNamespace + +import pytest +import torch + +from tests.worker.test_omni_connector_mixin import MixinHost, _make_model_config, _make_request +from vllm_omni.config.stage_routing import StageRouting +from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector +from vllm_omni.engine.stage_init_utils import get_stage_connector_spec + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _host(mocker, stage_id, route, async_chunk): + config = _make_model_config(async_chunk=async_chunk) + config.stage_id = stage_id + config.stage_connector_config = get_stage_connector_spec(None, stage_id, async_chunk, route) + host = MixinHost() + host.vllm_config = SimpleNamespace(model_config=config) + # Initialize all production state without starting competing I/O threads. + mocker.patch.object(host, "_create_connector", return_value=None) + host.init_omni_connectors(config) + host._omni_connector = SharedMemoryConnector({"stage_id": stage_id}) + return host + + +@pytest.mark.parametrize(("source", "target"), [(0, 2), (2, 1)]) +@pytest.mark.parametrize("async_chunk", [False, True]) +def test_mrv2_payload_uses_configured_shm_edge(mocker, source, target, async_chunk): + route = StageRouting.from_transitions(3, ((0, 2), (2, 1))) + sender = _host(mocker, source, route, async_chunk) + receiver = _host(mocker, target, route, async_chunk) + request = _make_request("r", f"configured-mrv2-{uuid.uuid4().hex}") + receiver.register_chunk_recv(request) + try: + for index in range(3 if async_chunk else 1): + payload = {"codes": {"audio": torch.tensor([index])}, "meta": {"finished": torch.tensor(False)}} + if async_chunk: + assert sender._enqueue_chunk_payload(request, payload)[0] + else: + assert sender.send_full_payload_outputs(None, {"r": (payload, request)}) == ["r"] + send_key = request.external_req_id if async_chunk else request.request_id + task = sender._pending_save_reqs[send_key].popleft() + assert sender._send_single_request(task) + assert receiver._poll_single_request("r") + assert receiver._get_req_chunk["r"] == index + 1 + expected = list(range(index + 1)) if async_chunk else [index] + assert receiver.get_local_stage_payload("r")["codes"]["audio"].tolist() == expected + receiver.get_omni_connector_output() + finally: + sender.shutdown_omni_connectors() + receiver.shutdown_omni_connectors() + + +def test_mrv2_terminal_has_no_outgoing_transport(mocker): + host = _host(mocker, 2, StageRouting.from_transitions(3, ((0, 2),)), True) + request = _make_request("r") + try: + assert host.send_full_payload_outputs(None, {"r": ({"value": 1}, request)}) == [] + assert host._enqueue_chunk_payload(request, {"value": 1}) == (True, None) + assert not host._pending_save_reqs and not host._put_req_chunk + finally: + host.shutdown_omni_connectors() diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 0d0b0f44d49..0d4a972869b 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -20,6 +20,7 @@ from vllm_omni.config.endpoint_policy import EndpointRestriction from vllm_omni.config.speech_cache import SpeechCacheConfig +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.config.yaml_util import create_config, load_yaml_config, to_dict from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler @@ -461,6 +462,9 @@ class PipelineConfig: default_deploy_config_name: str | None = None # Global CLI spelling -> (stage id, stage-local spelling). stage_cli_aliases: dict[str, tuple[int, str]] = field(default_factory=dict) + # Developer-owned control flow. None preserves numeric sequential routing; + # an explicit tuple replaces the route completely (no implicit edges). + stage_transitions: tuple[tuple[int, int], ...] | None = None def __post_init__(self) -> None: errors = self.get_validation_errors() @@ -503,6 +507,23 @@ def get_validation_errors(self) -> list[str]: # pipeline without one can never emit a result, so every request hangs. if not any(s.final_output for s in self.stages): errors.append("No terminal stage (stage with final_output=True)") + if self.stage_transitions is not None: + if any(type(stage_id) is not int for stage_id in stage_ids) or stage_ids != list(range(len(self.stages))): + errors.append("Explicit transitions require stages ordered by contiguous IDs starting at 0") + return errors + try: + route = StageRouting.from_transitions(len(self.stages), self.stage_transitions) + except ValueError as exc: + errors.append(str(exc)) + return errors + if not self.stages[route.stage_order[-1]].final_output: + errors.append("The last stage of the explicit route must declare final_output=True") + visited: set[int] = set() + for stage_id in route.stage_order: + for source in self.stages[stage_id].input_sources: + if source not in visited: + errors.append(f"Stage {stage_id} input source {source} must precede it in the explicit route") + visited.add(stage_id) return errors diff --git a/vllm_omni/config/stage_routing.py b/vllm_omni/config/stage_routing.py new file mode 100644 index 00000000000..8f4837cca6b --- /dev/null +++ b/vllm_omni/config/stage_routing.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Validated static stage transitions, independent of data dependencies.""" + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class StageRouting: + """One ordered route starting at stage 0; omitted stages stay inactive. + + Stage IDs remain indices into engine configuration and sampling arrays. + Branching, joins and loops require a separate execution contract. + """ + + stage_order: tuple[int, ...] + + @classmethod + def from_transitions( + cls, num_stages: int, transitions: tuple[tuple[int, int], ...] | None = None + ) -> "StageRouting": + if transitions is None: + return cls(tuple(range(num_stages))) + if num_stages < 1: + raise ValueError("Explicit stage transitions require stage 0") + if not isinstance(transitions, tuple) or any( + not isinstance(edge, tuple) or len(edge) != 2 for edge in transitions + ): + raise ValueError("Stage transitions must be an immutable tuple of (source, target) tuples") + successors: dict[int, int] = {} + predecessors: set[int] = set() + for source, target in transitions: + if any(type(stage_id) is not int or not 0 <= stage_id < num_stages for stage_id in (source, target)): + raise ValueError(f"Transition {source!r} -> {target!r} references an invalid stage ID") + if source == target: + raise ValueError(f"Self transition at stage {source}") + if source in successors: + raise ValueError(f"Stage {source} has duplicate or branching transitions") + if target in predecessors: + raise ValueError(f"Stage {target} has multiple incoming transitions") + successors[source] = target + predecessors.add(target) + order = [0] + while order[-1] in successors: + target = successors[order[-1]] + if target in order: + raise ValueError("Stage transitions contain a cycle") + order.append(target) + if len(order) - 1 != len(successors): + raise ValueError("Stage transitions contain edges unreachable from stage 0") + return cls(tuple(order)) + + def path_to(self, final_stage_id: int) -> tuple[int, ...]: + """Resolve a request's route prefix, rejecting unreachable endpoints.""" + if type(final_stage_id) is not int or final_stage_id not in self.stage_order: + raise ValueError(f"Final stage {final_stage_id!r} is not reachable from stage 0") + return self.stage_order[: self.stage_order.index(final_stage_id) + 1] + + def next_stage(self, stage_id: int, final_stage_id: int) -> int | None: + path = self.path_to(final_stage_id) + if stage_id not in path or stage_id == final_stage_id: + return None + return path[path.index(stage_id) + 1] diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index 71dc554463f..2655a38a240 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -475,6 +475,7 @@ def _process_chunk_queue( OmniChunkRecvHandle( request_id=request.request_id, external_req_id=request.external_req_id, + payload_sender_info=request.payload_sender_info, ) ) request.status = RequestStatus.WAITING_FOR_CHUNK diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index 7a5b910670c..43653faf2d0 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -102,9 +102,9 @@ class OmniChunkRecvHandle: """Minimal identifier carried from scheduler to runner for input-receive registration. - The runner's ``register_chunk_recv`` only consumes ``request_id`` and - ``external_req_id`` from each pending request, so we ship just those - two fields instead of the full Request object. Concrete typing + The runner's ``register_chunk_recv`` consumes the request identifiers and + the bound producer endpoint, so we ship those fields instead of the full + Request object. Concrete typing keeps msgspec serialization deterministic across IPC (default, PD-disagg, multi-node executor variants) and avoids the ``list[Any]`` fallback path. diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py index 7487b12d5cf..7f0816beca9 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py @@ -22,6 +22,7 @@ logger, should_accumulate_full_payload_output, ) +from vllm_omni.distributed.omni_connectors.utils.config import get_stage_connector_peer from vllm_omni.outputs import OmniConnectorOutput @@ -771,6 +772,8 @@ def send_full_payload_outputs( return list(outputs.keys()) sent_ids: list[str] = [] next_stage_id = self._next_stage_id + if next_stage_id is None: + return [] for req_id, value in outputs.items(): if isinstance(value, tuple) and len(value) == 2: raw_output, request = value @@ -1132,7 +1135,9 @@ def _poll_single_request(self, req_id: str, *, publish_ready: bool = True) -> bo ) return False - target_stage_id = self._stage_id - 1 + target_stage_id = get_stage_connector_peer(self._get_model_config(), self._stage_id, "from_stage") + if target_stage_id is None: + return False chunk_id = self._get_req_chunk[req_id] external_req_id = self._request_ids_mapping.get(req_id, req_id) request = self._pending_load_reqs.get(req_id) @@ -1562,6 +1567,8 @@ def _enqueue_chunk_payload( """Enqueue an already-built payload under its per-request chunk key.""" raw_req_id = getattr(request, "request_id", None) or getattr(request, "req_id", None) request_id = self._resolve_external_req_id(request, raw_req_id) + if self._next_stage_id is None: + return True, None if raw_req_id and raw_req_id != request_id: self._request_ids_mapping.setdefault(raw_req_id, request_id) chunk_id = self._put_req_chunk[request_id] diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py index 6667c0e4b1e..1cbdbd5d6e1 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py @@ -18,6 +18,7 @@ from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory from vllm_omni.distributed.omni_connectors.utils.config import ( ConnectorSpec, + get_stage_connector_peer, get_stage_connector_role, ) from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec @@ -112,7 +113,7 @@ def _stage_payload_broadcast_groups(self) -> tuple[Any, ...]: ... _async_chunk: bool _model_mode: str _stage_id: int - _next_stage_id: int + _next_stage_id: int | None _from_tp: int _to_tp: int _local_rank: int @@ -218,7 +219,7 @@ def init_omni_connectors( ) # -- next stage ID (from connector config or default stage_id + 1) -- - self._next_stage_id: int = self._resolve_next_stage_id(model_config) + self._next_stage_id = self._resolve_next_stage_id(model_config) # -- heterogeneous TP rank support -- rank_cfg = self._parse_rank_mapping(model_config) @@ -678,23 +679,9 @@ def _resolve_external_req_id(self, request: Any, fallback_req_id: str | None) -> raise ValueError("Connector request has neither an external nor an internal request ID") return fallback_req_id - def _resolve_next_stage_id(self, model_config: Any) -> int: - """Determine the downstream stage ID from connector config. - - Falls back to ``stage_id + 1`` when the config does not specify - a ``to_stage`` explicitly. - """ - connector_config = getattr(model_config, "stage_connector_config", None) - if connector_config is not None: - if isinstance(connector_config, dict): - to_stage = connector_config.get("to_stage") - else: - to_stage = getattr(connector_config, "to_stage", None) - if isinstance(to_stage, int): - return to_stage - if isinstance(to_stage, str) and to_stage.strip(): - return int(to_stage) - return self._stage_id + 1 + def _resolve_next_stage_id(self, model_config: Any) -> int | None: + """Read the downstream endpoint projected from the pipeline route.""" + return get_stage_connector_peer(model_config, self._stage_id, "to_stage") @staticmethod def _parse_rank_mapping(model_config: Any) -> dict[str, int]: diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 6654c79a30f..c5113bbc8c6 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -20,7 +20,7 @@ from ..adapter import construct_next_stage_streaming_input_prompt from ..connectors.shm_connector import SharedMemoryConnector from ..factory import OmniConnectorFactory -from ..utils.config import ConnectorSpec, stage_receives_chunks +from ..utils.config import ConnectorSpec, get_stage_connector_peer, stage_receives_chunks from ..utils.initialization import resolve_connector_spec from ..utils.kv_utils import get_local_tp_rank, get_omni_replica_id from ..utils.logging import get_connector_logger @@ -555,7 +555,9 @@ def _admit_save( def _poll_single_request(self, entry: _LoadEntry): request = entry.request stage_id = self.connector.stage_id - target_stage_id = stage_id - 1 + target_stage_id = get_stage_connector_peer(self.config, stage_id, "from_stage") + if target_stage_id is None: + return True req_id = request.request_id with self._receiver_state_lock: if self._registered_load_entries.get(req_id) is not entry: @@ -909,7 +911,9 @@ def _send_single_request_for_generation( is_finished = task["is_finished"] is_segment_finished = task["is_segment_finished"] stage_id = self.connector.stage_id - next_stage_id = stage_id + 1 + next_stage_id = get_stage_connector_peer(self.config, stage_id, "to_stage") + if next_stage_id is None: + return external_req_id = request.external_req_id chunk_id = self.put_req_chunk[external_req_id] connector_put_key = f"{external_req_id}_{stage_id}_{chunk_id}" diff --git a/vllm_omni/distributed/omni_connectors/utils/config.py b/vllm_omni/distributed/omni_connectors/utils/config.py index a646a3a3090..3e8f384c042 100644 --- a/vllm_omni/distributed/omni_connectors/utils/config.py +++ b/vllm_omni/distributed/omni_connectors/utils/config.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project from dataclasses import dataclass, field -from typing import Any +from typing import Any, Literal from .logging import get_connector_logger @@ -46,6 +46,24 @@ def stage_sends_async_output(model_config: Any) -> bool: return getattr(model_config, "stage_id", None) != 0 +def get_stage_connector_peer( + model_config: Any, stage_id: int, direction: Literal["from_stage", "to_stage"] +) -> int | None: + """Read a projected route endpoint, retaining sequential legacy defaults.""" + config = getattr(model_config, "stage_connector_config", None) + if isinstance(config, dict) and direction in config: + value = config[direction] + if value is None: + return None + else: + value = getattr(config, direction, None) + if isinstance(value, int): + return value + if isinstance(value, str) and value.strip(): + return int(value) + return stage_id + (1 if direction == "to_stage" else -1) + + @dataclass class ConnectorSpec: """Specification for a connector instance.""" diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 32fe3759d4d..6a17adaf809 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -338,6 +338,9 @@ def create_model_config(self) -> OmniModelConfig: "extra": self.stage_connector_spec.get("extra", {}).copy(), } stage_connector_config["extra"]["stage_id"] = self.stage_id + for direction in ("from_stage", "to_stage"): + if direction in self.stage_connector_spec: + stage_connector_config[direction] = self.stage_connector_spec[direction] hf_overrides = cast(dict[str, Any] | Callable[[Any], Any] | None, getattr(self, "hf_overrides", None)) # If model_arch is specified, inject it into hf_overrides so vLLM can diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index de5429905c5..01fa7bddaae 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -41,7 +41,11 @@ class AsyncOmniEngine(OmniEngineBase): """Turn-based engine used by ``Omni`` / ``AsyncOmni``.""" def _create_orchestrator(self, **orchestrator_kwargs: Any) -> OrchestratorBase: - return Orchestrator(**orchestrator_kwargs) + pipeline = self.pipeline_config + return Orchestrator( + stage_transitions=pipeline.stage_transitions if pipeline is not None else None, + **orchestrator_kwargs, + ) # ---- request helpers ---- diff --git a/vllm_omni/engine/duplex_omni_engine.py b/vllm_omni/engine/duplex_omni_engine.py index 22c02cef9cd..5d2777f41e1 100644 --- a/vllm_omni/engine/duplex_omni_engine.py +++ b/vllm_omni/engine/duplex_omni_engine.py @@ -80,6 +80,13 @@ def _validate_deployment(self) -> None: ) if not plugin_path: raise ValueError(f"{self.model!r} is not a duplex model: the pipeline declares no duplex_plugin") + transitions = getattr(pipeline_config, "stage_transitions", None) + if transitions is not None: + from vllm_omni.config.stage_routing import StageRouting + + route = StageRouting.from_transitions(len(pipeline_config.stages), transitions) + if route.stage_order != tuple(range(len(pipeline_config.stages))): + raise ValueError("Non-sequential stage transitions require a route-aware duplex plugin contract") deploy_config = self.deploy_config if deploy_config is None: raise ValueError("DuplexOmniEngine requires a deploy config with session_mode: duplex (none resolved)") diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py index 0f71e8c3e6a..598c69afbf9 100644 --- a/vllm_omni/engine/omni_engine_base.py +++ b/vllm_omni/engine/omni_engine_base.py @@ -400,6 +400,7 @@ def _initialize_stages(self, stage_init_timeout: int) -> None: single_stage_mode=self.single_stage_mode, stage_init_timeout=stage_init_timeout, async_chunk=self.async_chunk, + stage_transitions=self.pipeline_config.stage_transitions if self.pipeline_config is not None else None, tokenizer=self.tokenizer, parallel_stage_init=self._parallel_stage_init, single_stage_id_filter=self._single_stage_id_filter, @@ -909,6 +910,7 @@ def _set_pipeline_runtime_config( config_path: str | None, ) -> None: """Initialize engine-wide settings resolved from pipeline metadata.""" + self.pipeline_config = pipeline_config self.endpoint_restrictions = pipeline_config.endpoint_restrictions if pipeline_config is not None else () # No duplex_runtime_extension / duplex_serving_adapter / duplex_control # here: the pre-framework wiring they named is gone, and design rule 3 diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 267baed4f11..d43e51e9dca 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -15,6 +15,7 @@ from __future__ import annotations import asyncio +import copy import os import shutil import threading @@ -23,7 +24,7 @@ from contextlib import asynccontextmanager, nullcontext from dataclasses import dataclass, field from http import HTTPStatus -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Literal, cast import janus import torch @@ -39,6 +40,7 @@ from vllm.v1.engine.exceptions import EngineDeadError from vllm.v1.metrics.stats import IterationStats +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.core.sched.omni_scheduling_coordinator import uses_native_mrv2_data_plane from vllm_omni.data_entry_keys import ASYNC_CHUNK_PREWARM_NS, FIRST_AUDIO_KEY, FIRST_AUDIO_REQUIRED_KEY from vllm_omni.diffusion.data import is_diffusion_request_started_output @@ -67,7 +69,7 @@ ) from vllm_omni.engine.orchestrator_monitor import create_orch_monitor, replica_key from vllm_omni.engine.serialization import serialize_additional_information -from vllm_omni.engine.stage_pool import StagePool, StageUnavailableError +from vllm_omni.engine.stage_pool import RequestClosedError, StagePool, StageUnavailableError from vllm_omni.errors import DEFAULT_CLIENT_ERROR_TYPE, OmniClientError from vllm_omni.metrics import definitions as metric_defs from vllm_omni.metrics.prometheus import OmniRequestCounter @@ -246,9 +248,10 @@ class OrchestratorRequestState: # Stage the request was admitted to; > 0 when the stages before it are bypassed. entry_stage_id: int = 0 final_output_stage_ids: set[int] = field(default_factory=set) - finished_final_output_stage_ids: set[int] = field(default_factory=set) finished_stage_ids: set[int] = field(default_factory=set) pending_final_output: OutputMessage | None = None + # Failed aborts stay closed to admission until their remaining stages drain. + closing: Literal["aborting", "draining"] | None = None # Wall-clock timestamp when the client-facing engine request was accepted. request_timestamp: float = 0.0 @@ -278,6 +281,15 @@ class OrchestratorRequestState: pending_upstream_first_audio: tuple[int, OmniEngineCoreOutput] | None = None pending_first_audio_outputs: list[tuple[int, int, Any]] = field(default_factory=list) + @property + def final_outputs_finished(self) -> bool: + selected = self.final_output_stage_ids or {self.final_stage_id} + return selected <= self.finished_stage_ids + + @property + def submitted_stages_finished(self) -> bool: + return self.stage_submit_ts.keys() <= self.finished_stage_ids + @dataclass class StreamingSegmentState: @@ -353,6 +365,7 @@ def __init__( stage_pools: list[StagePool], *, async_chunk: bool = False, + stage_transitions: tuple[tuple[int, int], ...] | None = None, pd_config: dict[str, Any] | None = None, membership_controller: MembershipController | None = None, running_counter: OmniRequestCounter | None = None, @@ -371,6 +384,7 @@ def __init__( self.async_chunk = bool(async_chunk) self.num_stages = len(stage_pools) self.stage_pools: list[StagePool] = stage_pools + self._stage_routing = StageRouting.from_transitions(self.num_stages, stage_transitions) self.log_stats = log_stats if log_stats and not collect_stage_metrics: raise ValueError("log_stats=True requires collect_stage_metrics=True") @@ -682,6 +696,8 @@ async def _handle_abort(self, msg: AbortRequestMessage) -> None: logger.exception("[Orchestrator] Abort failed for request(s) %s", request_ids) if msg.rpc_id is None: raise + for request_id in request_ids: + await self.output_async_queue.put(ErrorMessage(request_id=request_id, error=error)) if msg.rpc_id is not None: # Always emit a result when rpc_id is set so CorrelatedRpcClient # waiters are unblocked even on failure. @@ -707,8 +723,14 @@ async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage if not request_ids: return [] abort_outputs: list[OutputMessage] = [] + abort_error: Exception | None = None for pool in self.stage_pools: - stage_outputs = await pool.abort_requests(request_ids) or [] + try: + stage_outputs = await pool.abort_requests(request_ids) or [] + except Exception as exc: + if abort_error is None: + abort_error = exc + continue if bool(getattr(pool, "final_output", False)) and getattr(pool, "stage_type", None) != "diffusion": final_output_type = getattr(pool.stage_client, "final_output_type", None) or "text" for orch_req_id, request_output in stage_outputs: @@ -738,6 +760,9 @@ async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage ) ) pool.release_bindings(request_ids) + if abort_error is not None: + raise abort_error + abort_outputs.sort(key=lambda output: self._stage_routing.stage_order.index(output.stage_id)) last_index_by_req: dict[str, int] = {} for index, output_msg in enumerate(abort_outputs): last_index_by_req[output_msg.request_id] = index @@ -869,6 +894,24 @@ async def _process_llm_stage_outputs( ``_finish_raw_terminal_requests`` once routing is done. """ pool = self.stage_pools[stage_id] + kept = [] + for eco in raw_outputs.outputs: + state = self.request_states.get(eco.request_id) + if state is None: + processor_state = pool.output_processor.request_states.get(eco.request_id) + if processor_state is not None: + state = self.request_states.get(processor_state.external_req_id) + if state is not None and state.closing: + # Preserve the abort snapshot, but retain terminals that arrive + # before an in-flight abort reports failure. + if eco.finish_reason is not None and not getattr(eco, "is_segment_finished", False): + state.finished_stage_ids.add(stage_id) + raw_terminal_request_ids.add(state.request_id) + else: + kept.append(eco) + if len(kept) != len(raw_outputs.outputs): + raw_outputs = copy.copy(raw_outputs) + raw_outputs.outputs = kept await self._route_upstream_first_audio(stage_id, replica_id, raw_outputs) await self._handle_kv_ready_raw_outputs(stage_id, raw_outputs) @@ -995,7 +1038,12 @@ async def _route_upstream_first_audio(self, stage_id: int, replica_id: int, raw_ req_state is None or req_state.upstream_first_audio or req_state.pending_upstream_first_audio is not None - or (stage_id if final_output else stage_id + 1) > req_state.final_stage_id + or req_state.closing + or ( + stage_id not in self._request_stage_path(req_state) + if final_output + else self._next_request_stage(stage_id, req_state) is None + ) ): continue audio = mm.get("model_outputs") @@ -1032,7 +1080,13 @@ async def _flush_upstream_first_audio(self, req_state: OrchestratorRequestState) if pending is None or self.request_states.get(req_state.request_id) is not req_state: return source_stage, first_output = pending - codec_stage = source_stage if self.stage_pools[source_stage].final_output else source_stage + 1 + codec_stage = ( + source_stage + if self.stage_pools[source_stage].final_output + else self._next_request_stage(source_stage, req_state) + ) + if codec_stage is None: + return pool = self.stage_pools[codec_stage] replica_id = pool.get_bound_replica_id(req_state.request_id) if replica_id is None or req_state.request_id not in pool.output_processor.request_states: @@ -1665,6 +1719,8 @@ async def _dispatch_or_fail_request( try: await dispatch() return True + except RequestClosedError: + return False except NativeKVHandoffError as e: await self._fail_request_client_error( req_id, @@ -1814,9 +1870,19 @@ async def _cleanup_request_ids( if cid not in batch: batch.add(cid) cleanup_ids.append(cid) + for request_id in cleanup_ids: + if (state := self.request_states.get(request_id)) is not None: + state.closing = "aborting" abort_outputs: list[OutputMessage] = [] if abort: - abort_outputs = await self._abort_request_ids(cleanup_ids) + try: + abort_outputs = await self._abort_request_ids(cleanup_ids) + except Exception: + for request_id in cleanup_ids: + if (state := self.request_states.get(request_id)) is not None: + state.closing = "draining" + await self._finish_draining_request(state) + raise self._release_request_bindings(cleanup_ids) self._release_stage_transfer_resources(cleanup_ids) for request_id in cleanup_ids: @@ -1829,6 +1895,18 @@ async def _cleanup_request_ids( req_state.running_counter_registered = False return abort_outputs + async def _finish_draining_request(self, state: OrchestratorRequestState) -> None: + pending_stages = { + pool.stage_id for pool in self.stage_pools if pool.get_bound_replica_id(state.request_id) is not None + } & state.stage_submit_ts.keys() + if not pending_stages <= state.finished_stage_ids: + return + # A failed abort retains its processor snapshot until a natural terminal. + for pool in self.stage_pools: + if pool.stage_type == "llm" and pool.get_bound_replica_id(state.request_id) is not None: + pool.output_processor.abort_requests([state.request_id], internal=False) + await self._cleanup_request_ids([state.request_id]) + async def _apply_raw_terminal_stage_finish( self, stage_id: int, @@ -1853,10 +1931,6 @@ async def _apply_raw_terminal_stage_finish( return False req_state.finished_stage_ids.add(stage_id) - final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} - if stage_id not in final_output_stage_ids: - return True - req_state.finished_final_output_stage_ids.add(stage_id) return True async def _report_duplex_session_request_error( @@ -1910,48 +1984,26 @@ async def _finish_raw_terminal_requests( ) -> None: """Finish requests whose raw terminal had no processed output.""" pool = self.stage_pools[stage_id] - for request_id in request_ids: req_state = self.request_states.get(request_id) - if req_state is None or req_state.session_owned: + if req_state is None: + continue + if req_state.closing == "draining": + await self._finish_draining_request(req_state) + if req_state.closing or req_state.session_owned: continue - - # A processed terminal output may have arrived before every submitted - # stage reported completion. Keep that real output instead of replacing - # it with the empty raw-terminal fallback below. pending = req_state.pending_final_output if pending is not None: - submitted_stage_ids = set(req_state.stage_submit_ts) - if submitted_stage_ids.issubset(req_state.finished_stage_ids): - req_state.pending_final_output = None - await self.output_async_queue.put(pending) - await self._cleanup_request_ids( - [ - request_id, - *self._cfg_tracker.cleanup_parent(request_id), - ] - ) - + if self._request_completion_ready(req_state): + await self._publish_final_output(req_state, pending) # A real terminal output is pending; do not replace it with # the swallowed-output fallback before upstream completion. continue - - # Ordinary raw terminals are client completions only when they - # come from a configured final-output stage. Non-session ERROR - # terminals are handled earlier by _handle_stage_error; session - # errors are reported separately and skipped above. - if not pool.final_output: + if not pool.final_output or stage_id != req_state.final_stage_id: continue - - final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} - if not final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids): + if stage_id not in req_state.finished_stage_ids: continue - - final_output_type = getattr( - pool.stage_client, - "final_output_type", - None, - ) + final_output_type = getattr(pool.stage_client, "final_output_type", None) logger.info( "[Orchestrator] req=%s stage-%s raw terminal produced no processed terminal; returning empty %s output", request_id, @@ -1963,23 +2015,19 @@ async def _finish_raw_terminal_requests( final_output_type=final_output_type, audio_sample_rate=pool._infer_audio_sample_rate(), ) - await self.output_async_queue.put( - OutputMessage( - request_id=request_id, - stage_id=stage_id, - replica_id=replica_id, - engine_outputs=terminal_output, - metrics=None, - finished=True, - stage_submit_ts=req_state.stage_submit_ts.get(stage_id), - ) - ) - await self._cleanup_request_ids( - [ - request_id, - *self._cfg_tracker.cleanup_parent(request_id), - ] + message = OutputMessage( + request_id=request_id, + stage_id=stage_id, + replica_id=replica_id, + engine_outputs=terminal_output, + metrics=None, + finished=True, + stage_submit_ts=req_state.stage_submit_ts.get(stage_id), ) + if self._request_completion_ready(req_state): + await self._publish_final_output(req_state, message) + else: + req_state.pending_final_output = message def _maybe_clone_diffusion_params_for_cfg(self, request_id: str, params: Any) -> Any: """Attach CFG companion ids to diffusion sampling params when needed.""" @@ -2013,10 +2061,16 @@ async def _route_output( stage_metrics: Any, ) -> None: """Route a processed output: send to frontend and/or forward.""" + segment_finished = req_state.streaming.enabled and req_state.streaming.segment(stage_id).finished + if req_state.closing: + if output.finished and not segment_finished: + req_state.finished_stage_ids.add(stage_id) + if req_state.closing == "draining": + await self._finish_draining_request(req_state) + return req_id = output.request_id finished = output.finished submit_ts = req_state.stage_submit_ts.get(stage_id) - segment_finished = req_state.streaming.enabled and req_state.streaming.segment(stage_id).finished # CFG companion: stash output so parent can bundle [parent, *companions] # into source_outputs for the bridge (e.g. thinker2imagegen). if finished and self._cfg_tracker.is_companion(req_id): @@ -2030,9 +2084,7 @@ async def _route_output( request_finished = False if finished and self.stage_pools[stage_id].final_output and not segment_finished: - req_state.finished_final_output_stage_ids.add(stage_id) - final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id} - request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids) + request_finished = stage_id == req_state.final_stage_id if req_state.session_owned: # Session-owned outputs are delivered through the session interceptor. pass @@ -2046,12 +2098,7 @@ async def _route_output( finished=request_finished, stage_submit_ts=submit_ts, ) - if ( - request_finished - and self.async_chunk - and not req_state.streaming.enabled - and not set(req_state.stage_submit_ts).issubset(req_state.finished_stage_ids) - ): + if request_finished: # Direct chunk delivery can finish the decoder before the # upstream engine's terminal output/metrics reaches this loop. # Keep request state until those messages are routed, otherwise @@ -2080,12 +2127,13 @@ async def _route_output( if await self._intercept_stage_output(stage_id, replica_id, output, req_state, stage_metrics, submit_ts): return + next_stage_id = self._next_request_stage(stage_id, req_state) if ( (finished or segment_finished) - and stage_id < req_state.final_stage_id + and next_stage_id is not None and not req_state.skip_legacy_stage_forward - and (not self.async_chunk or not self._stage_receives_async_chunks(stage_id + 1)) - and (not self._next_stage_already_submitted(stage_id, req_state) or req_state.streaming.enabled) + and (not self.async_chunk or not self._stage_receives_async_chunks(next_stage_id)) + and (next_stage_id not in req_state.stage_submit_ts or req_state.streaming.enabled) ): if ( finished @@ -2129,16 +2177,28 @@ async def _route_output( ) pending = req_state.pending_final_output - if pending is not None and set(req_state.stage_submit_ts).issubset(req_state.finished_stage_ids): - req_state.pending_final_output = None - await self.output_async_queue.put(pending) - request_finished = True + if pending is not None and not req_state.closing and self._request_completion_ready(req_state): + await self._publish_final_output(req_state, pending) + + async def _publish_final_output(self, req_state: OrchestratorRequestState, message: OutputMessage) -> None: + req_state.pending_final_output = None + await self.output_async_queue.put(message) + await self._cleanup_request_ids([req_state.request_id]) + + def _request_completion_ready(self, req_state: OrchestratorRequestState) -> bool: + return req_state.final_outputs_finished and ( + not self.async_chunk or req_state.streaming.enabled or req_state.submitted_stages_finished + ) + + def _request_stage_path(self, req_state: OrchestratorRequestState) -> tuple[int, ...]: + path = self._stage_routing.path_to(req_state.final_stage_id) + return path[path.index(req_state.entry_stage_id) :] - if request_finished and not req_state.session_owned: - await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)]) + def _next_request_stage(self, stage_id: int, req_state: OrchestratorRequestState) -> int | None: + return self._stage_routing.next_stage(stage_id, req_state.final_stage_id) def _next_stage_already_submitted(self, stage_id: int, req_state: OrchestratorRequestState) -> bool: - return (stage_id + 1) in req_state.stage_submit_ts + return self._next_request_stage(stage_id, req_state) in req_state.stage_submit_ts def _stage_receives_async_chunks(self, stage_id: int) -> bool: """Whether a stage's connector supplies its runtime inputs.""" @@ -2455,11 +2515,12 @@ async def _handle_cfg_companion_ready(self, req_id: str) -> None: return parent_state = self.request_states.get(parent_id) - if parent_state is None: + if parent_state is None or parent_state.closing: return stage_id = deferred["stage_id"] - if (stage_id + 1) in parent_state.stage_submit_ts: + next_stage_id = self._next_request_stage(stage_id, parent_state) + if next_stage_id is None or next_stage_id in parent_state.stage_submit_ts: return await self._forward_to_next_stage( @@ -2495,9 +2556,8 @@ async def _handle_kv_ready_raw_outputs( # Done-marking now happens solely on the processed-output path # in _route_output, after set_companion_output. continue - if stage_id >= req_state.final_stage_id: - continue - if (stage_id + 1) in req_state.stage_submit_ts: + next_stage_id = self._next_request_stage(stage_id, req_state) + if next_stage_id is None or next_stage_id in req_state.stage_submit_ts: continue if self._cfg_tracker.has_companions(req_id) and not self._cfg_tracker.all_companions_done(req_id): @@ -2597,18 +2657,22 @@ async def _forward_to_next_stage( is_final_update: bool = False, ) -> None: """Forward output to the next stage; a dead stage fails only this request.""" + next_stage_id = self._next_request_stage(src_stage_id, req_state) + if req_state.closing or next_stage_id is None: + return await self._dispatch_or_fail_request( lambda: self._forward_to_next_stage_unguarded( req_id, src_stage_id, output, req_state, + next_logical=next_stage_id, src_replica_id=src_replica_id, is_streaming_session=is_streaming_session, is_final_update=is_final_update, ), req_id=req_id, - stage_id=src_stage_id + 1, + stage_id=next_stage_id, operation="inter-stage forward", ) @@ -2619,12 +2683,12 @@ async def _forward_to_next_stage_unguarded( output: Any, req_state: OrchestratorRequestState, *, + next_logical: int, src_replica_id: int | None = None, is_streaming_session: bool = False, is_final_update: bool = False, ) -> None: """Forward output from the current logical stage to the next one.""" - next_logical = src_stage_id + 1 next_pool = self.stage_pools[next_logical] pinned_replica_id = self._session_route(next_logical, req_state) if pinned_replica_id is not None: @@ -2926,8 +2990,8 @@ async def _forward_to_next_stage_unguarded( final_output_type=final_output_type, audio_sample_rate=final_pool._infer_audio_sample_rate(), ) - submit_ts = _time.time() - req_state.stage_submit_ts[final_stage_id] = submit_ts + submit_ts = req_state.stage_submit_ts.get(final_stage_id) + final_replica_id = final_pool.get_bound_replica_id(req_id) logger.info( "[Orchestrator] req=%s stage-%s produced no terminal inputs for stage-%s; " "returning empty %s output from final stage-%s", @@ -2937,18 +3001,18 @@ async def _forward_to_next_stage_unguarded( final_output_type or "text", final_stage_id, ) + await self._cleanup_request_ids([req_id], abort=True) await self.output_async_queue.put( OutputMessage( request_id=req_id, stage_id=final_stage_id, - replica_id=0, + replica_id=final_replica_id, engine_outputs=terminal_output, metrics=None, finished=True, stage_submit_ts=submit_ts, ) ) - await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)]) return # Build and submit requests for each input @@ -2986,6 +3050,8 @@ async def _forward_to_next_stage_unguarded( # Rendering may have cached media before validation or # request conversion failed, without returning mm_features. await asyncio.to_thread(self._clear_stage_input_cache, next_logical) + if req_state.closing or self.request_states.get(req_id) is not req_state: + return # A forwarded input rejected by this stage fails only its request. logger.warning( "[Orchestrator] req=%s: stage-%s rejected the forwarded input: %s", req_id, next_logical, exc @@ -3050,7 +3116,10 @@ async def _prewarm_async_chunk_stages( an object the cleanup already popped from ``request_states``. """ entry_stage_id = req_state.entry_stage_id - if req_state.final_stage_id <= entry_stage_id: + if req_state.closing or self.request_states.get(request_id) is not req_state: + return False + stage_path = self._request_stage_path(req_state) + if len(stage_path) <= 1: return True prompt_token_ids = getattr(entry_request, "prompt_token_ids", None) @@ -3082,23 +3151,28 @@ async def _prewarm_async_chunk_stages( ) return False - for next_stage_id in range(entry_stage_id + 1, req_state.final_stage_id + 1): + for source_stage_id, next_stage_id in zip(stage_path, stage_path[1:]): + if req_state.closing or self.request_states.get(request_id) is not req_state: + return False next_pool = self.stage_pools[next_stage_id] params = req_state.sampling_params_list[next_stage_id] if not self._stage_receives_async_chunks(next_stage_id): - # Outgoing-only stages receive their first real input from the - # orchestrator. Pre-submitting a placeholder lets it race and - # execute before that conditioning payload arrives. continue - - for route_stage_id in (next_stage_id - 1, next_stage_id): + for route_stage_id in (source_stage_id, next_stage_id): try: pinned_replica_id = self._session_route(route_stage_id, req_state) + pool = self.stage_pools[route_stage_id] if pinned_replica_id is not None: - self.stage_pools[route_stage_id].bind_request_to_replica(request_id, pinned_replica_id) + pool.bind_request_to_replica(request_id, pinned_replica_id) + else: + await pool.pick(request_id, request_state=req_state) + except RequestClosedError: + return False except StageUnavailableError: await self._fail_request_dead_stage(request_id, route_stage_id) return False + if req_state.closing or self.request_states.get(request_id) is not req_state: + return False route_kwargs: dict[str, Any] = {} if pinned_replica_id is None else {"pinned_replica_id": pinned_replica_id} req_state.stage_submit_ts[next_stage_id] = _time.time() _t_submit_start = _time.perf_counter() @@ -3106,7 +3180,7 @@ async def _prewarm_async_chunk_stages( if next_pool.stage_type == "diffusion": submit_kwargs = self._diffusion_submit_kwargs( request_id, - next_stage_id - 1, + source_stage_id, next_pool.stage_client, req_state, ) @@ -3189,7 +3263,7 @@ async def _prewarm_async_chunk_stages( resumable=downstream_resumable, ) request.payload_sender_info = self._build_payload_sender_info( - next_stage_id - 1, request_id=request_id, req_state=req_state + source_stage_id, request_id=request_id, req_state=req_state ) request.external_req_id = request.request_id submitted = await self._dispatch_or_fail_request( @@ -3207,7 +3281,7 @@ async def _prewarm_async_chunk_stages( # Attribute the failure to the stage that failed, not stage 0; the # request is already cleaned up, so stop prewarming. - if not submitted: + if not submitted or req_state.closing or self.request_states.get(request_id) is not req_state: return False # The guard swallows submit_initial's return value; the pool binding @@ -3225,13 +3299,11 @@ async def _prewarm_async_chunk_stages( if self.request_states.get(request_id) is not req_state: return False - # async_chunk pre-submit fires per stage edge (N-1 -> N). Source - # replica is stage 0's bound replica (single-replica thinker in - # all current configs); fall back to 0 if unknown. + # Attribute transfer time to the sender on this configured edge. _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 - src_replica = self.stage_pools[next_stage_id - 1].get_bound_replica_id(request_id) + src_replica = self.stage_pools[source_stage_id].get_bound_replica_id(request_id) self._emit_tx_edge( - from_stage=next_stage_id - 1, + from_stage=source_stage_id, from_replica=src_replica if src_replica is not None else 0, to_stage=next_stage_id, to_pool=next_pool, @@ -3246,13 +3318,14 @@ def _maybe_attach_native_kv_transfer_params( req_state: OrchestratorRequestState, prompt: Any, ) -> None: - source_stage_id = 0 - if source_stage_id + 1 > req_state.final_stage_id: + source_stage_id = req_state.entry_stage_id + target_stage_id = self._next_request_stage(source_stage_id, req_state) + if target_stage_id is None: return config = getattr(self.stage_pools[source_stage_id].stage_vllm_config, "kv_transfer_config", None) if getattr(config, "kv_role", None) != "kv_producer": return - if self.stage_pools[source_stage_id + 1].stage_type != "diffusion": + if self.stage_pools[target_stage_id].stage_type != "diffusion": return from vllm_omni.diffusion.diffusion_kv.kv_connector import build_source_kv_transfer_params, mint_transfer_id @@ -3316,7 +3389,7 @@ def _build_kv_sender_info( sender_pool = self.stage_pools[sender_stage_id] sender_stage = sender_pool.get_bound_client(request_id) if request_id is not None else None - if sender_stage is None: + if request_id is None: sender_stage = sender_pool.stage_client get_sender_info = getattr(sender_stage, "get_kv_sender_info", None) if not callable(get_sender_info): @@ -3349,7 +3422,7 @@ def _build_payload_sender_info( if pinned_replica_id is not None: sender_pool.bind_request_to_replica(request_id, pinned_replica_id) sender_stage = sender_pool.get_bound_client(request_id) if request_id is not None else None - if sender_stage is None: + if request_id is None: sender_stage = sender_pool.stage_client get_sender_info = getattr(sender_stage, "get_payload_sender_info", None) if not callable(get_sender_info): @@ -3423,9 +3496,10 @@ async def _dispatch_message(self, msg: EngineQueueMessage) -> bool: return False return True - def _native_mrv2_receiver_stage(self, final_stage_id: int) -> int | None: + def _native_mrv2_receiver_stage(self, final_stage_id: int, entry_stage_id: int = 0) -> int | None: """First downstream stage up to ``final_stage_id`` that receives on MRv2's native data plane.""" - for stage_id in range(1, min(final_stage_id, len(self.stage_pools) - 1) + 1): + path = self._stage_routing.path_to(final_stage_id) + for stage_id in path[path.index(entry_stage_id) + 1 :]: vllm_config = getattr(self.stage_pools[stage_id], "stage_vllm_config", None) model_config = getattr(vllm_config, "model_config", None) if uses_native_mrv2_data_plane( @@ -3442,10 +3516,43 @@ async def _handle_add_request(self, msg: StageSubmissionMessage) -> None: prompt = msg.prompt original_prompt = msg.original_prompt sampling_params_list = msg.sampling_params_list - if not sampling_params_list: - raise ValueError(f"Missing sampling params for stage 0. Got {len(sampling_params_list)} stage params.") final_stage_id = msg.final_stage_id - final_output_stage_ids = set(msg.final_output_stage_ids or [final_stage_id]) + output_stage_ids = msg.final_output_stage_ids or [final_stage_id] + if request_id in self.request_states: + await self.output_async_queue.put( + ErrorMessage( + request_id=request_id, + stage_id=stage_id, + error="Request ID is already in use; wait for completion or abort acknowledgement before reuse", + status_code=HTTPStatus.CONFLICT.value, + ) + ) + return + try: + stage_path = self._stage_routing.path_to(final_stage_id) + if type(stage_id) is not int or stage_id not in stage_path: + raise ValueError("The entry stage must belong to the selected route") + stage_path = stage_path[stage_path.index(stage_id) :] + if any(type(sid) is not int for sid in output_stage_ids): + raise ValueError("Requested output stage IDs must be integers") + final_output_stage_ids = set(output_stage_ids) + if not final_output_stage_ids.issubset(stage_path): + raise ValueError("Requested output stages must belong to the request's configured route") + if final_stage_id not in final_output_stage_ids: + raise ValueError("The final stage must be a requested output stage") + # A distributed pool can temporarily be empty. Its registered + # pipeline metadata was checked by the frontend; do not turn + # replica unavailability into a false output-contract error. + if any( + self.stage_pools[sid].stage_client is not None and not self.stage_pools[sid].final_output + for sid in final_output_stage_ids + ): + raise ValueError("Requested output stages must declare final_output=True") + if len(sampling_params_list) <= max(stage_path): + raise ValueError("Missing sampling params for a stage in the configured route") + except ValueError as exc: + await self._fail_request_client_error(request_id, stage_id, str(exc)) + return if not self.stage_pools[stage_id].live_replica_ids(): # The entry stage lost all replicas between the HTTP-layer errored check and @@ -3455,7 +3562,7 @@ async def _handle_add_request(self, msg: StageSubmissionMessage) -> None: return if getattr(prompt, "resumable", False): - mrv2_stage_id = self._native_mrv2_receiver_stage(final_stage_id) + mrv2_stage_id = self._native_mrv2_receiver_stage(final_stage_id, stage_id) if mrv2_stage_id is not None: # Streaming-input sessions replace downstream prompts through # the V1 chunk adapter, which MRv2's native data plane lacks. @@ -3548,7 +3655,7 @@ async def _handle_add_request(self, msg: StageSubmissionMessage) -> None: await asyncio.to_thread(self._clear_stage_input_cache, stage_id) return - if self.async_chunk and final_stage_id > stage_id: + if self.async_chunk and final_stage_id != stage_id: await self._prewarm_async_chunk_stages(request_id, prompt, req_state, attach_prewarm_payload=True) async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: @@ -3557,7 +3664,7 @@ async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: request = msg.prompt final_stage_id = msg.final_stage_id req_state = self.request_states.get(request_id) - if req_state is None: + if req_state is None or req_state.closing: # Streaming updates always follow the first-chunk add_request # through the same ordered submission queue, so an unknown id # here can only mean the request already finished or was @@ -3571,6 +3678,19 @@ async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: return stage_id = req_state.entry_stage_id + try: + if type(final_stage_id) is not int or final_stage_id != req_state.final_stage_id: + raise ValueError("The final stage cannot change during a streaming request") + if msg.final_output_stage_ids and any(type(sid) is not int for sid in msg.final_output_stage_ids): + raise ValueError("Requested output stage IDs must be integers") + if msg.final_output_stage_ids and set(msg.final_output_stage_ids) != req_state.final_output_stage_ids: + raise ValueError("Output stages cannot change during a streaming request") + if msg.sampling_params_list and len(msg.sampling_params_list) <= max(self._request_stage_path(req_state)): + raise ValueError("Missing sampling params for a stage in the configured route") + except ValueError as exc: + await self._fail_request_client_error(request_id, stage_id, str(exc)) + return + if msg.sampling_params_list: req_state.sampling_params_list = msg.sampling_params_list @@ -3589,7 +3709,7 @@ async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: ): return - if self.async_chunk and final_stage_id > stage_id: + if self.async_chunk and final_stage_id != stage_id: await self._prewarm_async_chunk_stages(request_id, request, req_state) async def _handle_add_companion(self, msg: AddCompanionRequestMessage) -> None: @@ -3601,7 +3721,7 @@ async def _handle_add_companion(self, msg: AddCompanionRequestMessage) -> None: sampling_params_list = msg.sampling_params_list parent_state = self.request_states.get(parent_id) - if parent_state is None: + if parent_state is None or parent_state.closing: logger.info( "[Orchestrator] Dropping CFG companion %s (role=%s): parent %s is no longer active", companion_id, diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index e79e7e77098..0af485a516e 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -50,9 +50,11 @@ VllmOmniDiffusionStageConfig, ) from vllm_omni.config.stage_config import StageType +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.distributed.omni_connectors.utils.config import ( TRANSFER_ENGINE_CONNECTOR_NAMES, + OmniTransferConfig, ) from vllm_omni.engine.arg_utils import OmniEngineArgs from vllm_omni.entrypoints.stage_utils import _to_dict, set_stage_devices @@ -1982,9 +1984,48 @@ def load_omni_transfer_config_for_model(model: str, config_path: str | None) -> def get_stage_connector_spec( - omni_transfer_config: Any, + omni_transfer_config: OmniTransferConfig | None, stage_id: int, async_chunk: bool, + stage_routing: StageRouting | None = None, +) -> dict[str, Any]: + """Project the control route onto payload connector endpoints.""" + if stage_routing is None: + return _select_stage_connector_spec(omni_transfer_config, stage_id) + order = stage_routing.stage_order + if stage_id not in order: + return {"from_stage": None, "to_stage": None, "extra": {"role": "sender"}} + position = order.index(stage_id) + source = order[position - 1] if position else None + target = order[position + 1] if position + 1 < len(order) else None + transfer_config = omni_transfer_config + if omni_transfer_config is not None: + incoming = ( + omni_transfer_config.get_connector_for_edge(str(source), str(stage_id)) if source is not None else None + ) + outgoing = ( + omni_transfer_config.get_connector_for_edge(str(stage_id), str(target)) if target is not None else None + ) + if incoming is not None and outgoing is not None and incoming.name != outgoing.name: + raise ValueError(f"Stage {stage_id} requires the same connector backend on its incoming and outgoing edges") + edges = tuple((str(src), str(dst)) for src, dst in zip(order, order[1:])) + transfer_config = replace( + omni_transfer_config, + connectors={ + edge: backend_spec + for edge in edges + if (backend_spec := omni_transfer_config.get_connector_for_edge(*edge)) is not None + }, + ) + spec = {**_select_stage_connector_spec(transfer_config, stage_id), "from_stage": source, "to_stage": target} + if source is None: + spec["extra"] = {**spec.get("extra", {}), "role": "sender"} + return spec + + +def _select_stage_connector_spec( + omni_transfer_config: OmniTransferConfig | None, + stage_id: int, ) -> dict[str, Any]: """Return the first connector spec for a stage data-plane edge.""" from vllm_omni.distributed.omni_connectors import get_stage_connector_config diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index ca529ca0872..e806842da89 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -62,6 +62,10 @@ class StageUnavailableError(RuntimeError): """ +class RequestClosedError(RuntimeError): + """The request entered cleanup before stage admission completed.""" + + @dataclass class _ReplicaMetrics: """Per-replica metrics accumulators owned by a stage pool.""" @@ -336,6 +340,7 @@ async def pick( task: Task | None = None, *, affinity_request_id: str | None = None, + request_state: OrchestratorRequestState | None = None, ) -> int: """Return a replica id for ``request_id``. @@ -344,34 +349,41 @@ async def pick( ``request_id`` return the same replica. Bounded wait up to ``DISPATCH_WAIT_TIMEOUT_S`` when no UP replica is currently usable. - In non-distributed (legacy) mode: delegates to - :meth:`select_replica_id`. + With a request state, cleanup stops selection and a lost binding fails + the request: receivers may already hold that replica's sender address. + Without one, retain the legacy re-selection behavior. """ + if request_state is not None and request_state.closing: + raise RequestClosedError(request_id) if self._hub is None or self._lb is None: + bound = self.get_bound_replica_id(request_id) + if request_state is not None and bound is not None and not self.is_replica_available(bound): + raise StageUnavailableError(f"stage {self.stage_id} bound replica {bound} is unavailable") return self.select_replica_id(request_id, affinity_request_id=affinity_request_id) - # 1. Sticky: previously bound and still serviceable? - 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 - # Bound replica is gone or DOWN — fall through to re-select. - self._affinity.pop(request_id, None) - - # 2. Inherited affinity (CFG companion sharing a parent request_id). - 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 - - # 3. Fresh pick: poll hub + LB with bounded wait. task = task or Task(request_id=request_id) deadline = _time.monotonic() + self.DISPATCH_WAIT_TIMEOUT_S while True: + if request_state is not None and request_state.closing: + raise RequestClosedError(request_id) + # Another admission can establish affinity while this pick waits. + 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 + if request_state is not None: + raise StageUnavailableError(f"stage {self.stage_id} bound replica is no longer serviceable") + 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 + candidates = self._collect_serviceable_replicas() if candidates: # LB chose an index *into our candidates list*. @@ -1048,6 +1060,8 @@ async def submit_initial( expected_replica_id: int | None = None, ) -> int: """Submit a stage-entry request into this pool.""" + if req_state.closing: + raise RequestClosedError(request_id) params = params_override if params_override is not None else req_state.sampling_params_list[self.stage_id] # Direct engine callers may provide plain vLLM SamplingParams. if self.stage_type == "diffusion": @@ -1065,9 +1079,10 @@ async def submit_initial( if pinned_replica_id is not None: replica_id = self.bind_request_to_replica(request_id, pinned_replica_id) else: - replica_id = await self._pick_or_select( + replica_id = await self.pick( request_id, affinity_request_id=affinity_request_id, + request_state=req_state, ) client = self._diffusion_client(replica_id) await client.add_request_async(request_id, request, params, **submit_kwargs) @@ -1076,9 +1091,10 @@ async def submit_initial( if pinned_replica_id is not None: replica_id = self.bind_request_to_replica(request_id, pinned_replica_id) else: - replica_id = await self._pick_or_select( + replica_id = await self.pick( request_id, affinity_request_id=affinity_request_id, + request_state=req_state, ) # A request whose multimodal cache keys name a replica must not move: that # replica's sender-cache entries may omit media the new one never received. @@ -1146,17 +1162,15 @@ async def submit_update( ) -> int: """Submit a streaming update to an already admitted request.""" submit_kwargs = submit_kwargs or {} + if req_state.closing: + raise RequestClosedError(request_id) params = req_state.sampling_params_list[self.stage_id] if self.stage_type == "diffusion": params = OmniDiffusionSamplingParams.from_params(params) if pinned_replica_id is not None: replica_id = self.bind_request_to_replica(request_id, pinned_replica_id) else: - bound_replica_id = self.get_bound_replica_id(request_id) - if bound_replica_id is None or self.clients[bound_replica_id] is None: - replica_id = await self._pick_or_select(request_id) - else: - replica_id = bound_replica_id + replica_id = await self.pick(request_id, request_state=req_state) client = self.clients[replica_id] if client is None: @@ -1205,7 +1219,7 @@ async def submit_interaction( """Submit a midway interaction to an active diffusion (typically video generation) request.""" replica_id = self.get_bound_replica_id(request_id) if replica_id is None or self.clients[replica_id] is None: - replica_id = await self._pick_or_select(request_id) + replica_id = await self.pick(request_id) client = self._diffusion_client(replica_id) result = await client.submit_interaction_async(request_id, interaction) @@ -1215,17 +1229,6 @@ async def submit_interaction( raise ValueError(str(reason)) return replica_id - async def _pick_or_select( - self, - request_id: str, - *, - affinity_request_id: str | None = None, - ) -> int: - """Bridge to ``pick`` in distributed mode or ``select_replica_id`` legacy.""" - if self.is_distributed: - return await self.pick(request_id, affinity_request_id=affinity_request_id) - return self.select_replica_id(request_id, affinity_request_id=affinity_request_id) - # ---- Stage-local polling ---- async def read_llm_raw_output(self, client: StagePoolLLMClient) -> EngineCoreOutputs | None: @@ -1359,6 +1362,8 @@ async def abort_requests(self, request_ids: list[str]) -> list[tuple[str, Any]]: request_ids_by_replica.setdefault(replica_id, []).append(request_id) abort_outputs: list[tuple[str, Any]] = [] + aborted_request_ids: list[str] = [] + abort_error: Exception | None = None is_diffusion = self.stage_type == "diffusion" for replica_id, replica_request_ids in request_ids_by_replica.items(): # Orchestrator ids are OP external ids; EngineCore may use a @@ -1389,14 +1394,24 @@ async def abort_requests(self, request_ids: list[str]) -> list[tuple[str, Any]]: engine_abort_ids = list(aborted) client = self.clients[replica_id] if client is not None: - await client.abort_requests_async(engine_abort_ids) + try: + await client.abort_requests_async(engine_abort_ids) + except Exception as exc: + logger.warning("[StagePool-%s] abort failed on replica %s: %r", self.stage_id, replica_id, exc) + if abort_error is None: + abort_error = exc + continue # An evicted engine cannot produce further output. Release its # retained processor state without sending another abort RPC. if not is_diffusion and self._output_processor is not None: commit = getattr(self._output_processor, "commit_aborted_request_state", None) if callable(commit): commit(replica_request_ids, internal=False) + aborted_request_ids.extend(replica_request_ids) + if abort_error is not None: + self.release_bindings(aborted_request_ids) + raise abort_error return abort_outputs def schedule_release_request_resources(self, request_ids: list[str]) -> None: diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index b5cd739776a..0e4509da48c 100644 --- a/vllm_omni/engine/stage_runtime.py +++ b/vllm_omni/engine/stage_runtime.py @@ -20,6 +20,7 @@ from vllm.logger import init_logger from vllm_omni.config.omni_config import BaseVllmOmniStageConfig, VllmOmniDiffusionStageConfig +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.distributed.omni_connectors.utils.initialization import ( resolve_omni_kv_config_for_stage, ) @@ -171,12 +172,18 @@ def __init__( *, stage_init_timeout: int, async_chunk: bool, + stage_transitions: tuple[tuple[int, int], ...] | None = None, tokenizer: str | None = None, parallel_stage_init: bool = False, log_stats: bool = False, client_config: OmniClientConfig | None = None, ) -> None: self._stage_configs = stage_configs + self._stage_routing = ( + StageRouting.from_transitions(len(stage_configs), stage_transitions) + if stage_transitions is not None + else None + ) self._model = model self._config_path = config_path self._stage_init_timeout = stage_init_timeout @@ -815,6 +822,7 @@ def _build_logical_stage_init_plans( omni_transfer_config=omni_transfer_config, stage_id=stage_id, async_chunk=self._async_chunk, + stage_routing=self._stage_routing, ) omni_kv_connector = resolve_omni_kv_config_for_stage(omni_transfer_config, stage_id) num_replicas = replicas_per_stage[stage_idx] @@ -1414,6 +1422,7 @@ def __init__( omni_lb_policy: str = "random", request_queue: janus.Queue[EngineQueueMessage] | None = None, parallel_stage_init: bool = False, + stage_transitions: tuple[tuple[int, int], ...] | None = None, ) -> None: super().__init__( stage_configs=stage_configs, @@ -1421,6 +1430,7 @@ def __init__( config_path=config_path, stage_init_timeout=stage_init_timeout, async_chunk=async_chunk, + stage_transitions=stage_transitions, tokenizer=tokenizer, parallel_stage_init=parallel_stage_init, log_stats=log_stats, @@ -1743,6 +1753,7 @@ def create_stage_runtime( single_stage_mode: bool, stage_init_timeout: int, async_chunk: bool, + stage_transitions: tuple[tuple[int, int], ...] | None = None, tokenizer: str | None = None, parallel_stage_init: bool = False, # Distributed-only params: @@ -1766,6 +1777,7 @@ def create_stage_runtime( config_path=config_path, stage_init_timeout=stage_init_timeout, async_chunk=async_chunk, + stage_transitions=stage_transitions, tokenizer=tokenizer, parallel_stage_init=parallel_stage_init, log_stats=log_stats, @@ -1783,6 +1795,7 @@ def create_stage_runtime( config_path=config_path, stage_init_timeout=stage_init_timeout, async_chunk=async_chunk, + stage_transitions=stage_transitions, tokenizer=tokenizer, parallel_stage_init=parallel_stage_init, log_stats=log_stats, diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 7300953c74f..81d0f8e3531 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -25,6 +25,7 @@ from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG from vllm.logger import init_logger +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.diffusion.registry import resolve_native_single_file from vllm_omni.engine.arg_utils import FRONTEND_ONLY_ID_FILTER, is_frontend_only from vllm_omni.entrypoints.api_server_shared_dir import ( @@ -1243,6 +1244,7 @@ def _build_multi_api_stage_runtime(args: TrackingNamespace, num_api_servers: int stage_init_timeout=int(getattr(args, "stage_init_timeout", 300)), parallel_stage_init=bool(getattr(args, "parallel_stage_init", False)), async_chunk=async_chunk, + stage_transitions=resolved.pipeline_config.stage_transitions if resolved.pipeline_config is not None else None, tokenizer=getattr(args, "tokenizer", None), log_stats=not bool(getattr(args, "disable_log_stats", False)), ) @@ -1495,6 +1497,11 @@ def run_headless(args: TrackingNamespace) -> None: stage_connector_spec = get_stage_connector_spec( omni_transfer_config=omni_transfer_config, stage_id=stage_id, + stage_routing=( + StageRouting.from_transitions(len(stage_configs), resolved.pipeline_config.stage_transitions) + if resolved.pipeline_config is not None and resolved.pipeline_config.stage_transitions is not None + else None + ), async_chunk=bool( getattr(getattr(stage_cfg, "connector_config", None), "async_chunk", None) if hasattr(stage_cfg, "connector_config") diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index dff2a2d0758..4145b3da858 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -19,6 +19,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError from vllm_omni.config.stage_config import merge_sampling_constraints +from vllm_omni.config.stage_routing import StageRouting from vllm_omni.engine.messages import ( EngineQueueMessage, ErrorMessage, @@ -458,11 +459,18 @@ def _log_summary_and_cleanup(self, request_id: str, reason: str = "stage_error") # is still in self.request_states) is corrected after the pop. self._publish_request_gauges(len(self.request_states)) + def _configured_stage_routing(self) -> StageRouting: + pipeline = self.engine.pipeline_config + return StageRouting.from_transitions( + len(self._stage_meta_list), pipeline.stage_transitions if pipeline is not None else None + ) + def _compute_final_stage_id(self, output_modalities: list[str] | None) -> int: return get_final_stage_id_for_e2e( output_modalities, self.output_modalities, self._stage_meta_list, + stage_order=self._configured_stage_routing().stage_order, ) def _compute_final_output_stage_ids(self, output_modalities: list[str] | None) -> list[int]: @@ -472,8 +480,9 @@ def _compute_final_output_stage_ids(self, output_modalities: list[str] | None) - requested_modalities = self.output_modalities return [ sid - for sid, stage in enumerate(self._stage_meta_list) - if getattr(stage, "final_output", False) and stage.final_output_type in requested_modalities + for sid in self._configured_stage_routing().stage_order + if self._stage_meta_list[sid].final_output + and self._stage_meta_list[sid].final_output_type in requested_modalities ] def _process_stage_metrics_message(self, msg: StageMetricsMessage) -> None: diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index c091ae5e8c9..2bc538802ba 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -147,7 +147,10 @@ def prepare_stage_config_inputs( def get_final_stage_id_for_e2e( - output_modalities: list[str] | None, default_modalities: list[str], stage_list: list + output_modalities: list[str] | None, + default_modalities: list[str], + stage_list: list, + stage_order: tuple[int, ...] | None = None, ) -> int: """Get the final stage id for e2e. @@ -157,8 +160,14 @@ def get_final_stage_id_for_e2e( Returns: Final stage id for e2e """ - last_stage_id = len(stage_list) - 1 + order = stage_order if stage_order is not None else tuple(range(len(stage_list))) + last_stage_id = order[-1] if output_modalities is not None: + active_outputs = {stage_list[sid].final_output_type for sid in order if stage_list[sid].final_output} + configured_outputs = {stage.final_output_type for stage in stage_list if stage.final_output} + unreachable = (set(output_modalities) & configured_outputs) - active_outputs + if unreachable: + raise ValueError(f"Requested output modalities are unreachable on the stage route: {sorted(unreachable)}") prompt_modalities = [] for modality in output_modalities: if modality not in default_modalities: @@ -170,25 +179,10 @@ def get_final_stage_id_for_e2e( else: output_modalities = default_modalities - try: - final_stage_id_for_e2e = last_stage_id - for _sid in range(last_stage_id, -1, -1): - if ( - getattr(stage_list[_sid], "final_output", False) - and stage_list[_sid].final_output_type in output_modalities - ): - final_stage_id_for_e2e = _sid - break - except Exception as e: - logger.debug( - "[Orchestrator] Failed to determine final stage for E2E; \ - falling back to last: %s", - e, - exc_info=True, - ) - final_stage_id_for_e2e = last_stage_id - - return final_stage_id_for_e2e + for sid in reversed(order): + if stage_list[sid].final_output and stage_list[sid].final_output_type in output_modalities: + return sid + return last_stage_id def filter_dataclass_kwargs(cls: Any, kwargs: dict) -> dict: diff --git a/vllm_omni/utils/forced_aligner.py b/vllm_omni/utils/forced_aligner.py index e6dcecb57e1..a45bbcc9a34 100644 --- a/vllm_omni/utils/forced_aligner.py +++ b/vllm_omni/utils/forced_aligner.py @@ -101,6 +101,7 @@ def inject_forced_aligner_stage( StageExecutionType, StagePipelineConfig, ) + from vllm_omni.config.stage_routing import StageRouting from vllm_omni.model_executor.stage_input_processors.forced_aligner import ( POOLING_OUTPUT_DECODER_PATH, TIMESTAMPS_MODALITY, @@ -115,13 +116,17 @@ def inject_forced_aligner_stage( _FA = "vllm_omni.model_executor.stage_input_processors.forced_aligner" new_id = len(pipeline.stages) + source_id = StageRouting.from_transitions(new_id, pipeline.stage_transitions).stage_order[-1] + transitions = pipeline.stage_transitions + if transitions is not None: + transitions += ((source_id, new_id),) aligner_ps = StagePipelineConfig( stage_id=new_id, model_stage="forced_aligner", # Pooling stage = LLM_AR stage run with runner="pooling" (vLLM derives # is_pooling_model from runner_type). execution_type=StageExecutionType.LLM_AR, - input_sources=(new_id - 1,), + input_sources=(source_id,), final_output=True, final_output_type=TIMESTAMPS_MODALITY, owns_tokenizer=True, @@ -129,7 +134,7 @@ def inject_forced_aligner_stage( model_arch=fa.architecture, custom_process_input_func=f"{_FA}.code2wav2aligner", ) - extended = replace(pipeline, stages=pipeline.stages + (aligner_ps,)) + extended = replace(pipeline, stages=pipeline.stages + (aligner_ps,), stage_transitions=transitions) engine_extras: dict[str, Any] = { "model": fa.model,