[Bugfix] Fix active_stream_window silently stalling audio generation (#5349) (#5373)

Signed-off-by: Shenglei Fu <sfu@confluent.io>
Co-authored-by: Claude <noreply@anthropic.com>

tests/core/sched/test_omni_ar_scheduler_free_request_cleanup.py

@@ -0,0 +1,74 @@
+"""Regression tests for vllm-project/vllm-omni#5349 (P1): normal completion
+goes through _free_request(), not finish_requests() (the external
+abort/cancel entry point). Without a cleanup_receiver() call there,
+_active_streams never releases on ordinary completion, so the bounded-K
+window fills with stale entries after K completions.
+"""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock
+
+import pytest
+
+# Imports must run in this order: vllm_omni applies patches to vllm.v1.request before
+# Request / RequestStatus are bound in this module. Ruff isort would reorder them.
+# isort: off
+import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+# isort: on
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_scheduler(*, chunk_transfer_adapter=None) -> OmniARScheduler:
+    """Minimal OmniARScheduler exercising _free_request()'s no-KV-transfer
+    happy path."""
+    sched = OmniARScheduler.__new__(OmniARScheduler)
+    sched._omits_kv_transfer_cache = {}
+    sched._connector_finished = lambda request: (False, None)
+    sched.encoder_cache_manager = MagicMock()
+    sched.finished_req_ids = set()
+    sched._new_prompt_len_snapshot = {}
+    sched.finished_req_ids_dict = None
+    sched._should_transfer_kv_for_request = lambda req_id: False
+    sched._free_blocks = MagicMock()
+    sched._free_input_coordinator_request = MagicMock()
+    sched.chunk_transfer_adapter = chunk_transfer_adapter
+    return sched
+
+
+class _FakeFinishedRequest:
+    """Not a SimpleNamespace: SimpleNamespace defines __eq__, which makes it
+    unhashable, and _free_request() puts the request in a set."""
+
+    def __init__(self, request_id: str) -> None:
+        self.request_id = request_id
+
+    def is_finished(self) -> bool:
+        return True
+
+
+def _make_finished_request(request_id: str = "req-1"):
+    return _FakeFinishedRequest(request_id)
+
+
+def test_free_request_releases_chunk_transfer_adapter_receiver_state():
+    adapter = MagicMock()
+    sched = _make_scheduler(chunk_transfer_adapter=adapter)
+    request = _make_finished_request("req-1")
+
+    sched._free_request(request)
+
+    adapter.cleanup_receiver.assert_called_once_with("req-1")
+
+
+def test_free_request_is_safe_without_a_chunk_transfer_adapter():
+    """Most stages have no chunk transfer adapter configured."""
+    sched = _make_scheduler(chunk_transfer_adapter=None)
+    request = _make_finished_request("req-1")
+
+    sched._free_request(request)  # must not raise
+
+    sched._free_input_coordinator_request.assert_called_once_with("req-1")

tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py

@@ -75,6 +75,7 @@ def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None:
     coordinator = FakeInputCoordinator()
     scheduler = OmniARScheduler.__new__(OmniARScheduler)
     scheduler.input_coordinator = coordinator
+    scheduler.chunk_transfer_adapter = None
     scheduler._omits_kv_transfer_cache = {"req-free": True}
     scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None)
     scheduler.finished_req_ids = set()

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -192,7 +192,7 @@ def test_load_poll(build_adapter):
     assert request.additional_information == payload
     assert adapter.get_req_chunk["req-1"] == 1
     assert "req-1" in adapter._finished_load_reqs
-    assert "req-1" in adapter.finished_requests
+    assert "req-1" in adapter.upstream_exhausted_requests
     assert "req-1" not in adapter._pending_load_reqs
 
 
@@ -459,7 +459,7 @@ def test_load_poll_non_ar_merges_into_existing_additional_information(build_adap
     assert request.additional_information["meta"]["phase"] == "decode"
     assert request.additional_information["kv_metadata"] == {"foo": "bar"}
     assert "req-non-ar" in adapter._finished_load_reqs
-    assert "req-non-ar" in adapter.finished_requests
+    assert "req-non-ar" in adapter.upstream_exhausted_requests
 
 
 def test_load_poll_ar_request_additional_information_concats_tensors(build_adapter):
@@ -527,7 +527,7 @@ def test_process_and_restore_queues(build_adapter):
 
 
 def test_fifo_promotion(build_adapter):
-    adapter, _ = build_adapter(stage_id=1, max_num_seqs=2, active_stream_window=2)
+    adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=2, active_stream_window=2)
     reqs = [_req(f"req-{idx}", RequestStatus.WAITING) for idx in range(1, 5)]
     waiting_queue = DummyWaitingQueue(reqs)
     running_queue = []
