[Feature] Support Omni serving abort request (#486)

Signed-off-by: ZeldaHuang <hzm414167@alibaba-inc.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

.buildkite/pipeline.yml

@@ -96,11 +96,11 @@ steps:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Omni Model Test with H100"
-    timeout_in_minutes: 20
+    timeout_in_minutes: 30
     depends_on: image-build
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py
+      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_async_omni.py
     agents:
       queue: "mithril-h100-pool"
     plugins:

docs/contributing/model/adding_omni_model.md

@@ -361,6 +361,7 @@ Stage transitions happen automatically in the orchestrator (`OmniLLM` class) dur
 
            # Submit to next stage
            task = {
+               "type": OmniStageTaskType.GENERATE,
                "request_id": req_id,
                "engine_inputs": next_inputs[0],
                "sampling_params": sampling_params_list[next_stage_id],

tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml

@@ -0,0 +1,31 @@
+# The following config has been verified on 2x H100-80G GPUs.
+stage_args:
+  - stage_id: 0
+    runtime:
+      devices: "0,1"
+      max_batch_size: 5
+    engine_args:
+      model_stage: thinker
+      model_arch: Qwen3OmniMoeForConditionalGeneration
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      gpu_memory_utilization: 0.6
+      enforce_eager: true
+      trust_remote_code: true
+      engine_output_type: latent  # Output hidden states for talker
+      distributed_executor_backend: "mp"
+      enable_prefix_caching: false
+      hf_config_name: thinker_config
+      tensor_parallel_size: 2
+      load_format: dummy
+    final_output: true
+    final_output_type: text
+    is_comprehension: true
+    default_sampling_params:
+      temperature: 0.4
+      top_p: 0.9
+      top_k: 1
+      max_tokens: 100
+      seed: 42
+      detokenize: True
+      repetition_penalty: 1.05

tests/e2e/online_serving/test_async_omni.py

@@ -0,0 +1,100 @@
+import asyncio
+from contextlib import ExitStack
+from pathlib import Path
+
+import pytest
+from vllm.inputs import PromptType
+
+from vllm_omni.entrypoints.async_omni import AsyncOmni
+
+SEED = 42
+
+stage_config = str(Path(__file__).parent / "stage_configs" / "qwen3_omni_thinker_ci.yaml")
+model = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+
+
+async def generate(
+    engine: AsyncOmni,
+    request_id: str,
+    prompt: PromptType,
+    max_tokens: int,
+) -> tuple[int, str]:
+    # Ensure generate doesn't complete too fast for cancellation test.
+    await asyncio.sleep(0.2)
+    thinker_sampling_params = {
+        "temperature": 0.4,  # Deterministic
+        "top_p": 0.9,
+        "top_k": 1,
+        "max_tokens": max_tokens,
+        "repetition_penalty": 1.05,
+        "stop_token_ids": [151645],  # Qwen EOS token <|im_end|>
+        "seed": SEED,
+    }
+
+    sampling_params_list = [
+        thinker_sampling_params,
+    ]
+    count = 0
+    async for omni_output in engine.generate(
+        prompt=prompt,
+        request_id=request_id,
+        sampling_params_list=sampling_params_list,
+        output_modalities=["text"],
+    ):
+        stage_id = omni_output.stage_id
+        out = omni_output.request_output
+        if stage_id == 0:
+            num_tokens = sum(len(output.token_ids) for output in out.outputs)
+            count = num_tokens
+
+        await asyncio.sleep(0.0)
+
+    return count, request_id
+
+
+@pytest.mark.asyncio
+async def test_abort():
+    with ExitStack() as after:
+        engine = AsyncOmni(model=model, stage_configs_path=stage_config)
+        after.callback(engine.shutdown)
+
+        NUM_REQUESTS = 5
+        NUM_EXPECTED_TOKENS = 100
+        NUM_EXPECTED_TOKENS_LONG = 1000
+        REQUEST_IDS_TO_ABORT = [1, 2, 3]
+
+        prompt = "Hello my name is Robert and "
+
+        request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)]
+
+        # Create concurrent requests.
+        tasks: list[asyncio.Task] = []
+        for idx, request_id in enumerate(request_ids):
+            max_tokens = NUM_EXPECTED_TOKENS_LONG if (idx in REQUEST_IDS_TO_ABORT) else NUM_EXPECTED_TOKENS
+            tasks.append(asyncio.create_task(generate(engine, request_id, prompt, max_tokens)))
+
+        # API server cancels requests when they disconnect.
+        for idx in REQUEST_IDS_TO_ABORT:
+            tasks[idx].cancel()
+            await asyncio.sleep(0.1)
+
+        # Confirm the other requests are okay.
+        for idx, task in enumerate(tasks):
+            # Confirm that it was actually canceled.
+            if idx in REQUEST_IDS_TO_ABORT:
+                with pytest.raises((asyncio.CancelledError, GeneratorExit)):
+                    await task
+            else:
+                # Otherwise, make sure the request was not impacted.
+                num_generated_tokens, request_id = await task
+                expected_tokens = NUM_EXPECTED_TOKENS
+                assert num_generated_tokens == expected_tokens, (
+                    f"{request_id} generated {num_generated_tokens} but expected {expected_tokens}"
+                )
+
+        # Confirm we can do another generation.
+        request_id = f"request-{REQUEST_IDS_TO_ABORT[0]}"
+        task = asyncio.create_task(generate(engine, request_id, prompt, NUM_EXPECTED_TOKENS))
+        num_generated_tokens, request_id = await task
+        assert num_generated_tokens == NUM_EXPECTED_TOKENS
+    await asyncio.sleep(5)

