diff --git a/tests/core/sched/test_generation_scheduler_restore.py b/tests/core/sched/test_generation_scheduler_restore.py index 04d560f87..ab7c65aa7 100644 --- a/tests/core/sched/test_generation_scheduler_restore.py +++ b/tests/core/sched/test_generation_scheduler_restore.py @@ -133,7 +133,9 @@ def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) -> monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.create_request_queue", create_request_queue) scheduler._native_data_plane = True scheduler.chunk_transfer_adapter = None - scheduler.input_coordinator = SimpleNamespace(finished_requests=set(), restore_queues=lambda *args, **kwargs: None) + scheduler.input_coordinator = SimpleNamespace( + _async_chunk=True, finished_requests=set(), restore_queues=lambda *args, **kwargs: None + ) scheduler.running = [] scheduler.waiting = create_request_queue(scheduler.policy) scheduler.skipped_waiting = create_request_queue(scheduler.policy) diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index 717fc9672..94d796fd6 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -1527,7 +1527,9 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati sched.use_v2_model_runner = native sched.max_num_running_reqs = 8 sched.input_coordinator = ( - SimpleNamespace(_waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None) if native else None + SimpleNamespace(_async_chunk=True, _waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None) + if native + else None ) sched.chunk_transfer_adapter = ( None diff --git a/tests/core/sched/test_omni_scheduler_mixin_shared.py b/tests/core/sched/test_omni_scheduler_mixin_shared.py index ba4b2454c..ed18de0b9 100644 --- a/tests/core/sched/test_omni_scheduler_mixin_shared.py +++ b/tests/core/sched/test_omni_scheduler_mixin_shared.py @@ -6,8 +6,10 @@ from types import SimpleNamespace from typing import Any import pytest +from vllm.config import SchedulerConfig, VllmConfig from vllm.v1.engine import FinishReason +from vllm_omni.config.model import OmniModelConfig from vllm_omni.core.sched import omni_scheduler_mixin from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin from vllm_omni.core.sched.output import OmniChunkRecvHandle @@ -90,6 +92,31 @@ def test_full_payload_coordinator_matches_legacy_gate(monkeypatch, stage_id, asy assert (scheduler.input_coordinator is not None) is enabled +@pytest.mark.parametrize("async_chunk", [False, True]) +@pytest.mark.parametrize(("stage_id", "required"), [(0, False), (1, False), (1, True)]) +def test_native_data_plane_uses_the_selected_input_protocol(async_chunk, stage_id, required): + scheduler = _Scheduler() + model_config = object.__new__(OmniModelConfig) + model_config.stage_id = stage_id + model_config.async_chunk = async_chunk + model_config.requires_full_payload_input = required + model_config.supports_native_mrv2_data_plane = True + model_config.use_v2_model_runner = True + scheduler.vllm_config = object.__new__(VllmConfig) + scheduler.vllm_config.model_config = model_config + scheduler.vllm_config.scheduler_config = object.__new__(SchedulerConfig) + scheduler.vllm_config.scheduler_config.max_num_seqs = 1 + scheduler._init_omni_io_scheduling_state() + + assert scheduler._native_data_plane + if async_chunk or (stage_id > 0 and required): + assert scheduler.input_coordinator._async_chunk is async_chunk + else: + assert scheduler.input_coordinator is None + assert scheduler._async_chunk_transport_enabled() is async_chunk + assert scheduler.chunk_transfer_adapter is None + + def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths(): calls: list[tuple[Any, ...]] = [] scheduler = _Scheduler() diff --git a/tests/core/sched/test_omni_scheduler_mixin_timeouts.py b/tests/core/sched/test_omni_scheduler_mixin_timeouts.py index e7d126cd9..25aa74e2b 100644 --- a/tests/core/sched/test_omni_scheduler_mixin_timeouts.py +++ b/tests/core/sched/test_omni_scheduler_mixin_timeouts.py @@ -6,12 +6,8 @@ ``_process_pending_chunk_timeouts`` covers async-chunk. Both share ``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S``. -Verifies that the mixin correctly *delegates* timed-out requests to the -base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``. -The end-to-end effect (queue removal + status set + per-request cleanup + -client-facing FINISHED_ERROR emission) is the responsibility of upstream -vLLM's ``finish_requests`` implementation and is covered by upstream tests; -this file only asserts the wiring from the mixin to that API. +Verifies scheduler cleanup and explicit client-facing ERROR outputs. Upstream +``finish_requests`` frees requests but does not emit their terminal output. """ from __future__ import annotations @@ -23,6 +19,9 @@ import os from types import SimpleNamespace import pytest +from vllm import SamplingParams +from vllm.v1.engine import EngineCoreOutputs +from vllm.v1.request import Request from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin @@ -52,7 +51,7 @@ class _FakeScheduler(OmniSchedulerMixin): def test_process_pending_input_timeouts_delegates_to_finish_requests(): """Timed-out request present in self.requests is forwarded to finish_requests.""" req_id = "stuck-req" - requests = {req_id: SimpleNamespace(request_id=req_id)} + requests = {req_id: SimpleNamespace(request_id=req_id, client_index=0)} coord = _FakeCoordinator(timed_out_ids={req_id}) scheduler = _FakeScheduler(requests, coord) @@ -142,7 +141,7 @@ def test_process_pending_chunk_timeouts_delegates_to_finish_requests(): """ req_id = "stalled-stream" adapter = _FakeChunkAdapter(timed_out_ids={req_id}) - scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id)}, adapter) + scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id, client_index=0)}, adapter) scheduler._process_pending_chunk_timeouts() @@ -154,6 +153,41 @@ def test_process_pending_chunk_timeouts_delegates_to_finish_requests(): assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR") +@pytest.mark.parametrize("chunked", [False, True]) +@pytest.mark.parametrize("synthesize_abort_outputs", [False, True]) +def test_timeout_emits_error_once_even_without_model_outputs(chunked, synthesize_abort_outputs): + from vllm.v1.engine import FinishReason + + request = Request( + request_id="stuck", + prompt_token_ids=[1], + sampling_params=SamplingParams(max_tokens=1), + pooling_params=None, + client_index=3, + ) + if chunked: + scheduler = _FakeChunkScheduler({"stuck": request}, _FakeChunkAdapter({"stuck"})) + scheduler._process_pending_chunk_timeouts() + else: + scheduler = _FakeScheduler({"stuck": request}, _FakeCoordinator({"stuck"})) + scheduler._process_pending_input_timeouts() + # Model upstream's finish_requests cleanup. Generation must still emit + # ERROR, and AR's synthetic ABORT must not replace or duplicate it. + scheduler.requests.clear() + scheduler.finished_req_ids_dict = {3: {"stuck"}} + outputs: dict[int, EngineCoreOutputs] = {} + scheduler._attach_finished_request_sets(outputs, synthesize_abort_outputs=synthesize_abort_outputs) + + [error] = outputs[3].outputs + assert error.request_id == "stuck" + assert error.finish_reason == FinishReason.ERROR + assert "connector input" in error.stop_reason + assert outputs[3].finished_requests == {"stuck"} + outputs = {} + scheduler._attach_finished_request_sets(outputs, synthesize_abort_outputs=synthesize_abort_outputs) + assert outputs == {} + + def test_process_pending_chunk_timeouts_skips_already_freed_request(): adapter = _FakeChunkAdapter(timed_out_ids={"already-freed"}) scheduler = _FakeChunkScheduler(requests={}, adapter=adapter) diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 4bbca2515..4a5bc435c 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -2000,6 +2000,47 @@ async def test_abort_retry_does_not_repeat_successful_stage_abort(): assert second.bound == set() +@pytest.mark.asyncio +@pytest.mark.parametrize("stage_id", [0, 1]) +async def test_raw_stage_error_reaches_caller_before_normal_terminal_routing(mocker, stage_id): + pools = _build_stage_pools([[FakeStageClient()], [FakeStageClient()]]) + for pool in pools: + mocker.patch.object(pool, "process_llm_raw_outputs", new_callable=mocker.AsyncMock, return_value=[]) + output_queue: asyncio.Queue[ErrorMessage] = asyncio.Queue() + orchestrator = Orchestrator( + request_async_queue=asyncio.Queue(), + output_async_queue=output_queue, + rpc_async_queue=asyncio.Queue(), + stage_pools=[], + ) + orchestrator.stage_pools = 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) + normal_terminal = mocker.patch.object( + orchestrator, "_apply_raw_terminal_stage_finish", new_callable=mocker.AsyncMock + ) + reason = "Timed out waiting for connector input after 5s" + raw = EngineCoreOutputs( + outputs=[ + OmniEngineCoreOutput( + request_id="stuck", new_token_ids=[], finish_reason=FinishReason.ERROR, stop_reason=reason + ) + ] + ) + terminal_ids: set[str] = set() + + await orchestrator._process_llm_stage_outputs(stage_id, 0, raw, terminal_ids) + + error = output_queue.get_nowait() + assert isinstance(error, ErrorMessage) + assert (error.request_id, error.stage_id, error.error) == ("stuck", stage_id, reason) + cleanup.assert_awaited_once_with(["stuck"], abort=True, release_owners=True) + normal_terminal.assert_not_awaited() + assert not terminal_ids + assert output_queue.empty() + + @pytest.mark.asyncio async def test_duplex_session_request_error_finish_is_delivered_as_request_error() -> None: """A session-owned request the scheduler finished with FinishReason.ERROR diff --git a/tests/worker_v2/test_omni_data_plane.py b/tests/worker_v2/test_omni_data_plane.py index a1629e0bf..ccb51c460 100644 --- a/tests/worker_v2/test_omni_data_plane.py +++ b/tests/worker_v2/test_omni_data_plane.py @@ -48,6 +48,8 @@ def _bare_plane(*, delivery_timeout_s=1.0, shutdown_timeout_s=1.0): _native_output_closed=False, _native_outputs_in_flight=defaultdict(int), _native_terminal_pending=set(), + _async_chunk=True, + _pending_full_payload_send={}, _put_req_chunk=defaultdict(int), _ramp_chunk_count=defaultdict(int), _delivery_manager=OmniDeliveryManager( @@ -82,6 +84,76 @@ def plane(): p._stop_output_worker() +def test_full_payload_waits_for_terminal_and_last_deferred_frame(plane): + plane._async_chunk = False + plane._full_payload_replace_keys_cached = frozenset({"codes.ref"}) + request = _new_request() + request.resumable = False + plane.register_request(request) + first = torch.tensor([[1, 2]]) + last = torch.tensor([[3, 4]]) + ref = torch.tensor([[5, 6]]) + plane.reserve_outputs(["internal"]) + assert _complete(plane, [{"codes.audio": first, "codes.ref": ref}], token=21) == 0 + assert not plane.record.batches + + plane.reserve_outputs(["internal"]) + assert plane.request_terminal({"internal"}) == 0 + assert not plane.record.cleaned + assert _complete(plane, [{"codes.audio": last, "codes.ref": ref}], token=2150) == 1 + + [(snapshot, payload)] = plane.record.batches[0] + assert snapshot.is_finished() + assert snapshot.output_token_ids == [21, 2150] + torch.testing.assert_close(payload["codes.audio"], torch.cat([first, last])) + torch.testing.assert_close(payload["codes.ref"], ref) + assert plane.record.cleaned == ["internal"] + assert not plane._pending_full_payload_send + assert plane.request_terminal({"internal"}) == 0 + + +def test_full_payload_abort_discards_partial_and_late_outputs(plane): + plane._async_chunk = False + plane._full_payload_replace_keys_cached = frozenset() + plane.register_request(_new_request()) + assert _complete(plane, [{"codes.audio": torch.tensor([[1, 2]])}]) == 0 + plane.reserve_outputs(["internal"]) + assert plane.abort_requests({"internal"}) == 1 + [(snapshot, payload)] = plane.record.batches[0] + assert snapshot.is_finished() and payload is None + assert _complete(plane, [{"codes.audio": torch.tensor([[3, 4]])}]) == 0 + assert len(plane.record.batches) == 1 + assert not plane._pending_full_payload_send + + +def test_full_payload_qwen_builder_receives_complete_codec(raw_plane, monkeypatch): + from vllm_omni.model_executor.stage_input_processors.qwen3_tts import talker2code2wav_full_payload + + raw_plane._async_chunk = False + raw_plane._omni_connector = object() + raw_plane._request_ids_mapping = {} + raw_plane._custom_process_func = talker2code2wav_full_payload + raw_plane._custom_process_batch_func = None + monkeypatch.setattr(raw_plane, "is_data_transfer_rank", lambda: True) + sent = [] + + def publish(entries, **kwargs): + sent.extend(entries) + return len(entries) + + monkeypatch.setattr(raw_plane, "_publish_chunk_cohort", publish) + raw_plane.register_request(_new_request()) + for token in (21, 22, 2150): + _complete(raw_plane, [{"codes.audio": torch.full((1, 16), token)}], token=token) + assert not sent + assert raw_plane.request_terminal({"internal"}) == 1 + [(request, payload)] = sent + assert request.output_token_ids == [21, 22, 2150] + assert payload["meta"]["finished"].item() + # EOS is filtered, and the two remaining frames are codebook-major. + assert payload["codes"]["audio"].tolist() == [21, 22] * 16 + + def _gated_connector(): # put() blocks its first call until released (or forever without release). conn = SimpleNamespace( diff --git a/tests/worker_v2/test_omni_gpu_model_runner.py b/tests/worker_v2/test_omni_gpu_model_runner.py index 5caf8009a..835f2d5f5 100644 --- a/tests/worker_v2/test_omni_gpu_model_runner.py +++ b/tests/worker_v2/test_omni_gpu_model_runner.py @@ -15,6 +15,7 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.worker.gpu.model_runner import GPUModelRunner from vllm.v1.worker.gpu.sample.sampler import Sampler +from vllm_omni.config.model import OmniModelConfig from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.worker_v2.model_states import init_omni_model_state from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState @@ -53,6 +54,8 @@ def test_add_requests_empty_admission_and_stop_id_sanitization(): def test_prepare_native_data_plane_terminal_abort_split_and_warmup_skip(): runner = _make_runner() + runner.model_config = object.__new__(OmniModelConfig) + runner.model_config.async_chunk = True plane = SimpleNamespace( register_request=MagicMock(), register_receivers=MagicMock(), @@ -78,6 +81,19 @@ def test_prepare_native_data_plane_terminal_abort_split_and_warmup_skip(): plane.abort_requests.assert_called_once_with({"aborted"}) +def test_full_payload_receive_is_polled_without_scheduled_tokens(mocker): + runner = object.__new__(OmniGPUModelRunner) + runner.model_config = object.__new__(OmniModelConfig) + runner.model_config.async_chunk = False + plane = mocker.Mock() + runner._omni_data_plane = plane + scheduler_output = SchedulerOutput.make_empty() + + runner._prepare_native_data_plane(scheduler_output) + + plane.recv_full_payload_inputs.assert_called_once_with(scheduler_output) + + @pytest.mark.parametrize("output_form", ["tuple", "omni"]) def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form): runner = object.__new__(OmniGPUModelRunner) diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 756fd6e07..179781c90 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -189,7 +189,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): Stateful codecs retain lifetime admission until their decoder state is released, even when no chunk is currently executing. """ - if not getattr(self, "_native_data_plane", False) or self._retains_state_across_chunks: + if ( + not getattr(self, "_native_data_plane", False) + or not self._async_chunk_transport_enabled() + or self._retains_state_across_chunks + ): return in_flight: list[Request] = [] for request in self.running: @@ -236,7 +240,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self._drop_aborted_queued_requests() self._resync_streaming_input_counter() async_chunk_transport = self._async_chunk_transport_enabled() - native_chunks = bool(getattr(self, "_native_data_plane", False)) + native_chunks = bool(getattr(self, "_native_data_plane", False) and async_chunk_transport) # Parking releases an execution slot, but stateful codecs retain their # request-owned state until completion or abort, in either runner. reserved_running_slots = ( diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 78a6766bd..0f17fb54d 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -112,6 +112,7 @@ class OmniSchedulerMixin: requests: dict[str, Request] waiting: RequestQueue running: list[Request] + _pending_input_timeout_outputs: dict[str, tuple[int, OmniEngineCoreOutput]] def _init_omni_connector_output_inbox(self) -> None: self._omni_connector_output_inbox: queue.SimpleQueue[OmniConnectorOutput] = queue.SimpleQueue() @@ -149,7 +150,7 @@ class OmniSchedulerMixin: else None ) self.input_coordinator: OmniSchedulingCoordinator | None = None - if self._native_data_plane: + if self._native_data_plane and getattr(model_config, "async_chunk", False): self.input_coordinator = OmniSchedulingCoordinator( scheduler_max_num_seqs=self.vllm_config.scheduler_config.max_num_seqs, stage_id=getattr(model_config, "stage_id", 0), @@ -251,6 +252,7 @@ class OmniSchedulerMixin: def _async_chunk_transport_enabled(self) -> bool: return getattr(self, "chunk_transfer_adapter", None) is not None or bool( getattr(self, "_native_data_plane", False) + and getattr(getattr(self, "input_coordinator", None), "_async_chunk", False) ) def _get_async_chunk_reserved_running_slots(self) -> int: @@ -453,7 +455,29 @@ class OmniSchedulerMixin: DEFAULT_INPUT_WAIT_TIMEOUT_S, sorted(present_ids), ) - self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR) + self._finish_input_timeout_requests(present_ids) + + def _finish_input_timeout_requests(self, request_ids: set[str]) -> None: + # Upstream finish_requests frees scheduler state, but does not emit + # EngineCoreOutput. Keep an explicit ERROR for both AR and generation + # callers, including a tick with no scheduled model work. + pending = getattr(self, "_pending_input_timeout_outputs", None) + if pending is None: + pending = self._pending_input_timeout_outputs = {} + for request_id in request_ids: + request = self.requests[request_id] + reason = f"Timed out waiting for connector input after {DEFAULT_INPUT_WAIT_TIMEOUT_S:g}s" + request.stop_reason = reason + pending[request_id] = ( + request.client_index, + OmniEngineCoreOutput( + request_id=request_id, + new_token_ids=[], + finish_reason=FinishReason.ERROR, + stop_reason=reason, + ), + ) + self.finish_requests(request_ids, RequestStatus.FINISHED_ERROR) def _log_failed_chunk_sends(self) -> None: """Surface chunks the sender gave up on (R1.2 of #4855). @@ -531,7 +555,7 @@ class OmniSchedulerMixin: DEFAULT_INPUT_WAIT_TIMEOUT_S, sorted(present_ids), ) - self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR) + self._finish_input_timeout_requests(present_ids) def _capture_omni_connector_output(self, model_runner_output: Any) -> None: """Stash the model runner's omni_connector_output for next schedule(). @@ -712,6 +736,13 @@ class OmniSchedulerMixin: synthesize_abort_outputs: bool, ) -> None: """Attach finished IDs while keeping AR's synthetic-abort policy explicit.""" + pending = getattr(self, "_pending_input_timeout_outputs", None) + if pending: + for client_index, error_output in pending.values(): + output = engine_core_outputs.setdefault(client_index, EngineCoreOutputs()) + if not any(item.request_id == error_output.request_id for item in output.outputs): + output.outputs.append(error_output) + pending.clear() finished_req_ids = self.finished_req_ids_dict if not finished_req_ids: return diff --git a/vllm_omni/core/sched/omni_scheduling_coordinator.py b/vllm_omni/core/sched/omni_scheduling_coordinator.py index 70b682ac0..dee3a9b43 100644 --- a/vllm_omni/core/sched/omni_scheduling_coordinator.py +++ b/vllm_omni/core/sched/omni_scheduling_coordinator.py @@ -25,11 +25,7 @@ def uses_native_mrv2_data_plane( *, use_v2_model_runner: bool, ) -> bool: - return bool( - use_v2_model_runner - and getattr(model_config, "async_chunk", False) - and getattr(model_config, "supports_native_mrv2_data_plane", False) - ) + return bool(use_v2_model_runner and getattr(model_config, "supports_native_mrv2_data_plane", False)) def uses_full_payload_input_coordinator(model_config: Any) -> bool: diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 4c5ca1e65..5cd42c35e 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -789,6 +789,18 @@ class OrchestratorBase: req_state = self.request_states.get(getattr(eco, "request_id", None)) if req_state is None: continue + if ( + getattr(eco, "finish_reason", None) == FinishReason.ERROR + and not getattr(eco, "is_segment_finished", False) + and not req_state.session_owned + ): + reason = getattr(eco, "stop_reason", None) + await self._handle_stage_error( + stage_id, + eco, + error=reason if isinstance(reason, str) and reason else "Stage request failed", + ) + continue if req_state.streaming.enabled: segment_finished = bool(getattr(eco, "is_segment_finished", False)) raw_mm = self._completion_multimodal_output(eco, None) @@ -1200,7 +1212,7 @@ class OrchestratorBase: await self._route_output(stage_id, replica_id, output, req_state, stage_metrics) - async def _handle_stage_error(self, stage_id: int, output: Any) -> None: + async def _handle_stage_error(self, stage_id: int, output: Any, *, error: str | None = None) -> None: """Emit a frontend-visible error and clean up request state.""" if self._cfg_tracker.is_companion(output.request_id): parent_id = self._cfg_tracker.get_parent_id(output.request_id) or output.request_id @@ -1210,7 +1222,7 @@ class OrchestratorBase: ErrorMessage( request_id=parent_id, stage_id=stage_id, - error=output.error, + error=error if error is not None else output.error, status_code=getattr(output, "error_status_code", None), error_type=getattr(output, "error_type", None), ) diff --git a/vllm_omni/worker_v2/omni_data_plane.py b/vllm_omni/worker_v2/omni_data_plane.py index c8e5e57cb..50a449b2d 100644 --- a/vllm_omni/worker_v2/omni_data_plane.py +++ b/vllm_omni/worker_v2/omni_data_plane.py @@ -15,9 +15,10 @@ from typing import Any, cast import torch from vllm.logger import init_logger -from vllm_omni.data_entry_keys import OmniPayload, unflatten_payload +from vllm_omni.data_entry_keys import OmniPayload, flatten_payload, unflatten_payload from vllm_omni.worker.omni_connector_model_runner_mixin import ( OmniConnectorModelRunnerMixin, + should_accumulate_full_payload_output, ) from vllm_omni.worker.payload_span import get_tensor_span, merge_tensor_spans from vllm_omni.worker_v2.delivery import ( @@ -87,7 +88,11 @@ class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin): delivery_timeout_s=self._connector_delivery_timeout(), shutdown_timeout_s=_DEFAULT_SHUTDOWN_TIMEOUT_S, ) - self._can_send = self._custom_process_func is not None + self._can_send = ( + self._custom_process_func is not None + if self._async_chunk + else should_accumulate_full_payload_output(model_config, self._custom_process_func) + ) self._start_output_worker(max_pending_batches=_NATIVE_OUTPUT_QUEUE_DEPTH) def _connector_delivery_timeout(self) -> float: @@ -294,6 +299,8 @@ class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin): for req_id in active_req_ids: self._native_outputs_in_flight.pop(req_id, None) self._native_terminal_pending.discard(req_id) + # Cancellation must not flush a partially accumulated utterance. + self._pending_full_payload_send.pop(req_id, None) return self.emit_chunks( req_ids=[], inter_stage_outputs=None, @@ -494,11 +501,23 @@ class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin): state.output_token_ids.extend(int(token_id) for token_id in sampled_by_req.get(req_id, [])) state.finished = req_id in terminal_req_ids payload = payload_by_req.get(req_id) - if isinstance(payload, dict): + if not self._async_chunk: + # Full-payload consumers execute once, after the producer has + # finished. Reuse V1's concat/replace contract for flat outputs + # and retain the final token history for codec-row cropping. + if isinstance(payload, dict): + self.accumulate_full_payload_output(req_id, flatten_payload(payload), state) + if not state.finished: + continue + accumulated = self._pending_full_payload_send.pop(req_id, None) + payload = self._materialize_full_payload_entry(accumulated)[0] if accumulated is not None else None + if self._async_chunk and isinstance(payload, dict): payload = unflatten_payload(payload) if payload is None and not state.finished: continue - include_token_history = bool(state.resumable or self._put_req_chunk.get(state.external_req_id, 0) == 0) + include_token_history = bool( + not self._async_chunk or state.resumable or self._put_req_chunk.get(state.external_req_id, 0) == 0 + ) entries.append( ( state.snapshot(include_token_history=include_token_history), diff --git a/vllm_omni/worker_v2/omni_model_runner.py b/vllm_omni/worker_v2/omni_model_runner.py index 58d1b1259..33545df36 100644 --- a/vllm_omni/worker_v2/omni_model_runner.py +++ b/vllm_omni/worker_v2/omni_model_runner.py @@ -208,6 +208,8 @@ class OmniGPUModelRunner(GPUModelRunner): continue plane.register_request(request_data) plane.register_receivers(list(getattr(scheduler_output, "pending_input_registrations", []))) + if not getattr(self.model_config, "async_chunk", False): + plane.recv_full_payload_inputs(scheduler_output) natural_terminal_req_ids = set(getattr(scheduler_output, "data_plane_terminal_req_ids", set())) aborted_req_ids = set(getattr(scheduler_output, "finished_req_ids", set())).difference(natural_terminal_req_ids) if natural_terminal_req_ids: