From 21dd4b33ac46e71479e839a9de3ed1b801046296 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Sat, 29 Aug 2026 13:05:24 -0700 Subject: [PATCH] [Bugfix][MiniCPM-o] Stabilize native duplex streaming and Stage-1 handoffs (#6529) Signed-off-by: Sy03 <1370724210@qq.com> Co-authored-by: Nathan Price Co-authored-by: Eric(Qun) <186020953+0xtoward@users.noreply.github.com> --- ...t_omni_ar_scheduler_aborted_queue_sweep.py | 89 ++++++ .../test_omni_ar_scheduler_stale_drain.py | 3 + .../sched/test_omni_ar_scheduler_streaming.py | 142 ++++++++- ...ni_ar_scheduler_unified_decode_deferral.py | 3 + .../test_omni_sched_deferred_free_fence.py | 12 + ..._omni_scheduler_streaming_input_counter.py | 168 ++++++++++ .../test_chunk_transfer_adapter.py | 52 ++- tests/engine/test_orchestrator.py | 3 +- .../test_orchestrator_segment_stage_keying.py | 130 ++++++++ .../test_orchestrator_stage_input_bridge.py | 23 +- .../openai/test_duplex_protocol.py | 9 + .../openai_api/test_duplex_handler.py | 296 ++++++++++++++++++ vllm_omni/core/sched/omni_ar_scheduler.py | 49 ++- .../core/sched/omni_generation_scheduler.py | 3 + vllm_omni/core/sched/omni_scheduler_mixin.py | 24 ++ .../chunk_transfer_adapter.py | 34 +- vllm_omni/engine/orchestrator.py | 94 ++++-- .../fullduplex/openai/protocol.py | 109 ++++++- .../fullduplex/openai/realtime_state.py | 3 + .../fullduplex/openai/runtime_bridge.py | 13 + .../experimental/fullduplex/openai/serving.py | 60 +++- 21 files changed, 1259 insertions(+), 60 deletions(-) create mode 100644 tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py create mode 100644 tests/core/sched/test_omni_scheduler_streaming_input_counter.py create mode 100644 tests/engine/test_orchestrator_segment_stage_keying.py diff --git a/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py b/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py new file mode 100644 index 000000000..d9fbc452b --- /dev/null +++ b/tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py @@ -0,0 +1,89 @@ +"""Tests for ``OmniARScheduler._drop_aborted_queued_requests``. + +``schedule()`` sweeps ``FINISHED_ABORTED`` requests out of the queues before +handing control to upstream, because upstream ``Scheduler.schedule()`` raises +``RuntimeError: Invalid request status: FINISHED_ABORTED`` for any request it +admits in a finished state -- which kills the stage's engine core, not just +the request. + +Regression for the sweep missing ``skipped_waiting``, the third queue +upstream admits from. An aborted duplex session parked there was re-selected +by ``_select_waiting_queue_for_scheduling`` on a later tick and crashed the +stage. +""" + +from __future__ import annotations + +import pytest +from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue +from vllm.v1.request import RequestStatus + +import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_POLICIES = [ + pytest.param(SchedulingPolicy.FCFS, id="fcfs"), + pytest.param(SchedulingPolicy.PRIORITY, id="priority"), +] + + +class _StubRequest: + """Minimal ``Request`` stub with the surface the sweep exercises. + + Carries the ``(priority, arrival_time)`` ordering ``PriorityRequestQueue`` + relies on so the same stub works under both scheduling policies. + """ + + def __init__(self, request_id: str, status: RequestStatus, arrival_time: float = 0.0) -> None: + self.request_id = request_id + self.status = status + self.priority = 0 + self.arrival_time = arrival_time + + def __lt__(self, other: _StubRequest) -> bool: + return (self.priority, self.arrival_time) < (other.priority, other.arrival_time) + + +def _make_scheduler(policy: SchedulingPolicy, *, waiting=(), skipped_waiting=(), running=()): + """Bare scheduler with just the surface the sweep reads.""" + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.waiting = create_request_queue(policy) + for req in waiting: + scheduler.waiting.add_request(req) + scheduler.skipped_waiting = create_request_queue(policy) + for req in skipped_waiting: + scheduler.skipped_waiting.add_request(req) + scheduler.running = list(running) + return scheduler + + +@pytest.mark.parametrize("policy", _POLICIES) +def test_schedule_sweeps_skipped_waiting_before_upstream_selection( + monkeypatch: pytest.MonkeyPatch, + policy: SchedulingPolicy, +) -> None: + """The public schedule path must remove aborted skipped requests first.""" + aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED) + scheduler = _make_scheduler(policy, skipped_waiting=[aborted]) + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = None + scheduler.num_waiting_for_streaming_input = 0 + scheduler._process_pending_omni_inputs = lambda model_mode: None + scheduler._resync_streaming_input_counter = lambda: None + scheduler._should_defer_waiting_admission = lambda: False + scheduler._restore_omni_wait_queues = lambda: None + scheduler._postprocess_omni_schedule_output = lambda *args, **kwargs: None + scheduler.get_finished_requests_needing_kv_transfer = lambda: {} + scheduler._wrap_omni_scheduler_output = lambda output, **kwargs: output + + sentinel = object() + + def upstream_schedule(_self, _throttle_prefills: bool = False): + assert list(scheduler.skipped_waiting) == [] + return sentinel + + monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", upstream_schedule) + + assert OmniARScheduler.schedule(scheduler) is sentinel diff --git a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py index 7c72602cf..5c1290fc4 100644 --- a/tests/core/sched/test_omni_ar_scheduler_stale_drain.py +++ b/tests/core/sched/test_omni_ar_scheduler_stale_drain.py @@ -142,13 +142,16 @@ def test_exact_drain_delivers_new_segment_frame() -> None: session.num_computed_tokens = 6 session.num_output_placeholders = 1 session.num_in_flight_tokens = 1 + session.async_tokens_to_discard = 1 _replace_streaming_session(session) assert session.num_stale_output_tokens == 1 + assert session.async_tokens_to_discard == 1 sched = _make_drain_sched(session) assert _run_step(sched, session, num_scheduled=1, token=42) is False # late frame dropped assert session.num_stale_output_tokens == 0 + assert session.async_tokens_to_discard == 0 assert _run_step(sched, session, num_scheduled=1, token=43) is True # new segment survives diff --git a/tests/core/sched/test_omni_ar_scheduler_streaming.py b/tests/core/sched/test_omni_ar_scheduler_streaming.py index 438ebc72a..a1aa7351f 100644 --- a/tests/core/sched/test_omni_ar_scheduler_streaming.py +++ b/tests/core/sched/test_omni_ar_scheduler_streaming.py @@ -18,7 +18,7 @@ from vllm.sampling_params import SamplingParams from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus, StreamingUpdate -from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler, OmniARScheduler # isort: on @@ -28,7 +28,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] def _make_scheduler(*, stage_id: int = 0, session_mode: str = "turn") -> OmniARScheduler: sched = OmniARScheduler.__new__(OmniARScheduler) sched._new_prompt_len_snapshot = {} - sched.vllm_config = SimpleNamespace(model_config=SimpleNamespace(stage_id=stage_id, session_mode=session_mode)) + sched.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(stage_id=stage_id, session_mode=session_mode), + ) sched.num_waiting_for_streaming_input = 0 sched.log_stats = False sched.chunk_transfer_adapter = None @@ -185,6 +187,7 @@ def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> Non def replace_with_next_segment(request: Request) -> bool: request.num_computed_tokens = 0 + request._omni_segment_generation = 1 request.status = RequestStatus.WAITING return False @@ -201,6 +204,7 @@ def test_resumable_segment_boundary_keeps_pre_transition_send_watermark() -> Non True, new_token_ids=[42], confirmed_num_computed_tokens=26, + segment_generation=0, ) @@ -286,6 +290,140 @@ def test_queued_streaming_update_on_async_stop_fences_in_flight_once() -> None: assert session.num_stale_output_tokens == session.num_in_flight_tokens == 1 +def test_stale_async_frame_is_dropped_before_output_processing() -> None: + session = _make_request() + session.status = RequestStatus.RUNNING + session.num_in_flight_tokens = 2 + session.num_computed_tokens = session.num_tokens + 1 + session.num_output_placeholders = 1 + session.async_tokens_to_discard = 1 + session.sampling_params = SimpleNamespace(num_logprobs=1) + num_computed_tokens = session.num_computed_tokens + num_output_placeholders = session.num_output_placeholders + + sched = MagicMock() + sched.requests = {session.request_id: session} + sched.perf_metrics = None + sched.structured_output_manager.should_advance.return_value = False + + def discard_stale_output(request: Request, token_ids: list[int]) -> tuple[list[int], bool]: + request.async_tokens_to_discard = 0 + return token_ids, False + + sched._update_request_with_output.side_effect = discard_stale_output + sched._process_kv_transfer_trigger.return_value = False + sched.chunk_transfer_adapter = MagicMock() + sched.running = [session] + sched.waiting_for_transfer_free = set() + sched.transfer_triggered_requests = set() + sched.active_kv_transfers = set() + sched.pending_stop_after_extraction = set() + sched.connector = None + sched.kv_cache_manager.take_events.return_value = None + sched.kv_cache_manager.estimate_cached_tokens.return_value = 0 + sched.finished_req_ids_dict = {} + sched.make_stats.return_value = None + + scheduler_output = MagicMock(spec=SchedulerOutput) + scheduler_output.num_scheduled_tokens = {session.request_id: 1} + scheduler_output.scheduled_spec_decode_tokens = {} + scheduler_output.num_invalid_spec_tokens = 0 + + model_runner_output = MagicMock(spec=ModelRunnerOutput) + model_runner_output.sampled_token_ids = [[42]] + model_runner_output.logprobs = None + model_runner_output.prompt_logprobs_dict = {} + model_runner_output.pooler_output = None + model_runner_output.multimodal_outputs = None + model_runner_output.inter_stage_outputs = [{"hidden": object()}] + model_runner_output.num_nans_in_logits = None + model_runner_output.kv_connector_output = None + model_runner_output.cudagraph_stats = None + model_runner_output.req_id_to_index = {session.request_id: 0} + model_runner_output.routed_experts = None + + OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output) + + assert session.async_tokens_to_discard == 0 + assert session.status == RequestStatus.RUNNING + assert session.num_computed_tokens == num_computed_tokens + assert session.num_output_placeholders == num_output_placeholders + sched.chunk_transfer_adapter.save_async.assert_not_called() + + session.sampling_params = SamplingParams(max_tokens=8) + next_payload = {"hidden": object()} + sched._update_request_with_output.side_effect = None + sched._update_request_with_output.return_value = ([43], False) + model_runner_output.sampled_token_ids = [[43]] + model_runner_output.inter_stage_outputs = [next_payload] + + OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output) + + sched.chunk_transfer_adapter.save_async.assert_called_once_with( + next_payload, + session, + False, + new_token_ids=[43], + confirmed_num_computed_tokens=None, + ) + + +def test_legacy_stale_async_marker_bypasses_real_placeholder_accounting() -> None: + request = _make_request() + request.status = RequestStatus.RUNNING + request.num_in_flight_tokens = 1 + request.num_output_placeholders = 0 + request.async_tokens_to_discard = 1 + request.num_stale_output_tokens = 0 + + sched = OmniARAsyncScheduler.__new__(OmniARAsyncScheduler) + sched.requests = {request.request_id: request} + sched.perf_metrics = None + sched.chunk_transfer_adapter = None + sched.connector = None + sched.ec_connector = None + sched.waiting_for_transfer_free = set() + sched.transfer_triggered_requests = set() + sched.active_kv_transfers = set() + sched.pending_stop_after_extraction = set() + sched.finished_req_ids_dict = {} + sched._new_prompt_len_snapshot = {} + sched.kv_cache_manager = MagicMock() + sched.kv_cache_manager.take_events.return_value = None + sched._remove_stopped_requests_from_queues = MagicMock() + sched._handle_failed_kv_load_outputs = MagicMock(return_value=[]) + sched._cleanup_kv_tracking = MagicMock() + sched._aggregate_kv_connector_stats = MagicMock(return_value=None) + sched._publish_kv_cache_events = MagicMock() + sched._attach_finished_request_sets = MagicMock() + sched._attach_scheduler_stats = MagicMock() + sched._capture_omni_connector_output = MagicMock() + + scheduler_output = MagicMock(spec=SchedulerOutput) + scheduler_output.num_scheduled_tokens = {request.request_id: 1} + scheduler_output.scheduled_spec_decode_tokens = {} + scheduler_output.num_invalid_spec_tokens = 0 + + model_runner_output = MagicMock(spec=ModelRunnerOutput) + model_runner_output.sampled_token_ids = [[42]] + model_runner_output.logprobs = None + model_runner_output.prompt_logprobs_dict = {} + model_runner_output.pooler_output = None + model_runner_output.multimodal_outputs = None + model_runner_output.inter_stage_outputs = None + model_runner_output.num_nans_in_logits = None + model_runner_output.kv_connector_output = None + model_runner_output.cudagraph_stats = None + model_runner_output.req_id_to_index = {request.request_id: 0} + model_runner_output.routed_experts = None + + OmniARScheduler.update_from_output(sched, scheduler_output, model_runner_output) + + assert request.async_tokens_to_discard == 0 + assert request.num_output_placeholders == 0 + assert request._output_token_ids == [] + + def test_stage0_streaming_update_discards_outstanding_async_placeholder_token() -> None: sched = _make_scheduler(stage_id=0) session = _make_request() diff --git a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py index 27087d253..0e0ed04ad 100644 --- a/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py +++ b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py @@ -76,6 +76,9 @@ def _make_scheduler( else: hf_config.voxcpm2_runtime_config = runtime_config sched.vllm_config = SimpleNamespace(model_config=model_config) + # Upstream ``Scheduler.__init__`` always creates ``skipped_waiting``; the + # abort sweep at the top of ``schedule()`` reads it. + sched.skipped_waiting = _MockQueue() return sched diff --git a/tests/core/sched/test_omni_sched_deferred_free_fence.py b/tests/core/sched/test_omni_sched_deferred_free_fence.py index 376c64cb3..36b084f7c 100644 --- a/tests/core/sched/test_omni_sched_deferred_free_fence.py +++ b/tests/core/sched/test_omni_sched_deferred_free_fence.py @@ -166,3 +166,15 @@ def test_fast_path_schedule_defer_disabled_leaves_fence() -> None: OmniGenerationScheduler.schedule(sched) assert sched.sched_step_seq == 0 + + +def test_generation_schedule_sweeps_and_resyncs_around_pending_inputs() -> None: + sched, request = _make_fast_path_sched(defer_block_free=False) + calls: list[str] = [] + sched._drop_aborted_queued_requests.side_effect = lambda: calls.append("sweep") + sched._process_pending_omni_inputs.side_effect = lambda model_mode: calls.append(model_mode) + sched._resync_streaming_input_counter.side_effect = lambda: calls.append("resync") + + OmniGenerationScheduler.schedule(sched) + + assert calls == ["sweep", "generation", "sweep", "resync"] diff --git a/tests/core/sched/test_omni_scheduler_streaming_input_counter.py b/tests/core/sched/test_omni_scheduler_streaming_input_counter.py new file mode 100644 index 000000000..6f44aadac --- /dev/null +++ b/tests/core/sched/test_omni_scheduler_streaming_input_counter.py @@ -0,0 +1,168 @@ +"""Tests for the ``num_waiting_for_streaming_input`` resync placed in +``OmniARScheduler.finish_requests`` / ``OmniGenerationScheduler.finish_requests`` +and in ``OmniARScheduler.schedule``. + +Upstream keeps that counter incrementally (``+1`` when a resumable request +parks in a waiting queue as ``WAITING_FOR_STREAMING_REQ``, ``-1`` when the +next update arrives or the request is finished) and +``Scheduler.get_num_unfinished_requests`` *subtracts* it from the waiting +queues. Omni rewrites ``request.status`` outside both hooks -- the +chunk-transfer adapter's park/restore and +``_realign_request_status_to_queues`` -- so a counted request can leave +``WAITING_FOR_STREAMING_REQ`` without the matching decrement and inflate the +counter for the lifetime of the process. + +An inflated counter does not degrade throughput, it hangs the stage: +``EngineCore.has_work()`` reads false while a live request sits in +``waiting``, so the engine blocks in ``input_queue.get()`` and never calls +``schedule()`` again. Regression for the Qwen3-Omni duplex symptom where the +server answered the first conversation after a boot and was silent from the +second on while still reporting healthy. +""" + +from __future__ import annotations + +import pytest +from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue +from vllm.v1.request import RequestStatus + +import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod +import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod +from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler +from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _StubRequest: + def __init__(self, request_id: str, status: RequestStatus) -> None: + self.request_id = request_id + self.status = status + + def is_finished(self) -> bool: + return RequestStatus.is_finished(self.status) + + +_SCHEDULER_PARAMS = [ + pytest.param(OmniARScheduler, ar_sched_mod, id="ar"), + pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"), +] + + +def _make_scheduler(scheduler_cls, *, requests, running, waiting, counter, skipped=None): + scheduler = scheduler_cls.__new__(scheduler_cls) + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = None + scheduler.requests = requests + scheduler.running = running + scheduler.waiting = waiting + scheduler.skipped_waiting = skipped if skipped is not None else [] + scheduler.num_waiting_for_streaming_input = counter + return scheduler + + +@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS) +def test_finish_requests_clears_counter_left_by_a_stomped_status( + monkeypatch: pytest.MonkeyPatch, + scheduler_cls, + scheduler_mod, +) -> None: + """The abort case that wedged the talker. + + The request was counted while parked for streaming input, then had its + status rewritten to ``RUNNING`` by omni's queue surgery, so upstream's + ``finish_requests`` took the running branch and never decremented. With + no parked requests left, the counter must land at 0. + """ + aborted = _StubRequest("req-aborted", RequestStatus.RUNNING) + scheduler = _make_scheduler( + scheduler_cls, + requests={}, # upstream already dropped it + running=[], + waiting=[], + counter=1, + ) + + monkeypatch.setattr( + scheduler_mod.VLLMScheduler, + "finish_requests", + lambda self, request_ids, finished_status: [aborted], + ) + + scheduler_cls.finish_requests(scheduler, [aborted.request_id], RequestStatus.FINISHED_ABORTED) + + assert scheduler.num_waiting_for_streaming_input == 0 + + +def test_resync_counts_parked_requests_in_skipped_waiting() -> None: + """``get_num_unfinished_requests`` sums ``waiting`` and ``skipped_waiting``, + so the resync has to look at both or it under-counts and reintroduces the + spin the counter exists to prevent.""" + parked = _StubRequest("req-skipped", RequestStatus.WAITING_FOR_STREAMING_REQ) + scheduler = _make_scheduler( + OmniARScheduler, + requests={parked.request_id: parked}, + running=[], + waiting=[], + counter=0, + skipped=[parked], + ) + + scheduler._resync_streaming_input_counter() + + assert scheduler.num_waiting_for_streaming_input == 1 + + +def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.MonkeyPatch) -> None: + """Pins placement in ``schedule``: the chunk-transfer adapter rewrites + status during ``process_pending_chunks``, so the resync has to run after + that and before upstream schedules, while ``self.waiting`` is still the + real queue.""" + parked = _StubRequest("req-resumed", RequestStatus.RUNNING) + scheduler = OmniARScheduler.__new__(OmniARScheduler) + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = None + scheduler.requests = {parked.request_id: parked} + scheduler.running = [parked] + scheduler.waiting = [] + scheduler.skipped_waiting = [] + scheduler.num_waiting_for_streaming_input = 1 # leaked by the status rewrite + scheduler.policy = "fcfs" + + seen: dict[str, int] = {} + + monkeypatch.setattr( + ar_sched_mod.OmniSchedulerMixin, "_consume_pending_connector_output", lambda self, model_mode: None + ) + monkeypatch.setattr(ar_sched_mod.OmniSchedulerMixin, "_process_pending_input_timeouts", lambda self: None) + monkeypatch.setattr(OmniARScheduler, "_should_defer_waiting_admission", lambda self: False) + + def fake_schedule(self, throttle_prefills=False): + seen["counter"] = self.num_waiting_for_streaming_input + raise _StopScheduleError + + class _StopScheduleError(Exception): + pass + + monkeypatch.setattr(ar_sched_mod.VLLMScheduler, "schedule", fake_schedule) + + with pytest.raises(_StopScheduleError): + OmniARScheduler.schedule(scheduler) + + assert seen["counter"] == 0 + + +def test_generation_scheduler_sweep_removes_aborted_requests_from_all_queues() -> None: + aborted = _StubRequest("req-generation-aborted", RequestStatus.FINISHED_ABORTED) + scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler) + scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.waiting.add_request(aborted) + scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS) + scheduler.skipped_waiting.add_request(aborted) + scheduler.running = [aborted] + + scheduler._drop_aborted_queued_requests() + + assert list(scheduler.waiting) == [] + assert list(scheduler.skipped_waiting) == [] + assert scheduler.running == [] diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 59780d5d1..def66971f 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -787,6 +787,7 @@ def test_segment_boundary_starts_new_send_watermark_before_background_flush(buil adapter, _ = build_adapter(stage_id=1) request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream") request.resumable = True + request._omni_segment_generation = 1 request.num_computed_tokens = 0 adapter.requests_num_chunks_sent["ext-stream"] = 26 @@ -795,13 +796,16 @@ def test_segment_boundary_starts_new_send_watermark_before_background_flush(buil request=request, is_segment_finished=True, confirmed_num_computed_tokens=26, + segment_generation=0, ) assert len(adapter._pending_save_reqs) == 1 boundary_task = adapter._pending_save_reqs.popleft() + assert boundary_task["segment_generation"] == 0 assert "ext-stream" not in adapter.requests_num_chunks_sent request.num_computed_tokens = 3 + request._omni_segment_generation = 1 adapter.save_async( multimodal_output=None, request=request, @@ -854,6 +858,7 @@ def test_background_send_uses_enqueued_request_snapshot(build_adapter): request.additional_information["meta"]["segment"] = "next" request.prompt_token_ids.append(3) request.append_output_token_ids([8]) + request._omni_segment_generation = 1 adapter.save_async( multimodal_output=None, request=request, @@ -1102,7 +1107,7 @@ def test_personaplex_sender_cleanup_drops_delayed_frame_state(build_adapter): assert second.codes is None -def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapter): +def test_save_async_skips_stale_resumable_chunk_within_segment(build_adapter): adapter, _ = build_adapter(stage_id=1) request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream") request.resumable = True @@ -1121,6 +1126,51 @@ def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapt assert adapter.requests_num_chunks_sent["ext-stream"] == 0 +def test_save_async_drops_late_previous_segment_after_boundary_reset(build_adapter): + adapter, _ = build_adapter(stage_id=1) + request = _req("req-stream", RequestStatus.WAITING, external_req_id="ext-stream") + request.resumable = True + request._omni_segment_generation = 0 + adapter.requests_num_chunks_sent["ext-stream"] = 111 + + adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True) + + late = _req("req-stream-late", RequestStatus.WAITING, external_req_id="ext-stream") + late.resumable = True + late.num_computed_tokens = 111 + late._omni_segment_generation = 0 + adapter.save_async(multimodal_output=None, request=late, is_segment_finished=False) + + next_segment = _req("req-stream-next", RequestStatus.WAITING, external_req_id="ext-stream") + next_segment.resumable = True + next_segment.num_computed_tokens = 3 + next_segment._omni_segment_generation = 1 + adapter.save_async(multimodal_output=None, request=next_segment, is_segment_finished=False) + + assert len(adapter._pending_save_reqs) == 2 + assert adapter.requests_num_chunks_sent["ext-stream"] == 3 + assert adapter._pending_save_reqs[-1]["request"] is next_segment + + +def test_send_single_request_preserves_queued_preboundary_chunk(build_adapter): + adapter, connector = build_adapter(stage_id=1) + adapter.custom_process_next_stage_input_func = lambda **kwargs: OmniPayloadStruct(meta=MetaStruct()) + request = _req("req-queued-old", RequestStatus.WAITING, external_req_id="ext-queued-old") + request.resumable = True + request._omni_segment_generation = 0 + + adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False) + adapter.save_async(multimodal_output=None, request=request, is_segment_finished=True) + + old_task = adapter._pending_save_reqs.popleft() + adapter._send_single_request(old_task) + connector.put.assert_called_once() + + boundary_task = adapter._pending_save_reqs.popleft() + adapter._send_single_request(boundary_task) + assert connector.put.call_count == 2 + + def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch): adapter, _ = build_adapter(stage_id=1) request = _req("req-finished", RequestStatus.FINISHED_STOPPED, external_req_id="ext-finished") diff --git a/tests/engine/test_orchestrator.py b/tests/engine/test_orchestrator.py index 20f12e8bf..9dcf066cd 100644 --- a/tests/engine/test_orchestrator.py +++ b/tests/engine/test_orchestrator.py @@ -35,6 +35,7 @@ from vllm_omni.engine.messages import ( from vllm_omni.engine.orchestrator import ( Orchestrator, OrchestratorRequestState, + StreamingSegmentState, _build_terminal_empty_output, ) from vllm_omni.engine.stage_pool import StagePool @@ -2259,7 +2260,7 @@ async def test_resumable_segment_boundary_builds_stage_metrics() -> None: final_stage_id=0, ) req_state.streaming.enabled = True - req_state.streaming.segment_finished = True + req_state.streaming.segments[0] = StreamingSegmentState(finished=True) req_state.stage_submit_ts[0] = time.time() orchestrator.request_states = {"req-stream": req_state} orchestrator.stage_pools = [pool] diff --git a/tests/engine/test_orchestrator_segment_stage_keying.py b/tests/engine/test_orchestrator_segment_stage_keying.py new file mode 100644 index 000000000..02bb9a338 --- /dev/null +++ b/tests/engine/test_orchestrator_segment_stage_keying.py @@ -0,0 +1,130 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Streaming segment-boundary state must be keyed by stage. + +``Orchestrator._orchestration_loop`` polls every stage into the same +``OrchestratorRequestState``. While the segment-boundary fields lived in a single +flat slot per request, whichever stage polled most recently overwrote the others, +and the per-stage consumers then read back a different stage's boundary. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import pytest +from vllm.sampling_params import SamplingParams + +from vllm_omni.engine.orchestrator import ( + Orchestrator, + OrchestratorRequestState, + StreamingSegmentState, +) +from vllm_omni.experimental.fullduplex.engine.contracts import DuplexRequestIdentity +from vllm_omni.experimental.fullduplex.engine.messages import DuplexFence + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +THINKER_STAGE = 0 +TALKER_STAGE = 1 + + +@dataclass +class _StageMetrics: + pipeline_timings: dict[str, float] = field(default_factory=dict) + + +@dataclass +class _StageOutput: + """Stand-in for a processed stage output, with the fields the router reads.""" + + request_id: str + finished: bool = False + error: object | None = None + + +class _RecordingPool: + """Stage pool that records the requests it was asked to build metrics for.""" + + def __init__(self) -> None: + self.built_for: list[str] = [] + + def build_stage_metrics(self, outputs: list[_StageOutput], **_kwargs: object) -> _StageMetrics: + self.built_for.extend(output.request_id for output in outputs) + return _StageMetrics() + + +def _duplex_req_state(request_id: str = "req-duplex") -> OrchestratorRequestState: + req_state = OrchestratorRequestState( + request_id=request_id, + sampling_params_list=[SamplingParams(max_tokens=1), SamplingParams(max_tokens=1)], + final_stage_id=TALKER_STAGE, + ) + req_state.streaming.enabled = True + req_state.duplex_identity = DuplexRequestIdentity( + session_id="session-1", + fence=DuplexFence(session_id="session-1"), + ) + return req_state + + +def test_duplex_output_context_reads_only_its_own_stage_segment() -> None: + req_state = _duplex_req_state() + req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState( + finished=True, + token_ids=[11, 22], + output_metadata={"stage": "thinker"}, + ) + # The talker polls after the thinker and is still mid-segment. + req_state.streaming.segments[TALKER_STAGE] = StreamingSegmentState(finished=False) + + thinker = Orchestrator._duplex_output_context(req_state, stage_id=THINKER_STAGE) + talker = Orchestrator._duplex_output_context(req_state, stage_id=TALKER_STAGE) + + assert thinker is not None + assert talker is not None + + # These are the assertions that fail on the flat field: the talker is written + # second, so the thinker loses its boundary to whichever stage polled last. + assert thinker.segment_finished is True + assert thinker.segment_token_ids == (11, 22) + assert thinker.segment_output_metadata == {"stage": "thinker"} + + assert talker.segment_finished is False + assert talker.segment_token_ids == () + assert talker.segment_output_metadata == {} + + +@pytest.mark.asyncio +async def test_segment_boundary_on_one_stage_does_not_build_metrics_on_another() -> None: + req_state = _duplex_req_state() + req_state.streaming.segments[THINKER_STAGE] = StreamingSegmentState(finished=True) + + talker_pool = _RecordingPool() + orchestrator = object.__new__(Orchestrator) + orchestrator.request_states = {req_state.request_id: req_state} + orchestrator.stage_pools = [_RecordingPool(), talker_pool] + + routed: list[_StageMetrics | None] = [] + + async def record_route( + _stage: int, + _replica: int, + _output: _StageOutput, + _state: OrchestratorRequestState, + metrics: _StageMetrics | None, + ) -> None: + routed.append(metrics) + + orchestrator._route_output = record_route + + # The talker's own output is neither finished nor at a segment boundary, so + # it must not be treated as one just because the thinker reached one. + await orchestrator._handle_processed_outputs( + TALKER_STAGE, + 0, + [_StageOutput(request_id=req_state.request_id)], + ) + + assert talker_pool.built_for == [] + assert routed == [None] diff --git a/tests/engine/test_orchestrator_stage_input_bridge.py b/tests/engine/test_orchestrator_stage_input_bridge.py index 25637913b..52cee4657 100644 --- a/tests/engine/test_orchestrator_stage_input_bridge.py +++ b/tests/engine/test_orchestrator_stage_input_bridge.py @@ -17,6 +17,7 @@ from vllm.sampling_params import SamplingParams from vllm_omni.engine.orchestrator import ( Orchestrator, OrchestratorRequestState, + StreamingSegmentState, _OrchestratorDuplexStagePort, ) from vllm_omni.engine.stage_pool import StagePool @@ -160,6 +161,26 @@ def _duplex_stage_port_submission(): return port, stage_pools, request_states, prewarm, submission +def test_duplex_bridge_state_catches_up_after_serving_turn_boundary() -> None: + port, _, request_states, _, submission = _duplex_stage_port_submission() + request_state = request_states[submission.context.request_id] + bridge_state = request_state.streaming.bridge_states["duplex"] + bridge_state["model_turn_id"] = 0 + + next_context = DuplexStageRequestContext( + request_id=submission.context.request_id, + session_id=submission.context.session_id, + fence=DuplexFence(submission.context.session_id, turn_id=1), + stage_id=submission.context.stage_id, + final_stage_id=submission.context.final_stage_id, + config_generation=submission.context.config_generation, + sampling_params=submission.context.sampling_params, + ) + port.ensure_request(next_context) + + assert bridge_state["model_turn_id"] == 1 + + def _request_output(request_id: str) -> RequestOutput: completion = CompletionOutput( index=0, @@ -358,7 +379,7 @@ async def test_streaming_segment_does_not_complete_final_output_stage() -> None: final_output_stage_ids={0}, ) req_state.streaming.enabled = True - req_state.streaming.segment_finished = True + req_state.streaming.segments[0] = StreamingSegmentState(finished=True) output = SimpleNamespace( request_id=req_state.request_id, finished=True, diff --git a/tests/entrypoints/openai/test_duplex_protocol.py b/tests/entrypoints/openai/test_duplex_protocol.py index 89ae03f6f..9fd61095b 100644 --- a/tests/entrypoints/openai/test_duplex_protocol.py +++ b/tests/entrypoints/openai/test_duplex_protocol.py @@ -462,6 +462,15 @@ def test_minicpmo_native_capabilities_separate_model_state_from_core_kv_lease(): assert caps["stage_handoff_transport"] == "scheduler_data_plane" +def test_duplex_overlap_policy_auto_falls_back_to_listen_only(): + from vllm_omni.experimental.fullduplex.openai.protocol import ( + DuplexOverlapPolicy, + DuplexSessionConfig, + ) + + assert DuplexSessionConfig._normalize_overlap_policy("auto") == DuplexOverlapPolicy.LISTEN_ONLY.value + + def test_minicpmo_native_capabilities_do_not_overclaim_single_session_deployment(): caps = DuplexCapabilities.minicpmo45_native(max_sessions=1).as_dict() diff --git a/tests/entrypoints/openai_api/test_duplex_handler.py b/tests/entrypoints/openai_api/test_duplex_handler.py index 475d58bd5..b42ace688 100644 --- a/tests/entrypoints/openai_api/test_duplex_handler.py +++ b/tests/entrypoints/openai_api/test_duplex_handler.py @@ -1811,6 +1811,27 @@ def test_auto_response_playback_overlap_keeps_model_owned_listen_speak_decision( assert handler._should_force_listen_for_auto_response_overlap(session, {}, speech_payload) is False +def test_legacy_auto_overlap_policy_falls_back_to_listen_only(): + handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) + session = DuplexSession( + session_id="sid-auto-policy-fallback", + config=DuplexSessionConfig( + overlap_policy=DuplexSessionConfig._normalize_overlap_policy("auto"), + ), + ) + session.capabilities = DuplexCapabilities.minicpmo45_native() + + decision = handler._overlap_decision( + session, + {"duration_ms": session.config.overlap_barge_in_ms, "is_speech": True}, + _native_audio_payload(is_speech=True), + ) + + assert session.config.overlap_policy == DuplexOverlapPolicy.LISTEN_ONLY.value + assert decision["action"] == "listen" + assert decision["reason"] == "policy_listen_only" + + def test_auto_response_overlap_silence_advances_model_unit_and_preserves_realtime_input(): handler, session = _auto_response_context( "sid-auto-silent-wire-buffer", @@ -3250,6 +3271,7 @@ async def test_minicpmo_auto_response_continuation_stops_at_large_safety_boundar assert ws.sent_types()[-2:] == ["response.listen", "response.done"] assert session.active_response_id is None assert session.active_request_id == request_id + assert session.turn_id == 1 @pytest.mark.asyncio @@ -3293,9 +3315,65 @@ async def test_minicpmo_auto_response_boundary_listen_closes_response(): assert ws.sent_types() == ["response.listen", "response.done"] assert session.active_response_id is None assert session.active_request_id == request_id + assert session.turn_id == 1 assert native.continuation_owner_id is None +@pytest.mark.asyncio +async def test_minicpmo_auto_response_boundary_allows_next_model_turn_response(): + request_id = "duplex-sid-boundary-next-turn-e0-stage0" + handler = OmniDuplexSessionHandler(chat_service=FakeChatService(FakeEngineClient())) + session = DuplexSession( + session_id="sid-boundary-next-turn", + config=DuplexSessionConfig(extra_body={"auto_response": True}), + ) + session.capabilities = DuplexCapabilities.minicpmo45_native() + session.begin_response(turn_id=0) + session.bind_request(request_id) + native = handler._minicpmo_session_state(session) + native.continuation_owner_id = f"response:{session.active_response_id}" + native.continuation_units = handler._NATIVE_AUTO_RESPONSE_MAX_CONTINUATION_UNITS + ws = TimedWebSocket() + + await handler._maybe_continue_native_response( + ws.send_json, + session=session, + expected_epoch=session.epoch, + ) + + assert session.active_response_id is None + assert session.turn_id == 1 + + close_reason, emitted = await handler._send_one_native_duplex_event( + ws.send_json, + { + "supported": True, + "stage_role": "tts", + "is_listen": False, + "data_plane_request_id": request_id, + "text": "next response", + "audio_data": "next-audio", + "audio_format": "pcm16", + "audio_duration_ms": 100, + "end_of_turn": False, + "model_turn_id": 1, + }, + session=session, + expected_epoch=session.epoch, + ) + + assert close_reason is None + assert emitted is True + assert ws.sent_types() == [ + "response.listen", + "response.done", + "response.created", + "response.speak", + "response.output_audio.delta", + ] + assert session.active_response_turn_id == 1 + + @pytest.mark.asyncio async def test_minicpmo_auto_response_boundary_does_not_close_new_epoch_response(): request_id = "duplex-sid-boundary-race-e0-stage0" @@ -4269,6 +4347,224 @@ async def test_playback_ack_response_id_selects_matching_pending_history_item(): assert second_item_id not in session.pending_history_item_ids +@pytest.mark.asyncio +async def test_playback_ack_recovers_missing_response_history_registration(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession( + session_id="sid-playback-missing-registration", + config=DuplexSessionConfig(playback_commit_policy="ack_only"), + ) + response_id = session.begin_response() + session.append_assistant_text("response text") + session.mark_audio_sent(1000, text_chars=len("response text")) + session.end_response(commit_text=True, preserve_request=True) + assert not session.history_item_ids + assert f"item_{response_id}" in session.pending_history_item_ids + + ws = TimedWebSocket() + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": response_id, + "item_id": f"item_{response_id}", + "played_ms": 1000, + "committed_ms": 1000, + }, + ws.send_json, + ) + + ack = next(message for message in ws.sent if message.get("type") == "playback.acknowledged") + assert ack["history_committed"] is True + assert f"item_{response_id}" in session.history_item_ids + assert session.history[-1] == {"role": "assistant", "content": "response text"} + + +@pytest.mark.asyncio +async def test_playback_ack_rejects_response_after_followup_user_commit(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession( + session_id="sid-playback-after-followup-user", + config=DuplexSessionConfig(playback_commit_policy="ack_only"), + ) + session.append_text("user A") + session.commit_user_input() + response_id = session.begin_response() + session.append_assistant_text("assistant A") + session.mark_audio_sent(1000, text_chars=len("assistant A")) + session.end_response(commit_text=False, preserve_request=True) + session.append_text("user B") + session.commit_user_input() + ws = TimedWebSocket() + + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": response_id, + "item_id": f"item_{response_id}", + "played_ms": 1000, + "committed_ms": 1000, + }, + ws.send_json, + ) + + assert ws.sent_types() == ["error"] + assert ws.sent[0]["code"] == "playback_ack_too_late" + assert session.history == ( + {"role": "user", "content": "user A"}, + {"role": "user", "content": "user B"}, + ) + + +@pytest.mark.asyncio +async def test_cancel_does_not_append_old_assistant_after_followup_user(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession( + session_id="sid-cancel-after-followup-user", + config=DuplexSessionConfig(playback_commit_policy="ack_only"), + ) + session.append_text("user A") + session.commit_user_input() + session.begin_response() + session.append_assistant_text("assistant A") + session.mark_audio_sent(1000, text_chars=len("assistant A")) + session.acknowledge_playback(500, 500) + session.append_text("user B") + session.commit_user_input() + ws = TimedWebSocket() + + cancelled = await handler._cancel_active_response(session, None, ws.send_json, reason="new_response") + + assert cancelled is True + assert session.history == ( + {"role": "user", "content": "user A"}, + {"role": "user", "content": "user B"}, + ) + + +@pytest.mark.asyncio +async def test_playback_ack_rejects_item_from_another_response(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-playback-mismatch", config=DuplexSessionConfig()) + first_response_id = session.begin_response() + session.append_assistant_text("first response") + session.mark_audio_sent(1000) + session.end_response(commit_text=False, preserve_request=True) + session.register_history_item(f"item_{first_response_id}", None) + second_response_id = session.begin_response() + session.mark_audio_sent(1000) + ws = TimedWebSocket() + + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": second_response_id, + "item_id": f"item_{first_response_id}", + "played_ms": 500, + "committed_ms": 500, + }, + ws.send_json, + ) + + assert ws.sent_types() == ["error"] + assert ws.sent[0]["code"] == "playback_item_mismatch" + assert session.active_response_id == second_response_id + assert session.playback.as_dict() == { + "generated_ms": 1000, + "sent_ms": 1000, + "played_ms": 0, + "committed_ms": 0, + } + + +@pytest.mark.asyncio +async def test_playback_ack_rejects_user_history_item(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-playback-user-item", config=DuplexSessionConfig()) + user_message = {"role": "user", "content": "user input"} + session.append_history_message(user_message) + session.register_history_item("item_user", user_message) + ws = TimedWebSocket() + + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": "user", + "item_id": "item_user", + "played_ms": 500, + "committed_ms": 500, + }, + ws.send_json, + ) + + assert ws.sent_types() == ["error"] + assert ws.sent[0]["code"] == "playback_item_not_found" + assert session.history == ({"role": "user", "content": "user input"},) + + +@pytest.mark.asyncio +async def test_late_playback_ack_recovers_response_local_history_after_next_response_begins(): + handler = OmniDuplexSessionHandler( + chat_service=FakeChatService(FakeEngineClient()), + config_timeout_s=0.1, + idle_timeout_s=1, + ) + session = DuplexSession(session_id="sid-playback-response-local", config=DuplexSessionConfig()) + first_response_id = session.begin_response() + session.append_assistant_text("first response") + session.mark_audio_sent(1000, text_chars=len("first response")) + session.end_response(commit_text=False, preserve_request=True) + second_response_id = session.begin_response() + session.mark_audio_sent(2000, text_chars=len("second response")) + ws = TimedWebSocket() + + await handler._handle_playback_ack( + session, + { + "type": "playback.ack", + "response_id": first_response_id, + "item_id": f"item_{first_response_id}", + "played_ms": 1000, + "committed_ms": 1000, + }, + ws.send_json, + ) + + ack = next(message for message in ws.sent if message.get("type") == "playback.acknowledged") + assert ack["history_committed"] is True + assert session.history[-1] == {"role": "assistant", "content": "first response"} + assert session.active_response_id == second_response_id + assert session.playback.as_dict() == { + "generated_ms": 2000, + "sent_ms": 2000, + "played_ms": 0, + "committed_ms": 0, + } + + @pytest.mark.asyncio async def test_late_playback_ack_does_not_advance_new_response_cursor(): handler = OmniDuplexSessionHandler( diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 480758d82..ad7baf764 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -296,12 +296,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # async abort (e.g. client disconnect during TTS streaming) to leave # requests in the waiting/running queues temporarily. waiting = getattr(self, "waiting") - running = getattr(self, "running") - for queue in (waiting, running): - for req in list(queue): - if getattr(req, "status", None) == RequestStatus.FINISHED_ABORTED: - queue.remove(req) + self._drop_aborted_queued_requests() self._process_pending_omni_inputs(model_mode="ar") + self._drop_aborted_queued_requests() + self._resync_streaming_input_counter() original_waiting = None if self._should_defer_waiting_admission(): @@ -441,15 +439,24 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # request is aborted while the model is executing it (e.g., # in pipeline parallelism or async scheduling). continue - if output_is_stale: + req_index = model_runner_output.req_id_to_index[req_id] + generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else [] + + stale_async_tokens = int(getattr(request, "async_tokens_to_discard", 0) or 0) + async_output_is_stale = bool(generated_token_ids and stale_async_tokens > 0) + if async_output_is_stale: + # Drain this marker even when the same frame also belongs to + # the scheduled-token stale window below. Both accounting + # domains must consume the old frame before new output passes. + request.async_tokens_to_discard = max(0, stale_async_tokens - len(generated_token_ids)) + + if output_is_stale or async_output_is_stale: # Output of a step scheduled before the request's in-flight # tokens were discarded (segment stop / session replacement). # num_computed_tokens was rolled back at the discard site, so # this output must not be appended or emitted. continue - req_index = model_runner_output.req_id_to_index[req_id] - generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else [] status_before_stop = request.status new_logprobs = None logprob_validation_failed = False @@ -572,6 +579,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): prefill_stats.finalize(self.kv_cache_manager.estimate_cached_tokens(request)) confirmed_num_computed_tokens = None + boundary_generation = None if stopped: if self.chunk_transfer_adapter is not None: confirmed_num_computed_tokens = self.chunk_transfer_adapter._confirmed_num_computed_tokens(request) @@ -581,6 +589,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): # Capture finish_reason BEFORE _handle_stopped_request, which may # reset the status to WAITING for streaming requests that continue. finish_reason = request.get_finished_reason() + if self.chunk_transfer_adapter is not None: + try: + boundary_generation = int(getattr(request, "_omni_segment_generation", 0) or 0) + except (TypeError, ValueError): + boundary_generation = 0 finished = self._handle_stopped_request(request) is_segment_finished = not finished if finished: @@ -668,12 +681,17 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if self.chunk_transfer_adapter is not None and ( inter_stage_output is not None or is_segment_finished or finished ): + save_kwargs = { + "new_token_ids": new_token_ids, + "confirmed_num_computed_tokens": confirmed_num_computed_tokens, + } + if is_segment_finished: + save_kwargs["segment_generation"] = boundary_generation self.chunk_transfer_adapter.save_async( inter_stage_output, request, is_segment_finished, - new_token_ids=new_token_ids, - confirmed_num_computed_tokens=confirmed_num_computed_tokens, + **save_kwargs, ) self._remove_stopped_requests_from_queues( @@ -774,9 +792,15 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): session.num_output_placeholders = 0 session.spec_token_ids = [] stage_id = self.vllm_config.model_config.stage_id + + update_infos = ( + getattr(update, "model_intermediate_buffer", None), + getattr(update, "additional_information", None), + ) if self.chunk_transfer_adapter and self.chunk_transfer_adapter.receives_chunks: self.chunk_transfer_adapter.requests_num_chunks_sent.pop(session.external_req_id, None) if stage_id != 0: + session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1 # Downstream async-chunk stages receive real payloads from the # connector. This update only resumes polling for the next segment. self.chunk_transfer_adapter.segment_finished_requests.discard(session.request_id) @@ -794,10 +818,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if self.log_stats: session.record_event(EngineCoreEventType.QUEUED) return - update_infos = ( - getattr(update, "model_intermediate_buffer", None), - getattr(update, "additional_information", None), - ) replace_streaming_prompt = any( isinstance(info, dict) and isinstance(info.get("meta"), dict) @@ -808,6 +828,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): self._release_replaced_streaming_prompt_cache(session) self._replace_streaming_session(session, update) return + session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1 super()._update_request_as_session(session, update) if hasattr(update, "model_intermediate_buffer"): session.model_intermediate_buffer = update.model_intermediate_buffer diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 533f48062..db0eedcf7 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -101,7 +101,10 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): # Temporary queue: preserve waiting order while requests await input. skipped_waiting_requests = create_request_queue(self.policy) req_index = 0 + self._drop_aborted_queued_requests() self._process_pending_omni_inputs(model_mode="generation") + self._drop_aborted_queued_requests() + self._resync_streaming_input_counter() # OMNI: Track requests that are already finished (e.g., marked by connector) # These should be removed from running and not scheduled diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 2832b54f6..58c3546cd 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -149,6 +149,7 @@ class OmniSchedulerMixin: def _replace_streaming_session(self, session: Request, update: StreamingUpdate) -> None: """Replace a downstream stage's placeholder with its next payload.""" + session._omni_segment_generation = int(getattr(session, "_omni_segment_generation", 0) or 0) + 1 adapter = getattr(self, "chunk_transfer_adapter", None) if adapter is not None: adapter.segment_finished_requests.discard(session.request_id) @@ -703,6 +704,7 @@ class OmniSchedulerMixin: ) finished = super().finish_requests(request_ids, finished_status) self._purge_finished_from_running(target_request_ids) + self._resync_streaming_input_counter() for request in finished: self._free_input_coordinator_request(request.request_id) @@ -868,3 +870,25 @@ class OmniSchedulerMixin: return resumable_segment_stop and req.request_id not in target_request_ids self.running[:] = [req for req in self.running if keep_running(req)] + + def _drop_aborted_queued_requests(self) -> None: + for queue in (self.waiting, self.skipped_waiting): + aborted = [req for req in queue if req.status == RequestStatus.FINISHED_ABORTED] + if aborted: + queue.remove_requests(aborted) + self.running[:] = [req for req in self.running if req.status != RequestStatus.FINISHED_ABORTED] + + def _resync_streaming_input_counter(self) -> None: + counter = getattr(self, "num_waiting_for_streaming_input", None) + if counter is None: + return + parked = sum( + 1 + for queue in (getattr(self, "waiting", None), getattr(self, "skipped_waiting", None)) + if queue + for request in queue + if request.status == RequestStatus.WAITING_FOR_STREAMING_REQ + ) + if parked != counter: + logger.debug("[Omni] resynced streaming-input counter %s -> %s", counter, parked) + self.num_waiting_for_streaming_input = parked 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 0979956e1..c19159e58 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 @@ -170,6 +170,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): # been read"). self._held_non_active: deque[Any] = deque() self.requests_num_chunks_sent: dict[str, int] = defaultdict(int) + # Boundary tasks advance this before the background sender runs, so a + # late old-segment frame cannot restore the previous watermark. + self._segment_generation: dict[str, int] = defaultdict(int) self._pending_streaming_prefills: dict[str, dict] = {} # The recv thread only records AR payloads. Prompt mutation is finalized # by the scheduler after it observes the ready marker, where the request @@ -309,6 +312,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): is_segment_finished: bool = False, new_token_ids: Iterable[int] | None = None, confirmed_num_computed_tokens: int | None = None, + segment_generation: int | None = None, ): """Build and enqueue one chunk for asynchronous sending. @@ -330,9 +334,32 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): scheduler step, before resumable cleanup mutates the request confirmed_num_computed_tokens: committed token count captured before a streaming transition can mutate ``request`` + segment_generation: generation captured before a resumable stop + can apply a queued update to the mutable request """ is_finished = request.is_finished() and not request.resumable + if not hasattr(self, "_segment_generation"): + self._segment_generation = defaultdict(int) external_req_id = request.external_req_id + if segment_generation is None: + raw_generation = getattr(request, "_omni_segment_generation", 0) + try: + generation = int(raw_generation) + except (TypeError, ValueError): + generation = 0 + else: + generation = int(segment_generation) + expected_generation = self._segment_generation.get(external_req_id, generation) + if generation < expected_generation: + logger.warning( + "Skip late save_async for request %s, segment_generation=%s, expected=%s", + external_req_id, + generation, + expected_generation, + ) + return + self._segment_generation[external_req_id] = generation + if confirmed_num_computed_tokens is None: confirmed_num_computed_tokens = self._confirmed_num_computed_tokens(request) processor_request = ( @@ -344,6 +371,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): "is_finished": is_finished, "is_segment_finished": is_segment_finished, "new_token_ids": tuple(int(token_id) for token_id in (new_token_ids or ())), + "segment_generation": generation, } reject_reason = None @@ -352,7 +380,9 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if sender_token is not None and sender_token.cancelled: reject_reason = "previous sender generation is still draining" # If the request is preempted, skip the already saved chunks. - elif confirmed_num_computed_tokens < self.requests_num_chunks_sent.get(external_req_id, 0): + elif not is_segment_finished and confirmed_num_computed_tokens < self.requests_num_chunks_sent.get( + external_req_id, 0 + ): logger.warning( f"Enqueue save_async for request {external_req_id}, " f"request.num_computed_tokens={request.num_computed_tokens}, " @@ -371,6 +401,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): if is_segment_finished: # The queued FIFO item now owns the old segment. Start the next # segment's deduplication watermark before the worker sends it. + self._segment_generation[external_req_id] = generation + 1 self.requests_num_chunks_sent.pop(external_req_id, None) if reject_reason is not None: logger.error("Cannot enqueue %s: %s", external_req_id, reject_reason) @@ -822,6 +853,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.request_payload.pop(external_req_id, None) self.code_prompt_token_ids.pop(external_req_id, None) self.requests_num_chunks_sent.pop(external_req_id, None) + self._segment_generation.pop(external_req_id, None) self.ramp_chunk_count.pop(external_req_id, None) self._adaptive_states.pop(external_req_id, None) self._pending_streaming_prefills.pop(external_req_id, None) diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 353ca7627..c0240c4bc 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -231,17 +231,37 @@ class OrchestratorRequestState: @dataclass -class StreamingInputState: - # Flag of streaming input request - enabled: bool = False +class StreamingSegmentState: + """Streaming segment-boundary state for one stage of a request. + + Every stage of a multi-stage pipeline reaches its own segment boundaries + independently, so this is tracked per stage rather than per request. + """ + # Flag of segment of streaming input finished - segment_finished: bool = False + finished: bool = False # Tokens from the current raw segment boundary. The vLLM output processor # does not guarantee that EngineCoreOutput.new_token_ids survives on the # processed RequestOutput used by the routing layer. - segment_token_ids: list[int] = field(default_factory=list) - segment_output_metadata: dict[str, Any] = field(default_factory=dict) - # Streaming update prompt length + token_ids: list[int] = field(default_factory=list) + output_metadata: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class StreamingInputState: + # Flag of streaming input request + enabled: bool = False + # Keyed by stage. ``_orchestration_loop`` records a boundary per stage poll + # into this one ``req_state``, while every consumer reads it for a specific + # ``stage_id``, so a flat slot is only correct as long as each read happens in + # the same iteration as its own stage's write. It is not for a stage that + # never records one (the ``poll_diffusion_output`` branch reaches + # ``_handle_processed_outputs`` without the raw-output loop) or for any + # consumer reading outside the writing iteration. + segments: dict[int, StreamingSegmentState] = field(default_factory=dict) + # Streaming update prompt length. NOT stage-keyed: its only consumer reads it + # through the streaming context on the stage-input-bridge path, not by stage + # id. Do not assume the rest of this struct is stage-safe. new_prompt_len_snapshot: int | None = None # Model/bridge-specific runtime states (e.g., thinker->talker) bridge_states: dict[str, Any] = field(default_factory=dict) @@ -249,6 +269,14 @@ class StreamingInputState: # while the downstream input processor consumes upstream token output. source_token_decoder: Callable[..., str] | None = None + def segment(self, stage_id: int | None) -> StreamingSegmentState: + """Return ``stage_id``'s segment state, or a fresh empty one if unreported. + + Fresh rather than shared: ``output_metadata`` is handed out by reference. + """ + segment = self.segments.get(stage_id) + return segment if segment is not None else StreamingSegmentState() + class _OrchestratorDuplexStagePort: """Adapts generic stage pools to the model-neutral duplex control plane.""" @@ -290,7 +318,15 @@ class _OrchestratorDuplexStagePort: duplex_state = {} request_state.streaming.bridge_states["duplex"] = duplex_state previous_epoch = duplex_state.get("epoch") - if not isinstance(duplex_state.get("model_turn_id"), int) or previous_epoch != context.fence.epoch: + current_model_turn_id = duplex_state.get("model_turn_id") + if ( + not isinstance(current_model_turn_id, int) + or previous_epoch != context.fence.epoch + or current_model_turn_id < context.fence.turn_id + ): + # A serving-side safety boundary can close a Realtime response + # before the model emits its normal turn_eos. Catch up the + # engine-owned identity when the next fenced append starts. duplex_state["model_turn_id"] = context.fence.turn_id duplex_state.update( { @@ -1093,15 +1129,12 @@ class Orchestrator: req_state = self.request_states.get(getattr(eco, "request_id", None)) if req_state is None or not req_state.streaming.enabled: continue - req_state.streaming.segment_finished = bool(getattr(eco, "is_segment_finished", False)) - req_state.streaming.segment_token_ids = ( - self._coerce_int_list(getattr(eco, "new_token_ids", None)) - if req_state.streaming.segment_finished - else [] - ) + segment_finished = bool(getattr(eco, "is_segment_finished", False)) raw_mm = self._completion_multimodal_output(eco, None) - req_state.streaming.segment_output_metadata = ( - dict(raw_mm) if req_state.streaming.segment_finished and isinstance(raw_mm, dict) else {} + req_state.streaming.segments[stage_id] = StreamingSegmentState( + finished=segment_finished, + token_ids=(self._coerce_int_list(getattr(eco, "new_token_ids", None)) if segment_finished else []), + output_metadata=(dict(raw_mm) if segment_finished and isinstance(raw_mm, dict) else {}), ) req_state.streaming.new_prompt_len_snapshot = getattr( eco, @@ -1479,7 +1512,7 @@ class Orchestrator: continue stage_metrics = None - segment_finished = req_state.streaming.enabled and req_state.streaming.segment_finished + segment_finished = req_state.streaming.enabled and req_state.streaming.segment(stage_id).finished if output.finished or segment_finished: stage_metrics = pool.build_stage_metrics( [output], @@ -1946,6 +1979,7 @@ class Orchestrator: 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): @@ -1955,11 +1989,7 @@ class Orchestrator: return request_finished = False - if ( - finished - and self.stage_pools[stage_id].final_output - and not (req_state.streaming.enabled and req_state.streaming.segment_finished) - ): + 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) @@ -1971,7 +2001,7 @@ class Orchestrator: # filter out again (the official implementation returns exactly one # result per audio chunk). is_duplex_stage0_segment = ( - stage_id == 0 and self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished + stage_id == 0 and self._is_duplex_session_request(req_state) and req_state.streaming.segment(0).finished ) if self.stage_pools[stage_id].final_output and not is_duplex_stage0_segment: await self.output_async_queue.put( @@ -1983,7 +2013,10 @@ class Orchestrator: metrics=stage_metrics, finished=( request_finished - or (self._is_duplex_session_request(req_state) and req_state.streaming.segment_finished) + or ( + self._is_duplex_session_request(req_state) + and req_state.streaming.segment(stage_id).finished + ) ), stage_submit_ts=submit_ts, ) @@ -2018,7 +2051,7 @@ class Orchestrator: return if ( - (finished or (req_state.streaming.enabled and req_state.streaming.segment_finished)) + (finished or segment_finished) and stage_id < req_state.final_stage_id 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) @@ -2157,7 +2190,9 @@ class Orchestrator: def _duplex_output_context( req_state: OrchestratorRequestState, *, - stage_id: int | None = None, + # Required: segment state is stage-keyed, so a caller that omitted this + # would silently read "no boundary" rather than the intended stage's. + stage_id: int | None, ) -> DuplexOutputContext | None: identity = req_state.duplex_identity if identity is None: @@ -2168,15 +2203,16 @@ class Orchestrator: ) fence = req_state.duplex_stage_fences.get(stage_id, identity.fence) if stage_id is not None else identity.fence + segment = req_state.streaming.segment(stage_id) return DuplexOutputContext( identity=DuplexRequestIdentity( session_id=identity.session_id, fence=fence, ), final_stage_id=req_state.final_stage_id, - segment_finished=req_state.streaming.enabled and req_state.streaming.segment_finished, - segment_token_ids=tuple(req_state.streaming.segment_token_ids), - segment_output_metadata=req_state.streaming.segment_output_metadata, + segment_finished=req_state.streaming.enabled and segment.finished, + segment_token_ids=tuple(segment.token_ids), + segment_output_metadata=segment.output_metadata, ) @staticmethod diff --git a/vllm_omni/experimental/fullduplex/openai/protocol.py b/vllm_omni/experimental/fullduplex/openai/protocol.py index bc259d59d..45e702d2b 100644 --- a/vllm_omni/experimental/fullduplex/openai/protocol.py +++ b/vllm_omni/experimental/fullduplex/openai/protocol.py @@ -389,6 +389,8 @@ class ResponseState: active_request_id: str | None = None active_response_id: str | None = None active_response_turn_id: int | None = None + active_response_input_commit_seq: int | None = None + active_response_awaits_input_commit: bool = False last_response_id: str | None = None assistant_text_buffer: list[str] = field(default_factory=list) assistant_audio_text_marks: list[DuplexAssistantAudioTextMark] = field(default_factory=list) @@ -413,9 +415,13 @@ class ConversationHistory: item_audio_text_marks: dict[str, list[DuplexAssistantAudioTextMark]] = field(default_factory=dict) pending_item_ids: dict[str, dict[str, object]] = field(default_factory=dict) pending_item_audio_text_marks: dict[str, list[DuplexAssistantAudioTextMark]] = field(default_factory=dict) + pending_item_input_commit_seqs: dict[str, int] = field(default_factory=dict) pending_truncations_ms: dict[str, int] = field(default_factory=dict) last_assistant_full_message: dict[str, object] | None = None last_assistant_audio_text_marks: list[DuplexAssistantAudioTextMark] = field(default_factory=list) + assistant_response_snapshots: dict[str, tuple[dict[str, object], tuple[DuplexAssistantAudioTextMark, ...], int]] = ( + field(default_factory=dict) + ) @dataclass @@ -504,6 +510,26 @@ class DuplexSession: def last_assistant_audio_text_marks(self) -> tuple[DuplexAssistantAudioTextMark, ...]: return tuple(self._conversation.last_assistant_audio_text_marks) + def has_assistant_response_item(self, response_id: str, item_id: str) -> bool: + if item_id != f"item_{response_id}": + return False + if self.active_response_id == response_id: + return True + message = self._conversation.item_ids.get(item_id) or self._conversation.pending_item_ids.get(item_id) + if isinstance(message, dict) and message.get("role") == "assistant": + return True + return response_id in self._conversation.assistant_response_snapshots + + def playback_ack_is_too_late(self, response_id: str, item_id: str) -> bool: + input_commit_seq = self._conversation.pending_item_input_commit_seqs.get(item_id) + if input_commit_seq is None: + snapshot = self._conversation.assistant_response_snapshots.get(response_id) + if snapshot is not None: + input_commit_seq = snapshot[2] + if input_commit_seq is None and self.active_response_id == response_id: + input_commit_seq = self._response.active_response_input_commit_seq + return input_commit_seq is not None and self.input_commit_seq > input_commit_seq + @property def playback(self) -> DuplexPlaybackView: return self._playback.current.snapshot() @@ -649,6 +675,7 @@ class DuplexSession: else: content = text + self._bind_active_response_to_input_commit(self._input.commit_seq + 1) self._input.commit_seq += 1 message = {"role": "user", "content": content} self._conversation.messages.append(message) @@ -675,6 +702,7 @@ class DuplexSession: } if transcript: input_audio_part["transcript"] = transcript + self._bind_active_response_to_input_commit(self._input.commit_seq + 1) self._input.commit_seq += 1 message = {"role": "user", "content": [input_audio_part]} if transcript: @@ -696,6 +724,11 @@ class DuplexSession: if completed_turn_id >= self.turn_id: self.turn_id = completed_turn_id + 1 + def _bind_active_response_to_input_commit(self, input_commit_seq: int) -> None: + if self.active_response_id is not None and self._response.active_response_awaits_input_commit: + self._response.active_response_input_commit_seq = int(input_commit_seq) + self._response.active_response_awaits_input_commit = False + def reserve_response_options(self, options: ResponseCreateOptions) -> None: if self._response.active_response_id is not None: raise RuntimeError("response options cannot be reserved while a response is active") @@ -724,6 +757,8 @@ class DuplexSession: response_id = f"resp-{self.session_id}-{self.epoch}-{uuid4().hex[:8]}" self._response.active_response_id = response_id self._response.active_response_turn_id = self.turn_id if turn_id is None else int(turn_id) + self._response.active_response_input_commit_seq = self.input_commit_seq + self._response.active_response_awaits_input_commit = self.turn_state == DuplexTurnState.USER_SPEAKING self._response.last_response_id = response_id self._response.assistant_text_buffer.clear() self._response.assistant_audio_text_marks.clear() @@ -937,12 +972,23 @@ class DuplexSession: playback_commit_policy: str | None = None, preserve_request: bool = False, ) -> dict[str, object] | None: + response_id = self._response.active_response_id + response_input_commit_seq = self._response.active_response_input_commit_seq + if response_input_commit_seq is None: + response_input_commit_seq = self.input_commit_seq + response_history_is_late = self.input_commit_seq > response_input_commit_seq assistant_text = "".join(self._response.assistant_text_buffer).strip() message = None if assistant_text: self._conversation.last_assistant_full_message = {"role": "assistant", "content": assistant_text} self._conversation.last_assistant_audio_text_marks = list(self._response.assistant_audio_text_marks) - if commit_text and assistant_text: + if response_id is not None: + self._conversation.assistant_response_snapshots[response_id] = ( + copy.deepcopy(self._conversation.last_assistant_full_message), + tuple(copy.deepcopy(self._response.assistant_audio_text_marks)), + response_input_commit_seq, + ) + if commit_text and assistant_text and not response_history_is_late: committed_text = self._playback_committed_text( assistant_text, playback_commit_policy=playback_commit_policy, @@ -952,11 +998,21 @@ class DuplexSession: if commit_text and committed_text: message = {"role": "assistant", "content": committed_text} self._conversation.messages.append(message) + effective_playback_policy = playback_commit_policy or self.config.playback_commit_policy + if ( + response_id is not None + and assistant_text + and message is None + and effective_playback_policy == DuplexPlaybackCommitPolicy.ACK_ONLY.value + ): + self.register_history_item(f"item_{response_id}", None) self._response.assistant_text_buffer.clear() if not preserve_request: self._response.active_request_id = None self._response.active_response_id = None self._response.active_response_turn_id = None + self._response.active_response_input_commit_seq = None + self._response.active_response_awaits_input_commit = False self._clear_response_metrics() self.turn_state = DuplexTurnState.IDLE self._restore_response_config() @@ -965,15 +1021,19 @@ class DuplexSession: def register_history_item(self, item_id: str | None, message: dict[str, object] | None) -> None: if not item_id: return + response_id = item_id.removeprefix("item_") if item_id.startswith("item_") else None + response_snapshot = ( + self._conversation.assistant_response_snapshots.get(response_id) if response_id is not None else None + ) + response_audio_text_marks = list(response_snapshot[1]) if response_snapshot is not None else None if message is None: - last_message = self._conversation.last_assistant_full_message - if last_message is None: + if response_snapshot is None: return - self._conversation.pending_item_ids[item_id] = dict(last_message) - if self._conversation.last_assistant_audio_text_marks: - self._conversation.pending_item_audio_text_marks[item_id] = list( - self._conversation.last_assistant_audio_text_marks - ) + last_message, audio_text_marks, _ = response_snapshot + self._conversation.pending_item_ids[item_id] = copy.deepcopy(last_message) + self._conversation.pending_item_input_commit_seqs[item_id] = response_snapshot[2] + if audio_text_marks: + self._conversation.pending_item_audio_text_marks[item_id] = list(copy.deepcopy(audio_text_marks)) pending_audio_ms = self._conversation.pending_truncations_ms.pop(item_id, None) if pending_audio_ms is not None: self.truncate_history_item( @@ -987,7 +1047,11 @@ class DuplexSession: self._truncate_message_to_audio_ms( message, audio_end_ms=pending_audio_ms, - marks=self._response.assistant_audio_text_marks or self._conversation.last_assistant_audio_text_marks, + marks=( + response_audio_text_marks + if response_audio_text_marks is not None + else self._response.assistant_audio_text_marks or self._conversation.last_assistant_audio_text_marks + ), playback=self._playback_cursor_for_item_id(item_id), ) if self._message_text_len(message) <= 0: @@ -999,17 +1063,28 @@ class DuplexSession: self._conversation.item_ids[item_id] = message self._conversation.pending_item_ids.pop(item_id, None) self._conversation.pending_item_audio_text_marks.pop(item_id, None) + self._conversation.pending_item_input_commit_seqs.pop(item_id, None) + if response_id is not None: + self._conversation.assistant_response_snapshots.pop(response_id, None) if message.get("role") == "assistant": - marks = self._response.assistant_audio_text_marks or self._conversation.last_assistant_audio_text_marks + marks = ( + response_audio_text_marks + if response_audio_text_marks is not None + else self._response.assistant_audio_text_marks or self._conversation.last_assistant_audio_text_marks + ) if marks: self._conversation.item_audio_text_marks[item_id] = list(marks) def delete_history_item(self, item_id: str) -> bool: + response_id = item_id.removeprefix("item_") if item_id.startswith("item_") else None message = self._conversation.item_ids.pop(item_id, None) self._conversation.item_audio_text_marks.pop(item_id, None) pending = self._conversation.pending_item_ids.pop(item_id, None) self._conversation.pending_item_audio_text_marks.pop(item_id, None) + self._conversation.pending_item_input_commit_seqs.pop(item_id, None) self._conversation.pending_truncations_ms.pop(item_id, None) + if response_id is not None: + self._conversation.assistant_response_snapshots.pop(response_id, None) if message is None: return pending is not None try: @@ -1043,7 +1118,10 @@ class DuplexSession: if changed: self._conversation.pending_item_ids.pop(item_id, None) self._conversation.pending_item_audio_text_marks.pop(item_id, None) + self._conversation.pending_item_input_commit_seqs.pop(item_id, None) self._conversation.pending_truncations_ms.pop(item_id, None) + if item_id.startswith("item_"): + self._conversation.assistant_response_snapshots.pop(item_id.removeprefix("item_"), None) return changed self._conversation.messages.append(message) self._conversation.item_ids[item_id] = message @@ -1053,7 +1131,10 @@ class DuplexSession: ) self._conversation.pending_item_ids.pop(item_id, None) self._conversation.pending_item_audio_text_marks.pop(item_id, None) + self._conversation.pending_item_input_commit_seqs.pop(item_id, None) self._conversation.pending_truncations_ms.pop(item_id, None) + if item_id.startswith("item_"): + self._conversation.assistant_response_snapshots.pop(item_id.removeprefix("item_"), None) return True changed = self._truncate_message_to_audio_ms( message, @@ -1068,6 +1149,10 @@ class DuplexSession: self._conversation.messages.remove(message) except ValueError: pass + if item_id.startswith("item_"): + self._conversation.assistant_response_snapshots.pop(item_id.removeprefix("item_"), None) + elif changed and item_id.startswith("item_"): + self._conversation.assistant_response_snapshots.pop(item_id.removeprefix("item_"), None) return changed def _truncate_message_to_audio_ms( @@ -1214,6 +1299,8 @@ class DuplexSession: self._response.active_request_id = None self._response.active_response_id = None self._response.active_response_turn_id = None + self._response.active_response_input_commit_seq = None + self._response.active_response_awaits_input_commit = False self._clear_response_metrics() self._restore_response_config() self.turn_state = DuplexTurnState.BARGE_IN @@ -1227,6 +1314,8 @@ class DuplexSession: self.state = DuplexSessionState.CLOSED self.turn_state = DuplexTurnState.IDLE self._response.active_response_turn_id = None + self._response.active_response_input_commit_seq = None + self._response.active_response_awaits_input_commit = False self._clear_response_metrics() self._restore_response_config() diff --git a/vllm_omni/experimental/fullduplex/openai/realtime_state.py b/vllm_omni/experimental/fullduplex/openai/realtime_state.py index de05cc71c..6a2c8f347 100644 --- a/vllm_omni/experimental/fullduplex/openai/realtime_state.py +++ b/vllm_omni/experimental/fullduplex/openai/realtime_state.py @@ -48,6 +48,9 @@ REALTIME_ERROR_TYPES_BY_CODE = { "input_audio_buffer_empty": "invalid_request_error", "missing_item_id": "invalid_request_error", "item_not_found": "invalid_request_error", + "playback_item_mismatch": "invalid_request_error", + "playback_item_not_found": "invalid_request_error", + "playback_ack_too_late": "invalid_request_error", "unsupported_audio_format": "invalid_request_error", "unsupported_ref_audio_path": "invalid_request_error", "ref_audio_required": "invalid_request_error", diff --git a/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py b/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py index 39778a012..98e0aaaf5 100644 --- a/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py +++ b/vllm_omni/experimental/fullduplex/openai/runtime_bridge.py @@ -362,11 +362,13 @@ class NativeRuntimeBridgeMixin: *, session: DuplexSession, expected_epoch: int | None, + model_turn_id: int | None = None, ) -> None: """Close an auto-response whose final silence unit did not end it.""" native = self._runtime_session_state(session) response_id = session.active_response_id response_epoch = session.epoch + response_turn_id = model_turn_id if model_turn_id is not None else session.active_response_turn_id if expected_epoch is not None and response_epoch != expected_epoch: return native.clear_continuation() @@ -380,9 +382,13 @@ class NativeRuntimeBridgeMixin: } ) if response_id is None: + if session.epoch == response_epoch and response_turn_id is not None: + session.complete_model_turn(response_turn_id) return if session.epoch != response_epoch or session.active_response_id != response_id: return + if response_turn_id is not None: + session.complete_model_turn(response_turn_id) session.end_response(commit_text=False, preserve_request=True) await send_json( { @@ -469,6 +475,7 @@ class NativeRuntimeBridgeMixin: send_json, session=session, expected_epoch=expected_epoch, + model_turn_id=payload_turn_id, ) return payload = self._native_silence_unit_payload() @@ -1003,6 +1010,12 @@ class NativeRuntimeBridgeMixin: return close_reason, emitted_response if auto_response: self._runtime_session_state(session).clear_continuation() + if not auto_continuations_remaining: + completed_turn_id = model_turn_id + if completed_turn_id is None: + completed_turn_id = session.active_response_turn_id + if completed_turn_id is not None: + session.complete_model_turn(completed_turn_id) session.end_response(commit_text=False, preserve_request=auto_response) await send_json( { diff --git a/vllm_omni/experimental/fullduplex/openai/serving.py b/vllm_omni/experimental/fullduplex/openai/serving.py index d0ddccb27..ce73684c0 100644 --- a/vllm_omni/experimental/fullduplex/openai/serving.py +++ b/vllm_omni/experimental/fullduplex/openai/serving.py @@ -1672,6 +1672,54 @@ class OmniDuplexSessionHandler( if response_id is None and item_id is None and session.active_response_id is not None: response_id = session.active_response_id item_id = f"item_{response_id}" + if response_id is not None: + expected_item_id = f"item_{response_id}" + if item_id is None: + item_id = expected_item_id + elif item_id != expected_item_id: + await send_json( + { + "type": "error", + "session_id": session.session_id, + "epoch": session.epoch, + "code": "playback_item_mismatch", + "error": "playback.ack item_id must match item_.", + } + ) + return + if not session.has_assistant_response_item(response_id, item_id): + await send_json( + { + "type": "error", + "session_id": session.session_id, + "epoch": session.epoch, + "code": "playback_item_not_found", + "error": f"No assistant response item is registered for {response_id}.", + } + ) + return + if session.playback_ack_is_too_late(response_id, item_id): + await send_json( + { + "type": "error", + "session_id": session.session_id, + "epoch": session.epoch, + "code": "playback_ack_too_late", + "error": "playback.ack arrived after a later user input was committed.", + } + ) + return + elif item_id is not None: + await send_json( + { + "type": "error", + "session_id": session.session_id, + "epoch": session.epoch, + "code": "playback_item_not_found", + "error": "playback.ack requires a response-owned assistant item.", + } + ) + return if event.get("truncate") is True: playback = session.acknowledge_playback( int(played_ms), @@ -1690,6 +1738,16 @@ class OmniDuplexSessionHandler( ) committed_history = False if isinstance(item_id, str) and item_id: + expected_item_id = f"item_{response_id}" if response_id is not None else None + if ( + expected_item_id == item_id + and item_id not in session.history_item_ids + and item_id not in session.pending_history_item_ids + ): + # ACK_ONLY responses are normally registered as pending when + # the response ends. Recover the response-local item if that + # registration was lost across the response/ack boundary. + session.register_history_item(item_id, None) committed_history = session.truncate_history_item( item_id, audio_end_ms=committed_cursor, @@ -1770,7 +1828,7 @@ class OmniDuplexSessionHandler( item_id = f"item_{old_response_id}" if committed_message is not None: session.register_history_item(item_id, committed_message) - elif committed_ms > 0: + elif committed_ms > 0 and not session.playback_ack_is_too_late(old_response_id, item_id): session.truncate_history_item(item_id, audio_end_ms=committed_ms) new_epoch, old_playback = self._advance_barge_in_epoch(session) if old_request_id is not None: