From aa689e6af0738660ec88324253fbe2aaaca3677e Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Mon, 8 Dec 2025 11:52:07 +0800 Subject: [PATCH] [Core] add clean up method for diffusion engine (#219) Signed-off-by: zjy0516 --- tests/conftest.py | 36 ++++++ tests/single_stage/test_diffusion_model.py | 15 ++- tests/utils.py | 107 ++++++++++++++++++ vllm_omni/diffusion/data.py | 32 ++++++ vllm_omni/diffusion/diffusion_engine.py | 35 +++++- .../models/qwen_image/pipeline_qwen_image.py | 2 +- .../qwen_image/qwen_image_transformer.py | 5 +- .../models/z_image/pipeline_z_image.py | 2 +- vllm_omni/diffusion/omni_diffusion.py | 16 ++- vllm_omni/diffusion/scheduler.py | 6 +- vllm_omni/diffusion/worker/gpu_worker.py | 26 ++++- vllm_omni/diffusion/worker/npu/npu_worker.py | 26 ++++- vllm_omni/entrypoints/omni.py | 11 ++ 13 files changed, 309 insertions(+), 10 deletions(-) create mode 100644 tests/conftest.py create mode 100644 tests/utils.py diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 000000000..82c959f07 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,36 @@ +import os + +import pytest +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +@pytest.fixture(autouse=True) +def clean_gpu_memory_between_tests(): + if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1": + yield + return + + # Wait for GPU memory to be cleared before starting the test + import gc + + from tests.utils import wait_for_gpu_memory_to_clear + + num_gpus = torch.cuda.device_count() + if num_gpus > 0: + try: + wait_for_gpu_memory_to_clear( + devices=list(range(num_gpus)), + threshold_ratio=0.1, + ) + except ValueError as e: + logger.info("Failed to clean GPU memory: %s", e) + + yield + + # Clean up GPU memory after the test + if torch.cuda.is_available(): + torch.cuda.empty_cache() + gc.collect() diff --git a/tests/single_stage/test_diffusion_model.py b/tests/single_stage/test_diffusion_model.py index dc00a36e5..770978366 100644 --- a/tests/single_stage/test_diffusion_model.py +++ b/tests/single_stage/test_diffusion_model.py @@ -1,9 +1,20 @@ +import os +import sys +from pathlib import Path + import pytest import torch +# ruff: noqa: E402 +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + from vllm_omni import Omni -models = ["Tongyi-MAI/Z-Image-Turbo"] +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" + +models = ["Tongyi-MAI/Z-Image-Turbo", "riverclouds/qwen_image_random"] @pytest.mark.parametrize("model_name", models) @@ -25,4 +36,4 @@ def test_diffusion_model(model_name: str): # check image size assert images[0].width == width assert images[0].height == height - images[0].save("z_image_output.png") + images[0].save("image_output.png") diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 000000000..aba734501 --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,107 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +import time +from contextlib import contextmanager + +from vllm.platforms import current_platform + +if current_platform.is_rocm(): + from amdsmi import ( + amdsmi_get_gpu_vram_usage, + amdsmi_get_processor_handles, + amdsmi_init, + amdsmi_shut_down, + ) + + @contextmanager + def _nvml(): + try: + amdsmi_init() + yield + finally: + amdsmi_shut_down() +elif current_platform.is_cuda(): + from vllm.third_party.pynvml import ( + nvmlDeviceGetHandleByIndex, + nvmlDeviceGetMemoryInfo, + nvmlInit, + nvmlShutdown, + ) + + @contextmanager + def _nvml(): + try: + nvmlInit() + yield + finally: + nvmlShutdown() +else: + + @contextmanager + def _nvml(): + yield + + +def get_physical_device_indices(devices): + visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES") + if visible_devices is None: + return devices + + visible_indices = [int(x) for x in visible_devices.split(",")] + index_mapping = {i: physical for i, physical in enumerate(visible_indices)} + return [index_mapping[i] for i in devices if i in index_mapping] + + +@_nvml() +def wait_for_gpu_memory_to_clear( + *, + devices: list[int], + threshold_bytes: int | None = None, + threshold_ratio: float | None = None, + timeout_s: float = 120, +) -> None: + assert threshold_bytes is not None or threshold_ratio is not None + # Use nvml instead of pytorch to reduce measurement error from torch cuda + # context. + devices = get_physical_device_indices(devices) + start_time = time.time() + while True: + output: dict[int, str] = {} + output_raw: dict[int, tuple[float, float]] = {} + for device in devices: + if current_platform.is_rocm(): + dev_handle = amdsmi_get_processor_handles()[device] + mem_info = amdsmi_get_gpu_vram_usage(dev_handle) + gb_used = mem_info["vram_used"] / 2**10 + gb_total = mem_info["vram_total"] / 2**10 + else: + dev_handle = nvmlDeviceGetHandleByIndex(device) + mem_info = nvmlDeviceGetMemoryInfo(dev_handle) + gb_used = mem_info.used / 2**30 + gb_total = mem_info.total / 2**30 + output_raw[device] = (gb_used, gb_total) + output[device] = f"{gb_used:.02f}/{gb_total:.02f}" + + print("gpu memory used/total (GiB): ", end="") + for k, v in output.items(): + print(f"{k}={v}; ", end="") + print("") + + if threshold_bytes is not None: + is_free = lambda used, total: used <= threshold_bytes / 2**30 # noqa E731 + threshold = f"{threshold_bytes / 2**30} GiB" + else: + is_free = lambda used, total: used / total <= threshold_ratio # noqa E731 + threshold = f"{threshold_ratio:.2f}" + + dur_s = time.time() - start_time + if all(is_free(used, total) for used, total in output_raw.values()): + print(f"Done waiting for free GPU memory on devices {devices=} ({threshold=}) {dur_s=:.02f}") + break + + if dur_s >= timeout_s: + raise ValueError(f"Memory of devices {devices=} not free after {dur_s=:.02f} ({threshold=})") + + time.sleep(5) diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index a413753b6..d3af974d2 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -14,6 +14,32 @@ from vllm_omni.diffusion.utils.network_utils import is_port_available logger = init_logger(__name__) +@dataclass +class TransformerConfig: + """Container for raw transformer configuration dictionaries.""" + + params: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> "TransformerConfig": + if not isinstance(data, dict): + raise TypeError(f"Expected transformer config dict, got {type(data)!r}") + return cls(params=dict(data)) + + def to_dict(self) -> dict[str, Any]: + return dict(self.params) + + def get(self, key: str, default: Any | None = None) -> Any: + return self.params.get(key, default) + + def __getattr__(self, item: str) -> Any: + params = object.__getattribute__(self, "params") + try: + return params[item] + except KeyError as exc: + raise AttributeError(item) from exc + + @dataclass class OmniDiffusionConfig: # Model and path configuration (for convenience) @@ -23,6 +49,8 @@ class OmniDiffusionConfig: dtype: torch.dtype = torch.bfloat16 + tf_model_config: TransformerConfig = field(default_factory=TransformerConfig) + # Attention # attention_backend: str = None @@ -214,3 +242,7 @@ class AttentionBackendEnum(enum.Enum): def __str__(self): return self.name.lower() + + +# Special message broadcast via scheduler queues to signal worker shutdown. +SHUTDOWN_MESSAGE = {"type": "shutdown"} diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 05f13afaf..f68247271 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/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() diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 2a27ba1e0..54700fc98 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -267,7 +267,7 @@ class QwenImagePipeline( self.vae = AutoencoderKLQwenImage.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.device ) - self.transformer = QwenImageTransformer2DModel() + self.transformer = QwenImageTransformer2DModel(od_config=od_config) self.tokenizer = Qwen2Tokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.stage = None diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index 5f568d0c6..e62400668 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -19,6 +19,7 @@ from vllm.model_executor.layers.linear import QKVParallelLinear, ReplicatedLinea from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.data import OmniDiffusionConfig logger = init_logger(__name__) @@ -486,6 +487,7 @@ class QwenImageTransformer2DModel(nn.Module): # _repeated_blocks = ["QwenImageTransformerBlock"] def __init__( self, + od_config: OmniDiffusionConfig, patch_size: int = 2, in_channels: int = 64, out_channels: Optional[int] = 16, @@ -497,7 +499,8 @@ class QwenImageTransformer2DModel(nn.Module): axes_dims_rope: tuple[int, int, int] = (16, 56, 56), ): super().__init__() - self.config = None + model_config = od_config.tf_model_config + num_layers = model_config.num_layers self.in_channels = in_channels self.out_channels = out_channels or in_channels self.inner_dim = num_attention_heads * attention_head_dim diff --git a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py index 2579360b4..e9da3690d 100644 --- a/vllm_omni/diffusion/models/z_image/pipeline_z_image.py +++ b/vllm_omni/diffusion/models/z_image/pipeline_z_image.py @@ -146,7 +146,7 @@ class ZImagePipeline(nn.Module): def __init__( self, *, - od_config: OmniDiffusionConfig = None, + od_config: OmniDiffusionConfig, prefix: str = "", ): super().__init__() diff --git a/vllm_omni/diffusion/omni_diffusion.py b/vllm_omni/diffusion/omni_diffusion.py index c75d7da22..9f4382222 100644 --- a/vllm_omni/diffusion/omni_diffusion.py +++ b/vllm_omni/diffusion/omni_diffusion.py @@ -7,7 +7,7 @@ from dataclasses import fields from vllm.logger import init_logger from vllm.transformers_utils.config import get_hf_file_to_dict -from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -53,6 +53,11 @@ class OmniDiffusion: od_config.model, ) od_config.model_class_name = config_dict.get("_class_name", None) + tf_config_dict = get_hf_file_to_dict( + "transformer/config.json", + od_config.model, + ) + od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict) self.engine: DiffusionEngine = DiffusionEngine.make_engine(od_config) @@ -82,3 +87,12 @@ class OmniDiffusion: def _run_engine(self, requests: list[OmniDiffusionRequest]): return self.engine.step(requests) + + def close(self) -> None: + self.engine.close() + + def __del__(self): # pragma: no cover - best effort cleanup + try: + self.close() + except Exception: + pass diff --git a/vllm_omni/diffusion/scheduler.py b/vllm_omni/diffusion/scheduler.py index 57c2c87df..24dbb77a9 100644 --- a/vllm_omni/diffusion/scheduler.py +++ b/vllm_omni/diffusion/scheduler.py @@ -20,7 +20,8 @@ class Scheduler: return cls._instance def initialize(self, od_config: OmniDiffusionConfig): - if hasattr(self, "context") and not self.context.closed: + existing_context = getattr(self, "context", None) + if existing_context is not None and not existing_context.closed: logger.warning("SyncSchedulerClient is already initialized. Re-initializing.") self.close() @@ -65,6 +66,9 @@ class Scheduler: """Closes the socket and terminates the context.""" if hasattr(self, "context"): self.context.term() + self.context = None + self.mq = None + self.result_mq = None # Singleton instance for easy access diff --git a/vllm_omni/diffusion/worker/gpu_worker.py b/vllm_omni/diffusion/worker/gpu_worker.py index 84bf175ad..a3a76a43d 100644 --- a/vllm_omni/diffusion/worker/gpu_worker.py +++ b/vllm_omni/diffusion/worker/gpu_worker.py @@ -15,7 +15,11 @@ from vllm.distributed.parallel_state import ( from vllm.logger import init_logger from vllm.utils import DeviceMemoryProfiler, GiB_bytes -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.data import ( + SHUTDOWN_MESSAGE, + DiffusionOutput, + OmniDiffusionConfig, +) from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -91,6 +95,14 @@ class GPUWorker: output = self.pipeline.forward(req) return output + def shutdown(self) -> None: + if torch.distributed.is_initialized(): + try: + torch.distributed.destroy_process_group() + logger.info("Worker %s: Destroyed process group", self.rank) + except Exception as exc: # pragma: no cover - best effort cleanup + logger.warning("Worker %s: Failed to destroy process group: %s", self.rank, exc) + class WorkerProc: """Wrapper that runs one Worker in a separate process.""" @@ -162,6 +174,14 @@ class WorkerProc: ) continue + if reqs == SHUTDOWN_MESSAGE: + logger.info("Worker %s: Received shutdown message", self.gpu_id) + self._running = False + continue + if reqs is None: + logger.warning("Worker %s: Received empty payload, ignoring", self.gpu_id) + continue + # 2: execute, make sure a reply is always sent try: output = self.worker.execute_model(reqs, self.od_config) @@ -180,6 +200,10 @@ class WorkerProc: continue logger.info("event loop terminated.") + try: + self.worker.shutdown() + except Exception as exc: # pragma: no cover - best effort cleanup + logger.warning("Worker %s: Shutdown encountered an error: %s", self.gpu_id, exc) # if self.result_sender is not None: # self.result_sender.close() self.context.term() diff --git a/vllm_omni/diffusion/worker/npu/npu_worker.py b/vllm_omni/diffusion/worker/npu/npu_worker.py index 79a7e9f87..1bdb076e5 100644 --- a/vllm_omni/diffusion/worker/npu/npu_worker.py +++ b/vllm_omni/diffusion/worker/npu/npu_worker.py @@ -16,7 +16,11 @@ from vllm.distributed.parallel_state import ( from vllm.logger import init_logger from vllm.utils import DeviceMemoryProfiler, GiB_bytes -from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.data import ( + SHUTDOWN_MESSAGE, + DiffusionOutput, + OmniDiffusionConfig, +) from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.request import OmniDiffusionRequest @@ -95,6 +99,14 @@ class NPUWorker: output = self.pipeline.forward(req) return output + def shutdown(self) -> None: + if torch.distributed.is_initialized(): + try: + torch.distributed.destroy_process_group() + logger.info("Worker %s: Destroyed process group", self.rank) + except Exception as exc: # pragma: no cover - best effort cleanup + logger.warning("Worker %s: Failed to destroy process group: %s", self.rank, exc) + class NPUWorkerProc: """Wrapper that runs one Worker in a separate process.""" @@ -166,6 +178,14 @@ class NPUWorkerProc: ) continue + if reqs == SHUTDOWN_MESSAGE: + logger.info("Worker %s: Received shutdown message", self.gpu_id) + self._running = False + continue + if reqs is None: + logger.warning("Worker %s: Received empty payload, ignoring", self.gpu_id) + continue + # 2: execute, make sure a reply is always sent try: output = self.worker.execute_model(reqs, self.od_config) @@ -184,6 +204,10 @@ class NPUWorkerProc: continue logger.info("event loop terminated.") + try: + self.worker.shutdown() + except Exception as exc: # pragma: no cover - best effort cleanup + logger.warning("Worker %s: Shutdown encountered an error: %s", self.gpu_id, exc) # if self.result_sender is not None: # self.result_sender.close() self.context.term() diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index e5ac27aef..320a85805 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -47,3 +47,14 @@ class Omni: if hasattr(self.instance, "generate"): return getattr(self.instance, "generate")(*args, **kwargs) raise AttributeError(f"'{self.instance.__class__.__name__}' has no attribute 'generate'") + + def close(self) -> None: + close_method = getattr(self.instance, "close", None) + if callable(close_method): + close_method() + + def __del__(self): # pragma: no cover - best effort cleanup + try: + self.close() + except Exception: + pass