@@ -540,11 +540,12 @@ def test_fifo_promotion(build_adapter):
     assert reqs[1].status == RequestStatus.WAITING_FOR_CHUNK
     assert reqs[2].status == RequestStatus.WAITING
 
-    adapter.finished_requests.add("req-1")
+    # Ordinary completion calls cleanup_receiver(), not finish_requests()
+    # (the abort path) -- see test_omni_ar_scheduler_free_request_cleanup.py.
     # Eviction is deferred to postprocess_scheduler_output in the runtime path
-    # (commit c4d95fd9 — otherwise the terminal chunk deadlocks at c=8 K=2).
-    # Simulate it here so promotion can pick up the freed slot.
-    adapter._evict_finished_active_streams({"req-1"})
+    # (commit c4d95fd9 -- otherwise the terminal chunk deadlocks at c=8 K=2).
+    # Simulate the restore step here so promotion can pick up the freed slot.
+    adapter.cleanup_receiver("req-1")
     adapter.restore_queues(waiting_queue, running_queue)
     adapter.process_pending_chunks(waiting_queue, running_queue)
 
@@ -590,7 +591,7 @@ def test_legacy_k0(build_adapter):
 
 
 def test_finished_releases_slot(build_adapter):
-    adapter, _ = build_adapter(stage_id=1, max_num_seqs=1, active_stream_window=1)
+    adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=1, active_stream_window=1)
     req_1 = _req("req-1", RequestStatus.WAITING)
     req_2 = _req("req-2", RequestStatus.WAITING)
     waiting_queue = DummyWaitingQueue([req_1, req_2])
@@ -600,10 +601,11 @@ def test_finished_releases_slot(build_adapter):
     assert list(adapter._active_streams) == ["req-1"]
     assert waiting_queue == []
 
-    adapter.finished_requests.add("req-1")
-    # Eviction is deferred to postprocess_scheduler_output in the runtime path
-    # (commit c4d95fd9). Simulate it so promotion can pick up the freed slot.
-    adapter._evict_finished_active_streams({"req-1"})
+    # req-1 finishes ordinarily (see test_fifo_promotion for why this is
+    # cleanup_receiver(), not finish_requests()). Eviction is deferred to
+    # postprocess_scheduler_output in the runtime path (commit c4d95fd9);
+    # simulate the restore step so promotion can pick up the freed slot.
+    adapter.cleanup_receiver("req-1")
     adapter.restore_queues(waiting_queue, running_queue)
     adapter.process_pending_chunks(waiting_queue, running_queue)
 
@@ -631,6 +633,79 @@ def test_postprocess_scheduler_output(build_adapter):
     assert adapter.requests_with_ready_chunks == {"leftover"}
 
 
