From ad0224f2e6c23cc3830842907ced1e33e64fcedb Mon Sep 17 00:00:00 2001 From: Ziming Huang Date: Wed, 31 Dec 2025 16:14:03 +0800 Subject: [PATCH] [Feature] Support Omni serving abort request (#486) Signed-off-by: ZeldaHuang Signed-off-by: Hongsheng Liu Co-authored-by: Hongsheng Liu --- .buildkite/pipeline.yml | 4 +- docs/contributing/model/adding_omni_model.md | 1 + .../stage_configs/qwen3_omni_thinker_ci.yaml | 31 ++ tests/e2e/online_serving/test_async_omni.py | 100 +++++ tests/entrypoints/test_omni_llm.py | 6 +- .../distributed/omni_connectors/adapter.py | 3 + vllm_omni/entrypoints/async_omni.py | 360 +++++++++--------- vllm_omni/entrypoints/omni.py | 3 +- vllm_omni/entrypoints/omni_stage.py | 23 +- vllm_omni/entrypoints/stage_utils.py | 10 + 10 files changed, 354 insertions(+), 187 deletions(-) create mode 100644 tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml create mode 100644 tests/e2e/online_serving/test_async_omni.py diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 02e728b65..a01460424 100644 --- a/.buildkite/pipeline.yml +++ b/.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: diff --git a/docs/contributing/model/adding_omni_model.md b/docs/contributing/model/adding_omni_model.md index 0a66c1709..2a91a3050 100644 --- a/docs/contributing/model/adding_omni_model.md +++ b/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], diff --git a/tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml b/tests/e2e/online_serving/stage_configs/qwen3_omni_thinker_ci.yaml new file mode 100644 index 000000000..64c6b6ee3 --- /dev/null +++ b/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 diff --git a/tests/e2e/online_serving/test_async_omni.py b/tests/e2e/online_serving/test_async_omni.py new file mode 100644 index 000000000..83f4c4d1c --- /dev/null +++ b/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) diff --git a/tests/entrypoints/test_omni_llm.py b/tests/entrypoints/test_omni_llm.py index c77b3b179..381c53284 100644 --- a/tests/entrypoints/test_omni_llm.py +++ b/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 diff --git a/vllm_omni/distributed/omni_connectors/adapter.py b/vllm_omni/distributed/omni_connectors/adapter.py index 97d9765b1..cfc959017 100644 --- a/vllm_omni/distributed/omni_connectors/adapter.py +++ b/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, diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 597248c7e..d133c9b1b 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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: diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index 8b565802c..378caf378 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/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: diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index f5534290f..205deff58 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/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] = [] diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 768d1f473..986cdc0b1 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/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,