[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))