# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import multiprocessing as mp import time import weakref from collections.abc import Callable, Iterable from dataclasses import dataclass 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.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__) def supports_image_input(model_class_name: str) -> bool: model_cls = DiffusionModelRegistry._try_load_model_cls(model_class_name) if model_cls is None: return False 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.""" def __init__(self, od_config: OmniDiffusionConfig): """Initialize the diffusion engine. Args: config: The configuration for the diffusion engine. """ self.od_config = od_config 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() try: self._dummy_run() except Exception as e: logger.error(f"Dummy run failed: {e}") self.close() raise e def step(self, requests: list[OmniDiffusionRequest]): try: # Apply pre-processing if available if self.pre_process_func is not None: preprocess_start_time = time.time() requests = self.pre_process_func(requests) preprocess_time = time.time() - preprocess_start_time logger.info(f"Pre-processing completed in {preprocess_time:.4f} seconds") output = self.add_req_and_wait_for_response(requests) if output.error: raise Exception(f"{output.error}") logger.info("Generation completed successfully.") if output.output is None: logger.warning("Output is None, returning empty OmniRequestOutput") # Return empty output for the first request if len(requests) > 0: request = requests[0] request_id = request.request_id or "" prompt = request.prompt if isinstance(prompt, list): prompt = prompt[0] if prompt else None return OmniRequestOutput.from_diffusion( request_id=request_id, images=[], prompt=prompt, metrics={}, latents=None, ) return None postprocess_start_time = time.time() images = self.post_process_func(output.output) if self.post_process_func is not None else output.output postprocess_time = time.time() - postprocess_start_time logger.info(f"Post-processing completed in {postprocess_time:.4f} seconds") # Convert to OmniRequestOutput format # Ensure images is a list if not isinstance(images, list): images = [images] if images is not None else [] # Handle single request or multiple requests if len(requests) == 1: # Single request: return single OmniRequestOutput request = requests[0] request_id = request.request_id or "" prompt = request.prompt if isinstance(prompt, list): prompt = prompt[0] if prompt else None metrics = {} if output.trajectory_timesteps is not None: metrics["trajectory_timesteps"] = output.trajectory_timesteps return OmniRequestOutput.from_diffusion( request_id=request_id, images=images, prompt=prompt, metrics=metrics, latents=output.trajectory_latents, ) else: # Multiple requests: return list of OmniRequestOutput # Split images based on num_outputs_per_prompt for each request results = [] image_idx = 0 for request in requests: request_id = request.request_id or "" prompt = request.prompt if isinstance(prompt, list): prompt = prompt[0] if prompt else None # Get images for this request num_outputs = request.num_outputs_per_prompt request_images = images[image_idx : image_idx + num_outputs] if image_idx < len(images) else [] image_idx += num_outputs metrics = {} if output.trajectory_timesteps is not None: metrics["trajectory_timesteps"] = output.trajectory_timesteps results.append( OmniRequestOutput.from_diffusion( request_id=request_id, images=request_images, prompt=prompt, metrics=metrics, latents=output.trajectory_latents, ) ) return results except Exception as e: logger.error(f"Generation failed: {e}") return None @staticmethod def make_engine(config: OmniDiffusionConfig) -> "DiffusionEngine": """Factory method to create a DiffusionEngine instance. Args: config: The configuration for the diffusion engine. Returns: An instance of 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) def _dummy_run(self): """A dummy run to warm up the model.""" prompt = "dummy run" # note that num_inference_steps=1 will cause timestep and temb None in the pipeline num_inference_steps = 1 height = 1024 width = 1024 if supports_image_input(self.od_config.model_class_name): # Provide a dummy image input if the model supports it dummy_image = PIL.Image.new("RGB", (width, height), color=(0, 0, 0)) else: dummy_image = None req = OmniDiffusionRequest( prompt=prompt, height=height, width=width, pil_image=dummy_image, num_inference_steps=num_inference_steps, num_outputs_per_prompt=1, ) logger.info("dummy run to warm up the model") requests = self.pre_process_func([req]) if self.pre_process_func is not None else [req] self.add_req_and_wait_for_response(requests) def collective_rpc( self, method: str | Callable, timeout: float | None = None, args: tuple = (), kwargs: dict | None = None, unique_reply_rank: int | None = None, ) -> Any: """Call a method on worker processes and get results immediately. Args: method: The method name (str) or callable to execute on workers timeout: Optional timeout in seconds args: Positional arguments for the method kwargs: Keyword arguments for the method unique_reply_rank: If set, only get reply from this rank 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 def close(self) -> None: self._finalizer() def abort(self, request_id: str | Iterable[str]) -> None: # TODO implement it logger.warning("DiffusionEngine abort is not implemented yet") pass