[Bugfix] Reclaim resumable async-chunk requests on finish (#6360)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: Ruirui Yang | Rein <73573651+R2-Y@users.noreply.github.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/core/sched/test_omni_generation_scheduler_update_session.py

@@ -410,12 +410,13 @@ class TestRealignRequestStatusToQueues:
         assert stale.status == RequestStatus.RUNNING
         assert clean.status == RequestStatus.RUNNING
 
-    def test_finished_request_is_skipped(self) -> None:
+    def test_non_resumable_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.resumable = False
         req.status = RequestStatus.FINISHED_STOPPED
 
         sched = _RealignSchedulerStub(

tests/core/sched/test_omni_scheduler_finish_requests_purge.py

@@ -157,3 +157,71 @@ def test_finish_requests_leaves_healthy_running_intact(
     )
 
     assert scheduler.running == [alive]
+
+
+@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
+def test_finish_requests_preserves_unrelated_resumable_segment(
+    monkeypatch: pytest.MonkeyPatch,
+    scheduler_cls,
+    scheduler_mod,
+) -> None:
+    paused = _StubRequest("req-paused", RequestStatus.FINISHED_STOPPED)
+    paused.resumable = True
+    leaving = _StubRequest("req-leaving", RequestStatus.FINISHED_STOPPED)
+    scheduler = _make_scheduler(
+        scheduler_cls,
+        requests={paused.request_id: paused},
+        running=[paused],
+        waiting=[],
+    )
+
+    def fake_finish_requests(self, request_ids, finished_status):
+        return [leaving]
+
+    monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
+
+    scheduler_cls.finish_requests(scheduler, [leaving.request_id], RequestStatus.FINISHED_ABORTED)
+
+    assert scheduler.running == [paused]
+    assert scheduler.requests == {paused.request_id: paused}
+
+
+@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
+def test_finish_requests_materializes_single_pass_request_ids(
+    monkeypatch: pytest.MonkeyPatch,
+    scheduler_cls,
+    scheduler_mod,
+) -> None:
+    request = _StubRequest("req-generator", RequestStatus.RUNNING)
+    scheduler = _make_scheduler(
+        scheduler_cls,
+        requests={request.request_id: request},
+        running=[request],
+        waiting=[],
+    )
+
+    def fake_finish_requests(self, request_ids, finished_status):
+        assert tuple(request_ids) == (request.request_id,)
+        return [request]
+
+    monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
+
+    result = scheduler_cls.finish_requests(
+        scheduler,
+        (request_id for request_id in [request.request_id]),
+        RequestStatus.FINISHED_ABORTED,
+    )
+
+    assert result == [request]
+
+
+@pytest.mark.parametrize("scheduler_cls", [OmniARScheduler, OmniGenerationScheduler])
+def test_finish_requests_does_not_reopen_off_queue_deferred_free_terminal(scheduler_cls) -> None:
+    terminal = _StubRequest("req-deferred-free", RequestStatus.FINISHED_STOPPED)
+    terminal.resumable = True
+    scheduler = _make_scheduler(scheduler_cls, requests={terminal.request_id: terminal}, running=[], waiting=[])
+    scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed")
+
+    assert scheduler_cls.finish_requests(scheduler, [terminal.request_id], RequestStatus.FINISHED_ABORTED) == []
+    assert scheduler.requests == {terminal.request_id: terminal}
+    assert terminal.status == RequestStatus.FINISHED_STOPPED

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -3,8 +3,8 @@
 
 import threading
 from collections import deque
-from types import SimpleNamespace
-from unittest.mock import patch
+from types import MethodType, SimpleNamespace
+from unittest.mock import Mock, patch
 
 import pytest
 import torch
@@ -13,6 +13,7 @@ from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.metrics.stats import PrefillStats, PromptTokenStats
 from vllm.v1.request import RequestStatus
 
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
 from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
 from vllm_omni.data_entry_keys import CodesStruct, MetaStruct, OmniPayload, OmniPayloadStruct
 from vllm_omni.distributed.omni_connectors.adapter import construct_next_stage_streaming_input_prompt
@@ -32,9 +33,14 @@ class DummyWaitingQueue(list):
     def add_request(self, request):
         self.append(request)
 
+    def remove_requests(self, requests):
+        remove = set(requests)
+        self[:] = [request for request in self if request not in remove]
+
 
 def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None):
-    return SimpleNamespace(
+    request = Mock(
+        client_index=0,
         request_id=req_id,
         external_req_id=external_req_id or req_id,
         status=status,
@@ -44,8 +50,10 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None)
         num_output_placeholders=0,
         prefill_stats=None,
         additional_information=None,
-        is_finished=lambda: status == RequestStatus.FINISHED_STOPPED,
+        resumable=False,
     )
+    request.is_finished = lambda: RequestStatus.is_finished(request.status)
+    return request
 
 
 def test_streaming_payload_can_replace_placeholder_prompt(mocker: MockerFixture) -> None:
@@ -1006,12 +1014,30 @@ def test_finish_requests_restores_status(build_adapter):
     prior = RequestStatus.RUNNING
     request = _req(req_id, RequestStatus.WAITING_FOR_CHUNK)
     adapter.requests_origin_status[req_id] = prior
+    adapter.waiting_for_chunk_running_requests.append(request)
     requests_map = {req_id: request}
 
     adapter.finish_requests([req_id], RequestStatus.FINISHED_ABORTED, requests_map)
 
     assert request.status == prior
     assert req_id not in adapter.requests_origin_status
+    assert not adapter.waiting_for_chunk_running_requests
+
+
+def test_finish_requests_does_not_restore_stale_status_without_connector_ownership(build_adapter):
+    adapter, _ = build_adapter(stage_id=1)
+    request = _req("req-restored", RequestStatus.WAITING_FOR_STREAMING_REQ)
+    request.resumable = True
+    adapter.requests_origin_status[request.request_id] = RequestStatus.RUNNING
+
+    adapter.finish_requests(
+        [request.request_id],
+        RequestStatus.FINISHED_ABORTED,
+        {request.request_id: request},
+    )
+
+    assert request.status == RequestStatus.WAITING_FOR_STREAMING_REQ
+    assert request.request_id not in adapter.requests_origin_status
 
 
 def test_finish_requests_removes_zombies_from_chunk_waiting_deques(build_adapter):
@@ -1061,6 +1087,108 @@ def test_finish_requests_releases_active_stream_slot(build_adapter):
     assert waiting.status == RequestStatus.WAITING_FOR_CHUNK
 
 
+@pytest.mark.parametrize("scheduler_cls", [OmniGenerationScheduler, OmniARScheduler])
+@pytest.mark.parametrize(
+    ("placement", "origin_status", "initial_status", "streaming_counter_owned"),
+    [
+        ("hidden", RequestStatus.RUNNING, RequestStatus.FINISHED_STOPPED, False),
+        ("running", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, False),
+        ("waiting", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, False),
+        ("skipped", RequestStatus.WAITING, RequestStatus.FINISHED_STOPPED, True),
+        ("skipped_streaming", RequestStatus.RUNNING, RequestStatus.WAITING_FOR_STREAMING_REQ, True),
+        ("skipped_hidden", RequestStatus.RUNNING, RequestStatus.FINISHED_STOPPED, True),
+        ("skipped_non_streaming", RequestStatus.WAITING, RequestStatus.WAITING, False),
+    ],
+)
+def test_finish_requests_reclaims_resumable_segment_and_reuses_capacity(
+    build_adapter,
+    scheduler_cls,
+    placement,
+    origin_status,
+    initial_status,
+    streaming_counter_owned,
+):
+    adapter, _ = build_adapter(stage_id=1, active_stream_window=2)
+    request = _req(
+        "req-segment-stop",
+        initial_status,
+        external_req_id="ext-segment-stop",
+    )
+    request.resumable = True
+    adapter.requests_origin_status[request.request_id] = origin_status
+    if placement in {"hidden", "skipped_hidden"}:
+        adapter.waiting_for_chunk_running_requests.append(request)
+    adapter._active_streams[request.request_id] = request
+
+    scheduler = scheduler_cls.__new__(scheduler_cls)
+    scheduler.chunk_transfer_adapter = adapter
+    scheduler.input_coordinator = None
+    scheduler.requests = {request.request_id: request}
+    scheduler.running = [request] if placement == "running" else []
+    scheduler.waiting = DummyWaitingQueue([request] if placement == "waiting" else [])
+    scheduler.skipped_waiting = DummyWaitingQueue([request] if placement.startswith("skipped") else [])
+    scheduler.num_waiting_for_streaming_input = int(streaming_counter_owned)
+
+    freed = []
+
+    def fake_free_request(self, live, delay_free_blocks=False):
+        del delay_free_blocks
+        freed.append(live.request_id)
+        self.requests.pop(live.request_id)
+        return None, None
+
+    scheduler._free_request = MethodType(fake_free_request, scheduler)
+
+    first = scheduler_cls.finish_requests(scheduler, [request.request_id], RequestStatus.FINISHED_ABORTED)
+    second = scheduler_cls.finish_requests(scheduler, [request.request_id], RequestStatus.FINISHED_ABORTED)
+
+    assert len(first) == 1
+    assert second == []
+    assert freed == [request.request_id]
+    assert request.request_id not in scheduler.requests
+    assert request.request_id not in adapter._active_streams
+    assert not adapter.waiting_for_chunk_running_requests
+    assert scheduler.running == []
+    assert list(scheduler.waiting) == []
+    assert list(scheduler.skipped_waiting) == []
+    assert scheduler.num_waiting_for_streaming_input == 0
+
+    fresh_a = _req("req-fresh-a", RequestStatus.WAITING)
+    fresh_b = _req("req-fresh-b", RequestStatus.WAITING)
+    assert adapter._ensure_active_stream(fresh_a)
+    assert adapter._ensure_active_stream(fresh_b)
+    assert set(adapter._active_streams) == {fresh_a.request_id, fresh_b.request_id}
+
+
+@pytest.mark.parametrize("scheduler_cls", [OmniGenerationScheduler, OmniARScheduler])
+def test_finish_requests_does_not_reopen_off_queue_resumable_segment(build_adapter, scheduler_cls):
+    adapter, _ = build_adapter(stage_id=1, active_stream_window=2)
+    request = _req("req-off-queue", RequestStatus.FINISHED_STOPPED)
+    request.resumable = True
+    adapter.requests_origin_status[request.request_id] = RequestStatus.RUNNING
+
+    scheduler = scheduler_cls.__new__(scheduler_cls)
+    scheduler.chunk_transfer_adapter = adapter
+    scheduler.input_coordinator = None
+    scheduler.requests = {request.request_id: request}
+    scheduler.running = []
+    scheduler.waiting = DummyWaitingQueue()
+    scheduler.skipped_waiting = DummyWaitingQueue()
+    scheduler._free_request = lambda *args, **kwargs: pytest.fail("off-queue request was freed")
+
+    assert (
+        scheduler_cls.finish_requests(
+            scheduler,
+            [request.request_id],
+            RequestStatus.FINISHED_ABORTED,
+        )
+        == []
+    )
+    assert scheduler.requests == {request.request_id: request}
+    assert request.status == RequestStatus.FINISHED_STOPPED
+    assert request.request_id not in adapter.requests_origin_status
+
+
 def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_adapter):
     adapter, _ = build_adapter(stage_id=1)
     zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import math
 import os
 import time