tests/entrypoints/test_omni_llm.py

@@ -6,6 +6,8 @@ from unittest.mock import MagicMock
 
 import pytest
 
+from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK
+
 # Suppress noisy DeprecationWarnings from optional Swig bindings imported by vLLM dependencies.
 warnings.filterwarnings(
     "ignore",
@@ -127,7 +129,7 @@ class _FakeStage:
         """Mock stop_stage_worker: clean up queue references."""
         if self._in_q is not None:
             try:
-                self._in_q.put_nowait(None)
+                self._in_q.put_nowait(SHUTDOWN_TASK)
             except Exception:
                 pass
 
@@ -977,7 +979,7 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
     # Use get_nowait to avoid blocking (close() uses put_nowait, so should be safe)
     try:
         shutdown_signal = omni.stage_list[0]._in_q.get_nowait()
-        assert shutdown_signal is None
+        assert shutdown_signal == SHUTDOWN_TASK
     except Empty:
         # If queue was already empty or only had stage_ready, that's also acceptable
         # The important thing is that close() was called without error

vllm_omni/distributed/omni_connectors/adapter.py

@@ -7,6 +7,8 @@ import time
 from collections.abc import Callable
 from typing import Any
 
+from vllm_omni.entrypoints.stage_utils import OmniStageTaskType
+
 from .utils.logging import get_connector_logger
 
 logger = get_connector_logger(__name__)
@@ -49,6 +51,7 @@ def try_send_via_connector(
         if success:
             # Send lightweight notification via queue
             notify_payload = {
+                "type": OmniStageTaskType.GENERATE,
                 "request_id": req_id,
                 "sampling_params": sampling_params,
                 "from_connector": True,

vllm_omni/entrypoints/async_omni.py

@@ -29,6 +29,7 @@ from vllm_omni.entrypoints.log_utils import (
 )
 from vllm_omni.entrypoints.omni import OmniBase
 from vllm_omni.entrypoints.omni_stage import OmniStage
+from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK, OmniStageTaskType
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
     get_final_stage_id_for_e2e,
@@ -43,7 +44,7 @@ def _weak_close_cleanup_async(stage_list, stage_in_queues, ray_pg, output_handle
     if stage_list:
         for q in stage_in_queues:
             try:
-                q.put_nowait(None)
+                q.put_nowait(SHUTDOWN_TASK)
             except Exception as e:
                 logger.warning(f"Failed to send shutdown signal to stage input queue: {e}")
         for stage in stage_list:
@@ -266,197 +267,203 @@ class AsyncOmni(OmniBase):
             await self._pause_cond.wait_for(lambda: not self._paused)
 
         logger.debug(f"[{self._name}] generate() called")
+        try:
+            # Start output handler on the first call to generate()
+            self._run_output_handler()
 
-        # Start output handler on the first call to generate()
-        self._run_output_handler()
+            prompt = args[0] if args else kwargs.get("prompt")
+            request_id = args[1] if len(args) > 1 else kwargs.get("request_id")
+            sampling_params_list = args[2] if len(args) > 2 else kwargs.get("sampling_params_list")
+            output_modalities = kwargs.get("output_modalities", None)
+            # TODO: lora_request, trace_headers, priority are not supported yet
 
-        prompt = args[0] if args else kwargs.get("prompt")
-        request_id = args[1] if len(args) > 1 else kwargs.get("request_id")
-        sampling_params_list = args[2] if len(args) > 2 else kwargs.get("sampling_params_list")
-        output_modalities = kwargs.get("output_modalities", None)
-        # TODO: lora_request, trace_headers, priority are not supported yet
+            if sampling_params_list is None:
+                # For Omni LLM, the params are parsed via the yaml file. For the current version,
+                # diffusion params can parsed via the command line.
+                omni_params_kwargs = {
+                    k: v for k, v in kwargs.items() if k not in ["prompt", "request_id", "output_modalities"]
+                }
 
-        if sampling_params_list is None:
-            # For Omni LLM, the params are parsed via the yaml file. For the current version,
-            # diffusion params can parsed via the command line.
-            omni_params_kwargs = {
-                k: v for k, v in kwargs.items() if k not in ["prompt", "request_id", "output_modalities"]
-            }
+                per_stage_params: list[Any] = []
+                for stage_id, stage in enumerate(self.stage_list):
+                    stage_type = getattr(stage, "stage_type", "llm")
+                    if stage_type == "diffusion":
+                        default_dict = self.default_sampling_params_list[stage_id]
+                        # Merge user-provided kwargs
+                        merged = {**default_dict, **omni_params_kwargs}
+                        # Diffusion only needs to keep diff params, will be used via OmniDiffusionRequest
+                        per_stage_params.append(merged)
+                    else:
+                        # LLM directly constructs SamplingParams, don't use the merged params
+                        per_stage_params.append(self.default_sampling_params_list[stage_id])
 
-            per_stage_params: list[Any] = []
-            for stage_id, stage in enumerate(self.stage_list):
-                stage_type = getattr(stage, "stage_type", "llm")
-                if stage_type == "diffusion":
-                    default_dict = self.default_sampling_params_list[stage_id]
-                    # Merge user-provided kwargs
-                    merged = {**default_dict, **omni_params_kwargs}
-                    # Diffusion only needs to keep diff params, will be used via OmniDiffusionRequest
-                    per_stage_params.append(merged)
-                else:
-                    # LLM directly constructs SamplingParams, don't use the merged params
-                    per_stage_params.append(self.default_sampling_params_list[stage_id])
+                sampling_params_list = per_stage_params
 
-            sampling_params_list = per_stage_params
+            if len(sampling_params_list) != len(self.stage_list):
+                raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
 
-        if len(sampling_params_list) != len(self.stage_list):
-            raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
+            # Orchestrator keeps stage objects for input derivation
+            num_stages = len(self.stage_list)
+            # Track per-request start time for end-to-end timing
+            _req_start_ts: dict[int, float] = {}
+            _wall_start_ts: float = time.time()
+            # _last_finish_ts: float = _wall_start_ts
 
-        # Orchestrator keeps stage objects for input derivation
-        num_stages = len(self.stage_list)
-        # Track per-request start time for end-to-end timing
-        _req_start_ts: dict[int, float] = {}
-        _wall_start_ts: float = time.time()
-        # _last_finish_ts: float = _wall_start_ts
-
-        # Determine the final stage for E2E stats (highest stage_id with
-        # final_output=True; fallback to last stage)
-        final_stage_id_for_e2e = get_final_stage_id_for_e2e(output_modalities, self.output_modalities, self.stage_list)
-
-        # Metrics/aggregation helper
-        metrics = OrchestratorMetrics(
-            num_stages,
-            self._enable_stats,
-            _wall_start_ts,
-        )
-        # Seed stage-0 queue with all requests
-        logger.debug(f"[{self._name}] Seeding request into stage-0")
-        req_state = ClientRequestState(request_id)
-        self.request_states[request_id] = req_state
-
-        # Mark first input time for stage-0
-        metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
-
-        sp0: SamplingParams = sampling_params_list[0]  # type: ignore[index]
-        task = {
-            "request_id": request_id,
-            "engine_inputs": prompt,
-            "sampling_params": sp0,
-        }
-        self.stage_list[0].submit(task)
-        _req_start_ts[request_id] = time.time()
-        logger.debug(f"[{self._name}] Enqueued request {request_id} to stage-0")
-
-        logger.debug(f"[{self._name}] Entering scheduling loop: stages={num_stages}")
-        for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]):
-            result = await req_state.queue.get()
-            assert stage_id == req_state.stage_id
-
-            req_id = result.get("request_id")
-            if "error" in result:
-                logger.error(
-                    f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}",
-                )
-                raise RuntimeError(result)  # Request Finished due to error
-
-            engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
-            # Mark last output time for this stage whenever we receive outputs
-            metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
-            try:
-                _m = asdict(result.get("metrics"))
-                if _m is not None:
-                    metrics.on_stage_metrics(stage_id, req_id, _m)
-            except Exception as e:
-                logger.exception(
-                    f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}",
-                )
-            logger.debug(
-                f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing",
+            # Determine the final stage for E2E stats (highest stage_id with
+            # final_output=True; fallback to last stage)
+            final_stage_id_for_e2e = get_final_stage_id_for_e2e(
+                output_modalities, self.output_modalities, self.stage_list
             )
-            stage.set_engine_outputs(engine_outputs)
 
-            if getattr(stage, "final_output", False):
-                logger.debug(
-                    f"[{self._name}] Request {req_id} finalized at stage-{stage_id}",
-                )
+            # Metrics/aggregation helper
+            metrics = OrchestratorMetrics(
+                num_stages,
+                self._enable_stats,
+                _wall_start_ts,
+            )
+            # Seed stage-0 queue with all requests
+            logger.debug(f"[{self._name}] Seeding request into stage-0")
+            req_state = ClientRequestState(request_id)
+            self.request_states[request_id] = req_state
 
-                # End-to-end timing and time-per-token for final output
-                # (only once per request at the designated final stage)
+            # Mark first input time for stage-0
+            metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
+
+            sp0: SamplingParams = sampling_params_list[0]  # type: ignore[index]
+            task = {
+                "request_id": request_id,
+                "engine_inputs": prompt,
+                "sampling_params": sp0,
+            }
+            self.stage_list[0].submit(task)
+            _req_start_ts[request_id] = time.time()
+            logger.debug(f"[{self._name}] Enqueued request {request_id} to stage-0")
+
+            logger.debug(f"[{self._name}] Entering scheduling loop: stages={num_stages}")
+            for stage_id, stage in enumerate(self.stage_list[: final_stage_id_for_e2e + 1]):
+                result = await req_state.queue.get()
+                assert stage_id == req_state.stage_id
+
+                req_id = result.get("request_id")
+                if "error" in result:
+                    logger.error(
+                        f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}",
+                    )
+                    raise RuntimeError(result)  # Request Finished due to error
+
+                engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
+                # Mark last output time for this stage whenever we receive outputs
+                metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
                 try:
