# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """ Unit tests for WorkerWrapperBase class. This module tests the WorkerWrapperBase implementation: - Initialization with and without worker extensions - Custom pipeline initialization - Method delegation via execute_method - Attribute delegation via __getattr__ - Dynamic worker class extension """ import gc import weakref from types import SimpleNamespace from typing import Any import pytest import torch from pytest_mock import MockerFixture from vllm_omni.diffusion.worker.diffusion_worker import ( CustomPipelineWorkerExtension, DiffusionWorker, WorkerWrapperBase, ) pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] # ------------------------------------------------------------------------- # Fixtures # ------------------------------------------------------------------------- @pytest.fixture def mock_od_config(mocker: MockerFixture): """Create a mock OmniDiffusionConfig for use in tests.""" config = mocker.Mock() config.num_gpus = 1 config.master_port = 12345 config.enable_sleep_mode = False config.cache_backend = None config.cache_config = None config.model = "test-model" config.diffusion_load_format = "default" config.dtype = "float32" config.max_cpu_loras = 0 config.lora_path = None config.lora_scale = 1.0 return config class TestExtension: """Simple test extension adding one custom method.""" def custom_method(self): return "extension_method" class MockCustomPipeline: """Mock custom pipeline for testing.""" def __init__(self, *args, **kwargs): self.args = args self.kwargs = kwargs def __call__(self, *args, **kwargs): return "pipeline_output" # ------------------------------------------------------------------------- # Tests: Initialization # ------------------------------------------------------------------------- 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) wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker, ) assert wrapper.gpu_id == 0 assert wrapper.od_config == mock_od_config assert wrapper.base_worker_class == DiffusionWorker assert wrapper.worker_extension_cls is None assert wrapper.custom_pipeline_args is None assert wrapper.worker is not None mock_worker_init.assert_called_once_with( local_rank=0, rank=0, od_config=mock_od_config, skip_load_model=False, ) def test_custom_pipeline_args_without_pipeline_class_loads_normally(self, mocker: MockerFixture, mock_od_config): """Test native pipeline args do not trigger custom pipeline re-initialization.""" custom_args = {"components_path": "/tmp/anima-components"} mock_worker_init = mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker, custom_pipeline_args=custom_args, ) assert wrapper.worker_extension_cls is None assert CustomPipelineWorkerExtension not in wrapper.worker.__class__.__bases__ assert not hasattr(wrapper.worker, "re_init_pipeline") mock_worker_init.assert_called_once_with( local_rank=0, rank=0, od_config=mock_od_config, 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 # ------------------------------------------------------------------------- class TestWorkerWrapperBaseExtension: """Test WorkerWrapperBase worker extension functionality.""" def test_prepare_worker_class_without_extension(self, mocker: MockerFixture, mock_od_config): """Test _prepare_worker_class without a worker extension.""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker, ) worker_class = wrapper._prepare_worker_class() assert worker_class == DiffusionWorker def test_prepare_worker_class_with_extension_class(self, mocker: MockerFixture, mock_od_config): """Test _prepare_worker_class with an explicit extension class.""" class TestExtension: def custom_method(self): return "extension_method" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker, worker_extension_cls=TestExtension, ) assert hasattr(wrapper.worker.__class__, "custom_method") assert TestExtension in wrapper.worker.__class__.__bases__ def test_prepare_worker_class_with_extension_string(self, mocker: MockerFixture, mock_od_config): """Test _prepare_worker_class with worker extension as string.""" mock_resolve = mocker.patch("vllm.utils.import_utils.resolve_obj_by_qualname") mock_resolve.return_value = TestExtension mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker, worker_extension_cls="tests.diffusion.test_worker_wrapper_base.TestExtension", ) assert hasattr(wrapper.worker.__class__, "custom_method") # ------------------------------------------------------------------------- # Tests: Method Delegation # ------------------------------------------------------------------------- class TestWorkerWrapperBaseDelegation: """Test WorkerWrapperBase delegation to wrapped worker.""" def test_execute_model_delegation(self, mocker: MockerFixture, mock_od_config): """Test that execute_model() delegates to worker.execute_model().""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) mock_output = mocker.Mock() wrapper.worker.execute_model = mocker.Mock(return_value=mock_output) mock_req = mocker.Mock() result = wrapper.execute_model(mock_req, mock_od_config) wrapper.worker.execute_model.assert_called_once_with(mock_req, mock_od_config, kv_prefetch_job=None) assert result == mock_output def test_sleep_delegation(self, mocker: MockerFixture, mock_od_config): """Test that sleep() delegates to worker.sleep().""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) wrapper.worker.sleep = mocker.Mock(return_value=True) result = wrapper.sleep(level=1) wrapper.worker.sleep.assert_called_once_with(1) assert result is True def test_wake_up_delegation(self, mocker: MockerFixture, mock_od_config): """Test that wake_up() delegates to worker.wake_up().""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) wrapper.worker.wake_up = mocker.Mock(return_value=True) result = wrapper.wake_up(tags=["weights"]) wrapper.worker.wake_up.assert_called_once_with(["weights"]) assert result is True def test_shutdown_delegation(self, mocker: MockerFixture, mock_od_config): """Test that shutdown() delegates to worker.shutdown().""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) wrapper.worker.shutdown = mocker.Mock(return_value=None) result = wrapper.shutdown() wrapper.worker.shutdown.assert_called_once() assert result is None @pytest.mark.parametrize("native_kv", [False, True]) def test_worker_shutdown_releases_model_before_cumem_pools(self, mocker: MockerFixture, native_kv): events: list[str] = [] offload_backend = mocker.Mock() offload_backend.shutdown.side_effect = lambda: events.append("offload") kv_manager = mocker.Mock() kv_manager.close.side_effect = lambda: events.append("kv") 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", "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() 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): mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env") mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.shutdown_kv_connector") offload_backend = mocker.Mock() if cleanup_fails: offload_backend.shutdown.side_effect = RuntimeError("offload cleanup failed") worker = DiffusionWorker.__new__(DiffusionWorker) weights = torch.ones(1) weights_ref = weakref.ref(weights) pipeline = SimpleNamespace(weights=weights) worker.model_runner = SimpleNamespace( pipeline=pipeline, offload_backend=offload_backend, kv_transfer_manager=None ) worker.lora_manager = SimpleNamespace(pipeline=pipeline) worker._sleep_saved_buffers = {"buffer": weights} del weights, pipeline worker.shutdown() gc.collect() assert worker.model_runner is None assert worker.lora_manager is None assert worker._sleep_saved_buffers == {} assert weights_ref() is None @pytest.mark.parametrize("owns_pool", [False, True]) @pytest.mark.parametrize("cleanup_fails", [False, True]) def test_worker_shutdown_releases_only_owned_sleep_pool(self, mocker: MockerFixture, owns_pool, cleanup_fails): destroy = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env") shutdown_connector = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.shutdown_kv_connector") allocator_class = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker._get_cumem_allocator_class") release = allocator_class.return_value.get_instance.return_value.release_pools if cleanup_fails: release.side_effect = RuntimeError("pool release failed") worker = DiffusionWorker.__new__(DiffusionWorker) worker.model_runner = None worker._owns_sleep_pool = owns_pool worker.shutdown() assert release.call_count == int(owns_pool) shutdown_connector.assert_called_once_with() destroy.assert_called_once_with() if not cleanup_fails: assert not worker._owns_sleep_pool # ------------------------------------------------------------------------- # Tests: execute_method # ------------------------------------------------------------------------- class TestWorkerWrapperBaseExecuteMethod: """Test WorkerWrapperBase.execute_method functionality.""" def test_execute_method_success(self, mocker: MockerFixture, mock_od_config): """Test execute_method successfully calls worker method.""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) wrapper.worker.test_method = mocker.Mock(return_value="method_result") result = wrapper.execute_method("test_method", "arg1", kwarg1="value1") wrapper.worker.test_method.assert_called_once_with("arg1", kwarg1="value1") assert result == "method_result" def test_execute_method_with_no_args(self, mocker: MockerFixture, mock_od_config): """Test execute_method with no arguments.""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) wrapper.worker.no_args_method = mocker.Mock(return_value="no_args_result") result = wrapper.execute_method("no_args_method") wrapper.worker.no_args_method.assert_called_once_with() assert result == "no_args_result" def test_execute_method_error(self, mocker: MockerFixture, mock_od_config): """Test execute_method raises exception on error.""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) wrapper.worker.error_method = mocker.Mock(side_effect=RuntimeError("Test error")) with pytest.raises(RuntimeError, match="Test error"): wrapper.execute_method("error_method") def test_execute_method_invalid_type(self, mocker: MockerFixture, mock_od_config): """Test execute_method with invalid method type.""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) with pytest.raises(AssertionError, match="Method must be str"): wrapper.execute_method(b"bytes_method") # ------------------------------------------------------------------------- # Tests: __getattr__ delegation # ------------------------------------------------------------------------- class TestWorkerWrapperBaseGetAttr: """Test WorkerWrapperBase.__getattr__ delegation.""" def test_getattr_delegation(self, mocker: MockerFixture, mock_od_config): """Test __getattr__ delegates to worker attributes.""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) wrapper.worker.custom_attribute = "test_value" assert wrapper.custom_attribute == "test_value" def test_getattr_method_access(self, mocker: MockerFixture, mock_od_config): """Test __getattr__ delegates to worker methods.""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) wrapper.worker.custom_method = mocker.Mock(return_value="method_result") result = wrapper.custom_method() wrapper.worker.custom_method.assert_called_once() assert result == "method_result" def test_getattr_missing_attribute(self, mocker: MockerFixture, mock_od_config): """Test __getattr__ raises AttributeError for missing attributes.""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker) with pytest.raises(AttributeError): _ = wrapper.nonexistent_attribute # ------------------------------------------------------------------------- # Tests: Edge Cases # ------------------------------------------------------------------------- class TestWorkerWrapperBaseEdgeCases: """Test WorkerWrapperBase edge cases and special scenarios.""" def test_extension_conflict_warning(self, mocker: MockerFixture, mock_od_config, caplog): """Test a warning is logged when an extension conflicts with worker.""" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) class ConflictExtension: def load_model(self): return "extension_load_model" mocker.patch.object(DiffusionWorker, "load_model") wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker, worker_extension_cls=ConflictExtension, ) assert wrapper.worker is not None def test_multiple_extensions_same_class(self, mocker: MockerFixture, mock_od_config): """Test that applying same extension twice doesn't duplicate it.""" class TestExtension: def custom_method(self): return "extension" mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper1 = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker, worker_extension_cls=TestExtension, ) wrapper2 = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker, worker_extension_cls=TestExtension, ) assert hasattr(wrapper1.worker, "custom_method") assert hasattr(wrapper2.worker, "custom_method") # ------------------------------------------------------------------------- # Tests: CustomPipelineWorkerExtension # ------------------------------------------------------------------------- class TestCustomPipelineWorkerExtension: """Test CustomPipelineWorkerExtension functionality.""" def test_re_init_pipeline_basic(self, mocker: MockerFixture, mock_od_config): """Test basic re_init_pipeline functionality.""" mocker.patch("torch.accelerator.empty_cache") mocker.patch("gc.collect") mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker, worker_extension_cls=CustomPipelineWorkerExtension, ) # Setup mock model_runner and pipeline mock_model_runner = mocker.Mock() mock_pipeline = mocker.Mock() mock_model_runner.pipeline = mock_pipeline wrapper.worker.model_runner = mock_model_runner wrapper.worker.init_lora_manager = mocker.Mock() wrapper.worker.load_model = mocker.Mock() custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} # Call re_init_pipeline wrapper.worker.re_init_pipeline(custom_args) # Verify load_model was called with correct arguments wrapper.worker.load_model.assert_called_once_with( load_format="custom_pipeline", custom_pipeline_name="tests.diffusion.test_worker_wrapper_base.MockCustomPipeline", ) wrapper.worker.init_lora_manager.assert_called_once() def test_re_init_pipeline_cleanup(self, mocker: MockerFixture, mock_od_config): """Test that re_init_pipeline properly cleans up old pipeline.""" mock_gc_collect = mocker.patch("gc.collect") mock_empty_cache = mocker.patch("torch.accelerator.empty_cache") mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker, worker_extension_cls=CustomPipelineWorkerExtension, ) # Setup mock model_runner with pipeline mock_model_runner = mocker.Mock() mock_pipeline = mocker.Mock() mock_model_runner.pipeline = mock_pipeline wrapper.worker.model_runner = mock_model_runner wrapper.worker.init_lora_manager = mocker.Mock() wrapper.worker.load_model = mocker.Mock() custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} # Call re_init_pipeline wrapper.worker.re_init_pipeline(custom_args) # Verify cleanup was performed mock_gc_collect.assert_called_once() mock_empty_cache.assert_called_once() def test_re_init_pipeline_none_pipeline(self, mocker: MockerFixture, mock_od_config): """Test re_init_pipeline when pipeline is None.""" mocker.patch("torch.accelerator.empty_cache") mocker.patch("gc.collect") mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker, worker_extension_cls=CustomPipelineWorkerExtension, ) # Setup mock model_runner with None pipeline mock_model_runner = mocker.Mock() mock_model_runner.pipeline = None wrapper.worker.model_runner = mock_model_runner wrapper.worker.init_lora_manager = mocker.Mock() wrapper.worker.load_model = mocker.Mock() custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} # Should not raise an error wrapper.worker.re_init_pipeline(custom_args) # Verify load_model was still called wrapper.worker.load_model.assert_called_once() def test_custom_pipeline_args_initialization(self, mocker: MockerFixture, mock_od_config): """Test initialization with custom_pipeline_args calls re_init_pipeline.""" custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} mocker.patch.object(DiffusionWorker, "__init__", return_value=None) mock_prepare = mocker.patch.object(WorkerWrapperBase, "_prepare_worker_class") # Create a mock worker class with re_init_pipeline mock_worker_class = mocker.Mock() mock_worker_instance = mocker.Mock() mock_worker_instance.re_init_pipeline = mocker.Mock() mock_worker_class.return_value = mock_worker_instance mock_prepare.return_value = mock_worker_class _ = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker, custom_pipeline_args=custom_args, ) mock_worker_class.assert_called_once_with( local_rank=0, rank=0, od_config=mock_od_config, skip_load_model=True, ) # Verify re_init_pipeline was called with custom_pipeline_args mock_worker_instance.re_init_pipeline.assert_called_once_with(custom_args) @pytest.mark.parametrize("cleanup_fails", [False, True]) def test_custom_pipeline_initialization_failure_cleans_worker( self, mocker: MockerFixture, mock_od_config, cleanup_fails ): worker = mocker.Mock() worker.re_init_pipeline.side_effect = RuntimeError("custom pipeline load failed") if cleanup_fails: worker.shutdown.side_effect = RuntimeError("cleanup failed") worker_class = mocker.Mock(return_value=worker) mocker.patch.object(WorkerWrapperBase, "_prepare_worker_class", return_value=worker_class) with pytest.raises(RuntimeError, match="custom pipeline load failed"): WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, custom_pipeline_args={"pipeline_class": "unused"}, ) worker.shutdown.assert_called_once_with() def test_custom_pipeline_with_explicit_extension(self, mocker: MockerFixture, mock_od_config): """Test that explicit worker_extension_cls is preserved when custom_pipeline_args is provided.""" class CustomExtension: def re_init_pipeline(self, custom_pipeline_args: dict[str, Any]): return "custom_re_init_pipeline" def custom_extension_method(self): return "custom_extension_method" custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker, worker_extension_cls=CustomExtension, custom_pipeline_args=custom_args, ) # Should still have the explicitly provided extension assert CustomExtension in wrapper.worker.__class__.__bases__ assert hasattr(wrapper.worker, "custom_extension_method") def test_re_init_pipeline_multiple_calls(self, mocker: MockerFixture, mock_od_config): """Test calling re_init_pipeline multiple times.""" mocker.patch("torch.accelerator.empty_cache") mocker.patch("gc.collect") mocker.patch.object(DiffusionWorker, "__init__", return_value=None) wrapper = WorkerWrapperBase( gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker, worker_extension_cls=CustomPipelineWorkerExtension, ) # Setup mock model_runner mock_model_runner = mocker.Mock() mock_pipeline1 = mocker.Mock() mock_pipeline2 = mocker.Mock() mock_model_runner.pipeline = mock_pipeline1 wrapper.worker.model_runner = mock_model_runner wrapper.worker.init_lora_manager = mocker.Mock() wrapper.worker.load_model = mocker.Mock() # First call custom_args1 = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} wrapper.worker.re_init_pipeline(custom_args1) # Update pipeline for second call mock_model_runner.pipeline = mock_pipeline2 # Second call custom_args2 = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"} wrapper.worker.re_init_pipeline(custom_args2) # Verify load_model was called twice with different pipelines assert wrapper.worker.load_model.call_count == 2 assert wrapper.worker.init_lora_manager.call_count == 2 @pytest.mark.parametrize( ("init_method", "rank", "local_rank", "expected_addr", "expected_port"), [ (None, 1, 1, "localhost", "12345"), ("tcp://10.0.0.2:23456", 7, 0, "unrelated-host", "9999"), ], ) def test_worker_rendezvous_preserves_local_device_rank( monkeypatch, init_method, rank, local_rank, expected_addr, expected_port ): """TCP rendezvous must not confuse a global rank with the actor's GPU index.""" import os from contextlib import nullcontext from unittest.mock import Mock from vllm_omni.diffusion.worker import diffusion_worker as worker_module class ReachedRendezvousError(Exception): pass worker = object.__new__(DiffusionWorker) worker.rank = rank worker.local_rank = local_rank worker.distributed_init_method = init_method worker.od_config = SimpleNamespace(num_gpus=8, master_port=12345, moe_backend="auto") config = SimpleNamespace(kernel_config=SimpleNamespace(ir_op_priority=None)) platform = Mock() platform.get_torch_device.return_value = "cuda:0" rendezvous = Mock(side_effect=ReachedRendezvousError) monkeypatch.setenv("MASTER_ADDR", "unrelated-host") monkeypatch.setenv("MASTER_PORT", "9999") # Restore these variables after init_device publishes the worker identity. for key in ("LOCAL_RANK", "RANK", "WORLD_SIZE"): monkeypatch.setenv(key, "stale") monkeypatch.setattr(worker_module, "current_omni_platform", platform) monkeypatch.setattr(worker_module, "create_diffusion_vllm_config", Mock(return_value=config)) monkeypatch.setattr(worker_module, "_resolve_ir_op_priority", Mock(return_value=[])) monkeypatch.setattr(worker_module, "set_forward_context", lambda **kwargs: nullcontext()) monkeypatch.setattr(worker_module, "set_current_vllm_config", lambda config: nullcontext()) monkeypatch.setattr(worker_module, "init_distributed_environment", rendezvous) with pytest.raises(ReachedRendezvousError): worker.init_device() rendezvous.assert_called_once_with( world_size=8, rank=rank, distributed_init_method=init_method or "env://", local_rank=local_rank, ) platform.get_torch_device.assert_called_once_with(local_rank) assert os.environ["MASTER_ADDR"] == expected_addr assert os.environ["MASTER_PORT"] == expected_port assert os.environ["LOCAL_RANK"] == str(local_rank) assert os.environ["RANK"] == str(rank)