anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[ROCm] [CI] Fix diffusion test (#4574)

Browse Source 
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
This commit is contained in:
TJian authored and GitHub committed 2026-06-20 08:31:39 +08:00
1 parent 6ff246bf3c
commit e70822ed54
2 files changed
+26 -14
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/test_diffusion_worker.py
+17 -8
View File
@@ -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()
vllm_omni/diffusion/worker/diffusion_worker.py
+9 -6
View File
@@ -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."