-                    rid_key = str(req_id)
-                    if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done:
-                        metrics.on_finalize_request(
-                            stage_id,
-                            req_id,
-                            engine_outputs,
-                            _req_start_ts.get(req_id, _wall_start_ts),
-                        )
+                    _m = asdict(result.get("metrics"))
+                    if _m is not None:
+                        metrics.on_stage_metrics(stage_id, req_id, _m)
                 except Exception as e:
                     logger.exception(
-                        f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}",
+                        f"[{self._name}] Failed to process metrics for stage {stage_id}, req {req_id}: {e}",
+                    )
+                logger.debug(
+                    f"[{self._name}] Stage-{stage_id} completed request {req_id}; forwarding or finalizing",
+                )
+                stage.set_engine_outputs(engine_outputs)
+
+                if getattr(stage, "final_output", False):
+                    logger.debug(
+                        f"[{self._name}] Request {req_id} finalized at stage-{stage_id}",
                     )
 
-                if isinstance(engine_outputs, list):
-                    engine_outputs = engine_outputs[0]
-                # Handle diffusion outputs that already contain images
-                if stage.final_output_type == "image":
-                    images = []
-                    if isinstance(engine_outputs, OmniRequestOutput) and engine_outputs.images:
-                        images = engine_outputs.images
-                    elif hasattr(engine_outputs, "images") and engine_outputs.images:
-                        images = engine_outputs.images
-                    yield OmniRequestOutput(
-                        stage_id=stage_id,
-                        final_output_type=stage.final_output_type,
-                        request_output=engine_outputs,
-                        images=images,
-                    )
+                    # End-to-end timing and time-per-token for final output
+                    # (only once per request at the designated final stage)
+                    try:
+                        rid_key = str(req_id)
+                        if stage_id == final_stage_id_for_e2e and rid_key not in metrics.e2e_done:
+                            metrics.on_finalize_request(
+                                stage_id,
+                                req_id,
+                                engine_outputs,
+                                _req_start_ts.get(req_id, _wall_start_ts),
+                            )
+                    except Exception as e:
+                        logger.exception(
+                            f"[{self._name}] Finalize request handling error for req {req_id} at stage {stage_id}: {e}",
+                        )
+
+                    if isinstance(engine_outputs, list):
+                        engine_outputs = engine_outputs[0]
+                    # Handle diffusion outputs that already contain images
+                    if stage.final_output_type == "image":
+                        images = []
+                        if isinstance(engine_outputs, OmniRequestOutput) and engine_outputs.images:
+                            images = engine_outputs.images
+                        elif hasattr(engine_outputs, "images") and engine_outputs.images:
+                            images = engine_outputs.images
+                        yield OmniRequestOutput(
+                            stage_id=stage_id,
+                            final_output_type=stage.final_output_type,
+                            request_output=engine_outputs,
+                            images=images,
+                        )
+                    else:
+                        yield OmniRequestOutput(
+                            stage_id=stage_id,
+                            final_output_type=stage.final_output_type,
+                            request_output=engine_outputs,
+                        )
+
+                # Forward to next stage if there is one
+                next_stage_id = stage_id + 1
+                if next_stage_id <= final_stage_id_for_e2e:
+                    next_stage: OmniStage = self.stage_list[next_stage_id]
+                    next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
+                    sp_next: SamplingParams = sampling_params_list[next_stage_id]
+
+                    # Check if we have a connector for this edge
+                    connector_key = (str(stage_id), str(next_stage_id))
+                    connector = self.connectors.get(connector_key)
+
+                    sent_via_connector = False
+                    if connector:
+                        sent_via_connector = try_send_via_connector(
+                            connector=connector,
+                            stage_id=stage_id,
+                            next_stage_id=next_stage_id,
+                            req_id=req_id,
+                            next_inputs=next_inputs,
+                            sampling_params=sp_next,
+                            original_prompt=prompt,
+                            next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
+                            metrics=metrics,
+                        )
+
+                    if not sent_via_connector:
+                        # Fallback logic removed as we now enforce connector usage.
+                        # If no connector is found or send fails, we log an error and raise,
+                        # because continuing would cause the request to be silently dropped
+                        # and the orchestrator to hang waiting for completion.
+                        error_msg = (
+                            f"[{self._name}] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
+                            "Configure a connector for this edge or inspect connector logs for details."
+                        )
+                        logger.error(error_msg)
+                        raise RuntimeError(error_msg)
+                    logger.debug(f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}")
                 else:
