diff --git a/tests/core/sched/test_omni_ar_scheduler_free_request_cleanup.py b/tests/core/sched/test_omni_ar_scheduler_free_request_cleanup.py new file mode 100644 index 000000000..73c71fd63 --- /dev/null +++ b/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") diff --git a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py index cb0f50a5d..008e52636 100644 --- a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py +++ b/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() diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index bc8f9c30b..97b1618b5 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/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)) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index ecb4c4278..b8a2fbd9b 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/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.""" 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 5005ec280..cce5db472 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 @@ -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)