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

[Misc] Cleanup: use consistent pytest-mock in unit tests (#2698)

Browse Source 
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
This commit is contained in:
Yuanheng Zhao authored and GitHub committed 2026-04-13 19:21:42 +08:00
1 parent 2c67c30550
commit e0cdbe9a5d
26 files changed
+1598 -1369
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/diffusion/distributed/test_distributed_vae_executor.py
+16 -25
View File
@@ -1,4 +1,4 @@
from unittest.mock import MagicMock, patch
from types import SimpleNamespace
import pytest
import torch
@@ -61,40 +61,31 @@ class E2EOperator:
class DummyMixin(DistributedVaeMixin):
def __init__(self):
self.use_tiling = True
self.distributed_executor = MagicMock()
self.distributed_executor.parallel_size = 2
self.distributed_executor.group = None
self.distributed_executor = SimpleNamespace(parallel_size=2, group=None)
@pytest.fixture(autouse=True)
def mock_dist():
with (
patch.object(dist, "get_world_size", return_value=2),
patch.object(dist, "get_rank", return_value=0),
patch.object(dist, "is_initialized", return_value=True),
patch.object(dist, "all_reduce", return_value=None),
patch.object(dist, "gather", return_value=None),
patch.object(dist, "broadcast", return_value=None),
):
yield
def mock_dist(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(dist, "get_world_size", lambda *args, **kwargs: 2)
monkeypatch.setattr(dist, "get_rank", lambda *args, **kwargs: 0)
monkeypatch.setattr(dist, "is_initialized", lambda: True)
monkeypatch.setattr(dist, "all_reduce", lambda *args, **kwargs: None)
monkeypatch.setattr(dist, "gather", lambda *args, **kwargs: None)
monkeypatch.setattr(dist, "broadcast", lambda *args, **kwargs: None)
@pytest.fixture(autouse=True)
def mock_dit_group():
with patch(
def mock_dit_group(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(
"vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor.get_dit_group",
new=MagicMock(return_value=None),
):
yield
lambda: None,
)
@pytest.fixture(autouse=True)
def mock_dist_vae_executor():
with (
patch.object(DistributedVaeExecutor, "gather_tensors", side_effect=lambda x: [x]),
patch.object(DistributedVaeExecutor, "broadcast_tensor", side_effect=lambda x: x),
):
yield
def mock_dist_vae_executor(monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(DistributedVaeExecutor, "gather_tensors", lambda self, x: [x])
monkeypatch.setattr(DistributedVaeExecutor, "broadcast_tensor", lambda self, x: x)
# ============================
tests/diffusion/models/bagel/test_trajectory_recording.py
+21 -13
View File
@@ -4,10 +4,10 @@
import types
from dataclasses import dataclass
from unittest.mock import MagicMock, patch
import pytest
import torch
from pytest_mock import MockerFixture
from vllm_omni.diffusion.models.bagel.bagel_transformer import (
Bagel,
@@ -23,9 +23,9 @@ NUM_TIMESTEPS = 5
EXPECTED_STEPS = NUM_TIMESTEPS - 1
def _make_mock_bagel():
def _make_mock_bagel(mocker: MockerFixture):
"""Create a mock Bagel with forward returning constant velocity."""
mock = MagicMock(spec=Bagel)
mock = mocker.MagicMock(spec=Bagel)
mock._sp_size = 1
# forward returns a small constant velocity so x_t changes each step
@@ -78,18 +78,22 @@ def _make_generate_args(num_tokens=NUM_TOKENS, hidden_dim=HIDDEN_DIM, cfg=False)
@pytest.fixture(params=[False, True], ids=["no_cfg", "batched_cfg"])
def bagel_and_args(request):
def bagel_and_args(
request,
monkeypatch: pytest.MonkeyPatch,
mocker: MockerFixture,
):
"""Mock Bagel instance and generate_image arguments.
Parametrized over CFG mode so every test runs on both the no-CFG
and batched-CFG code paths.
"""
cfg = request.param
with patch(
monkeypatch.setattr(
"vllm_omni.diffusion.models.bagel.bagel_transformer.get_classifier_free_guidance_world_size",
return_value=1,
):
yield _make_mock_bagel(), _make_generate_args(cfg=cfg)
lambda: 1,
)
yield _make_mock_bagel(mocker), _make_generate_args(cfg=cfg)
class TestTrajectoryRecording:
@@ -188,12 +192,16 @@ class TestTrajectoryLogProbs:
"""Tests for log-prob recording when a scheduler is provided."""
@pytest.fixture()
def bagel_scheduler_args(self):
with patch(
def bagel_scheduler_args(
self,
monkeypatch: pytest.MonkeyPatch,
mocker: MockerFixture,
):
monkeypatch.setattr(
"vllm_omni.diffusion.models.bagel.bagel_transformer.get_classifier_free_guidance_world_size",
return_value=1,
):
yield _make_mock_bagel(), _make_generate_args(), _MockScheduler()
lambda: 1,
)
yield _make_mock_bagel(mocker), _make_generate_args(), _MockScheduler()
def test_log_probs_recorded_with_scheduler(self, bagel_scheduler_args):
bagel, args, scheduler = bagel_scheduler_args
tests/diffusion/models/flux2/test_flux2_transformer_tp.py
+11 -9
View File
@@ -1,7 +1,6 @@
from unittest.mock import MagicMock, patch
import pytest
import torch
from pytest_mock import MockerFixture
from tests.utils import hardware_test
from vllm_omni.diffusion.models.flux2.flux2_transformer import (
@@ -12,14 +11,17 @@ from vllm_omni.diffusion.models.flux2.flux2_transformer import (
# Initialize TP group before tests
@pytest.fixture(scope="function", autouse=True)
def setup_tp_group():
def setup_tp_group(mocker: MockerFixture):
"""Set up TP group for each test function"""
with patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=2):
with patch("vllm.distributed.parallel_state.get_tp_group") as mock_get_tp_group:
mock_tp_group = MagicMock()
mock_tp_group.world_size = 2
mock_get_tp_group.return_value = mock_tp_group
yield
mocker.patch(
"vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size",
return_value=2,
)
mock_get_tp_group = mocker.patch("vllm.distributed.parallel_state.get_tp_group")
mock_tp_group = mocker.MagicMock()
mock_tp_group.world_size = 2
mock_get_tp_group.return_value = mock_tp_group
yield
class TestFlux2TransformerWeightLoading:
tests/diffusion/offloader/test_sequential_backend.py
+56 -58
View File
@@ -3,8 +3,6 @@
"""Unit tests for SequentialOffloadBackend."""
from unittest.mock import patch
import pytest
import torch
from torch import nn
@@ -44,7 +42,7 @@ def _track_pin_memory_calls():
class TestMoveParamsPinMemory:
def test_dtensor_skips_pin_memory(self, accelerator_device):
def test_dtensor_skips_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
"""DTensor should skip pin_memory to avoid RuntimeError."""
module = _create_simple_module().to(accelerator_device)
tracker, mock_pin = _track_pin_memory_calls()
@@ -56,73 +54,73 @@ class TestMoveParamsPinMemory:
return True
return original_isinstance(obj, cls)
with patch.object(torch.Tensor, "pin_memory", mock_pin):
with patch("builtins.isinstance", fake_isinstance):
hook = SequentialOffloadHook(
offload_targets=[],
device=accelerator_device,
pin_memory=True,
use_hsdp=False,
)
hook._move_params(
module,
torch.device("cpu"),
non_blocking=False,
pin_memory=True,
)
assert not tracker["called"], "pin_memory should not be called for DTensor"
monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin)
monkeypatch.setattr("builtins.isinstance", fake_isinstance)
hook = SequentialOffloadHook(
offload_targets=[],
device=accelerator_device,
pin_memory=True,
use_hsdp=False,
)
hook._move_params(
module,
torch.device("cpu"),
non_blocking=False,
pin_memory=True,
)
assert not tracker["called"], "pin_memory should not be called for DTensor"
def test_regular_tensor_calls_pin_memory(self, accelerator_device):
def test_regular_tensor_calls_pin_memory(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
"""Regular tensor should call pin_memory when moving to CPU."""
module = _create_simple_module().to(accelerator_device)
tracker, mock_pin = _track_pin_memory_calls()
with patch.object(torch.Tensor, "pin_memory", mock_pin):
hook = SequentialOffloadHook(
offload_targets=[],
device=accelerator_device,
pin_memory=True,
use_hsdp=False,
)
hook._move_params(
module,
torch.device("cpu"),
non_blocking=False,
pin_memory=True,
)
assert tracker["called"], "pin_memory should be called for regular tensors"
monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin)
hook = SequentialOffloadHook(
offload_targets=[],
device=accelerator_device,
pin_memory=True,
use_hsdp=False,
)
hook._move_params(
module,
torch.device("cpu"),
non_blocking=False,
pin_memory=True,
)
assert tracker["called"], "pin_memory should be called for regular tensors"
def test_pin_memory_skipped_when_disabled(self, accelerator_device):
def test_pin_memory_skipped_when_disabled(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
"""pin_memory should not be called when pin_memory=False."""
module = _create_simple_module().to(accelerator_device)
tracker, mock_pin = _track_pin_memory_calls()
with patch.object(torch.Tensor, "pin_memory", mock_pin):
hook = SequentialOffloadHook(
offload_targets=[],
device=accelerator_device,
pin_memory=False,
use_hsdp=False,
)
hook._move_params(
module,
torch.device("cpu"),
non_blocking=False,
pin_memory=False,
)
assert not tracker["called"], "pin_memory should not be called when disabled"
monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin)
hook = SequentialOffloadHook(
offload_targets=[],
device=accelerator_device,
pin_memory=False,
use_hsdp=False,
)
hook._move_params(
module,
torch.device("cpu"),
non_blocking=False,
pin_memory=False,
)
assert not tracker["called"], "pin_memory should not be called when disabled"
def test_pin_memory_skipped_for_non_cpu_target(self, accelerator_device):
def test_pin_memory_skipped_for_non_cpu_target(self, accelerator_device, monkeypatch: pytest.MonkeyPatch):
"""pin_memory should not be called for non-CPU targets."""
module = _create_simple_module().to("cpu")
tracker, mock_pin = _track_pin_memory_calls()
with patch.object(torch.Tensor, "pin_memory", mock_pin):
hook = SequentialOffloadHook(
offload_targets=[],
device=torch.device("cpu"),
pin_memory=True,
use_hsdp=False,
)
hook._move_params(module, accelerator_device, non_blocking=False, pin_memory=True)
assert not tracker["called"], "pin_memory should not be called for non-CPU target"
monkeypatch.setattr(torch.Tensor, "pin_memory", mock_pin)
hook = SequentialOffloadHook(
offload_targets=[],
device=torch.device("cpu"),
pin_memory=True,
use_hsdp=False,
)
hook._move_params(module, accelerator_device, non_blocking=False, pin_memory=True)
assert not tracker["called"], "pin_memory should not be called for non-CPU target"
tests/diffusion/quantization/test_int8_config.py
+13 -19
View File
@@ -2,8 +2,6 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for Int8 quantization config."""
from unittest.mock import MagicMock, patch
import pytest
import torch
from pytest_mock import MockerFixture
@@ -102,7 +100,7 @@ def test_quantization_config_string_and_dict_equivalent():
assert config_str.quantization_config.activation_scheme == config_dict.quantization_config.activation_scheme
def test_get_quant_method(mocker: MockerFixture):
def test_get_quant_method(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
"""Test for get_quant_method method for GPU"""
from vllm_omni.quantization.int8_config import Int8OnlineLinearMethod
@@ -111,18 +109,16 @@ def test_get_quant_method(mocker: MockerFixture):
def _fake_init(self, quant_config):
pass
layer = MagicMock(spec=LinearBase)
layer = mocker.Mock(spec=LinearBase)
mocker.patch.object(Int8OnlineLinearMethod, "__init__", _fake_init)
prefix = "test_layer"
# Mock the platform to be GPU
with (
patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=True),
patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=False),
):
method = config.get_quant_method(layer, prefix)
assert isinstance(method, Int8OnlineLinearMethod)
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: False)
method = config.get_quant_method(layer, prefix)
assert isinstance(method, Int8OnlineLinearMethod)
# Test skipping quantization for a layer
config.ignored_layers = [prefix]
@@ -130,22 +126,20 @@ def test_get_quant_method(mocker: MockerFixture):
assert isinstance(method, UnquantizedLinearMethod)
def test_get_npu_quant_method():
def test_get_npu_quant_method(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
"""Test for get_quant_method method for NPU"""
from vllm_omni.quantization.int8_config import NPUInt8OnlineLinearMethod
config = build_quant_config("int8")
layer = MagicMock(spec=LinearBase)
layer = mocker.Mock(spec=LinearBase)
prefix = "test_layer"
# Mock the platform to be NPU
with (
patch("vllm_omni.platforms.current_omni_platform.is_cuda", return_value=False),
patch("vllm_omni.platforms.current_omni_platform.is_npu", return_value=True),
):
method = config.get_quant_method(layer, prefix)
assert isinstance(method, NPUInt8OnlineLinearMethod)
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
monkeypatch.setattr(current_omni_platform, "is_npu", lambda: True)
method = config.get_quant_method(layer, prefix)
assert isinstance(method, NPUInt8OnlineLinearMethod)
# Test skipping quantization for a layer
config.ignored_layers = [prefix]
@@ -245,7 +239,7 @@ class TestNPUInt8LinearMethod:
@pytest.fixture
def mock_torch_npu(self, mocker):
torch_npu = MagicMock()
torch_npu = mocker.MagicMock()
mocker.patch("vllm_omni.quantization.int8_config.torch_npu", return_value=torch_npu)
mocker.patch(
tests/diffusion/test_diffusion_scheduler.py
+63 -40
View File
@@ -4,10 +4,10 @@
import queue
import threading
from types import SimpleNamespace
from unittest.mock import Mock, patch
import pytest
import torch
from pytest_mock import MockerFixture
from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
@@ -97,19 +97,19 @@ class _StubScheduler(SchedulerInterface):
def add_request(self, request: OmniDiffusionRequest) -> str:
assert request is self._request
self._state = Mock(sched_req_id=self._sched_req_id, req=request)
self._state = SimpleNamespace(sched_req_id=self._sched_req_id, req=request)
return self._sched_req_id
def schedule(self):
if self._scheduled or self._state is None:
return Mock(
return SimpleNamespace(
scheduled_new_reqs=[],
scheduled_cached_reqs=CachedRequestData.make_empty(),
scheduled_req_ids=[],
is_empty=True,
)
self._scheduled = True
return Mock(
return SimpleNamespace(
scheduled_new_reqs=[NewRequestData.from_state(self._state)],
scheduled_cached_reqs=CachedRequestData.make_empty(),
scheduled_req_ids=[self._state.sched_req_id],
@@ -153,7 +153,7 @@ class _StubScheduler(SchedulerInterface):
class TestRequestScheduler:
def setup_method(self) -> None:
self.scheduler: RequestScheduler = RequestScheduler()
self.scheduler.initialize(Mock())
self.scheduler.initialize(SimpleNamespace())
def test_single_request_success_lifecycle(self) -> None:
req_id = self.scheduler.add_request(_make_request("a"))
@@ -276,23 +276,23 @@ class TestRequestScheduler:
class TestDiffusionEngine:
def test_add_req_and_wait_for_response_single_path(self) -> None:
def test_add_req_and_wait_for_response_single_path(self, mocker: MockerFixture) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.scheduler = RequestScheduler()
engine.scheduler.initialize(Mock())
engine.scheduler.initialize(SimpleNamespace())
engine._rpc_lock = threading.RLock()
engine.abort_queue = queue.Queue()
request = _make_request("engine")
runner_output = _make_request_output("engine")
engine.execute_fn = Mock(return_value=runner_output)
engine.execute_fn = mocker.Mock(return_value=runner_output)
output = engine.add_req_and_wait_for_response(request)
assert output is runner_output.result
engine.execute_fn.assert_called_once()
def test_supports_scheduler_interface_injection(self) -> None:
def test_supports_scheduler_interface_injection(self, mocker: MockerFixture) -> None:
request = _make_request("engine_iface")
runner_output = _make_request_output("engine_iface")
scheduler = _StubScheduler(request, runner_output)
@@ -301,33 +301,45 @@ class TestDiffusionEngine:
engine.scheduler = scheduler
engine._rpc_lock = threading.RLock()
engine.abort_queue = queue.Queue()
engine.execute_fn = Mock(return_value=runner_output)
engine.execute_fn = mocker.Mock(return_value=runner_output)
output = engine.add_req_and_wait_for_response(request)
assert output is runner_output.result
engine.execute_fn.assert_called_once()
def test_initializes_injected_scheduler(self) -> None:
def test_initializes_injected_scheduler(
self,
monkeypatch: pytest.MonkeyPatch,
mocker: MockerFixture,
) -> None:
request = _make_request("init")
scheduler = _StubScheduler(request, DiffusionOutput(output=None))
od_config = Mock(model_class_name="mock_model")
fake_executor_cls = Mock(return_value=Mock())
od_config = SimpleNamespace(model_class_name="mock_model")
fake_executor_cls = mocker.Mock(return_value=mocker.Mock())
with (
patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func", return_value=None),
patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func", return_value=None),
patch("vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class", return_value=fake_executor_cls),
patch.object(DiffusionEngine, "_dummy_run", return_value=None),
):
DiffusionEngine(od_config, scheduler=scheduler)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
lambda *args, **kwargs: None,
)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
lambda *args, **kwargs: None,
)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
lambda *args, **kwargs: fake_executor_cls,
)
monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None)
DiffusionEngine(od_config, scheduler=scheduler)
assert scheduler.initialized_with is od_config
fake_executor_cls.assert_called_once_with(od_config)
def test_scheduler_alias_keeps_default_request_scheduler(self) -> None:
scheduler = Scheduler()
scheduler.initialize(Mock())
scheduler.initialize(SimpleNamespace())
req_id = scheduler.add_request(_make_request("alias"))
sched_output = scheduler.schedule()
@@ -336,10 +348,10 @@ class TestDiffusionEngine:
assert req_id in finished
assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
def test_step_raises_aborted_error(self) -> None:
def test_step_raises_aborted_error(self, mocker: MockerFixture) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.pre_process_func = None
engine.add_req_and_wait_for_response = Mock(
engine.add_req_and_wait_for_response = mocker.Mock(
return_value=DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
)
@@ -349,7 +361,7 @@ class TestDiffusionEngine:
def test_abort_queue_marks_request_finished_aborted(self) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.scheduler = RequestScheduler()
engine.scheduler.initialize(Mock())
engine.scheduler.initialize(SimpleNamespace())
engine.abort_queue = queue.Queue()
req_id = engine.scheduler.add_request(_make_request("req-abort"))
@@ -361,7 +373,7 @@ class TestDiffusionEngine:
def test_finalize_finished_request_returns_aborted_output(self) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.scheduler = RequestScheduler()
engine.scheduler.initialize(Mock())
engine.scheduler.initialize(SimpleNamespace())
req_id = engine.scheduler.add_request(_make_request("req-finalize"))
engine.scheduler.finish_requests(req_id, DiffusionRequestStatus.FINISHED_ABORTED)
@@ -371,29 +383,40 @@ class TestDiffusionEngine:
assert output.aborted is True
assert output.abort_message == "Request req-finalize aborted."
def test_initializes_step_scheduler_when_step_execution_enabled(self) -> None:
od_config = Mock(model_class_name="mock_model")
def test_initializes_step_scheduler_when_step_execution_enabled(
self,
monkeypatch: pytest.MonkeyPatch,
mocker: MockerFixture,
) -> None:
od_config = SimpleNamespace(model_class_name="mock_model")
od_config.step_execution = True
fake_executor = Mock()
fake_executor_cls = Mock(return_value=fake_executor)
fake_executor = mocker.Mock()
fake_executor_cls = mocker.Mock(return_value=fake_executor)
with (
patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func", return_value=None),
patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func", return_value=None),
patch("vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class", return_value=fake_executor_cls),
patch.object(DiffusionEngine, "_dummy_run", return_value=None),
):
engine = DiffusionEngine(od_config)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func",
lambda *args, **kwargs: None,
)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func",
lambda *args, **kwargs: None,
)
monkeypatch.setattr(
"vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class",
lambda *args, **kwargs: fake_executor_cls,
)
monkeypatch.setattr(DiffusionEngine, "_dummy_run", lambda self: None)
engine = DiffusionEngine(od_config)
assert isinstance(engine.scheduler, StepScheduler)
assert engine.execute_fn is fake_executor.execute_step
fake_executor_cls.assert_called_once_with(od_config)
def test_dummy_run_raises_on_output_error(self) -> None:
def test_dummy_run_raises_on_output_error(self, mocker: MockerFixture) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.od_config = Mock(model_class_name="mock_model")
engine.od_config = SimpleNamespace(model_class_name="mock_model")
engine.pre_process_func = None
engine.add_req_and_wait_for_response = Mock(return_value=DiffusionOutput(error="boom"))
engine.add_req_and_wait_for_response = mocker.Mock(return_value=DiffusionOutput(error="boom"))
with pytest.raises(RuntimeError, match="Dummy run failed: boom"):
engine._dummy_run()
@@ -402,7 +425,7 @@ class TestDiffusionEngine:
class TestStepScheduler:
def setup_method(self) -> None:
self.scheduler: StepScheduler = StepScheduler()
self.scheduler.initialize(Mock())
self.scheduler.initialize(SimpleNamespace())
def test_single_request_step_lifecycle(self) -> None:
request = _make_step_request("step", num_inference_steps=3)
tests/diffusion/test_diffusion_step_pipeline.py
+13 -13
View File
@@ -7,10 +7,10 @@ import queue
import threading
from contextlib import contextmanager
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
from pytest_mock import MockerFixture
import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
from tests.utils import hardware_test
@@ -542,11 +542,11 @@ class TestWorker:
class TestExecutor:
"""MultiprocDiffusionExecutor.execute_step"""
def test_execute_step_passes_through_runner_output(self):
def test_execute_step_passes_through_runner_output(self, mocker: MockerFixture):
executor = object.__new__(MultiprocDiffusionExecutor)
executor._ensure_open = lambda: None
expected = RunnerOutput(req_id="req-step", step_index=1, finished=False, result=None)
executor.collective_rpc = Mock(return_value=expected)
executor.collective_rpc = mocker.Mock(return_value=expected)
request = _make_engine_request("req-step", num_inference_steps=2)
scheduler_output = _make_scheduler_output(request, sched_req_id="req-step")
@@ -578,9 +578,9 @@ class TestEngine:
),
],
)
def test_step_engine_returns_error(self, execute_fn, expected_error):
def test_step_engine_returns_error(self, execute_fn, expected_error, mocker: MockerFixture):
scheduler = StepScheduler()
scheduler.initialize(Mock())
scheduler.initialize(mocker.Mock())
engine = _make_engine(scheduler, execute_fn=execute_fn)
output = engine.add_req_and_wait_for_response(_make_engine_request("req-error", num_inference_steps=2))
@@ -588,9 +588,9 @@ class TestEngine:
assert output.output is None
assert expected_error in output.error
def test_step_execution_completes(self):
def test_step_execution_completes(self, mocker: MockerFixture):
scheduler = StepScheduler()
scheduler.initialize(Mock())
scheduler.initialize(mocker.Mock())
engine = _make_engine(scheduler)
request = _make_engine_request("req-step", num_inference_steps=2)
@@ -614,9 +614,9 @@ class TestEngine:
assert output.error is None
assert torch.equal(output.output, torch.tensor([2.0]))
def test_step_abort_stops_rescheduling_after_first_step(self):
def test_step_abort_stops_rescheduling_after_first_step(self, mocker: MockerFixture):
scheduler = StepScheduler()
scheduler.initialize(Mock())
scheduler.initialize(mocker.Mock())
engine = _make_engine(scheduler)
request = _make_engine_request("req-stop", num_inference_steps=4)
@@ -639,9 +639,9 @@ class TestEngine:
assert step["n"] == 1
_assert_aborted_output(output, "req-stop")
def test_step_abort_after_reschedule_returns_aborted_output(self):
def test_step_abort_after_reschedule_returns_aborted_output(self, mocker: MockerFixture):
scheduler = StepScheduler()
scheduler.initialize(Mock())
scheduler.initialize(mocker.Mock())
engine = _make_engine(scheduler)
request = _make_engine_request("req-mid", num_inference_steps=4)
@@ -666,9 +666,9 @@ class TestEngine:
assert step["n"] == 2
_assert_aborted_output(output, "req-mid")
def test_finished_step_without_result_returns_error(self):
def test_finished_step_without_result_returns_error(self, mocker: MockerFixture):
scheduler = StepScheduler()
scheduler.initialize(Mock())
scheduler.initialize(mocker.Mock())
engine = _make_engine(
scheduler,
execute_fn=lambda _: RunnerOutput(
tests/diffusion/test_diffusion_worker_cuda_profiler.py
+2 -4
View File
@@ -1,8 +1,6 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import MagicMock
import pytest
from pytest_mock import MockerFixture
@@ -55,8 +53,8 @@ class TestDiffusionWorkerCudaProfiler:
mock_diffusion_worker_dependencies,
):
fake_profiler = mocker.Mock()
fake_profiler.start = MagicMock()
fake_profiler.stop = MagicMock()
fake_profiler.start = mocker.Mock()
fake_profiler.stop = mocker.Mock()
mocker.patch(
"vllm_omni.diffusion.worker.diffusion_worker.CudaProfilerWrapper",
return_value=fake_profiler,
tests/diffusion/test_multiproc_engine_concurrency.py
+12 -16
View File
@@ -3,7 +3,7 @@
import queue
import threading
from unittest.mock import Mock, patch
from types import SimpleNamespace
import pytest
import torch
@@ -24,11 +24,9 @@ def _tagged_output(tag: str) -> DiffusionOutput:
return DiffusionOutput(output=torch.tensor([0]), error=tag)
def _mock_request(tag: str) -> Mock:
"""Return a mock ``OmniDiffusionRequest`` identifiable by *tag*."""
req = Mock()
req.request_ids = [tag]
return req
def _mock_request(tag: str):
"""Return a lightweight request object identifiable by *tag*."""
return SimpleNamespace(request_ids=[tag])
def _make_executor(num_gpus: int = 1):
@@ -36,20 +34,18 @@ def _make_executor(num_gpus: int = 1):
Returns ``(executor, request_queue, result_queue)``.
"""
od_cfg = Mock()
od_cfg.num_gpus = num_gpus
with patch.object(MultiprocDiffusionExecutor, "_init_executor"):
executor = MultiprocDiffusionExecutor(od_cfg)
od_cfg = SimpleNamespace(num_gpus=num_gpus)
monkeypatch = pytest.MonkeyPatch()
monkeypatch.setattr(MultiprocDiffusionExecutor, "_init_executor", lambda self: None)
executor = MultiprocDiffusionExecutor(od_cfg)
monkeypatch.undo()
req_q: queue.Queue = queue.Queue()
res_q: queue.Queue = queue.Queue()
mock_broadcast_mq = Mock()
mock_broadcast_mq.enqueue = req_q.put
mock_broadcast_mq = SimpleNamespace(enqueue=req_q.put)
mock_rmq = Mock()
mock_rmq.dequeue = lambda timeout=None: res_q.get(timeout=timeout if timeout is not None else 10)
mock_rmq = SimpleNamespace(dequeue=lambda timeout=None: res_q.get(timeout=timeout if timeout is not None else 10))
executor._broadcast_mq = mock_broadcast_mq
executor._result_mq = mock_rmq
@@ -63,7 +59,7 @@ def _make_engine(num_gpus: int = 1):
executor, req_q, res_q = _make_executor(num_gpus)
engine = DiffusionEngine.__new__(DiffusionEngine)
sched = RequestScheduler()
sched.initialize(Mock())
sched.initialize(SimpleNamespace())
engine.scheduler = sched
engine.executor = executor
engine._rpc_lock = threading.RLock()