diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index b31a8aa0a..d49f3304b 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/tests/diffusion/test_worker_wrapper_base.py @@ -79,6 +79,33 @@ class MockCustomPipeline: class TestWorkerWrapperBaseInitialization: """Test WorkerWrapperBase initialization behavior.""" + def test_worker_init_failure_after_device_setup_shuts_down(self, mocker: MockerFixture): + """A post-device constructor failure must unwind distributed state.""" + + class FakeRunner: + def __init__(self, **kwargs): + self.kwargs = kwargs + + config = SimpleNamespace( + stage_id=0, + diffusion_model_runner_cls="tests.FakeRunner", + diffusion_load_format="default", + ) + init_device = mocker.patch.object(DiffusionWorker, "init_device") + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.resolve_obj_by_qualname", + return_value=FakeRunner, + ) + mocker.patch.object(DiffusionWorker, "_create_profiler", return_value=None) + mocker.patch.object(DiffusionWorker, "load_model", side_effect=RuntimeError("load failed")) + shutdown = mocker.patch.object(DiffusionWorker, "shutdown") + + with pytest.raises(RuntimeError, match="load failed"): + DiffusionWorker(local_rank=0, rank=0, od_config=config) + + init_device.assert_called_once_with() + shutdown.assert_called_once_with() + def test_basic_initialization(self, mocker: MockerFixture, mock_od_config): """Test basic initialization without extensions.""" mock_worker_init = mocker.patch.object(DiffusionWorker, "__init__", return_value=None) @@ -125,6 +152,35 @@ class TestWorkerWrapperBaseInitialization: skip_load_model=False, ) + def test_custom_pipeline_init_failure_shuts_down_worker(self, mocker: MockerFixture, mock_od_config): + """A failed custom reload must not abandon an initialized worker.""" + instances = [] + + class FailingWorker: + def __init__(self, **kwargs): + self.kwargs = kwargs + self.shutdown_calls = 0 + instances.append(self) + + def re_init_pipeline(self, custom_pipeline_args): + raise RuntimeError(f"failed to load {custom_pipeline_args['pipeline_class']}") + + def shutdown(self): + self.shutdown_calls += 1 + + mocker.patch.object(WorkerWrapperBase, "_prepare_worker_class", return_value=FailingWorker) + + with pytest.raises(RuntimeError, match="failed to load test.Pipeline"): + WorkerWrapperBase( + gpu_id=0, + od_config=mock_od_config, + base_worker_class=DiffusionWorker, + custom_pipeline_args={"pipeline_class": "test.Pipeline"}, + ) + + assert len(instances) == 1 + assert instances[0].shutdown_calls == 1 + # ------------------------------------------------------------------------- # Tests: Worker Extension Functionality @@ -231,28 +287,101 @@ class TestWorkerWrapperBaseDelegation: assert result is None @pytest.mark.parametrize("native_kv", [False, True]) - def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture, native_kv): + def test_worker_shutdown_releases_model_before_cumem_pools(self, mocker: MockerFixture, native_kv): events: list[str] = [] offload_backend = mocker.Mock() offload_backend.disable.side_effect = lambda: events.append("offload") kv_manager = mocker.Mock() kv_manager.close.side_effect = lambda: events.append("kv") - destroy = mocker.patch( - "vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env", - side_effect=lambda: events.append("distributed"), - ) worker = DiffusionWorker.__new__(DiffusionWorker) + worker.od_config = SimpleNamespace(enable_sleep_mode=True) worker.model_runner = SimpleNamespace( offload_backend=offload_backend, kv_transfer_manager=None if native_kv else kv_manager, _kv_transfer_manager=kv_manager, ) + worker.lora_manager = object() + worker.profiler = object() + worker._sleep_saved_buffers = {"buffer": object()} + worker._step_lora_state = {"request": (None, 1.0)} + worker._shutdown_complete = False + worker._owns_sleep_pool = True + + class FakeAllocator: + def __init__(self): + self.pointer_to_data = {1: SimpleNamespace(handle=(0, 1024, 1, []))} + self.allocator_and_pools = {"weights": object()} + self.release_calls = 0 + + def get_current_usage(self): + return sum(data.handle[1] for data in self.pointer_to_data.values()) + + def release_pools(self): + assert worker.model_runner is None + assert worker.lora_manager is None + assert worker.profiler is None + assert worker._sleep_saved_buffers == {} + assert worker._step_lora_state == {} + assert self.get_current_usage() == 0 + self.allocator_and_pools.clear() + self.release_calls += 1 + events.append("pools") + + allocator = FakeAllocator() + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker._get_cumem_allocator_class", + return_value=SimpleNamespace(get_instance=lambda: allocator), + ) + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.shutdown_kv_connector", + side_effect=lambda: events.append("connector"), + ) + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env", + side_effect=lambda: events.append("distributed"), + ) + synchronize = mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform.synchronize", + ) + device_synchronized = False + + def synchronize_device(): + nonlocal device_synchronized + device_synchronized = True + events.append("synchronize") + + synchronize.side_effect = synchronize_device + + def collect_garbage(): + events.append("gc") + if device_synchronized and worker.model_runner is None: + allocator.pointer_to_data.clear() + + collect = mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker.gc.collect", + side_effect=collect_garbage, + ) + worker.shutdown() worker.shutdown() - assert events == ["offload", "kv", "distributed"] + assert events == [ + "offload", + "kv", + "connector", + "distributed", + "synchronize", + "gc", + "pools", + "gc", + ] + assert allocator.release_calls == 1 + assert allocator.get_current_usage() == 0 + assert allocator.allocator_and_pools == {} + assert not worker._owns_sleep_pool kv_manager.close.assert_called_once_with() - destroy.assert_called_once_with() + synchronize.assert_called_once_with() + assert collect.call_count == 2 @pytest.mark.parametrize("cleanup_fails", [False, True]) def test_worker_shutdown_releases_model_references(self, mocker: MockerFixture, cleanup_fails): @@ -272,11 +401,7 @@ class TestWorkerWrapperBaseDelegation: worker._sleep_saved_buffers = {"buffer": weights} del weights, pipeline - if cleanup_fails: - with pytest.raises(RuntimeError, match="offload cleanup failed"): - worker.shutdown() - else: - worker.shutdown() + worker.shutdown() gc.collect() assert worker.model_runner is None @@ -297,11 +422,7 @@ class TestWorkerWrapperBaseDelegation: worker.model_runner = None worker._owns_sleep_pool = owns_pool - if owns_pool and cleanup_fails: - with pytest.raises(RuntimeError, match="pool release failed"): - worker.shutdown() - else: - worker.shutdown() + worker.shutdown() assert release.call_count == int(owns_pool) shutdown_connector.assert_called_once_with() diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 95db7453f..3f4a39ac8 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -267,44 +267,52 @@ class DiffusionWorker: # request id. Used by step mode to recover LoRA identity for cached # requests, which only carry their request_id in subsequent ticks. self._step_lora_state: dict[str, tuple[LoRARequest | None, float]] = {} + self._shutdown_complete = False self.stage_id = getattr(od_config, "stage_id", 0) self.init_device() - # Create model runner — one decision chain, in precedence order: - # 1. explicit od_config.diffusion_model_runner_cls (user override), - # 2. the runner declared by the engine class that engine_backend - # selects (e.g. ARDiffusionEngine -> ARDiffusionModelRunner), - # 3. the platform default. - # Routing policy therefore lives on the engine class / config surface; - # engines never mutate od_config. Overrides must be import-path - # strings — guard with isinstance so a non-string (e.g. a Mock - # od_config in tests) doesn't shadow the platform hook. - runner_override = getattr(self.od_config, "diffusion_model_runner_cls", None) - engine_runner = None - if not (isinstance(runner_override, str) and runner_override): - try: - from vllm_omni.diffusion.diffusion_engine import DiffusionEngine + try: + # Create model runner — one decision chain, in precedence order: + # 1. explicit od_config.diffusion_model_runner_cls (user override), + # 2. the runner declared by the engine class that engine_backend + # selects (e.g. ARDiffusionEngine -> ARDiffusionModelRunner), + # 3. the platform default. + # Routing policy therefore lives on the engine class / config surface; + # engines never mutate od_config. Overrides must be import-path + # strings — guard with isinstance so a non-string (e.g. a Mock + # od_config in tests) doesn't shadow the platform hook. + runner_override = getattr(self.od_config, "diffusion_model_runner_cls", None) + engine_runner = None + if not (isinstance(runner_override, str) and runner_override): + try: + from vllm_omni.diffusion.diffusion_engine import DiffusionEngine - engine_cls = DiffusionEngine.resolve_engine_class(self.od_config) - engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None) - except Exception: - logger.warning("Worker %s: engine_backend resolution failed; using platform runner", self.rank) - engine_runner = None - if isinstance(runner_override, str) and runner_override: - model_runner_cls_path = runner_override - elif isinstance(engine_runner, str) and engine_runner: - model_runner_cls_path = engine_runner - else: - model_runner_cls_path = current_omni_platform.get_diffusion_model_runner_cls() - model_runner_cls = resolve_obj_by_qualname(model_runner_cls_path) - self.model_runner = model_runner_cls( - vllm_config=self.vllm_config, - od_config=self.od_config, - device=self.device, - ) - self.profiler: WorkerProfiler | None = self._create_profiler() - if not skip_load_model: - self.load_model(load_format=self.od_config.diffusion_load_format) - self.init_lora_manager() + engine_cls = DiffusionEngine.resolve_engine_class(self.od_config) + engine_runner = getattr(engine_cls, "default_diffusion_model_runner_cls", None) + except Exception: + logger.warning("Worker %s: engine_backend resolution failed; using platform runner", self.rank) + engine_runner = None + if isinstance(runner_override, str) and runner_override: + model_runner_cls_path = runner_override + elif isinstance(engine_runner, str) and engine_runner: + model_runner_cls_path = engine_runner + else: + model_runner_cls_path = current_omni_platform.get_diffusion_model_runner_cls() + model_runner_cls = resolve_obj_by_qualname(model_runner_cls_path) + self.model_runner = model_runner_cls( + vllm_config=self.vllm_config, + od_config=self.od_config, + device=self.device, + ) + self.profiler: WorkerProfiler | None = self._create_profiler() + if not skip_load_model: + self.load_model(load_format=self.od_config.diffusion_load_format) + self.init_lora_manager() + except Exception: + # init_device() owns process-global distributed state. If anything + # after it fails, unwind that state before another inline worker is + # allowed to initialize in this process. + self.shutdown() + raise logger.info(f"Worker {self.rank}: Initialization complete.") def init_device(self) -> None: @@ -1017,40 +1025,89 @@ class DiffusionWorker: return nullcontext() def shutdown(self) -> None: - """Shutdown the worker and cleanup distributed environment.""" - try: - if self.model_runner is not None: - mgr = getattr(self.model_runner, "kv_transfer_manager", None) - if mgr is None: - mgr = getattr(self.model_runner, "_kv_transfer_manager", None) - try: - offload_backend = getattr(self.model_runner, "offload_backend", None) - if offload_backend is not None: - offload_backend.disable() - finally: - if mgr is not None: - mgr.close() - finally: - self.model_runner = None - self.lora_manager = None - self._sleep_saved_buffers = {} + """Shutdown the worker and release process-global resources.""" + if getattr(self, "_shutdown_complete", False): + return + self._shutdown_complete = True + + # Detach every model-owned reference before releasing the CuMem pools. + # Inline executors stay in the parent process, so relying on their + # wrapper becoming unreachable is not sufficient for deterministic + # teardown between sequential engine instances. + model_runner = getattr(self, "model_runner", None) + self.model_runner = None + self.lora_manager = None + self.profiler = None + self._sleep_saved_buffers = {} + self._step_lora_state = {} + + if model_runner is not None: + mgr = getattr(model_runner, "kv_transfer_manager", None) + if mgr is None: + mgr = getattr(model_runner, "_kv_transfer_manager", None) try: - if getattr(self, "_owns_sleep_pool", False): + offload_backend = getattr(model_runner, "offload_backend", None) + if offload_backend is not None: + offload_backend.disable() + except Exception: + logger.exception("Failed to disable diffusion offload backend during shutdown") + try: + if mgr is not None: + mgr.close() + except Exception: + logger.exception("Failed to close diffusion KV transfer manager during shutdown") + + try: + shutdown_kv_connector() + except Exception: + logger.exception("Failed to shutdown diffusion KV connector") + + try: + a2a_permute = sys.modules.get("vllm_omni.diffusion.distributed.a2a_permute") + if a2a_permute is not None: + a2a_permute.clear_a2a_permute_workspaces() + except Exception: + logger.exception("Failed to release fused Ulysses symmetric-memory workspaces") + + try: + destroy_distributed_env() + except Exception: + logger.exception("Failed to destroy diffusion distributed environment") + + # Drop the local runner only after its background services are stopped. + # Tensor finalizers must run before release_pools(), otherwise the + # singleton keeps allocations visible to the next sleep-enabled worker. + del model_runner + owns_sleep_pool = getattr(self, "_owns_sleep_pool", False) + if owns_sleep_pool: + try: + current_omni_platform.synchronize() + except Exception: + logger.exception("Failed to synchronize diffusion device during shutdown") + try: + gc.collect() + except Exception: + logger.exception("Failed to collect diffusion model resources during shutdown") + if owns_sleep_pool: + try: + CuMemAllocator = _get_cumem_allocator_class() + allocator = CuMemAllocator.get_instance() + if allocator.get_current_usage(): gc.collect() - _get_cumem_allocator_class().get_instance().release_pools() - self._owns_sleep_pool = False - finally: - try: - shutdown_kv_connector() - finally: - try: - a2a_permute = sys.modules.get("vllm_omni.diffusion.distributed.a2a_permute") - if a2a_permute is not None: - a2a_permute.clear_a2a_permute_workspaces() - except Exception: - logger.exception("Failed to release fused Ulysses symmetric-memory workspaces") - finally: - destroy_distributed_env() + allocator.release_pools() + self._owns_sleep_pool = False + # Re-collect after dropping the pool references so their + # finalizers remove every tracked allocation before the next + # inline worker checks get_current_usage(). + gc.collect() + remaining_usage = allocator.get_current_usage() + if remaining_usage: + logger.warning( + "Diffusion CuMem allocator still tracks %s bytes after shutdown", + remaining_usage, + ) + except Exception: + logger.exception("Failed to release diffusion CuMem pools during shutdown") class CustomPipelineWorkerExtension: @@ -1643,26 +1700,29 @@ class WorkerWrapperBase: # Prepare worker class with extension support worker_class = self._prepare_worker_class() - # Create the actual worker instance - # Only dynamic custom pipelines skip initial loading; native pipelines - # may also use custom_pipeline_args for model-specific component paths. - self.worker = worker_class( + # Create the actual worker instance. Only dynamic custom pipelines skip + # initial loading; native pipelines may also use custom_pipeline_args + # for model-specific component paths. + worker = worker_class( local_rank=gpu_id, rank=gpu_id, od_config=od_config, skip_load_model=self.uses_custom_pipeline, ) - - # Re-initialize pipeline with custom pipeline if provided - if self.uses_custom_pipeline: + try: + # Re-initialize pipeline with custom pipeline if provided. + if self.uses_custom_pipeline: + worker.re_init_pipeline(self.custom_pipeline_args) + except Exception: + # The executor never receives a wrapper whose constructor raises. + # Unwind a successfully created worker here so its distributed + # groups do not poison the next inline engine initialization. try: - self.worker.re_init_pipeline(self.custom_pipeline_args) + worker.shutdown() except Exception: - try: - self.worker.shutdown() - except Exception: - logger.exception("Failed to clean up worker after custom pipeline initialization failure") - raise + logger.exception("Failed to clean up worker after custom pipeline initialization failure") + raise + self.worker = worker def _prepare_worker_class(self) -> type: """