anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Shenglei Fu <sfu@confluent.io>
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Shenglei FuandClaude authored and GitHub committed 2026-07-30 14:56:18 +00:00
1 parent ce3d8e4bac
commit 3b0718f22e
5 files changed
+211 -62
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py
+97 -22
View File
@@ -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))