-                    yield OmniRequestOutput(
-                        stage_id=stage_id,
-                        final_output_type=stage.final_output_type,
-                        request_output=engine_outputs,
-                    )
+                    logger.debug(f"[{self._name}] Request {req_id} fully completed")
 
-            # Forward to next stage if there is one
-            next_stage_id = stage_id + 1
-            if next_stage_id <= final_stage_id_for_e2e:
-                next_stage: OmniStage = self.stage_list[next_stage_id]
-                next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
-                sp_next: SamplingParams = sampling_params_list[next_stage_id]
+            logger.debug(f"[{self._name}] All requests completed")
 
-                # Check if we have a connector for this edge
-                connector_key = (str(stage_id), str(next_stage_id))
-                connector = self.connectors.get(connector_key)
-
-                sent_via_connector = False
-                if connector:
-                    sent_via_connector = try_send_via_connector(
-                        connector=connector,
-                        stage_id=stage_id,
-                        next_stage_id=next_stage_id,
-                        req_id=req_id,
-                        next_inputs=next_inputs,
-                        sampling_params=sp_next,
-                        original_prompt=prompt,
-                        next_stage_queue_submit_fn=self.stage_list[next_stage_id].submit,
-                        metrics=metrics,
-                    )
-
-                if not sent_via_connector:
-                    # Fallback logic removed as we now enforce connector usage.
-                    # If no connector is found or send fails, we log an error and raise,
-                    # because continuing would cause the request to be silently dropped
-                    # and the orchestrator to hang waiting for completion.
-                    error_msg = (
-                        f"[{self._name}] Failed to send request {req_id} to stage-{next_stage_id} via connector. "
-                        "Configure a connector for this edge or inspect connector logs for details."
-                    )
-                    logger.error(error_msg)
-                    raise RuntimeError(error_msg)
-                logger.debug(f"[{self._name}] Forwarded request {req_id} to stage-{next_stage_id}")
-            else:
-                logger.debug(f"[{self._name}] Request {req_id} fully completed")
-
-        logger.debug(f"[{self._name}] All requests completed")
-
-        # Summarize and print stats
-        try:
-            summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
-            logger.info("[Summary] %s", pformat(summary, sort_dicts=False))
-        except Exception as e:
-            logger.exception(f"[{self._name}] Failed to build/log summary: {e}")
-        finally:
-            self.request_states.pop(request_id, None)
+            # Summarize and print stats
+            try:
+                summary = metrics.build_and_log_summary(final_stage_id_for_e2e)
+                logger.info("[Summary] %s", pformat(summary, sort_dicts=False))
+            except Exception as e:
+                logger.exception(f"[{self._name}] Failed to build/log summary: {e}")
+            finally:
+                self.request_states.pop(request_id, None)
+        except (asyncio.CancelledError, GeneratorExit):
+            await self.abort(request_id)
+            logger.info("[AsyncOrchestrator] Request %s aborted.", request_id)
+            raise
 
     def _run_output_handler(self) -> None:
         if self.output_handler is not None:
