diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 8e65a5b2d..41dda9e6f 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -3,18 +3,53 @@ import multiprocessing as mp import time +import weakref +from dataclasses import dataclass from vllm.logger import init_logger from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, OmniDiffusionConfig from vllm_omni.diffusion.registry import get_diffusion_post_process_func, get_diffusion_pre_process_func from vllm_omni.diffusion.request import OmniDiffusionRequest -from vllm_omni.diffusion.scheduler import scheduler +from vllm_omni.diffusion.scheduler import Scheduler, 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. + 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.""" @@ -87,6 +122,15 @@ class DiffusionEngine: 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...") @@ -167,31 +211,5 @@ class DiffusionEngine: logger.info("dummy run to warm up the model") self.add_req_and_wait_for_response([req]) - def close(self, *, timeout_s: float = 30.0) -> None: - if self._closed: - return - self._closed = True - - # Send shutdown signal to worker processes via broadcast queue - try: - if getattr(scheduler, "mq", None) is not None: - for _ in range(self.od_config.num_gpus or 1): - scheduler.mq.enqueue(SHUTDOWN_MESSAGE) - except Exception as exc: # pragma: no cover - best effort cleanup - logger.warning("Failed to send shutdown signal: %s", exc) - - # Join all worker processes, terminate if they refuse to exit - for proc in self._processes: - if not proc.is_alive(): - continue - proc.join(timeout_s) - if proc.is_alive(): - logger.warning("Terminating diffusion worker %s after timeout", proc.name) - proc.terminate() - proc.join(timeout_s) - - scheduler.close() - self._processes = [] - - def __del__(self): # pragma: no cover - best effort cleanup - self.close() + def close(self) -> None: + self._finalizer() diff --git a/vllm_omni/diffusion/scheduler.py b/vllm_omni/diffusion/scheduler.py index 24dbb77a9..4d45812b5 100644 --- a/vllm_omni/diffusion/scheduler.py +++ b/vllm_omni/diffusion/scheduler.py @@ -25,15 +25,16 @@ class Scheduler: logger.warning("SyncSchedulerClient is already initialized. Re-initializing.") self.close() + self.num_workers = od_config.num_gpus self.od_config = od_config self.context = zmq.Context() # Standard synchronous context # Initialize MessageQueue for broadcasting requests # Assuming all readers are local for now as per current launch_engine implementation self.mq = MessageQueue( - n_reader=od_config.num_gpus, - n_local_reader=od_config.num_gpus, - local_reader_ranks=list(range(od_config.num_gpus)), + n_reader=self.num_workers, + n_local_reader=self.num_workers, + local_reader_ranks=list(range(self.num_workers)), ) self.result_mq = None