[Refactor] Refactor Diffusion Scheduler/Executor Boundaries and Request State Flow (#1625)

Signed-off-by: jader <yjader@foxmail.com>
Signed-off-by: JiangJie Zhang <76905040+yJader@users.noreply.github.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -4,12 +4,14 @@ import weakref
 from dataclasses import dataclass
 from typing import Any
 
+import zmq
+from vllm.distributed.device_communicators.shm_broadcast import MessageQueue
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, DiffusionOutput
 from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
+from vllm_omni.diffusion.ipc import unpack_diffusion_output_shm
 from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.scheduler import Scheduler
 from vllm_omni.diffusion.worker import WorkerProc
 
 logger = init_logger(__name__)
@@ -21,18 +23,23 @@ class BackgroundResources:
     Used as a finalizer for clean shutdown.
     """
 
-    scheduler: Scheduler | None = None
+    broadcast_mq: MessageQueue | None = None
+    result_mq: MessageQueue | None = None
+    num_workers: int = 0
     processes: list[mp.Process] | None = None
 
     def __call__(self):
         """Clean up background resources."""
-        if self.scheduler is not None:
+        if self.broadcast_mq is not None:
             try:
-                for _ in range(self.scheduler.num_workers):
-                    self.scheduler.mq.enqueue(SHUTDOWN_MESSAGE)
-                self.scheduler.close()
+                for _ in range(self.num_workers):
+                    self.broadcast_mq.enqueue(SHUTDOWN_MESSAGE)
+
+                self.broadcast_mq = None
+                self.result_mq = None
             except Exception as exc:
                 logger.warning("Failed to send shutdown signal: %s", exc)
+
         if self.processes:
             for proc in self.processes:
                 if not proc.is_alive():
@@ -51,24 +58,42 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         self._processes: list[mp.Process] = []
         self._closed = False
 
-        # Initialize scheduler
-        self.scheduler = Scheduler()
-        self.scheduler.initialize(self.od_config)
-        broadcast_handle = self.scheduler.get_broadcast_handle()
+        num_workers = self.od_config.num_gpus
+        self._broadcast_mq = self._init_broadcast_queue(num_workers)
+        broadcast_handle = self._broadcast_mq.export_handle()
 
         # Launch workers
         processes, result_handle = self._launch_workers(broadcast_handle)
-
-        if result_handle is not None:
-            self.scheduler.initialize_result_queue(result_handle)
-        else:
-            logger.error("Failed to get result queue handle from workers")
-
+        self._result_mq = self._init_result_queue(result_handle)
         self._processes = processes
 
-        self.resources = BackgroundResources(scheduler=self.scheduler, processes=self._processes)
+        self.resources = BackgroundResources(
+            broadcast_mq=self._broadcast_mq,
+            result_mq=self._result_mq,
+            num_workers=num_workers,
+            processes=self._processes,
+        )
         self._finalizer = weakref.finalize(self, self.resources)
 
+    def _init_broadcast_queue(self, num_workers: int) -> MessageQueue:
+        return MessageQueue(
+            n_reader=num_workers,
+            n_local_reader=num_workers,
+            local_reader_ranks=list(range(num_workers)),
+        )
+
+    def _init_result_queue(self, result_handle) -> MessageQueue | None:
+        if result_handle is None:
+            logger.error("Failed to get result queue handle from workers")
+            return None
+        return MessageQueue.create_from_handle(result_handle, 0)
+
+    def _ensure_open(self) -> None:
+        if self._closed:
+            raise RuntimeError("DiffusionExecutor is closed.")
+        if self._result_mq is None:
+            raise RuntimeError("Result queue not initialized")
+
     def _launch_workers(self, broadcast_handle):
         od_config = self.od_config
         logger.info("Starting server...")
@@ -134,7 +159,36 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         return processes, result_handle
 
     def add_req(self, request: OmniDiffusionRequest) -> DiffusionOutput:
-        return self.scheduler.add_req(request)
+        self._ensure_open()
+        rpc_request = {
+            "type": "rpc",
+            "method": "generate",
+            "args": (request,),
+            "kwargs": {},
+            "output_rank": 0,
+            "exec_all_ranks": True,
+        }
+
+        try:
+            self._broadcast_mq.enqueue(rpc_request)
+            response = self._result_mq.dequeue()
+
+            try:
+                unpack_diffusion_output_shm(response)
+            except Exception as e:
+                logger.warning("SHM unpack failed (data may already be inline): %s", e)
+
+            if isinstance(response, dict) and response.get("status") == "error":
+                raise RuntimeError(
+                    f"Worker failed with error '{response.get('error')}', "
+                    "please check the stack trace above for the root cause"
+                )
+            if not isinstance(response, DiffusionOutput):
+                raise RuntimeError(f"Unexpected response type for generate: {type(response)!r}")
+            return response
+        except Exception as e:
+            logger.error(f"Generate call failed: {e}")
+            raise
 
     def collective_rpc(
         self,
@@ -144,8 +198,7 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         kwargs: dict | None = None,
         unique_reply_rank: int | None = None,
     ) -> Any:
-        if self._closed:
-            raise RuntimeError("DiffusionExecutor is closed.")
+        self._ensure_open()
 
         deadline = None if timeout is None else time.monotonic() + timeout
         kwargs = kwargs or {}
@@ -160,44 +213,32 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
         }
 
         try:
-            # Acquire lock with timeout awareness so that a stalled add_req
-            # (holding the lock while blocked on dequeue) does not prevent
-            # this RPC from honouring its own timeout.
-            lock_timeout = None if deadline is None else max(0, deadline - time.monotonic())
-            acquired = self.scheduler._lock.acquire(timeout=lock_timeout if lock_timeout is not None else -1)
-            if not acquired:
-                raise TimeoutError(f"RPC call to {method} timed out waiting for scheduler lock.")
-            try:
-                # Broadcast RPC request to all workers via unified message queue
-                self.scheduler.mq.enqueue(rpc_request)
+            # Broadcast RPC request to all workers via unified message queue
+            self._broadcast_mq.enqueue(rpc_request)
 
-                # Determine which workers we expect responses from
-                num_responses = 1 if unique_reply_rank is not None else self.od_config.num_gpus
+            # Determine which workers we expect responses from
+            num_responses = 1 if unique_reply_rank is not None else self.od_config.num_gpus
 
-                responses = []
-                for _ in range(num_responses):
-                    dequeue_timeout = None if deadline is None else max(0, deadline - time.monotonic())
-                    try:
-                        if self.scheduler.result_mq is None:
-                            raise RuntimeError("Result queue not initialized")
+            responses = []
+            for _ in range(num_responses):
+                dequeue_timeout = None if deadline is None else max(0, deadline - time.monotonic())
+                try:
+                    response = self._result_mq.dequeue(timeout=dequeue_timeout)
 
-                        response = self.scheduler.result_mq.dequeue(timeout=dequeue_timeout)
+                    # Check if response indicates an error
+                    if isinstance(response, dict) and response.get("status") == "error":
+                        raise RuntimeError(
+                            f"Worker failed with error '{response.get('error')}', "
+                            "please check the stack trace above for the root cause"
+                        )
 
-                        # Check if response indicates an error
-                        if isinstance(response, dict) and response.get("status") == "error":
-                            raise RuntimeError(
-                                f"Worker failed with error '{response.get('error')}', "
-                                "please check the stack trace above for the root cause"
-                            )
-
-                        responses.append(response)
-                    except TimeoutError as e:
-                        raise TimeoutError(f"RPC call to {method} timed out.") from e
-
-                return responses[0] if unique_reply_rank is not None else responses
-            finally:
-                self.scheduler._lock.release()
+                    responses.append(response)
+                except zmq.error.Again as e:
+                    raise TimeoutError(f"RPC call to {method} timed out.") from e
+                except TimeoutError as e:
+                    raise TimeoutError(f"RPC call to {method} timed out.") from e
 
+            return responses[0] if unique_reply_rank is not None else responses
         except Exception as e:
             logger.error(f"RPC call failed: {e}")
             raise
@@ -210,4 +251,10 @@ class MultiprocDiffusionExecutor(DiffusionExecutor):
 
     def shutdown(self) -> None:
         self._closed = True
-        self._finalizer()
+        try:
+            self._finalizer()
+        finally:
+            self._broadcast_mq = None
+            self._result_mq = None
+            self.resources = None
+            self._processes = []