+@pytest.mark.parametrize("model_mode", ["ar", "generation"])
+def test_active_stream_window_stalls_lone_running_request_after_upstream_finishes(build_adapter, model_mode):
+    """Regression test for vllm-project/vllm-omni#5349.
+
+    upstream_exhausted_requests means the upstream sent its terminal chunk,
+    not that this stage's own generation is done. A downstream stage can
+    still have a long decode ahead after its upstream finishes. Evicting on
+    that signal and then permanently denying re-admission wedges the
+    request out of running_queue forever, with no error.
+
+    requests_with_ready_chunks is set up front so the request stays in
+    running_queue this tick instead of legitimately parking in
+    waiting_for_chunk_running_requests (a separate mechanism this test must
+    not conflate with the bug).
+
+    Parametrized over both worker types: the real issue hit an "ar" stage
+    (Qwen3-Omni's talker); must not regress "generation" (Qwen3-TTS's
+    Code2Wav), which this feature was designed for.
+    """
+    adapter, _ = build_adapter(stage_id=1, model_mode=model_mode, max_num_seqs=2, active_stream_window=2)
+    running_req = _req("req-1", RequestStatus.RUNNING)
+    running_queue = [running_req]
+    waiting_queue = DummyWaitingQueue([])
+
+    adapter._active_streams["req-1"] = running_req
+    adapter.requests_with_ready_chunks.add("req-1")
+
+    # Upstream terminal chunk arrives; this stage hasn't finished.
+    adapter.upstream_exhausted_requests.add("req-1")
+    scheduler_output = SimpleNamespace(
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=SimpleNamespace(req_ids=["req-1"]),
+    )
+    adapter.postprocess_scheduler_output(scheduler_output)
+    assert not running_req.is_finished()  # sanity: not done
+
+    # No competition for the slot; req-1 must stay schedulable.
+    adapter.process_pending_chunks(waiting_queue, running_queue)
+
+    assert running_queue == [running_req], "req-1 must still be schedulable -- losing it here is #5349"
+    assert not adapter.waiting_for_chunk_running_requests
+    assert running_req not in adapter._held_non_active
+    assert list(adapter._active_streams) == ["req-1"]
+
+
+def test_cleanup_receiver_releases_multiple_slots_in_sequence(build_adapter):
+    """Companion to test_fifo_promotion: after EACH of the K active streams
+    finishes (not just one), promotion of a new K-sized batch must still
+    work -- the window must not stay exhausted by stale entries. See
+    test_omni_ar_scheduler_free_request_cleanup.py for the scheduler-level
+    proof that ordinary completion calls cleanup_receiver().
+    """
+    adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=2, active_stream_window=2)
+    reqs = [_req(f"req-{idx}", RequestStatus.WAITING) for idx in range(1, 5)]
+    waiting_queue = DummyWaitingQueue(reqs)
+    running_queue = []
+
+    adapter.process_pending_chunks(waiting_queue, running_queue)
+    assert list(adapter._active_streams) == ["req-1", "req-2"]
+
+    adapter.cleanup_receiver("req-1")
+    adapter.cleanup_receiver("req-2")
+    # Eviction is deferred to postprocess_scheduler_output in the runtime
+    # path; simulate the restore step so promotion can pick up the freed
+    # slots (see test_fifo_promotion).
+    adapter.restore_queues(waiting_queue, running_queue)
+    adapter.process_pending_chunks(waiting_queue, running_queue)
+
+    assert list(adapter._active_streams) == ["req-3", "req-4"], (
+        "both freed slots must be reusable after ordinary completion"
+    )
+
+
 # ---------------------------------------------------------------
 # Cleanup tests
 # ---------------------------------------------------------------
@@ -638,7 +713,7 @@ def test_postprocess_scheduler_output(build_adapter):
 
 def _populate_adapter_state(adapter, req_id="req-1", ext_id="ext-1"):
     """Fill every per-request structure so cleanup can be verified."""
-    adapter.finished_requests.add(req_id)
+    adapter.upstream_exhausted_requests.add(req_id)
     adapter._active_streams[req_id] = SimpleNamespace(request_id=req_id)
     adapter.get_req_chunk[req_id] = 3
     adapter.requests_with_ready_chunks.add(req_id)
@@ -659,7 +734,7 @@ def test_cleanup_clears_all_state(build_adapter):
 
     adapter.cleanup(req_id, ext_id)
 
-    assert req_id not in adapter.finished_requests
+    assert req_id not in adapter.upstream_exhausted_requests
     assert req_id not in adapter._active_streams
     assert req_id not in adapter.get_req_chunk
     assert req_id not in adapter.requests_with_ready_chunks
@@ -712,7 +787,7 @@ def test_cleanup_request_id_reuse_not_polluted(build_adapter):
 
     adapter.cleanup(req_id, ext_id)
 
-    assert req_id not in adapter.finished_requests
+    assert req_id not in adapter.upstream_exhausted_requests
     assert req_id not in adapter.get_req_chunk
 
 
@@ -738,7 +813,7 @@ def test_cleanup_only_affects_target_request(build_adapter):
 
     adapter.cleanup("req-a", "ext-a")
 
-    assert "req-b" in adapter.finished_requests
+    assert "req-b" in adapter.upstream_exhausted_requests
     assert "req-b" in adapter.get_req_chunk
     assert "ext-b" in adapter.put_req_chunk
     assert "ext-b" in adapter.request_payload
@@ -761,13 +836,13 @@ def test_cleanup_after_poll_flow(build_adapter):
     connector.get.return_value = (payload, 8)
     adapter._poll_single_request(request)
 
-    assert "req-flow" in adapter.finished_requests
+    assert "req-flow" in adapter.upstream_exhausted_requests
     assert adapter.get_req_chunk["req-flow"] == 1
     assert "req-flow" in adapter.request_ids_mapping
 
     adapter.cleanup("req-flow", "ext-flow")
 
-    assert "req-flow" not in adapter.finished_requests
+    assert "req-flow" not in adapter.upstream_exhausted_requests
     assert "req-flow" not in adapter.get_req_chunk
     assert "req-flow" not in adapter.request_ids_mapping
     assert "ext-flow" not in adapter.request_payload
