[Core] add clean up method for diffusion engine (#219)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>

vllm_omni/diffusion/diffusion_engine.py

@@ -5,7 +5,7 @@ import multiprocessing as mp
 
 from vllm.logger import init_logger
 
-from vllm_omni.diffusion.data import OmniDiffusionConfig
+from vllm_omni.diffusion.data import SHUTDOWN_MESSAGE, OmniDiffusionConfig
 from vllm_omni.diffusion.registry import get_diffusion_post_process_func
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.diffusion.scheduler import scheduler
@@ -27,6 +27,8 @@ class DiffusionEngine:
 
         self.post_process_func = get_diffusion_post_process_func(od_config)
 
+        self._processes: list[mp.Process] = []
+        self._closed = False
         self._make_client()
 
     def step(self, requests: list[OmniDiffusionRequest]):
@@ -68,6 +70,8 @@ class DiffusionEngine:
         else:
             logger.error("Failed to get result queue handle from workers")
 
+        self._processes = processes
+
     def _launch_workers(self, broadcast_handle):
         od_config = self.od_config
         logger.info("Starting server...")
@@ -131,3 +135,32 @@ class DiffusionEngine:
 
     def add_req_and_wait_for_response(self, requests: list[OmniDiffusionRequest]):
         return scheduler.add_req(requests)
+
+    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()