[BugFix] Use weakref.finalize to cleanup DiffusionEngine (#461)

Signed-off-by: iwzbi <iwzbi@zju.edu.cn>
Co-authored-by: iwzbi <iwzbi@zju.edu.cn>

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()