import uuid import warnings from queue import Empty, Queue from typing import Any import pytest import torch from pytest_mock import MockerFixture from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK from vllm_omni.inputs.data import OmniDiffusionSamplingParams pytestmark = [pytest.mark.core_model, pytest.mark.cpu] # Suppress noisy DeprecationWarnings from optional Swig bindings imported by vLLM dependencies. warnings.filterwarnings( "ignore", message=r"builtin type SwigPy.*has no __module__ attribute", category=DeprecationWarning, ) MODEL = "riverclouds/qwen_image_random" class _FakeStageRequestStats: """Fake StageRequestStats object with necessary attributes aligned with real StageRequestStats.""" def __init__(self, **kwargs): # Required fields (with defaults for testing) self.batch_id = kwargs.get("batch_id", 0) self.batch_size = kwargs.get("batch_size", 1) self.num_tokens_in = kwargs.get("num_tokens_in", 0) self.num_tokens_out = kwargs.get("num_tokens_out", 1) self.stage_gen_time_ms = kwargs.get("stage_gen_time_ms", 10.0) self.rx_transfer_bytes = kwargs.get("rx_transfer_bytes", 0) self.rx_decode_time_ms = kwargs.get("rx_decode_time_ms", 0.0) self.rx_in_flight_time_ms = kwargs.get("rx_in_flight_time_ms", 0.0) self.stage_stats = kwargs.get("stage_stats", None) # Optional fields self.stage_id = kwargs.get("stage_id", None) self.final_output_type = kwargs.get("final_output_type", None) self.request_id = kwargs.get("request_id", None) self.postprocess_time_ms = kwargs.get("postprocess_time_ms", 0.0) self.diffusion_metrics = kwargs.get("diffusion_metrics", None) self.audio_generated_frames = kwargs.get("audio_generated_frames", 0) # Allow additional attributes for flexibility for key, value in kwargs.items(): if not hasattr(self, key): setattr(self, key, value) class _FakeEngineArgs(dict): """Fake engine args that can be used both as object attributes and as **kwargs.""" def __init__(self, args_dict: dict[str, Any]): super().__init__(args_dict) # Add required attributes if not present if "model_stage" not in self: self["model_stage"] = None if "engine_output_type" not in self: self["engine_output_type"] = None # Also set as attributes for object-style access for key, value in self.items(): setattr(self, key, value) class _FakeStageConfig: """Fake stage config object that mimics the real stage config structure.""" def __init__(self, config_dict: dict[str, Any]): # engine_args needs to work both as object (for OmniStage) and as dict (for **kwargs) engine_args_dict = config_dict.get("engine_args", {}) self.engine_args = _FakeEngineArgs(engine_args_dict) self.final_output = config_dict.get("final_output", False) self.final_output_type = config_dict.get("final_output_type", None) self.stage_id = config_dict.get("stage_id", 0) # Store original dict for reference self._config_dict = config_dict class _FakeQueue: """Fake queue using standard library Queue to replace mp.Queue.""" def __init__(self, maxsize=0): self._queue = Queue(maxsize=maxsize) def put(self, item): self._queue.put(item) def put_nowait(self, item): self._queue.put_nowait(item) def get(self, timeout=None): if timeout is None: return self._queue.get() return self._queue.get(timeout=timeout) def get_nowait(self): return self._queue.get_nowait() def empty(self): return self._queue.empty() class _FakeZmqQueue: """Fake ZmqQueue that wraps _FakeQueue and matches ZmqQueue interface.""" def __init__( self, ctx=None, socket_type=None, *, bind: str | None = None, connect: str | None = None, recv_timeout_ms: int | None = None, send_timeout_ms: int | None = None, ): """Initialize fake ZMQ queue with same signature as real ZmqQueue.""" self._queue = _FakeQueue(maxsize=0) # Determine endpoint from bind or connect path = bind if bind is not None else connect self.endpoint = path or f"fake://zmq-endpoint-{id(self)}" self._recv_timeout_ms = recv_timeout_ms self._send_timeout_ms = send_timeout_ms def put(self, obj: Any) -> None: """Send an object to the queue.""" self._queue.put(obj) def put_nowait(self, obj: Any) -> None: """Send an object to the queue without blocking.""" self._queue.put_nowait(obj) def get(self, timeout: float | None = None) -> Any: """Receive an object from the queue with optional timeout in seconds.""" return self._queue.get(timeout=timeout) def get_nowait(self) -> Any: """Receive an object from the queue without blocking.""" return self._queue.get_nowait() def empty(self) -> bool: """Check if the queue is empty without blocking.""" return self._queue.empty() def close(self) -> None: """Close the queue.""" pass class _FakeStage: """Lightweight Stage stub for multi-process pipeline version with queue support.""" def __init__(self, mocker: MockerFixture, config, stage_init_timeout: int = 300): # Handle both dict and object configs if isinstance(config, dict): config = _FakeStageConfig(config) self.config = config self.stage_config = config self.engine = None self.engine_outputs = None # Set attributes that OmniStage expects self.stage_id = getattr(config, "stage_id", 0) self.engine_args = config.engine_args self.model_stage = getattr(config.engine_args, "model_stage", None) self.stage_type = "diffusion" # set default sampling params self.default_sampling_params = OmniDiffusionSamplingParams(num_inference_steps=1) # Allow configuring final_output and final_output_type self.final_output = config.final_output if hasattr(config, "final_output") else False self.final_output_type = getattr(config, "final_output_type", None) # Configurable processing logic, default returns placeholder processed_input = getattr(config, "_config_dict", {}).get("processed_input", ["processed"]) self._processed_input = processed_input # Queue references (set by attach_queues) self._in_q = None self._out_q = None self._proc = None # Mock process reference self._stage_init_timeout = max(0, int(stage_init_timeout)) # NOTE: pytest mocker fixture object is injected into test functions and classes # with specific naming (e.g. `class TestXXX`, `def test_XXX` ) automatically, # but **must be provided** when initializing this helper cls self._mocker = mocker def attach_queues(self, in_q, out_q): """Attach input and output queues.""" self._in_q = in_q self._out_q = out_q def init_stage_worker( self, model: str, *, is_async: bool = False, shm_threshold_bytes: int = 65536, ctx=None, batch_timeout: int = 10, **kwargs, ): """Mock init_stage_worker: don't start real process, just send stage_ready message.""" # Create a mock process object self._proc = self._mocker.MagicMock() self._proc.start = self._mocker.MagicMock() self._proc.join = self._mocker.MagicMock() self._proc.is_alive = self._mocker.MagicMock(return_value=False) self._proc.terminate = self._mocker.MagicMock() # Send stage_ready message to output queue if self._out_q is not None: try: self._out_q.put_nowait({"type": "stage_ready", "stage_id": self.stage_id}) except Exception: pass def stop_stage_worker(self): """Mock stop_stage_worker: clean up queue references.""" if self._in_q is not None: try: self._in_q.put_nowait(SHUTDOWN_TASK) except Exception: pass def submit(self, payload: dict[str, Any]): """Submit task to input queue.""" if self._in_q is not None: self._in_q.put(payload) def try_collect(self) -> Any: """Non-blocking collect from output queue.""" if self._out_q is None: return None try: return self._out_q.get_nowait() except Empty: return None def set_engine_outputs(self, outputs): """Set engine outputs for the stage.""" self.engine_outputs = outputs def process_engine_inputs(self, stage_list, prompts): """Process engine inputs: return preset processed result.""" return self._processed_input class _FakeEngine: """Lightweight Engine stub: provides generate iterator output.""" def __init__(self, outputs: list[Any]): self._outputs = outputs def generate(self, prompts, sampling_params): # Record the most recent prompts for outer assertions self._last_prompts = prompts # Simplified: return preset list at once, ensuring iterability yield from self._outputs @pytest.fixture def fake_stage_config(): return { # Don't include 'model' in engine_args since it's passed separately "engine_args": {}, "final_output": True, "final_output_type": "text", # Second stage will use processed_input to verify the chain "processed_input": ["processed-by-stage"], } def _setup_engine_mocks(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture): """Helper function to set up common engine mocks. fixture objects `monkeypatch` and `mocker` must be passed in. """ fake_engine = mocker.MagicMock() # Add necessary attributes to fake_engine fake_engine.tokenizer = mocker.MagicMock() fake_engine.log_stats = False fake_engine.vllm_config = mocker.MagicMock() fake_engine.vllm_config.model_config = mocker.MagicMock() fake_engine.vllm_config.model_config.io_processor_plugin = None fake_engine.get_supported_tasks = mocker.MagicMock(return_value=[]) fake_engine.model_config = mocker.MagicMock() fake_engine.model_config.io_processor_plugin = None # Add registry with resolve_model_cls method fake_registry = mocker.MagicMock() fake_registry.resolve_model_cls = mocker.MagicMock(return_value=(mocker.MagicMock(), "test_arch")) fake_engine.model_config.registry = fake_registry fake_engine.vllm_config.model_config.registry = fake_registry monkeypatch.setattr( "vllm.v1.engine.llm_engine.LLMEngine.from_engine_args", lambda **kw: fake_engine, raising=False, ) # Mock model_config.registry.resolve_model_cls to return a tuple # Use a real class instead of MagicMock to avoid inspect.getsource issues class FakeModelClass: pass monkeypatch.setattr( "vllm.model_executor.model_loader.utils.get_model_architecture", lambda model_config: (FakeModelClass, "test_arch"), raising=False, ) monkeypatch.setattr( "vllm.model_executor.model_loader.utils._get_model_architecture", lambda model_config: (FakeModelClass, "test_arch"), raising=False, ) # Mock try_create_mm_pooling_model_cls to return the class as-is monkeypatch.setattr( "vllm.model_executor.models.adapters.try_create_mm_pooling_model_cls", lambda model_cls: model_cls, raising=False, ) # Mock _enable_processor_cache to return False monkeypatch.setattr( "vllm.multimodal.cache._enable_processor_cache", lambda model_config, mm_registry: False, raising=False, ) # Mock get_io_processor to return None monkeypatch.setattr( "vllm.plugins.io_processors.get_io_processor", lambda vllm_config, io_processor_plugin: None, raising=False, ) def _setup_multiprocessing_mocks(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture): """Helper function to set up multiprocessing mocks. fixture objects `monkeypatch` and `mocker` must be passed in. """ import multiprocessing as mp # Mock Process fake_process_class = mocker.MagicMock() fake_process_instance = mocker.MagicMock() fake_process_instance.start = mocker.MagicMock() fake_process_instance.join = mocker.MagicMock() fake_process_instance.is_alive = mocker.MagicMock(return_value=True) fake_process_instance.terminate = mocker.MagicMock() fake_process_class.return_value = fake_process_instance # Mock get_context to return a context with Queue that returns _FakeQueue fake_ctx = mocker.MagicMock() fake_ctx.Queue = lambda maxsize=0: _FakeQueue(maxsize=maxsize) fake_ctx.Process = fake_process_class def _mock_get_context(method): return fake_ctx monkeypatch.setattr(mp, "get_context", _mock_get_context, raising=False) monkeypatch.setattr(mp, "Process", fake_process_class, raising=False) # Mock ZmqQueue to use _FakeZmqQueue monkeypatch.setattr( "vllm_omni.entrypoints.zmq_utils.ZmqQueue", _FakeZmqQueue, raising=False, ) # Also mock where ZmqQueue is imported/used monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.ZmqQueue", _FakeZmqQueue, raising=False, ) def _setup_ipc_mocks(monkeypatch: pytest.MonkeyPatch): """Helper function to set up IPC function mocks.""" # Mock _encode: simple serialization def _fake_encode(obj, threshold, obj_key, shm_key): return {obj_key: obj} # Mock _load: extract object from result def _fake_load(result, obj_key, shm_key): return result.get(obj_key) # Mock _set: calculate serialization size def _fake_set(obj): return str(obj).encode() monkeypatch.setattr("vllm_omni.entrypoints.omni._encode", _fake_encode, raising=False) monkeypatch.setattr("vllm_omni.entrypoints.omni._load", _fake_load, raising=False) monkeypatch.setattr("vllm_omni.entrypoints.omni._set", _fake_set, raising=False) def _setup_log_mocks(monkeypatch: pytest.MonkeyPatch): """Helper function to set up logging and stats mocks.""" # Mock OrchestratorMetrics to be a simple class that doesn't require file operations class _FakeOrchestratorMetrics: def __init__(self, num_stages, enable_stats, wall_start_ts): self.num_stages = num_stages self.enable_stats = enable_stats self.stage_first_ts = [None] * num_stages self.stage_last_ts = [None] * num_stages self.e2e_done = set() def on_stage_metrics(self, stage_id, req_id, metrics): pass def on_finalize_request(self, stage_id, req_id, start_ts): self.e2e_done.add(req_id) def on_forward(self, from_stage, to_stage, req_id, size_bytes, tx_ms, use_shm): pass def build_and_log_summary(self, final_stage_id): return "Fake summary" monkeypatch.setattr( "vllm_omni.entrypoints.omni.OrchestratorMetrics", _FakeOrchestratorMetrics, raising=False, ) def _setup_connector_mocks(monkeypatch, mocker, omni_module=None): """Helper function to set up connector mocks for stage-to-stage forwarding. If omni_module is provided, mocks directly on the module. Otherwise, uses string path. """ # Mock initialize_orchestrator_connectors to return fake connectors def _fake_initialize_orchestrator_connectors(config_path, worker_backend=None, shm_threshold_bytes=None): # Create fake connectors for all stage-to-stage edges # Each connector is just a mock object that will be passed to try_send_via_connector fake_connectors = {} # Add connectors for common edges (0->1, 1->2, etc.) for i in range(10): # Support up to 10 stages fake_connectors[(str(i), str(i + 1))] = mocker.MagicMock() return None, fake_connectors if omni_module is not None: # Mock directly on the omni module where it's used (after import) monkeypatch.setattr(omni_module, "initialize_orchestrator_connectors", _fake_initialize_orchestrator_connectors) else: # Mock via string path (before import) monkeypatch.setattr( "vllm_omni.entrypoints.omni.initialize_orchestrator_connectors", _fake_initialize_orchestrator_connectors, raising=False, ) def _setup_connector_adapter_mock(monkeypatch, omni_module): """Helper function to mock try_send_via_connector on the omni module. This must be called AFTER importing omni module, to mock the function where it's actually used. """ # Mock try_send_via_connector to always succeed def _fake_try_send_via_connector( connector, stage_id, next_stage_id, req_id, next_inputs, sampling_params, original_prompt, next_stage_queue_submit_fn, metrics, ): # Simulate successful send by calling the submit function task = { "request_id": req_id, "engine_inputs": next_inputs, "sampling_params": sampling_params, } next_stage_queue_submit_fn(task) return True # Mock directly on the omni module where it's used monkeypatch.setattr(omni_module, "try_send_via_connector", _fake_try_send_via_connector) @pytest.fixture(autouse=True) def mock_get_config(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture): """Auto-mock get_config and related model loading functions to avoid model path validation.""" # CRITICAL: Mock tokenizer-related imports FIRST, before any module imports # This prevents ImportError when async_omni is imported (which happens via omni_stage) import sys fake_tokenizer = mocker.MagicMock() fake_tokenizer.encode = mocker.MagicMock(return_value=[1, 2, 3]) fake_tokenizer.decode = mocker.MagicMock(return_value="test") # Mock init_tokenizer_from_configs (used in async_omni) def _mock_init_tokenizer_from_configs(model_config=None, **kwargs): return fake_tokenizer # Strategy 1: Mock in the original location (vllm.transformers_utils.tokenizer) # This works if the module hasn't been imported yet monkeypatch.setattr( "vllm.transformers_utils.tokenizer.init_tokenizer_from_configs", _mock_init_tokenizer_from_configs, raising=False, ) # Strategy 2: If the module is already in sys.modules, patch it directly tokenizer_module_path = "vllm.transformers_utils.tokenizer" if tokenizer_module_path in sys.modules: tokenizer_module = sys.modules[tokenizer_module_path] setattr(tokenizer_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs) # CRITICAL: Mock length_from_prompt_token_ids_or_embeds BEFORE trying to mock async_omni # This is because async_omni imports processor.py, which imports this function at module level # Mock length_from_prompt_token_ids_or_embeds (used in processor.py) def _mock_length_from_prompt_token_ids_or_embeds(prompt_token_ids=None, prompt_embeds=None): # Return a reasonable default length if prompt_token_ids is not None: if isinstance(prompt_token_ids, list): return len(prompt_token_ids) elif hasattr(prompt_token_ids, "shape"): return prompt_token_ids.shape[-1] if len(prompt_token_ids.shape) > 0 else 1 if prompt_embeds is not None: if hasattr(prompt_embeds, "shape"): return prompt_embeds.shape[-2] if len(prompt_embeds.shape) > 1 else 1 return 10 # Default length # Mock in vllm.utils monkeypatch.setattr( "vllm.utils.length_from_prompt_token_ids_or_embeds", _mock_length_from_prompt_token_ids_or_embeds, raising=False, ) # Also mock in processor module if it's imported monkeypatch.setattr( "vllm_omni.engine.input_processor.length_from_prompt_token_ids_or_embeds", _mock_length_from_prompt_token_ids_or_embeds, raising=False, ) # If processor module is already imported, patch it directly processor_module_path = "vllm_omni.engine.input_processor" if processor_module_path in sys.modules: processor_module = sys.modules[processor_module_path] setattr( processor_module, "length_from_prompt_token_ids_or_embeds", _mock_length_from_prompt_token_ids_or_embeds ) # Strategy 3: Now mock async_omni AFTER length_from_prompt_token_ids_or_embeds is mocked # This prevents ImportError when async_omni imports processor.py monkeypatch.setattr( "vllm_omni.entrypoints.async_omni.init_tokenizer_from_configs", _mock_init_tokenizer_from_configs, raising=False, ) # Strategy 4: If async_omni is already imported, patch it directly async_omni_path = "vllm_omni.entrypoints.async_omni" if async_omni_path in sys.modules: async_omni_module = sys.modules[async_omni_path] setattr(async_omni_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs) # Now mock get_config and other functions fake_hf_config = mocker.MagicMock() fake_hf_config.model_type = "qwen2_5_omni" def _mock_get_config(model, **kwargs): return fake_hf_config monkeypatch.setattr( "vllm.transformers_utils.config.get_config", _mock_get_config, raising=False, ) monkeypatch.setattr( "vllm_omni.entrypoints.utils.get_config", _mock_get_config, raising=False, ) # Mock transformers' cached_file to avoid downloading model configs def _mock_cached_file(path_or_repo_id, *args, **kwargs): import os import tempfile fake_config_file = os.path.join(tempfile.gettempdir(), "fake_config.json") if not os.path.exists(fake_config_file): with open(fake_config_file, "w") as f: f.write('{"model_type": "qwen2_5_omni"}') return fake_config_file monkeypatch.setattr( "transformers.utils.hub.cached_file", _mock_cached_file, raising=False, ) monkeypatch.setattr( "transformers.utils.hub.cached_files", lambda path_or_repo_id, filenames, **kwargs: ( [_mock_cached_file(path_or_repo_id, filenames[0])] if filenames else None ), raising=False, ) def test_initialize_stage_configs_called_when_none( monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config ): """Test that stage configs are auto-loaded when stage_configs_path is None.""" def _fake_loader( model: str, stage_configs_path: str | None = None, base_engine_args: dict | None = None, default_stage_cfg_factory=None, ): cfg0 = dict(fake_stage_config) cfg0["stage_id"] = 0 cfg1 = dict(fake_stage_config) cfg1["stage_id"] = 1 return None, [ _FakeStageConfig(cfg0), _FakeStageConfig(cfg1), ] # Remove modules from cache BEFORE setting mocks import sys for module_name in [ "vllm_omni.entrypoints.utils", "vllm_omni.entrypoints.omni", "vllm_omni.entrypoints.omni_stage", ]: if module_name in sys.modules: del sys.modules[module_name] # Set up mocks _setup_engine_mocks(monkeypatch, mocker) _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) _setup_connector_mocks(monkeypatch, mocker) # Mock load_and_resolve_stage_configs monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", _fake_loader, raising=False, ) # Replace OmniStage monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) # Import the module after mocks are set import vllm_omni.entrypoints.omni as omni_module # Patch the imported class in the module monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) from vllm_omni.entrypoints.omni import Omni omni = Omni(model=MODEL, init_timeout=1) # Verify: auto-loaded stage_configs and stage_list have consistent count assert isinstance(omni.stage_configs, list) assert len(omni.stage_configs) == 2 assert len(omni.stage_list) == 2 # Verify: each Stage is _FakeStage instance for st in omni.stage_list: assert isinstance(st, _FakeStage) # Verify: queues are attached for st in omni.stage_list: assert st._in_q is not None assert st._out_q is not None # Verify: all stages are ready assert len(omni._stages_ready) == 2 def test_generate_raises_on_length_mismatch(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test that generate raises ValueError when sampling_params_list length doesn't match.""" def _fake_loader( model: str, stage_configs_path: str | None = None, base_engine_args: dict | None = None, default_stage_cfg_factory=None, ): return None, [_FakeStageConfig(fake_stage_config)] import sys for module_name in [ "vllm_omni.entrypoints.utils", "vllm_omni.entrypoints.omni", "vllm_omni.entrypoints.omni_stage", ]: if module_name in sys.modules: del sys.modules[module_name] _setup_engine_mocks(monkeypatch, mocker) _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", _fake_loader, raising=False, ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) from vllm_omni.entrypoints.omni import Omni omni = Omni(model=MODEL, init_timeout=1) with pytest.raises(ValueError): omni.generate(prompts=["hi"], sampling_params_list=[]) def test_generate_pipeline_and_final_outputs(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test multi-stage generation pipeline with queue polling.""" stage_cfg0 = dict(fake_stage_config) stage_cfg0["stage_id"] = 0 stage_cfg1 = dict(fake_stage_config) stage_cfg1["stage_id"] = 1 stage_cfg1["processed_input"] = ["processed-for-stage-1"] def _fake_loader( model: str, stage_configs_path: str | None = None, base_engine_args: dict | None = None, default_stage_cfg_factory=None, ): return None, [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)] import sys for module_name in [ "vllm_omni.entrypoints.utils", "vllm_omni.entrypoints.omni", "vllm_omni.entrypoints.omni_stage", "vllm_omni.distributed.omni_connectors.adapter", "vllm_omni.distributed.omni_connectors", ]: if module_name in sys.modules: del sys.modules[module_name] _setup_engine_mocks(monkeypatch, mocker) _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", _fake_loader, raising=False, ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) # Apply connector and adapter mocks after importing omni module _setup_connector_mocks(monkeypatch, mocker, omni_module) _setup_connector_adapter_mock(monkeypatch, omni_module) # Mock uuid.uuid4() to return a predictable value for request ID generation test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000") monkeypatch.setattr(uuid, "uuid4", lambda: test_uuid) monkeypatch.setattr(omni_module, "uuid", uuid) from vllm_omni.entrypoints.omni import Omni omni = Omni(model=MODEL, init_timeout=1) # Generate the expected request ID format: "0_" expected_request_id = f"0_{test_uuid}" # Simulate worker behavior: manually put results into output queues # Note: We put results before calling generate, which simulates worker processes # that have already completed. The polling loop will collect them in stage order. # Stage 0 output (will be collected first) omni.stage_list[0]._out_q.put_nowait( { "request_id": expected_request_id, "engine_outputs": [{"stage": 0, "text": "s0"}], "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0), } ) # Stage 1 output (will be collected after stage 0 forwards to it) # Note: In real flow, stage 1 result would appear after stage 0 forwards, # but for testing we pre-populate it. The polling loop processes stages # in order, so stage 0 result will be collected first, then forwarded, # then stage 1 result will be collected. omni.stage_list[1]._out_q.put_nowait( { "request_id": expected_request_id, "engine_outputs": [{"stage": 1, "text": "s1"}], "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0), } ) sampling_params_list = [ OmniDiffusionSamplingParams(num_inference_steps=1), OmniDiffusionSamplingParams(num_inference_steps=1, max_sequence_length=10), ] prompts = ["hi"] outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list) # Both stages have final_output=True, so should aggregate two OmniRequestOutput assert len(outputs) == 2 # Verify stage outputs are set assert omni.stage_list[0].engine_outputs == [{"stage": 0, "text": "s0"}] assert omni.stage_list[1].engine_outputs == [{"stage": 1, "text": "s1"}] # Verify stage 0 input queue received the task assert not omni.stage_list[0]._in_q.empty() # Verify stage 1 received forwarded task (process_engine_inputs was called) assert omni.stage_list[1].process_engine_inputs([], []) is not None def test_generate_pipeline_with_batch_input(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test single-stage generation pipeline with multiple inputs in one batch.""" stage_cfg0 = dict(fake_stage_config) stage_cfg0["stage_id"] = 0 stage_cfg0["final_output"] = False stage_cfg1 = dict(fake_stage_config) stage_cfg1["stage_id"] = 1 def _fake_loader( model: str, stage_configs_path: str | None = None, base_engine_args: dict | None = None, default_stage_cfg_factory=None, ): return None, [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)] import sys for module_name in [ "vllm_omni.entrypoints.utils", "vllm_omni.entrypoints.omni", "vllm_omni.entrypoints.omni_stage", "vllm_omni.distributed.omni_connectors.adapter", "vllm_omni.distributed.omni_connectors", ]: if module_name in sys.modules: del sys.modules[module_name] _setup_engine_mocks(monkeypatch, mocker) _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", _fake_loader, raising=False, ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) # Apply connector and adapter mocks after importing omni module _setup_connector_mocks(monkeypatch, mocker, omni_module) _setup_connector_adapter_mock(monkeypatch, omni_module) # Mock uuid.uuid4() to return a predictable value for request ID generation test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000") monkeypatch.setattr(uuid, "uuid4", lambda: test_uuid) monkeypatch.setattr(omni_module, "uuid", uuid) from vllm_omni.entrypoints.omni import Omni omni = Omni(model=MODEL, init_timeout=1) # Generate the expected request ID format: "0_" expected_request_id = f"0_{test_uuid}" # Simulate worker behavior: manually put results into output queues # Note: We put results before calling generate, which simulates worker processes # that have already completed. The polling loop will collect them in stage order. omni.stage_list[0]._out_q.put_nowait( { "request_id": expected_request_id, "engine_outputs": [{"stage": 0, "text": "s0"}], "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0), } ) omni.stage_list[0]._out_q.put_nowait( { "request_id": expected_request_id, "engine_outputs": [{"stage": 0, "text": "s0"}], "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0), } ) omni.stage_list[1]._out_q.put_nowait( { "request_id": expected_request_id, "engine_outputs": [{"stage": 1}], "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0), } ) omni.stage_list[1]._out_q.put_nowait( { "request_id": expected_request_id, "engine_outputs": [{"stage": 1}], "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0), } ) outputs = omni.generate( prompts=[ { "prompt": "hi", "negative_prompt": "hi", "multi_modal_data": {"image": ["dog.jpg", "cat.jpg"]}, }, { "prompt": "hi", "negative_prompt": "hi", "multi_modal_data": {"image": ["dog.jpg", "cat.jpg"]}, }, ], sampling_params_list=[ OmniDiffusionSamplingParams(num_inference_steps=1), OmniDiffusionSamplingParams(num_inference_steps=1), ], ) assert len(outputs) == 2 def test_generate_no_final_output_returns_empty( monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config ): """Test that generate returns empty list when all stages have final_output=False.""" stage_cfg0 = dict(fake_stage_config) stage_cfg0["stage_id"] = 0 stage_cfg0["final_output"] = False stage_cfg1 = dict(fake_stage_config) stage_cfg1["stage_id"] = 1 stage_cfg1["final_output"] = False def _fake_loader( model: str, stage_configs_path: str | None = None, base_engine_args: dict | None = None, default_stage_cfg_factory=None, ): return None, [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)] import sys for module_name in [ "vllm_omni.entrypoints.utils", "vllm_omni.entrypoints.omni", "vllm_omni.entrypoints.omni_stage", "vllm_omni.distributed.omni_connectors.adapter", "vllm_omni.distributed.omni_connectors", ]: if module_name in sys.modules: del sys.modules[module_name] _setup_engine_mocks(monkeypatch, mocker) _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", _fake_loader, raising=False, ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) # Apply connector and adapter mocks after importing omni module _setup_connector_mocks(monkeypatch, mocker, omni_module) _setup_connector_adapter_mock(monkeypatch, omni_module) # Mock uuid.uuid4() to return a predictable value for request ID generation test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000") monkeypatch.setattr(uuid, "uuid4", lambda: test_uuid) monkeypatch.setattr(omni_module, "uuid", uuid) from vllm_omni.entrypoints.omni import Omni omni = Omni(model=MODEL, init_timeout=1) # Generate the expected request ID format: "0_" expected_request_id = f"0_{test_uuid}" # Simulate worker behavior: put results into output queues omni.stage_list[0]._out_q.put_nowait( { "request_id": expected_request_id, "engine_outputs": [{"stage": 0}], "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0}, } ) omni.stage_list[1]._out_q.put_nowait( { "request_id": expected_request_id, "engine_outputs": [{"stage": 1}], "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0}, } ) outputs = omni.generate( prompts=["p"], sampling_params_list=[ OmniDiffusionSamplingParams(num_inference_steps=1), OmniDiffusionSamplingParams(num_inference_steps=1, max_sequence_length=10), ], ) assert outputs == [] def test_generate_sampling_params_none_use_default( monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config ): """Test that generate uses default sampling params when sampling_params_list is None.""" stage_cfg0 = dict(fake_stage_config) stage_cfg0["stage_id"] = 0 stage_cfg0["final_output"] = False stage_cfg1 = dict(fake_stage_config) stage_cfg1["stage_id"] = 1 stage_cfg1["final_output"] = False def _fake_loader( model: str, stage_configs_path: str | None = None, base_engine_args: dict | None = None, default_stage_cfg_factory=None, ): return None, [_FakeStageConfig(stage_cfg0), _FakeStageConfig(stage_cfg1)] import sys for module_name in [ "vllm_omni.entrypoints.utils", "vllm_omni.entrypoints.omni", "vllm_omni.entrypoints.omni_stage", "vllm_omni.distributed.omni_connectors.adapter", "vllm_omni.distributed.omni_connectors", ]: if module_name in sys.modules: del sys.modules[module_name] _setup_engine_mocks(monkeypatch, mocker) _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", _fake_loader, raising=False, ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) # Apply connector and adapter mocks after importing omni module _setup_connector_mocks(monkeypatch, mocker, omni_module) _setup_connector_adapter_mock(monkeypatch, omni_module) # Mock uuid.uuid4() to return a predictable value for request ID generation test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000") monkeypatch.setattr(uuid, "uuid4", lambda: test_uuid) monkeypatch.setattr(omni_module, "uuid", uuid) from vllm_omni.entrypoints.omni import Omni omni = Omni(model=MODEL, init_timeout=1) # Generate the expected request ID format: "0_" expected_request_id = f"0_{test_uuid}" # Simulate worker behavior: put results into output queues omni.stage_list[0]._out_q.put_nowait( { "request_id": expected_request_id, "engine_outputs": [{"stage": 0}], "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0), } ) omni.stage_list[1]._out_q.put_nowait( { "request_id": expected_request_id, "engine_outputs": [{"stage": 1}], "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0), } ) # Use the default sampling params omni.generate(prompts=["p"], sampling_params_list=None) def test_wait_for_stages_ready_timeout(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test that _wait_for_stages_ready handles timeout correctly.""" def _fake_loader( model: str, stage_configs_path: str | None = None, base_engine_args: dict | None = None, default_stage_cfg_factory=None, ): return None, [_FakeStageConfig(fake_stage_config)] import sys for module_name in [ "vllm_omni.entrypoints.utils", "vllm_omni.entrypoints.omni", "vllm_omni.entrypoints.omni_stage", ]: if module_name in sys.modules: del sys.modules[module_name] _setup_engine_mocks(monkeypatch, mocker) _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", _fake_loader, raising=False, ) # Create a stage that doesn't send stage_ready message class _FakeStageNoReady(_FakeStage): def init_stage_worker(self, *args, **kwargs): # Don't send stage_ready message self._proc = self._mocker.MagicMock() self._proc.start = self._mocker.MagicMock() self._proc.join = self._mocker.MagicMock() self._proc.is_alive = self._mocker.MagicMock(return_value=False) self._proc.terminate = self._mocker.MagicMock() monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", lambda cfg, **kwargs: _FakeStageNoReady(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStageNoReady(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) from vllm_omni.entrypoints.omni import Omni # Use very short timeout omni = Omni(model=MODEL, init_timeout=0.01) # Verify that no stages are ready assert len(omni._stages_ready) == 0 def test_generate_handles_error_messages(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test that generate handles error messages from stages correctly.""" def _fake_loader( model: str, stage_configs_path: str | None = None, base_engine_args: dict | None = None, default_stage_cfg_factory=None, ): return None, [_FakeStageConfig(fake_stage_config)] import sys for module_name in [ "vllm_omni.entrypoints.utils", "vllm_omni.entrypoints.omni", "vllm_omni.entrypoints.omni_stage", ]: if module_name in sys.modules: del sys.modules[module_name] _setup_engine_mocks(monkeypatch, mocker) _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", _fake_loader, raising=False, ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) # Mock uuid.uuid4() to return a predictable value for request ID generation test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000") monkeypatch.setattr(uuid, "uuid4", lambda: test_uuid) monkeypatch.setattr(omni_module, "uuid", uuid) from vllm_omni.entrypoints.omni import Omni omni = Omni(model=MODEL, init_timeout=1) # Generate the expected request ID format: "0_" expected_request_id = f"0_{test_uuid}" # Put error message in output queue omni.stage_list[0]._out_q.put_nowait( { "request_id": expected_request_id, "error": "test error", } ) # Also put a valid result after error to allow the loop to complete # (error handling continues the loop, so we need a valid result to finish) omni.stage_list[0]._out_q.put_nowait( { "request_id": expected_request_id, "engine_outputs": [{"stage": 0, "text": "result"}], "metrics": _FakeStageRequestStats(num_tokens_out=1, stage_gen_time_ms=10.0), } ) # Generate should handle error gracefully (log but continue) sampling_params_list = [OmniDiffusionSamplingParams(num_inference_steps=1)] outputs = omni.generate(prompts=["hi"], sampling_params_list=sampling_params_list) # Should return final output (error was logged but didn't stop processing) assert isinstance(outputs, list) # Since final_output=True, should have one output assert len(outputs) == 1 def test_close_sends_shutdown_signal(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture, fake_stage_config): """Test that close() sends shutdown signal to all input queues.""" def _fake_loader( model: str, stage_configs_path: str | None = None, base_engine_args: dict | None = None, default_stage_cfg_factory=None, ): return None, [_FakeStageConfig(fake_stage_config)] import sys for module_name in [ "vllm_omni.entrypoints.utils", "vllm_omni.entrypoints.omni", "vllm_omni.entrypoints.omni_stage", ]: if module_name in sys.modules: del sys.modules[module_name] _setup_engine_mocks(monkeypatch, mocker) _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) _setup_connector_mocks(monkeypatch, mocker) monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_and_resolve_stage_configs", _fake_loader, raising=False, ) monkeypatch.setattr( "vllm_omni.entrypoints.omni_stage.OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs), raising=False, ) import vllm_omni.entrypoints.omni as omni_module monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader) monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs)) from vllm_omni.entrypoints.omni import Omni omni = Omni(model=MODEL, init_timeout=1) # Call close omni.close() # Verify shutdown signal (None) was sent to input queue # Use get_nowait to avoid blocking (close() uses put_nowait, so should be safe) try: shutdown_signal = omni.stage_list[0]._in_q.get_nowait() assert shutdown_signal == SHUTDOWN_TASK except Empty: # If queue was already empty or only had stage_ready, that's also acceptable # The important thing is that close() was called without error pass # Verify stop_stage_worker was called (process should be set) assert omni.stage_list[0]._proc is not None # Tests below are for diffusion dtype normalization fix from the following: # https://github.com/vllm-project/vllm-omni/pull/1391 # In the future we should ensure dtypes are parsed in less hacky way and make # these tests more atomic. @pytest.mark.parametrize("dtype", ["float16", torch.float16]) def test_dtype_normalization_valid_types( monkeypatch, dtype: str | torch.dtype, mocker: MockerFixture, fake_stage_config ): """Ensure Diffusion Config builder coerces valid types correctly.""" def _fake_loader(model: str, base_engine_args=None): # Return not stage configs to fall back to the diffusion cfg builder return [] import sys for module_name in [ "vllm_omni.entrypoints.utils", "vllm_omni.entrypoints.omni", "vllm_omni.entrypoints.omni_stage", ]: if module_name in sys.modules: del sys.modules[module_name] _setup_engine_mocks(monkeypatch, mocker) _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) from vllm_omni.entrypoints.omni import Omni monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_stage_configs_from_model", _fake_loader, raising=False, ) omni = Omni(model="any", init_timeout=1, dtype=dtype) # Dtype parsing being checked is on the diffusion path assert len(omni.stage_configs) == 1 assert omni.stage_configs[0]["stage_type"] == "diffusion" # Regardless of whether a str / dtype is passed, it resolves correctly engine_args = omni.stage_configs[0]["engine_args"] assert "dtype" in engine_args assert isinstance(engine_args["dtype"], str) assert engine_args["dtype"] == "float16" def test_dtype_normalization_invalid_types(monkeypatch, mocker: MockerFixture, fake_stage_config): """Ensure Diffusion Config builder correctly handles bad dtype overrides.""" def _fake_loader(model: str, base_engine_args=None): # Return not stage configs to fall back to the diffusion cfg builder return [] class NotATorchDtype: pass import sys for module_name in [ "vllm_omni.entrypoints.utils", "vllm_omni.entrypoints.omni", "vllm_omni.entrypoints.omni_stage", ]: if module_name in sys.modules: del sys.modules[module_name] _setup_engine_mocks(monkeypatch, mocker) _setup_multiprocessing_mocks(monkeypatch, mocker) _setup_ipc_mocks(monkeypatch) _setup_log_mocks(monkeypatch) from vllm_omni.entrypoints.omni import Omni monkeypatch.setattr( "vllm_omni.entrypoints.utils.load_stage_configs_from_model", _fake_loader, raising=False, ) # Raise TypeError if we get an unrecognized type with pytest.raises(TypeError): Omni(model="any", init_timeout=1, dtype=NotATorchDtype)