@@ -527,7 +534,10 @@ class AsyncOmni(OmniBase):
         return EngineDeadError()
 
     async def abort(self, request_id: str | Iterable[str]) -> None:
-        pass
+        abort_task = {"type": OmniStageTaskType.ABORT, "request_id": request_id}
+        for stage in self.stage_list:
+            stage.submit(abort_task)
+        return None
 
     async def get_vllm_config(self) -> VllmConfig:
         for stage in self.stage_list:

vllm_omni/entrypoints/omni.py

@@ -30,6 +30,7 @@ from vllm_omni.distributed.ray_utils.utils import (
 )
 from vllm_omni.entrypoints.log_utils import OrchestratorMetrics
 from vllm_omni.entrypoints.omni_stage import OmniStage
+from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK
 from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
 from vllm_omni.entrypoints.utils import (
     get_final_stage_id_for_e2e,
@@ -47,7 +48,7 @@ def _weak_close_cleanup(stage_list, stage_in_queues, ray_pg):
     if stage_list:
         for q in stage_in_queues:
             try:
-                q.put_nowait(None)
+                q.put_nowait(SHUTDOWN_TASK)
             except Exception as e:
                 logger.warning(f"Failed to send shutdown signal to stage input queue: {e}")
         for stage in stage_list:

vllm_omni/entrypoints/omni_stage.py

@@ -38,6 +38,8 @@ from vllm_omni.entrypoints.log_utils import count_tokens_from_outputs
 from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
 from vllm_omni.entrypoints.omni_llm import OmniLLM
 from vllm_omni.entrypoints.stage_utils import (
+    SHUTDOWN_TASK,
+    OmniStageTaskType,
     _to_dict,
     maybe_dump_to_shm,
     set_stage_devices,
@@ -327,7 +329,7 @@ class OmniStage:
         """
         if self._in_q is not None:
             try:
-                self._in_q.put_nowait(None)
+                self._in_q.put_nowait(SHUTDOWN_TASK)
             except Exception as e:
                 logger.warning("Failed to send shutdown to in_q: %s", e)
 
@@ -629,8 +631,9 @@ def _stage_worker(
         task = in_q.get()
 
         _recv_dequeue_ts = _time.time()
-        if task is None:
-            logger.info("Received shutdown signal")
+        task_type = task.get("type", OmniStageTaskType.GENERATE)
+        if task_type == OmniStageTaskType.SHUTDOWN:
+            logger.error("Received shutdown signal")
             break
 
         batch_tasks: list[dict[str, Any]] = [task]
@@ -639,8 +642,8 @@ def _stage_worker(
             while len(batch_tasks) < max_batch_size:
                 if not in_q.empty():
                     extra = in_q.get_nowait()
-                    if extra is None:
-                        in_q.put(None)
+                    if extra == SHUTDOWN_TASK:
+                        in_q.put(SHUTDOWN_TASK)
                         break
                     batch_tasks.append(extra)
                     end_time = _time.time()
@@ -1170,10 +1173,16 @@ async def _stage_worker_async(
     while True:
         try:
             task = in_q.get_nowait()
-            if task is None:
+            task_type = task.get("type", OmniStageTaskType.GENERATE)
+            if task_type == OmniStageTaskType.SHUTDOWN:
                 logger.debug("Received shutdown signal")
+                stage_engine.shutdown()
                 break
-            asyncio.create_task(generation_single_request(task))
+            elif task_type == OmniStageTaskType.ABORT:
+                rid = task["request_id"]
+                asyncio.create_task(stage_engine.abort(rid))
+            else:
+                asyncio.create_task(generation_single_request(task))
         except queue.Empty:
             await asyncio.sleep(0.001)
         batch_request_outputs: list[Any] = []

vllm_omni/entrypoints/stage_utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import enum
 import json
 import logging
 import os
@@ -11,6 +12,15 @@ from omegaconf import OmegaConf
 logger = logging.getLogger(__name__)
 
 
+class OmniStageTaskType(enum.Enum):
+    GENERATE = "generate"
+    ABORT = "abort"
+    SHUTDOWN = "shutdown"
+
+
+SHUTDOWN_TASK = {"type": OmniStageTaskType.SHUTDOWN}
+
+
 def set_stage_devices(
     stage_id: int,
     devices: str | int | None,