[Bugfix] Fix colocate-async sleep admission race and deliver AR abort tokens (#6367)

Signed-off-by: knlnguyen1802 <knlnguyen1802@gmail.com>
Signed-off-by: Cursor <cursoragent@cursor.com>
Signed-off-by: Samit <285365963@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Samit <285365963@qq.com>

vllm_omni/outputs/output_processor.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
+from collections.abc import Iterable
 from dataclasses import fields as dataclass_fields
 from typing import Any
 
@@ -432,6 +433,116 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         return self._native_text_metrics_by_request.pop(request_id, {})
 
     def abort_requests(self, request_ids, internal: bool) -> list[str]:
+        aborted_ids, _outputs = self.abort_requests_collecting_outputs(request_ids, internal=internal)
+        return aborted_ids
+
+    def abort_requests_collecting_outputs(
+        self,
+        request_ids: Iterable[str],
+        *,
+        internal: bool,
+        commit_state: bool = True,
+    ) -> tuple[list[str], list[RequestOutput | PoolingRequestOutput]]:
+        """Abort requests and return terminal abort outputs with partial tokens.
+
+        Mirrors upstream ``OutputProcessor.abort_requests``, but also returns
+        abort ``RequestOutput`` objects when ``queue`` is ``None`` (Omni stage
+        pools register OP state without a collector queue).
+
+        Terminal abort outputs use ``RequestOutputKind.CUMULATIVE`` so DELTA
+        streaming requests still surface the full prefix generated so far.
+        ``new_token_ids`` is the detokenizer prefix: Omni's
+        ``make_request_output`` writes that list onto ``CompletionOutput.token_ids``.
+
+        When ``commit_state`` is False, processor mappings stay in place so a
+        failed physical EngineCore abort can retry with the same prefix.
+        """
+        request_ids = list(request_ids)
+        if commit_state:
+            for request_id in request_ids:
+                if internal:
+                    req_state = self.request_states.get(request_id)
+                    if req_state is not None:
+                        self._native_text_metrics_by_request.pop(req_state.external_req_id, None)
+                else:
+                    self._native_text_metrics_by_request.pop(request_id, None)
+
+        internal_req_ids: list[str] = []
+        for request_id in request_ids:
+            if internal:
+                internal_req_ids.append(request_id)
+                if commit_state:
+                    if req_state := self.request_states.get(request_id):
+                        external_req_id = req_state.external_req_id
+                        internal_ids = self.external_req_ids.get(external_req_id)
+                        if internal_ids is not None:
+                            try:
+                                internal_ids.remove(request_id)
+                            except ValueError:
+                                pass
+                            if not internal_ids:
+                                self.external_req_ids.pop(external_req_id, None)
+            elif commit_state:
+                if internal_ids := self.external_req_ids.pop(request_id, []):
+                    internal_req_ids.extend(internal_ids)
+            else:
+                internal_req_ids.extend(self.external_req_ids.get(request_id, []))
+
+        request_ids_to_abort: list[str] = []
+        abort_outputs: list[RequestOutput | PoolingRequestOutput] = []
+        for request_id in internal_req_ids:
+            req_state = (
+                self.request_states.pop(request_id, None) if commit_state else self.request_states.get(request_id)
+            )
+            if req_state is not None:
+                if commit_state:
+                    self.lora_states.request_finished(request_id, req_state.lora_name)
+                request_ids_to_abort.append(request_id)
+                original_kind = req_state.output_kind
+                req_state.output_kind = RequestOutputKind.CUMULATIVE
+                detok = req_state.detokenizer
+                new_token_ids = list(detok.output_token_ids) if detok is not None else []
+                saved_children = None
+                if not commit_state and req_state.parent_req is not None:
+                    saved_children = set(req_state.parent_req.child_requests)
+                try:
+                    request_output = req_state.make_request_output(
+                        new_token_ids=new_token_ids,
+                        pooling_output=None,
+                        finish_reason=FinishReason.ABORT,
+                        stop_reason=None,
+                        kv_transfer_params=None,
+                    )
+                finally:
+                    req_state.output_kind = original_kind
+                    if saved_children is not None and req_state.parent_req is not None:
+                        req_state.parent_req.child_requests = saved_children
+                if request_output is not None:
+                    abort_outputs.append(request_output)
+                    if commit_state and req_state.queue is not None:
+                        req_state.queue.put(request_output)
+                if commit_state:
+                    parent_req = req_state.parent_req
+                    if parent_req is not None:
+                        parent_req.child_requests.discard(request_id)
+                        if not parent_req.child_requests:
+                            self.parent_requests.pop(parent_req.request_id, None)
+            elif parent := self.parent_requests.get(request_id):
+                if parent.child_requests:
+                    child_reqs = list(parent.child_requests)
+                    child_aborted, child_outputs = self.abort_requests_collecting_outputs(
+                        child_reqs,
+                        internal=True,
+                        commit_state=commit_state,
+                    )
+                    request_ids_to_abort.extend(child_aborted)
+                    abort_outputs.extend(child_outputs)
+                if commit_state:
+                    self.parent_requests.pop(request_id, None)
+        return request_ids_to_abort, abort_outputs
+
+    def commit_aborted_request_state(self, request_ids: Iterable[str], *, internal: bool) -> None:
+        """Drop processor mappings after a successful physical EngineCore abort."""
         request_ids = list(request_ids)
         for request_id in request_ids:
             if internal:
@@ -440,7 +551,37 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                     self._native_text_metrics_by_request.pop(req_state.external_req_id, None)
             else:
                 self._native_text_metrics_by_request.pop(request_id, None)
-        return super().abort_requests(request_ids, internal)
+
+        internal_req_ids: list[str] = []
+        for request_id in request_ids:
+            if internal:
+                internal_req_ids.append(request_id)
+                if req_state := self.request_states.get(request_id):
+                    external_req_id = req_state.external_req_id
+                    internal_ids = self.external_req_ids.get(external_req_id)
+                    if internal_ids is not None:
+                        try:
+                            internal_ids.remove(request_id)
+                        except ValueError:
+                            pass
+                        if not internal_ids:
+                            self.external_req_ids.pop(external_req_id, None)
+            elif internal_ids := self.external_req_ids.pop(request_id, []):
+                internal_req_ids.extend(internal_ids)
+
+        for request_id in internal_req_ids:
+            req_state = self.request_states.pop(request_id, None)
+            if req_state is not None:
+                self.lora_states.request_finished(request_id, req_state.lora_name)
+                parent_req = req_state.parent_req
+                if parent_req is not None:
+                    parent_req.child_requests.discard(request_id)
+                    if not parent_req.child_requests:
+                        self.parent_requests.pop(parent_req.request_id, None)
+            elif parent := self.parent_requests.get(request_id):
+                if parent.child_requests:
+                    self.commit_aborted_request_state(list(parent.child_requests), internal=True)
+                self.parent_requests.pop(request_id, None)
 
     def add_request(
         self,