[BugFix] Fixing occasional engine crashes caused by abort requests (#2871)

Signed-off-by: amy-why-3459 <wuhaiyan17@huawei.com>

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

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

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.

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)