-from collections.abc import Iterable
+from collections.abc import Iterable, Iterator
 from typing import Any
 
 import torch
@@ -571,12 +571,35 @@ class OmniSchedulerMixin:
         finished_status: RequestStatus,
     ) -> list[Request]:
         """Finish requests and clean all Omni-owned queue/coordinator state."""
+        if isinstance(request_ids, str):
+            target_request_ids = {request_ids}
+        elif request_ids is None:
+            target_request_ids = set(self.requests)
+        else:
+            if isinstance(request_ids, Iterator):
+                request_ids = tuple(request_ids)
+            target_request_ids = set(request_ids)
+
+        skipped_waiting_ids = {r.request_id for r in getattr(self, "skipped_waiting", ())}
+        pre_adapter_streaming_wait_ids = {
+            rid
+            for rid in target_request_ids & skipped_waiting_ids
+            if (req := self.requests.get(rid)) is not None
+            and (
+                req.status == RequestStatus.WAITING_FOR_STREAMING_REQ
+                or (getattr(req, "resumable", False) and req.status == RequestStatus.FINISHED_STOPPED)
+            )
+        }
+
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
 
-        self._realign_request_status_to_queues(request_ids)
+        self._realign_request_status_to_queues(
+            request_ids,
+            pre_adapter_streaming_wait_ids=pre_adapter_streaming_wait_ids,
+        )
         finished = super().finish_requests(request_ids, finished_status)
