[Misc] Migrate L1 tests to use pytest-mock (#1315)
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com> Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
This commit is contained in:
Yuanheng Zhao authored and GitHub committed
1 parent
9913c00144
commit
ed08a216c1
20 files changed
+977
-954
No files matched your search
@@ -10,10 +10,9 @@ This module tests the DiffusionWorker implementation:
|
||||
- wake_up: Waking worker from sleep mode
|
||||
"""
|
||||
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from pytest_mock import MockerFixture
|
||||
|
||||
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
|
||||
|
||||
@@ -21,9 +20,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_od_config():
|
||||
def mock_od_config(mocker: MockerFixture):
|
||||
"""Create a mock OmniDiffusionConfig."""
|
||||
config = Mock()
|
||||
config = mocker.Mock()
|
||||
config.num_gpus = 1
|
||||
config.master_port = 12345
|
||||
config.enable_sleep_mode = False
|
||||
@@ -34,22 +33,23 @@ def mock_od_config():
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_gpu_worker(mock_od_config):
|
||||
def mock_gpu_worker(mocker: MockerFixture, mock_od_config):
|
||||
"""Create a DiffusionWorker with mocked initialization."""
|
||||
with patch.object(DiffusionWorker, "init_device"), patch.object(DiffusionWorker, "load_model"):
|
||||
worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
|
||||
# Mock the model_runner with pipeline
|
||||
worker.model_runner = Mock()
|
||||
worker.model_runner.pipeline = Mock()
|
||||
worker.device = torch.device("cuda", 0)
|
||||
worker._sleep_saved_buffers = {}
|
||||
return worker
|
||||
mocker.patch.object(DiffusionWorker, "init_device")
|
||||
mocker.patch.object(DiffusionWorker, "load_model")
|
||||
worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
|
||||
# Mock the model_runner with pipeline
|
||||
worker.model_runner = mocker.Mock()
|
||||
worker.model_runner.pipeline = mocker.Mock()
|
||||
worker.device = torch.device("cuda", 0)
|
||||
worker._sleep_saved_buffers = {}
|
||||
return worker
|
||||
|
||||
|
||||
class TestDiffusionWorkerLoadWeights:
|
||||
"""Test DiffusionWorker.load_weights method."""
|
||||
|
||||
def test_load_weights_calls_pipeline(self, mock_gpu_worker):
|
||||
def test_load_weights_calls_pipeline(self, mocker: MockerFixture, mock_gpu_worker):
|
||||
"""Test that load_weights delegates to model_runner.load_weights."""
|
||||
# Setup mock weights
|
||||
mock_weights = [
|
||||
@@ -59,7 +59,7 @@ class TestDiffusionWorkerLoadWeights:
|
||||
expected_loaded = {"layer1.weight", "layer2.weight"}
|
||||
|
||||
# Configure model_runner mock
|
||||
mock_gpu_worker.model_runner.load_weights = Mock(return_value=expected_loaded)
|
||||
mock_gpu_worker.model_runner.load_weights = mocker.Mock(return_value=expected_loaded)
|
||||
|
||||
# Call load_weights
|
||||
result = mock_gpu_worker.load_weights(mock_weights)
|
||||
@@ -68,9 +68,9 @@ class TestDiffusionWorkerLoadWeights:
|
||||
mock_gpu_worker.model_runner.load_weights.assert_called_once_with(mock_weights)
|
||||
assert result == expected_loaded
|
||||
|
||||
def test_load_weights_empty_iterable(self, mock_gpu_worker):
|
||||
def test_load_weights_empty_iterable(self, mocker: MockerFixture, mock_gpu_worker):
|
||||
"""Test load_weights with empty weights iterable."""
|
||||
mock_gpu_worker.model_runner.load_weights = Mock(return_value=set())
|
||||
mock_gpu_worker.model_runner.load_weights = mocker.Mock(return_value=set())
|
||||
|
||||
result = mock_gpu_worker.load_weights([])
|
||||
|
||||
@@ -81,11 +81,12 @@ class TestDiffusionWorkerLoadWeights:
|
||||
class TestDiffusionWorkerSleep:
|
||||
"""Test DiffusionWorker.sleep method."""
|
||||
|
||||
@patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
|
||||
@patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
|
||||
@patch("vllm.device_allocator.cumem.CuMemAllocator")
|
||||
def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_get_process_memory, mock_gpu_worker):
|
||||
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")
|
||||
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")
|
||||
|
||||
# Setup process-scoped memory mocks
|
||||
# Before sleep: 3GB used
|
||||
# After sleep: 1GB used (freed 2GB)
|
||||
@@ -95,9 +96,9 @@ class TestDiffusionWorkerSleep:
|
||||
]
|
||||
|
||||
# Setup allocator mock
|
||||
mock_allocator = Mock()
|
||||
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
|
||||
mock_allocator.sleep = Mock()
|
||||
mock_allocator = mocker.Mock()
|
||||
mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
|
||||
mock_allocator.sleep = mocker.Mock()
|
||||
|
||||
# Call sleep with level 1
|
||||
result = mock_gpu_worker.sleep(level=1)
|
||||
@@ -108,11 +109,12 @@ class TestDiffusionWorkerSleep:
|
||||
# Verify buffers were NOT saved (level 1 doesn't save buffers)
|
||||
assert len(mock_gpu_worker._sleep_saved_buffers) == 0
|
||||
|
||||
@patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
|
||||
@patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
|
||||
@patch("vllm.device_allocator.cumem.CuMemAllocator")
|
||||
def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_get_process_memory, mock_gpu_worker):
|
||||
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")
|
||||
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")
|
||||
|
||||
# Setup process-scoped memory mocks
|
||||
mock_get_process_memory.side_effect = [
|
||||
5 * 1024**3, # Before sleep
|
||||
@@ -120,14 +122,14 @@ class TestDiffusionWorkerSleep:
|
||||
]
|
||||
|
||||
# Setup allocator mock
|
||||
mock_allocator = Mock()
|
||||
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
|
||||
mock_allocator.sleep = Mock()
|
||||
mock_allocator = mocker.Mock()
|
||||
mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
|
||||
mock_allocator.sleep = mocker.Mock()
|
||||
|
||||
# Mock pipeline buffers
|
||||
mock_buffer1 = torch.randn(10, 10)
|
||||
mock_buffer2 = torch.randn(20, 20)
|
||||
mock_gpu_worker.model_runner.pipeline.named_buffers = Mock(
|
||||
mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(
|
||||
return_value=[
|
||||
("buffer1", mock_buffer1),
|
||||
("buffer2", mock_buffer2),
|
||||
@@ -146,42 +148,32 @@ class TestDiffusionWorkerSleep:
|
||||
assert "buffer1" in mock_gpu_worker._sleep_saved_buffers
|
||||
assert "buffer2" in mock_gpu_worker._sleep_saved_buffers
|
||||
|
||||
@patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
|
||||
@patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
|
||||
@patch("vllm.device_allocator.cumem.CuMemAllocator")
|
||||
def test_sleep_memory_freed_validation(
|
||||
self,
|
||||
mock_allocator_class,
|
||||
mock_platform,
|
||||
mock_get_process_memory,
|
||||
mock_gpu_worker,
|
||||
):
|
||||
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")
|
||||
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")
|
||||
|
||||
# Simulate process memory increase (should trigger assertion error)
|
||||
mock_get_process_memory.side_effect = [
|
||||
1 * 1024**3, # Before sleep: 1GB used
|
||||
3 * 1024**3, # After sleep: 3GB used (negative freed)
|
||||
]
|
||||
|
||||
mock_allocator = Mock()
|
||||
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
|
||||
mock_allocator.sleep = Mock()
|
||||
mock_allocator = mocker.Mock()
|
||||
mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
|
||||
mock_allocator.sleep = mocker.Mock()
|
||||
|
||||
# This should raise an assertion error
|
||||
with pytest.raises(AssertionError, match="Memory usage increased after sleeping"):
|
||||
mock_gpu_worker.sleep(level=1)
|
||||
|
||||
@patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
|
||||
@patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
|
||||
@patch("vllm.device_allocator.cumem.CuMemAllocator")
|
||||
def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(
|
||||
self,
|
||||
mock_allocator_class,
|
||||
mock_platform,
|
||||
mock_get_process_memory,
|
||||
mock_gpu_worker,
|
||||
):
|
||||
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_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]
|
||||
mock_platform.get_free_memory.side_effect = [
|
||||
1 * 1024**3, # Before sleep
|
||||
@@ -189,9 +181,9 @@ class TestDiffusionWorkerSleep:
|
||||
]
|
||||
mock_platform.get_device_total_memory.return_value = 8 * 1024**3
|
||||
|
||||
mock_allocator = Mock()
|
||||
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
|
||||
mock_allocator.sleep = Mock()
|
||||
mock_allocator = mocker.Mock()
|
||||
mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
|
||||
mock_allocator.sleep = mocker.Mock()
|
||||
|
||||
result = mock_gpu_worker.sleep(level=1)
|
||||
|
||||
@@ -202,13 +194,14 @@ class TestDiffusionWorkerSleep:
|
||||
class TestDiffusionWorkerWakeUp:
|
||||
"""Test DiffusionWorker.wake_up method."""
|
||||
|
||||
@patch("vllm.device_allocator.cumem.CuMemAllocator")
|
||||
def test_wake_up_without_buffers(self, mock_allocator_class, mock_gpu_worker):
|
||||
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")
|
||||
|
||||
# Setup allocator mock
|
||||
mock_allocator = Mock()
|
||||
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
|
||||
mock_allocator.wake_up = Mock()
|
||||
mock_allocator = mocker.Mock()
|
||||
mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
|
||||
mock_allocator.wake_up = mocker.Mock()
|
||||
|
||||
# Ensure no saved buffers
|
||||
mock_gpu_worker._sleep_saved_buffers = {}
|
||||
@@ -220,13 +213,14 @@ class TestDiffusionWorkerWakeUp:
|
||||
mock_allocator.wake_up.assert_called_once_with(["weights"])
|
||||
assert result is True
|
||||
|
||||
@patch("vllm.device_allocator.cumem.CuMemAllocator")
|
||||
def test_wake_up_with_buffers(self, mock_allocator_class, mock_gpu_worker):
|
||||
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")
|
||||
|
||||
# Setup allocator mock
|
||||
mock_allocator = Mock()
|
||||
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
|
||||
mock_allocator.wake_up = Mock()
|
||||
mock_allocator = mocker.Mock()
|
||||
mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
|
||||
mock_allocator.wake_up = mocker.Mock()
|
||||
|
||||
# Create saved buffers
|
||||
saved_buffer1 = torch.randn(10, 10)
|
||||
@@ -237,12 +231,12 @@ class TestDiffusionWorkerWakeUp:
|
||||
}
|
||||
|
||||
# Mock pipeline buffers (these will be restored)
|
||||
mock_buffer1 = Mock()
|
||||
mock_buffer1.data = Mock()
|
||||
mock_buffer2 = Mock()
|
||||
mock_buffer2.data = Mock()
|
||||
mock_buffer1 = mocker.Mock()
|
||||
mock_buffer1.data = mocker.Mock()
|
||||
mock_buffer2 = mocker.Mock()
|
||||
mock_buffer2.data = mocker.Mock()
|
||||
|
||||
mock_gpu_worker.model_runner.pipeline.named_buffers = Mock(
|
||||
mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(
|
||||
return_value=[
|
||||
("buffer1", mock_buffer1),
|
||||
("buffer2", mock_buffer2),
|
||||
@@ -263,13 +257,14 @@ class TestDiffusionWorkerWakeUp:
|
||||
assert len(mock_gpu_worker._sleep_saved_buffers) == 0
|
||||
assert result is True
|
||||
|
||||
@patch("vllm.device_allocator.cumem.CuMemAllocator")
|
||||
def test_wake_up_partial_buffer_restore(self, mock_allocator_class, mock_gpu_worker):
|
||||
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")
|
||||
|
||||
# Setup allocator mock
|
||||
mock_allocator = Mock()
|
||||
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
|
||||
mock_allocator.wake_up = Mock()
|
||||
mock_allocator = mocker.Mock()
|
||||
mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
|
||||
mock_allocator.wake_up = mocker.Mock()
|
||||
|
||||
# Only save buffer1, not buffer2
|
||||
saved_buffer1 = torch.randn(10, 10)
|
||||
@@ -278,12 +273,12 @@ class TestDiffusionWorkerWakeUp:
|
||||
}
|
||||
|
||||
# Mock pipeline has both buffers
|
||||
mock_buffer1 = Mock()
|
||||
mock_buffer1.data = Mock()
|
||||
mock_buffer2 = Mock()
|
||||
mock_buffer2.data = Mock()
|
||||
mock_buffer1 = mocker.Mock()
|
||||
mock_buffer1.data = mocker.Mock()
|
||||
mock_buffer2 = mocker.Mock()
|
||||
mock_buffer2.data = mocker.Mock()
|
||||
|
||||
mock_gpu_worker.model_runner.pipeline.named_buffers = Mock(
|
||||
mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(
|
||||
return_value=[
|
||||
("buffer1", mock_buffer1),
|
||||
("buffer2", mock_buffer2),
|
||||
|
||||
Reference in new issue
Block a user