diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 7a3caba11..256e3e0a3 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -4,10 +4,12 @@ import threading from collections import deque from types import SimpleNamespace +from unittest.mock import patch import pytest import torch from pytest_mock import MockerFixture +from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.request import RequestStatus from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase @@ -335,6 +337,27 @@ def test_cleanup_after_poll_flow(build_adapter): assert "ext-flow" not in adapter.request_payload +def test_finish_requests_restores_status(build_adapter): + """Abort path must pop ``requests_origin_status`` and restore pre-wait status. + + While ``process_pending_chunks`` holds a request off the scheduler queues, the + adapter records the prior status (WAITING or RUNNING). ``finish_requests`` must + put that status back on the live ``Request`` so base ``Scheduler.finish_requests`` + can finish bookkeeping without inconsistent state / crashes. + """ + adapter, _ = build_adapter(stage_id=1) + req_id = "req-abort-during-chunk" + prior = RequestStatus.RUNNING + request = _req(req_id, RequestStatus.WAITING_FOR_CHUNK) + adapter.requests_origin_status[req_id] = prior + 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 + + # --------------------------------------------------------------- # Scheduler trigger tests # --------------------------------------------------------------- @@ -508,3 +531,31 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF assert len(cleanup_calls) == 0 assert len(save_calls) == 1 + + +def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture): + """``OmniARScheduler.finish_requests`` must run chunk adapter hook before vLLM base.""" + from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + + order: list[str] = [] + + adapter = mocker.MagicMock() + + def _adapter_finish(request_ids, finished_status, requests): + order.append("adapter") + return [] + + adapter.finish_requests.side_effect = _adapter_finish + + def _super_finish(_self, request_ids, finished_status): + order.append("super") + return [] + + sched = OmniARScheduler.__new__(OmniARScheduler) + sched.chunk_transfer_adapter = adapter + sched.requests = {} + + with patch.object(VLLMScheduler, "finish_requests", _super_finish): + OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED) + + assert order == ["adapter", "super"] diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 0ee8cd16a..c5031ad11 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -539,6 +539,24 @@ class OmniARScheduler(VLLMScheduler): return engine_core_outputs + def finish_requests(self, request_ids: Any, finished_status: RequestStatus) -> list[tuple[str, int]]: + """Handles the finish signal from outside the scheduler. + + For example, the API server can abort a request when the client + disconnects. + + If request_ids is None, all requests will be finished. + + Returns: + Tuple of (req_id, client_index) for requests that were aborted. Will not + include any that were already finished. + """ + + if self.chunk_transfer_adapter: + self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests) + + return super().finish_requests(request_ids, finished_status) + def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None: # TODO(wzliu)! for offline mode, we should not end process until all data is transferred """Mark a request as finished and free its resources.""" diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 1c4356d4f..87f2d282a 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -324,6 +324,24 @@ class OmniGenerationScheduler(VLLMScheduler): return scheduler_output + def finish_requests(self, request_ids, finished_status: RequestStatus) -> list[tuple[str, int]]: + """Handles the finish signal from outside the scheduler. + + For example, the API server can abort a request when the client + disconnects. + + If request_ids is None, all requests will be finished. + + Returns: + Tuple of (req_id, client_index) for requests that were aborted. Will not + include any that were already finished. + """ + + if self.chunk_transfer_adapter: + self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests) + + return super().finish_requests(request_ids, finished_status) + """ Scheduler for the diffusion model. This scheduler is modified to stop the request immediately for the diffusion model. 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 393d0e801..1d5fc398a 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 @@ -58,6 +58,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.waiting_for_chunk_waiting_requests: deque[Any] = deque() self.waiting_for_chunk_running_requests: deque[Any] = deque() self.requests_with_ready_chunks = set() + self.requests_origin_status = {} @classmethod def create_connector(cls, model_config: Any): @@ -279,6 +280,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.get_req_chunk.pop(request_id, None) self.requests_with_ready_chunks.discard(request_id) self.request_ids_mapping.pop(request_id, None) + self.requests_origin_status.pop(request_id, None) self._cancelled_load_reqs.add(request_id) self._finished_load_reqs.discard(request_id) @@ -408,6 +410,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.requests_with_ready_chunks.add(request.request_id) continue queue.remove(request) + self.requests_origin_status[request.request_id] = target_status waiting_for_chunk_list.append(request) def _clear_chunk_ready(self, scheduler_output: Any) -> None: @@ -420,3 +423,23 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): for req_id in scheduler_output.scheduled_cached_reqs.req_ids: if req_id in self.requests_with_ready_chunks: self.requests_with_ready_chunks.remove(req_id) + + def finish_requests( + self, request_ids: Any, finished_status: RequestStatus, requests: dict[str, Request] | None = None + ) -> list[tuple[str, int]]: + assert RequestStatus.is_finished(finished_status) + if isinstance(request_ids, str): + request_ids = (request_ids,) + elif request_ids is not None: + request_ids = set(request_ids) + else: + request_ids = requests.keys() + + # 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(): + # Invalid request ID. + continue + if req_id in self.requests_origin_status: + request.status = self.requests_origin_status.pop(req_id)