-        self._purge_finished_from_running()
+        self._purge_finished_from_running(target_request_ids)
 
         for request in finished:
             self._free_input_coordinator_request(request.request_id)
@@ -592,6 +615,8 @@ class OmniSchedulerMixin:
     def _realign_request_status_to_queues(
         self,
         request_ids: str | Iterable[str] | None,
+        *,
+        pre_adapter_streaming_wait_ids: set[str] | None = None,
     ) -> None:
         """Realign ``request.status`` to actual queue membership.
 
@@ -618,7 +643,13 @@ class OmniSchedulerMixin:
         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
+        ``self.waiting``. A resumable segment stop still held in
+        ``self.skipped_waiting`` is restored to
+        ``WAITING_FOR_STREAMING_REQ`` so upstream also balances its
+        paused-session counter. Because adapter cleanup may restore a
+        connector-owned request's prior status first, ``finish_requests``
+        snapshots these counter-bearing rows before invoking the adapter.
+        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
@@ -651,17 +682,36 @@ class OmniSchedulerMixin:
 
         running_ids = {r.request_id for r in self.running}
         waiting_ids = {r.request_id for r in self.waiting}
+        skipped_waiting_ids = {r.request_id for r in getattr(self, "skipped_waiting", ())}
 
         for rid in ids_to_align:
             req = self.requests.get(rid)
-            if req is None or req.is_finished():
+            if req is None:
                 continue
-            if rid in running_ids and req.status != RequestStatus.RUNNING:
+            # A persistent Session may be closed after its current segment
+            # reached FINISHED_STOPPED but while it is still owned by a live
+            # scheduler/connector queue. vLLM skips already-finished requests
+            # in finish_requests(), leaking the KV and worker slot. Only
+            # recover requests with positive queue ownership; an off-queue
+            # terminal can legitimately be waiting for deferred block free.
+            resumable_segment_stop = bool(
+                getattr(req, "resumable", False) and req.status == RequestStatus.FINISHED_STOPPED
+            )
+            streaming_wait = (
+                rid in pre_adapter_streaming_wait_ids
+                if pre_adapter_streaming_wait_ids is not None
+                else resumable_segment_stop
+            )
+            if req.is_finished() and not (resumable_segment_stop or streaming_wait):
+                continue
+            if rid in skipped_waiting_ids and streaming_wait:
+                req.status = RequestStatus.WAITING_FOR_STREAMING_REQ
+            elif rid in running_ids and req.status != RequestStatus.RUNNING:
                 req.status = RequestStatus.RUNNING
-            elif rid in waiting_ids and req.status == RequestStatus.RUNNING:
+            elif rid in waiting_ids and (req.status == RequestStatus.RUNNING or resumable_segment_stop):
                 req.status = RequestStatus.WAITING
 
-    def _purge_finished_from_running(self) -> None:
+    def _purge_finished_from_running(self, target_request_ids: set[str] | None = None) -> None:
         """Defensive post-finish sweep of ``self.running``.
 
         Belt-and-suspenders to ``_realign_request_status_to_queues``:
