[Core]Add Diffusion executor (#865)

Signed-off-by: wzliu <wzliu@connect.hku.hk>

vllm_omni/diffusion/diffusion_engine.py

@@ -1,27 +1,23 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-import multiprocessing as mp
 import os
 import time
-import weakref
-from collections.abc import Callable, Iterable
-from dataclasses import dataclass
+from collections.abc import Iterable
 from typing import Any
 
 import PIL.Image
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, OmniDiffusionConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
 from vllm_omni.diffusion.registry import (
     DiffusionModelRegistry,
     get_diffusion_post_process_func,
     get_diffusion_pre_process_func,
 )
 from vllm_omni.diffusion.request import OmniDiffusionRequest
-from vllm_omni.diffusion.scheduler import Scheduler, scheduler
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import get_diffusion_worker_class
 
 logger = init_logger(__name__)
 
@@ -33,39 +29,6 @@ def supports_image_input(model_class_name: str) -> bool:
     return bool(getattr(model_cls, "support_image_input", False))
 
 
-@dataclass
-class BackgroundResources:
-    """
-    Used as a finalizer for clean shutdown.
-    Create a BackgroundResources instance to encapsulate all background resources
-    (e.g., the scheduler and worker processes) that need explicit cleanup.
-    This object holds references to external system resources that are not managed
-    by Python's garbage collector (like OS processes, message queues, etc.),
-    so they must be cleaned up manually to avoid resource leaks or zombie processes.
-    """
-
-    scheduler: Scheduler | None = None
-    processes: list[mp.Process] | None = None
-
-    def __call__(self):
-        """Clean up background resources."""
-        if scheduler is not None:
-            try:
-                for _ in range(scheduler.num_workers):
-                    scheduler.mq.enqueue(SHUTDOWN_MESSAGE)
-                scheduler.close()
-            except Exception as exc:
-                logger.warning("Failed to send shutdown signal: %s", exc)
-        for proc in self.processes:
-            if not proc.is_alive():
-                continue
-            proc.join(30)
-            if proc.is_alive():
-                logger.warning("Terminating diffusion worker %s after timeout", proc.name)
-                proc.terminate()
-                proc.join(30)
-
-
 class DiffusionEngine:
     """The diffusion engine for vLLM-Omni diffusion models."""
 
@@ -80,9 +43,9 @@ class DiffusionEngine:
         self.post_process_func = get_diffusion_post_process_func(od_config)
         self.pre_process_func = get_diffusion_pre_process_func(od_config)
 
-        self._processes: list[mp.Process] = []
-        self._closed = False
-        self._make_client()
+        executor_class = DiffusionExecutor.get_class(od_config)
+        self.executor = executor_class(od_config)
+
         try:
             self._dummy_run()
         except Exception as e:
@@ -200,96 +163,8 @@ class DiffusionEngine:
         """
         return DiffusionEngine(config)
 
-    def _make_client(self):
-        # TODO rename it
-        scheduler.initialize(self.od_config)
-
-        # Get the broadcast handle from the initialized scheduler
-        broadcast_handle = scheduler.get_broadcast_handle()
-
-        processes, result_handle = self._launch_workers(
-            broadcast_handle=broadcast_handle,
-        )
-
-        if result_handle is not None:
-            scheduler.initialize_result_queue(result_handle)
-        else:
-            logger.error("Failed to get result queue handle from workers")
-
-        self._processes = processes
-
-        self.resources = BackgroundResources(scheduler=scheduler, processes=self._processes)
-        # Use weakref.finalize instead of __del__ or relying on self.close() at shutdown.
-        # During interpreter shutdown, global state (e.g., modules, built-ins) may already
-        # be cleared (set to None), so calling normal cleanup methods can fail with
-        # AttributeError: 'NoneType' object has no attribute '...'.
-        # weakref.finalize schedules cleanup *before* such destruction begins,
-        # ensuring resources are released while the runtime environment is still intact.
-        self._finalizer = weakref.finalize(self, self.resources)
-
-    def _launch_workers(self, broadcast_handle):
-        od_config = self.od_config
-        logger.info("Starting server...")
-
-        num_gpus = od_config.num_gpus
-        mp.set_start_method("spawn", force=True)
-        processes = []
-
-        # Get the appropriate worker class for current device
-        worker_proc = get_diffusion_worker_class()
-
-        # Launch all worker processes
-        scheduler_pipe_readers = []
-        scheduler_pipe_writers = []
-
-        for i in range(num_gpus):
-            reader, writer = mp.Pipe(duplex=False)
-            scheduler_pipe_writers.append(writer)
-            process = mp.Process(
-                target=worker_proc.worker_main,
-                args=(
-                    i,  # rank
-                    od_config,
-                    writer,
-                    broadcast_handle,
-                ),
-                name=f"DiffusionWorker-{i}",
-                daemon=True,
-            )
-            scheduler_pipe_readers.append(reader)
-            process.start()
-            processes.append(process)
-
-        # Wait for all workers to be ready
-        scheduler_infos = []
-        result_handle = None
-        for writer in scheduler_pipe_writers:
-            writer.close()
-
-        for i, reader in enumerate(scheduler_pipe_readers):
-            try:
-                data = reader.recv()
-            except EOFError:
-                logger.error(f"Rank {i} scheduler is dead. Please check if there are relevant logs.")
-                processes[i].join()
-                logger.error(f"Exit code: {processes[i].exitcode}")
-                raise
-
-            if data["status"] != "ready":
-                raise RuntimeError("Initialization failed. Please see the error messages above.")
-
-            if i == 0:
-                result_handle = data.get("result_handle")
-
-            scheduler_infos.append(data)
-            reader.close()
-
-        logger.debug("All workers are ready")
-
-        return processes, result_handle
-
     def add_req_and_wait_for_response(self, requests: list[OmniDiffusionRequest]):
-        return scheduler.add_req(requests)
+        return self.executor.add_req(requests)
 
     def start_profile(self, trace_filename: str | None = None) -> None:
         """
@@ -437,7 +312,7 @@ class DiffusionEngine:
 
     def collective_rpc(
         self,
-        method: str | Callable,
+        method: str,
         timeout: float | None = None,
         args: tuple = (),
         kwargs: dict | None = None,
@@ -446,7 +321,7 @@ class DiffusionEngine:
         """Call a method on worker processes and get results immediately.
 
         Args:
-            method: The method name (str) or callable to execute on workers
+            method: The method name (str) to execute on workers
             timeout: Optional timeout in seconds
             args: Positional arguments for the method
             kwargs: Keyword arguments for the method
@@ -455,59 +330,18 @@ class DiffusionEngine:
         Returns:
             Single result if unique_reply_rank is provided, otherwise list of results
         """
-        if self._closed:
-            raise RuntimeError("DiffusionEngine is closed.")
-
-        deadline = None if timeout is None else time.monotonic() + timeout
-        kwargs = kwargs or {}
-
-        assert isinstance(method, str)
-        send_method = method
-
-        # Prepare RPC request message
-        rpc_request = {
-            "type": "rpc",
-            "method": send_method,
-            "args": args,
-            "kwargs": kwargs,
-            "output_rank": unique_reply_rank,
-        }
-
-        try:
-            # Broadcast RPC request to all workers via unified message queue
-            scheduler.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
-
-            responses = []
-            for _ in range(num_responses):
-                dequeue_timeout = None if deadline is None else (deadline - time.monotonic())
-                try:
-                    if scheduler.result_mq is None:
-                        raise RuntimeError("Result queue not initialized")
-
-                    response = 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"
-                        )
-
-                    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
-
-        except Exception as e:
-            logger.error(f"RPC call failed: {e}")
-            raise
+        assert isinstance(method, str), "Only string method names are supported for now"
+        return self.executor.collective_rpc(
+            method=method,
+            timeout=timeout,
+            args=args,
+            kwargs=kwargs,
+            unique_reply_rank=unique_reply_rank,
+        )
 
     def close(self) -> None:
-        self._finalizer()
+        if hasattr(self, "executor"):
+            self.executor.shutdown()
 
     def abort(self, request_id: str | Iterable[str]) -> None:
         # TODO implement it

vllm_omni/diffusion/executor/abstract.py

@@ -0,0 +1,86 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+from vllm.utils.import_utils import resolve_obj_by_qualname
+
+from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+
+
+class DiffusionExecutor(ABC):
+    """Abstract base class for Diffusion executors."""
+
+    uses_multiproc: bool = False
+
+    @staticmethod
+    def get_class(od_config: OmniDiffusionConfig) -> type["DiffusionExecutor"]:
+        executor_class: type[DiffusionExecutor]
+        distributed_executor_backend = od_config.distributed_executor_backend
+
+        if isinstance(distributed_executor_backend, type):
+            if not issubclass(distributed_executor_backend, DiffusionExecutor):
+                raise TypeError(
+                    "distributed_executor_backend must be a subclass of "
+                    f"DiffusionExecutor. Got {distributed_executor_backend}."
+                )
+            executor_class = distributed_executor_backend
+        elif distributed_executor_backend == "ray":
+            raise NotImplementedError("ray backend is not yet supported.")
+        elif distributed_executor_backend == "mp":
+            from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
+
+            executor_class = MultiprocDiffusionExecutor
+        elif distributed_executor_backend == "external_launcher":
+            raise NotImplementedError("external_launcher backend is not yet supported.")
+        elif isinstance(distributed_executor_backend, str):
+            try:
+                executor_class = resolve_obj_by_qualname(distributed_executor_backend)
+            except (ImportError, ValueError) as e:
+                raise ValueError(
+                    f"Failed to load executor backend '{distributed_executor_backend}'. "
+                    f"Ensure it is a valid python path. Error: {e}"
+                ) from e
+
+            if not issubclass(executor_class, DiffusionExecutor):
+                raise TypeError(
+                    f"distributed_executor_backend must be a subclass of DiffusionExecutor. Got {executor_class}."
+                )
+        else:
+            raise ValueError(f"Unknown distributed executor backend: {distributed_executor_backend}")
+        return executor_class
+
+    def __init__(self, od_config: OmniDiffusionConfig):
+        self.od_config = od_config
+        self._init_executor()
+
+    @abstractmethod
+    def _init_executor(self) -> None:
+        """Initialize the executor (e.g., launch workers, setup IPC)."""
+        pass
+
+    @abstractmethod
+    def add_req(self, requests: list[OmniDiffusionRequest]):
+        """Add requests to the execution queue."""
+        pass
+
+    @abstractmethod
+    def collective_rpc(
+        self,
+        method: str,
+        timeout: float | None = None,
+        args: tuple = (),
+        kwargs: dict | None = None,
+        unique_reply_rank: int | None = None,
+    ) -> Any:
+        """Execute a method on workers."""
+        pass
+
+    @abstractmethod
+    def check_health(self) -> None:
+        """Check if the executor and workers are healthy."""
+        pass
+
+    @abstractmethod
+    def shutdown(self) -> None:
+        """Shutdown the executor and release resources."""
+        pass

vllm_omni/diffusion/executor/multiproc_executor.py

@@ -0,0 +1,200 @@
+import multiprocessing as mp
+import time
+import weakref
+from dataclasses import dataclass
+from typing import Any
+
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE
+from vllm_omni.diffusion.executor.abstract import DiffusionExecutor
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.scheduler import Scheduler
+from vllm_omni.utils.platform_utils import get_diffusion_worker_class
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class BackgroundResources:
+    """
+    Used as a finalizer for clean shutdown.
+    """
+
+    scheduler: Scheduler | None = None
+    processes: list[mp.Process] | None = None
+
+    def __call__(self):
+        """Clean up background resources."""
+        if self.scheduler is not None:
+            try:
+                for _ in range(self.scheduler.num_workers):
+                    self.scheduler.mq.enqueue(SHUTDOWN_MESSAGE)
+                self.scheduler.close()
+            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():
+                    continue
+                proc.join(30)
+                if proc.is_alive():
+                    logger.warning("Terminating diffusion worker %s after timeout", proc.name)
+                    proc.terminate()
+                    proc.join(30)
+
+
+class MultiprocDiffusionExecutor(DiffusionExecutor):
+    uses_multiproc: bool = True
+
+    def _init_executor(self) -> None:
+        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()
+
+        # 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._processes = processes
+
+        self.resources = BackgroundResources(scheduler=self.scheduler, processes=self._processes)
+        self._finalizer = weakref.finalize(self, self.resources)
+
+    def _launch_workers(self, broadcast_handle):
+        od_config = self.od_config
+        logger.info("Starting server...")
+
+        num_gpus = od_config.num_gpus
+        mp.set_start_method("spawn", force=True)
+        processes = []
+
+        # Get the appropriate worker class for current device
+        worker_proc = get_diffusion_worker_class()
+
+        # Launch all worker processes
+        scheduler_pipe_readers = []
+        scheduler_pipe_writers = []
+
+        for i in range(num_gpus):
+            reader, writer = mp.Pipe(duplex=False)
+            scheduler_pipe_writers.append(writer)
+            process = mp.Process(
+                target=worker_proc.worker_main,
+                args=(
+                    i,  # rank
+                    od_config,
+                    writer,
+                    broadcast_handle,
+                ),
+                name=f"DiffusionWorker-{i}",
+                daemon=True,
+            )
+            scheduler_pipe_readers.append(reader)
+            process.start()
+            processes.append(process)
+
+        # Wait for all workers to be ready
+        scheduler_infos = []
+        result_handle = None
+        for writer in scheduler_pipe_writers:
+            writer.close()
+
+        for i, reader in enumerate(scheduler_pipe_readers):
+            try:
+                data = reader.recv()
+            except EOFError:
+                logger.error(f"Rank {i} scheduler is dead. Please check if there are relevant logs.")
+                processes[i].join()
+                logger.error(f"Exit code: {processes[i].exitcode}")
+                raise
+
+            if data["status"] != "ready":
+                raise RuntimeError("Initialization failed. Please see the error messages above.")
+
+            if i == 0:
+                result_handle = data.get("result_handle")
+
+            scheduler_infos.append(data)
+            reader.close()
+
+        logger.debug("All workers are ready")
+
+        return processes, result_handle
+
+    def add_req(self, requests: list[OmniDiffusionRequest]):
+        return self.scheduler.add_req(requests)
+
+    def collective_rpc(
+        self,
+        method: str,
+        timeout: float | None = None,
+        args: tuple = (),
+        kwargs: dict | None = None,
+        unique_reply_rank: int | None = None,
+    ) -> Any:
+        if self._closed:
+            raise RuntimeError("DiffusionExecutor is closed.")
+
+        deadline = None if timeout is None else time.monotonic() + timeout
+        kwargs = kwargs or {}
+
+        # Prepare RPC request message
+        rpc_request = {
+            "type": "rpc",
+            "method": method,
+            "args": args,
+            "kwargs": kwargs,
+            "output_rank": unique_reply_rank,
+        }
+
+        try:
+            # Broadcast RPC request to all workers via unified message queue
+            self.scheduler.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
+
+            responses = []
+            for _ in range(num_responses):
+                dequeue_timeout = None if deadline is None else (deadline - time.monotonic())
+                try:
+                    if self.scheduler.result_mq is None:
+                        raise RuntimeError("Result queue not initialized")
+
+                    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"
+                        )
+
+                    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
+
+        except Exception as e:
+            logger.error(f"RPC call failed: {e}")
+            raise
+
+    def check_health(self) -> None:
+        # Simple check if processes are alive
+        for p in self._processes:
+            if not p.is_alive():
+                raise RuntimeError(f"Worker process {p.name} is dead")
+
+    def shutdown(self) -> None:
+        self._closed = True
+        self._finalizer()

vllm_omni/diffusion/scheduler.py

@@ -12,13 +12,6 @@ logger = init_logger(__name__)
 
 
 class Scheduler:
-    _instance = None
-
-    def __new__(cls, *args, **kwargs):
-        if not cls._instance:
-            cls._instance = super().__new__(cls)
-        return cls._instance
-
     def initialize(self, od_config: OmniDiffusionConfig):
         existing_context = getattr(self, "context", None)
         if existing_context is not None and not existing_context.closed:
@@ -81,7 +74,3 @@ class Scheduler:
         self.context = None
         self.mq = None
         self.result_mq = None
-
-
-# Singleton instance for easy access
-scheduler = Scheduler()