diff --git a/tests/diffusion/test_diffusion_worker.py b/tests/diffusion/test_diffusion_worker.py index 8ad6c2734..fb45cd41d 100644 --- a/tests/diffusion/test_diffusion_worker.py +++ b/tests/diffusion/test_diffusion_worker.py @@ -23,6 +23,15 @@ from vllm_omni.diffusion.worker.diffusion_worker import ( pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.gpu] +def patch_cumem_allocator(mocker: MockerFixture): + mock_allocator_class = mocker.Mock() + mocker.patch( + "vllm_omni.diffusion.worker.diffusion_worker._get_cumem_allocator_class", + return_value=mock_allocator_class, + ) + return mock_allocator_class + + @pytest.fixture def mock_od_config(mocker: MockerFixture): """Create a mock OmniDiffusionConfig.""" @@ -118,7 +127,7 @@ class TestDiffusionWorkerSleep: """ Unified interception of Allocators, and provision of default security values. """ - self.mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + self.mock_allocator_class = patch_cumem_allocator(mocker) self.mock_allocator = mocker.Mock() self.mock_allocator_class.get_instance.return_value = self.mock_allocator self.mock_allocator.get_current_usage.return_value = 4 * 1024**3 @@ -126,7 +135,7 @@ class TestDiffusionWorkerSleep: def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep mode level 1 (offload weights only).""" - mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_allocator_class = patch_cumem_allocator(mocker) mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") mock_platform.get_free_memory.side_effect = [10 * 1024**3, 12 * 1024**3] mock_platform.get_device_total_memory.return_value = 80 * 1024**3 @@ -158,7 +167,7 @@ class TestDiffusionWorkerSleep: def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep mode level 2 (offload all, save buffers).""" - mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_allocator_class = patch_cumem_allocator(mocker) mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") mock_platform.get_free_memory.side_effect = [5 * 1024**3, 10 * 1024**3] mock_platform.get_device_total_memory.return_value = 80 * 1024**3 @@ -201,7 +210,7 @@ class TestDiffusionWorkerSleep: def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker): """Test that sleep validates memory was actually freed.""" - mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_allocator_class = patch_cumem_allocator(mocker) mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") mock_platform.get_free_memory.return_value = 10 * 1024**3 mock_platform.get_device_total_memory.return_value = 80 * 1024**3 @@ -226,7 +235,7 @@ class TestDiffusionWorkerSleep: def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(self, mocker: MockerFixture, mock_gpu_worker): """Test sleep uses device-scoped fallback when NVML is unavailable.""" - mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_allocator_class = patch_cumem_allocator(mocker) mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform") mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory") mock_get_process_memory.side_effect = [None, None] @@ -252,7 +261,7 @@ class TestDiffusionWorkerWakeUp: def test_wake_up_without_buffers(self, mocker: MockerFixture, mock_gpu_worker): """Test wake_up without saved buffers (level 1 sleep).""" - mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_allocator_class = patch_cumem_allocator(mocker) # Setup allocator mock mock_allocator = mocker.Mock() @@ -272,7 +281,7 @@ class TestDiffusionWorkerWakeUp: def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker): """Test wake_up with saved buffers (level 2 sleep).""" - mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_allocator_class = patch_cumem_allocator(mocker) # Setup allocator mock mock_allocator = mocker.Mock() @@ -317,7 +326,7 @@ class TestDiffusionWorkerWakeUp: def test_wake_up_partial_buffer_restore(self, mocker: MockerFixture, mock_gpu_worker): """Test wake_up only restores buffers that were saved.""" - mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator") + mock_allocator_class = patch_cumem_allocator(mocker) # Setup allocator mock mock_allocator = mocker.Mock() diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 0b4e87718..62cddd4a6 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -161,6 +161,12 @@ def _create_diffusion_worker_vllm_config(device: torch.device, od_config: OmniDi return vllm_config +def _get_cumem_allocator_class() -> type: + from vllm.device_allocator.cumem import CuMemAllocator + + return CuMemAllocator + + def _resolve_ir_op_priority(od_config: OmniDiffusionConfig, vllm_config: VllmConfig) -> Any: ir_op_priority = current_omni_platform.get_default_ir_op_priority(vllm_config) ir_op_priority_func = get_diffusion_ir_op_priority_func(od_config) @@ -455,8 +461,7 @@ class DiffusionWorker: Args: level: Sleep level. Level 1 offloads weights, level 2 also saves buffers. """ - from vllm.device_allocator.cumem import CuMemAllocator - + CuMemAllocator = _get_cumem_allocator_class() allocator = CuMemAllocator.get_instance() usage_before = allocator.get_current_usage() @@ -508,8 +513,7 @@ class DiffusionWorker: tags: List of memory pool tags to re-activate (e.g., ["weights"] to match Level 1 sleep). If None, all pools are re-activated. """ - from vllm.device_allocator.cumem import CuMemAllocator - + CuMemAllocator = _get_cumem_allocator_class() allocator = CuMemAllocator.get_instance() allocator.wake_up(tags) current_omni_platform.synchronize() @@ -631,8 +635,7 @@ class DiffusionWorker: if is_sleep_enabled: current_omni_platform.synchronize() gc.collect() - from vllm.device_allocator.cumem import CuMemAllocator - + CuMemAllocator = _get_cumem_allocator_class() allocator = CuMemAllocator.get_instance() if tag == "weights": assert allocator.get_current_usage() == 0, "Sleep mode can only be used for one instance per process."