@@ -678,18 +728,14 @@ class OmniSchedulerMixin:
         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.
+        A resumable ``FINISHED_STOPPED`` request may legitimately remain
+        in ``self.running`` between realtime segments. Preserve such a
+        request unless it belongs to this finish call. Other finished or
+        untracked entries are stale and can be swept defensively.
+
+        When ``target_request_ids`` is ``None`` or empty, every resumable
+        ``FINISHED_STOPPED`` entry is preserved. Production
+        ``finish_requests`` always passes its resolved finish set.
 
         In-place via ``self.running[:] = ...`` for minor consistency
         with idiomatic vLLM scheduler mutation; upstream
@@ -706,4 +752,16 @@ class OmniSchedulerMixin:
         """
         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]
+        target_request_ids = target_request_ids or set()
+
+        def keep_running(req: Request) -> bool:
+            if req.request_id not in self.requests:
+                return False
+            if not req.is_finished():
+                return True
+            resumable_segment_stop = bool(
+                getattr(req, "resumable", False) and req.status == RequestStatus.FINISHED_STOPPED
+            )
+            return resumable_segment_stop and req.request_id not in target_request_ids
+
+        self.running[:] = [req for req in self.running if keep_running(req)]

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -869,13 +869,30 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         else:
             request_ids = requests.keys()
 
+        connector_owned_ids = {
+            request.request_id
+            for queue in (
+                self.waiting_for_chunk_waiting_requests,
+                self.waiting_for_chunk_running_requests,
+                self._held_non_active,
+            )
+            for request in queue
+        }
+
         # First pass: collect requests to remove from queues
         for req_id in request_ids:
             request = requests.get(req_id) if requests else None
-            if request is None or request.is_finished():
+            if request is None:
                 # Invalid request ID.
                 continue
-            if req_id in self.requests_origin_status:
+            resumable_segment_stop = bool(
+                getattr(request, "resumable", False) and request.status == RequestStatus.FINISHED_STOPPED
+            )
+            if request.is_finished() and not resumable_segment_stop:
+                continue
+            # Once restored to a scheduler queue, the saved origin is stale and
+            # must not overwrite statuses such as WAITING_FOR_STREAMING_REQ.
+            if req_id in self.requests_origin_status and req_id in connector_owned_ids:
                 request.status = self.requests_origin_status.pop(req_id)
 
         request_ids = set(request_ids)
@@ -891,11 +908,6 @@ 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.upstream_exhausted_requests.discard(req_id)
-            self._finished_load_reqs.discard(req_id)
-            self._cancelled_load_reqs.add(req_id)
-            self._waiting_since.pop(req_id, None)
+            self.cleanup_receiver(req_id)
 
         return []