diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 0f21d4b28..731ed425e 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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 diff --git a/vllm_omni/diffusion/executor/__init__.py b/vllm_omni/diffusion/executor/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/diffusion/executor/abstract.py b/vllm_omni/diffusion/executor/abstract.py new file mode 100644 index 000000000..021c210ef --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/executor/multiproc_executor.py b/vllm_omni/diffusion/executor/multiproc_executor.py new file mode 100644 index 000000000..a04a40e86 --- /dev/null +++ b/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() diff --git a/vllm_omni/diffusion/scheduler.py b/vllm_omni/diffusion/scheduler.py index 7f7902b80..386218a61 100644 --- a/vllm_omni/diffusion/scheduler.py +++ b/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()