[Bugfix][Diffusion] Clean up inline worker lifecycle (#7986)

Signed-off-by: andyluo7 <andy.luo@amd.com>

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()