@@ -801,7 +876,7 @@ def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter
     adapter.waiting_for_chunk_waiting_requests = deque([zombie, other])
     adapter.waiting_for_chunk_running_requests = deque([other, zombie])
     adapter.requests_with_ready_chunks.add("req-zombie")
-    adapter.finished_requests.add("req-zombie")
+    adapter.upstream_exhausted_requests.add("req-zombie")
     requests_map = {
         "req-zombie": zombie,
         "req-live": other,
@@ -816,7 +891,7 @@ def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter
     assert [req.request_id for req in adapter.waiting_for_chunk_waiting_requests] == ["req-live"]
     assert [req.request_id for req in adapter.waiting_for_chunk_running_requests] == ["req-live"]
     assert "req-zombie" not in adapter.requests_with_ready_chunks
-    assert "req-zombie" not in adapter.finished_requests
+    assert "req-zombie" not in adapter.upstream_exhausted_requests
 
 
 def test_finish_requests_releases_active_stream_slot(build_adapter):
@@ -886,7 +961,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
     cleanup_calls = []
 
     adapter_mock = mocker.MagicMock()
-    adapter_mock.finished_requests = {"req-s1"}
+    adapter_mock.upstream_exhausted_requests = {"req-s1"}
     adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
 
     from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
@@ -1122,7 +1197,7 @@ def test_wire_round_trip_struct_to_dict_contract():
 def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs):
     """Build a mock scheduler with requests queued for deferred finish."""
     adapter_mock = mocker.MagicMock()
-    adapter_mock.finished_requests = {r.request_id for r in pending_finish_reqs}
+    adapter_mock.upstream_exhausted_requests = {r.request_id for r in pending_finish_reqs}
     cleanup_calls = []
     adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
 

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -917,6 +917,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             return kv_xfer_params, None
         finally:
             self._free_input_coordinator_request(request_id)
+            # Normal completion runs through here, not finish_requests()
+            # (the abort path) -- see vllm-project/vllm-omni#5349.
+            if self.chunk_transfer_adapter is not None:
+                self.chunk_transfer_adapter.cleanup_receiver(request_id)
 
     def _mark_request_for_kv_transfer(self, req_id: str, seq_len: int) -> None:
         """Mark a request as needing KV cache transfer when it finishes."""

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -73,7 +73,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         # but popped at segment boundaries (unlike put_req_chunk which is
         # request-global for connector key continuity).
         self.ramp_chunk_count: dict[str, int] = defaultdict(int)
-        self.finished_requests: set[str] = set()
+        self.upstream_exhausted_requests: set[str] = set()
         self.segment_finished_requests: set[str] = set()
         self.request_payload = {}
         self.code_prompt_token_ids: dict[str, list[torch.Tensor]] = defaultdict(list)
@@ -235,13 +235,13 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                     construct_next_stage_streaming_input_prompt(payload_data, request)
 
                 if payload_finished:
-                    self.finished_requests.add(req_id)
+                    self.upstream_exhausted_requests.add(req_id)
                     request.resumable = False
                 if payload_segment_finished:
                     self.segment_finished_requests.add(req_id)
             else:
                 if payload_finished:
-                    self.finished_requests.add(req_id)
+                    self.upstream_exhausted_requests.add(req_id)
                     request.resumable = False
                 if payload_segment_finished:
                     self.segment_finished_requests.add(req_id)
@@ -384,11 +384,12 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
     def is_done_receiving_chunks(self, request_id: str) -> bool:
         """Return True if the request should stop polling upstream chunks.
 
-        Covers both the whole-request finish marker (``finished_requests``) and
-        the per-segment finish marker (``segment_finished_requests``) used while
-        waiting for the next streaming input slice.
+        Covers both the whole-request marker (``upstream_exhausted_requests``)
+        and the per-segment marker (``segment_finished_requests``) used while
+        waiting for the next streaming input slice. Neither means this
+        stage's own generation is done -- see vllm-project/vllm-omni#5349.
         """
-        return request_id in self.finished_requests or request_id in self.segment_finished_requests
+        return request_id in self.upstream_exhausted_requests or request_id in self.segment_finished_requests
 
     ########################################################################
     # Cleanup
@@ -401,22 +402,41 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         enqueued work that the background thread has not yet processed,
         because it only touches receiver-side dictionaries.
 
