From 2eebd22e3227734d17565a9fcd5d5b20630c5729 Mon Sep 17 00:00:00 2001 From: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com> Date: Wed, 31 Dec 2025 03:59:25 +0800 Subject: [PATCH] [BugFix] __del__ exit don't join the spawned process (#555) Signed-off-by: yinpeiqi Co-authored-by: Hongsheng Liu --- tests/e2e/offline_inference/test_cache_dit.py | 2 - tests/e2e/offline_inference/test_t2i_model.py | 2 - tests/e2e/offline_inference/test_t2v_model.py | 2 - tests/e2e/offline_inference/test_teacache.py | 2 - vllm_omni/entrypoints/async_omni.py | 44 ++++++++++++---- vllm_omni/entrypoints/omni.py | 50 +++++++++++-------- 6 files changed, 61 insertions(+), 41 deletions(-) diff --git a/tests/e2e/offline_inference/test_cache_dit.py b/tests/e2e/offline_inference/test_cache_dit.py index 9f1a6ee07..01b41342e 100644 --- a/tests/e2e/offline_inference/test_cache_dit.py +++ b/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() diff --git a/tests/e2e/offline_inference/test_t2i_model.py b/tests/e2e/offline_inference/test_t2i_model.py index d448b920c..ffdd7a463 100644 --- a/tests/e2e/offline_inference/test_t2i_model.py +++ b/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() diff --git a/tests/e2e/offline_inference/test_t2v_model.py b/tests/e2e/offline_inference/test_t2v_model.py index 2e81c87fa..1dc2e1ff1 100644 --- a/tests/e2e/offline_inference/test_t2v_model.py +++ b/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() diff --git a/tests/e2e/offline_inference/test_teacache.py b/tests/e2e/offline_inference/test_teacache.py index 55ed072a2..2b737911a 100644 --- a/tests/e2e/offline_inference/test_teacache.py +++ b/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() diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 2a7bbd3c4..597248c7e 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/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. diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index e1af5cf6c..ba0099f2f 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/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.