From dc45f43a6ad22954265e6807e0b3caacda15d723 Mon Sep 17 00:00:00 2001
From: dengyunyang <584797741@qq.com>
Date: Tue, 3 Feb 2026 15:49:58 +0800
Subject: [PATCH] [Feature] Abort diffusion request (#1156)

Signed-off-by: dengyunyang <584797741@qq.com>
---
 vllm_omni/diffusion/diffusion_engine.py       | 16 +++--
 vllm_omni/diffusion/executor/abstract.py      |  5 ++
 .../diffusion/executor/multiproc_executor.py  |  3 +
 .../diffusion/worker/diffusion_worker.py      | 16 +++++
 vllm_omni/entrypoints/async_omni_diffusion.py | 63 ++++++++++++++++---
 vllm_omni/entrypoints/openai/api_server.py    |  3 +-
 6 files changed, 90 insertions(+), 16 deletions(-)

diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py
index a44596983..05a461361 100644
--- a/vllm_omni/diffusion/diffusion_engine.py
+++ b/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}")
diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py
index e41f41d11..150f2057a 100644
--- a/vllm_omni/diffusion/executor/abstract.py
+++ b/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
diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py
index a421a01ed..cce4bdad0 100644
--- a/vllm_omni/diffusion/executor/multiproc_executor.py
+++ b/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,
diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py
index 0fad90a5b..65906467b 100644
--- a/vllm_omni/diffusion/worker/diffusion_worker.py
+++ b/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."""
diff --git a/vllm_omni/entrypoints/async_omni_diffusion.py b/vllm_omni/entrypoints/async_omni_diffusion.py
index 6f447c332..3d2f3edeb 100644
--- a/vllm_omni/entrypoints/async_omni_diffusion.py
+++ b/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,
@@ -148,12 +169,17 @@ class AsyncOmniDiffusion:
         # Run engine in thread pool
         loop = asyncio.get_event_loop()
         # In async mode, only a single request is submitted at a time
-        result = await loop.run_in_executor(
-            self._executor,
-            self.engine.step,
-            request,
-        )
-        result = result[0]
+        try:
+            # In async mode, only a single request is submitted at a time
+            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]
+        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:
@@ -215,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:
diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py
index 2e057901a..f9983b31e 100644
--- a/vllm_omni/entrypoints/openai/api_server.py
+++ b/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.
     """