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 0f654715a..247f0ce84 100644 --- a/tests/core/sched/test_omni_generation_scheduler_update_session.py +++ b/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( diff --git a/tests/core/sched/test_omni_scheduler_finish_requests_purge.py b/tests/core/sched/test_omni_scheduler_finish_requests_purge.py index fbfeb384f..770c85578 100644 --- a/tests/core/sched/test_omni_scheduler_finish_requests_purge.py +++ b/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 diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 063751b01..fcfd284bf 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/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) diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 7ce5f4470..05ab26e06 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/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)] 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 ae1cb74b9..4a5abcbe5 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 @@ -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 []