[BugFix] __del__ exit don't join the spawned process (#555)

Signed-off-by: yinpeiqi <yinpeiqi809@gmail.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/e2e/offline_inference/test_cache_dit.py

@@ -78,5 +78,3 @@ def test_cache_dit(model_name: str):
     # Check image size
     assert images[0].width == width
     assert images[0].height == height
-    # manually close the Omni instance
-    m.close()

tests/e2e/offline_inference/test_t2i_model.py

@@ -63,5 +63,3 @@ def test_diffusion_model(model_name: str):
     assert images[0].width == width
     assert images[0].height == height
     images[0].save("image_output.png")
-    # manually close the Omni instance
-    m.close()

tests/e2e/offline_inference/test_t2v_model.py

@@ -56,5 +56,3 @@ def test_video_diffusion_model(model_name: str):
     assert frames.shape[1] == num_frames
     assert frames.shape[2] == height
     assert frames.shape[3] == width
-    # manually close the Omni instance
-    m.close()

tests/e2e/offline_inference/test_teacache.py

@@ -74,5 +74,3 @@ def test_teacache(model_name: str):
     # Check image size
     assert images[0].width == width
     assert images[0].height == height
-    # manually close the Omni instance
-    m.close()

vllm_omni/entrypoints/async_omni.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import asyncio
 import time
+import weakref
 from collections.abc import AsyncGenerator, Iterable
 from dataclasses import asdict
 from pprint import pformat
@@ -20,6 +21,7 @@ from vllm.v1.engine.exceptions import EngineDeadError
 from vllm_omni.config import OmniModelConfig
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.distributed.omni_connectors.adapter import try_send_via_connector
+from vllm_omni.distributed.ray_utils.utils import try_close_ray
 from vllm_omni.engine.input_processor import OmniInputProcessor
 from vllm_omni.entrypoints.client_request_state import ClientRequestState
 from vllm_omni.entrypoints.log_utils import (
@@ -36,6 +38,25 @@ from vllm_omni.outputs import OmniRequestOutput
 logger = init_logger(__name__)
 
 
+def _weak_close_cleanup_async(stage_list, stage_in_queues, ray_pg, output_handler):
+    """Weak reference cleanup function for AsyncOmni instances."""
+    if stage_list:
+        for q in stage_in_queues:
+            try:
+                q.put_nowait(None)
+            except Exception as e:
+                logger.warning(f"Failed to send shutdown signal to stage input queue: {e}")
+        for stage in stage_list:
+            try:
+                stage.stop_stage_worker()
+            except Exception as e:
+                logger.warning(f"Failed to stop stage worker: {e}")
+    try_close_ray(ray_pg)
+    # Cancel output handler
+    if output_handler is not None:
+        output_handler.cancel()
+
+
 class AsyncOmni(OmniBase):
     """Asynchronous unified entry point supporting multi-stage pipelines for LLM and Diffusion models.
 
@@ -83,6 +104,16 @@ class AsyncOmni(OmniBase):
 
         super().__init__(*args, **kwargs)
 
+        # Register weak reference cleanup (called on garbage collection)
+        self._weak_finalizer = weakref.finalize(
+            self,
+            _weak_close_cleanup_async,
+            self.stage_list,
+            self._stage_in_queues,
+            self._ray_pg,
+            self.output_handler,
+        )
+
     def _create_default_diffusion_stage_cfg(self, kwargs: dict[str, Any]) -> dict[str, Any]:
         """Create default diffusion stage configuration."""
         # TODO: here is different from the Omni class. We should merge the two in the future.
@@ -187,23 +218,14 @@ class AsyncOmni(OmniBase):
             self.io_processor = None
             self.model_config = None
 
-    def close(self) -> None:
-        """Close all stage processes and clean up resources."""
-        if self.output_handler is not None:
-            self.output_handler.cancel()
-            self.output_handler = None
-        super().close()
-
     def shutdown(self):
         """Shutdown, cleaning up the background proc and IPC.
 
         Alias for close() method. Cleans up all stage processes
         and inter-process communication resources.
         """
-        try:
-            self.close()
-        except Exception as e:
-            logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True)
+        if hasattr(self, "_weak_finalizer"):
+            self._weak_finalizer()
 
     async def generate(self, *args: Any, **kwargs: dict[str, Any]) -> AsyncGenerator[OmniRequestOutput, None]:
         """Generate outputs for the given prompt asynchronously.

vllm_omni/entrypoints/omni.py

@@ -5,6 +5,7 @@ import multiprocessing as mp
 import os
 import time
 import uuid
+import weakref
 from collections.abc import Sequence
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import asdict
@@ -40,6 +41,22 @@ from vllm_omni.outputs import OmniRequestOutput
 logger = init_logger(__name__)
 
 
+def _weak_close_cleanup(stage_list, stage_in_queues, ray_pg):
+    """Weak reference cleanup function for OmniBase instances."""
+    if stage_list:
+        for q in stage_in_queues:
+            try:
+                q.put_nowait(None)
+            except Exception as e:
+                logger.warning(f"Failed to send shutdown signal to stage input queue: {e}")
+        for stage in stage_list:
+            try:
+                stage.stop_stage_worker()
+            except Exception as e:
+                logger.warning(f"Failed to stop stage worker: {e}")
+    try_close_ray(ray_pg)
+
+
 def _dummy_snapshot_download(model_id):
     return model_id
 
@@ -316,28 +333,8 @@ class OmniBase:
 
     def close(self) -> None:
         """Close all stage processes and clean up resources."""
-        # Close stages if they exist (for LLM models)
-        if self.stage_list:
-            for q in self._stage_in_queues:
-                try:
-                    q.put_nowait(None)
-                except Exception as e:
-                    logger.warning(
-                        f"[{self._name}] Failed to send shutdown signal to stage input queue: {e}",
-                    )
-            for stage in self.stage_list:
-                try:
-                    stage.stop_stage_worker()
-                except Exception as e:
-                    logger.warning(f"[{self._name}] Failed to stop stage worker: {e}")
-
-            try_close_ray(self._ray_pg)
-
-    def __del__(self):  # pragma: no cover - best effort cleanup
-        try:
-            self.close()
-        except Exception:
-            logger.debug(f"[{self._name}] __del__ close() raised", exc_info=True)
+        if hasattr(self, "_weak_finalizer"):
+            self._weak_finalizer()
 
     @property
     def _name(self) -> str:
@@ -380,6 +377,15 @@ class Omni(OmniBase):
     def __init__(self, *args: Any, **kwargs: dict[str, Any]) -> None:
         super().__init__(*args, **kwargs)
 
+        # Register weak reference cleanup (called on garbage collection)
+        self._weak_finalizer = weakref.finalize(
+            self,
+            _weak_close_cleanup,
+            self.stage_list,
+            self._stage_in_queues,
+            self._ray_pg,
+        )
+
     def generate(self, *args: Any, **kwargs: dict[str, Any]) -> list[OmniRequestOutput]:
         """Generate outputs for the given prompts.