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

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]

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]

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"]

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 == []

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:

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)

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]

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: