From 926857899a8a84af56a16c8ebee5e43221c6924d Mon Sep 17 00:00:00 2001 From: Pinoeer-kingxi <13022943007@163.com> Date: Sat, 10 Oct 2026 08:54:53 -0700 Subject: [PATCH 1/9] [Core] Route stages through configured pipeline transitions Signed-off-by: Pinoeer-kingxi <13022943007@163.com> --- docs/configuration/README.md | 2 + docs/configuration/stage_routing.md | 114 +++ tests/config/test_forced_aligner_injection.py | 48 ++ tests/config/test_stage_routing.py | 181 +++++ .../sched/test_omni_scheduling_coordinator.py | 2 + .../test_chunk_transfer_adapter.py | 92 ++- tests/engine/test_arg_utils.py | 27 +- .../test_async_omni_engine_stage_init.py | 17 +- tests/engine/test_cfg_companion_lifecycle.py | 12 +- tests/engine/test_configured_stage_routing.py | 690 ++++++++++++++++++ tests/engine/test_duplex_omni_engine.py | 14 + tests/engine/test_orchestrator.py | 156 +++- .../test_orchestrator_async_chunk_prewarm.py | 21 +- .../test_orchestrator_error_handling.py | 6 +- tests/engine/test_orchestrator_first_audio.py | 39 +- .../test_orchestrator_kv_sender_info.py | 36 +- .../test_orchestrator_stage_input_bridge.py | 21 +- tests/engine/test_stage_engine_core_proc.py | 33 + .../engine/test_stage_pool_collective_rpc.py | 23 +- .../test_terminal_stage_metrics_order.py | 2 + tests/entrypoints/test_serve.py | 44 +- .../test_configured_connector_routing.py | 68 ++ tests/worker/test_omni_connector_mixin.py | 19 + tests/worker_v2/test_omni_data_plane.py | 103 +++ vllm_omni/config/stage_config.py | 21 + vllm_omni/config/stage_routing.py | 63 ++ .../core/sched/omni_scheduling_coordinator.py | 1 + vllm_omni/core/sched/output.py | 6 +- .../omni_connector_payload_transport.py | 30 +- .../model_runner/omni_connector_runtime.py | 26 +- .../chunk_transfer_adapter.py | 97 ++- .../omni_connectors/utils/config.py | 23 +- vllm_omni/engine/arg_utils.py | 3 + vllm_omni/engine/duplex_omni_engine.py | 7 + vllm_omni/engine/duplex_orchestrator.py | 2 +- vllm_omni/engine/errors.py | 4 + vllm_omni/engine/omni_engine_base.py | 223 +----- vllm_omni/engine/orchestrator.py | 547 +++++++++----- vllm_omni/engine/stage_engine_core_proc.py | 19 +- vllm_omni/engine/stage_init_utils.py | 43 +- vllm_omni/engine/stage_pool.py | 77 +- vllm_omni/engine/stage_runtime.py | 13 + vllm_omni/entrypoints/cli/serve.py | 7 + vllm_omni/entrypoints/omni_base.py | 13 +- vllm_omni/entrypoints/utils.py | 36 +- vllm_omni/utils/forced_aligner.py | 9 +- vllm_omni/worker/mixins.py | 8 + vllm_omni/worker_v2/omni_data_plane.py | 35 +- 48 files changed, 2511 insertions(+), 572 deletions(-) create mode 100644 docs/configuration/stage_routing.md create mode 100644 tests/config/test_stage_routing.py create mode 100644 tests/engine/test_configured_stage_routing.py create mode 100644 tests/worker/test_configured_connector_routing.py create mode 100644 vllm_omni/config/stage_routing.py 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..340fbd44c0f --- /dev/null +++ b/docs/configuration/stage_routing.md @@ -0,0 +1,114 @@ +# 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`. +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 closes admission before waiting for aborts. For async-chunk transport +and native MRv2 payloads, it also waits for each backend's cleanup contract +across all live replicas before terminal success or successful abort acknowledgement. +Release markers use the existing transport queue after +pending puts; native workers also apply their existing cancellation fence. +V1 uses the scheduler's chunk adapter as its transport owner, while native MRv2 +uses the worker data plane. +Legacy V1 full-payload and native KV transfer use their existing cleanup paths. +SHM implements per-prefix release. A backend exposing `wait_for_cleanup` also +confirms its release before the operation succeeds. Other backends retain their +existing per-payload lifecycle after the same queued publication fence; their +acknowledgement does not promise physical release. NIXL ownership changes are +reviewed separately from configured routing. +Native stages without a payload connector have no transport owner to release. + +Concurrent cleanup calls share the same operation; cancelling a duplicate waiter +does not cancel its owner. A parent joins a companion's running cleanup without +replacing its operation. CFG relationships remain available until cleanup +succeeds, so retrying a failed parent also releases its companions. Replica waits +cannot submit a closed request. A failed +release reports a request-scoped error and retains closed request state for a +cleanup retry. That request ID remains unavailable for reuse until cleanup +succeeds. Chunk cancellation fences and per-request chunk counters remain owned +by the existing transport. + +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/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 f0b8d45b59a..b3bc90be4f5 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -21,6 +21,7 @@ from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct from vllm_omni.distributed.omni_connectors.adapter import construct_next_stage_streaming_input_prompt +from vllm_omni.distributed.omni_connectors.connectors.base import OmniConnectorBase from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( @@ -31,6 +32,63 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_pipeline_release_preserves_legacy_backend_and_fences_queued_puts(build_adapter, mocker): + adapter, _ = build_adapter(stage_id=0) + connector = mocker.Mock(spec=OmniConnectorBase) + connector.stage_id = 0 + connector.put.return_value = (True, 1, {}) + adapter.connector = connector + adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct() + request = _req("internal", RequestStatus.RUNNING, external_req_id="external") + adapter.save_async(None, request) + completion = adapter.release_request_resources("external") + assert not completion.done() + connector.put.assert_not_called() + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + connector.put.assert_called_once() + assert not completion.done() + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + completion.result(timeout=0) + assert not adapter._pending_save_reqs + connector.cleanup.assert_not_called() + + +@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() + + @pytest.fixture def shm_sender(build_adapter): adapter, _ = build_adapter(stage_id=0) @@ -97,6 +155,36 @@ def blocking_put(**kwargs): assert not adapter.code_prompt_token_ids +def test_connector_release_is_queued_after_payload_put(build_adapter): + adapter, connector = build_adapter(stage_id=0) + events = [] + + def put(**kwargs): + events.append("put") + return True, 1, {} + + def cleanup_prefix(prefix): + events.append(prefix) + return 1 + + connector.put.side_effect = put + connector.cleanup_prefix.side_effect = cleanup_prefix + adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct( + codes=CodesStruct(audio=torch.tensor([1])) + ) + request = _req("internal", RequestStatus.RUNNING, external_req_id="external") + adapter.request_ids_mapping[request.request_id] = request.external_req_id + adapter.save_async(None, request) + completion = adapter.release_request_resources(request.request_id) + assert completion is not None and not completion.done() + assert events == [] + while adapter._pending_save_reqs: + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + assert events == ["put", "external_0_"] + completion.result(timeout=0) + assert adapter._reclaimed_payload_total == 1 + + def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender): adapter, connector = shm_sender ext_id = f"release_{uuid.uuid4().hex}" @@ -109,13 +197,13 @@ def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender): # The consumer read only the first chunk before its own request ended. assert connector.get("0", "1", f"{ext_id}_0_0") is not None - adapter.release_shm_resources(ext_id) + adapter.release_request_resources(ext_id) # The release is queued behind sends; nothing is unlinked on the caller. assert f"{ext_id}_0_1" in connector._pending_keys while adapter._pending_save_reqs: adapter._send_single_request(adapter._pending_save_reqs.popleft()) - assert adapter._reclaimed_shm_total == 2 + assert adapter._reclaimed_payload_total == 2 assert not any(k.startswith(f"{ext_id}_0_") for k in connector._pending_keys) assert connector.get("0", "1", f"{ext_id}_0_1") is None assert connector.get("0", "1", sibling_key)[0] == "sibling" 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_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index befa22857f7..d17ff6f7e28 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -17,7 +17,8 @@ from vllm.v1.engine.utils import EngineZmqAddresses 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 PipelineConfig, StagePipelineConfig 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 @@ -313,11 +314,19 @@ 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),)]) +def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch, transitions): 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, + ) + 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 @@ -348,12 +357,14 @@ def _capture_create_stage_runtime(**kwargs): assert captured["stage_init_timeout"] == 7 assert captured["log_stats"] is True assert captured["client_config"] is engine._client_config + assert captured["stage_transitions"] == transitions 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 7fa04d1912c..27556f1da99 100644 --- a/tests/engine/test_cfg_companion_lifecycle.py +++ b/tests/engine/test_cfg_companion_lifecycle.py @@ -19,9 +19,10 @@ import pytest +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 +104,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,12 +115,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, request_id: str = "p") -> OrchestratorRequestState: + return OrchestratorRequestState( + request_id=request_id, 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, prompt={"prompt": "x"}, pipeline_timings={}, @@ -131,7 +133,7 @@ async def test_kv_ready_does_not_mark_companion_done(): output (set_companion_output) may. Guards the original race.""" orch = _make_orchestrator() orch._cfg_tracker.register_companion("p", "neg", "p__neg") - orch.request_states["p__neg"] = _req_state() + orch.request_states["p__neg"] = _req_state(request_id="p__neg") raw = SimpleNamespace(request_id="p__neg", kv_transfer_params={"kv_ready": True}) await orch._handle_kv_ready_raw_outputs(0, SimpleNamespace(outputs=[raw])) diff --git a/tests/engine/test_configured_stage_routing.py b/tests/engine/test_configured_stage_routing.py new file mode 100644 index 00000000000..5431bb3e9ab --- /dev/null +++ b/tests/engine/test_configured_stage_routing.py @@ -0,0 +1,690 @@ +# 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 +from types import SimpleNamespace + +import pytest +from vllm.sampling_params import SamplingParams +from vllm.v1.engine import EngineCoreOutputs, FinishReason + +from tests.engine.test_orchestrator import ( + FakeOutputProcessor, + FakePromptRequest, + FakeStageClient, + _build_request_output, + _build_stage_pools, + _engine_core_outputs, + _enqueue_add_request, + _get_output_message, + _shutdown_orchestrator, + _wait_for, + orchestrator_factory, +) +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.errors import ResourceReleaseError +from vllm_omni.engine.messages import ErrorMessage, StageSubmissionMessage +from vllm_omni.engine.omni_engine_base import OmniEngineBase +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState +from vllm_omni.engine.stage_pool import RequestClosedError, StageUnavailableError +from vllm_omni.entrypoints.omni_base import OmniBase +from vllm_omni.entrypoints.utils import get_final_stage_id_for_e2e + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +configured_orchestrator_factory = orchestrator_factory + + +@pytest.mark.parametrize("selected", [set(), {0, 1}]) +def test_completion_derives_selected_outputs_from_stage_termination(selected): + state = OrchestratorRequestState( + request_id="r", final_stage_id=1, final_output_stage_ids=selected, finished_stage_ids={2} + ) + assert not state.final_outputs_finished + state.finished_stage_ids.add(0) + assert not state.final_outputs_finished + state.finished_stage_ids.add(1) + assert state.final_outputs_finished + + +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 = [ + 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 +async def test_cancel_while_waiting_for_replica_stops_old_request_admission(mocker): + 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_duplicate_cleanup_waiter_does_not_cancel_abort_owner(mocker): + orch, state, clients = _direct_orchestrator(async_chunk=True) + orch.stage_pools[0].select_replica_id("r") + started, finish = asyncio.Event(), asyncio.Event() + + async def abort(_ids): + started.set() + await finish.wait() + + abort_call = mocker.patch.object(clients[0], "abort_requests_async", side_effect=abort) + owner = asyncio.create_task(orch._cleanup_request_ids(["r"], abort=True)) + await asyncio.wait_for(started.wait(), timeout=1) + waiter = asyncio.create_task(orch._cleanup_request_ids(["r"], abort=True)) + await asyncio.sleep(0) + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + assert state.closing and not state.cleanup_future.done() + with pytest.raises(RequestClosedError): + await orch.stage_pools[0].submit_update("r", state, SimpleNamespace()) + await orch._handle_add_request(_submission()) + error = orch.output_async_queue.get_nowait() + assert isinstance(error, ErrorMessage) and error.status_code == 409 + assert orch.request_states["r"] is state + finish.set() + await asyncio.wait_for(owner, timeout=1) + assert abort_call.await_count == 1 + assert "r" not in orch.request_states + + +@pytest.mark.asyncio +@pytest.mark.parametrize("own_release_fails", [False, True]) +async def test_automatic_cleanup_joins_failed_owner_and_reports_only_failed_requests(mocker, own_release_fails): + orch, state, _ = _direct_orchestrator() + other_state = OrchestratorRequestState(request_id="s") + orch.request_states["s"] = other_state + entered, second_entered, finish = asyncio.Event(), asyncio.Event(), asyncio.Event() + failures = {"r", "s"} if own_release_fails else {"r"} + + async def release(ids): + if ids == ["r"]: + entered.set() + await finish.wait() + else: + second_entered.set() + if failures.intersection(ids): + raise ResourceReleaseError("READ ownership remains active") + + mocker.patch.object(orch.stage_pools[0], "release_request_resources", side_effect=release) + owner = asyncio.create_task(orch._cleanup_request_ids(["r"], abort=True)) + await asyncio.wait_for(entered.wait(), timeout=1) + automatic = asyncio.create_task(orch._cleanup_request_ids(["r", "s"])) + await asyncio.wait_for(second_entered.wait(), timeout=1) + await asyncio.sleep(0) + assert not automatic.done() + finish.set() + with pytest.raises(ResourceReleaseError): + await asyncio.wait_for(owner, timeout=1) + await asyncio.wait_for(automatic, timeout=1) + errors = [] + while not orch.output_async_queue.empty(): + errors.append(orch.output_async_queue.get_nowait()) + assert {error.request_id for error in errors} == failures + assert all(isinstance(error, ErrorMessage) and not error.fatal for error in errors) + assert state.closing and orch.request_states["r"] is state + assert ("s" in orch.request_states) is own_release_fails + failures.clear() + await orch._cleanup_request_ids(["r", "s"], abort=True) + assert not orch.request_states + + +@pytest.mark.asyncio +@pytest.mark.parametrize("release_fails", [False, True]) +async def test_parent_cleanup_joins_completed_companion_cleanup(mocker, release_fails): + orch, state, _ = _direct_orchestrator() + companion_id = "r__neg" + companion = OrchestratorRequestState(request_id=companion_id) + orch.request_states[companion_id] = companion + orch._cfg_tracker.register_companion("r", "neg", companion_id) + orch._cfg_tracker.set_companion_output(companion_id, {"conditioning": 1}) + orch._cfg_tracker.on_companion_completed(companion_id) + entered, parent_entered, finish = asyncio.Event(), asyncio.Event(), asyncio.Event() + + async def release(ids): + if ids == [companion_id]: + entered.set() + await finish.wait() + if release_fails: + raise ResourceReleaseError("READ ownership remains active") + else: + parent_entered.set() + + release_call = mocker.patch.object(orch.stage_pools[0], "release_request_resources", side_effect=release) + owner = asyncio.create_task(orch._cleanup_request_ids([companion_id])) + await asyncio.wait_for(entered.wait(), timeout=1) + companion_future = companion.cleanup_future + parent = asyncio.create_task(orch._cleanup_request_ids(["r"], abort=True)) + try: + await asyncio.wait_for(parent_entered.wait(), timeout=1) + assert companion.cleanup_future is companion_future + assert orch.request_states["r"] is state and state.closing + assert not parent.done() + finally: + finish.set() + results = await asyncio.wait_for(asyncio.gather(owner, parent, return_exceptions=True), timeout=1) + assert [call.args[0] for call in release_call.await_args_list] == [[companion_id], ["r"]] + assert results[0] == [] + if release_fails: + assert isinstance(results[1], ResourceReleaseError) + assert orch.request_states["r"] is state and state.closing + assert companion.closing and orch._cfg_tracker.has_companions("r") + release_fails = False + await orch._cleanup_request_ids(["r"], abort=True) + assert set(release_call.await_args_list[-1].args[0]) == {"r", companion_id} + else: + assert results[1] == [] + assert not orch.request_states + + +@pytest.mark.asyncio +@pytest.mark.parametrize("retry_id", ["r", "r__neg"]) +async def test_failed_cfg_cleanup_retry_releases_the_complete_owner_batch(mocker, retry_id): + orch, state, _ = _direct_orchestrator() + companion_id = "r__neg" + companion = OrchestratorRequestState(request_id=companion_id) + orch.request_states[companion_id] = companion + orch._cfg_tracker.register_companion("r", "neg", companion_id) + release_call = mocker.patch.object( + orch.stage_pools[0], + "release_request_resources", + side_effect=ResourceReleaseError("READ ownership remains active"), + ) + with pytest.raises(ResourceReleaseError): + await orch._cleanup_request_ids(["r"], abort=True) + assert state.closing and companion.closing + assert state.cleanup_future is companion.cleanup_future + assert orch._cfg_tracker.has_companions("r") + release_call.side_effect = None + await orch._cleanup_request_ids([retry_id], abort=True) + assert not orch.request_states + assert all(set(call.args[0]) == {"r", companion_id} for call in release_call.await_args_list) + + +@pytest.mark.asyncio +async def test_failed_abort_can_be_retried_without_reopening_admission(mocker): + orch, state, clients = _direct_orchestrator() + orch.stage_pools[0].select_replica_id("r") + abort = mocker.patch.object(clients[0], "abort_requests_async", side_effect=[RuntimeError("RPC unavailable"), None]) + with pytest.raises(RuntimeError, match="RPC unavailable"): + await orch._cleanup_request_ids(["r"], abort=True) + assert state.closing and orch.request_states["r"] is state + with pytest.raises(RequestClosedError): + await orch.stage_pools[0].pick("r", request_state=state) + await orch._cleanup_request_ids(["r"], abort=True) + assert abort.await_count == 2 + assert "r" not in orch.request_states + + +@pytest.mark.asyncio +async def test_lost_preselected_replica_cannot_be_reselected(mocker): + 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) + with pytest.raises(StageUnavailableError): + await pool.submit_initial("r", state, SimpleNamespace()) + await orch._handle_dead_replica(2, 0, RuntimeError("replica exited")) + 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("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 +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("empty_inputs", [None, []]) +async def test_empty_diffusion_cleanup_failure_keeps_request_closed(mocker, empty_inputs): + orch, state, clients = _direct_orchestrator() + clients[2].stage_type = "diffusion" + clients[2].custom_process_input_func = mocker.Mock(return_value=empty_inputs) + mocker.patch.object(orch, "_release_stage_transfer_resources", side_effect=ResourceReleaseError("release failed")) + + await orch._forward_to_next_stage("r", 0, _build_request_output("r"), state) + + assert orch.output_async_queue.get_nowait().finished + error = orch.output_async_queue.get_nowait() + assert isinstance(error, ErrorMessage) and not error.fatal and error.request_id == "r" + assert orch.request_states["r"] is state and state.closing + assert not clients[2].add_request_calls + assert orch.rpc_async_queue.empty() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event_driven", ["0", "1"]) +@pytest.mark.parametrize("transitions", [((0, 2),), ((0, 2), (2, 1))]) +async def test_configured_route_runs_through_real_loop( + monkeypatch, configured_orchestrator_factory, event_driven, transitions +): + monkeypatch.setenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", event_driven) + 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") + 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_cancelled_prewarm_releases_bindings_without_resurrection(): + orch, state, clients = _direct_orchestrator(async_chunk=True) + await orch._prewarm_async_chunk_stages( + "r", FakePromptRequest(request_id="r", prompt_token_ids=[1, 2], resumable=False), state + ) + await orch._cleanup_request_ids(["r"], abort=True) + assert "r" not in orch.request_states + assert all(pool.get_bound_replica_id("r") is None for pool in orch.stage_pools) + assert not await orch._prewarm_async_chunk_stages( + "r", FakePromptRequest(request_id="r", prompt_token_ids=[1, 2], resumable=False), state + ) + assert len(clients[2].add_request_calls) == 1 + + +@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) diff --git a/tests/engine/test_duplex_omni_engine.py b/tests/engine/test_duplex_omni_engine.py index 90297273556..577801a26e5 100644 --- a/tests/engine/test_duplex_omni_engine.py +++ b/tests/engine/test_duplex_omni_engine.py @@ -12,6 +12,7 @@ import pytest +from vllm_omni.config.stage_config import 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 @@ -29,6 +30,19 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +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() + + class _FakeRpcClient: def __init__(self, result: Any) -> None: self.result = result diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 4a5bc435c8b..7ede13ce5d0 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -21,6 +21,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 ( @@ -85,6 +86,36 @@ async def fail_outputs() -> None: assert "stage engine died" in fatal.error +@pytest.mark.asyncio +async def test_shutdown_cancels_membership_watcher_and_closes_stages(mocker) -> None: + watcher = asyncio.create_task(asyncio.Event().wait()) + membership = mocker.Mock() + membership.start.return_value = watcher + membership.drain_tasks = mocker.AsyncMock() + requests: asyncio.Queue = asyncio.Queue() + requests.put_nowait(ShutdownRequestMessage()) + orchestrator = Orchestrator( + request_async_queue=requests, + output_async_queue=asyncio.Queue(), + rpc_async_queue=asyncio.Queue(), + stage_pools=[], + membership_controller=membership, + ) + close_stages = mocker.spy(orchestrator, "_shutdown_stages") + # run() owns its loop's tasks; keep the guard on this same task. + task = asyncio.current_task() + assert task is not None + guard = asyncio.get_running_loop().call_later(2, task.cancel) + try: + await orchestrator.run() + finally: + guard.cancel() + assert watcher.cancelled() + membership.drain_tasks.assert_awaited_once() + membership.shutdown.assert_called_once() + close_stages.assert_called_once() + + @dataclass class OrchestratorFixture: orchestrator: Orchestrator @@ -132,10 +163,14 @@ def __init__( self.abort_calls: list[list[str]] = [] self.collective_rpc_calls: list[tuple[str, float | None, tuple[Any, ...], dict[str, Any]]] = [] self.shutdown_calls = 0 - self._engine_core_outputs = queue.Queue() - self._diffusion_outputs = queue.Queue() + self.sample_rate: int | None = None + self._engine_core_outputs: queue.Queue[Any] = queue.Queue() + self._diffusion_outputs: queue.Queue[Any] = queue.Queue() # Orchestrator-facing interface. + async def call_utility_async(self, method, request_ids): + assert method == "omni_release_request_resources" + async def add_request_async(self, *args, **kwargs) -> None: self.add_request_calls.append(args) @@ -228,7 +263,7 @@ async def collective_rpc_async( class FakeOutputProcessor: - def __init__(self, *, request_outputs: list[object] | None = None) -> None: + def __init__(self, *, request_outputs: list[Any] | None = None) -> None: self.request_outputs = list(request_outputs or []) self.add_request_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] self.abort_calls: list[list[str]] = [] @@ -359,7 +394,7 @@ def _build_stage_pools( stage_clients: list[list[FakeStageClient]], *, output_processors: list[FakeOutputProcessor] | None = None, - stage_vllm_configs: list[object] | None = None, + stage_vllm_configs: list[Any] | None = None, ) -> list[StagePool]: """Build StagePool list from per-stage replica lists. @@ -389,13 +424,14 @@ def _build_stage_pools( def _build_harness( - stage_clients: list[object], + stage_clients: list[Any], *, - output_processors: list[object] | None = None, + output_processors: list[Any] | None = None, stage_vllm_configs: list[object] | None = None, 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. @@ -430,6 +466,7 @@ async def _run() -> None: rpc_async_queue=rpc_queue.async_q, stage_pools=stage_pools, async_chunk=async_chunk, + stage_transitions=stage_transitions, log_stats=log_stats, ) ready_future.set_result((orchestrator, request_queue, output_queue, rpc_queue)) @@ -1053,6 +1090,103 @@ async def test_run_abort_emits_result_when_rpc_id_set(orchestrator_factory) -> N await _shutdown_orchestrator(orchestrator_fixture) +@pytest.mark.asyncio +@pytest.mark.parametrize("release_fails", [False, True]) +async def test_abort_ack_waits_for_resource_release_and_failed_cleanup_can_retry( + orchestrator_factory, monkeypatch, release_fails +): + fixture = orchestrator_factory([FakeStageClient(final_output=True)]) + obj = fixture.orchestrator + entered = threading.Event() + released = threading.Event() + calls = [] + + async def release(ids): + calls.append(ids) + entered.set() + while not released.is_set(): + await asyncio.sleep(0.001) + if release_fails and len(calls) == 1: + raise TimeoutError("READ ownership remains active") + + monkeypatch.setattr(obj.stage_pools[0], "release_request_resources", release) + try: + await _enqueue_add_request( + fixture, + request_id="owned", + prompt=FakePromptRequest(request_id="owned", prompt_token_ids=[1], resumable=False), + original_prompt={"prompt": "ownership"}, + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + await _wait_for(lambda: len(obj.stage_pools[0].clients[0].add_request_calls) == 1) + fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["owned"], rpc_id="first")) + await _wait_for(entered.is_set) + assert fixture.queues[2].sync_q.empty() + assert obj.request_states["owned"].closing + released.set() + result = await _get_rpc_message(fixture) + assert result.success is (not release_fails) + if release_fails: + assert "READ ownership" in result.error + assert obj.request_states["owned"].closing + fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["owned"], rpc_id="retry")) + result = await _get_rpc_message(fixture) + assert result.rpc_id == "retry" and result.success + assert calls == [["owned"], ["owned"]] + assert "owned" not in obj.request_states + finally: + released.set() + await _shutdown_orchestrator(fixture) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cleanup_trigger", ["completion", "uncorrelated_abort"]) +async def test_automatic_release_failure_is_request_scoped_and_does_not_publish_success( + orchestrator_factory, monkeypatch, cleanup_trigger +): + stage = FakeStageClient(final_output=True) + fixture = orchestrator_factory( + [stage], + output_processors=[FakeOutputProcessor(request_outputs=[_build_request_output("owned")])], + stage_vllm_configs=[SimpleNamespace(model_config=SimpleNamespace(async_chunk=True))], + ) + obj = fixture.orchestrator + original = stage.call_utility_async + + async def fail_release(method, ids): + raise TimeoutError("READ ownership remains active") + + monkeypatch.setattr(stage, "call_utility_async", fail_release) + try: + await _enqueue_add_request( + fixture, + request_id="owned", + prompt=FakePromptRequest(request_id="owned", prompt_token_ids=[1], resumable=False), + original_prompt={"prompt": "ownership"}, + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + await _wait_for(lambda: len(stage.add_request_calls) == 1) + if cleanup_trigger == "completion": + stage.push_engine_core_outputs(_engine_core_outputs("owned", 1.0)) + else: + fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["owned"])) + await _wait_for(lambda: not fixture.output_sync_q.empty()) + error = fixture.output_sync_q.get_nowait() + assert isinstance(error, ErrorMessage) + assert error.request_id == "owned" and not error.fatal + assert fixture.queues[1].sync_q.empty() + assert obj.request_states["owned"].closing + assert fixture.thread.is_alive() + monkeypatch.setattr(stage, "call_utility_async", original) + fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["owned"], rpc_id="retry")) + result = await _get_rpc_message(fixture) + assert result.success and "owned" not in obj.request_states + finally: + await _shutdown_orchestrator(fixture) + + @pytest.mark.asyncio async def test_abort_marks_only_last_final_stage_output_finished() -> None: """A request with two final AR stages must keep earlier abort outputs unfinished.""" @@ -1394,6 +1528,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: @@ -1476,6 +1611,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}) @@ -1540,6 +1676,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( @@ -1584,6 +1721,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): @@ -1769,8 +1907,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]], @@ -1985,6 +2123,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"]) @@ -2014,6 +2153,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 76a25f15a43..2e190b69fbe 100644 --- a/tests/engine/test_orchestrator_async_chunk_prewarm.py +++ b/tests/engine/test_orchestrator_async_chunk_prewarm.py @@ -10,6 +10,7 @@ from __future__ import annotations +import asyncio import time from types import SimpleNamespace from typing import Any @@ -50,6 +51,8 @@ class _FakePool: """Single-replica LLM pool that records the requests it is sent.""" stage_type = "llm" + final_output = True + num_replicas = 1 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 +64,12 @@ def __init__(self, role: str, stage_client: Any = None) -> None: def live_replica_ids(self) -> list[int]: return [0] + def available_replica_ids(self) -> list[int]: + return [0] + + async def pick(self, request_id, *, request_state=None): + return 0 + async def submit_initial(self, request_id, _req_state, request, prompt_text=None): self.submitted.append(request) return 0 @@ -79,9 +88,14 @@ def _orchestrator(payload_func: Any = None, cls: type[Orchestrator] = Orchestrat """Thinker -> Talker (no payload hook) -> Code2Wav (``payload_func``, if any).""" 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.async_chunk = True - orchestrator.request_states = {} + Orchestrator.__init__( + orchestrator, + request_async_queue=asyncio.Queue(), + output_async_queue=asyncio.Queue(), + rpc_async_queue=asyncio.Queue(), + stage_pools=[_FakePool("sender"), _FakePool("receiver"), _FakePool("receiver", code2wav_client)], + async_chunk=True, + ) orchestrator._emit_tx_edge = lambda **_kwargs: None return orchestrator @@ -194,6 +208,7 @@ async def test_prewarm_submits_plain_placeholder(case: dict[str, Any], expected_ final_stage_id=2, session_owned=case.get("session_owned", False), ) + orchestrator.request_states["req"] = req_state streaming = case.get("streaming", False) req_state.streaming.enabled = streaming stage0_kwargs = case.get("stage0", {"resumable": False}) diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index 7e1635d191b..e74aaedf479 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -573,7 +573,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 @@ -1182,7 +1182,7 @@ async def test_diffusion_client_error_output_propagates_non_400_status( @pytest.mark.asyncio -async def test_control_rpc_failure_is_reported_to_the_caller_not_fatal() -> None: +async def test_control_rpc_failure_is_reported_to_the_caller_not_fatal(mocker) -> None: """A control RPC that fails on one replica must come back as that replica's error result instead of raising out of _request_handler. """ @@ -1192,7 +1192,7 @@ async def collective_rpc_async(self, *args, **kwargs): raise TimeoutError("outputs did not drain") ok = FakeStageClient(stage_type="diffusion") - ok.collective_rpc_async = AsyncMock(return_value=None) + mocker.patch.object(ok, "collective_rpc_async", new_callable=AsyncMock, return_value=None) orchestrator, queues = _build_bare_orchestrator( _build_stage_pools([[ok], [_FailingClient(stage_type="diffusion")]]) ) diff --git a/tests/engine/test_orchestrator_first_audio.py b/tests/engine/test_orchestrator_first_audio.py index bc99cf7fbef..5437341c3c2 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,10 @@ def register(self): def get_bound_replica_id(self, request_id): return self.bound_replica + async def pick(self, request_id, *, request_state=None): + self.bound_replica = self.replica_id + return self.replica_id + async def process_llm_raw_outputs(self, replica_id, raw, **kwargs): assert replica_id == self.replica_id return self.output_processor.process_outputs(raw.outputs, raw.timestamp).request_outputs @@ -65,6 +70,7 @@ async def process_llm_raw_outputs(self, replica_id, raw, **kwargs): def _orchestrator(output_kind=RequestOutputKind.DELTA, *, registered=True, replica_id=2, final_stage_id=1): obj = Orchestrator.__new__(Orchestrator) obj.stage_pools = [SimpleNamespace(final_output=False), _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() @@ -72,6 +78,7 @@ def _orchestrator(output_kind=RequestOutputKind.DELTA, *, registered=True, repli obj._running_counter = None obj._abort_request_ids = AsyncMock(return_value=[]) obj._release_request_bindings = Mock() + obj._release_stage_transfer_resources = AsyncMock() obj._finish_raw_terminal_requests = AsyncMock() async def process(stage, replica, raw, terminals): @@ -132,12 +139,34 @@ def _audio(output): return torch.cat(audio) if isinstance(audio, list) else audio +@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 + + @pytest.mark.asyncio @pytest.mark.parametrize("kind", list(RequestOutputKind)) @pytest.mark.parametrize("suffix_overtook", [False, True]) 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: @@ -164,6 +193,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) @@ -175,6 +205,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)) @@ -238,6 +269,10 @@ async def test_first_audio_waits_for_codec_registration_and_ignores_duplicates() @pytest.mark.parametrize("processing_failed", [False, True]) async def test_prewarm_flushes_first_audio_after_successful_submission(mocker, processing_failed): obj = _orchestrator(registered=False, final_stage_id=2 if processing_failed else 1) + if processing_failed: + later_pool = SimpleNamespace(submit_initial=AsyncMock()) + obj.stage_pools.append(later_pool) + obj._stage_routing = StageRouting.from_transitions(len(obj.stage_pools)) req_state = obj.request_states["r"] req_state.sampling_params_list = [None] * (req_state.final_stage_id + 1) await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11))) @@ -260,8 +295,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()) - 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) @@ -314,6 +347,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) diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index 5f3d1147145..a0d5df12b4a 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -5,11 +5,13 @@ from collections.abc import Callable from types import SimpleNamespace from typing import Any +from unittest.mock import AsyncMock import pytest 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 +357,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() @@ -366,6 +369,7 @@ def test_forward_to_diffusion_attaches_kv_sender_info(): ) output = SimpleNamespace(request_id="req-1", finished=True) + sender_pool.select_replica_id("req-1") asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-1", sender_pool.stage_id, output, req_state)) assert diffusion_stage.calls[0]["request_id"] == "req-1" @@ -384,6 +388,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", @@ -392,6 +397,7 @@ def test_diffusion_resubmission_preserves_sender_endpoints(): final_stage_id=1, ) output = SimpleNamespace(request_id="streaming", finished=True) + sender_pool.select_replica_id("streaming") async def forward_twice(): await orchestrator._forward_to_next_stage("streaming", 0, output, state) @@ -414,6 +420,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() @@ -425,6 +432,8 @@ def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info(): ) output = SimpleNamespace(request_id="req-3", finished=True) + source_pool.select_replica_id("req-3") + previous_pool.select_replica_id("req-3") asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-3", previous_pool.stage_id, output, req_state)) assert diffusion_stage.calls[0]["kv_sender_info"] == { @@ -448,6 +457,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 +494,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( @@ -494,6 +505,8 @@ def test_prewarm_diffusion_attaches_kv_sender_info(): ) stage0_request = SimpleNamespace(prompt_token_ids=[1, 2, 3]) + orchestrator.request_states = {req_state.request_id: req_state} + sender_pool.select_replica_id(req_state.request_id) asyncio.run(Orchestrator._prewarm_async_chunk_stages(orchestrator, "req-2", stage0_request, req_state)) assert diffusion_stage.calls[0]["request_id"] == "req-2" @@ -522,8 +535,10 @@ async def submit(key, state, request, **kwargs): stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(hf_config=None, max_model_len=64)), submit_initial=submit, get_bound_replica_id=lambda key: {"a": 7, "b": 1}[key], + pick=AsyncMock(side_effect=lambda key, **kwargs: {"a": 7, "b": 1}[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 +550,21 @@ async def dispatch(callback, **kwargs): orchestrator._dispatch_or_fail_request = dispatch async def run(): + orchestrator.request_states = { + key: OrchestratorRequestState( + request_id=key, + prompt={}, + sampling_params_list=[SamplingParams(), SamplingParams()], + final_stage_id=1, + ) + for key in endpoints + } 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, - ), + orchestrator.request_states[key], ) for key in endpoints ] @@ -587,6 +606,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"]}}, @@ -595,6 +615,10 @@ async def test_async_pipeline_defers_sync_stage_until_audio_finishes(mocker): final_output_stage_ids={0, 1}, ) + orchestrator.request_states = {state.request_id: state} + orchestrator._cleanup_request_ids.side_effect = lambda ids: [ + orchestrator.request_states.pop(request_id, None) for request_id in ids + ] 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 3c1dd6d2be3..8f014dc38e5 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, @@ -46,7 +47,7 @@ def __init__( self.next_inputs = list(next_inputs or []) self.add_request_calls: list[tuple[Any, ...]] = [] self.decoded_source_tokens: str | None = None - self._engine_core_outputs = queue.Queue() + self._engine_core_outputs: queue.Queue[Any] = queue.Queue() async def add_request_async(self, *args, **_kwargs) -> None: self.add_request_calls.append(args) @@ -99,7 +100,7 @@ def process_inputs(self, **kwargs): class FakePrewarmPool: - """Single-replica pool double that binds requests only on submission.""" + """Single-replica pool double that binds on selection or submission.""" stage_type = "llm" @@ -119,6 +120,10 @@ async def submit_initial(self, request_id, _req_state, request, prompt_text=None self._bound_request_ids.add(request_id) return 0 + async def pick(self, request_id, *, request_state=None): + self._bound_request_ids.add(request_id) + return 0 + def get_bound_replica_id(self, request_id): return 0 if request_id in self._bound_request_ids else None @@ -215,6 +220,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() req_state = OrchestratorRequestState( @@ -224,6 +230,8 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N final_stage_id=2, ) + orchestrator.request_states = {req_state.request_id: req_state} + prewarmed = await orchestrator._prewarm_async_chunk_stages( "req-prewarm", SimpleNamespace(prompt_token_ids=[1, 2], resumable=True), @@ -232,7 +240,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 @@ -263,6 +271,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( @@ -272,6 +281,8 @@ async def test_async_prewarm_skips_stage_with_custom_process_input_func() -> Non final_stage_id=2, ) + orchestrator.request_states = {req_state.request_id: req_state} + prewarmed = await orchestrator._prewarm_async_chunk_stages( "req-prewarm-custom", SimpleNamespace(prompt_token_ids=[1, 2], resumable=True), @@ -297,6 +308,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", @@ -322,6 +334,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() @@ -340,7 +353,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_stage_engine_core_proc.py b/tests/engine/test_stage_engine_core_proc.py index fdb6424817d..3ed3fa0e163 100644 --- a/tests/engine/test_stage_engine_core_proc.py +++ b/tests/engine/test_stage_engine_core_proc.py @@ -47,6 +47,39 @@ def test_decoder_capability_matches_sink_binding(mocker, backend, executor_class assert runner.model_state.set_first_audio_sink.call_count == int(expected) +@pytest.mark.parametrize("has_plane", [False, True]) +def test_native_resource_release_reaches_worker_owned_plane(mocker, has_plane): + from vllm_omni.worker.mixins import OmniWorkerMixin + + plane = mocker.Mock() if has_plane else None + worker = SimpleNamespace(model_runner=SimpleNamespace(_omni_data_plane=plane)) + engine = StageEngineCoreProc.__new__(StageEngineCoreProc) + engine.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(use_v2_model_runner=True, supports_native_mrv2_data_plane=True) + ) + engine.model_executor = mocker.Mock() + engine.model_executor.collective_rpc.side_effect = lambda method, args: getattr(OmniWorkerMixin, method)( + worker, *args + ) + engine.omni_release_request_resources(["external"]) + engine.model_executor.collective_rpc.assert_called_once_with("omni_release_request_resources", args=(["external"],)) + if has_plane: + plane.release_request_resources.assert_called_once_with(["external"]) + + +def test_legacy_resource_release_reaches_scheduler_adapter(mocker): + engine = StageEngineCoreProc.__new__(StageEngineCoreProc) + engine.vllm_config = SimpleNamespace(model_config=SimpleNamespace(use_v2_model_runner=False)) + adapter = mocker.Mock() + engine.scheduler = SimpleNamespace(chunk_transfer_adapter=adapter) + engine.omni_release_request_resources(["external"]) + adapter.release_request_resources.assert_called_once() + args, kwargs = adapter.release_request_resources.call_args + assert args == ("external",) + assert kwargs["deadline"] > 0 + adapter.release_request_resources.return_value.result.assert_called_once() + + @pytest.mark.parametrize( "first_decoder,stream_decoder,stream_first_audio,expected", [(True, False, False, True), (False, True, False, False), (False, True, True, True), (False, False, True, False)], diff --git a/tests/engine/test_stage_pool_collective_rpc.py b/tests/engine/test_stage_pool_collective_rpc.py index aa7bbcb397c..3e2e8f59d5a 100644 --- a/tests/engine/test_stage_pool_collective_rpc.py +++ b/tests/engine/test_stage_pool_collective_rpc.py @@ -191,6 +191,25 @@ async def run() -> None: asyncio.run(run()) +@pytest.mark.cpu +def test_release_request_resources_includes_native_full_payload(): + async def run() -> None: + call = AsyncMock() + pool = StagePool( + 0, + [SimpleNamespace(call_utility_async=call)], # type: ignore[list-item] + stage_vllm_config=SimpleNamespace( + model_config=SimpleNamespace( + async_chunk=False, use_v2_model_runner=True, supports_native_mrv2_data_plane=True + ) + ), + ) + await pool.release_request_resources(["req-1"]) + call.assert_awaited_once_with("omni_release_request_resources", ["req-1"]) + + asyncio.run(run()) + + @pytest.mark.cpu def test_release_request_resources_times_out_hung_replica_without_blocking_others(monkeypatch): async def run() -> None: @@ -210,7 +229,9 @@ async def hang(*_args): live.assert_awaited_once_with("omni_release_request_resources", ["req-1"]) assert not release.done() - await asyncio.wait_for(release, timeout=1.0) + with pytest.raises(RuntimeError, match="replica 0 resource release failed") as error: + await asyncio.wait_for(release, timeout=1.0) + assert isinstance(error.value.__cause__, TimeoutError) asyncio.run(run()) 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_serve.py b/tests/entrypoints/test_serve.py index 48382d98bcd..129a82b32b9 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.stage_engine_startup import StageReplicaResources from vllm_omni.engine.stage_runtime import StageEngineLaunch from vllm_omni.entrypoints.cli.serve import ( @@ -327,8 +328,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 @@ -342,14 +344,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 = [ @@ -359,7 +361,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( @@ -369,7 +371,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) @@ -377,6 +383,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]) @@ -852,7 +861,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.""" @@ -871,7 +883,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) @@ -919,6 +939,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/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index ea042d4e846..f4dd9f11138 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -79,12 +79,31 @@ def _make_model_config( ) -> SimpleNamespace: return SimpleNamespace( stage_connector_config=None, + use_v2_model_runner=True, + supports_native_mrv2_data_plane=True, async_chunk=async_chunk, worker_type=worker_type, custom_process_next_stage_input_func=custom_func, ) +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("role", ["sender", "receiver"]) +def test_async_transport_has_one_owner_for_runner_and_scheduler(mocker, native, role): + config = _make_model_config(async_chunk=True, custom_func="stage.payload") + config.stage_connector_config = {"role": role} + config.use_v2_model_runner = native + host = MixinHost() + create = mocker.patch.object(host, "_create_connector", return_value=None) + mocker.patch.object(host, "_load_custom_func", return_value=(None, None)) + host.init_omni_connectors(config) + try: + assert create.call_count == int(native) + assert host._omni_connector is None + finally: + host.shutdown_omni_connectors() + + def _make_request(req_id: str, external_req_id: str | None = None): r = SimpleNamespace( request_id=req_id, diff --git a/tests/worker_v2/test_omni_data_plane.py b/tests/worker_v2/test_omni_data_plane.py index 62898a456e2..2a97f081e08 100644 --- a/tests/worker_v2/test_omni_data_plane.py +++ b/tests/worker_v2/test_omni_data_plane.py @@ -47,6 +47,7 @@ def _bare_plane(*, delivery_timeout_s=1.0, shutdown_timeout_s=1.0): _native_output_error=None, _native_output_error_lock=threading.Lock(), _native_output_closed=False, + model_config=SimpleNamespace(stage_connector_config={"from_stage": None}), _native_outputs_in_flight=defaultdict(int), _native_terminal_pending=set(), _async_chunk=True, @@ -260,6 +261,108 @@ def _stop_save_thread(plane): assert not plane._save_thread.is_alive() +def test_pipeline_release_preserves_legacy_backend_and_retires_request(plane): + connector = SimpleNamespace(cleanup=lambda request_id: None) + plane.register_request(_new_request("external", "external")) + _start_save_thread(plane, connector) + try: + plane.release_request_resources(["external"]) + assert "external" not in plane._native_requests + assert not plane._pending_save_reqs + assert not plane._pending_save_counts + finally: + _stop_save_thread(plane) + + +def test_pipeline_release_without_a_transport_owner(plane): + plane._omni_connector = None + plane.release_request_resources(["external"]) + + +def test_pipeline_release_fences_scheduler_terminal_arriving_later(plane): + plane._omni_connector = SimpleNamespace(cleanup_prefix=lambda prefix: 0) + plane._stage_id = 0 + plane._lock = threading.Lock() + plane._pending_save_reqs = {} + plane._pending_save_counts = defaultdict(int) + plane._save_work_available = threading.Event() + plane.register_request(_new_request("external", "external")) + plane.reserve_outputs(["external"]) + + _start_save_thread(plane, plane._omni_connector) + try: + plane.release_request_resources(["external"]) + finally: + _stop_save_thread(plane) + + assert not plane._native_requests + assert plane.record.cleaned == ["external"] + assert len(plane.record.batches) == 1 + assert plane.record.batches[0][0][0].is_finished() + assert plane.abort_requests({"external"}) == 0 + assert _complete(plane, [{"codes.audio": "late"}], req="external") == 0 + assert len(plane.record.batches) == 1 + assert not plane._pending_save_reqs + assert not plane._pending_save_counts + + +def test_pipeline_release_waits_for_queued_payloads_and_uses_external_prefix(raw_plane): + connector = _gated_connector() + reclaimed = threading.Event() + prefixes = [] + + def cleanup_prefix(prefix): + prefixes.append(prefix) + reclaimed.set() + + connector.cleanup_prefix = cleanup_prefix + _start_save_thread(raw_plane, connector) + try: + request = SimpleNamespace(request_id="external") + raw_plane._enqueue_chunk_payload(request, {"c": 1}, wait_for_delivery=False) + raw_plane._enqueue_chunk_payload(request, {"c": 2}, wait_for_delivery=False) + assert connector.put_started.wait(timeout=1) + with ThreadPoolExecutor(max_workers=1) as executor: + release = executor.submit(raw_plane.release_request_resources, ["external"]) + assert not reclaimed.is_set() + assert not release.done() + connector.release.set() + release.result(timeout=2) + assert reclaimed.is_set() + assert connector.put_keys == ["external_0_0", "external_0_1"] + assert prefixes == ["external_0_"] + finally: + connector.release.set() + _stop_save_thread(raw_plane) + assert not raw_plane._pending_save_reqs + assert not raw_plane._pending_save_counts + + +@pytest.mark.parametrize("failure_method", ["cleanup_prefix", "wait_for_cleanup"]) +def test_pipeline_release_failure_is_reported_and_retry_is_request_scoped(raw_plane, failure_method): + connector = _gated_connector() + connector.cleanup_prefix = lambda prefix: None + connector.wait_for_cleanup = lambda prefix, **kwargs: None + + def fail(*args, **kwargs): + raise TimeoutError("READ claim remains active") + + original = getattr(connector, failure_method) + setattr(connector, failure_method, fail) + _start_save_thread(raw_plane, connector) + try: + with pytest.raises(TimeoutError, match="READ claim remains active"): + raw_plane.release_request_resources(["external"]) + assert not raw_plane._pending_save_reqs + assert not raw_plane._pending_save_counts + assert not raw_plane._delivery_manager.is_quarantined + setattr(connector, failure_method, original) + raw_plane.release_request_resources(["external"]) + assert not raw_plane._pending_save_counts + finally: + _stop_save_thread(raw_plane) + + class _Output: copy_event = None diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 2b59f4ed153..e304e065a50 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -19,6 +19,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 @@ -345,6 +346,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() @@ -387,6 +391,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 9ad4ec5a9fc..48f2a9f5182 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -457,6 +457,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 60195e7d339..54f704d5946 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -99,9 +99,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..ce322e73cc1 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) @@ -1380,6 +1385,27 @@ def _send_single_request(self, task: dict) -> bool: return True request_id = task["request_id"] + if prefix := task.get("release_prefix"): + completion = task["release_completion"] + release_error = None + try: + cleanup = getattr(connector, "cleanup_prefix", None) + if callable(cleanup): + cleanup(prefix) + # Legacy backends keep their existing per-payload lifecycle. + # Completing this marker still fences the local send queue. + wait = getattr(connector, "wait_for_cleanup", None) + if callable(wait): + wait(prefix, receive_prefix=task["receive_prefix"], deadline=task["release_deadline"]) + except Exception as error: + release_error = error + finally: + self._decrement_pending_save_count(request_id) + if release_error is not None: + completion.set_exception(release_error) + else: + completion.set_result(None) + return True payload_data = task.get("data") if payload_data is None and task.get("request") is not None: payload_data = self._build_custom_process_payload( @@ -1562,6 +1588,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..46547be743c 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 @@ -15,9 +15,11 @@ import torch from vllm.logger import init_logger +from vllm_omni.core.sched.omni_scheduling_coordinator import uses_native_mrv2_data_plane 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 @@ -51,6 +53,12 @@ def _should_create_payload_connector(model_config: Any) -> bool: Sender edges may instead be owned solely by KV transfer. Receivers still need a connector even though they do not declare a downstream payload hook. """ + if getattr(model_config, "async_chunk", False) and not uses_native_mrv2_data_plane( + model_config, use_v2_model_runner=getattr(model_config, "use_v2_model_runner", False) + ): + # The scheduler's chunk adapter owns V1 transport. Creating another + # runner connector duplicates its endpoint and payload ownership. + return False if get_stage_connector_role(model_config) != "sender": return True @@ -112,7 +120,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 +226,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 +686,13 @@ 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: + def _resolve_next_stage_id(self, model_config: Any) -> int | None: """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 + 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 5e4ef6820ba..1319d6fc4a1 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 @@ -8,7 +8,8 @@ import time from collections import defaultdict, deque from collections.abc import Callable, Iterable, Mapping -from typing import Any +from concurrent.futures import Future +from typing import Any, cast import torch from vllm.v1.metrics.stats import PrefillStats @@ -20,7 +21,12 @@ 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 ( + REQUEST_RESOURCE_RELEASE_TIMEOUT_S, + 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 @@ -137,6 +143,10 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): work and return immediately. """ + # Backends expose transport-specific capabilities in addition to the + # shared connector interface; this subclass always creates one. + connector: Any + def __init__(self, vllm_config: Any): model_config = vllm_config.model_config # The base constructor starts the save thread, so sender-generation @@ -195,17 +205,17 @@ def __init__(self, vllm_config: Any): self._adaptive_states: dict[str, Any] = {} self.upstream_exhausted_requests: set[str] = set() self.segment_finished_requests: set[str] = set() - self.request_payload = {} + self.request_payload: dict[str, Any] = {} self.code_prompt_token_ids: dict[str, list[torch.Tensor]] = defaultdict(list) self.request_ids_mapping: dict[str, str] = {} - # Save-thread only: running count of segments reclaimed after finish. - self._reclaimed_shm_total = 0 + # Save-thread only: running count of payloads reclaimed after finish. + self._reclaimed_payload_total = 0 self.waiting_for_chunk_waiting_requests: deque[Any] = deque() self.waiting_for_chunk_running_requests: deque[Any] = deque() - self.requests_with_ready_chunks = set() + self.requests_with_ready_chunks: set[str] = set() self.replaced_streaming_prompt_ids: set[str] = set() - self.requests_origin_status = {} + self.requests_origin_status: dict[str, RequestStatus] = {} self._active_streams: dict[str, Any] = {} # Private hold-queue for non-active running requests. Restored to # running_queue inside restore_queues(). Avoids calling @@ -410,6 +420,7 @@ def save_async( because ``finish_requests`` skips requests that already report ``is_finished()``. See ``_finish_parked_streaming_session``. """ + request = cast(Request, request) is_finished = force_request_finished or (request.is_finished() and not request.resumable) if not hasattr(self, "_segment_generation"): self._segment_generation = defaultdict(int) @@ -498,7 +509,9 @@ def save_async( 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: @@ -692,8 +705,21 @@ def _send_single_request(self, task: dict): for chunk_id in range(num_puts): self.connector.cleanup(f"{key_prefix}_{chunk_id}") return - if "release_shm" in task: - self._release_shm_prefix(task["release_shm"]) + if "release_prefix" in task: + completion = task["release_completion"] + try: + self._release_connector_prefix(task["release_prefix"]) + wait = getattr(self.connector, "wait_for_cleanup", None) + if callable(wait): + wait( + task["release_prefix"], + receive_prefix=task["receive_prefix"], + deadline=task["release_deadline"], + ) + except Exception as error: + completion.set_exception(error) + else: + completion.set_result(None) return request = task["request"] external_req_id = request.external_req_id @@ -774,12 +800,14 @@ def _send_single_request_for_generation( sender_token: _SenderGeneration | None = None, ): raw_mm = task["multimodal_output"] - multimodal_output = unflatten_payload(raw_mm) if isinstance(raw_mm, Mapping) else raw_mm + multimodal_output = unflatten_payload(dict(raw_mm)) if isinstance(raw_mm, Mapping) else raw_mm request = task["request"] 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}" @@ -1050,8 +1078,8 @@ def _clear_sender_state_locked(self, external_req_id: str) -> None: if cached_ic is not None: cached_ic.pop(external_req_id, None) - def release_shm_resources(self, request_id: str) -> None: - """Queue reclaim of inter-stage segments this request left unconsumed. + def release_request_resources(self, request_id: str, *, deadline: float | None = None) -> Future[None]: + """Finish queued publications and apply the backend's request cleanup. Safe only once every stage has finished with the request: the producer keeps writing until its own request ends, and a consumer that stopped @@ -1060,30 +1088,45 @@ def release_shm_resources(self, request_id: str) -> None: orchestrator -- the only party that knows all stages are done -- drives this (see ``Orchestrator._cleanup_request_ids``). - Like abort cleanup, the unlink runs on the save thread, so the - scheduler never performs SHM I/O. + Cleanup runs behind puts on the save thread, so the scheduler never + performs connector I/O. Legacy backends without prefix cleanup retain + their existing resource lifecycle; their completion confirms the local + publication fence, not physical reclamation by the backend. """ - if not isinstance(self.connector, SharedMemoryConnector): - return external_req_id = self.request_ids_mapping.get(request_id, request_id) - self._pending_save_reqs.append({"release_shm": f"{external_req_id}_{self.connector.stage_id}_"}) + from_stage = get_stage_connector_peer(self.config, self.connector.stage_id, "from_stage") + completion: Future[None] = Future() + self._pending_save_reqs.append( + { + "release_prefix": f"{external_req_id}_{self.connector.stage_id}_", + "receive_prefix": f"{external_req_id}_{from_stage}_" if from_stage is not None else None, + "release_completion": completion, + "release_deadline": deadline + if deadline is not None + else time.monotonic() + REQUEST_RESOURCE_RELEASE_TIMEOUT_S, + } + ) with self._save_cond: self._save_cond.notify() + return completion - def _release_shm_prefix(self, key_prefix: str) -> None: - reclaimed = self.connector.cleanup_prefix(key_prefix) + def _release_connector_prefix(self, key_prefix: str) -> None: + cleanup = getattr(self.connector, "cleanup_prefix", None) + if not callable(cleanup): + return + reclaimed = cleanup(key_prefix) if not reclaimed: return # Non-zero means a consumer stopped before draining the producer. # Common on audio requests, so warn only when the running total # crosses another _RECLAIM_WARN_INTERVAL boundary. - prev = self._reclaimed_shm_total - self._reclaimed_shm_total = prev + reclaimed - if prev // _RECLAIM_WARN_INTERVAL != self._reclaimed_shm_total // _RECLAIM_WARN_INTERVAL: + prev = self._reclaimed_payload_total + self._reclaimed_payload_total = prev + reclaimed + if prev // _RECLAIM_WARN_INTERVAL != self._reclaimed_payload_total // _RECLAIM_WARN_INTERVAL: logger.warning( - "Reclaimed %d unconsumed inter-stage segments so far; " + "Reclaimed %d unconsumed inter-stage payloads so far; " "a downstream stage finished before draining its producer", - self._reclaimed_shm_total, + self._reclaimed_payload_total, ) def cleanup( @@ -1627,7 +1670,7 @@ def finish_requests( elif request_ids is not None: request_ids = set(request_ids) else: - request_ids = requests.keys() + request_ids = cast(dict[str, Request], requests).keys() connector_owned_ids = { request.request_id diff --git a/vllm_omni/distributed/omni_connectors/utils/config.py b/vllm_omni/distributed/omni_connectors/utils/config.py index a646a3a3090..63ad077d4dc 100644 --- a/vllm_omni/distributed/omni_connectors/utils/config.py +++ b/vllm_omni/distributed/omni_connectors/utils/config.py @@ -2,12 +2,33 @@ # 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 logger = get_connector_logger(__name__) +REQUEST_RESOURCE_RELEASE_TIMEOUT_S = 5.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) + + TRANSFER_ENGINE_CONNECTOR_NAMES = frozenset( { "MooncakeTransferEngineConnector", diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index b3a23d8bd64..0e47c8f3c70 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -328,6 +328,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/duplex_omni_engine.py b/vllm_omni/engine/duplex_omni_engine.py index 98ea81fa0b5..d5e80ea70ea 100644 --- a/vllm_omni/engine/duplex_omni_engine.py +++ b/vllm_omni/engine/duplex_omni_engine.py @@ -72,6 +72,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/duplex_orchestrator.py b/vllm_omni/engine/duplex_orchestrator.py index 39719453aec..35f079999d4 100644 --- a/vllm_omni/engine/duplex_orchestrator.py +++ b/vllm_omni/engine/duplex_orchestrator.py @@ -208,7 +208,7 @@ async def _handle_forward_failure( if runner is not None: runner.on_stage_failure(next_stage_id, exc, request_id=req_id) await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + [req_id], abort=True, release_owners=close_session, ) diff --git a/vllm_omni/engine/errors.py b/vllm_omni/engine/errors.py index 0128176d660..abeecc64175 100644 --- a/vllm_omni/engine/errors.py +++ b/vllm_omni/engine/errors.py @@ -6,3 +6,7 @@ class NativeKVHandoffError(RuntimeError): """One request lost its producer binding or transfer metadata.""" + + +class ResourceReleaseError(RuntimeError): + """A live replica could not confirm request transfer ownership was released.""" diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py index 793711682d2..c2cfcda4093 100644 --- a/vllm_omni/engine/omni_engine_base.py +++ b/vllm_omni/engine/omni_engine_base.py @@ -14,12 +14,10 @@ import uuid import weakref from collections.abc import Mapping -from dataclasses import asdict from pathlib import Path from typing import Any, cast import janus -import torch from vllm import envs as vllm_envs from vllm.logger import init_logger from vllm.v1.engine.input_processor import InputProcessor @@ -32,11 +30,6 @@ PipelineConfig, load_deploy_config, ) -from vllm_omni.diffusion.data import ( - DiffusionParallelConfig, - parse_attention_config, -) -from vllm_omni.diffusion.io_support import get_diffusion_output_type from vllm_omni.engine.async_engine_utils import ( SHUTDOWN_ENQUEUE_TIMEOUT_S, SHUTDOWN_JOIN_TIMEOUT_S, @@ -377,6 +370,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, @@ -453,6 +447,7 @@ async def _run_orchestrator() -> None: output_async_queue=self.output_queue.async_q, rpc_async_queue=self.rpc_output_queue.async_q, stage_pools=self.stage_pools, + stage_transitions=self.pipeline_config.stage_transitions if self.pipeline_config is not None else None, async_chunk=self.async_chunk, pd_config=pd_config, membership_controller=membership_controller, @@ -629,219 +624,6 @@ def _detect_pd_config(self) -> dict[str, Any] | None: "prefill_engine_id": prefill_engine_id, } - @staticmethod - def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list: - """Create a default single-stage diffusion config from kwargs.""" - # We temporally create a default config for diffusion stage. - # In the future, we should merge the default config with the user-provided config. - normalized_kwargs = dict(kwargs) - default_sampling_params = normalized_kwargs.get("default_sampling_params") - if isinstance(default_sampling_params, str): - try: - default_sampling_params = json.loads(default_sampling_params) - except json.JSONDecodeError: - logger.warning("Invalid default_sampling_params JSON, ignoring stage defaults.") - default_sampling_params = None - if not isinstance(default_sampling_params, dict): - default_sampling_params = None - stage_default_sampling_params = default_sampling_params.get("0", {}) if default_sampling_params else {} - if normalized_kwargs.get("dtype") is None: - normalized_kwargs["dtype"] = "auto" - - # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error. - if "dtype" in normalized_kwargs and not isinstance(normalized_kwargs["dtype"], str): - if not isinstance(normalized_kwargs["dtype"], torch.dtype): - raise TypeError( - f"Provided dtype must be a string or torch.dtype, got {type(normalized_kwargs['dtype']).__name__}" - ) - normalized_kwargs["dtype"] = str(normalized_kwargs["dtype"]).removeprefix("torch.") - - cache_backend = normalized_kwargs.get("cache_backend", "none") - cache_config = OmniEngineBase._normalize_cache_config( - cache_backend, - normalized_kwargs.get("cache_config", None), - ) - - parallel_config = normalized_kwargs.get("parallel_config") - if isinstance(parallel_config, dict): - parallel_config = DiffusionParallelConfig.from_dict(parallel_config) - if parallel_config is None: - ulysses_degree = normalized_kwargs.get("ulysses_degree") or 1 - ring_degree = normalized_kwargs.get("ring_degree") or 1 - allgather_degree = normalized_kwargs.get("allgather_degree") or 1 - ulysses_mode = normalized_kwargs.get("ulysses_mode") or "strict" - ulysses_a2a_permute = bool(normalized_kwargs.get("ulysses_a2a_permute", False)) - sequence_parallel_size = normalized_kwargs.get("sequence_parallel_size") - pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 - data_parallel_size = normalized_kwargs.get("data_parallel_size") - tensor_parallel_size = normalized_kwargs.get("tensor_parallel_size") or 1 - cfg_parallel_size = normalized_kwargs.get("cfg_parallel_size") or 1 - pipeline_parallel_size = normalized_kwargs.get("pipeline_parallel_size") or 1 - vae_patch_parallel_size = normalized_kwargs.get("vae_patch_parallel_size") or 1 - vae_parallel_mode = normalized_kwargs.get("vae_parallel_mode") or "tile" - text_encoder_tp_size = normalized_kwargs.get("text_encoder_tp_size") or 1 - enable_expert_parallel = normalized_kwargs.get("enable_expert_parallel") or False - use_hsdp = normalized_kwargs.get("use_hsdp", False) - hsdp_shard_size = normalized_kwargs.get("hsdp_shard_size", -1) - hsdp_replicate_size = normalized_kwargs.get("hsdp_replicate_size", 1) - if sequence_parallel_size is None: - sequence_parallel_size = allgather_degree if allgather_degree > 1 else ulysses_degree * ring_degree - - parallel_config = DiffusionParallelConfig( - pipeline_parallel_size=pipeline_parallel_size, - data_parallel_size=data_parallel_size, - tensor_parallel_size=tensor_parallel_size, - enable_expert_parallel=enable_expert_parallel, - sequence_parallel_size=sequence_parallel_size, - ulysses_degree=ulysses_degree, - ring_degree=ring_degree, - allgather_degree=allgather_degree, - ulysses_mode=ulysses_mode, - ulysses_a2a_permute=ulysses_a2a_permute, - cfg_parallel_size=cfg_parallel_size, - vae_patch_parallel_size=vae_patch_parallel_size, - vae_parallel_mode=vae_parallel_mode, - text_encoder_tp_size=text_encoder_tp_size, - use_hsdp=use_hsdp, - hsdp_shard_size=hsdp_shard_size, - hsdp_replicate_size=hsdp_replicate_size, - ) - - num_gpus = normalized_kwargs.get("num_gpus") - if num_gpus is not None: - num_gpus = int(num_gpus) - parallel_config.resolve_data_parallel_size(num_gpus) - - num_devices = max(1, int(parallel_config.world_size)) - devices = ",".join(str(i) for i in range(num_devices)) - model_class_name = kwargs.get("model_class_name", None) - final_output_type = get_diffusion_output_type(model_class_name) - - attention_config = None - if ( - kwargs.get("diffusion_attention_config") is not None - or kwargs.get("diffusion_attention_backend") is not None - or kwargs.get("fastvideo_vsa_topk") is not None - ): - attention_config = parse_attention_config( - kwargs.get("diffusion_attention_config"), - attention_backend=kwargs.get("diffusion_attention_backend"), - fastvideo_vsa_topk=kwargs.get("fastvideo_vsa_topk"), - ) - - extras = dict(kwargs.get("extras") or {}) - for key, default in ( - ("auxiliary_text_encoder", None), - ("default_llama_model_id", "meta-llama/Meta-Llama-3.1-8B-Instruct"), - ): - top_level_value = kwargs.get(key) - if top_level_value is not None: - extras[key] = top_level_value - else: - extras.setdefault(key, default) - - stage_engine_args = { - "max_num_seqs": kwargs.get("max_num_seqs") or 1, - "parallel_config": parallel_config, - # Default-stage construction bypasses the structured projection. - # Runner selection remains owned by the selected engine/platform. - "engine_backend": kwargs.get("engine_backend", "default"), - "model_class_name": kwargs.get("model_class_name", None), - "task_type": kwargs.get("task_type", None), - "model_config": kwargs.get("model_config", None), - "additional_config": kwargs.get("additional_config", None), - "step_execution": kwargs.get("step_execution", False), - "request_batch_max_wait_ms": kwargs.get("request_batch_max_wait_ms", 0.0), - "vae_use_slicing": kwargs.get("vae_use_slicing", False), - "vae_use_tiling": kwargs.get("vae_use_tiling", False), - "cache_backend": cache_backend, - "cache_config": cache_config, - "enable_cache_dit_summary": kwargs.get("enable_cache_dit_summary", False), - "diffusion_offload_config": kwargs.get("diffusion_offload_config", None), - "enable_cpu_offload": kwargs.get("enable_cpu_offload", False), - "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False), - "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False), - "dlo_use_allgather": kwargs.get("dlo_use_allgather", True), - "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0), - "host_weight_runtime_mode": kwargs.get("host_weight_runtime_mode", "disabled"), - "host_weight_runtime_root": kwargs.get("host_weight_runtime_root"), - "dlo_host_registration_limit_gib": kwargs.get("dlo_host_registration_limit_gib", 0.0), - "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"), - "diffusion_compile_granularity": ( - "regional" - if kwargs.get("diffusion_compile_granularity") is None - else kwargs["diffusion_compile_granularity"] - ), - "diffusion_compile_dynamic": ( - True if kwargs.get("diffusion_compile_dynamic") is None else kwargs["diffusion_compile_dynamic"] - ), - "fa_deterministic": bool(kwargs.get("fa_deterministic", False)), - "boundary_ratio": kwargs.get("boundary_ratio", None), - "flow_shift": kwargs.get("flow_shift", None), - "diffusion_load_format": kwargs.get("diffusion_load_format", "default"), - "lora_path": kwargs.get("lora_path", None), - "lora_scale": kwargs.get("lora_scale", 1.0), - "lora_backend": kwargs.get("lora_backend", "peft"), - "custom_pipeline_args": kwargs.get("custom_pipeline_args", None), - "worker_extension_cls": kwargs.get("worker_extension_cls", None), - "trust_remote_code": (False if kwargs.get("trust_remote_code") is None else kwargs["trust_remote_code"]), - "distributed_executor_backend": kwargs.get("distributed_executor_backend"), - "enable_sleep_mode": kwargs.get("enable_sleep_mode", False), - "enable_prompt_embed_cache": kwargs.get("enable_prompt_embed_cache", False), - "prompt_embed_cache_size": kwargs.get("prompt_embed_cache_size", 32), - "enable_multithread_weight_load": kwargs.get("enable_multithread_weight_load", True), - "num_weight_load_threads": kwargs.get("num_weight_load_threads", 4), - "quantization": kwargs.get("quantization", None), - "quantization_config": kwargs.get("quantization_config", None), - "diffusion_kv_cache_dtype": kwargs.get("diffusion_kv_cache_dtype", None), - "diffusion_kv_cache_skip_steps": kwargs.get("diffusion_kv_cache_skip_steps", None), - "diffusion_kv_cache_skip_layers": kwargs.get("diffusion_kv_cache_skip_layers", None), - **({"diffusion_attention_config": attention_config} if attention_config is not None else {}), - "force_cutlass_fp8": bool(kwargs.get("force_cutlass_fp8", False)), - "enable_diffusion_pipeline_profiler": kwargs.get("enable_diffusion_pipeline_profiler", False), - "streaming_output": kwargs.get("diffusion_streaming_output", False), - "enable_ar_profiler": kwargs.get("enable_ar_profiler", False), - "extras": extras, - **( - { - "profiler_config": asdict(kwargs["profiler_config"]) - if hasattr(kwargs["profiler_config"], "__dataclass_fields__") - else kwargs["profiler_config"] - } - if kwargs.get("profiler_config") is not None - else {} - ), - } - if num_gpus is not None: - stage_engine_args["num_gpus"] = num_gpus - # Only set dtype if it was already explicitly passed and normalized - if "dtype" in normalized_kwargs: - stage_engine_args["dtype"] = normalized_kwargs["dtype"] - - # New split fields for diffusers adapter kwargs. - if kwargs.get("diffusers_load_kwargs") is not None: - stage_engine_args["diffusers_load_kwargs"] = kwargs["diffusers_load_kwargs"] - if kwargs.get("diffusers_call_kwargs") is not None: - stage_engine_args["diffusers_call_kwargs"] = kwargs["diffusers_call_kwargs"] - - default_stage_cfg = [ - { - "stage_id": 0, - "stage_type": "diffusion", - "runtime": { - "process": True, - "devices": devices, - }, - "engine_args": stage_engine_args, - "engine_input_source": [], - "default_sampling_params": stage_default_sampling_params, - "final_output": True, - "final_output_type": final_output_type, - } - ] - default_stage_cfg[0]["engine_args"]["model_stage"] = "diffusion" - return default_stage_cfg - def _apply_strategy_lb_policy(self, derived: str | None, kwargs: dict[str, Any]) -> None: """Apply a strategy-derived ``omni_lb_policy`` to the engine. @@ -882,6 +664,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 () self._event_driven_orch_default = _event_driven_orch_default_for_pipeline( pipeline_config.model_type if pipeline_config is not None else None diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index a135b60bede..980d3d34a3a 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -34,13 +34,14 @@ 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 from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker -from vllm_omni.engine.errors import NativeKVHandoffError +from vllm_omni.engine.errors import NativeKVHandoffError, ResourceReleaseError from vllm_omni.engine.membership_controller import MembershipController from vllm_omni.engine.messages import ( AbortRequestMessage, @@ -60,7 +61,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 @@ -239,7 +240,6 @@ class OrchestratorRequestState: sampling_params_list: list[Any] = field(default_factory=list) final_stage_id: int = -1 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 @@ -270,6 +270,22 @@ class OrchestratorRequestState: upstream_first_audio: bool = False pending_upstream_first_audio: tuple[int, OmniEngineCoreOutput] | None = None pending_first_audio_outputs: list[tuple[int, int, Any]] = field(default_factory=list) + # One cleanup operation owns the request until abort and binding release + # finish. Retain state for partial abort outputs, but forbid new admission. + cleanup_future: asyncio.Future[list[OutputMessage]] | None = field(default=None, repr=False) + + @property + def closing(self) -> bool: + return self.cleanup_future is not None + + @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 @@ -336,7 +352,6 @@ class OrchestratorBase: _transfer_emitter: Any = None _prom_metrics: Any = None _stat_logger: OmniPrometheusStatLogger | None = None - _transfer_release_tasks: set[asyncio.Task] = set() def __init__( self, @@ -346,6 +361,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, @@ -363,6 +379,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 self._prom_metrics = prom_metrics self._stage_replica_waiting: dict[tuple[int, int], int] = {} @@ -385,7 +402,6 @@ def __init__( self._pd_prefill_engine_id = pd_config.get("prefill_engine_id") self.request_states: dict[str, OrchestratorRequestState] = {} # Strong refs for in-flight releases; the loop only weak-refs tasks, so dropping these risks mid-flight GC. - self._transfer_release_tasks: set[asyncio.Task] = set() self._init_metrics_state( stage_pools, running_counter, @@ -489,7 +505,7 @@ async def run(self) -> None: if self._membership is not None: async def cleanup_member_requests(ids: list[str]) -> None: - await self._cleanup_request_ids(ids, abort=True) + await self._cleanup_request_ids(ids, abort=True, release_owners=True) self._membership.install_unregister_handlers( output_queue=self.output_async_queue, @@ -505,7 +521,11 @@ async def cleanup_member_requests(ids: list[str]) -> None: tasks.append(membership_watcher) try: - await asyncio.gather(*tasks) + # These loops live until shutdown or failure. Once one exits, + # cancel the others in finally, including the membership watcher. + done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + for task in done: + task.result() except asyncio.CancelledError: raise except EngineDeadError as e: @@ -664,7 +684,10 @@ async def _handle_abort(self, msg: AbortRequestMessage) -> None: error = str(exc) logger.exception("[Orchestrator] Abort failed for request(s) %s", request_ids) if msg.rpc_id is None: - raise + if isinstance(exc, ResourceReleaseError): + await self._report_release_failure(list(request_ids), exc) + else: + raise if msg.rpc_id is not None: # Always emit a result when rpc_id is set so CorrelatedRpcClient # waiters are unblocked even on failure. @@ -728,35 +751,26 @@ async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage output_msg.finished = index == last_index_by_req[output_msg.request_id] return abort_outputs - def _release_stage_transfer_resources(self, request_ids: list[str]) -> None: - """Drop each stage's inter-stage transfer resources for finished requests. + async def _release_stage_transfer_resources(self, request_ids: list[str]) -> None: + """Finish each stage's transfer lifecycle for terminated requests. This is the only point that knows every stage is done with the request, so it is the only safe place to reclaim segments a consumer never - drained. Scheduled rather than awaited: reclaim is best-effort and must - not add RPC latency to request teardown. + drained. Every live replica completes its local publication fence and + applicable backend cleanup. SHM/NIXL additionally confirm ownership + drain; legacy backends retain their existing resource contract. A + failed operation leaves the request closed for retry. """ if not request_ids: return - async def _run() -> None: - results = await asyncio.gather( - *(pool.release_request_resources(request_ids) for pool in self.stage_pools), - return_exceptions=True, - ) - for result in results: - if isinstance(result, Exception): - logger.warning("[Orchestrator] release transfer resources failed: %s", result) - - try: - task = asyncio.get_running_loop().create_task(_run()) - self._transfer_release_tasks.add(task) - task.add_done_callback(self._transfer_release_tasks.discard) - except RuntimeError: - logger.warning( - "[Orchestrator] no running event loop; skipped reclaim of transfer resources for %s", - request_ids, - ) + results = await asyncio.gather( + *(pool.release_request_resources(request_ids) for pool in self.stage_pools), + return_exceptions=True, + ) + for result in results: + if isinstance(result, BaseException): + raise result def _release_request_bindings(self, request_ids: list[str]) -> None: """Release all stage-local route bindings for the given request ids.""" @@ -964,9 +978,14 @@ async def _route_upstream_first_audio(self, stage_id: int, replica_id: int, raw_ # Only the dedicated first-frame channel is intercepted. if ( req_state is None + or req_state.closing 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 ( + 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") @@ -978,6 +997,7 @@ async def _route_upstream_first_audio(self, stage_id: int, replica_id: int, raw_ continue if ( req_state is not None + and not req_state.closing and final_output and not req_state.upstream_first_audio and (requires_first or req_state.pending_first_audio_outputs) @@ -1000,10 +1020,16 @@ async def _flush_upstream_first_audio(self, req_state: OrchestratorRequestState) from vllm_omni.engine import OmniEngineCoreOutputs pending = req_state.pending_upstream_first_audio - if pending is None or self.request_states.get(req_state.request_id) is not req_state: + if pending is None or req_state.closing 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: @@ -1022,11 +1048,11 @@ async def _flush_upstream_first_audio(self, req_state: OrchestratorRequestState) logger.exception("First-audio output processing failed for %s", req_state.request_id) await self._handle_stage_error(codec_stage, first_output, error=str(error)) return - if self.request_states.get(req_state.request_id) is not req_state: + if req_state.closing or self.request_states.get(req_state.request_id) is not req_state: return ready, req_state.pending_first_audio_outputs = req_state.pending_first_audio_outputs, [] for pending_stage, pending_replica, output in ready: - if self.request_states.get(req_state.request_id) is not req_state: + if req_state.closing or self.request_states.get(req_state.request_id) is not req_state: break terminal: set[str] = set() processed = await self._process_llm_stage_outputs( @@ -1416,7 +1442,7 @@ async def _handle_stage_error(self, stage_id: int, output: Any, *, error: str | ) ) await self._cleanup_request_ids( - [parent_id, *self._cfg_tracker.cleanup_parent(parent_id)], + [parent_id], abort=True, release_owners=True, ) @@ -1472,9 +1498,9 @@ async def _handle_dead_replica(self, stage_id: int, replica_id: int, error: Engi stage_has_live = bool(pool.live_replica_ids()) failed_ids: list[str] = [] for req_id, req_state in list(self.request_states.items()): - if stage_id not in req_state.stage_submit_ts: - continue bound = pool.get_bound_replica_id(req_id) + if stage_id not in req_state.stage_submit_ts and bound is None: + continue if bound == replica_id or not stage_has_live: await self.output_async_queue.put( ErrorMessage( @@ -1492,7 +1518,7 @@ async def _handle_dead_replica(self, stage_id: int, replica_id: int, error: Engi # distributed-membership eviction path. for req_id in failed_ids: await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + [req_id], abort=True, # A session-owned request has an owner that outlives the stage # request: without this the runner and its admission slot @@ -1526,7 +1552,7 @@ async def _fail_request_dead_stage(self, req_id: str, stage_id: int) -> None: ) ) await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + [req_id], abort=True, release_owners=True, ) @@ -1557,7 +1583,7 @@ async def _fail_request_client_error( ) ) await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + [req_id], abort=True, release_owners=release_owners, ) @@ -1596,6 +1622,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, @@ -1676,26 +1704,27 @@ async def _cleanup_request_ids( cannot outlive its requests. ``release_owners`` asks the subclass to also release any session that - owns these requests (error paths pass True); the base ignores it. + owns these requests (error paths pass True). It also marks automatic + error cleanup, whose release failure must remain request-scoped. Returns: Final-stage AR abort ``OutputMessage`` list when ``abort=True``; otherwise an empty list. """ - del release_owners if not request_ids: return [] - cleanup_ids = list(dict.fromkeys(request_ids)) - batch = set(cleanup_ids) + candidates = list(dict.fromkeys(request_ids)) + batch = set(candidates) orphaned_parents: dict[str, str] = {} - for rid in cleanup_ids: + for rid in candidates: pid = self._cfg_tracker.get_parent_id(rid) if pid is not None and pid not in batch and not self._cfg_tracker.is_companion_done(rid): orphaned_parents.setdefault(pid, rid) + errors: list[ErrorMessage] = [] for pid, cid in orphaned_parents.items(): deferred = self._cfg_tracker.pop_pending_parent(pid) - await self.output_async_queue.put( + errors.append( ErrorMessage( request_id=pid, stage_id=deferred["stage_id"] if deferred is not None else None, @@ -1703,25 +1732,107 @@ async def _cleanup_request_ids( ) ) batch.add(pid) - cleanup_ids.append(pid) - for rid in list(cleanup_ids): - for cid in self._cfg_tracker.cleanup_parent(rid): + candidates.append(pid) + for rid in list(candidates): + for cid in self._cfg_tracker.get_companion_request_ids(rid).values(): if cid not in batch: batch.add(cid) - cleanup_ids.append(cid) + candidates.append(cid) + + # Expand CFG ownership before selecting cleanup owners so a parent's + # cleanup joins a companion's running operation rather than replacing it. + cleanup_ids: list[str] = [] + pending: dict[asyncio.Future[list[OutputMessage]], list[str]] = {} + for request_id in candidates: + state = self.request_states.get(request_id) + future = state.cleanup_future if state is not None else None + if future is not None and (not future.done() or future.exception() is None): + pending.setdefault(future, []).append(request_id) + else: + cleanup_ids.append(request_id) + + # A parent must stay closed until its already-closing companions have + # released their resources. Unrelated joined requests retain their own + # outcome, so their failure cannot prevent this batch's cleanup. + dependencies = tuple( + future + for future, joined_ids in pending.items() + if any(self._cfg_tracker.get_parent_id(rid) in cleanup_ids for rid in joined_ids) + or any(self._cfg_tracker.get_parent_id(rid) in joined_ids for rid in cleanup_ids) + ) + for future in dependencies: + pending.pop(future) + abort_outputs: list[OutputMessage] = [] - if abort: - abort_outputs = await self._abort_request_ids(cleanup_ids) - self._release_request_bindings(cleanup_ids) - self._release_stage_transfer_resources(cleanup_ids) - for request_id in cleanup_ids: + if cleanup_ids: + future = asyncio.get_running_loop().create_future() + for request_id in cleanup_ids: + state = self.request_states.get(request_id) + if state is not None: + state.cleanup_future = future + try: + for error in errors: + await self.output_async_queue.put(error) + abort_outputs = await self._finish_request_cleanup(cleanup_ids, abort=abort, dependencies=dependencies) + except BaseException as exc: + future.set_exception(exc) + # The owner raises directly; joining callers receive the same + # failure, without an unobserved-Future warning if none join. + future.exception() + if isinstance(exc, ResourceReleaseError) and (not abort or release_owners): + await self._report_release_failure(cleanup_ids, exc) + else: + raise + else: + future.set_result(abort_outputs) + + for future, joined_ids in pending.items(): + # Cancelling a duplicate abort waiter must not cancel the owner. + try: + abort_outputs.extend(await asyncio.shield(future)) + except ResourceReleaseError as error: + if abort and not release_owners: + raise + await self._report_release_failure(joined_ids, error) + return abort_outputs if abort else [] + + async def _report_release_failure(self, request_ids: list[str], error: ResourceReleaseError) -> None: + for request_id in request_ids: + await self.output_async_queue.put( + ErrorMessage( + error=str(error), + fatal=False, + request_id=request_id, + status_code=HTTPStatus.BAD_GATEWAY.value, + error_type=type(error).__name__, + ) + ) + + async def _finish_request_cleanup( + self, + request_ids: list[str], + *, + abort: bool, + dependencies: tuple[asyncio.Future[list[OutputMessage]], ...] = (), + ) -> list[OutputMessage]: + """Finish the claimed cleanup while request state remains available.""" + abort_outputs = await self._abort_request_ids(request_ids) if abort else [] + await self._release_stage_transfer_resources(request_ids) + results = await asyncio.gather(*(asyncio.shield(future) for future in dependencies), return_exceptions=True) + for result in results: + if isinstance(result, BaseException): + raise result + abort_outputs.extend(result) + self._release_request_bindings(request_ids) + for request_id in request_ids: + self._cfg_tracker.cleanup_parent(request_id) self._pd_kv_params.pop(request_id, None) req_state = self.request_states.pop(request_id, None) if req_state is not None: - cleanup_request_artifact_dirs(getattr(req_state, "request_artifact_dirs", set())) - if req_state is not None and req_state.running_counter_registered and self._running_counter is not None: - self._running_counter.decrement() - req_state.running_counter_registered = False + cleanup_request_artifact_dirs(req_state.request_artifact_dirs) + if req_state.running_counter_registered and self._running_counter is not None: + self._running_counter.decrement() + req_state.running_counter_registered = False return abort_outputs async def _apply_raw_terminal_stage_finish( @@ -1748,10 +1859,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( @@ -1807,23 +1914,19 @@ async def _finish_raw_terminal_requests( 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 or req_state.closing or req_state.session_owned: continue pending = req_state.pending_final_output if pending is not None: - if 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) - 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 - 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) logger.info( "[Orchestrator] req=%s stage-%s raw terminal produced no processed terminal; returning empty %s output", @@ -1836,18 +1939,32 @@ 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), - ) + 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), ) - 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, message) + else: + req_state.pending_final_output = message + + async def _publish_final_output(self, req_state: OrchestratorRequestState, message: OutputMessage) -> None: + """Publish success after the request's transfer lifecycle finishes.""" + request_id = req_state.request_id + await self._cleanup_request_ids([request_id]) + if request_id not in self.request_states: + req_state.pending_final_output = None + await self.output_async_queue.put(message) + + 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 _maybe_clone_diffusion_params_for_cfg(self, request_id: str, params: Any) -> Any: """Attach CFG companion ids to diffusion sampling params when needed.""" @@ -1881,6 +1998,8 @@ async def _route_output( stage_metrics: Any, ) -> None: """Route a processed output: send to frontend and/or forward.""" + if req_state.closing: + return req_id = output.request_id finished = output.finished submit_ts = req_state.stage_submit_ts.get(stage_id) @@ -1898,9 +2017,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 @@ -1914,12 +2031,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 @@ -1948,12 +2060,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 @@ -1997,16 +2110,14 @@ 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) - if request_finished and not req_state.session_owned: - await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)]) + def _request_stage_path(self, req_state: OrchestratorRequestState) -> tuple[int, ...]: + return self._stage_routing.path_to(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 + 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 _stage_receives_async_chunks(self, stage_id: int) -> bool: """Whether a stage's connector supplies its runtime inputs.""" @@ -2148,11 +2259,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( @@ -2178,7 +2290,7 @@ async def _handle_kv_ready_raw_outputs( req_id = raw_output.request_id req_state = self.request_states.get(req_id) - if req_state is None: + if req_state is None or req_state.closing: continue if self._cfg_tracker.is_companion(req_id): # kv_ready only says the companion's KV hit the connector; its @@ -2188,9 +2300,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): @@ -2290,18 +2401,24 @@ 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.""" + if req_state.closing: + return + next_stage_id = self._next_request_stage(src_stage_id, req_state) + if 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_stage_id=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", ) @@ -2312,23 +2429,23 @@ async def _forward_to_next_stage_unguarded( output: Any, req_state: OrchestratorRequestState, *, + next_stage_id: 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] + next_pool = self.stage_pools[next_stage_id] if not next_pool.live_replica_ids(): # The downstream stage lost all of its replicas before this in-flight # request could be forwarded. - await self._fail_request_dead_stage(req_id, next_logical) + await self._fail_request_dead_stage(req_id, next_stage_id) return next_client = next_pool.stage_client - params = req_state.sampling_params_list[next_logical] + params = req_state.sampling_params_list[next_stage_id] source_outputs = [output] next_stage_resumable = is_streaming_session and not is_final_update - already_submitted = self._next_stage_already_submitted(src_stage_id, req_state) + already_submitted = next_stage_id in req_state.stage_submit_ts requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False) _t_submit_start = _time.perf_counter() @@ -2371,7 +2488,7 @@ async def _forward_to_next_stage_unguarded( ) ) await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + [req_id], abort=True, ) return @@ -2401,59 +2518,36 @@ async def _forward_to_next_stage_unguarded( req_id, _dt_ar2d, src_stage_id, - next_logical, + next_stage_id, ) - if diffusion_prompt is None: + if diffusion_prompt is None or (isinstance(diffusion_prompt, list) and not diffusion_prompt): error_output = OmniRequestOutput.from_error( req_id, - f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}", + f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_stage_id}", ) logger.warning( "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; " "routing terminal error output", req_id, src_stage_id, - next_logical, + next_stage_id, ) await self.output_async_queue.put( OutputMessage( request_id=req_id, - stage_id=next_logical, + stage_id=next_stage_id, engine_outputs=error_output, metrics=None, finished=True, ) ) await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + [req_id], + abort=True, + release_owners=True, ) return if isinstance(diffusion_prompt, list): - if not diffusion_prompt: - error_output = OmniRequestOutput.from_error( - req_id, - f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}", - ) - logger.warning( - "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; " - "routing terminal error output", - req_id, - src_stage_id, - next_logical, - ) - await self.output_async_queue.put( - OutputMessage( - request_id=req_id, - stage_id=next_logical, - engine_outputs=error_output, - metrics=None, - finished=True, - ) - ) - await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], - ) - return if len(diffusion_prompt) == 1: diffusion_prompt = diffusion_prompt[0] else: @@ -2475,18 +2569,20 @@ async def _forward_to_next_stage_unguarded( submit_kwargs=submit_kwargs, params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params), ) + if req_state.closing: + return self._on_stage_submitted( - next_logical, + next_stage_id, req_id, replica_id, req_state, ) - req_state.stage_submit_ts[next_logical] = _time.time() + req_state.stage_submit_ts[next_stage_id] = _time.time() _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 self._emit_tx_edge( from_stage=src_stage_id, from_replica=src_replica_id if src_replica_id is not None else 0, - to_stage=next_logical, + to_stage=next_stage_id, to_pool=next_pool, request_id=req_id, tx_ms=_tx_ms, @@ -2494,7 +2590,7 @@ async def _forward_to_next_stage_unguarded( return # PD disaggregation: prefill → decode routing uses original prompt + KV transfer params - if self._pd_pair is not None and (src_stage_id, next_logical) == self._pd_pair: + if self._pd_pair is not None and (src_stage_id, next_stage_id) == self._pd_pair: params = self._build_pd_decode_params(req_id, params) # Use the original user prompt for the decode stage (not processed embeddings) @@ -2528,19 +2624,21 @@ async def _forward_to_next_stage_unguarded( replica_id = await next_pool.submit_update(req_id, req_state, request) else: replica_id = await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) + if req_state.closing: + return self._on_stage_submitted( - next_logical, + next_stage_id, req_id, replica_id, req_state, ) - req_state.stage_submit_ts[next_logical] = _time.time() + req_state.stage_submit_ts[next_stage_id] = _time.time() _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 self._emit_tx_edge( from_stage=src_stage_id, from_replica=src_replica_id if src_replica_id is not None else 0, - to_stage=next_logical, + to_stage=next_stage_id, to_pool=next_pool, request_id=req_id, tx_ms=_tx_ms, @@ -2571,7 +2669,7 @@ async def _forward_to_next_stage_unguarded( except OmniClientError as exc: await self._fail_request_client_error( req_id, - next_logical, + next_stage_id, str(exc), status_code=exc.status_code, error_type=exc.error_type, @@ -2582,9 +2680,9 @@ async def _forward_to_next_stage_unguarded( logger.exception( "[Orchestrator] req=%s process_engine_inputs FAILED for stage-%s", req_id, - next_logical, + next_stage_id, ) - if not await self._handle_forward_failure(req_id, next_logical, req_state, exc): + if not await self._handle_forward_failure(req_id, next_stage_id, req_state, exc): raise return finally: @@ -2596,7 +2694,7 @@ async def _forward_to_next_stage_unguarded( "[Orchestrator] req=%s stage-%s produced no inputs for stage-%s; waiting for more outputs", req_id, src_stage_id, - next_logical, + next_stage_id, ) return @@ -2608,14 +2706,13 @@ 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) logger.info( "[Orchestrator] req=%s stage-%s produced no terminal inputs for stage-%s; " "returning empty %s output from final stage-%s", req_id, src_stage_id, - next_logical, + next_stage_id, final_output_type or "text", final_stage_id, ) @@ -2623,14 +2720,16 @@ async def _forward_to_next_stage_unguarded( OutputMessage( request_id=req_id, stage_id=final_stage_id, - replica_id=0, + replica_id=final_pool.get_bound_replica_id(req_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)]) + # Later chunk receivers may already be prewarmed. Empty bridge + # input means they will never receive a payload or terminal marker. + await self._cleanup_request_ids([req_id], abort=True, release_owners=True) return # Build and submit requests for each input @@ -2641,7 +2740,7 @@ async def _forward_to_next_stage_unguarded( mm_features = req_state.mm_features if model_stage == "thinker" else None request = self._build_next_stage_request( req_id, - next_logical, + next_stage_id, next_input, params=params, mm_features=mm_features, @@ -2653,19 +2752,21 @@ async def _forward_to_next_stage_unguarded( replica_id = await next_pool.submit_update(req_id, req_state, request) else: replica_id = await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) + if req_state.closing: + return self._on_stage_submitted( - next_logical, + next_stage_id, req_id, replica_id, req_state, ) - req_state.stage_submit_ts[next_logical] = _time.time() + req_state.stage_submit_ts[next_stage_id] = _time.time() _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 self._emit_tx_edge( from_stage=src_stage_id, from_replica=src_replica_id if src_replica_id is not None else 0, - to_stage=next_logical, + to_stage=next_stage_id, to_pool=next_pool, request_id=req_id, tx_ms=_tx_ms, @@ -2690,7 +2791,10 @@ async def _prewarm_async_chunk_stages( caller still holding ``req_state`` stops instead of recording state on an object the cleanup already popped from ``request_states``. """ - if req_state.final_stage_id <= 0: + stage_path = self._request_stage_path(req_state) + if req_state.closing or self.request_states.get(request_id) is not req_state: + return False + if len(stage_path) <= 1: return True prompt_token_ids = getattr(stage0_request, "prompt_token_ids", None) @@ -2722,9 +2826,23 @@ async def _prewarm_async_chunk_stages( ) return False - for next_stage_id in range(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] + # Reserve bridge-fed senders without submitting a placeholder. + # Their downstream receiver must use the same replica as the later + # real submission, including when the default replica is unavailable. + if not await self._dispatch_or_fail_request( + lambda: next_pool.pick(request_id, request_state=req_state), + req_id=request_id, + stage_id=next_stage_id, + operation="async-chunk replica selection", + ): + return False + if req_state.closing or self.request_states.get(request_id) is not req_state: + return False 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 @@ -2737,7 +2855,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, ) @@ -2804,7 +2922,7 @@ async def _prewarm_async_chunk_stages( model_config=next_pool.stage_vllm_config.model_config, resumable=downstream_resumable, ) - request.payload_sender_info = self._build_payload_sender_info(next_stage_id - 1, request_id=request_id) + request.payload_sender_info = self._build_payload_sender_info(source_stage_id, request_id=request_id) request.external_req_id = request.request_id submitted = await self._dispatch_or_fail_request( lambda: next_pool.submit_initial( @@ -2820,7 +2938,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 @@ -2835,16 +2953,15 @@ async def _prewarm_async_chunk_stages( if req_state.pending_upstream_first_audio is not None or req_state.upstream_first_audio: await self._flush_upstream_first_audio(req_state) - if self.request_states.get(request_id) is not req_state: + if req_state.closing or 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 prewarming to the actual route edge and bound source + # replica; fall back to replica 0 when no source binding exists. _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, @@ -2860,12 +2977,13 @@ def _maybe_attach_native_kv_transfer_params( prompt: Any, ) -> None: source_stage_id = 0 - if source_stage_id + 1 > req_state.final_stage_id: + next_stage_id = self._next_request_stage(source_stage_id, req_state) + if next_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[next_stage_id].stage_type != "diffusion": return from vllm_omni.diffusion.diffusion_kv.kv_connector import build_source_kv_transfer_params, mint_transfer_id @@ -2929,7 +3047,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): @@ -2957,7 +3075,7 @@ def _build_payload_sender_info( return None 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_payload_sender_info", None) if not callable(get_sender_info): @@ -2998,7 +3116,7 @@ async def _dispatch_message(self, msg: EngineQueueMessage) -> bool: def _native_mrv2_receiver_stage(self, final_stage_id: int) -> 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): + for stage_id in self._stage_routing.path_to(final_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( @@ -3015,10 +3133,40 @@ 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 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(): # Stage 0 lost all replicas between the HTTP-layer errored check and @@ -3087,7 +3235,7 @@ async def _handle_add_request(self, msg: StageSubmissionMessage) -> None: ): return - if self.async_chunk and stage_id == 0 and final_stage_id > 0: + if self.async_chunk and req_state.final_stage_id != 0: await self._prewarm_async_chunk_stages(request_id, prompt, req_state, attach_prewarm_payload=True) async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: @@ -3097,7 +3245,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 @@ -3110,6 +3258,19 @@ async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: ) return + 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 @@ -3128,7 +3289,7 @@ async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: ): return - if self.async_chunk and stage_id == 0 and final_stage_id > 0: + if self.async_chunk and req_state.final_stage_id != 0: await self._prewarm_async_chunk_stages(request_id, request, req_state) async def _handle_add_companion(self, msg: AddCompanionRequestMessage) -> None: @@ -3140,7 +3301,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, @@ -3192,7 +3353,7 @@ async def _handle_interaction(self, msg: InteractionMessage) -> None: request_id = msg.request_id event_id = msg.interaction.get("event_id") req_state = self.request_states.get(request_id) - if req_state is None: + if req_state is None or req_state.closing: logger.info("[Orchestrator] Dropping interaction for inactive req %s", request_id) await self.output_async_queue.put( ErrorMessage( @@ -3206,7 +3367,7 @@ async def _handle_interaction(self, msg: InteractionMessage) -> None: return try: - await self.stage_pools[stage_id].submit_interaction(request_id, msg.interaction) + await self.stage_pools[stage_id].submit_interaction(request_id, msg.interaction, request_state=req_state) except Exception as exc: logger.info( "[Orchestrator] Failed interaction for req %s: %s", diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index d109047a116..6e9c6593fe6 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -13,6 +13,7 @@ import contextlib import os import signal +import time from typing import Any import vllm.v1.engine.core as _vllm_engine_core_module @@ -33,6 +34,8 @@ ) from vllm.v1.executor.uniproc_executor import UniProcExecutor +from vllm_omni.core.sched.omni_scheduling_coordinator import uses_native_mrv2_data_plane +from vllm_omni.distributed.omni_connectors.utils.config import REQUEST_RESOURCE_RELEASE_TIMEOUT_S from vllm_omni.distributed.omni_coordinator import create_stage_coord_client from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.stage_init_utils import ( @@ -138,14 +141,20 @@ def omni_release_request_resources(self, request_ids: list[str]) -> None: Invoked over the UTILITY channel by the orchestrator once every stage has finished with the request. Idempotent and safe for unknown ids. """ + model_config = self.vllm_config.model_config + if uses_native_mrv2_data_plane( + model_config, use_v2_model_runner=getattr(model_config, "use_v2_model_runner", False) + ): + self.model_executor.collective_rpc("omni_release_request_resources", args=(request_ids,)) + return adapter = getattr(getattr(self, "scheduler", None), "chunk_transfer_adapter", None) if adapter is None: return - for request_id in request_ids or (): - try: - adapter.release_shm_resources(request_id) - except Exception as e: - logger.debug("omni_release_request_resources(%s) failed: %s", request_id, e) + deadline = time.monotonic() + REQUEST_RESOURCE_RELEASE_TIMEOUT_S + completions = [adapter.release_request_resources(request_id, deadline=deadline) for request_id in request_ids] + for completion in completions: + if completion is not None: + completion.result(timeout=max(0, deadline - time.monotonic())) def preprocess_add_request(self, request: OmniEngineCoreRequest) -> tuple[Any, int]: """Preserve omni payloads when vLLM builds its scheduler request.""" diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 4413cd01aad..d735a42ce36 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -49,9 +49,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 @@ -1933,9 +1935,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 b5d5bec6d01..e6ad1f7f489 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -15,6 +15,7 @@ from vllm.v1.engine import EngineCoreOutputs from vllm.v1.metrics.stats import IterationStats +from vllm_omni.core.sched.omni_scheduling_coordinator import uses_native_mrv2_data_plane from vllm_omni.data_entry_keys import flatten_payload from vllm_omni.distributed.omni_coordinator import ( LoadBalancer, @@ -23,6 +24,7 @@ ReplicaStatus, ) from vllm_omni.distributed.omni_coordinator.load_balancer import Task +from vllm_omni.engine.errors import ResourceReleaseError from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.engine.stage_client import ( StagePoolClient, @@ -62,6 +64,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.""" @@ -100,7 +106,7 @@ class StagePool: DISPATCH_WAIT_TIMEOUT_S: float = 10.0 DISPATCH_RETRY_INTERVAL_S: float = 0.1 # A replica that dies mid-release never answers; without a bound its background release never finishes. - RELEASE_RPC_TIMEOUT_S: float = 5.0 + RELEASE_RPC_TIMEOUT_S: float = 15.0 # Only these EngineCore helpers may skip collective_rpc_async. A generic # ``{method}_async`` on AsyncMPClient must not silently drop timeout. _CACHE_RESET_METHODS = frozenset({"reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"}) @@ -138,8 +144,12 @@ def __init__( self.clients: list[StagePoolClient | None] = list(normalized_clients) self._output_processor = output_processor self._stage_vllm_config = stage_vllm_config - self._has_chunk_transfer_adapter = bool( - getattr(getattr(stage_vllm_config, "model_config", None), "async_chunk", False) + model_config = getattr(stage_vllm_config, "model_config", None) + self._has_transfer_resources = bool( + getattr(model_config, "async_chunk", False) + or uses_native_mrv2_data_plane( + model_config, use_v2_model_runner=getattr(model_config, "use_v2_model_runner", False) + ) ) self._next_replica_id = 0 self._request_bindings: dict[str, int] = {} @@ -321,6 +331,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``. @@ -329,10 +340,16 @@ 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? @@ -341,6 +358,8 @@ async def pick( 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") # Bound replica is gone or DOWN — fall through to re-select. self._affinity.pop(request_id, None) @@ -357,6 +376,8 @@ async def pick( 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) candidates = self._collect_serviceable_replicas() if candidates: # LB chose an index *into our candidates list*. @@ -835,7 +856,7 @@ def _collect_audio_metrics( def _infer_audio_sample_rate( self, - mm_output: dict[str, Any] | None = None, + mm_output: Mapping[str, Any] | None = None, *, use_default: bool = True, ) -> int: @@ -1011,17 +1032,19 @@ async def submit_initial( payload_sender_info = getattr(request, "payload_sender_info", None) if payload_sender_info is not None: submit_kwargs.setdefault("payload_sender_info", payload_sender_info) - 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) + diffusion_client = self._diffusion_client(replica_id) + await diffusion_client.add_request_async(request_id, request, params, **submit_kwargs) return replica_id - 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.clients[replica_id] if client is None: @@ -1070,9 +1093,7 @@ async def submit_update( params = req_state.sampling_params_list[self.stage_id] if self.stage_type == "diffusion": params = OmniDiffusionSamplingParams.from_params(params) - 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, request_state=req_state) client = self.clients[replica_id] if client is None: @@ -1117,11 +1138,11 @@ async def submit_interaction( self, request_id: str, interaction: OmniInteractionPrompt, + *, + request_state: OrchestratorRequestState | None = None, ) -> int: """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, request_state=request_state) client = self._diffusion_client(replica_id) result = await client.submit_interaction_async(request_id, interaction) @@ -1131,17 +1152,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 _poll_stage_raw(self, client: StagePoolLLMClient) -> EngineCoreOutputs | None: @@ -1317,20 +1327,23 @@ async def release_request_resources(self, request_ids: list[str]) -> None: bindings as part of the same teardown, so a binding lookup here would race it. The engine-core handler is idempotent for unknown ids. """ - if not request_ids or not self._has_chunk_transfer_adapter: + if not request_ids or not self._has_transfer_resources: return ids = list(request_ids) async def release(replica_id: int, call: Any) -> None: try: + if call is None: + raise RuntimeError("Replica does not support resource-release RPC") await asyncio.wait_for(call("omni_release_request_resources", ids), timeout=self.RELEASE_RPC_TIMEOUT_S) except Exception as e: - logger.warning( - "[StagePool-%s] release_request_resources on replica %s failed: %r", self.stage_id, replica_id, e - ) + raise ResourceReleaseError(f"Stage {self.stage_id} replica {replica_id} resource release failed") from e calls = [(i, getattr(self.clients[i], "call_utility_async", None)) for i in self.live_replica_ids()] - await asyncio.gather(*(release(i, call) for i, call in calls if call is not None)) + results = await asyncio.gather(*(release(i, call) for i, call in calls), return_exceptions=True) + for result in results: + if isinstance(result, BaseException): + raise result async def collective_rpc( self, diff --git a/vllm_omni/engine/stage_runtime.py b/vllm_omni/engine/stage_runtime.py index dc43575a56c..1e59608aa13 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, ) @@ -170,12 +171,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 @@ -767,6 +774,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] @@ -1366,6 +1374,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, @@ -1373,6 +1382,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, @@ -1691,6 +1701,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: @@ -1714,6 +1725,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, @@ -1731,6 +1743,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 765d540a3d3..01aabb24378 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.entrypoints.cli.logo import log_logo from vllm_omni.entrypoints.openai.api_server import ( @@ -1187,6 +1188,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)), ) @@ -1435,6 +1437,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 ca86587e9fb..dc5c68eec42 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, @@ -455,11 +456,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]: @@ -469,8 +477,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, diff --git a/vllm_omni/worker/mixins.py b/vllm_omni/worker/mixins.py index 724cc87be18..f8bcd3cb573 100644 --- a/vllm_omni/worker/mixins.py +++ b/vllm_omni/worker/mixins.py @@ -9,6 +9,14 @@ class OmniWorkerMixin: """Shared Omni plugin and native KV connector setup for workers.""" + model_runner: Any + + def omni_release_request_resources(self, request_ids: list[str]) -> None: + """Release native runner payloads once the pipeline is done.""" + plane = self.model_runner._omni_data_plane + if plane is not None: + plane.release_request_resources(request_ids) + def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) from vllm_omni.plugins import load_omni_general_plugins diff --git a/vllm_omni/worker_v2/omni_data_plane.py b/vllm_omni/worker_v2/omni_data_plane.py index d82e913dc32..9423426e81b 100644 --- a/vllm_omni/worker_v2/omni_data_plane.py +++ b/vllm_omni/worker_v2/omni_data_plane.py @@ -5,7 +5,8 @@ import threading import time -from collections import defaultdict +from collections import defaultdict, deque +from concurrent.futures import Future from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass, field from queue import Full, Queue @@ -16,6 +17,10 @@ from vllm.logger import init_logger from vllm_omni.data_entry_keys import OmniPayload, flatten_payload, unflatten_payload +from vllm_omni.distributed.omni_connectors.utils.config import ( + REQUEST_RESOURCE_RELEASE_TIMEOUT_S, + get_stage_connector_peer, +) from vllm_omni.worker.omni_connector_model_runner_mixin import ( OmniConnectorModelRunnerMixin, should_accumulate_full_payload_output, @@ -338,6 +343,34 @@ def abort_requests(self, req_ids: set[str]) -> int: terminal_req_ids=active_req_ids, ) + def release_request_resources(self, request_ids: list[str]) -> None: + """Fence publications and apply the existing backend cleanup contract.""" + if self._omni_connector is None: + return + # The scheduler's final runner step may follow the release RPC. Fence + # late publications with the existing abort lifecycle before reclaiming. + self.abort_requests(set(request_ids)) + deadline = time.monotonic() + REQUEST_RESOURCE_RELEASE_TIMEOUT_S + from_stage = get_stage_connector_peer(self.model_config, self._stage_id, "from_stage") + completions: list[Future[None]] = [] + with self._native_send_lock, self._lock: + for request_id in request_ids: + completion: Future[None] = Future() + completions.append(completion) + self._pending_save_reqs.setdefault(request_id, deque()).append( + { + "request_id": request_id, + "release_prefix": f"{request_id}_{self._stage_id}_", + "receive_prefix": f"{request_id}_{from_stage}_" if from_stage is not None else None, + "release_completion": completion, + "release_deadline": deadline, + } + ) + self._pending_save_counts[request_id] += 1 + self._save_work_available.set() + for completion in completions: + completion.result(timeout=max(0, deadline - time.monotonic())) + def complete_outputs( self, *, From 01106ded8191fe90a664e818fba4c40587da7a48 Mon Sep 17 00:00:00 2001 From: Pinoeer-kingxi <13022943007@163.com> Date: Sat, 10 Oct 2026 19:35:55 -0700 Subject: [PATCH 2/9] [Core] Simplify payload conversion and release handling Signed-off-by: Pinoeer-kingxi <13022943007@163.com> --- tests/engine/test_data_entry_keys.py | 18 ++++++++++++------ vllm_omni/data_entry_keys.py | 2 +- .../transfer_adapter/chunk_transfer_adapter.py | 2 +- vllm_omni/engine/stage_engine_core_proc.py | 3 +-- 4 files changed, 15 insertions(+), 10 deletions(-) diff --git a/tests/engine/test_data_entry_keys.py b/tests/engine/test_data_entry_keys.py index dc9a196ca4a..6f75cb4aa95 100644 --- a/tests/engine/test_data_entry_keys.py +++ b/tests/engine/test_data_entry_keys.py @@ -3,6 +3,8 @@ """Tests for data_entry_keys.""" +from types import MappingProxyType + import msgspec import pytest import torch @@ -257,34 +259,38 @@ def test_mixed_nested_and_top_level(self): assert set(flat.keys()) == {"codes.audio", "latent", "meta.finished"} +@pytest.mark.parametrize("mapping_type", [dict, MappingProxyType]) class TestUnflattenPayload: - def test_basic_dotted_to_nested(self): + def test_basic_dotted_to_nested(self, mapping_type): flat = { "codes.audio": torch.tensor([1.0]), "meta.finished": torch.tensor(True, dtype=torch.bool), "meta.left_context_size": 5, } - nested = unflatten_payload(flat) + nested = unflatten_payload(mapping_type(flat)) assert torch.equal(nested["codes"]["audio"], torch.tensor([1.0])) + assert nested["codes"]["audio"] is flat["codes.audio"] assert nested["meta"]["finished"].item() is True assert nested["meta"]["left_context_size"] == 5 - def test_top_level_keys_preserved(self): + def test_top_level_keys_preserved(self, mapping_type): flat = {"latent": torch.tensor([9.0]), "generated_len": 42} - nested = unflatten_payload(flat) + nested = unflatten_payload(mapping_type(flat)) assert torch.equal(nested["latent"], torch.tensor([9.0])) + assert nested["latent"] is flat["latent"] assert nested["generated_len"] == 42 - def test_hidden_states_layers_collected(self): + def test_hidden_states_layers_collected(self, mapping_type): flat = { "hidden_states.output": torch.tensor([1.0]), "hidden_states.layer_0": torch.tensor([2.0]), "hidden_states.layer_24": torch.tensor([3.0]), } - nested = unflatten_payload(flat) + nested = unflatten_payload(mapping_type(flat)) assert torch.equal(nested["hidden_states"]["output"], torch.tensor([1.0])) assert torch.equal(nested["hidden_states"]["layers"][0], torch.tensor([2.0])) assert torch.equal(nested["hidden_states"]["layers"][24], torch.tensor([3.0])) + assert nested["hidden_states"]["layers"][24] is flat["hidden_states.layer_24"] class TestFlattenUnflattenRoundTrip: diff --git a/vllm_omni/data_entry_keys.py b/vllm_omni/data_entry_keys.py index 046f98e08c0..4f9e099abae 100644 --- a/vllm_omni/data_entry_keys.py +++ b/vllm_omni/data_entry_keys.py @@ -386,7 +386,7 @@ def flatten_payload(payload: Mapping[str, Any]) -> dict[str, Any]: return flat -def unflatten_payload(flat: dict[str, Any]) -> dict[str, Any]: +def unflatten_payload(flat: Mapping[str, Any]) -> dict[str, Any]: """Unflatten dotted keys back to nested dicts. Reverse of :func:`flatten_payload`. 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 1319d6fc4a1..ea0e96569ce 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 @@ -800,7 +800,7 @@ def _send_single_request_for_generation( sender_token: _SenderGeneration | None = None, ): raw_mm = task["multimodal_output"] - multimodal_output = unflatten_payload(dict(raw_mm)) if isinstance(raw_mm, Mapping) else raw_mm + multimodal_output = unflatten_payload(raw_mm) if isinstance(raw_mm, Mapping) else raw_mm request = task["request"] is_finished = task["is_finished"] is_segment_finished = task["is_segment_finished"] diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index 6e9c6593fe6..c370b32b9b4 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -153,8 +153,7 @@ def omni_release_request_resources(self, request_ids: list[str]) -> None: deadline = time.monotonic() + REQUEST_RESOURCE_RELEASE_TIMEOUT_S completions = [adapter.release_request_resources(request_id, deadline=deadline) for request_id in request_ids] for completion in completions: - if completion is not None: - completion.result(timeout=max(0, deadline - time.monotonic())) + completion.result(timeout=max(0, deadline - time.monotonic())) def preprocess_add_request(self, request: OmniEngineCoreRequest) -> tuple[Any, int]: """Preserve omni payloads when vLLM builds its scheduler request.""" From 6548f518e24cc3a8179f828e348294e1114941f6 Mon Sep 17 00:00:00 2001 From: Pinoeer-kingxi <13022943007@163.com> Date: Sat, 10 Oct 2026 21:08:05 -0700 Subject: [PATCH 3/9] [Core] Preserve route-aware cleanup outcomes Retain abort prefixes until transfer release succeeds and return them in configured route order. Publish empty-bridge success after cleanup and commit output processor abort state after all replicas stop. Keep public IDs retryable after failed, timed-out or repeated cancellation cleanup, and exclude closed requests from updates and audio metrics. Cover the failure paths and request isolation with CPU regressions. Signed-off-by: Pinoeer-kingxi <13022943007@163.com> --- tests/engine/test_configured_stage_routing.py | 101 +++++++++- tests/engine/test_orchestrator.py | 1 + .../engine/test_stage_pool_collective_rpc.py | 40 ++++ tests/entrypoints/test_async_omni.py | 190 ++++++++++++++++-- .../test_async_omni_prompt_update.py | 27 ++- tests/entrypoints/test_omni_entrypoints.py | 1 + .../entrypoints/test_stream_finish_reason.py | 9 +- vllm_omni/engine/orchestrator.py | 23 ++- vllm_omni/engine/stage_pool.py | 10 +- vllm_omni/entrypoints/async_omni.py | 14 +- vllm_omni/entrypoints/async_omni_base.py | 24 ++- vllm_omni/entrypoints/client_request_state.py | 3 + vllm_omni/entrypoints/omni_base.py | 3 +- vllm_omni/entrypoints/openai/serving_chat.py | 6 +- 14 files changed, 416 insertions(+), 36 deletions(-) diff --git a/tests/engine/test_configured_stage_routing.py b/tests/engine/test_configured_stage_routing.py index 5431bb3e9ab..b8402bd860a 100644 --- a/tests/engine/test_configured_stage_routing.py +++ b/tests/engine/test_configured_stage_routing.py @@ -26,7 +26,7 @@ from vllm_omni.config.stage_routing import StageRouting from vllm_omni.engine import OmniEngineCoreOutput from vllm_omni.engine.errors import ResourceReleaseError -from vllm_omni.engine.messages import ErrorMessage, StageSubmissionMessage +from vllm_omni.engine.messages import ErrorMessage, OutputMessage, StageSubmissionMessage from vllm_omni.engine.omni_engine_base import OmniEngineBase from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState from vllm_omni.engine.stage_pool import RequestClosedError, StageUnavailableError @@ -375,6 +375,105 @@ async def test_endpoint_completion_waits_for_other_selected_output(transitions, assert orch.output_async_queue.empty() and "r" not in orch.request_states +@pytest.mark.asyncio +@pytest.mark.parametrize("release_fails", [False, True]) +async def test_empty_bridge_waits_for_cleanup_before_publishing_success(mocker, release_fails): + orch, state, _ = _direct_orchestrator(async_chunk=True) + entered, resume = asyncio.Event(), asyncio.Event() + + async def release(_ids): + entered.set() + await resume.wait() + if release_fails: + raise ResourceReleaseError("receiver still owns the payload") + + mocker.patch.object(orch.stage_pools[0], "release_request_resources", side_effect=release) + task = asyncio.create_task(orch._forward_to_next_stage("r", 0, _build_request_output("r"), state)) + try: + await asyncio.wait_for(entered.wait(), timeout=1) + assert orch.output_async_queue.empty() + assert state.closing and not state.cleanup_future.done() + finally: + resume.set() + await asyncio.wait_for(task, timeout=1) + result = orch.output_async_queue.get_nowait() + if release_fails: + assert isinstance(result, ErrorMessage) + assert orch.request_states["r"] is state + else: + assert isinstance(result, OutputMessage) and result.finished + assert "r" not in orch.request_states + assert orch.output_async_queue.empty() + + +@pytest.mark.asyncio +async def test_abort_retry_preserves_tokens_after_transfer_release_failure(mocker): + orch, state, clients = _direct_orchestrator() + pool = orch.stage_pools[2] + pool.select_replica_id("r") + pool.output_processor.request_outputs = [_build_request_output("r", token_ids=[7, 8, 9], finished=False)] + release = mocker.patch.object( + orch.stage_pools[0], + "release_request_resources", + side_effect=[ResourceReleaseError("receiver still owns the payload"), None], + ) + with pytest.raises(ResourceReleaseError): + await orch._cleanup_request_ids(["r"], abort=True) + assert orch.request_states["r"] is state and state.closing + assert pool.get_bound_replica_id("r") is None + outputs = await orch._cleanup_request_ids(["r"], abort=True) + assert len(outputs) == 1 and outputs[0].finished + assert list(outputs[0].engine_outputs.outputs[0].token_ids) == [7, 8, 9] + assert outputs[0].replica_id == 0 + assert clients[2].abort_calls == [["r"]] + assert release.await_count == 2 and "r" not in orch.request_states + + +@pytest.mark.asyncio +async def test_abort_retry_preserves_earlier_stage_tokens_after_later_stage_abort_failure(mocker): + orch, state, clients = _direct_orchestrator(((0, 2), (2, 1))) + clients[2].final_output = True + state.final_output_stage_ids = {1, 2} + for sid in (1, 2): + pool = orch.stage_pools[sid] + pool.select_replica_id("r") + pool.output_processor.request_outputs = [_build_request_output("r", token_ids=[sid], finished=False)] + abort = mocker.patch.object( + clients[2], "abort_requests_async", side_effect=[RuntimeError("stage abort failed"), None] + ) + with pytest.raises(RuntimeError, match="stage abort failed"): + await orch._cleanup_request_ids(["r"], abort=True) + assert orch.stage_pools[1].get_bound_replica_id("r") is None + outputs = await orch._cleanup_request_ids(["r"], abort=True) + assert [(msg.stage_id, msg.finished) for msg in outputs] == [(2, False), (1, True)] + assert [list(msg.engine_outputs.outputs[0].token_ids) for msg in outputs] == [[2], [1]] + assert clients[1].abort_calls == [["r"]] and abort.await_count == 2 + assert "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) diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 7ede13ce5d0..f2ede54e31c 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -2124,6 +2124,7 @@ def release_bindings(self, request_ids: list[str]) -> None: orchestrator = object.__new__(Orchestrator) orchestrator.stage_pools = [first, second] orchestrator._stage_routing = StageRouting.from_transitions(len(orchestrator.stage_pools)) + orchestrator.request_states = {} with pytest.raises(RuntimeError, match="stage abort failed"): await orchestrator._abort_request_ids(["req-a"]) diff --git a/tests/engine/test_stage_pool_collective_rpc.py b/tests/engine/test_stage_pool_collective_rpc.py index 3e2e8f59d5a..e76bb2e77eb 100644 --- a/tests/engine/test_stage_pool_collective_rpc.py +++ b/tests/engine/test_stage_pool_collective_rpc.py @@ -11,6 +11,7 @@ import pytest +from tests.engine.test_orchestrator import FakeStageClient, _build_request_output from vllm_omni.engine.stage_pool import StagePool pytestmark = [pytest.mark.core_model] @@ -175,6 +176,45 @@ def commit_aborted_request_state(self, request_ids, *, internal=False): asyncio.run(run()) +@pytest.mark.cpu +def test_abort_retry_keeps_all_replica_prefixes_until_the_pool_stops(mocker): + async def run() -> None: + class RecordingOutputProcessor: + def __init__(self) -> None: + self.prefixes = {"first": [7, 8], "second": [9]} + + def abort_requests_collecting_outputs(self, request_ids, *, internal=False, commit_state=True): + assert internal is False and commit_state is False + outputs = [ + _build_request_output(rid, token_ids=self.prefixes[rid]) + for rid in request_ids + if rid in self.prefixes + ] + return list(request_ids), outputs + + def commit_aborted_request_state(self, request_ids, *, internal=False): + assert internal is False + for rid in request_ids: + self.prefixes.pop(rid) + + first, second = FakeStageClient(), FakeStageClient() + mocker.patch.object(second, "abort_requests_async", side_effect=[RuntimeError("lost ack"), None]) + processor = RecordingOutputProcessor() + pool = StagePool(0, [first, second], output_processor=processor) + pool._request_bindings.update({"first": 0, "second": 1}) + with pytest.raises(RuntimeError, match="lost ack"): + await pool.abort_requests(["first", "second"]) + assert processor.prefixes == {"first": [7, 8], "second": [9]} + outputs = await pool.abort_requests(["first", "second"]) + assert [(rid, list(output.outputs[0].token_ids)) for rid, output in outputs] == [ + ("first", [7, 8]), + ("second", [9]), + ] + assert not processor.prefixes + + asyncio.run(run()) + + @pytest.mark.cpu def test_release_request_resources_skips_without_async_chunk(): async def run() -> None: diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index 829e30145d1..7920e069881 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -12,11 +12,16 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.v1.kv_hints import KvHintsEnvelope +from tests.engine.test_configured_stage_routing import _direct_orchestrator from tests.helpers.mark import hardware_test from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.errors import ResourceReleaseError +from vllm_omni.engine.messages import ErrorMessage, OutputMessage +from vllm_omni.engine.orchestrator import OrchestratorRequestState from vllm_omni.entrypoints import async_omni as async_omni_mod from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.utils import coerce_param_message_types from vllm_omni.model_executor.models.qwen3_omni.pipeline import QWEN3_OMNI_PIPELINE from vllm_omni.outputs import OmniRequestOutput @@ -80,10 +85,124 @@ def get_async_omni_instance(fake_add_request=_noop, fake_abort_request=_noop) -> omni._compute_final_stage_id = lambda output_modalities: 0 omni._compute_final_output_stage_ids = lambda output_modalities: [0] omni._process_orchestrator_results = fake_process_results - omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None) return omni +@pytest.mark.cpu +def test_generate_retains_public_id_until_failed_cleanup_retry_succeeds(): + async def run(): + orch, _, _ = _direct_orchestrator() + orch.request_states.clear() + fail_release = True + abort_calls = [] + + async def release(_ids): + if fail_release: + raise ResourceReleaseError("receiver still owns the payload") + + orch.stage_pools[0].release_request_resources = release + + async def add(*, request_id, sampling_params_list, final_stage_id, final_output_stage_ids, **kwargs): + orch.request_states[request_id] = OrchestratorRequestState( + request_id=request_id, + sampling_params_list=sampling_params_list, + final_stage_id=final_stage_id, + final_output_stage_ids=set(final_output_stage_ids), + ) + await orch._cleanup_request_ids([request_id]) + await omni.request_states[request_id].queue.put(orch.output_async_queue.get_nowait()) + + async def abort(ids, timeout=None): + abort_calls.append(list(ids)) + return await orch._cleanup_request_ids(list(ids), abort=True) + + omni = get_async_omni_instance(fake_add_request=add, fake_abort_request=abort) + omni.engine.num_stages = len(orch.stage_pools) + omni._compute_final_stage_id = lambda _modalities: 2 + omni._compute_final_output_stage_ids = lambda _modalities: [2] + omni._process_orchestrator_results = AsyncOmni._process_orchestrator_results.__get__(omni) + with pytest.raises(RuntimeError, match="receiver still owns the payload"): + async for _ in omni.generate( + prompt={"prompt": "test"}, + request_id="public-r", + sampling_params_list=[SamplingParams(max_tokens=1) for _ in orch.stage_pools], + output_modalities=["text"], + ): + pass + (internal_id,) = orch.request_states + state = omni.request_states[internal_id] + assert state.cleanup_pending and state.external_request_id == "public-r" + assert orch.request_states[internal_id].closing + with pytest.raises(ResourceReleaseError): + await omni.abort("public-r") + assert omni.request_states[internal_id] is state + fail_release = False + await omni.abort("public-r") + assert abort_calls == [[internal_id], [internal_id], [internal_id]] + assert not omni.request_states and not orch.request_states + assert state.queue.empty() + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_abort_retry_keeps_active_requests_with_the_same_public_id(): + async def run(): + calls: list[list[str]] = [] + omni = get_async_omni_instance(fake_abort_request=get_fake_abort(calls)) + closed = ClientRequestState("closed", external_request_id="public") + closed.cleanup_pending = True + active = ClientRequestState("active", external_request_id="public") + unrelated = ClientRequestState("unrelated", external_request_id="other") + omni.request_states = {state.request_id: state for state in (closed, active, unrelated)} + await omni.abort("public") + assert calls == [["closed", "active"]] + assert omni.request_states == {"active": active, "unrelated": unrelated} + assert closed.queue.empty() and unrelated.queue.empty() + terminal = active.queue.get_nowait() + assert terminal.finished and terminal.engine_outputs.outputs[0].finish_reason == "abort" + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_output_dispatcher_drops_messages_for_requests_pending_cleanup(): + async def run(): + omni = get_async_omni_instance() + pending = ClientRequestState("pending") + pending.cleanup_pending = True + active = ClientRequestState("active") + omni.request_states = {"pending": pending, "active": active} + messages: asyncio.Queue = asyncio.Queue() + omni.engine.try_get_output_async = messages.get + omni.final_output_task = None + for _ in range(64): + messages.put_nowait(ErrorMessage(request_id="pending", error="late release failure")) + messages.put_nowait( + OutputMessage( + request_id="pending", + stage_id=0, + engine_outputs=OmniRequestOutput(request_id="pending", stage_id=0, finished=True, outputs=[]), + finished=True, + ) + ) + marker = ErrorMessage(request_id="active", error="active request failure") + messages.put_nowait(marker) + AsyncOmni._final_output_handler(omni) + task = omni.final_output_task + assert task is not None + try: + assert await asyncio.wait_for(active.queue.get(), timeout=1) is marker + assert pending.queue.empty() and not pending.consumed_metric_message_ids + assert omni.request_states["pending"] is pending + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(run()) + + @pytest.mark.cpu def test_generate_submits_randomized_id_to_engine(): """Ensure the engine receives a UUID-suffixed ID, not the raw request ID""" @@ -243,10 +362,9 @@ async def slow_abort_async(request_ids, timeout=None): await release.wait() omni = get_async_omni_instance(fake_abort_request=slow_abort_async) - omni.request_states["req-1-aaaa"] = SimpleNamespace( + omni.request_states["req-1-aaaa"] = ClientRequestState( request_id="req-1-aaaa", external_request_id="req-1", - input_stream_task=None, ) task = asyncio.create_task(omni.abort("req-1")) @@ -270,10 +388,9 @@ async def failing_abort_async(request_ids, timeout=None): raise RuntimeError("orchestrator abort failed") omni = get_async_omni_instance(fake_abort_request=failing_abort_async) - omni.request_states["req-1-bbbb"] = SimpleNamespace( + omni.request_states["req-1-bbbb"] = ClientRequestState( request_id="req-1-bbbb", external_request_id="req-1", - input_stream_task=None, ) with pytest.raises(RuntimeError, match="orchestrator abort failed"): @@ -327,10 +444,9 @@ async def abort_with_prefix(request_ids, timeout=None): ] omni = get_async_omni_instance(fake_abort_request=abort_with_prefix) - omni.request_states["req-1-cccc"] = SimpleNamespace( + omni.request_states["req-1-cccc"] = ClientRequestState( request_id="req-1-cccc", external_request_id="req-1", - input_stream_task=None, queue=queue, ) await omni._abort(["req-1-cccc"]) @@ -355,10 +471,9 @@ async def empty_abort_async(request_ids, timeout=None): return [] omni = get_async_omni_instance(fake_abort_request=empty_abort_async) - omni.request_states["req-1-dddd"] = SimpleNamespace( + omni.request_states["req-1-dddd"] = ClientRequestState( request_id="req-1-dddd", external_request_id="req-1", - input_stream_task=None, queue=queue, ) await omni._abort(["req-1-dddd"]) @@ -387,7 +502,6 @@ def test_abort_drops_consumed_metric_message_ids_with_state(): from collections.abc import Mapping from vllm_omni.engine.messages import OutputMessage - from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.metrics.stats import OrchestratorAggregator async def run(): @@ -595,7 +709,50 @@ async def collect(): await asyncio.wait_for(task, timeout=1.0) assert seen_timeouts assert seen_timeouts[-1] == 0.05 - assert omni.request_states == {} + (state,) = omni.request_states.values() + assert state.cleanup_pending and state.external_request_id == "cancel-hang-abort" + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_generate_keeps_public_id_if_cleanup_abort_is_cancelled_again(): + async def run(): + submitted, abort_started, resume = asyncio.Event(), asyncio.Event(), asyncio.Event() + calls: list[list[str]] = [] + + async def add(**kwargs): + submitted.set() + + async def abort(request_ids, timeout=None): + calls.append(list(request_ids)) + abort_started.set() + await resume.wait() + + omni = get_async_omni_instance(fake_add_request=add, fake_abort_request=abort) + + async def collect(): + async for _ in omni.generate( + prompt={"prompt": "test"}, + request_id="cancel-twice", + sampling_params_list=[SamplingParams(max_tokens=1)], + output_modalities=["text"], + ): + pass + + task = asyncio.create_task(collect()) + await asyncio.wait_for(submitted.wait(), timeout=1) + task.cancel() + await asyncio.wait_for(abort_started.wait(), timeout=1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1) + (state,) = omni.request_states.values() + assert state.cleanup_pending and state.external_request_id == "cancel-twice" + resume.set() + await omni.abort("cancel-twice") + assert calls == [[state.request_id], [state.request_id]] + assert not omni.request_states and state.queue.empty() asyncio.run(run()) @@ -713,7 +870,16 @@ async def cancel_request(): await asyncio.wait_for(cancel_request(), timeout=1.0) assert abort_started.is_set() assert abort_finished.is_set() - assert omni.request_states == {} + (state,) = omni.request_states.values() + assert state.cleanup_pending + assert state.external_request_id == "cancel-asgi-hanging-abort" + + async def recovered_abort(request_ids, timeout=None): + assert request_ids == [state.request_id] + + omni.engine.abort_async = recovered_abort + await omni.abort(state.external_request_id) + assert omni.request_states == {} and state.queue.empty() asyncio.run(run_test()) diff --git a/tests/entrypoints/test_async_omni_prompt_update.py b/tests/entrypoints/test_async_omni_prompt_update.py index f3edeeec72b..256e2f00d1f 100644 --- a/tests/entrypoints/test_async_omni_prompt_update.py +++ b/tests/entrypoints/test_async_omni_prompt_update.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Entrypoint contract tests for ``AsyncOmni.submit_interaction_async``.""" from __future__ import annotations @@ -68,6 +68,18 @@ async def test_submit_interaction_async_maps_external_to_internal_id() -> None: ) +@pytest.mark.asyncio +async def test_submit_interaction_async_ignores_requests_pending_cleanup() -> None: + omni = _make_async_omni() + closed = omni.request_states["external-abc-uuid-1"] + closed.cleanup_pending = True + await omni.submit_interaction_async("external-abc", interaction=_prompt_interaction()) + omni.engine.submit_interaction_async.assert_awaited_once_with( + "external-abc-uuid-2", interaction=_prompt_interaction() + ) + assert omni.request_states["external-abc-uuid-1"] is closed + + @pytest.mark.asyncio async def test_submit_interaction_async_passes_missing_transition_chunks() -> None: omni = _make_async_omni() @@ -99,11 +111,18 @@ async def test_submit_interaction_async_rejects_negative_transition_chunks() -> @pytest.mark.asyncio -async def test_submit_interaction_async_rejects_inactive_request() -> None: +@pytest.mark.parametrize("pending_cleanup", [False, True]) +async def test_submit_interaction_async_rejects_inactive_request(pending_cleanup) -> None: omni = _make_async_omni() - omni.request_states.clear() + if pending_cleanup: + for state in omni.request_states.values(): + state.cleanup_pending = True + request_id = "external-abc" + else: + omni.request_states.clear() + request_id = "missing" with pytest.raises(ValueError, match="No active request"): - await omni.submit_interaction_async("missing", interaction=_prompt_interaction(transition_chunks=2)) + await omni.submit_interaction_async(request_id, interaction=_prompt_interaction(transition_chunks=2)) @pytest.mark.asyncio diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index cf22c79f177..c1d822b3164 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -104,6 +104,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_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py index 87aeb787ebb..eb2628a4dae 100644 --- a/tests/entrypoints/test_stream_finish_reason.py +++ b/tests/entrypoints/test_stream_finish_reason.py @@ -228,12 +228,14 @@ async def result_generator(): ], ) @pytest.mark.asyncio +@pytest.mark.parametrize("pending_cleanup", [False, True]) async def test_streaming_audio_metrics_resolve_replica_id( monkeypatch, output_replica_id, pool_replica_id, expected_replica_id, pool_lookup, + pending_cleanup, ): import vllm_omni.entrypoints.openai.serving_chat as serving_chat_mod from vllm_omni.entrypoints.client_request_state import ClientRequestState @@ -245,7 +247,11 @@ async def test_streaming_audio_metrics_resolve_replica_id( external_request_id="test-req", ) req_state.request_arrival_ts = time.time() - 1.0 - serving_chat.engine_client.request_states = {"internal-req": req_state} + closed = ClientRequestState(request_id="old-req", external_request_id="test-req") + closed.cleanup_pending = True + serving_chat.engine_client.request_states = ( + {"old-req": closed, "internal-req": req_state} if pending_cleanup else {"internal-req": req_state} + ) serving_chat.engine_client.mod_metrics = object() serving_chat.engine_client.engine = MagicMock() serving_chat.engine_client.engine.stage_pools = [MagicMock(), MagicMock(), MagicMock()] @@ -286,6 +292,7 @@ async def result_generator(): assert first_packet_calls[0]["replica_id"] == expected_replica_id assert finalize_calls[0]["stage_id"] == 2 assert finalize_calls[0]["replica_id"] == expected_replica_id + assert req_state.first_audio_ts is not None and closed.first_audio_ts is None # --------------------------------------------------------------------------- diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 980d3d34a3a..c0b55becc79 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -242,6 +242,8 @@ class OrchestratorRequestState: 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 + # Keep collected tokens until cleanup succeeds, including across retries. + pending_abort_outputs: list[OutputMessage] = field(default_factory=list, repr=False) # Wall-clock timestamp when the client-facing engine request was accepted. request_timestamp: float = 0.0 @@ -712,7 +714,6 @@ async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage """ if not request_ids: return [] - abort_outputs: list[OutputMessage] = [] for pool in self.stage_pools: stage_outputs = await pool.abort_requests(request_ids) or [] if bool(getattr(pool, "final_output", False)) and getattr(pool, "stage_type", None) != "diffusion": @@ -732,7 +733,7 @@ async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage replica_id=pool.get_bound_replica_id(orch_req_id), final_output_type=final_output_type, ) - abort_outputs.append( + req_state.pending_abort_outputs.append( OutputMessage( request_id=orch_req_id, stage_id=pool.stage_id, @@ -744,6 +745,13 @@ async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage ) ) pool.release_bindings(request_ids) + abort_outputs = [ + output + for request_id in request_ids + if (state := self.request_states.get(request_id)) is not None + for output in state.pending_abort_outputs + ] + 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 @@ -2707,6 +2715,7 @@ async def _forward_to_next_stage_unguarded( audio_sample_rate=final_pool._infer_audio_sample_rate(), ) 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", @@ -2716,20 +2725,22 @@ async def _forward_to_next_stage_unguarded( final_output_type or "text", final_stage_id, ) + # Later chunk receivers may already be prewarmed. Release them + # before reporting success; failed cleanup reports a request error. + await self._cleanup_request_ids([req_id], abort=True, release_owners=True) + if req_id in self.request_states: + return await self.output_async_queue.put( OutputMessage( request_id=req_id, stage_id=final_stage_id, - replica_id=final_pool.get_bound_replica_id(req_id), + replica_id=final_replica_id, engine_outputs=terminal_output, metrics=None, finished=True, stage_submit_ts=submit_ts, ) ) - # Later chunk receivers may already be prewarmed. Empty bridge - # input means they will never receive a payload or terminal marker. - await self._cleanup_request_ids([req_id], abort=True, release_owners=True) return # Build and submit requests for each input diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index e6ad1f7f489..e432b3b46f4 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -1313,9 +1313,13 @@ async def abort_requests(self, request_ids: list[str]) -> list[tuple[str, Any]]: if client is None: continue await client.abort_requests_async(engine_abort_ids) - if not is_diffusion and self._output_processor is not None: - commit = getattr(self._output_processor, "commit_aborted_request_state", None) - if callable(commit): + + # Keep prefixes available if a later replica abort fails and this pool + # must retry. The caller receives outputs only after the whole pool stops. + if not is_diffusion and self._output_processor is not None: + commit = getattr(self._output_processor, "commit_aborted_request_state", None) + if callable(commit): + for replica_request_ids in request_ids_by_replica.values(): commit(replica_request_ids, internal=False) return abort_outputs diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 097c7aca8d6..8aa863b1d1c 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -187,6 +187,7 @@ async def generate( logger.debug(f"[AsyncOmni] generate() called for request {external_request_id}") input_stream_task: asyncio.Task | None = None + cleanup_complete = False try: _sleeping_tags = getattr(self, "_sleeping_tags", None) if _sleeping_tags: @@ -320,6 +321,7 @@ async def generate( ): yield output + cleanup_complete = True logger.debug(f"[AsyncOmni] Request {request_id} completed") except (asyncio.CancelledError, GeneratorExit): @@ -327,12 +329,12 @@ async def generate( # ASGI cancellation also affects subsequent awaits. Shield the # bounded engine abort before removing the local request state. with anyio.CancelScope(shield=True): - await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S) + cleanup_complete = await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S) logger.info(f"[AsyncOmni] Request {request_id} aborted.") raise except Exception as e: self._record_request_failure_once(request_id, reason="stage_error") - await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S) + cleanup_complete = await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S) logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}") raise finally: @@ -340,6 +342,8 @@ async def generate( input_stream_task.cancel() if admitting: await self._release_generate_admission() + if request_id in self.request_states: + self.request_states[request_id].cleanup_pending = not cleanup_complete self._log_summary_and_cleanup(request_id) async def _release_generate_admission(self) -> None: @@ -647,7 +651,11 @@ async def submit_interaction_async( if stage_meta.stage_type != "diffusion": raise ValueError("interaction requires a diffusion stage") - internal_ids = [s.request_id for s in self.request_states.values() if s.external_request_id == request_id] + internal_ids = [ + s.request_id + for s in self.request_states.values() + if s.external_request_id == request_id and not s.cleanup_pending + ] if not internal_ids: raise ValueError(f"No active request for interaction: {request_id!r}") if len(internal_ids) > 1: diff --git a/vllm_omni/entrypoints/async_omni_base.py b/vllm_omni/entrypoints/async_omni_base.py index c4c1bc4919c..86ca661c1be 100644 --- a/vllm_omni/entrypoints/async_omni_base.py +++ b/vllm_omni/entrypoints/async_omni_base.py @@ -66,7 +66,7 @@ def watch_task(self, task_id: str, expected_count: int = 1) -> asyncio.Future: } return fut - async def resolve(self, ack: OmniACK): + async def resolve(self, ack: OmniACK | dict): tid = getattr(ack, "task_id", None) if tid is None and isinstance(ack, dict): @@ -359,6 +359,11 @@ async def _final_output_loop(): await self.event_resolver.resolve(msg) continue + request_id = getattr(msg, "request_id", None) + req_state = self.request_states.get(request_id) if request_id is not None else None + if req_state is not None and req_state.cleanup_pending: + continue + if isinstance(msg, ErrorMessage): # Route request-scoped errors to that request's queue and # keep the loop alive. A request whose stage replica died @@ -396,6 +401,8 @@ async def _final_output_loop(): except OmniEngineDeadError as e: logger.error("[%s] Engine dead: %s", self._name, e) for req_state in list(self.request_states.values()): + if req_state.cleanup_pending: + continue error_msg = ErrorMessage( error=str(e), fatal=True, @@ -407,6 +414,8 @@ async def _final_output_loop(): except EngineDeadError as e: logger.error("[%s] Engine dead: %s", self._name, e) for req_state in list(self.request_states.values()): + if req_state.cleanup_pending: + continue error_msg = ErrorMessage( error=str(e), fatal=True, @@ -417,6 +426,8 @@ async def _final_output_loop(): except Exception as e: logger.exception("[%s] final_output_loop failed.", self._name) for req_state in list(self.request_states.values()): + if req_state.cleanup_pending: + continue error_msg = ErrorMessage( request_id=req_state.request_id, error=str(e), @@ -433,7 +444,7 @@ async def _abort_internal_requests( request_id: str | Iterable[str], *, timeout: float = ABORT_TIMEOUT_S, - ): + ) -> bool: """Abort request(s) via the Orchestrator given internal request IDs, which take the format -. """ @@ -443,6 +454,7 @@ async def _abort_internal_requests( try: # Unbind generate() if abort_async blocks in the executor. await asyncio.wait_for(self._abort(internal_req_ids, timeout=timeout), timeout=timeout) + return True except TimeoutError: logger.warning( "[AsyncOmni] Timed out aborting %s after %.1fs; " @@ -452,6 +464,7 @@ async def _abort_internal_requests( ) except Exception: logger.exception("[AsyncOmni] Cleanup abort failed for %s", ",".join(internal_req_ids)) + return False async def _abort(self, request_ids: list[str], *, timeout: float | None = None) -> None: """Abort request IDs via the engine and enqueue terminal abort outputs. @@ -473,14 +486,14 @@ async def _abort(self, request_ids: list[str], *, timeout: float | None = None) if req_id is None: continue state = self.request_states.get(req_id) - if state is None: + if state is None or state.cleanup_pending: logger.debug("[AsyncOmni] Dropping abort output for unknown req %s", req_id) continue await state.queue.put(output_msg) delivered.add(req_id) for rid in request_ids: state = self.request_states.get(rid) - if state is not None and rid not in delivered: + if state is not None and not state.cleanup_pending and rid not in delivered: queue = getattr(state, "queue", None) if queue is not None: final_stage_id = getattr(state, "final_stage_id", None) @@ -506,6 +519,9 @@ async def _abort(self, request_ids: list[str], *, timeout: float | None = None) input_stream_task = getattr(state, "input_stream_task", None) if input_stream_task is not None and not input_stream_task.done(): input_stream_task.cancel() + if state is not None and state.cleanup_pending: + self.request_states.pop(rid) + self._publish_request_gauges(len(self.request_states)) if self.log_stats: logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids)) diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index 9a83910521c..c16368902b0 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/vllm_omni/entrypoints/client_request_state.py @@ -23,6 +23,9 @@ def __init__( self.queue = queue if queue is not None else asyncio.Queue() self.metrics: OrchestratorAggregator | None = None self.input_stream_task: asyncio.Task | None = None + # generate() has exited, but engine cleanup still needs an acknowledgment. + # Retain the public ID mapping for abort retries without queuing outputs. + self.cleanup_pending = False # Request-scoped idempotency guard for Prometheus failure counters. self.failure_recorded = False # Wall-clock time at which the user's request arrived in the engine diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index dc5c68eec42..6647146b646 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -450,7 +450,8 @@ def _log_summary_and_cleanup(self, request_id: str, reason: str = "stage_error") request_id, ) finally: - self.request_states.pop(request_id, None) + if req_state is None or not req_state.cleanup_pending: + self.request_states.pop(request_id, None) # Republish gauges so any stale value left by the per-stage # publish in _process_single_result (which runs while the request # is still in self.request_states) is corrected after the pop. diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index aac3d6b8572..d9aed43692c 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -2234,7 +2234,11 @@ async def chat_completion_stream_generator( # `request_id` we have here is the external (user-visible) id, # so resolve via the external_request_id field. req_state = next( - (s for s in self.engine_client.request_states.values() if s.external_request_id == request_id), + ( + s + for s in self.engine_client.request_states.values() + if s.external_request_id == request_id and not s.cleanup_pending + ), None, ) if req_state is not None and req_state_audio_ref is None: From 60eb9f64dc5d83b862c9d1c3c103b699f8d80af4 Mon Sep 17 00:00:00 2001 From: Pinoeer-kingxi <13022943007@163.com> Date: Sat, 10 Oct 2026 21:37:27 -0700 Subject: [PATCH 4/9] [Core] Preserve abort snapshots during late output processing Signed-off-by: Pinoeer-kingxi <13022943007@163.com> --- .../engine/test_stage_pool_collective_rpc.py | 39 +++++++++++++++---- vllm_omni/engine/orchestrator.py | 11 +++++- 2 files changed, 42 insertions(+), 8 deletions(-) diff --git a/tests/engine/test_stage_pool_collective_rpc.py b/tests/engine/test_stage_pool_collective_rpc.py index e76bb2e77eb..832e1257fc9 100644 --- a/tests/engine/test_stage_pool_collective_rpc.py +++ b/tests/engine/test_stage_pool_collective_rpc.py @@ -11,7 +11,8 @@ import pytest -from tests.engine.test_orchestrator import FakeStageClient, _build_request_output +from tests.engine.test_orchestrator import FakeStageClient, _build_request_output, _terminal_engine_core_outputs +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState from vllm_omni.engine.stage_pool import StagePool pytestmark = [pytest.mark.core_model] @@ -177,7 +178,8 @@ def commit_aborted_request_state(self, request_ids, *, internal=False): @pytest.mark.cpu -def test_abort_retry_keeps_all_replica_prefixes_until_the_pool_stops(mocker): +@pytest.mark.parametrize("request_finishes_during_abort", [False, True]) +def test_abort_retry_keeps_all_replica_prefixes_until_the_pool_stops(mocker, request_finishes_during_abort): async def run() -> None: class RecordingOutputProcessor: def __init__(self) -> None: @@ -195,22 +197,45 @@ def abort_requests_collecting_outputs(self, request_ids, *, internal=False, comm def commit_aborted_request_state(self, request_ids, *, internal=False): assert internal is False for rid in request_ids: - self.prefixes.pop(rid) + self.prefixes.pop(rid, None) - first, second = FakeStageClient(), FakeStageClient() + def process_outputs(self, outputs, *_args): + for output in outputs: + if output.finish_reason is not None: + self.prefixes.pop(output.request_id, None) + return SimpleNamespace(request_outputs=[], reqs_to_abort=[]) + + first, second = FakeStageClient(final_output=True), FakeStageClient(final_output=True) mocker.patch.object(second, "abort_requests_async", side_effect=[RuntimeError("lost ack"), None]) processor = RecordingOutputProcessor() pool = StagePool(0, [first, second], output_processor=processor) pool._request_bindings.update({"first": 0, "second": 1}) + orch = Orchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=asyncio.Queue(), + rpc_async_queue=asyncio.Queue(), + stage_pools=[pool], + ) + for rid in ("first", "second"): + orch.request_states[rid] = OrchestratorRequestState( + request_id=rid, final_stage_id=0, final_output_stage_ids={0} + ) + if request_finishes_during_abort: + + async def abort_first(_ids): + await orch._process_llm_stage_outputs(0, 0, _terminal_engine_core_outputs("first"), set()) + + mocker.patch.object(first, "abort_requests_async", side_effect=abort_first) with pytest.raises(RuntimeError, match="lost ack"): - await pool.abort_requests(["first", "second"]) + await orch._cleanup_request_ids(["first", "second"], abort=True) assert processor.prefixes == {"first": [7, 8], "second": [9]} - outputs = await pool.abort_requests(["first", "second"]) - assert [(rid, list(output.outputs[0].token_ids)) for rid, output in outputs] == [ + outputs = await orch._cleanup_request_ids(["first", "second"], abort=True) + assert [(msg.request_id, list(msg.engine_outputs.outputs[0].token_ids)) for msg in outputs] == [ ("first", [7, 8]), ("second", [9]), ] assert not processor.prefixes + assert not orch.request_states asyncio.run(run()) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index c0b55becc79..486098bb7ec 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -906,7 +906,7 @@ async def _process_llm_stage_outputs( ) request_id = getattr(eco, "request_id", None) req_state = self.request_states.get(request_id) if request_id is not None else None - if req_state is None: + if req_state is None or req_state.closing: continue if ( getattr(eco, "finish_reason", None) == FinishReason.ERROR @@ -936,6 +936,15 @@ async def _process_llm_stage_outputs( if await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state): raw_terminal_request_ids.add(req_state.request_id) await self._report_duplex_session_request_error(stage_id, replica_id, eco, req_state) + # An in-flight abort owns the processor snapshot. Late terminal outputs + # must not delete it before a failed stage/replica abort can retry. + raw_outputs.outputs = [ + eco + for eco in raw_outputs.outputs + if (request_id := getattr(eco, "request_id", None)) is None + or (state := self.request_states.get(request_id)) is None + or not state.closing + ] iteration_stats = IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None processed = await pool.process_llm_raw_outputs( replica_id, From 486b3b32bb6f1bc2a6e3cd7a1456f814537ee34f Mon Sep 17 00:00:00 2001 From: Pinoeer-kingxi <13022943007@163.com> Date: Sat, 17 Oct 2026 19:21:06 -0700 Subject: [PATCH 5/9] [Core] Fence request admission and confirm transport cleanup Signed-off-by: Pinoeer-kingxi <13022943007@163.com> --- docs/design/module/engine_orchestration.md | 14 + .../test_chunk_transfer_adapter.py | 56 +- tests/engine/test_cfg_companion_lifecycle.py | 10 +- tests/engine/test_orchestrator.py | 148 ++++- .../test_orchestrator_async_chunk_prewarm.py | 21 +- .../test_orchestrator_error_handling.py | 6 +- tests/engine/test_orchestrator_first_audio.py | 11 +- .../test_orchestrator_kv_sender_info.py | 28 +- .../test_orchestrator_stage_input_bridge.py | 16 +- tests/engine/test_request_cleanup.py | 346 +++++++++++ tests/engine/test_stage_engine_core_proc.py | 33 + .../engine/test_stage_pool_collective_rpc.py | 88 ++- tests/entrypoints/test_async_omni.py | 190 +++++- .../test_async_omni_prompt_update.py | 27 +- .../entrypoints/test_stream_finish_reason.py | 9 +- tests/worker/test_omni_connector_mixin.py | 19 + tests/worker_v2/test_omni_data_plane.py | 103 ++++ .../omni_connector_payload_transport.py | 30 +- .../model_runner/omni_connector_runtime.py | 26 +- .../chunk_transfer_adapter.py | 95 ++- .../omni_connectors/utils/config.py | 23 +- vllm_omni/engine/errors.py | 4 + vllm_omni/engine/orchestrator.py | 570 ++++++++++++------ vllm_omni/engine/stage_engine_core_proc.py | 18 +- vllm_omni/engine/stage_pool.py | 87 +-- vllm_omni/entrypoints/async_omni.py | 14 +- vllm_omni/entrypoints/async_omni_base.py | 24 +- vllm_omni/entrypoints/client_request_state.py | 3 + vllm_omni/entrypoints/omni_base.py | 3 +- vllm_omni/entrypoints/openai/serving_chat.py | 6 +- vllm_omni/worker/mixins.py | 8 + vllm_omni/worker_v2/omni_data_plane.py | 35 +- 32 files changed, 1733 insertions(+), 338 deletions(-) create mode 100644 tests/engine/test_request_cleanup.py diff --git a/docs/design/module/engine_orchestration.md b/docs/design/module/engine_orchestration.md index c4dd4bb3699..63bc4ce2694 100644 --- a/docs/design/module/engine_orchestration.md +++ b/docs/design/module/engine_orchestration.md @@ -137,6 +137,20 @@ physical EngineCore abort succeeds. Changes that cross into stage runtime or public error behavior require review from that contract's owners. +### Request cleanup + +Cleanup closes request admission before stopping stages and releasing transport +resources. Queued release markers follow pending publications. Backends exposing +`wait_for_cleanup` confirm their ownership drain; other backends retain their +existing physical-release contract. V1 async chunks use the scheduler adapter, +while native MRv2 payloads use the worker data plane. + +Terminal success and abort acknowledgement wait for cleanup. A failed release +reports a request-scoped error and retains the closed state and public request-ID +mapping for retry. Concurrent callers join the same cleanup operation; cancelling +a joining caller does not cancel its owner. CFG ownership remains registered +until the parent and companions finish cleanup. + ## Promotion gate - Reconcile current names and boundaries after #5441 merges, closes, or is diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index f0b8d45b59a..950fd5f6548 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -21,6 +21,7 @@ from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct from vllm_omni.distributed.omni_connectors.adapter import construct_next_stage_streaming_input_prompt +from vllm_omni.distributed.omni_connectors.connectors.base import OmniConnectorBase from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import ( @@ -31,6 +32,27 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_pipeline_release_preserves_legacy_backend_and_fences_queued_puts(build_adapter, mocker): + adapter, _ = build_adapter(stage_id=0) + connector = mocker.Mock(spec=OmniConnectorBase) + connector.stage_id = 0 + connector.put.return_value = (True, 1, {}) + adapter.connector = connector + adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct() + request = _req("internal", RequestStatus.RUNNING, external_req_id="external") + adapter.save_async(None, request) + completion = adapter.release_request_resources("external") + assert not completion.done() + connector.put.assert_not_called() + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + connector.put.assert_called_once() + assert not completion.done() + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + completion.result(timeout=0) + assert not adapter._pending_save_reqs + connector.cleanup.assert_not_called() + + @pytest.fixture def shm_sender(build_adapter): adapter, _ = build_adapter(stage_id=0) @@ -97,6 +119,36 @@ def blocking_put(**kwargs): assert not adapter.code_prompt_token_ids +def test_connector_release_is_queued_after_payload_put(build_adapter): + adapter, connector = build_adapter(stage_id=0) + events = [] + + def put(**kwargs): + events.append("put") + return True, 1, {} + + def cleanup_prefix(prefix): + events.append(prefix) + return 1 + + connector.put.side_effect = put + connector.cleanup_prefix.side_effect = cleanup_prefix + adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct( + codes=CodesStruct(audio=torch.tensor([1])) + ) + request = _req("internal", RequestStatus.RUNNING, external_req_id="external") + adapter.request_ids_mapping[request.request_id] = request.external_req_id + adapter.save_async(None, request) + completion = adapter.release_request_resources(request.request_id) + assert completion is not None and not completion.done() + assert events == [] + while adapter._pending_save_reqs: + adapter._send_single_request(adapter._pending_save_reqs.popleft()) + assert events == ["put", "external_0_"] + completion.result(timeout=0) + assert adapter._reclaimed_payload_total == 1 + + def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender): adapter, connector = shm_sender ext_id = f"release_{uuid.uuid4().hex}" @@ -109,13 +161,13 @@ def test_shm_release_reclaims_undrained_chunks_after_finish(shm_sender): # The consumer read only the first chunk before its own request ended. assert connector.get("0", "1", f"{ext_id}_0_0") is not None - adapter.release_shm_resources(ext_id) + adapter.release_request_resources(ext_id) # The release is queued behind sends; nothing is unlinked on the caller. assert f"{ext_id}_0_1" in connector._pending_keys while adapter._pending_save_reqs: adapter._send_single_request(adapter._pending_save_reqs.popleft()) - assert adapter._reclaimed_shm_total == 2 + assert adapter._reclaimed_payload_total == 2 assert not any(k.startswith(f"{ext_id}_0_") for k in connector._pending_keys) assert connector.get("0", "1", f"{ext_id}_0_1") is None assert connector.get("0", "1", sibling_key)[0] == "sibling" diff --git a/tests/engine/test_cfg_companion_lifecycle.py b/tests/engine/test_cfg_companion_lifecycle.py index 7fa04d1912c..bd64c6719d7 100644 --- a/tests/engine/test_cfg_companion_lifecycle.py +++ b/tests/engine/test_cfg_companion_lifecycle.py @@ -21,7 +21,7 @@ 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] @@ -113,12 +113,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, request_id: str = "p") -> OrchestratorRequestState: + return OrchestratorRequestState( + request_id=request_id, 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, prompt={"prompt": "x"}, pipeline_timings={}, @@ -131,7 +131,7 @@ async def test_kv_ready_does_not_mark_companion_done(): output (set_companion_output) may. Guards the original race.""" orch = _make_orchestrator() orch._cfg_tracker.register_companion("p", "neg", "p__neg") - orch.request_states["p__neg"] = _req_state() + orch.request_states["p__neg"] = _req_state(request_id="p__neg") raw = SimpleNamespace(request_id="p__neg", kv_transfer_params={"kv_ready": True}) await orch._handle_kv_ready_raw_outputs(0, SimpleNamespace(outputs=[raw])) diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 4a5bc435c8b..d33b9dbc497 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -85,6 +85,36 @@ async def fail_outputs() -> None: assert "stage engine died" in fatal.error +@pytest.mark.asyncio +async def test_shutdown_cancels_membership_watcher_and_closes_stages(mocker) -> None: + watcher = asyncio.create_task(asyncio.Event().wait()) + membership = mocker.Mock() + membership.start.return_value = watcher + membership.drain_tasks = mocker.AsyncMock() + requests: asyncio.Queue = asyncio.Queue() + requests.put_nowait(ShutdownRequestMessage()) + orchestrator = Orchestrator( + request_async_queue=requests, + output_async_queue=asyncio.Queue(), + rpc_async_queue=asyncio.Queue(), + stage_pools=[], + membership_controller=membership, + ) + close_stages = mocker.spy(orchestrator, "_shutdown_stages") + # run() owns its loop's tasks; keep the guard on this same task. + task = asyncio.current_task() + assert task is not None + guard = asyncio.get_running_loop().call_later(2, task.cancel) + try: + await orchestrator.run() + finally: + guard.cancel() + assert watcher.cancelled() + membership.drain_tasks.assert_awaited_once() + membership.shutdown.assert_called_once() + close_stages.assert_called_once() + + @dataclass class OrchestratorFixture: orchestrator: Orchestrator @@ -132,10 +162,14 @@ def __init__( self.abort_calls: list[list[str]] = [] self.collective_rpc_calls: list[tuple[str, float | None, tuple[Any, ...], dict[str, Any]]] = [] self.shutdown_calls = 0 - self._engine_core_outputs = queue.Queue() - self._diffusion_outputs = queue.Queue() + self.sample_rate: int | None = None + self._engine_core_outputs: queue.Queue[Any] = queue.Queue() + self._diffusion_outputs: queue.Queue[Any] = queue.Queue() # Orchestrator-facing interface. + async def call_utility_async(self, method, request_ids): + assert method == "omni_release_request_resources" + async def add_request_async(self, *args, **kwargs) -> None: self.add_request_calls.append(args) @@ -228,7 +262,7 @@ async def collective_rpc_async( class FakeOutputProcessor: - def __init__(self, *, request_outputs: list[object] | None = None) -> None: + def __init__(self, *, request_outputs: list[Any] | None = None) -> None: self.request_outputs = list(request_outputs or []) self.add_request_calls: list[tuple[tuple[Any, ...], dict[str, Any]]] = [] self.abort_calls: list[list[str]] = [] @@ -359,7 +393,7 @@ def _build_stage_pools( stage_clients: list[list[FakeStageClient]], *, output_processors: list[FakeOutputProcessor] | None = None, - stage_vllm_configs: list[object] | None = None, + stage_vllm_configs: list[Any] | None = None, ) -> list[StagePool]: """Build StagePool list from per-stage replica lists. @@ -389,9 +423,9 @@ def _build_stage_pools( def _build_harness( - stage_clients: list[object], + stage_clients: list[Any], *, - output_processors: list[object] | None = None, + output_processors: list[Any] | None = None, stage_vllm_configs: list[object] | None = None, async_chunk: bool = False, log_stats: bool = False, @@ -1053,6 +1087,103 @@ async def test_run_abort_emits_result_when_rpc_id_set(orchestrator_factory) -> N await _shutdown_orchestrator(orchestrator_fixture) +@pytest.mark.asyncio +@pytest.mark.parametrize("release_fails", [False, True]) +async def test_abort_ack_waits_for_resource_release_and_failed_cleanup_can_retry( + orchestrator_factory, monkeypatch, release_fails +): + fixture = orchestrator_factory([FakeStageClient(final_output=True)]) + obj = fixture.orchestrator + entered = threading.Event() + released = threading.Event() + calls = [] + + async def release(ids): + calls.append(ids) + entered.set() + while not released.is_set(): + await asyncio.sleep(0.001) + if release_fails and len(calls) == 1: + raise TimeoutError("READ ownership remains active") + + monkeypatch.setattr(obj.stage_pools[0], "release_request_resources", release) + try: + await _enqueue_add_request( + fixture, + request_id="owned", + prompt=FakePromptRequest(request_id="owned", prompt_token_ids=[1], resumable=False), + original_prompt={"prompt": "ownership"}, + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + await _wait_for(lambda: len(obj.stage_pools[0].clients[0].add_request_calls) == 1) + fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["owned"], rpc_id="first")) + await _wait_for(entered.is_set) + assert fixture.queues[2].sync_q.empty() + assert obj.request_states["owned"].closing + released.set() + result = await _get_rpc_message(fixture) + assert result.success is (not release_fails) + if release_fails: + assert "READ ownership" in result.error + assert obj.request_states["owned"].closing + fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["owned"], rpc_id="retry")) + result = await _get_rpc_message(fixture) + assert result.rpc_id == "retry" and result.success + assert calls == [["owned"], ["owned"]] + assert "owned" not in obj.request_states + finally: + released.set() + await _shutdown_orchestrator(fixture) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cleanup_trigger", ["completion", "uncorrelated_abort"]) +async def test_automatic_release_failure_is_request_scoped_and_does_not_publish_success( + orchestrator_factory, monkeypatch, cleanup_trigger +): + stage = FakeStageClient(final_output=True) + fixture = orchestrator_factory( + [stage], + output_processors=[FakeOutputProcessor(request_outputs=[_build_request_output("owned")])], + stage_vllm_configs=[SimpleNamespace(model_config=SimpleNamespace(async_chunk=True))], + ) + obj = fixture.orchestrator + original = stage.call_utility_async + + async def fail_release(method, ids): + raise TimeoutError("READ ownership remains active") + + monkeypatch.setattr(stage, "call_utility_async", fail_release) + try: + await _enqueue_add_request( + fixture, + request_id="owned", + prompt=FakePromptRequest(request_id="owned", prompt_token_ids=[1], resumable=False), + original_prompt={"prompt": "ownership"}, + sampling_params_list=[_sampling_params()], + final_stage_id=0, + ) + await _wait_for(lambda: len(stage.add_request_calls) == 1) + if cleanup_trigger == "completion": + stage.push_engine_core_outputs(_engine_core_outputs("owned", 1.0)) + else: + fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["owned"])) + await _wait_for(lambda: not fixture.output_sync_q.empty()) + error = fixture.output_sync_q.get_nowait() + assert isinstance(error, ErrorMessage) + assert error.request_id == "owned" and not error.fatal + assert fixture.queues[1].sync_q.empty() + assert obj.request_states["owned"].closing + assert fixture.thread.is_alive() + monkeypatch.setattr(stage, "call_utility_async", original) + fixture.request_sync_q.put_nowait(AbortRequestMessage(request_ids=["owned"], rpc_id="retry")) + result = await _get_rpc_message(fixture) + assert result.success and "owned" not in obj.request_states + finally: + await _shutdown_orchestrator(fixture) + + @pytest.mark.asyncio async def test_abort_marks_only_last_final_stage_output_finished() -> None: """A request with two final AR stages must keep earlier abort outputs unfinished.""" @@ -1769,8 +1900,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]], @@ -1985,6 +2116,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.request_states = {} with pytest.raises(RuntimeError, match="stage abort failed"): await orchestrator._abort_request_ids(["req-a"]) diff --git a/tests/engine/test_orchestrator_async_chunk_prewarm.py b/tests/engine/test_orchestrator_async_chunk_prewarm.py index 76a25f15a43..2e190b69fbe 100644 --- a/tests/engine/test_orchestrator_async_chunk_prewarm.py +++ b/tests/engine/test_orchestrator_async_chunk_prewarm.py @@ -10,6 +10,7 @@ from __future__ import annotations +import asyncio import time from types import SimpleNamespace from typing import Any @@ -50,6 +51,8 @@ class _FakePool: """Single-replica LLM pool that records the requests it is sent.""" stage_type = "llm" + final_output = True + num_replicas = 1 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 +64,12 @@ def __init__(self, role: str, stage_client: Any = None) -> None: def live_replica_ids(self) -> list[int]: return [0] + def available_replica_ids(self) -> list[int]: + return [0] + + async def pick(self, request_id, *, request_state=None): + return 0 + async def submit_initial(self, request_id, _req_state, request, prompt_text=None): self.submitted.append(request) return 0 @@ -79,9 +88,14 @@ def _orchestrator(payload_func: Any = None, cls: type[Orchestrator] = Orchestrat """Thinker -> Talker (no payload hook) -> Code2Wav (``payload_func``, if any).""" 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.async_chunk = True - orchestrator.request_states = {} + Orchestrator.__init__( + orchestrator, + request_async_queue=asyncio.Queue(), + output_async_queue=asyncio.Queue(), + rpc_async_queue=asyncio.Queue(), + stage_pools=[_FakePool("sender"), _FakePool("receiver"), _FakePool("receiver", code2wav_client)], + async_chunk=True, + ) orchestrator._emit_tx_edge = lambda **_kwargs: None return orchestrator @@ -194,6 +208,7 @@ async def test_prewarm_submits_plain_placeholder(case: dict[str, Any], expected_ final_stage_id=2, session_owned=case.get("session_owned", False), ) + orchestrator.request_states["req"] = req_state streaming = case.get("streaming", False) req_state.streaming.enabled = streaming stage0_kwargs = case.get("stage0", {"resumable": False}) diff --git a/tests/engine/test_orchestrator_error_handling.py b/tests/engine/test_orchestrator_error_handling.py index 7e1635d191b..e74aaedf479 100644 --- a/tests/engine/test_orchestrator_error_handling.py +++ b/tests/engine/test_orchestrator_error_handling.py @@ -573,7 +573,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 @@ -1182,7 +1182,7 @@ async def test_diffusion_client_error_output_propagates_non_400_status( @pytest.mark.asyncio -async def test_control_rpc_failure_is_reported_to_the_caller_not_fatal() -> None: +async def test_control_rpc_failure_is_reported_to_the_caller_not_fatal(mocker) -> None: """A control RPC that fails on one replica must come back as that replica's error result instead of raising out of _request_handler. """ @@ -1192,7 +1192,7 @@ async def collective_rpc_async(self, *args, **kwargs): raise TimeoutError("outputs did not drain") ok = FakeStageClient(stage_type="diffusion") - ok.collective_rpc_async = AsyncMock(return_value=None) + mocker.patch.object(ok, "collective_rpc_async", new_callable=AsyncMock, return_value=None) orchestrator, queues = _build_bare_orchestrator( _build_stage_pools([[ok], [_FailingClient(stage_type="diffusion")]]) ) diff --git a/tests/engine/test_orchestrator_first_audio.py b/tests/engine/test_orchestrator_first_audio.py index bc99cf7fbef..1e73e2331b3 100644 --- a/tests/engine/test_orchestrator_first_audio.py +++ b/tests/engine/test_orchestrator_first_audio.py @@ -57,6 +57,10 @@ def register(self): def get_bound_replica_id(self, request_id): return self.bound_replica + async def pick(self, request_id, *, request_state=None): + self.bound_replica = self.replica_id + return self.replica_id + async def process_llm_raw_outputs(self, replica_id, raw, **kwargs): assert replica_id == self.replica_id return self.output_processor.process_outputs(raw.outputs, raw.timestamp).request_outputs @@ -72,6 +76,7 @@ def _orchestrator(output_kind=RequestOutputKind.DELTA, *, registered=True, repli obj._running_counter = None obj._abort_request_ids = AsyncMock(return_value=[]) obj._release_request_bindings = Mock() + obj._release_stage_transfer_resources = AsyncMock() obj._finish_raw_terminal_requests = AsyncMock() async def process(stage, replica, raw, terminals): @@ -238,6 +243,9 @@ async def test_first_audio_waits_for_codec_registration_and_ignores_duplicates() @pytest.mark.parametrize("processing_failed", [False, True]) async def test_prewarm_flushes_first_audio_after_successful_submission(mocker, processing_failed): obj = _orchestrator(registered=False, final_stage_id=2 if processing_failed else 1) + if processing_failed: + later_pool = SimpleNamespace(submit_initial=AsyncMock()) + obj.stage_pools.append(later_pool) req_state = obj.request_states["r"] req_state.sampling_params_list = [None] * (req_state.final_stage_id + 1) await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11))) @@ -260,8 +268,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()) - 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) @@ -314,6 +320,7 @@ 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)) 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) diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index 5f3d1147145..00460640501 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -5,6 +5,7 @@ from collections.abc import Callable from types import SimpleNamespace from typing import Any +from unittest.mock import AsyncMock import pytest from vllm import SamplingParams @@ -366,6 +367,7 @@ def test_forward_to_diffusion_attaches_kv_sender_info(): ) output = SimpleNamespace(request_id="req-1", finished=True) + sender_pool.select_replica_id("req-1") asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-1", sender_pool.stage_id, output, req_state)) assert diffusion_stage.calls[0]["request_id"] == "req-1" @@ -392,6 +394,7 @@ def test_diffusion_resubmission_preserves_sender_endpoints(): final_stage_id=1, ) output = SimpleNamespace(request_id="streaming", finished=True) + sender_pool.select_replica_id("streaming") async def forward_twice(): await orchestrator._forward_to_next_stage("streaming", 0, output, state) @@ -425,6 +428,8 @@ def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info(): ) output = SimpleNamespace(request_id="req-3", finished=True) + source_pool.select_replica_id("req-3") + previous_pool.select_replica_id("req-3") asyncio.run(Orchestrator._forward_to_next_stage(orchestrator, "req-3", previous_pool.stage_id, output, req_state)) assert diffusion_stage.calls[0]["kv_sender_info"] == { @@ -494,6 +499,8 @@ def test_prewarm_diffusion_attaches_kv_sender_info(): ) stage0_request = SimpleNamespace(prompt_token_ids=[1, 2, 3]) + orchestrator.request_states = {req_state.request_id: req_state} + sender_pool.select_replica_id(req_state.request_id) asyncio.run(Orchestrator._prewarm_async_chunk_stages(orchestrator, "req-2", stage0_request, req_state)) assert diffusion_stage.calls[0]["request_id"] == "req-2" @@ -522,6 +529,7 @@ async def submit(key, state, request, **kwargs): stage_vllm_config=SimpleNamespace(model_config=SimpleNamespace(hf_config=None, max_model_len=64)), submit_initial=submit, get_bound_replica_id=lambda key: {"a": 7, "b": 1}[key], + pick=AsyncMock(side_effect=lambda key, **kwargs: {"a": 7, "b": 1}[key]), ) orchestrator.stage_pools = [source, target] orchestrator._stage_receives_async_chunks = lambda stage: True @@ -535,17 +543,21 @@ async def dispatch(callback, **kwargs): orchestrator._dispatch_or_fail_request = dispatch async def run(): + orchestrator.request_states = { + key: OrchestratorRequestState( + request_id=key, + prompt={}, + sampling_params_list=[SamplingParams(), SamplingParams()], + final_stage_id=1, + ) + for key in endpoints + } 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, - ), + orchestrator.request_states[key], ) for key in endpoints ] @@ -595,6 +607,10 @@ async def test_async_pipeline_defers_sync_stage_until_audio_finishes(mocker): final_output_stage_ids={0, 1}, ) + orchestrator.request_states = {state.request_id: state} + orchestrator._cleanup_request_ids.side_effect = lambda ids: [ + orchestrator.request_states.pop(request_id, None) for request_id in ids + ] 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 3c1dd6d2be3..9c49f577b8d 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -46,7 +46,7 @@ def __init__( self.next_inputs = list(next_inputs or []) self.add_request_calls: list[tuple[Any, ...]] = [] self.decoded_source_tokens: str | None = None - self._engine_core_outputs = queue.Queue() + self._engine_core_outputs: queue.Queue[Any] = queue.Queue() async def add_request_async(self, *args, **_kwargs) -> None: self.add_request_calls.append(args) @@ -99,7 +99,7 @@ def process_inputs(self, **kwargs): class FakePrewarmPool: - """Single-replica pool double that binds requests only on submission.""" + """Single-replica pool double that binds on selection or submission.""" stage_type = "llm" @@ -119,6 +119,10 @@ async def submit_initial(self, request_id, _req_state, request, prompt_text=None self._bound_request_ids.add(request_id) return 0 + async def pick(self, request_id, *, request_state=None): + self._bound_request_ids.add(request_id) + return 0 + def get_bound_replica_id(self, request_id): return 0 if request_id in self._bound_request_ids else None @@ -224,6 +228,8 @@ async def test_async_prewarm_skips_outgoing_only_stage(payload_sender_info) -> N final_stage_id=2, ) + orchestrator.request_states = {req_state.request_id: req_state} + prewarmed = await orchestrator._prewarm_async_chunk_stages( "req-prewarm", SimpleNamespace(prompt_token_ids=[1, 2], resumable=True), @@ -232,7 +238,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 @@ -272,6 +278,8 @@ async def test_async_prewarm_skips_stage_with_custom_process_input_func() -> Non final_stage_id=2, ) + orchestrator.request_states = {req_state.request_id: req_state} + prewarmed = await orchestrator._prewarm_async_chunk_stages( "req-prewarm-custom", SimpleNamespace(prompt_token_ids=[1, 2], resumable=True), @@ -340,7 +348,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_request_cleanup.py b/tests/engine/test_request_cleanup.py new file mode 100644 index 00000000000..1f3831d9192 --- /dev/null +++ b/tests/engine/test_request_cleanup.py @@ -0,0 +1,346 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Request admission, abort ownership and cleanup retries on the default route.""" + +import asyncio +from types import SimpleNamespace + +import pytest +from vllm.sampling_params import SamplingParams + +from tests.engine.test_orchestrator import ( + FakePromptRequest, + FakeStageClient, + _build_request_output, + _build_stage_pools, +) +from vllm_omni.engine.errors import ResourceReleaseError +from vllm_omni.engine.messages import ErrorMessage, OutputMessage, StageSubmissionMessage +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState +from vllm_omni.engine.stage_pool import RequestClosedError, StageUnavailableError + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _direct_orchestrator(*, async_chunk=False): + final_stage_id = 2 + 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, + 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.parametrize("selected", [set(), {0, 1}]) +def test_completion_derives_selected_outputs_from_stage_termination(selected): + state = OrchestratorRequestState( + request_id="r", final_stage_id=1, final_output_stage_ids=selected, finished_stage_ids={2} + ) + assert not state.final_outputs_finished + state.finished_stage_ids.add(0) + assert not state.final_outputs_finished + state.finished_stage_ids.add(1) + assert state.final_outputs_finished + + +@pytest.mark.asyncio +async def test_cancel_while_waiting_for_replica_stops_old_request_admission(mocker): + 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_duplicate_cleanup_waiter_does_not_cancel_abort_owner(mocker): + orch, state, clients = _direct_orchestrator(async_chunk=True) + orch.stage_pools[0].select_replica_id("r") + started, finish = asyncio.Event(), asyncio.Event() + + async def abort(_ids): + started.set() + await finish.wait() + + abort_call = mocker.patch.object(clients[0], "abort_requests_async", side_effect=abort) + owner = asyncio.create_task(orch._cleanup_request_ids(["r"], abort=True)) + await asyncio.wait_for(started.wait(), timeout=1) + waiter = asyncio.create_task(orch._cleanup_request_ids(["r"], abort=True)) + await asyncio.sleep(0) + waiter.cancel() + with pytest.raises(asyncio.CancelledError): + await waiter + assert state.closing and not state.cleanup_future.done() + with pytest.raises(RequestClosedError): + await orch.stage_pools[0].submit_update("r", state, SimpleNamespace()) + await orch._handle_add_request(_submission()) + error = orch.output_async_queue.get_nowait() + assert isinstance(error, ErrorMessage) and error.status_code == 409 + assert orch.request_states["r"] is state + finish.set() + await asyncio.wait_for(owner, timeout=1) + assert abort_call.await_count == 1 + assert "r" not in orch.request_states + + +@pytest.mark.asyncio +@pytest.mark.parametrize("own_release_fails", [False, True]) +async def test_automatic_cleanup_joins_failed_owner_and_reports_only_failed_requests(mocker, own_release_fails): + orch, state, _ = _direct_orchestrator() + other_state = OrchestratorRequestState(request_id="s") + orch.request_states["s"] = other_state + entered, second_entered, finish = asyncio.Event(), asyncio.Event(), asyncio.Event() + failures = {"r", "s"} if own_release_fails else {"r"} + + async def release(ids): + if ids == ["r"]: + entered.set() + await finish.wait() + else: + second_entered.set() + if failures.intersection(ids): + raise ResourceReleaseError("READ ownership remains active") + + mocker.patch.object(orch.stage_pools[0], "release_request_resources", side_effect=release) + owner = asyncio.create_task(orch._cleanup_request_ids(["r"], abort=True)) + await asyncio.wait_for(entered.wait(), timeout=1) + automatic = asyncio.create_task(orch._cleanup_request_ids(["r", "s"])) + await asyncio.wait_for(second_entered.wait(), timeout=1) + await asyncio.sleep(0) + assert not automatic.done() + finish.set() + with pytest.raises(ResourceReleaseError): + await asyncio.wait_for(owner, timeout=1) + await asyncio.wait_for(automatic, timeout=1) + errors = [] + while not orch.output_async_queue.empty(): + errors.append(orch.output_async_queue.get_nowait()) + assert {error.request_id for error in errors} == failures + assert all(isinstance(error, ErrorMessage) and not error.fatal for error in errors) + assert state.closing and orch.request_states["r"] is state + assert ("s" in orch.request_states) is own_release_fails + failures.clear() + await orch._cleanup_request_ids(["r", "s"], abort=True) + assert not orch.request_states + + +@pytest.mark.asyncio +@pytest.mark.parametrize("release_fails", [False, True]) +async def test_parent_cleanup_joins_completed_companion_cleanup(mocker, release_fails): + orch, state, _ = _direct_orchestrator() + companion_id = "r__neg" + companion = OrchestratorRequestState(request_id=companion_id) + orch.request_states[companion_id] = companion + orch._cfg_tracker.register_companion("r", "neg", companion_id) + orch._cfg_tracker.set_companion_output(companion_id, {"conditioning": 1}) + orch._cfg_tracker.on_companion_completed(companion_id) + entered, parent_entered, finish = asyncio.Event(), asyncio.Event(), asyncio.Event() + + async def release(ids): + if ids == [companion_id]: + entered.set() + await finish.wait() + if release_fails: + raise ResourceReleaseError("READ ownership remains active") + else: + parent_entered.set() + + release_call = mocker.patch.object(orch.stage_pools[0], "release_request_resources", side_effect=release) + owner = asyncio.create_task(orch._cleanup_request_ids([companion_id])) + await asyncio.wait_for(entered.wait(), timeout=1) + companion_future = companion.cleanup_future + parent = asyncio.create_task(orch._cleanup_request_ids(["r"], abort=True)) + try: + await asyncio.wait_for(parent_entered.wait(), timeout=1) + assert companion.cleanup_future is companion_future + assert orch.request_states["r"] is state and state.closing + assert not parent.done() + finally: + finish.set() + results = await asyncio.wait_for(asyncio.gather(owner, parent, return_exceptions=True), timeout=1) + assert [call.args[0] for call in release_call.await_args_list] == [[companion_id], ["r"]] + assert results[0] == [] + if release_fails: + assert isinstance(results[1], ResourceReleaseError) + assert orch.request_states["r"] is state and state.closing + assert companion.closing and orch._cfg_tracker.has_companions("r") + release_fails = False + await orch._cleanup_request_ids(["r"], abort=True) + assert set(release_call.await_args_list[-1].args[0]) == {"r", companion_id} + else: + assert results[1] == [] + assert not orch.request_states + + +@pytest.mark.asyncio +@pytest.mark.parametrize("retry_id", ["r", "r__neg"]) +async def test_failed_cfg_cleanup_retry_releases_the_complete_owner_batch(mocker, retry_id): + orch, state, _ = _direct_orchestrator() + companion_id = "r__neg" + companion = OrchestratorRequestState(request_id=companion_id) + orch.request_states[companion_id] = companion + orch._cfg_tracker.register_companion("r", "neg", companion_id) + release_call = mocker.patch.object( + orch.stage_pools[0], + "release_request_resources", + side_effect=ResourceReleaseError("READ ownership remains active"), + ) + with pytest.raises(ResourceReleaseError): + await orch._cleanup_request_ids(["r"], abort=True) + assert state.closing and companion.closing + assert state.cleanup_future is companion.cleanup_future + assert orch._cfg_tracker.has_companions("r") + release_call.side_effect = None + await orch._cleanup_request_ids([retry_id], abort=True) + assert not orch.request_states + assert all(set(call.args[0]) == {"r", companion_id} for call in release_call.await_args_list) + + +@pytest.mark.asyncio +async def test_failed_abort_can_be_retried_without_reopening_admission(mocker): + orch, state, clients = _direct_orchestrator() + orch.stage_pools[0].select_replica_id("r") + abort = mocker.patch.object(clients[0], "abort_requests_async", side_effect=[RuntimeError("RPC unavailable"), None]) + with pytest.raises(RuntimeError, match="RPC unavailable"): + await orch._cleanup_request_ids(["r"], abort=True) + assert state.closing and orch.request_states["r"] is state + with pytest.raises(RequestClosedError): + await orch.stage_pools[0].pick("r", request_state=state) + await orch._cleanup_request_ids(["r"], abort=True) + assert abort.await_count == 2 + assert "r" not in orch.request_states + + +@pytest.mark.asyncio +async def test_lost_preselected_replica_cannot_be_reselected(mocker): + 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) + with pytest.raises(StageUnavailableError): + await pool.submit_initial("r", state, SimpleNamespace()) + await orch._handle_dead_replica(2, 0, RuntimeError("replica exited")) + 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("release_fails", [False, True]) +async def test_empty_bridge_waits_for_cleanup_before_publishing_success(mocker, release_fails): + orch, state, _ = _direct_orchestrator(async_chunk=True) + entered, resume = asyncio.Event(), asyncio.Event() + + async def release(_ids): + entered.set() + await resume.wait() + if release_fails: + raise ResourceReleaseError("receiver still owns the payload") + + mocker.patch.object(orch.stage_pools[0], "release_request_resources", side_effect=release) + task = asyncio.create_task(orch._forward_to_next_stage("r", 0, _build_request_output("r"), state)) + try: + await asyncio.wait_for(entered.wait(), timeout=1) + assert orch.output_async_queue.empty() + assert state.closing and not state.cleanup_future.done() + finally: + resume.set() + await asyncio.wait_for(task, timeout=1) + result = orch.output_async_queue.get_nowait() + if release_fails: + assert isinstance(result, ErrorMessage) + assert orch.request_states["r"] is state + else: + assert isinstance(result, OutputMessage) and result.finished + assert "r" not in orch.request_states + assert orch.output_async_queue.empty() + + +@pytest.mark.asyncio +async def test_abort_retry_preserves_tokens_after_transfer_release_failure(mocker): + orch, state, clients = _direct_orchestrator() + pool = orch.stage_pools[2] + pool.select_replica_id("r") + pool.output_processor.request_outputs = [_build_request_output("r", token_ids=[7, 8, 9], finished=False)] + release = mocker.patch.object( + orch.stage_pools[0], + "release_request_resources", + side_effect=[ResourceReleaseError("receiver still owns the payload"), None], + ) + with pytest.raises(ResourceReleaseError): + await orch._cleanup_request_ids(["r"], abort=True) + assert orch.request_states["r"] is state and state.closing + assert pool.get_bound_replica_id("r") is None + outputs = await orch._cleanup_request_ids(["r"], abort=True) + assert len(outputs) == 1 and outputs[0].finished + assert list(outputs[0].engine_outputs.outputs[0].token_ids) == [7, 8, 9] + assert outputs[0].replica_id == 0 + assert clients[2].abort_calls == [["r"]] + assert release.await_count == 2 and "r" not in orch.request_states + + +@pytest.mark.asyncio +async def test_cancelled_prewarm_releases_bindings_without_resurrection(): + orch, state, clients = _direct_orchestrator(async_chunk=True) + await orch._prewarm_async_chunk_stages( + "r", FakePromptRequest(request_id="r", prompt_token_ids=[1, 2], resumable=False), state + ) + await orch._cleanup_request_ids(["r"], abort=True) + assert "r" not in orch.request_states + assert all(pool.get_bound_replica_id("r") is None for pool in orch.stage_pools) + assert not await orch._prewarm_async_chunk_stages( + "r", FakePromptRequest(request_id="r", prompt_token_ids=[1, 2], resumable=False), state + ) + assert len(clients[2].add_request_calls) == 1 diff --git a/tests/engine/test_stage_engine_core_proc.py b/tests/engine/test_stage_engine_core_proc.py index fdb6424817d..3ed3fa0e163 100644 --- a/tests/engine/test_stage_engine_core_proc.py +++ b/tests/engine/test_stage_engine_core_proc.py @@ -47,6 +47,39 @@ def test_decoder_capability_matches_sink_binding(mocker, backend, executor_class assert runner.model_state.set_first_audio_sink.call_count == int(expected) +@pytest.mark.parametrize("has_plane", [False, True]) +def test_native_resource_release_reaches_worker_owned_plane(mocker, has_plane): + from vllm_omni.worker.mixins import OmniWorkerMixin + + plane = mocker.Mock() if has_plane else None + worker = SimpleNamespace(model_runner=SimpleNamespace(_omni_data_plane=plane)) + engine = StageEngineCoreProc.__new__(StageEngineCoreProc) + engine.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(use_v2_model_runner=True, supports_native_mrv2_data_plane=True) + ) + engine.model_executor = mocker.Mock() + engine.model_executor.collective_rpc.side_effect = lambda method, args: getattr(OmniWorkerMixin, method)( + worker, *args + ) + engine.omni_release_request_resources(["external"]) + engine.model_executor.collective_rpc.assert_called_once_with("omni_release_request_resources", args=(["external"],)) + if has_plane: + plane.release_request_resources.assert_called_once_with(["external"]) + + +def test_legacy_resource_release_reaches_scheduler_adapter(mocker): + engine = StageEngineCoreProc.__new__(StageEngineCoreProc) + engine.vllm_config = SimpleNamespace(model_config=SimpleNamespace(use_v2_model_runner=False)) + adapter = mocker.Mock() + engine.scheduler = SimpleNamespace(chunk_transfer_adapter=adapter) + engine.omni_release_request_resources(["external"]) + adapter.release_request_resources.assert_called_once() + args, kwargs = adapter.release_request_resources.call_args + assert args == ("external",) + assert kwargs["deadline"] > 0 + adapter.release_request_resources.return_value.result.assert_called_once() + + @pytest.mark.parametrize( "first_decoder,stream_decoder,stream_first_audio,expected", [(True, False, False, True), (False, True, False, False), (False, True, True, True), (False, False, True, False)], diff --git a/tests/engine/test_stage_pool_collective_rpc.py b/tests/engine/test_stage_pool_collective_rpc.py index aa7bbcb397c..832e1257fc9 100644 --- a/tests/engine/test_stage_pool_collective_rpc.py +++ b/tests/engine/test_stage_pool_collective_rpc.py @@ -11,6 +11,8 @@ import pytest +from tests.engine.test_orchestrator import FakeStageClient, _build_request_output, _terminal_engine_core_outputs +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState from vllm_omni.engine.stage_pool import StagePool pytestmark = [pytest.mark.core_model] @@ -175,6 +177,69 @@ def commit_aborted_request_state(self, request_ids, *, internal=False): asyncio.run(run()) +@pytest.mark.cpu +@pytest.mark.parametrize("request_finishes_during_abort", [False, True]) +def test_abort_retry_keeps_all_replica_prefixes_until_the_pool_stops(mocker, request_finishes_during_abort): + async def run() -> None: + class RecordingOutputProcessor: + def __init__(self) -> None: + self.prefixes = {"first": [7, 8], "second": [9]} + + def abort_requests_collecting_outputs(self, request_ids, *, internal=False, commit_state=True): + assert internal is False and commit_state is False + outputs = [ + _build_request_output(rid, token_ids=self.prefixes[rid]) + for rid in request_ids + if rid in self.prefixes + ] + return list(request_ids), outputs + + def commit_aborted_request_state(self, request_ids, *, internal=False): + assert internal is False + for rid in request_ids: + self.prefixes.pop(rid, None) + + def process_outputs(self, outputs, *_args): + for output in outputs: + if output.finish_reason is not None: + self.prefixes.pop(output.request_id, None) + return SimpleNamespace(request_outputs=[], reqs_to_abort=[]) + + first, second = FakeStageClient(final_output=True), FakeStageClient(final_output=True) + mocker.patch.object(second, "abort_requests_async", side_effect=[RuntimeError("lost ack"), None]) + processor = RecordingOutputProcessor() + pool = StagePool(0, [first, second], output_processor=processor) + pool._request_bindings.update({"first": 0, "second": 1}) + orch = Orchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=asyncio.Queue(), + rpc_async_queue=asyncio.Queue(), + stage_pools=[pool], + ) + for rid in ("first", "second"): + orch.request_states[rid] = OrchestratorRequestState( + request_id=rid, final_stage_id=0, final_output_stage_ids={0} + ) + if request_finishes_during_abort: + + async def abort_first(_ids): + await orch._process_llm_stage_outputs(0, 0, _terminal_engine_core_outputs("first"), set()) + + mocker.patch.object(first, "abort_requests_async", side_effect=abort_first) + with pytest.raises(RuntimeError, match="lost ack"): + await orch._cleanup_request_ids(["first", "second"], abort=True) + assert processor.prefixes == {"first": [7, 8], "second": [9]} + outputs = await orch._cleanup_request_ids(["first", "second"], abort=True) + assert [(msg.request_id, list(msg.engine_outputs.outputs[0].token_ids)) for msg in outputs] == [ + ("first", [7, 8]), + ("second", [9]), + ] + assert not processor.prefixes + assert not orch.request_states + + asyncio.run(run()) + + @pytest.mark.cpu def test_release_request_resources_skips_without_async_chunk(): async def run() -> None: @@ -191,6 +256,25 @@ async def run() -> None: asyncio.run(run()) +@pytest.mark.cpu +def test_release_request_resources_includes_native_full_payload(): + async def run() -> None: + call = AsyncMock() + pool = StagePool( + 0, + [SimpleNamespace(call_utility_async=call)], # type: ignore[list-item] + stage_vllm_config=SimpleNamespace( + model_config=SimpleNamespace( + async_chunk=False, use_v2_model_runner=True, supports_native_mrv2_data_plane=True + ) + ), + ) + await pool.release_request_resources(["req-1"]) + call.assert_awaited_once_with("omni_release_request_resources", ["req-1"]) + + asyncio.run(run()) + + @pytest.mark.cpu def test_release_request_resources_times_out_hung_replica_without_blocking_others(monkeypatch): async def run() -> None: @@ -210,7 +294,9 @@ async def hang(*_args): live.assert_awaited_once_with("omni_release_request_resources", ["req-1"]) assert not release.done() - await asyncio.wait_for(release, timeout=1.0) + with pytest.raises(RuntimeError, match="replica 0 resource release failed") as error: + await asyncio.wait_for(release, timeout=1.0) + assert isinstance(error.value.__cause__, TimeoutError) asyncio.run(run()) diff --git a/tests/entrypoints/test_async_omni.py b/tests/entrypoints/test_async_omni.py index 829e30145d1..a32a8b6880c 100644 --- a/tests/entrypoints/test_async_omni.py +++ b/tests/entrypoints/test_async_omni.py @@ -12,11 +12,16 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from vllm.v1.kv_hints import KvHintsEnvelope +from tests.engine.test_request_cleanup import _direct_orchestrator from tests.helpers.mark import hardware_test from tests.helpers.stage_config import get_deploy_config_path from vllm_omni.engine.async_omni_engine import AsyncOmniEngine +from vllm_omni.engine.errors import ResourceReleaseError +from vllm_omni.engine.messages import ErrorMessage, OutputMessage +from vllm_omni.engine.orchestrator import OrchestratorRequestState from vllm_omni.entrypoints import async_omni as async_omni_mod from vllm_omni.entrypoints.async_omni import AsyncOmni +from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.utils import coerce_param_message_types from vllm_omni.model_executor.models.qwen3_omni.pipeline import QWEN3_OMNI_PIPELINE from vllm_omni.outputs import OmniRequestOutput @@ -80,10 +85,124 @@ def get_async_omni_instance(fake_add_request=_noop, fake_abort_request=_noop) -> omni._compute_final_stage_id = lambda output_modalities: 0 omni._compute_final_output_stage_ids = lambda output_modalities: [0] omni._process_orchestrator_results = fake_process_results - omni._log_summary_and_cleanup = lambda request_id: omni.request_states.pop(request_id, None) return omni +@pytest.mark.cpu +def test_generate_retains_public_id_until_failed_cleanup_retry_succeeds(): + async def run(): + orch, _, _ = _direct_orchestrator() + orch.request_states.clear() + fail_release = True + abort_calls = [] + + async def release(_ids): + if fail_release: + raise ResourceReleaseError("receiver still owns the payload") + + orch.stage_pools[0].release_request_resources = release + + async def add(*, request_id, sampling_params_list, final_stage_id, final_output_stage_ids, **kwargs): + orch.request_states[request_id] = OrchestratorRequestState( + request_id=request_id, + sampling_params_list=sampling_params_list, + final_stage_id=final_stage_id, + final_output_stage_ids=set(final_output_stage_ids), + ) + await orch._cleanup_request_ids([request_id]) + await omni.request_states[request_id].queue.put(orch.output_async_queue.get_nowait()) + + async def abort(ids, timeout=None): + abort_calls.append(list(ids)) + return await orch._cleanup_request_ids(list(ids), abort=True) + + omni = get_async_omni_instance(fake_add_request=add, fake_abort_request=abort) + omni.engine.num_stages = len(orch.stage_pools) + omni._compute_final_stage_id = lambda _modalities: 2 + omni._compute_final_output_stage_ids = lambda _modalities: [2] + omni._process_orchestrator_results = AsyncOmni._process_orchestrator_results.__get__(omni) + with pytest.raises(RuntimeError, match="receiver still owns the payload"): + async for _ in omni.generate( + prompt={"prompt": "test"}, + request_id="public-r", + sampling_params_list=[SamplingParams(max_tokens=1) for _ in orch.stage_pools], + output_modalities=["text"], + ): + pass + (internal_id,) = orch.request_states + state = omni.request_states[internal_id] + assert state.cleanup_pending and state.external_request_id == "public-r" + assert orch.request_states[internal_id].closing + with pytest.raises(ResourceReleaseError): + await omni.abort("public-r") + assert omni.request_states[internal_id] is state + fail_release = False + await omni.abort("public-r") + assert abort_calls == [[internal_id], [internal_id], [internal_id]] + assert not omni.request_states and not orch.request_states + assert state.queue.empty() + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_abort_retry_keeps_active_requests_with_the_same_public_id(): + async def run(): + calls: list[list[str]] = [] + omni = get_async_omni_instance(fake_abort_request=get_fake_abort(calls)) + closed = ClientRequestState("closed", external_request_id="public") + closed.cleanup_pending = True + active = ClientRequestState("active", external_request_id="public") + unrelated = ClientRequestState("unrelated", external_request_id="other") + omni.request_states = {state.request_id: state for state in (closed, active, unrelated)} + await omni.abort("public") + assert calls == [["closed", "active"]] + assert omni.request_states == {"active": active, "unrelated": unrelated} + assert closed.queue.empty() and unrelated.queue.empty() + terminal = active.queue.get_nowait() + assert terminal.finished and terminal.engine_outputs.outputs[0].finish_reason == "abort" + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_output_dispatcher_drops_messages_for_requests_pending_cleanup(): + async def run(): + omni = get_async_omni_instance() + pending = ClientRequestState("pending") + pending.cleanup_pending = True + active = ClientRequestState("active") + omni.request_states = {"pending": pending, "active": active} + messages: asyncio.Queue = asyncio.Queue() + omni.engine.try_get_output_async = messages.get + omni.final_output_task = None + for _ in range(64): + messages.put_nowait(ErrorMessage(request_id="pending", error="late release failure")) + messages.put_nowait( + OutputMessage( + request_id="pending", + stage_id=0, + engine_outputs=OmniRequestOutput(request_id="pending", stage_id=0, finished=True, outputs=[]), + finished=True, + ) + ) + marker = ErrorMessage(request_id="active", error="active request failure") + messages.put_nowait(marker) + AsyncOmni._final_output_handler(omni) + task = omni.final_output_task + assert task is not None + try: + assert await asyncio.wait_for(active.queue.get(), timeout=1) is marker + assert pending.queue.empty() and not pending.consumed_metric_message_ids + assert omni.request_states["pending"] is pending + finally: + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(run()) + + @pytest.mark.cpu def test_generate_submits_randomized_id_to_engine(): """Ensure the engine receives a UUID-suffixed ID, not the raw request ID""" @@ -243,10 +362,9 @@ async def slow_abort_async(request_ids, timeout=None): await release.wait() omni = get_async_omni_instance(fake_abort_request=slow_abort_async) - omni.request_states["req-1-aaaa"] = SimpleNamespace( + omni.request_states["req-1-aaaa"] = ClientRequestState( request_id="req-1-aaaa", external_request_id="req-1", - input_stream_task=None, ) task = asyncio.create_task(omni.abort("req-1")) @@ -270,10 +388,9 @@ async def failing_abort_async(request_ids, timeout=None): raise RuntimeError("orchestrator abort failed") omni = get_async_omni_instance(fake_abort_request=failing_abort_async) - omni.request_states["req-1-bbbb"] = SimpleNamespace( + omni.request_states["req-1-bbbb"] = ClientRequestState( request_id="req-1-bbbb", external_request_id="req-1", - input_stream_task=None, ) with pytest.raises(RuntimeError, match="orchestrator abort failed"): @@ -327,10 +444,9 @@ async def abort_with_prefix(request_ids, timeout=None): ] omni = get_async_omni_instance(fake_abort_request=abort_with_prefix) - omni.request_states["req-1-cccc"] = SimpleNamespace( + omni.request_states["req-1-cccc"] = ClientRequestState( request_id="req-1-cccc", external_request_id="req-1", - input_stream_task=None, queue=queue, ) await omni._abort(["req-1-cccc"]) @@ -355,10 +471,9 @@ async def empty_abort_async(request_ids, timeout=None): return [] omni = get_async_omni_instance(fake_abort_request=empty_abort_async) - omni.request_states["req-1-dddd"] = SimpleNamespace( + omni.request_states["req-1-dddd"] = ClientRequestState( request_id="req-1-dddd", external_request_id="req-1", - input_stream_task=None, queue=queue, ) await omni._abort(["req-1-dddd"]) @@ -387,7 +502,6 @@ def test_abort_drops_consumed_metric_message_ids_with_state(): from collections.abc import Mapping from vllm_omni.engine.messages import OutputMessage - from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.metrics.stats import OrchestratorAggregator async def run(): @@ -595,7 +709,50 @@ async def collect(): await asyncio.wait_for(task, timeout=1.0) assert seen_timeouts assert seen_timeouts[-1] == 0.05 - assert omni.request_states == {} + (state,) = omni.request_states.values() + assert state.cleanup_pending and state.external_request_id == "cancel-hang-abort" + + asyncio.run(run()) + + +@pytest.mark.cpu +def test_generate_keeps_public_id_if_cleanup_abort_is_cancelled_again(): + async def run(): + submitted, abort_started, resume = asyncio.Event(), asyncio.Event(), asyncio.Event() + calls: list[list[str]] = [] + + async def add(**kwargs): + submitted.set() + + async def abort(request_ids, timeout=None): + calls.append(list(request_ids)) + abort_started.set() + await resume.wait() + + omni = get_async_omni_instance(fake_add_request=add, fake_abort_request=abort) + + async def collect(): + async for _ in omni.generate( + prompt={"prompt": "test"}, + request_id="cancel-twice", + sampling_params_list=[SamplingParams(max_tokens=1)], + output_modalities=["text"], + ): + pass + + task = asyncio.create_task(collect()) + await asyncio.wait_for(submitted.wait(), timeout=1) + task.cancel() + await asyncio.wait_for(abort_started.wait(), timeout=1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, timeout=1) + (state,) = omni.request_states.values() + assert state.cleanup_pending and state.external_request_id == "cancel-twice" + resume.set() + await omni.abort("cancel-twice") + assert calls == [[state.request_id], [state.request_id]] + assert not omni.request_states and state.queue.empty() asyncio.run(run()) @@ -713,7 +870,16 @@ async def cancel_request(): await asyncio.wait_for(cancel_request(), timeout=1.0) assert abort_started.is_set() assert abort_finished.is_set() - assert omni.request_states == {} + (state,) = omni.request_states.values() + assert state.cleanup_pending + assert state.external_request_id == "cancel-asgi-hanging-abort" + + async def recovered_abort(request_ids, timeout=None): + assert request_ids == [state.request_id] + + omni.engine.abort_async = recovered_abort + await omni.abort(state.external_request_id) + assert omni.request_states == {} and state.queue.empty() asyncio.run(run_test()) diff --git a/tests/entrypoints/test_async_omni_prompt_update.py b/tests/entrypoints/test_async_omni_prompt_update.py index f3edeeec72b..256e2f00d1f 100644 --- a/tests/entrypoints/test_async_omni_prompt_update.py +++ b/tests/entrypoints/test_async_omni_prompt_update.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Entrypoint contract tests for ``AsyncOmni.submit_interaction_async``.""" from __future__ import annotations @@ -68,6 +68,18 @@ async def test_submit_interaction_async_maps_external_to_internal_id() -> None: ) +@pytest.mark.asyncio +async def test_submit_interaction_async_ignores_requests_pending_cleanup() -> None: + omni = _make_async_omni() + closed = omni.request_states["external-abc-uuid-1"] + closed.cleanup_pending = True + await omni.submit_interaction_async("external-abc", interaction=_prompt_interaction()) + omni.engine.submit_interaction_async.assert_awaited_once_with( + "external-abc-uuid-2", interaction=_prompt_interaction() + ) + assert omni.request_states["external-abc-uuid-1"] is closed + + @pytest.mark.asyncio async def test_submit_interaction_async_passes_missing_transition_chunks() -> None: omni = _make_async_omni() @@ -99,11 +111,18 @@ async def test_submit_interaction_async_rejects_negative_transition_chunks() -> @pytest.mark.asyncio -async def test_submit_interaction_async_rejects_inactive_request() -> None: +@pytest.mark.parametrize("pending_cleanup", [False, True]) +async def test_submit_interaction_async_rejects_inactive_request(pending_cleanup) -> None: omni = _make_async_omni() - omni.request_states.clear() + if pending_cleanup: + for state in omni.request_states.values(): + state.cleanup_pending = True + request_id = "external-abc" + else: + omni.request_states.clear() + request_id = "missing" with pytest.raises(ValueError, match="No active request"): - await omni.submit_interaction_async("missing", interaction=_prompt_interaction(transition_chunks=2)) + await omni.submit_interaction_async(request_id, interaction=_prompt_interaction(transition_chunks=2)) @pytest.mark.asyncio diff --git a/tests/entrypoints/test_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py index 87aeb787ebb..eb2628a4dae 100644 --- a/tests/entrypoints/test_stream_finish_reason.py +++ b/tests/entrypoints/test_stream_finish_reason.py @@ -228,12 +228,14 @@ async def result_generator(): ], ) @pytest.mark.asyncio +@pytest.mark.parametrize("pending_cleanup", [False, True]) async def test_streaming_audio_metrics_resolve_replica_id( monkeypatch, output_replica_id, pool_replica_id, expected_replica_id, pool_lookup, + pending_cleanup, ): import vllm_omni.entrypoints.openai.serving_chat as serving_chat_mod from vllm_omni.entrypoints.client_request_state import ClientRequestState @@ -245,7 +247,11 @@ async def test_streaming_audio_metrics_resolve_replica_id( external_request_id="test-req", ) req_state.request_arrival_ts = time.time() - 1.0 - serving_chat.engine_client.request_states = {"internal-req": req_state} + closed = ClientRequestState(request_id="old-req", external_request_id="test-req") + closed.cleanup_pending = True + serving_chat.engine_client.request_states = ( + {"old-req": closed, "internal-req": req_state} if pending_cleanup else {"internal-req": req_state} + ) serving_chat.engine_client.mod_metrics = object() serving_chat.engine_client.engine = MagicMock() serving_chat.engine_client.engine.stage_pools = [MagicMock(), MagicMock(), MagicMock()] @@ -286,6 +292,7 @@ async def result_generator(): assert first_packet_calls[0]["replica_id"] == expected_replica_id assert finalize_calls[0]["stage_id"] == 2 assert finalize_calls[0]["replica_id"] == expected_replica_id + assert req_state.first_audio_ts is not None and closed.first_audio_ts is None # --------------------------------------------------------------------------- diff --git a/tests/worker/test_omni_connector_mixin.py b/tests/worker/test_omni_connector_mixin.py index ea042d4e846..f4dd9f11138 100644 --- a/tests/worker/test_omni_connector_mixin.py +++ b/tests/worker/test_omni_connector_mixin.py @@ -79,12 +79,31 @@ def _make_model_config( ) -> SimpleNamespace: return SimpleNamespace( stage_connector_config=None, + use_v2_model_runner=True, + supports_native_mrv2_data_plane=True, async_chunk=async_chunk, worker_type=worker_type, custom_process_next_stage_input_func=custom_func, ) +@pytest.mark.parametrize("native", [False, True]) +@pytest.mark.parametrize("role", ["sender", "receiver"]) +def test_async_transport_has_one_owner_for_runner_and_scheduler(mocker, native, role): + config = _make_model_config(async_chunk=True, custom_func="stage.payload") + config.stage_connector_config = {"role": role} + config.use_v2_model_runner = native + host = MixinHost() + create = mocker.patch.object(host, "_create_connector", return_value=None) + mocker.patch.object(host, "_load_custom_func", return_value=(None, None)) + host.init_omni_connectors(config) + try: + assert create.call_count == int(native) + assert host._omni_connector is None + finally: + host.shutdown_omni_connectors() + + def _make_request(req_id: str, external_req_id: str | None = None): r = SimpleNamespace( request_id=req_id, diff --git a/tests/worker_v2/test_omni_data_plane.py b/tests/worker_v2/test_omni_data_plane.py index 62898a456e2..2a97f081e08 100644 --- a/tests/worker_v2/test_omni_data_plane.py +++ b/tests/worker_v2/test_omni_data_plane.py @@ -47,6 +47,7 @@ def _bare_plane(*, delivery_timeout_s=1.0, shutdown_timeout_s=1.0): _native_output_error=None, _native_output_error_lock=threading.Lock(), _native_output_closed=False, + model_config=SimpleNamespace(stage_connector_config={"from_stage": None}), _native_outputs_in_flight=defaultdict(int), _native_terminal_pending=set(), _async_chunk=True, @@ -260,6 +261,108 @@ def _stop_save_thread(plane): assert not plane._save_thread.is_alive() +def test_pipeline_release_preserves_legacy_backend_and_retires_request(plane): + connector = SimpleNamespace(cleanup=lambda request_id: None) + plane.register_request(_new_request("external", "external")) + _start_save_thread(plane, connector) + try: + plane.release_request_resources(["external"]) + assert "external" not in plane._native_requests + assert not plane._pending_save_reqs + assert not plane._pending_save_counts + finally: + _stop_save_thread(plane) + + +def test_pipeline_release_without_a_transport_owner(plane): + plane._omni_connector = None + plane.release_request_resources(["external"]) + + +def test_pipeline_release_fences_scheduler_terminal_arriving_later(plane): + plane._omni_connector = SimpleNamespace(cleanup_prefix=lambda prefix: 0) + plane._stage_id = 0 + plane._lock = threading.Lock() + plane._pending_save_reqs = {} + plane._pending_save_counts = defaultdict(int) + plane._save_work_available = threading.Event() + plane.register_request(_new_request("external", "external")) + plane.reserve_outputs(["external"]) + + _start_save_thread(plane, plane._omni_connector) + try: + plane.release_request_resources(["external"]) + finally: + _stop_save_thread(plane) + + assert not plane._native_requests + assert plane.record.cleaned == ["external"] + assert len(plane.record.batches) == 1 + assert plane.record.batches[0][0][0].is_finished() + assert plane.abort_requests({"external"}) == 0 + assert _complete(plane, [{"codes.audio": "late"}], req="external") == 0 + assert len(plane.record.batches) == 1 + assert not plane._pending_save_reqs + assert not plane._pending_save_counts + + +def test_pipeline_release_waits_for_queued_payloads_and_uses_external_prefix(raw_plane): + connector = _gated_connector() + reclaimed = threading.Event() + prefixes = [] + + def cleanup_prefix(prefix): + prefixes.append(prefix) + reclaimed.set() + + connector.cleanup_prefix = cleanup_prefix + _start_save_thread(raw_plane, connector) + try: + request = SimpleNamespace(request_id="external") + raw_plane._enqueue_chunk_payload(request, {"c": 1}, wait_for_delivery=False) + raw_plane._enqueue_chunk_payload(request, {"c": 2}, wait_for_delivery=False) + assert connector.put_started.wait(timeout=1) + with ThreadPoolExecutor(max_workers=1) as executor: + release = executor.submit(raw_plane.release_request_resources, ["external"]) + assert not reclaimed.is_set() + assert not release.done() + connector.release.set() + release.result(timeout=2) + assert reclaimed.is_set() + assert connector.put_keys == ["external_0_0", "external_0_1"] + assert prefixes == ["external_0_"] + finally: + connector.release.set() + _stop_save_thread(raw_plane) + assert not raw_plane._pending_save_reqs + assert not raw_plane._pending_save_counts + + +@pytest.mark.parametrize("failure_method", ["cleanup_prefix", "wait_for_cleanup"]) +def test_pipeline_release_failure_is_reported_and_retry_is_request_scoped(raw_plane, failure_method): + connector = _gated_connector() + connector.cleanup_prefix = lambda prefix: None + connector.wait_for_cleanup = lambda prefix, **kwargs: None + + def fail(*args, **kwargs): + raise TimeoutError("READ claim remains active") + + original = getattr(connector, failure_method) + setattr(connector, failure_method, fail) + _start_save_thread(raw_plane, connector) + try: + with pytest.raises(TimeoutError, match="READ claim remains active"): + raw_plane.release_request_resources(["external"]) + assert not raw_plane._pending_save_reqs + assert not raw_plane._pending_save_counts + assert not raw_plane._delivery_manager.is_quarantined + setattr(connector, failure_method, original) + raw_plane.release_request_resources(["external"]) + assert not raw_plane._pending_save_counts + finally: + _stop_save_thread(raw_plane) + + class _Output: copy_event = None 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..ce322e73cc1 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) @@ -1380,6 +1385,27 @@ def _send_single_request(self, task: dict) -> bool: return True request_id = task["request_id"] + if prefix := task.get("release_prefix"): + completion = task["release_completion"] + release_error = None + try: + cleanup = getattr(connector, "cleanup_prefix", None) + if callable(cleanup): + cleanup(prefix) + # Legacy backends keep their existing per-payload lifecycle. + # Completing this marker still fences the local send queue. + wait = getattr(connector, "wait_for_cleanup", None) + if callable(wait): + wait(prefix, receive_prefix=task["receive_prefix"], deadline=task["release_deadline"]) + except Exception as error: + release_error = error + finally: + self._decrement_pending_save_count(request_id) + if release_error is not None: + completion.set_exception(release_error) + else: + completion.set_result(None) + return True payload_data = task.get("data") if payload_data is None and task.get("request") is not None: payload_data = self._build_custom_process_payload( @@ -1562,6 +1588,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..46547be743c 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 @@ -15,9 +15,11 @@ import torch from vllm.logger import init_logger +from vllm_omni.core.sched.omni_scheduling_coordinator import uses_native_mrv2_data_plane 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 @@ -51,6 +53,12 @@ def _should_create_payload_connector(model_config: Any) -> bool: Sender edges may instead be owned solely by KV transfer. Receivers still need a connector even though they do not declare a downstream payload hook. """ + if getattr(model_config, "async_chunk", False) and not uses_native_mrv2_data_plane( + model_config, use_v2_model_runner=getattr(model_config, "use_v2_model_runner", False) + ): + # The scheduler's chunk adapter owns V1 transport. Creating another + # runner connector duplicates its endpoint and payload ownership. + return False if get_stage_connector_role(model_config) != "sender": return True @@ -112,7 +120,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 +226,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 +686,13 @@ 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: + def _resolve_next_stage_id(self, model_config: Any) -> int | None: """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 + 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 5e4ef6820ba..ea0e96569ce 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 @@ -8,7 +8,8 @@ import time from collections import defaultdict, deque from collections.abc import Callable, Iterable, Mapping -from typing import Any +from concurrent.futures import Future +from typing import Any, cast import torch from vllm.v1.metrics.stats import PrefillStats @@ -20,7 +21,12 @@ 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 ( + REQUEST_RESOURCE_RELEASE_TIMEOUT_S, + 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 @@ -137,6 +143,10 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): work and return immediately. """ + # Backends expose transport-specific capabilities in addition to the + # shared connector interface; this subclass always creates one. + connector: Any + def __init__(self, vllm_config: Any): model_config = vllm_config.model_config # The base constructor starts the save thread, so sender-generation @@ -195,17 +205,17 @@ def __init__(self, vllm_config: Any): self._adaptive_states: dict[str, Any] = {} self.upstream_exhausted_requests: set[str] = set() self.segment_finished_requests: set[str] = set() - self.request_payload = {} + self.request_payload: dict[str, Any] = {} self.code_prompt_token_ids: dict[str, list[torch.Tensor]] = defaultdict(list) self.request_ids_mapping: dict[str, str] = {} - # Save-thread only: running count of segments reclaimed after finish. - self._reclaimed_shm_total = 0 + # Save-thread only: running count of payloads reclaimed after finish. + self._reclaimed_payload_total = 0 self.waiting_for_chunk_waiting_requests: deque[Any] = deque() self.waiting_for_chunk_running_requests: deque[Any] = deque() - self.requests_with_ready_chunks = set() + self.requests_with_ready_chunks: set[str] = set() self.replaced_streaming_prompt_ids: set[str] = set() - self.requests_origin_status = {} + self.requests_origin_status: dict[str, RequestStatus] = {} self._active_streams: dict[str, Any] = {} # Private hold-queue for non-active running requests. Restored to # running_queue inside restore_queues(). Avoids calling @@ -410,6 +420,7 @@ def save_async( because ``finish_requests`` skips requests that already report ``is_finished()``. See ``_finish_parked_streaming_session``. """ + request = cast(Request, request) is_finished = force_request_finished or (request.is_finished() and not request.resumable) if not hasattr(self, "_segment_generation"): self._segment_generation = defaultdict(int) @@ -498,7 +509,9 @@ def save_async( 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: @@ -692,8 +705,21 @@ def _send_single_request(self, task: dict): for chunk_id in range(num_puts): self.connector.cleanup(f"{key_prefix}_{chunk_id}") return - if "release_shm" in task: - self._release_shm_prefix(task["release_shm"]) + if "release_prefix" in task: + completion = task["release_completion"] + try: + self._release_connector_prefix(task["release_prefix"]) + wait = getattr(self.connector, "wait_for_cleanup", None) + if callable(wait): + wait( + task["release_prefix"], + receive_prefix=task["receive_prefix"], + deadline=task["release_deadline"], + ) + except Exception as error: + completion.set_exception(error) + else: + completion.set_result(None) return request = task["request"] external_req_id = request.external_req_id @@ -779,7 +805,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}" @@ -1050,8 +1078,8 @@ def _clear_sender_state_locked(self, external_req_id: str) -> None: if cached_ic is not None: cached_ic.pop(external_req_id, None) - def release_shm_resources(self, request_id: str) -> None: - """Queue reclaim of inter-stage segments this request left unconsumed. + def release_request_resources(self, request_id: str, *, deadline: float | None = None) -> Future[None]: + """Finish queued publications and apply the backend's request cleanup. Safe only once every stage has finished with the request: the producer keeps writing until its own request ends, and a consumer that stopped @@ -1060,30 +1088,45 @@ def release_shm_resources(self, request_id: str) -> None: orchestrator -- the only party that knows all stages are done -- drives this (see ``Orchestrator._cleanup_request_ids``). - Like abort cleanup, the unlink runs on the save thread, so the - scheduler never performs SHM I/O. + Cleanup runs behind puts on the save thread, so the scheduler never + performs connector I/O. Legacy backends without prefix cleanup retain + their existing resource lifecycle; their completion confirms the local + publication fence, not physical reclamation by the backend. """ - if not isinstance(self.connector, SharedMemoryConnector): - return external_req_id = self.request_ids_mapping.get(request_id, request_id) - self._pending_save_reqs.append({"release_shm": f"{external_req_id}_{self.connector.stage_id}_"}) + from_stage = get_stage_connector_peer(self.config, self.connector.stage_id, "from_stage") + completion: Future[None] = Future() + self._pending_save_reqs.append( + { + "release_prefix": f"{external_req_id}_{self.connector.stage_id}_", + "receive_prefix": f"{external_req_id}_{from_stage}_" if from_stage is not None else None, + "release_completion": completion, + "release_deadline": deadline + if deadline is not None + else time.monotonic() + REQUEST_RESOURCE_RELEASE_TIMEOUT_S, + } + ) with self._save_cond: self._save_cond.notify() + return completion - def _release_shm_prefix(self, key_prefix: str) -> None: - reclaimed = self.connector.cleanup_prefix(key_prefix) + def _release_connector_prefix(self, key_prefix: str) -> None: + cleanup = getattr(self.connector, "cleanup_prefix", None) + if not callable(cleanup): + return + reclaimed = cleanup(key_prefix) if not reclaimed: return # Non-zero means a consumer stopped before draining the producer. # Common on audio requests, so warn only when the running total # crosses another _RECLAIM_WARN_INTERVAL boundary. - prev = self._reclaimed_shm_total - self._reclaimed_shm_total = prev + reclaimed - if prev // _RECLAIM_WARN_INTERVAL != self._reclaimed_shm_total // _RECLAIM_WARN_INTERVAL: + prev = self._reclaimed_payload_total + self._reclaimed_payload_total = prev + reclaimed + if prev // _RECLAIM_WARN_INTERVAL != self._reclaimed_payload_total // _RECLAIM_WARN_INTERVAL: logger.warning( - "Reclaimed %d unconsumed inter-stage segments so far; " + "Reclaimed %d unconsumed inter-stage payloads so far; " "a downstream stage finished before draining its producer", - self._reclaimed_shm_total, + self._reclaimed_payload_total, ) def cleanup( @@ -1627,7 +1670,7 @@ def finish_requests( elif request_ids is not None: request_ids = set(request_ids) else: - request_ids = requests.keys() + request_ids = cast(dict[str, Request], requests).keys() connector_owned_ids = { request.request_id diff --git a/vllm_omni/distributed/omni_connectors/utils/config.py b/vllm_omni/distributed/omni_connectors/utils/config.py index a646a3a3090..63ad077d4dc 100644 --- a/vllm_omni/distributed/omni_connectors/utils/config.py +++ b/vllm_omni/distributed/omni_connectors/utils/config.py @@ -2,12 +2,33 @@ # 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 logger = get_connector_logger(__name__) +REQUEST_RESOURCE_RELEASE_TIMEOUT_S = 5.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) + + TRANSFER_ENGINE_CONNECTOR_NAMES = frozenset( { "MooncakeTransferEngineConnector", diff --git a/vllm_omni/engine/errors.py b/vllm_omni/engine/errors.py index 0128176d660..abeecc64175 100644 --- a/vllm_omni/engine/errors.py +++ b/vllm_omni/engine/errors.py @@ -6,3 +6,7 @@ class NativeKVHandoffError(RuntimeError): """One request lost its producer binding or transfer metadata.""" + + +class ResourceReleaseError(RuntimeError): + """A live replica could not confirm request transfer ownership was released.""" diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index a135b60bede..79caf3f3a34 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -40,7 +40,7 @@ from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker -from vllm_omni.engine.errors import NativeKVHandoffError +from vllm_omni.engine.errors import NativeKVHandoffError, ResourceReleaseError from vllm_omni.engine.membership_controller import MembershipController from vllm_omni.engine.messages import ( AbortRequestMessage, @@ -60,7 +60,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 @@ -239,9 +239,10 @@ class OrchestratorRequestState: sampling_params_list: list[Any] = field(default_factory=list) final_stage_id: int = -1 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 + # Keep collected tokens until cleanup succeeds, including across retries. + pending_abort_outputs: list[OutputMessage] = field(default_factory=list, repr=False) # Wall-clock timestamp when the client-facing engine request was accepted. request_timestamp: float = 0.0 @@ -270,6 +271,22 @@ class OrchestratorRequestState: upstream_first_audio: bool = False pending_upstream_first_audio: tuple[int, OmniEngineCoreOutput] | None = None pending_first_audio_outputs: list[tuple[int, int, Any]] = field(default_factory=list) + # One cleanup operation owns the request until abort and binding release + # finish. Retain state for partial abort outputs, but forbid new admission. + cleanup_future: asyncio.Future[list[OutputMessage]] | None = field(default=None, repr=False) + + @property + def closing(self) -> bool: + return self.cleanup_future is not None + + @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 @@ -336,7 +353,6 @@ class OrchestratorBase: _transfer_emitter: Any = None _prom_metrics: Any = None _stat_logger: OmniPrometheusStatLogger | None = None - _transfer_release_tasks: set[asyncio.Task] = set() def __init__( self, @@ -385,7 +401,6 @@ def __init__( self._pd_prefill_engine_id = pd_config.get("prefill_engine_id") self.request_states: dict[str, OrchestratorRequestState] = {} # Strong refs for in-flight releases; the loop only weak-refs tasks, so dropping these risks mid-flight GC. - self._transfer_release_tasks: set[asyncio.Task] = set() self._init_metrics_state( stage_pools, running_counter, @@ -489,7 +504,7 @@ async def run(self) -> None: if self._membership is not None: async def cleanup_member_requests(ids: list[str]) -> None: - await self._cleanup_request_ids(ids, abort=True) + await self._cleanup_request_ids(ids, abort=True, release_owners=True) self._membership.install_unregister_handlers( output_queue=self.output_async_queue, @@ -505,7 +520,11 @@ async def cleanup_member_requests(ids: list[str]) -> None: tasks.append(membership_watcher) try: - await asyncio.gather(*tasks) + # These loops live until shutdown or failure. Once one exits, + # cancel the others in finally, including the membership watcher. + done, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + for task in done: + task.result() except asyncio.CancelledError: raise except EngineDeadError as e: @@ -664,7 +683,10 @@ async def _handle_abort(self, msg: AbortRequestMessage) -> None: error = str(exc) logger.exception("[Orchestrator] Abort failed for request(s) %s", request_ids) if msg.rpc_id is None: - raise + if isinstance(exc, ResourceReleaseError): + await self._report_release_failure(list(request_ids), exc) + else: + raise if msg.rpc_id is not None: # Always emit a result when rpc_id is set so CorrelatedRpcClient # waiters are unblocked even on failure. @@ -689,7 +711,6 @@ async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage """ if not request_ids: return [] - abort_outputs: list[OutputMessage] = [] for pool in self.stage_pools: stage_outputs = await pool.abort_requests(request_ids) or [] if bool(getattr(pool, "final_output", False)) and getattr(pool, "stage_type", None) != "diffusion": @@ -709,7 +730,7 @@ async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage replica_id=pool.get_bound_replica_id(orch_req_id), final_output_type=final_output_type, ) - abort_outputs.append( + req_state.pending_abort_outputs.append( OutputMessage( request_id=orch_req_id, stage_id=pool.stage_id, @@ -721,6 +742,13 @@ async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage ) ) pool.release_bindings(request_ids) + abort_outputs = [ + output + for request_id in request_ids + if (state := self.request_states.get(request_id)) is not None + for output in state.pending_abort_outputs + ] + abort_outputs.sort(key=lambda output: 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 @@ -728,35 +756,26 @@ async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage output_msg.finished = index == last_index_by_req[output_msg.request_id] return abort_outputs - def _release_stage_transfer_resources(self, request_ids: list[str]) -> None: - """Drop each stage's inter-stage transfer resources for finished requests. + async def _release_stage_transfer_resources(self, request_ids: list[str]) -> None: + """Finish each stage's transfer lifecycle for terminated requests. This is the only point that knows every stage is done with the request, so it is the only safe place to reclaim segments a consumer never - drained. Scheduled rather than awaited: reclaim is best-effort and must - not add RPC latency to request teardown. + drained. Every live replica completes its local publication fence and + applicable backend cleanup. SHM/NIXL additionally confirm ownership + drain; legacy backends retain their existing resource contract. A + failed operation leaves the request closed for retry. """ if not request_ids: return - async def _run() -> None: - results = await asyncio.gather( - *(pool.release_request_resources(request_ids) for pool in self.stage_pools), - return_exceptions=True, - ) - for result in results: - if isinstance(result, Exception): - logger.warning("[Orchestrator] release transfer resources failed: %s", result) - - try: - task = asyncio.get_running_loop().create_task(_run()) - self._transfer_release_tasks.add(task) - task.add_done_callback(self._transfer_release_tasks.discard) - except RuntimeError: - logger.warning( - "[Orchestrator] no running event loop; skipped reclaim of transfer resources for %s", - request_ids, - ) + results = await asyncio.gather( + *(pool.release_request_resources(request_ids) for pool in self.stage_pools), + return_exceptions=True, + ) + for result in results: + if isinstance(result, BaseException): + raise result def _release_request_bindings(self, request_ids: list[str]) -> None: """Release all stage-local route bindings for the given request ids.""" @@ -884,7 +903,7 @@ async def _process_llm_stage_outputs( ) request_id = getattr(eco, "request_id", None) req_state = self.request_states.get(request_id) if request_id is not None else None - if req_state is None: + if req_state is None or req_state.closing: continue if ( getattr(eco, "finish_reason", None) == FinishReason.ERROR @@ -914,6 +933,15 @@ async def _process_llm_stage_outputs( if await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state): raw_terminal_request_ids.add(req_state.request_id) await self._report_duplex_session_request_error(stage_id, replica_id, eco, req_state) + # An in-flight abort owns the processor snapshot. Late terminal outputs + # must not delete it before a failed stage/replica abort can retry. + raw_outputs.outputs = [ + eco + for eco in raw_outputs.outputs + if (request_id := getattr(eco, "request_id", None)) is None + or (state := self.request_states.get(request_id)) is None + or not state.closing + ] iteration_stats = IterationStats() if (self._stat_logger is not None and raw_outputs.outputs) else None processed = await pool.process_llm_raw_outputs( replica_id, @@ -964,9 +992,14 @@ async def _route_upstream_first_audio(self, stage_id: int, replica_id: int, raw_ # Only the dedicated first-frame channel is intercepted. if ( req_state is None + or req_state.closing 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 ( + 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") @@ -978,6 +1011,7 @@ async def _route_upstream_first_audio(self, stage_id: int, replica_id: int, raw_ continue if ( req_state is not None + and not req_state.closing and final_output and not req_state.upstream_first_audio and (requires_first or req_state.pending_first_audio_outputs) @@ -1000,10 +1034,16 @@ async def _flush_upstream_first_audio(self, req_state: OrchestratorRequestState) from vllm_omni.engine import OmniEngineCoreOutputs pending = req_state.pending_upstream_first_audio - if pending is None or self.request_states.get(req_state.request_id) is not req_state: + if pending is None or req_state.closing 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: @@ -1022,11 +1062,11 @@ async def _flush_upstream_first_audio(self, req_state: OrchestratorRequestState) logger.exception("First-audio output processing failed for %s", req_state.request_id) await self._handle_stage_error(codec_stage, first_output, error=str(error)) return - if self.request_states.get(req_state.request_id) is not req_state: + if req_state.closing or self.request_states.get(req_state.request_id) is not req_state: return ready, req_state.pending_first_audio_outputs = req_state.pending_first_audio_outputs, [] for pending_stage, pending_replica, output in ready: - if self.request_states.get(req_state.request_id) is not req_state: + if req_state.closing or self.request_states.get(req_state.request_id) is not req_state: break terminal: set[str] = set() processed = await self._process_llm_stage_outputs( @@ -1416,7 +1456,7 @@ async def _handle_stage_error(self, stage_id: int, output: Any, *, error: str | ) ) await self._cleanup_request_ids( - [parent_id, *self._cfg_tracker.cleanup_parent(parent_id)], + [parent_id], abort=True, release_owners=True, ) @@ -1472,9 +1512,9 @@ async def _handle_dead_replica(self, stage_id: int, replica_id: int, error: Engi stage_has_live = bool(pool.live_replica_ids()) failed_ids: list[str] = [] for req_id, req_state in list(self.request_states.items()): - if stage_id not in req_state.stage_submit_ts: - continue bound = pool.get_bound_replica_id(req_id) + if stage_id not in req_state.stage_submit_ts and bound is None: + continue if bound == replica_id or not stage_has_live: await self.output_async_queue.put( ErrorMessage( @@ -1492,7 +1532,7 @@ async def _handle_dead_replica(self, stage_id: int, replica_id: int, error: Engi # distributed-membership eviction path. for req_id in failed_ids: await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + [req_id], abort=True, # A session-owned request has an owner that outlives the stage # request: without this the runner and its admission slot @@ -1526,7 +1566,7 @@ async def _fail_request_dead_stage(self, req_id: str, stage_id: int) -> None: ) ) await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + [req_id], abort=True, release_owners=True, ) @@ -1557,7 +1597,7 @@ async def _fail_request_client_error( ) ) await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + [req_id], abort=True, release_owners=release_owners, ) @@ -1596,6 +1636,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, @@ -1676,26 +1718,27 @@ async def _cleanup_request_ids( cannot outlive its requests. ``release_owners`` asks the subclass to also release any session that - owns these requests (error paths pass True); the base ignores it. + owns these requests (error paths pass True). It also marks automatic + error cleanup, whose release failure must remain request-scoped. Returns: Final-stage AR abort ``OutputMessage`` list when ``abort=True``; otherwise an empty list. """ - del release_owners if not request_ids: return [] - cleanup_ids = list(dict.fromkeys(request_ids)) - batch = set(cleanup_ids) + candidates = list(dict.fromkeys(request_ids)) + batch = set(candidates) orphaned_parents: dict[str, str] = {} - for rid in cleanup_ids: + for rid in candidates: pid = self._cfg_tracker.get_parent_id(rid) if pid is not None and pid not in batch and not self._cfg_tracker.is_companion_done(rid): orphaned_parents.setdefault(pid, rid) + errors: list[ErrorMessage] = [] for pid, cid in orphaned_parents.items(): deferred = self._cfg_tracker.pop_pending_parent(pid) - await self.output_async_queue.put( + errors.append( ErrorMessage( request_id=pid, stage_id=deferred["stage_id"] if deferred is not None else None, @@ -1703,25 +1746,107 @@ async def _cleanup_request_ids( ) ) batch.add(pid) - cleanup_ids.append(pid) - for rid in list(cleanup_ids): - for cid in self._cfg_tracker.cleanup_parent(rid): + candidates.append(pid) + for rid in list(candidates): + for cid in self._cfg_tracker.get_companion_request_ids(rid).values(): if cid not in batch: batch.add(cid) - cleanup_ids.append(cid) + candidates.append(cid) + + # Expand CFG ownership before selecting cleanup owners so a parent's + # cleanup joins a companion's running operation rather than replacing it. + cleanup_ids: list[str] = [] + pending: dict[asyncio.Future[list[OutputMessage]], list[str]] = {} + for request_id in candidates: + state = self.request_states.get(request_id) + future = state.cleanup_future if state is not None else None + if future is not None and (not future.done() or future.exception() is None): + pending.setdefault(future, []).append(request_id) + else: + cleanup_ids.append(request_id) + + # A parent must stay closed until its already-closing companions have + # released their resources. Unrelated joined requests retain their own + # outcome, so their failure cannot prevent this batch's cleanup. + dependencies = tuple( + future + for future, joined_ids in pending.items() + if any(self._cfg_tracker.get_parent_id(rid) in cleanup_ids for rid in joined_ids) + or any(self._cfg_tracker.get_parent_id(rid) in joined_ids for rid in cleanup_ids) + ) + for future in dependencies: + pending.pop(future) + abort_outputs: list[OutputMessage] = [] - if abort: - abort_outputs = await self._abort_request_ids(cleanup_ids) - self._release_request_bindings(cleanup_ids) - self._release_stage_transfer_resources(cleanup_ids) - for request_id in cleanup_ids: + if cleanup_ids: + future = asyncio.get_running_loop().create_future() + for request_id in cleanup_ids: + state = self.request_states.get(request_id) + if state is not None: + state.cleanup_future = future + try: + for error in errors: + await self.output_async_queue.put(error) + abort_outputs = await self._finish_request_cleanup(cleanup_ids, abort=abort, dependencies=dependencies) + except BaseException as exc: + future.set_exception(exc) + # The owner raises directly; joining callers receive the same + # failure, without an unobserved-Future warning if none join. + future.exception() + if isinstance(exc, ResourceReleaseError) and (not abort or release_owners): + await self._report_release_failure(cleanup_ids, exc) + else: + raise + else: + future.set_result(abort_outputs) + + for future, joined_ids in pending.items(): + # Cancelling a duplicate abort waiter must not cancel the owner. + try: + abort_outputs.extend(await asyncio.shield(future)) + except ResourceReleaseError as error: + if abort and not release_owners: + raise + await self._report_release_failure(joined_ids, error) + return abort_outputs if abort else [] + + async def _report_release_failure(self, request_ids: list[str], error: ResourceReleaseError) -> None: + for request_id in request_ids: + await self.output_async_queue.put( + ErrorMessage( + error=str(error), + fatal=False, + request_id=request_id, + status_code=HTTPStatus.BAD_GATEWAY.value, + error_type=type(error).__name__, + ) + ) + + async def _finish_request_cleanup( + self, + request_ids: list[str], + *, + abort: bool, + dependencies: tuple[asyncio.Future[list[OutputMessage]], ...] = (), + ) -> list[OutputMessage]: + """Finish the claimed cleanup while request state remains available.""" + abort_outputs = await self._abort_request_ids(request_ids) if abort else [] + await self._release_stage_transfer_resources(request_ids) + results = await asyncio.gather(*(asyncio.shield(future) for future in dependencies), return_exceptions=True) + for result in results: + if isinstance(result, BaseException): + raise result + abort_outputs.extend(result) + self._release_request_bindings(request_ids) + for request_id in request_ids: + self._cfg_tracker.cleanup_parent(request_id) self._pd_kv_params.pop(request_id, None) req_state = self.request_states.pop(request_id, None) if req_state is not None: - cleanup_request_artifact_dirs(getattr(req_state, "request_artifact_dirs", set())) - if req_state is not None and req_state.running_counter_registered and self._running_counter is not None: - self._running_counter.decrement() - req_state.running_counter_registered = False + cleanup_request_artifact_dirs(req_state.request_artifact_dirs) + if req_state.running_counter_registered and self._running_counter is not None: + self._running_counter.decrement() + req_state.running_counter_registered = False return abort_outputs async def _apply_raw_terminal_stage_finish( @@ -1748,10 +1873,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( @@ -1807,23 +1928,19 @@ async def _finish_raw_terminal_requests( 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 or req_state.closing or req_state.session_owned: continue pending = req_state.pending_final_output if pending is not None: - if 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) - 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 - 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) logger.info( "[Orchestrator] req=%s stage-%s raw terminal produced no processed terminal; returning empty %s output", @@ -1836,18 +1953,32 @@ 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), - ) + 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), ) - 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, message) + else: + req_state.pending_final_output = message + + async def _publish_final_output(self, req_state: OrchestratorRequestState, message: OutputMessage) -> None: + """Publish success after the request's transfer lifecycle finishes.""" + request_id = req_state.request_id + await self._cleanup_request_ids([request_id]) + if request_id not in self.request_states: + req_state.pending_final_output = None + await self.output_async_queue.put(message) + + 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 _maybe_clone_diffusion_params_for_cfg(self, request_id: str, params: Any) -> Any: """Attach CFG companion ids to diffusion sampling params when needed.""" @@ -1881,6 +2012,8 @@ async def _route_output( stage_metrics: Any, ) -> None: """Route a processed output: send to frontend and/or forward.""" + if req_state.closing: + return req_id = output.request_id finished = output.finished submit_ts = req_state.stage_submit_ts.get(stage_id) @@ -1898,9 +2031,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 @@ -1914,12 +2045,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 @@ -1948,12 +2074,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 @@ -1997,16 +2124,14 @@ 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) - if request_finished and not req_state.session_owned: - await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)]) + def _request_stage_path(self, req_state: OrchestratorRequestState) -> tuple[int, ...]: + return tuple(range(req_state.final_stage_id + 1)) - def _next_stage_already_submitted(self, stage_id: int, req_state: OrchestratorRequestState) -> bool: - return (stage_id + 1) in req_state.stage_submit_ts + def _next_request_stage(self, stage_id: int, req_state: OrchestratorRequestState) -> int | None: + return stage_id + 1 if stage_id < req_state.final_stage_id else None def _stage_receives_async_chunks(self, stage_id: int) -> bool: """Whether a stage's connector supplies its runtime inputs.""" @@ -2148,11 +2273,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( @@ -2178,7 +2304,7 @@ async def _handle_kv_ready_raw_outputs( req_id = raw_output.request_id req_state = self.request_states.get(req_id) - if req_state is None: + if req_state is None or req_state.closing: continue if self._cfg_tracker.is_companion(req_id): # kv_ready only says the companion's KV hit the connector; its @@ -2188,9 +2314,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): @@ -2290,18 +2415,24 @@ 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.""" + if req_state.closing: + return + next_stage_id = self._next_request_stage(src_stage_id, req_state) + if 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_stage_id=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", ) @@ -2312,23 +2443,23 @@ async def _forward_to_next_stage_unguarded( output: Any, req_state: OrchestratorRequestState, *, + next_stage_id: 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] + next_pool = self.stage_pools[next_stage_id] if not next_pool.live_replica_ids(): # The downstream stage lost all of its replicas before this in-flight # request could be forwarded. - await self._fail_request_dead_stage(req_id, next_logical) + await self._fail_request_dead_stage(req_id, next_stage_id) return next_client = next_pool.stage_client - params = req_state.sampling_params_list[next_logical] + params = req_state.sampling_params_list[next_stage_id] source_outputs = [output] next_stage_resumable = is_streaming_session and not is_final_update - already_submitted = self._next_stage_already_submitted(src_stage_id, req_state) + already_submitted = next_stage_id in req_state.stage_submit_ts requires_multimodal_data = getattr(next_client, "requires_multimodal_data", False) _t_submit_start = _time.perf_counter() @@ -2371,7 +2502,7 @@ async def _forward_to_next_stage_unguarded( ) ) await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + [req_id], abort=True, ) return @@ -2401,59 +2532,36 @@ async def _forward_to_next_stage_unguarded( req_id, _dt_ar2d, src_stage_id, - next_logical, + next_stage_id, ) - if diffusion_prompt is None: + if diffusion_prompt is None or (isinstance(diffusion_prompt, list) and not diffusion_prompt): error_output = OmniRequestOutput.from_error( req_id, - f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}", + f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_stage_id}", ) logger.warning( "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; " "routing terminal error output", req_id, src_stage_id, - next_logical, + next_stage_id, ) await self.output_async_queue.put( OutputMessage( request_id=req_id, - stage_id=next_logical, + stage_id=next_stage_id, engine_outputs=error_output, metrics=None, finished=True, ) ) await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + [req_id], + abort=True, + release_owners=True, ) return if isinstance(diffusion_prompt, list): - if not diffusion_prompt: - error_output = OmniRequestOutput.from_error( - req_id, - f"Stage-{src_stage_id} produced no valid inputs for diffusion stage-{next_logical}", - ) - logger.warning( - "[Orchestrator] req=%s stage=%d produced empty diffusion inputs for stage=%d; " - "routing terminal error output", - req_id, - src_stage_id, - next_logical, - ) - await self.output_async_queue.put( - OutputMessage( - request_id=req_id, - stage_id=next_logical, - engine_outputs=error_output, - metrics=None, - finished=True, - ) - ) - await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], - ) - return if len(diffusion_prompt) == 1: diffusion_prompt = diffusion_prompt[0] else: @@ -2475,18 +2583,20 @@ async def _forward_to_next_stage_unguarded( submit_kwargs=submit_kwargs, params_override=self._maybe_clone_diffusion_params_for_cfg(req_id, params), ) + if req_state.closing: + return self._on_stage_submitted( - next_logical, + next_stage_id, req_id, replica_id, req_state, ) - req_state.stage_submit_ts[next_logical] = _time.time() + req_state.stage_submit_ts[next_stage_id] = _time.time() _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 self._emit_tx_edge( from_stage=src_stage_id, from_replica=src_replica_id if src_replica_id is not None else 0, - to_stage=next_logical, + to_stage=next_stage_id, to_pool=next_pool, request_id=req_id, tx_ms=_tx_ms, @@ -2494,7 +2604,7 @@ async def _forward_to_next_stage_unguarded( return # PD disaggregation: prefill → decode routing uses original prompt + KV transfer params - if self._pd_pair is not None and (src_stage_id, next_logical) == self._pd_pair: + if self._pd_pair is not None and (src_stage_id, next_stage_id) == self._pd_pair: params = self._build_pd_decode_params(req_id, params) # Use the original user prompt for the decode stage (not processed embeddings) @@ -2528,19 +2638,21 @@ async def _forward_to_next_stage_unguarded( replica_id = await next_pool.submit_update(req_id, req_state, request) else: replica_id = await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) + if req_state.closing: + return self._on_stage_submitted( - next_logical, + next_stage_id, req_id, replica_id, req_state, ) - req_state.stage_submit_ts[next_logical] = _time.time() + req_state.stage_submit_ts[next_stage_id] = _time.time() _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 self._emit_tx_edge( from_stage=src_stage_id, from_replica=src_replica_id if src_replica_id is not None else 0, - to_stage=next_logical, + to_stage=next_stage_id, to_pool=next_pool, request_id=req_id, tx_ms=_tx_ms, @@ -2571,7 +2683,7 @@ async def _forward_to_next_stage_unguarded( except OmniClientError as exc: await self._fail_request_client_error( req_id, - next_logical, + next_stage_id, str(exc), status_code=exc.status_code, error_type=exc.error_type, @@ -2582,9 +2694,9 @@ async def _forward_to_next_stage_unguarded( logger.exception( "[Orchestrator] req=%s process_engine_inputs FAILED for stage-%s", req_id, - next_logical, + next_stage_id, ) - if not await self._handle_forward_failure(req_id, next_logical, req_state, exc): + if not await self._handle_forward_failure(req_id, next_stage_id, req_state, exc): raise return finally: @@ -2596,7 +2708,7 @@ async def _forward_to_next_stage_unguarded( "[Orchestrator] req=%s stage-%s produced no inputs for stage-%s; waiting for more outputs", req_id, src_stage_id, - next_logical, + next_stage_id, ) return @@ -2608,29 +2720,33 @@ 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", req_id, src_stage_id, - next_logical, + next_stage_id, final_output_type or "text", final_stage_id, ) + # Later chunk receivers may already be prewarmed. Release them + # before reporting success; failed cleanup reports a request error. + await self._cleanup_request_ids([req_id], abort=True, release_owners=True) + if req_id in self.request_states: + return 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 @@ -2641,7 +2757,7 @@ async def _forward_to_next_stage_unguarded( mm_features = req_state.mm_features if model_stage == "thinker" else None request = self._build_next_stage_request( req_id, - next_logical, + next_stage_id, next_input, params=params, mm_features=mm_features, @@ -2653,19 +2769,21 @@ async def _forward_to_next_stage_unguarded( replica_id = await next_pool.submit_update(req_id, req_state, request) else: replica_id = await next_pool.submit_initial(req_id, req_state, request, prompt_text=None) + if req_state.closing: + return self._on_stage_submitted( - next_logical, + next_stage_id, req_id, replica_id, req_state, ) - req_state.stage_submit_ts[next_logical] = _time.time() + req_state.stage_submit_ts[next_stage_id] = _time.time() _tx_ms = (_time.perf_counter() - _t_submit_start) * 1000.0 self._emit_tx_edge( from_stage=src_stage_id, from_replica=src_replica_id if src_replica_id is not None else 0, - to_stage=next_logical, + to_stage=next_stage_id, to_pool=next_pool, request_id=req_id, tx_ms=_tx_ms, @@ -2690,7 +2808,10 @@ async def _prewarm_async_chunk_stages( caller still holding ``req_state`` stops instead of recording state on an object the cleanup already popped from ``request_states``. """ - if req_state.final_stage_id <= 0: + stage_path = self._request_stage_path(req_state) + if req_state.closing or self.request_states.get(request_id) is not req_state: + return False + if len(stage_path) <= 1: return True prompt_token_ids = getattr(stage0_request, "prompt_token_ids", None) @@ -2722,9 +2843,23 @@ async def _prewarm_async_chunk_stages( ) return False - for next_stage_id in range(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] + # Reserve bridge-fed senders without submitting a placeholder. + # Their downstream receiver must use the same replica as the later + # real submission, including when the default replica is unavailable. + if not await self._dispatch_or_fail_request( + lambda: next_pool.pick(request_id, request_state=req_state), + req_id=request_id, + stage_id=next_stage_id, + operation="async-chunk replica selection", + ): + return False + if req_state.closing or self.request_states.get(request_id) is not req_state: + return False 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 @@ -2737,7 +2872,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, ) @@ -2804,7 +2939,7 @@ async def _prewarm_async_chunk_stages( model_config=next_pool.stage_vllm_config.model_config, resumable=downstream_resumable, ) - request.payload_sender_info = self._build_payload_sender_info(next_stage_id - 1, request_id=request_id) + request.payload_sender_info = self._build_payload_sender_info(source_stage_id, request_id=request_id) request.external_req_id = request.request_id submitted = await self._dispatch_or_fail_request( lambda: next_pool.submit_initial( @@ -2820,7 +2955,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 @@ -2835,16 +2970,15 @@ async def _prewarm_async_chunk_stages( if req_state.pending_upstream_first_audio is not None or req_state.upstream_first_audio: await self._flush_upstream_first_audio(req_state) - if self.request_states.get(request_id) is not req_state: + if req_state.closing or 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 prewarming to the actual route edge and bound source + # replica; fall back to replica 0 when no source binding exists. _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, @@ -2860,12 +2994,13 @@ def _maybe_attach_native_kv_transfer_params( prompt: Any, ) -> None: source_stage_id = 0 - if source_stage_id + 1 > req_state.final_stage_id: + next_stage_id = self._next_request_stage(source_stage_id, req_state) + if next_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[next_stage_id].stage_type != "diffusion": return from vllm_omni.diffusion.diffusion_kv.kv_connector import build_source_kv_transfer_params, mint_transfer_id @@ -2929,7 +3064,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): @@ -2957,7 +3092,7 @@ def _build_payload_sender_info( return None 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_payload_sender_info", None) if not callable(get_sender_info): @@ -3015,10 +3150,42 @@ 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: + if type(final_stage_id) is not int or not 0 <= final_stage_id < self.num_stages: + raise ValueError(f"Final stage {final_stage_id!r} is not reachable from stage 0") + stage_path = tuple(range(final_stage_id + 1)) + 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(): # Stage 0 lost all replicas between the HTTP-layer errored check and @@ -3087,7 +3254,7 @@ async def _handle_add_request(self, msg: StageSubmissionMessage) -> None: ): return - if self.async_chunk and stage_id == 0 and final_stage_id > 0: + if self.async_chunk and req_state.final_stage_id != 0: await self._prewarm_async_chunk_stages(request_id, prompt, req_state, attach_prewarm_payload=True) async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: @@ -3097,7 +3264,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 @@ -3110,6 +3277,19 @@ async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: ) return + 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 @@ -3128,7 +3308,7 @@ async def _handle_streaming_update(self, msg: StageSubmissionMessage) -> None: ): return - if self.async_chunk and stage_id == 0 and final_stage_id > 0: + if self.async_chunk and req_state.final_stage_id != 0: await self._prewarm_async_chunk_stages(request_id, request, req_state) async def _handle_add_companion(self, msg: AddCompanionRequestMessage) -> None: @@ -3140,7 +3320,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, @@ -3192,7 +3372,7 @@ async def _handle_interaction(self, msg: InteractionMessage) -> None: request_id = msg.request_id event_id = msg.interaction.get("event_id") req_state = self.request_states.get(request_id) - if req_state is None: + if req_state is None or req_state.closing: logger.info("[Orchestrator] Dropping interaction for inactive req %s", request_id) await self.output_async_queue.put( ErrorMessage( @@ -3206,7 +3386,7 @@ async def _handle_interaction(self, msg: InteractionMessage) -> None: return try: - await self.stage_pools[stage_id].submit_interaction(request_id, msg.interaction) + await self.stage_pools[stage_id].submit_interaction(request_id, msg.interaction, request_state=req_state) except Exception as exc: logger.info( "[Orchestrator] Failed interaction for req %s: %s", diff --git a/vllm_omni/engine/stage_engine_core_proc.py b/vllm_omni/engine/stage_engine_core_proc.py index d109047a116..c370b32b9b4 100644 --- a/vllm_omni/engine/stage_engine_core_proc.py +++ b/vllm_omni/engine/stage_engine_core_proc.py @@ -13,6 +13,7 @@ import contextlib import os import signal +import time from typing import Any import vllm.v1.engine.core as _vllm_engine_core_module @@ -33,6 +34,8 @@ ) from vllm.v1.executor.uniproc_executor import UniProcExecutor +from vllm_omni.core.sched.omni_scheduling_coordinator import uses_native_mrv2_data_plane +from vllm_omni.distributed.omni_connectors.utils.config import REQUEST_RESOURCE_RELEASE_TIMEOUT_S from vllm_omni.distributed.omni_coordinator import create_stage_coord_client from vllm_omni.engine import OmniEngineCoreRequest from vllm_omni.engine.stage_init_utils import ( @@ -138,14 +141,19 @@ def omni_release_request_resources(self, request_ids: list[str]) -> None: Invoked over the UTILITY channel by the orchestrator once every stage has finished with the request. Idempotent and safe for unknown ids. """ + model_config = self.vllm_config.model_config + if uses_native_mrv2_data_plane( + model_config, use_v2_model_runner=getattr(model_config, "use_v2_model_runner", False) + ): + self.model_executor.collective_rpc("omni_release_request_resources", args=(request_ids,)) + return adapter = getattr(getattr(self, "scheduler", None), "chunk_transfer_adapter", None) if adapter is None: return - for request_id in request_ids or (): - try: - adapter.release_shm_resources(request_id) - except Exception as e: - logger.debug("omni_release_request_resources(%s) failed: %s", request_id, e) + deadline = time.monotonic() + REQUEST_RESOURCE_RELEASE_TIMEOUT_S + completions = [adapter.release_request_resources(request_id, deadline=deadline) for request_id in request_ids] + for completion in completions: + completion.result(timeout=max(0, deadline - time.monotonic())) def preprocess_add_request(self, request: OmniEngineCoreRequest) -> tuple[Any, int]: """Preserve omni payloads when vLLM builds its scheduler request.""" diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index b5d5bec6d01..e432b3b46f4 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -15,6 +15,7 @@ from vllm.v1.engine import EngineCoreOutputs from vllm.v1.metrics.stats import IterationStats +from vllm_omni.core.sched.omni_scheduling_coordinator import uses_native_mrv2_data_plane from vllm_omni.data_entry_keys import flatten_payload from vllm_omni.distributed.omni_coordinator import ( LoadBalancer, @@ -23,6 +24,7 @@ ReplicaStatus, ) from vllm_omni.distributed.omni_coordinator.load_balancer import Task +from vllm_omni.engine.errors import ResourceReleaseError from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.engine.stage_client import ( StagePoolClient, @@ -62,6 +64,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.""" @@ -100,7 +106,7 @@ class StagePool: DISPATCH_WAIT_TIMEOUT_S: float = 10.0 DISPATCH_RETRY_INTERVAL_S: float = 0.1 # A replica that dies mid-release never answers; without a bound its background release never finishes. - RELEASE_RPC_TIMEOUT_S: float = 5.0 + RELEASE_RPC_TIMEOUT_S: float = 15.0 # Only these EngineCore helpers may skip collective_rpc_async. A generic # ``{method}_async`` on AsyncMPClient must not silently drop timeout. _CACHE_RESET_METHODS = frozenset({"reset_prefix_cache", "reset_encoder_cache", "reset_mm_cache"}) @@ -138,8 +144,12 @@ def __init__( self.clients: list[StagePoolClient | None] = list(normalized_clients) self._output_processor = output_processor self._stage_vllm_config = stage_vllm_config - self._has_chunk_transfer_adapter = bool( - getattr(getattr(stage_vllm_config, "model_config", None), "async_chunk", False) + model_config = getattr(stage_vllm_config, "model_config", None) + self._has_transfer_resources = bool( + getattr(model_config, "async_chunk", False) + or uses_native_mrv2_data_plane( + model_config, use_v2_model_runner=getattr(model_config, "use_v2_model_runner", False) + ) ) self._next_replica_id = 0 self._request_bindings: dict[str, int] = {} @@ -321,6 +331,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``. @@ -329,10 +340,16 @@ 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? @@ -341,6 +358,8 @@ async def pick( 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") # Bound replica is gone or DOWN — fall through to re-select. self._affinity.pop(request_id, None) @@ -357,6 +376,8 @@ async def pick( 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) candidates = self._collect_serviceable_replicas() if candidates: # LB chose an index *into our candidates list*. @@ -835,7 +856,7 @@ def _collect_audio_metrics( def _infer_audio_sample_rate( self, - mm_output: dict[str, Any] | None = None, + mm_output: Mapping[str, Any] | None = None, *, use_default: bool = True, ) -> int: @@ -1011,17 +1032,19 @@ async def submit_initial( payload_sender_info = getattr(request, "payload_sender_info", None) if payload_sender_info is not None: submit_kwargs.setdefault("payload_sender_info", payload_sender_info) - 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) + diffusion_client = self._diffusion_client(replica_id) + await diffusion_client.add_request_async(request_id, request, params, **submit_kwargs) return replica_id - 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.clients[replica_id] if client is None: @@ -1070,9 +1093,7 @@ async def submit_update( params = req_state.sampling_params_list[self.stage_id] if self.stage_type == "diffusion": params = OmniDiffusionSamplingParams.from_params(params) - 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, request_state=req_state) client = self.clients[replica_id] if client is None: @@ -1117,11 +1138,11 @@ async def submit_interaction( self, request_id: str, interaction: OmniInteractionPrompt, + *, + request_state: OrchestratorRequestState | None = None, ) -> int: """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, request_state=request_state) client = self._diffusion_client(replica_id) result = await client.submit_interaction_async(request_id, interaction) @@ -1131,17 +1152,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 _poll_stage_raw(self, client: StagePoolLLMClient) -> EngineCoreOutputs | None: @@ -1303,9 +1313,13 @@ async def abort_requests(self, request_ids: list[str]) -> list[tuple[str, Any]]: if client is None: continue await client.abort_requests_async(engine_abort_ids) - if not is_diffusion and self._output_processor is not None: - commit = getattr(self._output_processor, "commit_aborted_request_state", None) - if callable(commit): + + # Keep prefixes available if a later replica abort fails and this pool + # must retry. The caller receives outputs only after the whole pool stops. + if not is_diffusion and self._output_processor is not None: + commit = getattr(self._output_processor, "commit_aborted_request_state", None) + if callable(commit): + for replica_request_ids in request_ids_by_replica.values(): commit(replica_request_ids, internal=False) return abort_outputs @@ -1317,20 +1331,23 @@ async def release_request_resources(self, request_ids: list[str]) -> None: bindings as part of the same teardown, so a binding lookup here would race it. The engine-core handler is idempotent for unknown ids. """ - if not request_ids or not self._has_chunk_transfer_adapter: + if not request_ids or not self._has_transfer_resources: return ids = list(request_ids) async def release(replica_id: int, call: Any) -> None: try: + if call is None: + raise RuntimeError("Replica does not support resource-release RPC") await asyncio.wait_for(call("omni_release_request_resources", ids), timeout=self.RELEASE_RPC_TIMEOUT_S) except Exception as e: - logger.warning( - "[StagePool-%s] release_request_resources on replica %s failed: %r", self.stage_id, replica_id, e - ) + raise ResourceReleaseError(f"Stage {self.stage_id} replica {replica_id} resource release failed") from e calls = [(i, getattr(self.clients[i], "call_utility_async", None)) for i in self.live_replica_ids()] - await asyncio.gather(*(release(i, call) for i, call in calls if call is not None)) + results = await asyncio.gather(*(release(i, call) for i, call in calls), return_exceptions=True) + for result in results: + if isinstance(result, BaseException): + raise result async def collective_rpc( self, diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 097c7aca8d6..8aa863b1d1c 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -187,6 +187,7 @@ async def generate( logger.debug(f"[AsyncOmni] generate() called for request {external_request_id}") input_stream_task: asyncio.Task | None = None + cleanup_complete = False try: _sleeping_tags = getattr(self, "_sleeping_tags", None) if _sleeping_tags: @@ -320,6 +321,7 @@ async def generate( ): yield output + cleanup_complete = True logger.debug(f"[AsyncOmni] Request {request_id} completed") except (asyncio.CancelledError, GeneratorExit): @@ -327,12 +329,12 @@ async def generate( # ASGI cancellation also affects subsequent awaits. Shield the # bounded engine abort before removing the local request state. with anyio.CancelScope(shield=True): - await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S) + cleanup_complete = await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S) logger.info(f"[AsyncOmni] Request {request_id} aborted.") raise except Exception as e: self._record_request_failure_once(request_id, reason="stage_error") - await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S) + cleanup_complete = await self._abort_internal_requests(request_id, timeout=ABORT_TIMEOUT_S) logger.info(f"[AsyncOmni] Request {request_id} failed (input error): {e}") raise finally: @@ -340,6 +342,8 @@ async def generate( input_stream_task.cancel() if admitting: await self._release_generate_admission() + if request_id in self.request_states: + self.request_states[request_id].cleanup_pending = not cleanup_complete self._log_summary_and_cleanup(request_id) async def _release_generate_admission(self) -> None: @@ -647,7 +651,11 @@ async def submit_interaction_async( if stage_meta.stage_type != "diffusion": raise ValueError("interaction requires a diffusion stage") - internal_ids = [s.request_id for s in self.request_states.values() if s.external_request_id == request_id] + internal_ids = [ + s.request_id + for s in self.request_states.values() + if s.external_request_id == request_id and not s.cleanup_pending + ] if not internal_ids: raise ValueError(f"No active request for interaction: {request_id!r}") if len(internal_ids) > 1: diff --git a/vllm_omni/entrypoints/async_omni_base.py b/vllm_omni/entrypoints/async_omni_base.py index c4c1bc4919c..86ca661c1be 100644 --- a/vllm_omni/entrypoints/async_omni_base.py +++ b/vllm_omni/entrypoints/async_omni_base.py @@ -66,7 +66,7 @@ def watch_task(self, task_id: str, expected_count: int = 1) -> asyncio.Future: } return fut - async def resolve(self, ack: OmniACK): + async def resolve(self, ack: OmniACK | dict): tid = getattr(ack, "task_id", None) if tid is None and isinstance(ack, dict): @@ -359,6 +359,11 @@ async def _final_output_loop(): await self.event_resolver.resolve(msg) continue + request_id = getattr(msg, "request_id", None) + req_state = self.request_states.get(request_id) if request_id is not None else None + if req_state is not None and req_state.cleanup_pending: + continue + if isinstance(msg, ErrorMessage): # Route request-scoped errors to that request's queue and # keep the loop alive. A request whose stage replica died @@ -396,6 +401,8 @@ async def _final_output_loop(): except OmniEngineDeadError as e: logger.error("[%s] Engine dead: %s", self._name, e) for req_state in list(self.request_states.values()): + if req_state.cleanup_pending: + continue error_msg = ErrorMessage( error=str(e), fatal=True, @@ -407,6 +414,8 @@ async def _final_output_loop(): except EngineDeadError as e: logger.error("[%s] Engine dead: %s", self._name, e) for req_state in list(self.request_states.values()): + if req_state.cleanup_pending: + continue error_msg = ErrorMessage( error=str(e), fatal=True, @@ -417,6 +426,8 @@ async def _final_output_loop(): except Exception as e: logger.exception("[%s] final_output_loop failed.", self._name) for req_state in list(self.request_states.values()): + if req_state.cleanup_pending: + continue error_msg = ErrorMessage( request_id=req_state.request_id, error=str(e), @@ -433,7 +444,7 @@ async def _abort_internal_requests( request_id: str | Iterable[str], *, timeout: float = ABORT_TIMEOUT_S, - ): + ) -> bool: """Abort request(s) via the Orchestrator given internal request IDs, which take the format -. """ @@ -443,6 +454,7 @@ async def _abort_internal_requests( try: # Unbind generate() if abort_async blocks in the executor. await asyncio.wait_for(self._abort(internal_req_ids, timeout=timeout), timeout=timeout) + return True except TimeoutError: logger.warning( "[AsyncOmni] Timed out aborting %s after %.1fs; " @@ -452,6 +464,7 @@ async def _abort_internal_requests( ) except Exception: logger.exception("[AsyncOmni] Cleanup abort failed for %s", ",".join(internal_req_ids)) + return False async def _abort(self, request_ids: list[str], *, timeout: float | None = None) -> None: """Abort request IDs via the engine and enqueue terminal abort outputs. @@ -473,14 +486,14 @@ async def _abort(self, request_ids: list[str], *, timeout: float | None = None) if req_id is None: continue state = self.request_states.get(req_id) - if state is None: + if state is None or state.cleanup_pending: logger.debug("[AsyncOmni] Dropping abort output for unknown req %s", req_id) continue await state.queue.put(output_msg) delivered.add(req_id) for rid in request_ids: state = self.request_states.get(rid) - if state is not None and rid not in delivered: + if state is not None and not state.cleanup_pending and rid not in delivered: queue = getattr(state, "queue", None) if queue is not None: final_stage_id = getattr(state, "final_stage_id", None) @@ -506,6 +519,9 @@ async def _abort(self, request_ids: list[str], *, timeout: float | None = None) input_stream_task = getattr(state, "input_stream_task", None) if input_stream_task is not None and not input_stream_task.done(): input_stream_task.cancel() + if state is not None and state.cleanup_pending: + self.request_states.pop(rid) + self._publish_request_gauges(len(self.request_states)) if self.log_stats: logger.info("[AsyncOmni] Aborted request(s) %s", ",".join(request_ids)) diff --git a/vllm_omni/entrypoints/client_request_state.py b/vllm_omni/entrypoints/client_request_state.py index 9a83910521c..c16368902b0 100644 --- a/vllm_omni/entrypoints/client_request_state.py +++ b/vllm_omni/entrypoints/client_request_state.py @@ -23,6 +23,9 @@ def __init__( self.queue = queue if queue is not None else asyncio.Queue() self.metrics: OrchestratorAggregator | None = None self.input_stream_task: asyncio.Task | None = None + # generate() has exited, but engine cleanup still needs an acknowledgment. + # Retain the public ID mapping for abort retries without queuing outputs. + self.cleanup_pending = False # Request-scoped idempotency guard for Prometheus failure counters. self.failure_recorded = False # Wall-clock time at which the user's request arrived in the engine diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index ca86587e9fb..b1e77c022b0 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -449,7 +449,8 @@ def _log_summary_and_cleanup(self, request_id: str, reason: str = "stage_error") request_id, ) finally: - self.request_states.pop(request_id, None) + if req_state is None or not req_state.cleanup_pending: + self.request_states.pop(request_id, None) # Republish gauges so any stale value left by the per-stage # publish in _process_single_result (which runs while the request # is still in self.request_states) is corrected after the pop. diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index aac3d6b8572..d9aed43692c 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -2234,7 +2234,11 @@ async def chat_completion_stream_generator( # `request_id` we have here is the external (user-visible) id, # so resolve via the external_request_id field. req_state = next( - (s for s in self.engine_client.request_states.values() if s.external_request_id == request_id), + ( + s + for s in self.engine_client.request_states.values() + if s.external_request_id == request_id and not s.cleanup_pending + ), None, ) if req_state is not None and req_state_audio_ref is None: diff --git a/vllm_omni/worker/mixins.py b/vllm_omni/worker/mixins.py index 724cc87be18..f8bcd3cb573 100644 --- a/vllm_omni/worker/mixins.py +++ b/vllm_omni/worker/mixins.py @@ -9,6 +9,14 @@ class OmniWorkerMixin: """Shared Omni plugin and native KV connector setup for workers.""" + model_runner: Any + + def omni_release_request_resources(self, request_ids: list[str]) -> None: + """Release native runner payloads once the pipeline is done.""" + plane = self.model_runner._omni_data_plane + if plane is not None: + plane.release_request_resources(request_ids) + def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) from vllm_omni.plugins import load_omni_general_plugins diff --git a/vllm_omni/worker_v2/omni_data_plane.py b/vllm_omni/worker_v2/omni_data_plane.py index d82e913dc32..9423426e81b 100644 --- a/vllm_omni/worker_v2/omni_data_plane.py +++ b/vllm_omni/worker_v2/omni_data_plane.py @@ -5,7 +5,8 @@ import threading import time -from collections import defaultdict +from collections import defaultdict, deque +from concurrent.futures import Future from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass, field from queue import Full, Queue @@ -16,6 +17,10 @@ from vllm.logger import init_logger from vllm_omni.data_entry_keys import OmniPayload, flatten_payload, unflatten_payload +from vllm_omni.distributed.omni_connectors.utils.config import ( + REQUEST_RESOURCE_RELEASE_TIMEOUT_S, + get_stage_connector_peer, +) from vllm_omni.worker.omni_connector_model_runner_mixin import ( OmniConnectorModelRunnerMixin, should_accumulate_full_payload_output, @@ -338,6 +343,34 @@ def abort_requests(self, req_ids: set[str]) -> int: terminal_req_ids=active_req_ids, ) + def release_request_resources(self, request_ids: list[str]) -> None: + """Fence publications and apply the existing backend cleanup contract.""" + if self._omni_connector is None: + return + # The scheduler's final runner step may follow the release RPC. Fence + # late publications with the existing abort lifecycle before reclaiming. + self.abort_requests(set(request_ids)) + deadline = time.monotonic() + REQUEST_RESOURCE_RELEASE_TIMEOUT_S + from_stage = get_stage_connector_peer(self.model_config, self._stage_id, "from_stage") + completions: list[Future[None]] = [] + with self._native_send_lock, self._lock: + for request_id in request_ids: + completion: Future[None] = Future() + completions.append(completion) + self._pending_save_reqs.setdefault(request_id, deque()).append( + { + "request_id": request_id, + "release_prefix": f"{request_id}_{self._stage_id}_", + "receive_prefix": f"{request_id}_{from_stage}_" if from_stage is not None else None, + "release_completion": completion, + "release_deadline": deadline, + } + ) + self._pending_save_counts[request_id] += 1 + self._save_work_available.set() + for completion in completions: + completion.result(timeout=max(0, deadline - time.monotonic())) + def complete_outputs( self, *, From 87cecae42cc1e522cda55014f44f8bb0bb681c61 Mon Sep 17 00:00:00 2001 From: Pinoeer-kingxi <13022943007@163.com> Date: Sat, 17 Oct 2026 19:24:10 -0700 Subject: [PATCH 6/9] [Core] Route stages through configured pipeline transitions Signed-off-by: Pinoeer-kingxi <13022943007@163.com> --- docs/configuration/README.md | 2 + docs/configuration/stage_routing.md | 95 ++++ tests/config/test_forced_aligner_injection.py | 48 ++ tests/config/test_stage_routing.py | 181 +++++++ .../sched/test_omni_scheduling_coordinator.py | 2 + .../test_chunk_transfer_adapter.py | 36 ++ tests/engine/test_arg_utils.py | 27 +- .../test_async_omni_engine_stage_init.py | 17 +- tests/engine/test_cfg_companion_lifecycle.py | 2 + tests/engine/test_configured_stage_routing.py | 503 ++++++++++++++++++ tests/engine/test_duplex_omni_engine.py | 14 + tests/engine/test_orchestrator.py | 9 + tests/engine/test_orchestrator_first_audio.py | 28 + .../test_orchestrator_kv_sender_info.py | 8 + .../test_orchestrator_stage_input_bridge.py | 5 + .../test_terminal_stage_metrics_order.py | 2 + tests/entrypoints/test_omni_entrypoints.py | 1 + tests/entrypoints/test_serve.py | 44 +- .../test_configured_connector_routing.py | 68 +++ vllm_omni/config/stage_config.py | 21 + vllm_omni/config/stage_routing.py | 63 +++ .../core/sched/omni_scheduling_coordinator.py | 1 + vllm_omni/core/sched/output.py | 6 +- vllm_omni/engine/arg_utils.py | 3 + vllm_omni/engine/duplex_omni_engine.py | 7 + vllm_omni/engine/duplex_orchestrator.py | 2 +- vllm_omni/engine/omni_engine_base.py | 3 + vllm_omni/engine/orchestrator.py | 15 +- vllm_omni/engine/stage_init_utils.py | 43 +- vllm_omni/engine/stage_runtime.py | 13 + vllm_omni/entrypoints/cli/serve.py | 7 + vllm_omni/entrypoints/omni_base.py | 13 +- vllm_omni/entrypoints/utils.py | 36 +- vllm_omni/utils/forced_aligner.py | 9 +- 34 files changed, 1283 insertions(+), 51 deletions(-) create mode 100644 docs/configuration/stage_routing.md create mode 100644 tests/config/test_stage_routing.py create mode 100644 tests/engine/test_configured_stage_routing.py create mode 100644 tests/worker/test_configured_connector_routing.py create mode 100644 vllm_omni/config/stage_routing.py 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..a35595ba2c6 --- /dev/null +++ b/docs/configuration/stage_routing.md @@ -0,0 +1,95 @@ +# 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`. +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. +Requests use the existing +[request cleanup contract](../design/module/engine_orchestration.md#request-cleanup). +Changing a route does not change the selected backend's physical-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/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 950fd5f6548..b3bc90be4f5 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -53,6 +53,42 @@ def test_pipeline_release_preserves_legacy_backend_and_fences_queued_puts(build_ connector.cleanup.assert_not_called() +@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() + + @pytest.fixture def shm_sender(build_adapter): adapter, _ = build_adapter(stage_id=0) 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_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index befa22857f7..d17ff6f7e28 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -17,7 +17,8 @@ from vllm.v1.engine.utils import EngineZmqAddresses 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 PipelineConfig, StagePipelineConfig 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 @@ -313,11 +314,19 @@ 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),)]) +def test_async_omni_engine_initialize_stages_passes_log_stats_and_client_config_to_runtime(monkeypatch, transitions): 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, + ) + 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 @@ -348,12 +357,14 @@ def _capture_create_stage_runtime(**kwargs): assert captured["stage_init_timeout"] == 7 assert captured["log_stats"] is True assert captured["client_config"] is engine._client_config + assert captured["stage_transitions"] == transitions 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 bd64c6719d7..27556f1da99 100644 --- a/tests/engine/test_cfg_companion_lifecycle.py +++ b/tests/engine/test_cfg_companion_lifecycle.py @@ -19,6 +19,7 @@ import pytest +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, OrchestratorRequestState @@ -103,6 +104,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 diff --git a/tests/engine/test_configured_stage_routing.py b/tests/engine/test_configured_stage_routing.py new file mode 100644 index 00000000000..eb66d18d72d --- /dev/null +++ b/tests/engine/test_configured_stage_routing.py @@ -0,0 +1,503 @@ +# 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 +from types import SimpleNamespace + +import pytest +from vllm.sampling_params import SamplingParams +from vllm.v1.engine import EngineCoreOutputs, FinishReason + +from tests.engine.test_orchestrator import ( + FakeOutputProcessor, + FakePromptRequest, + FakeStageClient, + _build_request_output, + _build_stage_pools, + _engine_core_outputs, + _enqueue_add_request, + _get_output_message, + _shutdown_orchestrator, + _wait_for, + orchestrator_factory, +) +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.errors import ResourceReleaseError +from vllm_omni.engine.messages import ErrorMessage, StageSubmissionMessage +from vllm_omni.engine.omni_engine_base import OmniEngineBase +from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState +from vllm_omni.entrypoints.omni_base import OmniBase +from vllm_omni.entrypoints.utils import get_final_stage_id_for_e2e + +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 = [ + 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 +async def test_abort_retry_preserves_earlier_stage_tokens_after_later_stage_abort_failure(mocker): + orch, state, clients = _direct_orchestrator(((0, 2), (2, 1))) + clients[2].final_output = True + state.final_output_stage_ids = {1, 2} + for sid in (1, 2): + pool = orch.stage_pools[sid] + pool.select_replica_id("r") + pool.output_processor.request_outputs = [_build_request_output("r", token_ids=[sid], finished=False)] + abort = mocker.patch.object( + clients[2], "abort_requests_async", side_effect=[RuntimeError("stage abort failed"), None] + ) + with pytest.raises(RuntimeError, match="stage abort failed"): + await orch._cleanup_request_ids(["r"], abort=True) + assert orch.stage_pools[1].get_bound_replica_id("r") is None + outputs = await orch._cleanup_request_ids(["r"], abort=True) + assert [(msg.stage_id, msg.finished) for msg in outputs] == [(2, False), (1, True)] + assert [list(msg.engine_outputs.outputs[0].token_ids) for msg in outputs] == [[2], [1]] + assert clients[1].abort_calls == [["r"]] and abort.await_count == 2 + assert "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("empty_inputs", [None, []]) +async def test_empty_diffusion_cleanup_failure_keeps_request_closed(mocker, empty_inputs): + orch, state, clients = _direct_orchestrator() + clients[2].stage_type = "diffusion" + clients[2].custom_process_input_func = mocker.Mock(return_value=empty_inputs) + mocker.patch.object(orch, "_release_stage_transfer_resources", side_effect=ResourceReleaseError("release failed")) + + await orch._forward_to_next_stage("r", 0, _build_request_output("r"), state) + + assert orch.output_async_queue.get_nowait().finished + error = orch.output_async_queue.get_nowait() + assert isinstance(error, ErrorMessage) and not error.fatal and error.request_id == "r" + assert orch.request_states["r"] is state and state.closing + assert not clients[2].add_request_calls + assert orch.rpc_async_queue.empty() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event_driven", ["0", "1"]) +@pytest.mark.parametrize("transitions", [((0, 2),), ((0, 2), (2, 1))]) +async def test_configured_route_runs_through_real_loop( + monkeypatch, configured_orchestrator_factory, event_driven, transitions +): + monkeypatch.setenv("VLLM_OMNI_EVENT_DRIVEN_ORCH", event_driven) + 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") + 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) diff --git a/tests/engine/test_duplex_omni_engine.py b/tests/engine/test_duplex_omni_engine.py index 90297273556..577801a26e5 100644 --- a/tests/engine/test_duplex_omni_engine.py +++ b/tests/engine/test_duplex_omni_engine.py @@ -12,6 +12,7 @@ import pytest +from vllm_omni.config.stage_config import 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 @@ -29,6 +30,19 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +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() + + class _FakeRpcClient: def __init__(self, result: Any) -> None: self.result = result diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index d33b9dbc497..f2ede54e31c 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -21,6 +21,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 ( @@ -430,6 +431,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. @@ -464,6 +466,7 @@ async def _run() -> None: rpc_async_queue=rpc_queue.async_q, stage_pools=stage_pools, async_chunk=async_chunk, + stage_transitions=stage_transitions, log_stats=log_stats, ) ready_future.set_result((orchestrator, request_queue, output_queue, rpc_queue)) @@ -1525,6 +1528,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: @@ -1607,6 +1611,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}) @@ -1671,6 +1676,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( @@ -1715,6 +1721,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): @@ -2116,6 +2123,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)) orchestrator.request_states = {} with pytest.raises(RuntimeError, match="stage abort failed"): @@ -2146,6 +2154,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_first_audio.py b/tests/engine/test_orchestrator_first_audio.py index 1e73e2331b3..5437341c3c2 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 @@ -69,6 +70,7 @@ async def process_llm_raw_outputs(self, replica_id, raw, **kwargs): def _orchestrator(output_kind=RequestOutputKind.DELTA, *, registered=True, replica_id=2, final_stage_id=1): obj = Orchestrator.__new__(Orchestrator) obj.stage_pools = [SimpleNamespace(final_output=False), _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() @@ -137,12 +139,34 @@ def _audio(output): return torch.cat(audio) if isinstance(audio, list) else audio +@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 + + @pytest.mark.asyncio @pytest.mark.parametrize("kind", list(RequestOutputKind)) @pytest.mark.parametrize("suffix_overtook", [False, True]) 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: @@ -169,6 +193,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) @@ -180,6 +205,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)) @@ -246,6 +272,7 @@ async def test_prewarm_flushes_first_audio_after_successful_submission(mocker, p if processing_failed: later_pool = SimpleNamespace(submit_initial=AsyncMock()) obj.stage_pools.append(later_pool) + obj._stage_routing = StageRouting.from_transitions(len(obj.stage_pools)) req_state = obj.request_states["r"] req_state.sampling_params_list = [None] * (req_state.final_stage_id + 1) await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11))) @@ -321,6 +348,7 @@ async def test_later_aligner_receives_full_codec_waveform(): 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) diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index 00460640501..a0d5df12b4a 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -11,6 +11,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 @@ -356,6 +357,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() @@ -386,6 +388,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", @@ -417,6 +420,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() @@ -453,6 +457,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 = {} @@ -489,6 +494,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( @@ -532,6 +538,7 @@ async def submit(key, state, request, **kwargs): pick=AsyncMock(side_effect=lambda key, **kwargs: {"a": 7, "b": 1}[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 @@ -599,6 +606,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"]}}, diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 9c49f577b8d..8f014dc38e5 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, @@ -219,6 +220,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() req_state = OrchestratorRequestState( @@ -269,6 +271,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( @@ -305,6 +308,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", @@ -330,6 +334,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() 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 cf22c79f177..c1d822b3164 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -104,6 +104,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 48382d98bcd..129a82b32b9 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.stage_engine_startup import StageReplicaResources from vllm_omni.engine.stage_runtime import StageEngineLaunch from vllm_omni.entrypoints.cli.serve import ( @@ -327,8 +328,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 @@ -342,14 +344,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 = [ @@ -359,7 +361,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( @@ -369,7 +371,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) @@ -377,6 +383,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]) @@ -852,7 +861,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.""" @@ -871,7 +883,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) @@ -919,6 +939,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 2b59f4ed153..e304e065a50 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -19,6 +19,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 @@ -345,6 +346,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() @@ -387,6 +391,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 9ad4ec5a9fc..48f2a9f5182 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -457,6 +457,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 60195e7d339..54f704d5946 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -99,9 +99,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/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index b3a23d8bd64..0e47c8f3c70 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -328,6 +328,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/duplex_omni_engine.py b/vllm_omni/engine/duplex_omni_engine.py index 98ea81fa0b5..d5e80ea70ea 100644 --- a/vllm_omni/engine/duplex_omni_engine.py +++ b/vllm_omni/engine/duplex_omni_engine.py @@ -72,6 +72,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/duplex_orchestrator.py b/vllm_omni/engine/duplex_orchestrator.py index 39719453aec..35f079999d4 100644 --- a/vllm_omni/engine/duplex_orchestrator.py +++ b/vllm_omni/engine/duplex_orchestrator.py @@ -208,7 +208,7 @@ async def _handle_forward_failure( if runner is not None: runner.on_stage_failure(next_stage_id, exc, request_id=req_id) await self._cleanup_request_ids( - [req_id, *self._cfg_tracker.cleanup_parent(req_id)], + [req_id], abort=True, release_owners=close_session, ) diff --git a/vllm_omni/engine/omni_engine_base.py b/vllm_omni/engine/omni_engine_base.py index 793711682d2..4e1f3da38b9 100644 --- a/vllm_omni/engine/omni_engine_base.py +++ b/vllm_omni/engine/omni_engine_base.py @@ -377,6 +377,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, @@ -453,6 +454,7 @@ async def _run_orchestrator() -> None: output_async_queue=self.output_queue.async_q, rpc_async_queue=self.rpc_output_queue.async_q, stage_pools=self.stage_pools, + stage_transitions=self.pipeline_config.stage_transitions if self.pipeline_config is not None else None, async_chunk=self.async_chunk, pd_config=pd_config, membership_controller=membership_controller, @@ -882,6 +884,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 () self._event_driven_orch_default = _event_driven_orch_default_for_pipeline( pipeline_config.model_type if pipeline_config is not None else None diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 79caf3f3a34..486098bb7ec 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -34,6 +34,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 @@ -362,6 +363,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, @@ -379,6 +381,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 self._prom_metrics = prom_metrics self._stage_replica_waiting: dict[tuple[int, int], int] = {} @@ -748,7 +751,7 @@ async def _abort_request_ids(self, request_ids: list[str]) -> list[OutputMessage if (state := self.request_states.get(request_id)) is not None for output in state.pending_abort_outputs ] - abort_outputs.sort(key=lambda output: output.stage_id) + 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 @@ -2128,10 +2131,10 @@ async def _route_output( await self._publish_final_output(req_state, pending) def _request_stage_path(self, req_state: OrchestratorRequestState) -> tuple[int, ...]: - return tuple(range(req_state.final_stage_id + 1)) + return self._stage_routing.path_to(req_state.final_stage_id) def _next_request_stage(self, stage_id: int, req_state: OrchestratorRequestState) -> int | None: - return stage_id + 1 if stage_id < req_state.final_stage_id else None + return self._stage_routing.next_stage(stage_id, req_state.final_stage_id) def _stage_receives_async_chunks(self, stage_id: int) -> bool: """Whether a stage's connector supplies its runtime inputs.""" @@ -3133,7 +3136,7 @@ async def _dispatch_message(self, msg: EngineQueueMessage) -> bool: def _native_mrv2_receiver_stage(self, final_stage_id: int) -> 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): + for stage_id in self._stage_routing.path_to(final_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( @@ -3163,9 +3166,7 @@ async def _handle_add_request(self, msg: StageSubmissionMessage) -> None: ) return try: - if type(final_stage_id) is not int or not 0 <= final_stage_id < self.num_stages: - raise ValueError(f"Final stage {final_stage_id!r} is not reachable from stage 0") - stage_path = tuple(range(final_stage_id + 1)) + stage_path = self._stage_routing.path_to(final_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) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 4413cd01aad..d735a42ce36 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -49,9 +49,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 @@ -1933,9 +1935,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_runtime.py b/vllm_omni/engine/stage_runtime.py index dc43575a56c..1e59608aa13 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, ) @@ -170,12 +171,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 @@ -767,6 +774,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] @@ -1366,6 +1374,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, @@ -1373,6 +1382,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, @@ -1691,6 +1701,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: @@ -1714,6 +1725,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, @@ -1731,6 +1743,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 765d540a3d3..01aabb24378 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.entrypoints.cli.logo import log_logo from vllm_omni.entrypoints.openai.api_server import ( @@ -1187,6 +1188,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)), ) @@ -1435,6 +1437,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 b1e77c022b0..6647146b646 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, @@ -456,11 +457,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]: @@ -470,8 +478,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, From f15372216dda0901ffb4ddcb680e4530f0847aaf Mon Sep 17 00:00:00 2001 From: Pinoeer-kingxi <13022943007@163.com> Date: Sat, 17 Oct 2026 21:47:52 -0700 Subject: [PATCH 7/9] [Bugfix] Preserve admission affinity and drain failed aborts Signed-off-by: Pinoeer-kingxi <13022943007@163.com> --- docs/design/module/engine_orchestration.md | 4 + tests/engine/test_configured_stage_routing.py | 294 +++++++++++++++++- tests/engine/test_duplex_orchestrator.py | 3 +- tests/engine/test_orchestrator.py | 1 + vllm_omni/engine/orchestrator.py | 84 +++-- vllm_omni/engine/stage_pool.py | 53 ++-- 6 files changed, 393 insertions(+), 46 deletions(-) diff --git a/docs/design/module/engine_orchestration.md b/docs/design/module/engine_orchestration.md index 26b5c78b75f..5c3334bb3f1 100644 --- a/docs/design/module/engine_orchestration.md +++ b/docs/design/module/engine_orchestration.md @@ -139,6 +139,10 @@ 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. +An abort failure still attempts the other stages and replicas. The request +remains closed to admission while unacknowledged stages 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/engine/test_configured_stage_routing.py b/tests/engine/test_configured_stage_routing.py index add90b6682b..69d713eb9bf 100644 --- a/tests/engine/test_configured_stage_routing.py +++ b/tests/engine/test_configured_stage_routing.py @@ -3,15 +3,18 @@ """Real orchestrator and StagePool logic with lightweight engine clients (L1).""" import asyncio +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, @@ -26,11 +29,16 @@ 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 ErrorMessage, StageSubmissionMessage +from vllm_omni.engine.messages import AbortRequestMessage, ErrorMessage, StageSubmissionMessage from vllm_omni.engine.omni_engine_base import OmniEngineBase -from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState +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 @@ -534,7 +542,7 @@ async def test_reordered_route_starts_at_admitted_nonzero_entry(mocker, async_ch @pytest.mark.asyncio async def test_closing_request_does_not_consume_its_abort_snapshot(mocker): orch, state, _ = _direct_orchestrator() - state.closing = True + 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")] @@ -545,3 +553,283 @@ async def test_closing_request_does_not_consume_its_abort_snapshot(mocker): 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) 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 203e4796baf..39b8a564e47 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -238,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]] = [] diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 47649e273b7..55f7d8a6ce6 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -24,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 @@ -250,8 +250,8 @@ class OrchestratorRequestState: 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 - # Stops admission while stage aborts yield to other orchestrator tasks. - closing: bool = False + # 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 @@ -696,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. @@ -721,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: @@ -752,6 +760,8 @@ 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): @@ -884,17 +894,24 @@ async def _process_llm_stage_outputs( ``_finish_raw_terminal_requests`` once routing is done. """ pool = self.stage_pools[stage_id] - if any( - (state := self.request_states.get(eco.request_id)) is not None and state.closing - for eco in raw_outputs.outputs - ): - # Leave the processor's abort snapshot intact while abort RPCs yield. + 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 = [ - eco - for eco in raw_outputs.outputs - if (state := self.request_states.get(eco.request_id)) is None or not state.closing - ] + 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) @@ -1855,10 +1872,17 @@ async def _cleanup_request_ids( cleanup_ids.append(cid) for request_id in cleanup_ids: if (state := self.request_states.get(request_id)) is not None: - state.closing = True + 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: @@ -1871,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, @@ -1950,7 +1986,11 @@ async def _finish_raw_terminal_requests( 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.closing 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 pending = req_state.pending_final_output if pending is not None: @@ -2021,12 +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): @@ -3110,6 +3154,8 @@ async def _prewarm_async_chunk_stages( 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): + continue for route_stage_id in (source_stage_id, next_stage_id): try: pinned_replica_id = self._session_route(route_stage_id, req_state) @@ -3125,8 +3171,6 @@ async def _prewarm_async_chunk_stages( return False if req_state.closing or self.request_states.get(request_id) is not req_state: return False - if not self._stage_receives_async_chunks(next_stage_id): - continue 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() diff --git a/vllm_omni/engine/stage_pool.py b/vllm_omni/engine/stage_pool.py index 75af4f3fb80..e806842da89 100644 --- a/vllm_omni/engine/stage_pool.py +++ b/vllm_omni/engine/stage_pool.py @@ -361,32 +361,29 @@ async def pick( 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 - if request_state is not None: - raise StageUnavailableError(f"stage {self.stage_id} bound replica is no longer serviceable") - # 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*. @@ -1365,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 @@ -1395,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: From df51893e89363c341ac6393f48a6b1dd4fcb0b28 Mon Sep 17 00:00:00 2001 From: Pinoeer-kingxi <13022943007@163.com> Date: Sat, 17 Oct 2026 21:50:21 -0700 Subject: [PATCH 8/9] [Doc] Clarify abort-state cleanup after natural completion Signed-off-by: Pinoeer-kingxi <13022943007@163.com> --- docs/design/module/engine_orchestration.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/design/module/engine_orchestration.md b/docs/design/module/engine_orchestration.md index 5c3334bb3f1..3479ea45243 100644 --- a/docs/design/module/engine_orchestration.md +++ b/docs/design/module/engine_orchestration.md @@ -138,7 +138,7 @@ 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. +physical EngineCore abort succeeds or a natural terminal confirms stage completion. An abort failure still attempts the other stages and replicas. The request remains closed to admission while unacknowledged stages drain; their natural terminals permit cleanup without publishing a normal completion. Correlated From 539648e794954acfe549ef26bbefe32f4d960ac1 Mon Sep 17 00:00:00 2001 From: Pinoeer-kingxi <13022943007@163.com> Date: Sat, 17 Oct 2026 22:36:37 -0700 Subject: [PATCH 9/9] [Bugfix] Fence late bridge validation errors after cancellation Signed-off-by: Pinoeer-kingxi <13022943007@163.com> --- docs/design/module/engine_orchestration.md | 8 +- tests/engine/test_configured_stage_routing.py | 110 +++++++++++++++++- vllm_omni/engine/orchestrator.py | 2 + 3 files changed, 116 insertions(+), 4 deletions(-) diff --git a/docs/design/module/engine_orchestration.md b/docs/design/module/engine_orchestration.md index 3479ea45243..28f5a3ee52d 100644 --- a/docs/design/module/engine_orchestration.md +++ b/docs/design/module/engine_orchestration.md @@ -137,10 +137,12 @@ 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 or a natural terminal confirms stage completion. +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 unacknowledged stages drain; their natural +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 diff --git a/tests/engine/test_configured_stage_routing.py b/tests/engine/test_configured_stage_routing.py index 69d713eb9bf..f9e33c79e06 100644 --- a/tests/engine/test_configured_stage_routing.py +++ b/tests/engine/test_configured_stage_routing.py @@ -3,6 +3,7 @@ """Real orchestrator and StagePool logic with lightweight engine clients (L1).""" import asyncio +import threading from dataclasses import replace from types import SimpleNamespace @@ -21,15 +22,17 @@ _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, ErrorMessage, StageSubmissionMessage +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 @@ -833,3 +836,108 @@ async def fail_abort(request_ids): 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/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 55f7d8a6ce6..d43e51e9dca 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -3050,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