[Feature] Abort diffusion request (#1156)

Signed-off-by: dengyunyang <584797741@qq.com>

vllm_omni/diffusion/diffusion_engine.py

@@ -2,8 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import os
+import signal
 import time
-from collections.abc import Iterable
 from typing import Any
 
 import PIL.Image
@@ -371,7 +371,13 @@ class DiffusionEngine:
         if hasattr(self, "executor"):
             self.executor.shutdown()
 
-    def abort(self, request_id: str | Iterable[str]) -> None:
-        # TODO implement it
-        logger.warning("DiffusionEngine abort is not implemented yet")
-        pass
+    def abort(self) -> None:
+        """
+        Abort diffusion generation requests.
+        """
+        try:
+            for worker_pid in self.executor.get_worker_pids():
+                os.kill(worker_pid, signal.SIGUSR1)
+                logger.info("Interrupt signal sent to worker.")
+        except Exception as e:
+            logger.error(f"Failed to interrupt current diffusion forward: {e}")

vllm_omni/diffusion/executor/abstract.py

@@ -84,3 +84,8 @@ class DiffusionExecutor(ABC):
     def shutdown(self) -> None:
         """Shutdown the executor and release resources."""
         pass
+
+    @abstractmethod
+    def get_worker_pids(self) -> None:
+        """Get worker pid."""
+        pass

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -130,6 +130,9 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
     def add_req(self, request: OmniDiffusionRequest) -> DiffusionOutput:
         return self.scheduler.add_req(request)
 
+    def get_worker_pids(self):
+        return [process.pid for process in self._processes]
+
     def collective_rpc(
         self,
         method: str,

vllm_omni/diffusion/models/longcat_image/pipeline_longcat_image_edit.py

@@ -72,11 +72,15 @@ def get_longcat_image_edit_pre_process_func(
             if "additional_information" not in prompt:
                 prompt["additional_information"] = {}
 
-            if raw_image is None or isinstance(raw_image, list):
-                raise ValueError(
-                    """Received no image or a list of image. Only a single image is supported by this model."""
-                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`"""
-                )
+            if not raw_image:  # None or empty list
+                raise ValueError("""Received no input image. This model requires one input image to run.""")
+            elif isinstance(raw_image, list):
+                if len(raw_image) > 1:
+                    raise ValueError(
+                        """Received multiple input images. Only a single image is supported by this model."""
+                    )
+                else:
+                    raw_image = raw_image[0]
 
             if isinstance(raw_image, str):
                 image = PIL.Image.open(raw_image)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py

@@ -75,11 +75,15 @@ def get_qwen_image_edit_pre_process_func(
                 prompt["additional_information"] = {}
 
             # Only handles single image
-            if raw_image is None or isinstance(raw_image, list):
-                raise ValueError(
-                    """Received no image or a list of image. Only a single image is supported by this model."""
-                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`"""
-                )
+            if not raw_image:  # None or empty list
+                raise ValueError("""Received no input image. This model requires one input image to run.""")
+            elif isinstance(raw_image, list):
+                if len(raw_image) > 1:
+                    raise ValueError(
+                        """Received multiple input images. Only a single image is supported by this model."""
+                    )
+                else:
+                    raw_image = raw_image[0]
 
             if isinstance(raw_image, str):
                 image = PIL.Image.open(raw_image)

vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py

@@ -79,11 +79,15 @@ def get_qwen_image_layered_pre_process_func(
             if "additional_information" not in prompt:
                 prompt["additional_information"] = {}
 
-            if raw_image is None or isinstance(raw_image, list):
-                raise ValueError(
-                    """Received no image or a list of image. Only a single image is supported by this model."""
-                    """Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`"""
-                )
+            if not raw_image:  # None or empty list
+                raise ValueError("""Received no input image. This model requires one input image to run.""")
+            elif isinstance(raw_image, list):
+                if len(raw_image) > 1:
+                    raise ValueError(
+                        """Received multiple input images. Only a single image is supported by this model."""
+                    )
+                else:
+                    raw_image = raw_image[0]
 
             if isinstance(raw_image, str):
                 image = PIL.Image.open(raw_image)

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -9,6 +9,7 @@ to DiffusionModelRunner.
 
 import multiprocessing as mp
 import os
+import signal
 from contextlib import AbstractContextManager, nullcontext
 
 import torch
@@ -66,6 +67,21 @@ class DiffusionWorker:
         self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
         self.lora_manager: DiffusionLoRAManager | None = None
         self.init_device()
+        # Registry signal
+        signal.signal(signal.SIGUSR1, self._on_interrupt)
+
+    def _on_interrupt(self, signum, frame):
+        # Set interrupt flag on the pipeline
+        # The pipeline checks this flag in the diffuse loop
+        if (
+            self.model_runner is not None
+            and self.model_runner.pipeline is not None
+            and hasattr(self.model_runner.pipeline, "_interrupt")
+        ):
+            self.model_runner.pipeline._interrupt = True
+            logger.info("Worker: Interrupt flag set for current Running Request")
+        else:
+            logger.warning("Worker: Cannot set interrupt flag - model_runner or pipeline is None")
 
     def init_device(self) -> None:
         """Initialize the device and distributed environment."""

vllm_omni/entrypoints/async_omni.py

@@ -471,6 +471,29 @@ class AsyncOmni(OmniBase):
             else:
                 logger.debug(f"[{self._name}] Request {req_id} fully completed")
 
+    def handle_oom(self, result: dict[str, Any], stage: OmniStage, stage_id: int) -> None:
+        if getattr(stage, "stage_type", None) == "diffusion":
+            error_text = str(result.get("error", ""))
+            error_lower = error_text.lower()
+            is_oom = "out of memory" in error_lower or "cuda oom" in error_lower
+            if is_oom:
+                logger.critical(
+                    f"[{self._name}] Diffusion stage {stage_id} reported OOM; stopping stage worker.",
+                )
+                try:
+                    stage.stop_stage_worker()
+                except Exception as e:
+                    logger.warning(
+                        f"[{self._name}] Failed to stop diffusion stage {stage_id}: {e}",
+                    )
+                logger.critical(
+                    f"[{self._name}] Exiting API process due to diffusion stage {stage_id} OOM.",
+                )
+
+                import os as _os
+
+                _os._exit(1)
+
     def _process_single_result(
         self,
         result: dict[str, Any],
@@ -493,6 +516,7 @@ class AsyncOmni(OmniBase):
             logger.error(
                 f"[{self._name}] Stage {stage_id} error on request {req_id}: {result['error']}",
             )
+            self.handle_oom(result, stage, stage_id)  # only temporary workaround for oom
             raise RuntimeError(result)
 
         engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")

vllm_omni/entrypoints/async_omni_diffusion.py

@@ -9,8 +9,9 @@ enabling concurrent request handling and streaming generation.
 """
 
 import asyncio
+import threading
 import uuid
-from collections.abc import AsyncGenerator, Iterable
+from collections.abc import AsyncGenerator
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any
 
@@ -106,8 +107,24 @@ class AsyncOmniDiffusion:
         self._executor = ThreadPoolExecutor(max_workers=1)
         self._closed = False
 
+        # Track active requests for abort functionality
+        self._active_generate_task: dict[str, asyncio.Task] = {}
+        self._active_engine_step_task: dict[str, asyncio.Task] = {}
+        self._running_lock = threading.Lock()
+        self._running_requests: set[str] = set()
+
         logger.info("AsyncOmniDiffusion initialized with model: %s", model)
 
+    def _engine_step_wrapper(self, request, request_id):
+        with self._running_lock:
+            self._running_requests.add(request_id)
+
+        try:
+            return self.engine.step(request)
+        finally:
+            with self._running_lock:
+                self._running_requests.discard(request_id)
+
     async def generate(
         self,
         prompt: OmniPromptType,
@@ -137,6 +154,10 @@ class AsyncOmniDiffusion:
         if lora_request is not None:
             sampling_params.lora_request = lora_request
 
+        # Register this task for tracking
+        current_task = asyncio.current_task()
+        self._active_generate_task[request_id] = current_task
+
         request = OmniDiffusionRequest(
             prompts=[prompt],
             sampling_params=sampling_params,
@@ -147,17 +168,18 @@ class AsyncOmniDiffusion:
 
         # Run engine in thread pool
         loop = asyncio.get_event_loop()
+        # In async mode, only a single request is submitted at a time
         try:
             # In async mode, only a single request is submitted at a time
-            result = await loop.run_in_executor(
-                self._executor,
-                self.engine.step,
-                request,
-            )
+            step_task = loop.run_in_executor(self._executor, self._engine_step_wrapper, request, request_id)
+            self._active_engine_step_task[request_id] = step_task
+            result = await step_task
             result = result[0]
-        except Exception as e:
-            logger.error("Generation failed for request %s: %s", request_id, e)
-            raise RuntimeError(f"Diffusion generation failed: {e}") from e
+        finally:
+            # Unregister the request
+            self._active_generate_task.pop(request_id, None)
+            self._active_engine_step_task.pop(request_id, None)
+            logger.debug(f"{request_id} finished")
 
         # Update request_id if needed
         if not result.request_id:
@@ -219,9 +241,26 @@ class AsyncOmniDiffusion:
         except Exception:
             pass
 
-    async def abort(self, request_id: str | Iterable[str]) -> None:
-        """Abort a request."""
-        self.engine.abort(request_id)
+    async def abort(self, request_id: str):
+        logger.info(f"Aborting {request_id}.")
+        with self._running_lock:
+            is_running = request_id in self._running_requests
+            if is_running:
+                self.engine.abort()
+            else:
+                step_task = self._active_engine_step_task.get(request_id)
+                if step_task and not step_task.done():
+                    step_task.cancel()
+        if is_running:
+            logger.info(f"{request_id} engine step is running, sending interrupt")
+        else:
+            logger.info(f"{request_id} engine step not running, cancel in queue")
+
+        gen_task = self._active_generate_task.get(request_id)
+        if gen_task and not gen_task.done():
+            gen_task.cancel()
+        self._active_engine_step_task.pop(request_id, None)
+        self._active_generate_task.pop(request_id, None)
 
     @property
     def is_running(self) -> bool:

vllm_omni/entrypoints/openai/api_server.py

@@ -991,6 +991,7 @@ async def generate_images(request: ImageGenerationRequest, raw_request: Request)
         HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
     },
 )
+@with_cancellation
 async def edit_images(
     raw_request: Request,
     image: list[UploadFile] | None = File(None),
@@ -1014,7 +1015,7 @@ async def edit_images(
     seed: int | None = Form(None),
     # vllm-omni extension for per-request LoRA.
     lora: str | None = Form(None),  # Json string
-) -> ImageGenerationResponse:
+) -> ImageGenerationResponse | None:
     """
     OpenAI-compatible image edit endpoint.
     """