+        Must also purge the request from the chunk-parking deques
+        (``waiting_for_chunk_waiting_requests`` / ``_running_requests`` /
+        ``_held_non_active``): otherwise a caller that calls
+        ``restore_queues()`` without ``scheduler_requests`` (e.g. a unit
+        test, or any future caller not synced with the scheduler's own
+        request-removal timing) would re-admit an already-finished
+        request into the visible queue, which ``_promote_active_streams``
+        would then FIFO-promote ahead of genuinely-waiting requests. See
+        vllm-project/vllm-omni#5349's active-stream-window tests.
+
         Idempotent: calling with an already-cleaned or unknown id is safe.
         """
-        if request_id in self.finished_requests:
-            self._evict_finished_active_streams({request_id})
-        else:
-            self._active_streams.pop(request_id, None)
-        self.finished_requests.discard(request_id)
+        self._active_streams.pop(request_id, None)
+        self.upstream_exhausted_requests.discard(request_id)
         self.segment_finished_requests.discard(request_id)
         self.get_req_chunk.pop(request_id, None)
         self.requests_with_ready_chunks.discard(request_id)
         self.request_ids_mapping.pop(request_id, None)
         self.requests_origin_status.pop(request_id, None)
+        self._discard_from_chunk_deque(self.waiting_for_chunk_waiting_requests, request_id)
+        self._discard_from_chunk_deque(self.waiting_for_chunk_running_requests, request_id)
+        self._discard_from_chunk_deque(self._held_non_active, request_id)
 
         self._cancelled_load_reqs.add(request_id)
         self._finished_load_reqs.discard(request_id)
 
+    @staticmethod
+    def _discard_from_chunk_deque(deque_list: deque[Any], request_id: str) -> None:
+        if not deque_list:
+            return
+        for _ in range(len(deque_list)):
+            request = deque_list.popleft()
+            if request.request_id != request_id:
+                deque_list.append(request)
+
     def cleanup_sender(self, external_req_id: str) -> None:
         """Reclaim sender-side per-request state (keyed by external id).
 
@@ -524,13 +544,6 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self._promote_active_streams(waiting_queue)
         self._preempt_non_active_running(waiting_queue, running_queue)
 
-    def _evict_finished_active_streams(self, request_ids: set[str] | None = None) -> None:
-        for request_id in list(self._active_streams):
-            if request_ids is not None and request_id not in request_ids:
-                continue
-            if request_id in self.finished_requests:
-                self._active_streams.pop(request_id, None)
-
     def _promote_active_streams(self, queue: Any) -> None:
         if len(self._active_streams) >= self._active_window:
             return
@@ -538,7 +551,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
             if len(self._active_streams) >= self._active_window:
                 return
             request_id = request.request_id
-            if request_id in self._active_streams or request_id in self.finished_requests:
+            if request_id in self._active_streams:
                 continue
             # Iterating the existing queue preserves FIFO admission.
             self._active_streams[request_id] = request
@@ -550,7 +563,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         if request_id in self._active_streams:
             self._active_streams[request_id] = request
             return True
-        if request_id in self.finished_requests or len(self._active_streams) >= self._active_window:
+        if len(self._active_streams) >= self._active_window:
             return False
         self._active_streams[request_id] = request
         return True
@@ -697,25 +710,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
         if requests is not None:
             self.attach_cached_additional_information(scheduler_output, requests)
-        scheduled_req_ids = self._scheduled_request_ids(scheduler_output)
         self._clear_chunk_ready(scheduler_output)
-        if scheduled_req_ids:
-            # Terminal chunks must stay active until they are scheduled once.
-            self._evict_finished_active_streams(scheduled_req_ids)
-
-    @staticmethod
-    def _scheduled_request_ids(scheduler_output: Any) -> set[str]:
-        req_ids: set[str] = set()
-        if scheduler_output.scheduled_new_reqs:
-            for req_data in scheduler_output.scheduled_new_reqs:
-                req_id = getattr(req_data, "req_id", None)
-                if req_id:
-                    req_ids.add(req_id)
-        if scheduler_output.scheduled_cached_reqs:
-            for req_id in scheduler_output.scheduled_cached_reqs.req_ids:
-                if req_id:
-                    req_ids.add(req_id)
-        return req_ids
 
     @staticmethod
     def attach_cached_additional_information(scheduler_output: Any, requests: dict[str, Request]) -> None:
@@ -817,7 +812,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         for req_id in request_ids:
             self._active_streams.pop(req_id, None)
             self.requests_with_ready_chunks.discard(req_id)
-            self.finished_requests.discard(req_id)
+            self.upstream_exhausted_requests.discard(req_id)
             self._finished_load_reqs.discard(req_id)
             self._cancelled_load_reqs.add(req_id)