From 73df8326d76fe3bba0b7b5a6abf6ad68976f24e8 Mon Sep 17 00:00:00 2001 From: longguo <107740309+abinggo@users.noreply.github.com> Date: Wed, 10 Jun 2026 09:00:23 +0800 Subject: [PATCH] [Bugfix] purge chunk-transfer zombies on every schedule tick to keep engine-core alive on aborts (fixes #3736) (#3774) Signed-off-by: abinggo <107740309+abinggo@users.noreply.github.com> Co-authored-by: npuichigo <11533479+npuichigo@users.noreply.github.com> --- ...mni_generation_scheduler_update_session.py | 298 ++++++++++++++++++ ...st_omni_scheduler_finish_requests_purge.py | 158 ++++++++++ ...mni_scheduler_input_coordinator_cleanup.py | 3 + .../test_chunk_transfer_adapter.py | 168 +++++++++- vllm_omni/core/sched/omni_ar_scheduler.py | 31 +- .../core/sched/omni_generation_scheduler.py | 20 +- vllm_omni/core/sched/omni_scheduler_mixin.py | 120 +++++++ .../chunk_transfer_adapter.py | 64 ++++ 8 files changed, 858 insertions(+), 4 deletions(-) create mode 100644 tests/core/sched/test_omni_scheduler_finish_requests_purge.py diff --git a/tests/core/sched/test_omni_generation_scheduler_update_session.py b/tests/core/sched/test_omni_generation_scheduler_update_session.py index cf492e7cc..d94996014 100644 --- a/tests/core/sched/test_omni_generation_scheduler_update_session.py +++ b/tests/core/sched/test_omni_generation_scheduler_update_session.py @@ -20,6 +20,7 @@ from vllm.sampling_params import SamplingParams from vllm.v1.engine import EngineCoreEventType from vllm.v1.request import Request, RequestStatus, StreamingUpdate from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler +from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin # isort: on @@ -142,3 +143,300 @@ class TestReplaceSessionWithStreamingUpdate: assert session.events assert session.events[-1].type == EngineCoreEventType.QUEUED + + +class _RealignSchedulerStub(OmniSchedulerMixin): + """Minimal scheduler surface for ``_realign_request_status_to_queues``. + + Pins ``self.requests`` / ``self.running`` / ``self.waiting`` to the + same shapes the upstream ``Scheduler.finish_requests`` reads, so the + helper sees realistic state without spinning a real scheduler. + """ + + def __init__( + self, + *, + requests: dict, + running: list, + waiting: list, + ) -> None: + self.requests = requests + self.running = running + self.waiting = waiting + + +class TestRealignRequestStatusToQueues: + """Regression for the residual hang described in + https://github.com/vllm-project/vllm-omni/pull/3774 -- chunk-transfer + adapter's ``requests_origin_status`` table goes stale on the + ``waiting → running`` admit transition, and an abort that lands + before the next deque round-trip stomps stale ``WAITING`` onto a + request that actually lives in ``self.running``. After + ``max_num_seqs`` such aborts every ``input_batch`` slot is leaked and + new requests hang at ``chunks=0``. + """ + + def test_running_with_stale_waiting_status_is_realigned_to_running( + self, + ) -> None: + """The exact race the hang reproduces: a request lives in + ``self.running`` but its ``status`` is ``WAITING``. After + realign, status must be ``RUNNING`` so upstream + ``Scheduler.finish_requests`` removes it from ``self.running``. + """ + req = _make_request(request_id="req-stale") + req.status = RequestStatus.WAITING # stale -- actually in running + + sched = _RealignSchedulerStub( + requests={req.request_id: req}, + running=[req], + waiting=[], + ) + sched._realign_request_status_to_queues([req.request_id]) + + assert req.status == RequestStatus.RUNNING + + def test_waiting_with_stale_running_status_is_realigned_to_waiting( + self, + ) -> None: + """Symmetric case: request lives in ``self.waiting`` but status + is ``RUNNING``. Upstream's else branch should run, so realign to + ``WAITING``. + """ + req = _make_request(request_id="req-stale-2") + req.status = RequestStatus.RUNNING # stale -- actually in waiting + + sched = _RealignSchedulerStub( + requests={req.request_id: req}, + running=[], + waiting=[req], + ) + sched._realign_request_status_to_queues([req.request_id]) + + assert req.status == RequestStatus.WAITING + + def test_already_aligned_status_is_left_unchanged(self) -> None: + """Healthy case: status matches actual queue. No status mutation + means no spurious side effects. + """ + req_running = _make_request(request_id="req-r") + req_running.status = RequestStatus.RUNNING + req_waiting = _make_request(request_id="req-w") + req_waiting.status = RequestStatus.WAITING + + sched = _RealignSchedulerStub( + requests={ + req_running.request_id: req_running, + req_waiting.request_id: req_waiting, + }, + running=[req_running], + waiting=[req_waiting], + ) + sched._realign_request_status_to_queues([req_running.request_id, req_waiting.request_id]) + + assert req_running.status == RequestStatus.RUNNING + assert req_waiting.status == RequestStatus.WAITING + + def test_unknown_request_id_is_skipped_silently(self) -> None: + """Aligning ids that aren't in ``self.requests`` (already freed + upstream, or never existed) must be a no-op. + """ + sched = _RealignSchedulerStub(requests={}, running=[], waiting=[]) + sched._realign_request_status_to_queues(["never-existed"]) # no raise + + def test_request_in_neither_queue_is_left_unchanged(self) -> None: + """If a tracked request is in neither queue (e.g. parked in a + chunk-transfer deque), realign must not invent a status. The + adapter / deque purge owns that surface; the realign helper is + only for the admit-transition staleness. + """ + req = _make_request(request_id="req-parked") + req.status = RequestStatus.WAITING_FOR_CHUNK + + sched = _RealignSchedulerStub( + requests={req.request_id: req}, + running=[], + waiting=[], + ) + sched._realign_request_status_to_queues([req.request_id]) + + assert req.status == RequestStatus.WAITING_FOR_CHUNK + + def test_request_ids_str_is_treated_as_single_id(self) -> None: + """Match upstream ``Scheduler.finish_requests`` resolution: a + bare string is treated as one id, not iterated as characters. + """ + req = _make_request(request_id="req-s") + req.status = RequestStatus.WAITING + + sched = _RealignSchedulerStub( + requests={req.request_id: req}, + running=[req], + waiting=[], + ) + sched._realign_request_status_to_queues(req.request_id) + + assert req.status == RequestStatus.RUNNING + + def test_request_ids_none_aligns_every_known_request(self) -> None: + """``request_ids=None`` matches upstream's "all requests" path. + Realign must walk every entry in ``self.requests`` and fix any + stale status it finds. + """ + stale = _make_request(request_id="req-stale-none") + stale.status = RequestStatus.WAITING # but actually in running + + clean = _make_request(request_id="req-clean-none") + clean.status = RequestStatus.RUNNING + + sched = _RealignSchedulerStub( + requests={ + stale.request_id: stale, + clean.request_id: clean, + }, + running=[stale, clean], + waiting=[], + ) + sched._realign_request_status_to_queues(None) + + assert stale.status == RequestStatus.RUNNING + assert clean.status == RequestStatus.RUNNING + + def test_finished_request_is_skipped(self) -> None: + """Already-finished requests must not be touched -- they may + have legitimate finished statuses (FINISHED_STOPPED etc.) that + a status flip would corrupt. + """ + req = _make_request(request_id="req-finished") + req.status = RequestStatus.FINISHED_STOPPED + + sched = _RealignSchedulerStub( + requests={req.request_id: req}, + running=[req], # not realistic, but we want to prove the guard fires + waiting=[], + ) + sched._realign_request_status_to_queues([req.request_id]) + + assert req.status == RequestStatus.FINISHED_STOPPED + + +class TestPurgeFinishedFromRunning: + """Regression for the residual ``self.running`` slot leak surface + paired with ``_realign_request_status_to_queues``: even after + realign + ``super().finish_requests`` runs, corner cases can leave + already-finished or untracked entries in ``self.running`` -- e.g. + a connector cleanup that pops from ``self.requests`` without + unwinding ``self.running``, or a ``status`` mid-transition when + finish ran. The defensive post-finish purge sweeps those residues + so the worker's ``input_batch`` slot never pins a freed request. + + See https://github.com/vllm-project/vllm-omni/pull/3774 discussion + and the residual-hang reproduction in that PR. + """ + + def test_finished_request_is_purged_from_running(self) -> None: + """``is_finished()`` request lingering in ``self.running`` must + be swept so its ``input_batch`` slot is freed.""" + finished = _make_request(request_id="req-finished") + finished.status = RequestStatus.FINISHED_STOPPED # is_finished() True + + sched = _RealignSchedulerStub( + requests={finished.request_id: finished}, + running=[finished], + waiting=[], + ) + sched._purge_finished_from_running() + + assert sched.running == [] + + def test_untracked_request_is_purged_from_running(self) -> None: + """Request lingering in ``self.running`` but no longer present + in ``self.requests`` (already freed by upstream / connector) + must be swept.""" + untracked = _make_request(request_id="req-untracked") + untracked.status = RequestStatus.RUNNING + + sched = _RealignSchedulerStub( + requests={}, # already deleted from self.requests + running=[untracked], + waiting=[], + ) + sched._purge_finished_from_running() + + assert sched.running == [] + + def test_healthy_running_is_left_unchanged(self) -> None: + """Live tracked running requests must not be touched -- the + purge is defensive, not aggressive.""" + alive = _make_request(request_id="req-alive") + alive.status = RequestStatus.RUNNING + + sched = _RealignSchedulerStub( + requests={alive.request_id: alive}, + running=[alive], + waiting=[], + ) + sched._purge_finished_from_running() + + assert sched.running == [alive] + + def test_empty_running_is_noop(self) -> None: + """Empty ``self.running`` must short-circuit cleanly.""" + sched = _RealignSchedulerStub(requests={}, running=[], waiting=[]) + sched._purge_finished_from_running() + + assert sched.running == [] + + def test_mixed_alive_and_dead_keeps_only_alive_in_order(self) -> None: + """Mixed ``self.running`` -- some alive, some finished, some + untracked. Sweep keeps only the alive ones and preserves their + relative order.""" + alive_a = _make_request(request_id="req-alive-a") + alive_a.status = RequestStatus.RUNNING + finished_b = _make_request(request_id="req-finished-b") + finished_b.status = RequestStatus.FINISHED_STOPPED + untracked_c = _make_request(request_id="req-untracked-c") + untracked_c.status = RequestStatus.RUNNING + alive_d = _make_request(request_id="req-alive-d") + alive_d.status = RequestStatus.RUNNING + + sched = _RealignSchedulerStub( + requests={ + alive_a.request_id: alive_a, + finished_b.request_id: finished_b, # tracked but is_finished + # untracked_c is NOT in self.requests + alive_d.request_id: alive_d, + }, + running=[alive_a, finished_b, untracked_c, alive_d], + waiting=[], + ) + sched._purge_finished_from_running() + + assert sched.running == [alive_a, alive_d] + + def test_in_place_mutation_keeps_helper_local_list_identity(self) -> None: + """The slice-assign form (``self.running[:] = ...``) avoids an + extra rebind inside this helper. Note that across a full + ``finish_requests`` call upstream ``Scheduler.finish_requests`` + already rebinds ``self.running``, so global identity is not a + contract -- this test only pins the local helper behavior so a + future refactor away from in-place mutation is intentional. + """ + finished = _make_request(request_id="req-finished") + finished.status = RequestStatus.FINISHED_STOPPED + alive = _make_request(request_id="req-alive") + alive.status = RequestStatus.RUNNING + + sched = _RealignSchedulerStub( + requests={ + finished.request_id: finished, + alive.request_id: alive, + }, + running=[finished, alive], + waiting=[], + ) + held_ref = sched.running + sched._purge_finished_from_running() + + assert held_ref is sched.running + assert held_ref == [alive] diff --git a/tests/core/sched/test_omni_scheduler_finish_requests_purge.py b/tests/core/sched/test_omni_scheduler_finish_requests_purge.py new file mode 100644 index 000000000..9b564d7d5 --- /dev/null +++ b/tests/core/sched/test_omni_scheduler_finish_requests_purge.py @@ -0,0 +1,158 @@ +"""Integration tests for the post-finish defensive purge of +``self.running`` placed in +``OmniARScheduler.finish_requests`` / +``OmniGenerationScheduler.finish_requests``. + +Pins *placement* — that ``_purge_finished_from_running`` runs AFTER +``super().finish_requests`` so it catches entries upstream's removal +missed (e.g. status mid-transition, connector cleanup popping +``self.requests`` without unwinding ``self.running``). Direct helper +unit tests live in ``test_omni_scheduler_mixin.py``; these tests +exercise the full ``finish_requests`` path end-to-end. + +Regression for the residual self.running slot leak surface paired +with ``_realign_request_status_to_queues`` (vllm-omni#3774). +""" + +from __future__ import annotations + +import pytest +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: + """Minimal Request stub with the surface the helpers exercise.""" + + 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): + """Bare scheduler with just the surface ``finish_requests`` reads.""" + scheduler = scheduler_cls.__new__(scheduler_cls) + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = None + scheduler.requests = requests + scheduler.running = running + scheduler.waiting = waiting + return scheduler + + +@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS) +def test_finish_requests_purges_finished_request_from_running( + monkeypatch: pytest.MonkeyPatch, + scheduler_cls, + scheduler_mod, +) -> None: + """A request whose ``status`` is a finished value but still lives + in ``self.running`` after ``super().finish_requests`` returns must + be swept by the post-finish purge. Mocks ``super()`` to simulate + the corner case where upstream's status-driven removal missed the + entry. + """ + finished = _StubRequest("req-finished", RequestStatus.FINISHED_STOPPED) + scheduler = _make_scheduler( + scheduler_cls, + requests={finished.request_id: finished}, + running=[finished], + waiting=[], + ) + + def fake_finish_requests(self, request_ids, finished_status): + # super() returning the abort tuple but leaving self.running untouched + return [(finished.request_id, 0)] + + monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests) + + result = scheduler_cls.finish_requests( + scheduler, + [finished.request_id], + RequestStatus.FINISHED_STOPPED, + ) + + assert scheduler.running == [] + assert result == [(finished.request_id, 0)] + + +@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS) +def test_finish_requests_purges_untracked_request_from_running( + monkeypatch: pytest.MonkeyPatch, + scheduler_cls, + scheduler_mod, +) -> None: + """A request that has been deleted from ``self.requests`` (e.g. + after upstream's ``_free_request`` ran but a status mismatch + caused the ``self.running`` removal branch to no-op) must be + swept by the post-finish purge. + """ + untracked = _StubRequest("req-untracked", RequestStatus.RUNNING) + scheduler = _make_scheduler( + scheduler_cls, + requests={}, # already deleted from self.requests + running=[untracked], + waiting=[], + ) + + def fake_finish_requests(self, request_ids, finished_status): + return [(untracked.request_id, 0)] + + monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests) + + result = scheduler_cls.finish_requests( + scheduler, + [untracked.request_id], + RequestStatus.FINISHED_STOPPED, + ) + + assert scheduler.running == [] + assert result == [(untracked.request_id, 0)] + + +@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS) +def test_finish_requests_leaves_healthy_running_intact( + monkeypatch: pytest.MonkeyPatch, + scheduler_cls, + scheduler_mod, +) -> None: + """A live ``RUNNING`` request that is not in the finish set must + not be swept by the defensive purge -- this pins that the predicate + is conservative and won't aggressively drop healthy entries. + """ + alive = _StubRequest("req-alive", RequestStatus.RUNNING) + leaving_id = "req-leaving" + scheduler = _make_scheduler( + scheduler_cls, + requests={alive.request_id: alive}, + running=[alive], # leaving already removed by super() (typical happy path) + waiting=[], + ) + + def fake_finish_requests(self, request_ids, finished_status): + return [(leaving_id, 0)] + + monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests) + + scheduler_cls.finish_requests( + scheduler, + [leaving_id], + RequestStatus.FINISHED_STOPPED, + ) + + assert scheduler.running == [alive] 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 4088cbf78..fcac8172b 100644 --- a/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py +++ b/tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py @@ -46,6 +46,9 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids( scheduler = scheduler_cls.__new__(scheduler_cls) scheduler.chunk_transfer_adapter = None scheduler.input_coordinator = coordinator + scheduler.requests = {} + scheduler.running = [] + scheduler.waiting = [] def fake_finish_requests(self, request_ids, finished_status): assert request_ids == ["req-a", "req-b"] diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index f8eb6af8b..ab671ec66 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -387,14 +387,15 @@ def test_process_and_restore_queues(build_adapter): running_req = _req("r1", RequestStatus.RUNNING) waiting_queue = DummyWaitingQueue([waiting_req]) running_queue = [running_req] + scheduler_requests = {waiting_req.request_id: waiting_req, running_req.request_id: running_req} - adapter.process_pending_chunks(waiting_queue, running_queue) + adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests) assert waiting_req.status == RequestStatus.WAITING_FOR_CHUNK assert running_req.status == RequestStatus.WAITING_FOR_CHUNK assert waiting_queue == [] assert running_queue == [] - adapter.restore_queues(waiting_queue, running_queue) + adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests) assert waiting_queue == [waiting_req] assert running_queue == [running_req] assert adapter.waiting_for_chunk_waiting_requests == deque() @@ -897,6 +898,8 @@ def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture): sched = OmniARScheduler.__new__(OmniARScheduler) sched.chunk_transfer_adapter = adapter sched.requests = {} + sched.running = [] + sched.waiting = [] with patch.object(VLLMScheduler, "finish_requests", _super_finish): OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED) @@ -1170,3 +1173,164 @@ def test_deferred_finish_not_finished_still_emits_output(mocker: MockerFixture): assert eco.outputs[0].finish_reason == "stop" assert eco.outputs[0].kv_transfer_params is None assert scheduler._pending_finish_reqs == [] + + +# --------------------------------------------------------------- +# Zombie purge tests (regression for vllm-project/vllm-omni#3736) +# --------------------------------------------------------------- + + +def test_process_pending_chunks_purges_zombies_in_running_deque( + build_adapter, +): + """A request aborted while parked in ``waiting_for_chunk_running_requests`` + is no longer tracked by ``scheduler.requests`` once + ``Scheduler._free_request`` runs. ``process_pending_chunks`` must drop + that zombie *before* ``restore_queues`` would re-inject it onto the + running queue. + + Regression for https://github.com/vllm-project/vllm-omni/issues/3736 + (engine-core ``KeyError`` on aborts under chunk-transfer pressure). + """ + adapter, _ = build_adapter(stage_id=1, max_num_seqs=8) + + live_req = _req("live-1", RequestStatus.WAITING_FOR_CHUNK) + zombie_req = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK) + adapter.waiting_for_chunk_running_requests.append(live_req) + adapter.waiting_for_chunk_running_requests.append(zombie_req) + # Mirror state the adapter would carry for an in-flight request so we can + # later assert ``cleanup_receiver`` actually fired against the zombie. + adapter.requests_origin_status[zombie_req.request_id] = RequestStatus.RUNNING + adapter.requests_origin_status[live_req.request_id] = RequestStatus.RUNNING + + waiting_queue = DummyWaitingQueue() + running_queue: list = [] + # Simulate the post-abort scheduler: only ``live-1`` is still tracked. + scheduler_requests = {live_req.request_id: live_req} + + adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests) + + # 1. Zombie was popped from the deque. + assert zombie_req not in adapter.waiting_for_chunk_running_requests + # 2. Live request is still in the deque. + assert live_req in adapter.waiting_for_chunk_running_requests + # 3. ``cleanup_receiver`` ran for the zombie (drops origin-status mapping + # and registers the id as cancelled so a late load/poll is dropped too). + assert zombie_req.request_id not in adapter.requests_origin_status + assert zombie_req.request_id in adapter._cancelled_load_reqs + # 4. Live request's bookkeeping is untouched. + assert adapter.requests_origin_status[live_req.request_id] == RequestStatus.RUNNING + assert live_req.request_id not in adapter._cancelled_load_reqs + + # 5. ``restore_queues`` (which the scheduler runs in its ``finally`` + # clause) now only re-injects the live request -- the zombie is + # gone, so the worker's ``_update_states`` cannot crash on it. + adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests) + assert running_queue == [live_req] + + +def test_process_pending_chunks_purges_zombies_in_waiting_deque(build_adapter): + """Zombie purge applies symmetrically to the waiting-side deque. + + Regression for https://github.com/vllm-project/vllm-omni/issues/3736 -- + aborted requests that landed in ``waiting_for_chunk_waiting_requests`` + must also be removed before ``restore_queues`` re-injects them, since + ``restore_queues`` uses ``waiting_queue.add_request`` rather than + ``running_queue.extend`` for that path. + """ + adapter, _ = build_adapter(stage_id=1, max_num_seqs=8) + + live_req = _req("live-w", RequestStatus.WAITING_FOR_CHUNK) + zombie_req = _req("zombie-w", RequestStatus.WAITING_FOR_CHUNK) + adapter.waiting_for_chunk_waiting_requests.append(live_req) + adapter.waiting_for_chunk_waiting_requests.append(zombie_req) + + waiting_queue = DummyWaitingQueue() + running_queue: list = [] + scheduler_requests = {live_req.request_id: live_req} + + adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests) + + assert live_req in adapter.waiting_for_chunk_waiting_requests + assert zombie_req not in adapter.waiting_for_chunk_waiting_requests + assert zombie_req.request_id in adapter._cancelled_load_reqs + + adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests) + assert waiting_queue == [live_req] + + +def test_purge_preserves_live_order_with_interleaved_zombies(build_adapter): + """Interleaved live/zombie ordering -- pin the ``popleft`` + append-live + in-place filter at chunk_transfer_adapter._purge_untracked_chunk_requests: + survivor order must match insertion order, and ``cleanup_receiver`` must + run exactly once per zombie. + """ + adapter, _ = build_adapter(stage_id=1, max_num_seqs=8) + + live1 = _req("live-1", RequestStatus.WAITING_FOR_CHUNK) + zombie1 = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK) + live2 = _req("live-2", RequestStatus.WAITING_FOR_CHUNK) + zombie2 = _req("zombie-2", RequestStatus.WAITING_FOR_CHUNK) + for req in (live1, zombie1, live2, zombie2): + adapter.waiting_for_chunk_running_requests.append(req) + + waiting_queue = DummyWaitingQueue() + running_queue: list = [] + scheduler_requests = {live1.request_id: live1, live2.request_id: live2} + + adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests) + + assert list(adapter.waiting_for_chunk_running_requests) == [live1, live2] + assert zombie1.request_id in adapter._cancelled_load_reqs + assert zombie2.request_id in adapter._cancelled_load_reqs + + +def test_restore_queues_purges_late_aborts_after_process_pending_chunks( + build_adapter, +): + """Race window guard: an abort can fire between ``process_pending_chunks`` + and the scheduler's ``finally``-clause ``restore_queues`` call. The + second purge inside ``restore_queues`` must drop the now-untracked + request so it does not get re-injected onto ``running_queue``. + """ + adapter, _ = build_adapter(stage_id=1, max_num_seqs=8) + + req = _req("late-abort", RequestStatus.WAITING_FOR_CHUNK) + adapter.waiting_for_chunk_running_requests.append(req) + + waiting_queue = DummyWaitingQueue() + running_queue: list = [] + + # Tick N: process_pending_chunks runs while req is still tracked. + scheduler_requests: dict = {req.request_id: req} + adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests) + assert req in adapter.waiting_for_chunk_running_requests + + # Mid-tick: abort fires and the scheduler's free path deletes the entry. + del scheduler_requests[req.request_id] + + # finally: restore_queues sees the now-untracked req and must drop it + # instead of blindly extending it onto running_queue. + adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests) + assert running_queue == [] + assert req.request_id in adapter._cancelled_load_reqs + + +def test_purge_is_noop_on_empty_deques(build_adapter): + """Empty deques short-circuit -- guards against any accidental + ``IndexError`` from ``popleft`` on an empty deque if a future caller + shadows the empty check. + """ + adapter, _ = build_adapter(stage_id=1, max_num_seqs=8) + assert len(adapter.waiting_for_chunk_waiting_requests) == 0 + assert len(adapter.waiting_for_chunk_running_requests) == 0 + + waiting_queue = DummyWaitingQueue() + running_queue: list = [] + + adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests={}) + assert len(adapter.waiting_for_chunk_waiting_requests) == 0 + assert len(adapter.waiting_for_chunk_running_requests) == 0 + adapter.restore_queues(waiting_queue, running_queue, scheduler_requests={}) + assert running_queue == [] + assert waiting_queue == [] diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 67514994e..4c4ae361a 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -218,7 +218,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): self._process_pending_input_timeouts() if self.chunk_transfer_adapter: - self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running) + self.chunk_transfer_adapter.process_pending_chunks( + self.waiting, self.running, scheduler_requests=self.requests + ) try: scheduler_output = super().schedule() @@ -626,7 +628,34 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if self.chunk_transfer_adapter: self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests) + # Realign stale ``request.status`` (chunk-transfer-adapter's + # ``requests_origin_status`` table doesn't follow the + # ``waiting → running`` admit transition; without this, an abort + # arriving between admit and the next deque round-trip leaves + # the request in ``self.running`` with ``status=WAITING`` and + # upstream ``Scheduler.finish_requests`` silently fails to + # release the worker's ``input_batch`` slot -- after + # ``max_num_seqs`` such aborts new requests hang at + # ``chunks=0``). Only the ``async_chunk`` path triggers the + # staleness; with ``async_chunk`` disabled this is a cheap O(n) + # no-op over an already-aligned set, kept unconditional so the + # abort path stays uniform across configurations. See + # ``OmniSchedulerMixin._realign_request_status_to_queues`` and + # #3774 discussion. + self._realign_request_status_to_queues(request_ids) + finished = super().finish_requests(request_ids, finished_status) + + # Defensive post-finish purge: belt-and-suspenders to the + # realignment above. Even after realign + ``super()``, corner + # cases (mid-transition status, connector cleanups that pop + # from ``self.requests`` without unwinding ``self.running``) + # can leave already-finished or untracked entries in + # ``self.running``. Sweep them now so the worker's + # ``input_batch`` slot never pins a freed request and starves + # new admissions. See ``OmniSchedulerMixin._purge_finished_from_running``. + self._purge_finished_from_running() + input_coordinator = getattr(self, "input_coordinator", None) if input_coordinator is not None: for request_id, _ in finished: diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 1c43b958a..9f5d3e104 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -87,7 +87,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): self._consume_pending_connector_output(model_mode="generation") self._process_pending_input_timeouts() if self.chunk_transfer_adapter: - self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running) + self.chunk_transfer_adapter.process_pending_chunks( + self.waiting, self.running, scheduler_requests=self.requests + ) # OMNI: Track requests that are already finished (e.g., marked by connector) # These should be removed from running and not scheduled @@ -364,7 +366,23 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): if self.chunk_transfer_adapter: self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests) + # See ``OmniSchedulerMixin._realign_request_status_to_queues`` -- + # closes the residual hang on the + # ``waiting → running → abort-before-next-deque-round-trip`` race + # surfaced by #3774's reproduction matrix. Only the + # ``async_chunk`` path actually triggers the staleness; with + # ``async_chunk`` disabled this is a cheap O(n) no-op, kept + # unconditional so the abort path stays uniform. + self._realign_request_status_to_queues(request_ids) + finished = super().finish_requests(request_ids, finished_status) + + # See ``OmniSchedulerMixin._purge_finished_from_running`` -- + # defensive belt-and-suspenders sweep paired with the realign + # above. Closes residual ``self.running`` slot leaks if any + # corner case slips past upstream's status-driven removal. + self._purge_finished_from_running() + if self.input_coordinator is not None: for request_id, _ in finished: self._free_input_coordinator_request(request_id) diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index c33d16dac..165c04e96 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/vllm_omni/core/sched/omni_scheduler_mixin.py @@ -2,6 +2,7 @@ from __future__ import annotations import os import time +from collections.abc import Iterable from typing import Any from vllm.logger import init_logger @@ -156,3 +157,122 @@ class OmniSchedulerMixin: return None self._last_stats_time = now return super().make_stats(*args, **kwargs) + + def _realign_request_status_to_queues( + self, + request_ids: str | Iterable[str] | None, + ) -> None: + """Realign ``request.status`` to actual queue membership. + + ``OmniChunkTransferAdapter._process_chunk_queue`` stamps + ``requests_origin_status[req.id] = WAITING`` (or ``RUNNING``) when + first parking a request in a chunk-transfer deque. On the next + tick, when the chunk arrives, ``_process_chunk_queue`` sets + ``request.status = target_status`` and continues, but + ``requests_origin_status`` is left at its first-park value -- no + hook updates it on the ``waiting → running`` admit transition + that ``super().schedule()`` later performs. The table stays + stale until the request makes another deque round-trip. + + If an abort lands in the gap between admit and the next deque + round-trip, ``chunk_transfer_adapter.finish_requests`` reads the + stale ``WAITING`` from ``requests_origin_status``, stomps it + onto ``request.status``, and the upstream + ``Scheduler.finish_requests`` else branch silently fails to + remove from ``self.running`` -- the request stays alive in + ``self.running`` and the worker's ``input_batch`` slot leaks. + After ``max_num_seqs`` such aborts every new request hangs at + ``chunks=0`` until the client times out. + + Realign here: if a request lives in ``self.running`` but its + status is not ``RUNNING``, set it to ``RUNNING``; symmetrically + flip ``RUNNING → WAITING`` when the request is actually in + ``self.waiting``. This is a localized safety net for + ``requests_origin_status`` staleness on the admit transition; + it does not touch the adapter's invariants and is complementary + to the chunk-transfer-adapter deque purge that already runs + inside ``process_pending_chunks`` / ``restore_queues``. + + Note on scope: only the ``async_chunk`` path actually triggers + the ``requests_origin_status`` staleness this helper repairs. + When ``async_chunk`` is disabled, no chunk-transfer round-trip + occurs between admit and finish, so the realignment walk is a + cheap O(n) no-op over an already-aligned set. The call is kept + unconditional in ``finish_requests`` to (a) keep the abort path + uniform and (b) defend any future configuration that re-enables + chunk transfer from rediscovering the same regression. + + See https://github.com/vllm-project/vllm-omni/pull/3774 and the + residual-hang reproduction discussed in that PR. + """ + # Mirror the upstream Scheduler.finish_requests resolution of + # ``request_ids`` so realignment touches exactly the set that + # ``super().finish_requests`` will then walk. + if isinstance(request_ids, str): + ids_to_align: Iterable[str] = (request_ids,) + elif request_ids is None: + ids_to_align = list(self.requests.keys()) + else: + ids_to_align = list(request_ids) + + if not ids_to_align: + return + + running_ids = {r.request_id for r in self.running} + waiting_ids = {r.request_id for r in self.waiting} + + for rid in ids_to_align: + req = self.requests.get(rid) + if req is None or req.is_finished(): + continue + if rid in running_ids and req.status != RequestStatus.RUNNING: + req.status = RequestStatus.RUNNING + elif rid in waiting_ids and req.status == RequestStatus.RUNNING: + req.status = RequestStatus.WAITING + + def _purge_finished_from_running(self) -> None: + """Defensive post-finish sweep of ``self.running``. + + Belt-and-suspenders to ``_realign_request_status_to_queues``: + even after status realignment lets upstream + ``Scheduler.finish_requests`` pick the right removal branch, + a future regression or an unexpected ``status`` mid-transition + could still leave already-finished entries in ``self.running``. + Sweeping here guarantees the worker's ``input_batch`` slot is + not pinned by a freed request. + + Complementary to ``_realign_request_status_to_queues``: realign + is preventive (fix ``status`` before ``super().finish_requests`` + so the right branch fires); this purge is defensive (sweep the + residue after ``super().finish_requests`` so any stale entries + are reclaimed). + + Scope of the predicate. ``is_finished()`` covers entries the + upstream ``finish_requests`` already drained from ``self.requests`` + but failed to remove from ``self.running``; the + ``request_id not in self.requests`` arm catches the same surface + from a different angle and is the post-cleanup mirror of the + deque purge ``_purge_untracked_chunk_requests`` already runs at + the chunk-transfer-adapter layer. It does **not** by itself make + arbitrary direct deletions of ``self.requests`` safe -- callers + that pop ``self.requests`` outside the standard finish path + still have to go through ``_free_request`` (or equivalent) for + block / connector / coordinator cleanup. This sweep only + reclaims the ``self.running`` slot reference. + + In-place via ``self.running[:] = ...`` for minor consistency + with idiomatic vLLM scheduler mutation; upstream + ``Scheduler.finish_requests`` itself rebinds ``self.running``, + so list identity across the whole call is not preserved -- the + slice form is just to avoid an extra rebind inside this helper. + + Assumes the upstream V1 invariant that scheduler ticks are + serialized on a single thread; in-place mutation here is no more + racy than the rest of the scheduler under that assumption. + + See https://github.com/vllm-project/vllm-omni/pull/3774 + discussion. + """ + if not self.running: + return + self.running[:] = [req for req in self.running if not req.is_finished() and req.request_id in self.requests] 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 121fb79ec..eb1c3eec6 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 @@ -444,12 +444,37 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self, waiting_queue: Any, running_queue: list[Request], + *, + scheduler_requests: dict[str, Request] | None = None, ) -> None: """ Process pending chunks for waiting and running queues. + + When ``scheduler_requests`` is provided, purges any + ``waiting_for_chunk_*_requests`` deque entries whose + ``request_id`` is no longer tracked by it (e.g. after a + mid-flight abort that ran ``Scheduler._free_request``) before + processing chunks. Without this purge, ``restore_queues`` would + later re-inject the freed ``Request`` onto ``running_queue`` and + the worker's ``_update_states`` would crash with ``KeyError`` + reading ``self.requests[req_id]``. See vllm-project/vllm-omni#3736. + + ``scheduler_requests`` is keyword-only and optional; production + schedulers always pass their live request map, while legacy + callers that don't track aborts may omit it to keep the prior + (unguarded) behaviour. """ if self.connector.stage_id == 0: return + + # Purge deque entries whose request was freed mid-flight (abort → + # Scheduler._free_request) before any chunk processing, so neither + # the legacy nor the active-stream path can re-inject a zombie + # Request onto the queues. See vllm-project/vllm-omni#3736. + if scheduler_requests is not None: + self._purge_untracked_chunk_requests(self.waiting_for_chunk_waiting_requests, scheduler_requests) + self._purge_untracked_chunk_requests(self.waiting_for_chunk_running_requests, scheduler_requests) + if self._active_window <= 0: self._process_chunk_queue_legacy( waiting_queue, self.waiting_for_chunk_waiting_requests, RequestStatus.WAITING, self._finished_load_reqs @@ -558,6 +583,29 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.requests_origin_status[request.request_id] = target_status waiting_for_chunk_list.append(request) + def _purge_untracked_chunk_requests( + self, + deque_list: deque[Any], + scheduler_requests: dict[str, Request], + ) -> None: + """Drop deque entries whose ``request_id`` is not in + ``scheduler_requests`` and reclaim their receiver-side state. + + Handles requests that were aborted mid-flight while parked in a + chunk-transfer deque: ``Scheduler._free_request`` deleted the + entry from ``scheduler.requests`` but the deque still holds a + reference to the now-freed ``Request``. Order of survivors is + preserved. + """ + if not deque_list: + return + for _ in range(len(deque_list)): + request = deque_list.popleft() + if request.request_id in scheduler_requests: + deque_list.append(request) + else: + self.cleanup_receiver(request.request_id) + def restore_queues( self, waiting_queue: Any, @@ -566,7 +614,23 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): ) -> None: """ Restore requests waiting for chunk to the waiting and running queues. + + Re-runs the zombie purge first to close the race window where an + abort fires *between* ``process_pending_chunks`` and the + ``finally``-clause ``restore_queues`` call. Without the second + purge, ``running_queue.extend(...)`` would still re-inject a + freed ``Request`` and crash the worker on the next tick. + + ``scheduler_requests`` is optional for back-compat with legacy + callers (older tests pass only the two queue arguments). When + provided, it gates both the deque purge and the per-request + admit checks below; when ``None``, the purge is skipped and + every parked request is restored unconditionally (the + pre-purge behavior). """ + if scheduler_requests is not None: + self._purge_untracked_chunk_requests(self.waiting_for_chunk_waiting_requests, scheduler_requests) + self._purge_untracked_chunk_requests(self.waiting_for_chunk_running_requests, scheduler_requests) # Add request waiting for chunk to the waiting and running queue for request in self.waiting_for_chunk_waiting_requests: if scheduler_requests is None or request.request_id in scheduler_requests: