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

Merge origin/main into dev/vllm-align

Browse Source 
Resolved conflicts:
- vllm_omni/inputs/preprocess.py: took origin/main's elif mm_processor_kwargs branch, fixed undefined mm_uuids reference
- vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py: took origin/main's NPU support

Made-with: Cursor

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
This commit is contained in:
tzhouam committed 2026-03-02 02:10:32 +00:00
2 parents 8d0c46487e fec0182ac2
commit 26744bcb35
240 files changed
+25869 -3091
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

tests/entrypoints/openai_api/test_image_server.py
+5 -5
View File
@@ -10,11 +10,11 @@ OpenAI-compatible async text-to-image generation API endpoints in api_server.py.
import base64
import io
from argparse import Namespace
from unittest.mock import AsyncMock, Mock
import pytest
from fastapi.testclient import TestClient
from PIL import Image
from pytest_mock import MockerFixture
from vllm import SamplingParams
from vllm_omni.entrypoints.openai.image_api_utils import (
@@ -128,11 +128,11 @@ class FakeAsyncOmni:
@pytest.fixture
def mock_async_diffusion():
def mock_async_diffusion(mocker: MockerFixture):
"""Mock AsyncOmniDiffusion instance that returns fake images"""
mock = Mock()
mock = mocker.Mock()
mock.is_running = True # For health endpoint
mock.check_health = AsyncMock() # For LLM mode health check
mock.check_health = mocker.AsyncMock() # For LLM mode health check
async def generate(**kwargs):
# Return n PIL images wrapped in result object
@@ -143,7 +143,7 @@ def mock_async_diffusion():
images = [Image.new("RGB", (64, 64), color="blue") for _ in range(n)]
return MockGenerationResult(images)
mock.generate = AsyncMock(side_effect=generate)
mock.generate = mocker.AsyncMock(side_effect=generate)
return mock
tests/entrypoints/openai_api/test_serving_chat_sampling_params.py
+23 -18
View File
@@ -6,27 +6,26 @@ Tests that standard OpenAI API parameters (max_tokens, temperature, etc.)
are correctly applied to the comprehension stage while preserving YAML defaults.
"""
from unittest.mock import MagicMock
import pytest
from pytest_mock import MockerFixture
from vllm.sampling_params import SamplingParams
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.fixture
def mock_comprehension_stage():
def mock_comprehension_stage(mocker: MockerFixture):
"""Create a mock comprehension stage with is_comprehension=True."""
stage = MagicMock()
stage = mocker.MagicMock()
stage.is_comprehension = True
stage.model_stage = "comprehension"
return stage
@pytest.fixture
def mock_other_stage():
def mock_other_stage(mocker: MockerFixture):
"""Create a mock non-comprehension stage."""
stage = MagicMock()
stage = mocker.MagicMock()
stage.is_comprehension = False
stage.model_stage = "other"
return stage
@@ -57,9 +56,15 @@ def default_other_params():
@pytest.fixture
def mock_engine_client(mock_comprehension_stage, mock_other_stage, default_comprehension_params, default_other_params):
def mock_engine_client(
mock_comprehension_stage,
mock_other_stage,
default_comprehension_params,
default_other_params,
mocker: MockerFixture,
):
"""Create mock engine client with stage_list and default_sampling_params_list."""
engine_client = MagicMock()
engine_client = mocker.MagicMock()
engine_client.stage_list = [mock_comprehension_stage, mock_other_stage]
engine_client.default_sampling_params_list = [
default_comprehension_params,
@@ -80,9 +85,9 @@ def serving_chat(mock_engine_client):
@pytest.fixture
def mock_request():
def mock_request(mocker: MockerFixture):
"""Create a mock request with all OpenAI sampling params set to None."""
request = MagicMock()
request = mocker.MagicMock()
# OpenAI standard sampling fields
request.temperature = None
request.top_p = None
@@ -294,35 +299,35 @@ def test_get_comprehension_stage_index_finds_first_stage(mock_engine_client):
assert instance._get_comprehension_stage_index() == 0
def test_get_comprehension_stage_index_finds_second_stage():
def test_get_comprehension_stage_index_finds_second_stage(mocker: MockerFixture):
"""Test finding comprehension stage when it's at index 1."""
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
instance = object.__new__(OmniOpenAIServingChat)
other = MagicMock()
other = mocker.MagicMock()
other.is_comprehension = False
comprehension = MagicMock()
comprehension = mocker.MagicMock()
comprehension.is_comprehension = True
instance.engine_client = MagicMock()
instance.engine_client = mocker.MagicMock()
instance.engine_client.stage_list = [other, comprehension]
assert instance._get_comprehension_stage_index() == 1
def test_get_comprehension_stage_index_raises_when_not_found():
def test_get_comprehension_stage_index_raises_when_not_found(mocker: MockerFixture):
"""Test that ValueError is raised when no comprehension stage exists."""
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
instance = object.__new__(OmniOpenAIServingChat)
stage1 = MagicMock()
stage1 = mocker.MagicMock()
stage1.is_comprehension = False
stage2 = MagicMock()
stage2 = mocker.MagicMock()
stage2.is_comprehension = False
instance.engine_client = MagicMock()
instance.engine_client = mocker.MagicMock()
instance.engine_client.stage_list = [stage1, stage2]
with pytest.raises(ValueError, match="No comprehension stage"):
tests/entrypoints/openai_api/test_serving_speech.py
+193 -84
View File
@@ -1,13 +1,14 @@
# tests/entrypoints/openai/test_serving_speech.py
import logging
from inspect import Signature, signature
from unittest.mock import MagicMock, patch
import numpy as np
import pytest
import torch
from fastapi import FastAPI
from fastapi.testclient import TestClient
from pydantic import ValidationError
from pytest_mock import MockerFixture
from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin
from vllm_omni.entrypoints.openai.protocol.audio import CreateAudio, OpenAICreateSpeechRequest
@@ -26,25 +27,24 @@ class TestAudioMixin:
def audio_mixin(self):
return AudioMixin()
def test_stereo_to_mono_conversion(self, audio_mixin):
def test_stereo_to_mono_conversion(self, audio_mixin, mocker: MockerFixture):
stereo_tensor = np.random.rand(24000, 2).astype(np.float32)
audio_obj = CreateAudio(audio_tensor=stereo_tensor)
with (
patch.object(
audio_mixin, "_apply_speed_adjustment", side_effect=lambda tensor, speed, sr: (tensor, sr)
) as mock_speed,
patch("soundfile.write") as _,
):
audio_mixin.create_audio(audio_obj)
mock_speed = mocker.patch.object(
audio_mixin, "_apply_speed_adjustment", side_effect=lambda tensor, speed, sr: (tensor, sr)
)
mocker.patch("soundfile.write")
# Check that the tensor passed to speed adjustment is mono
mock_speed.assert_called_once()
adjusted_tensor = mock_speed.call_args[0][0]
assert len(adjusted_tensor) == 24000
audio_mixin.create_audio(audio_obj)
@patch("librosa.effects.time_stretch")
def test_speed_adjustment(self, mock_time_stretch, audio_mixin):
# Check that the tensor passed to speed adjustment is mono
mock_speed.assert_called_once()
adjusted_tensor = mock_speed.call_args[0][0]
assert len(adjusted_tensor) == 24000
def test_speed_adjustment(self, audio_mixin, mocker: MockerFixture):
mock_time_stretch = mocker.patch("librosa.effects.time_stretch")
mock_time_stretch.return_value = np.zeros(12000)
audio_tensor = np.random.rand(24000).astype(np.float32)
@@ -53,8 +53,8 @@ class TestAudioMixin:
mock_time_stretch.assert_called_with(y=audio_tensor, rate=2.0)
assert adjusted_audio.shape == (12000,)
@patch("soundfile.write")
def test_unsupported_format_fallback(self, mock_write, audio_mixin, caplog):
def test_unsupported_format_fallback(self, audio_mixin, caplog, mocker: MockerFixture):
mock_write = mocker.patch("soundfile.write")
audio_tensor = np.random.rand(24000).astype(np.float32)
# Use a format that is not in the list of supported formats
audio_obj = CreateAudio(audio_tensor=audio_tensor, response_format="vorbis")
@@ -66,50 +66,50 @@ class TestAudioMixin:
write_kwargs = mock_write.call_args.kwargs
assert write_kwargs["format"] == "WAV"
def test_mono_audio_preservation(self, audio_mixin):
def test_mono_audio_preservation(self, audio_mixin, mocker: MockerFixture):
"""Test that mono (1D) audio tensors are processed correctly and passed to writer."""
mono_tensor = np.random.rand(24000).astype(np.float32)
audio_obj = CreateAudio(audio_tensor=mono_tensor)
with patch("soundfile.write") as mock_write:
audio_mixin.create_audio(audio_obj)
mock_write = mocker.patch("soundfile.write")
audio_mixin.create_audio(audio_obj)
mock_write.assert_called_once()
# Verify the tensor passed to soundfile.write is the exact 1D tensor
output_tensor = mock_write.call_args[0][1]
assert output_tensor.ndim == 1
assert output_tensor.shape == (24000,)
assert np.array_equal(output_tensor, mono_tensor)
mock_write.assert_called_once()
# Verify the tensor passed to soundfile.write is the exact 1D tensor
output_tensor = mock_write.call_args[0][1]
assert output_tensor.ndim == 1
assert output_tensor.shape == (24000,)
assert np.array_equal(output_tensor, mono_tensor)
def test_stereo_audio_preservation(self, audio_mixin):
def test_stereo_audio_preservation(self, audio_mixin, mocker: MockerFixture):
"""Test that stereo (2D) audio tensors are processed correctly and preserved."""
stereo_tensor = np.random.rand(24000, 2).astype(np.float32)
audio_obj = CreateAudio(audio_tensor=stereo_tensor)
with patch("soundfile.write") as mock_write:
audio_mixin.create_audio(audio_obj)
mock_write = mocker.patch("soundfile.write")
audio_mixin.create_audio(audio_obj)
mock_write.assert_called_once()
# Verify the tensor passed to soundfile.write is the exact 2D tensor
output_tensor = mock_write.call_args[0][1]
assert output_tensor.ndim == 2
assert output_tensor.shape == (24000, 2)
assert np.array_equal(output_tensor, stereo_tensor)
mock_write.assert_called_once()
# Verify the tensor passed to soundfile.write is the exact 2D tensor
output_tensor = mock_write.call_args[0][1]
assert output_tensor.ndim == 2
assert output_tensor.shape == (24000, 2)
assert np.array_equal(output_tensor, stereo_tensor)
def test_speed_adjustment_bypass(self, audio_mixin):
def test_speed_adjustment_bypass(self, audio_mixin, mocker: MockerFixture):
"""Test that speed=1.0 bypasses the expensive librosa time stretching."""
audio_tensor = np.random.rand(24000).astype(np.float32)
with patch("librosa.effects.time_stretch") as mock_time_stretch:
# speed=1.0 should return immediately without calling librosa
result, _ = audio_mixin._apply_speed_adjustment(audio_tensor, speed=1.0, sample_rate=24000)
mock_time_stretch = mocker.patch("librosa.effects.time_stretch")
# speed=1.0 should return immediately without calling librosa
result, _ = audio_mixin._apply_speed_adjustment(audio_tensor, speed=1.0, sample_rate=24000)
mock_time_stretch.assert_not_called()
assert np.array_equal(result, audio_tensor)
mock_time_stretch.assert_not_called()
assert np.array_equal(result, audio_tensor)
@patch("librosa.effects.time_stretch")
def test_speed_adjustment_stereo_handling(self, mock_time_stretch, audio_mixin):
def test_speed_adjustment_stereo_handling(self, audio_mixin, mocker: MockerFixture):
"""Test that speed adjustment is attempted on stereo inputs."""
mock_time_stretch = mocker.patch("librosa.effects.time_stretch")
stereo_tensor = np.random.rand(24000, 2).astype(np.float32)
# Mock return value representing a sped-up version (half length)
mock_time_stretch.return_value = np.zeros((12000, 2), dtype=np.float32)
@@ -161,22 +161,22 @@ def create_mock_audio_output_for_test(
@pytest.fixture
def test_app():
def test_app(mocker: MockerFixture):
# Mock the engine client
mock_engine_client = MagicMock()
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
async def mock_generate_fn(*args, **kwargs):
yield create_mock_audio_output_for_test(request_id=kwargs.get("request_id"))
mock_engine_client.generate = MagicMock(side_effect=mock_generate_fn)
mock_engine_client.generate = mocker.MagicMock(side_effect=mock_generate_fn)
mock_engine_client.default_sampling_params_list = [{}]
# Mock models to have an is_base_model method
mock_models = MagicMock()
mock_models = mocker.MagicMock()
mock_models.is_base_model.return_value = True
mock_request_logger = MagicMock()
mock_request_logger = mocker.MagicMock()
speech_server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
@@ -186,7 +186,7 @@ def test_app():
# Patch the signature of create_speech to remove 'raw_request' for FastAPI route introspection
original_create_speech = speech_server.create_speech
_ = MagicMock(side_effect=original_create_speech)
_ = mocker.MagicMock(side_effect=original_create_speech)
sig = signature(original_create_speech)
@@ -253,11 +253,13 @@ class TestSpeechAPI:
response = client.post("/v1/audio/speech", json=payload)
assert response.status_code == 422 # Unprocessable Entity
@patch("vllm_omni.entrypoints.openai.serving_speech.OmniOpenAIServingSpeech.create_audio")
def test_speed_parameter_is_used(self, mock_create_audio, test_app):
def test_speed_parameter_is_used(self, test_app, mocker: MockerFixture):
mock_create_audio = mocker.patch(
"vllm_omni.entrypoints.openai.serving_speech.OmniOpenAIServingSpeech.create_audio"
)
client = TestClient(test_app)
mock_audio_response = MagicMock()
mock_audio_response = mocker.MagicMock()
mock_audio_response.audio_data = b"dummy_audio"
mock_audio_response.media_type = "audio/wav"
mock_create_audio.return_value = mock_audio_response
@@ -287,17 +289,17 @@ class TestTTSMethods:
"""Unit tests for TTS validation and parameter building."""
@pytest.fixture
def speech_server(self):
mock_engine_client = MagicMock()
def speech_server(self, mocker: MockerFixture):
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
mock_engine_client.stage_list = None
mock_engine_client.tts_max_instructions_length = None
mock_models = MagicMock()
mock_models = mocker.MagicMock()
mock_models.is_base_model.return_value = True
return OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=MagicMock(),
request_logger=mocker.MagicMock(),
)
def test_is_tts_detection_no_stage(self, speech_server):
@@ -306,25 +308,25 @@ class TestTTSMethods:
assert speech_server._is_tts is False
assert speech_server._tts_stage is None
def test_is_tts_detection_with_tts_stage(self):
def test_is_tts_detection_with_tts_stage(self, mocker: MockerFixture):
"""Test TTS model detection when TTS stage exists."""
mock_engine_client = MagicMock()
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
mock_engine_client.tts_max_instructions_length = None
# Create a TTS stage
mock_stage = MagicMock()
mock_stage = mocker.MagicMock()
mock_stage.model_stage = "qwen3_tts"
mock_stage.tts_args = {}
mock_engine_client.stage_list = [mock_stage]
mock_models = MagicMock()
mock_models = mocker.MagicMock()
mock_models.is_base_model.return_value = True
server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=MagicMock(),
request_logger=mocker.MagicMock(),
)
assert server._is_tts is True
@@ -379,24 +381,24 @@ class TestTTSMethods:
assert params["language"] == ["English"]
assert params["task_type"] == ["CustomVoice"]
def test_load_supported_speakers(self):
def test_load_supported_speakers(self, mocker: MockerFixture):
"""Test _load_supported_speakers."""
mock_engine_client = MagicMock()
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
mock_engine_client.stage_list = None
# Mock talker_config with mixed-case speaker names
mock_talker_config = MagicMock()
mock_talker_config = mocker.MagicMock()
mock_talker_config.spk_id = {"Ryan": 0, "Vivian": 1, "Aiden": 2}
mock_engine_client.model_config.hf_config.talker_config = mock_talker_config
mock_models = MagicMock()
mock_models = mocker.MagicMock()
mock_models.is_base_model.return_value = True
server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=MagicMock(),
request_logger=mocker.MagicMock(),
)
# Verify speakers are normalized to lowercase
@@ -407,34 +409,34 @@ class TestTTSMethods:
# Fixture creates server with no CLI override and no TTS stage
assert speech_server._max_instructions_length == 500
def test_max_instructions_length_cli_override(self):
def test_max_instructions_length_cli_override(self, mocker: MockerFixture):
"""Test CLI override (stored in engine_client) takes highest priority."""
mock_engine_client = MagicMock()
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
mock_engine_client.stage_list = None
# CLI override is stored in engine_client
mock_engine_client.tts_max_instructions_length = 1000
mock_models = MagicMock()
mock_models = mocker.MagicMock()
mock_models.is_base_model.return_value = True
server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=MagicMock(),
request_logger=mocker.MagicMock(),
)
# Value is cached during __init__
assert server._max_instructions_length == 1000
def test_max_instructions_length_stage_config(self):
def test_max_instructions_length_stage_config(self, mocker: MockerFixture):
"""Test stage config value is used when no CLI override."""
mock_engine_client = MagicMock()
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
mock_engine_client.tts_max_instructions_length = None # No CLI override
mock_models = MagicMock()
mock_models = mocker.MagicMock()
mock_models.is_base_model.return_value = True
# Mock stage with tts_args
mock_stage = MagicMock()
mock_stage = mocker.MagicMock()
mock_stage.model_stage = "qwen3_tts"
mock_stage.tts_args = {"max_instructions_length": 750}
mock_engine_client.stage_list = [mock_stage]
@@ -442,22 +444,22 @@ class TestTTSMethods:
server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=MagicMock(),
request_logger=mocker.MagicMock(),
)
# Value is cached during __init__
assert server._max_instructions_length == 750
def test_max_instructions_length_cli_overrides_stage_config(self):
def test_max_instructions_length_cli_overrides_stage_config(self, mocker: MockerFixture):
"""Test CLI override (in engine_client) takes precedence over stage config."""
mock_engine_client = MagicMock()
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
# CLI override stored in engine_client
mock_engine_client.tts_max_instructions_length = 2000
mock_models = MagicMock()
mock_models = mocker.MagicMock()
mock_models.is_base_model.return_value = True
# Mock stage with tts_args
mock_stage = MagicMock()
mock_stage = mocker.MagicMock()
mock_stage.model_stage = "qwen3_tts"
mock_stage.tts_args = {"max_instructions_length": 750}
mock_engine_client.stage_list = [mock_stage]
@@ -465,25 +467,25 @@ class TestTTSMethods:
server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=MagicMock(),
request_logger=mocker.MagicMock(),
)
# CLI value (2000) should override stage config (750)
assert server._max_instructions_length == 2000
def test_validate_instructions_length_uses_cached_value(self):
def test_validate_instructions_length_uses_cached_value(self, mocker: MockerFixture):
"""Test instructions length validation uses cached _max_instructions_length."""
mock_engine_client = MagicMock()
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
mock_engine_client.stage_list = None
# CLI override with max length of 10 characters
mock_engine_client.tts_max_instructions_length = 10
mock_models = MagicMock()
mock_models = mocker.MagicMock()
mock_models.is_base_model.return_value = True
server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=MagicMock(),
request_logger=mocker.MagicMock(),
)
# Verify cached value
@@ -506,3 +508,110 @@ class TestTTSMethods:
error = server._validate_tts_request(req)
assert error is not None
assert "max 10 characters" in error
class TestStreamingProtocolValidation:
"""Unit tests for the stream field validators in OpenAICreateSpeechRequest."""
def test_stream_validation_errors(self):
"""stream=True requires response_format='pcm' and speed=1.0."""
with pytest.raises(ValidationError, match="response_format='pcm'"):
OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="wav")
with pytest.raises(ValidationError, match="Speed adjustment is not supported"):
OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm", speed=2.0)
def test_stream_valid(self):
"""stream=True + response_format='pcm' + speed=1.0 is accepted."""
req = OpenAICreateSpeechRequest(input="Hello", stream=True, response_format="pcm")
assert req.stream is True
def test_sse_stream_format_is_blocked(self):
"""stream_format='sse' is blocked."""
with pytest.raises(ValidationError, match="sse"):
OpenAICreateSpeechRequest(input="Hello", stream_format="sse")
class TestStreamingResponse:
"""Integration tests for the streaming audio response path."""
@pytest.fixture
def streaming_app(self, mocker: MockerFixture):
"""Test app whose mock engine yields one intermediate chunk then a final chunk."""
mock_engine_client = mocker.MagicMock()
mock_engine_client.errored = False
def _make_output(finished: bool) -> OmniRequestOutput:
chunk = torch.sin(torch.linspace(0, 440 * 2 * torch.pi, 24000))
class MockCompletionOutput:
def __init__(self, index: int = 0):
self.index = index
self.text = ""
self.token_ids = []
self.finish_reason = "stop"
self.stop_reason = None
self.logprobs = None
class MockRequestOutput:
def __init__(self, audio_tensor: torch.Tensor):
self.request_id = "speech-stream-test"
self.outputs = [MockCompletionOutput(index=0)]
self.multimodal_output = {"audio": audio_tensor}
self.finished = finished
self.prompt_token_ids = None
self.encoder_prompt_token_ids = None
self.num_cached_tokens = None
self.prompt_logprobs = None
self.kv_transfer_params = None
return OmniRequestOutput(
stage_id=0,
final_output_type="audio",
request_output=MockRequestOutput(audio_tensor=chunk),
finished=finished,
)
async def mock_generate_streaming(*args, **kwargs):
yield _make_output(finished=False)
yield _make_output(finished=True)
mock_engine_client.generate = mocker.MagicMock(side_effect=mock_generate_streaming)
mock_engine_client.default_sampling_params_list = [{}]
mock_models = mocker.MagicMock()
mock_models.is_base_model.return_value = True
speech_server = OmniOpenAIServingSpeech(
engine_client=mock_engine_client,
models=mock_models,
request_logger=mocker.MagicMock(),
)
original_create_speech = speech_server.create_speech
sig = signature(original_create_speech)
new_parameters = [p for name, p in sig.parameters.items() if name != "raw_request"]
new_sig = Signature(parameters=new_parameters, return_annotation=sig.return_annotation)
async def awaitable_create_speech(*args, **kwargs):
return await original_create_speech(*args, **kwargs)
awaitable_create_speech.__signature__ = new_sig
speech_server.create_speech = awaitable_create_speech
app = FastAPI()
app.add_api_route("/v1/audio/speech", speech_server.create_speech, methods=["POST"], response_model=None)
return app
def test_streaming(self, streaming_app):
"""stream=True must return audio/pcm with non-empty body."""
client = TestClient(streaming_app)
response = client.post("/v1/audio/speech", json={"input": "Hello", "stream": True, "response_format": "pcm"})
assert response.status_code == 200
assert "audio/pcm" in response.headers["content-type"]
assert len(response.content) > 0
def test_non_streaming_unchanged(self, streaming_app):
"""Non-streaming path must still return audio/wav."""
client = TestClient(streaming_app)
response = client.post("/v1/audio/speech", json={"input": "Hello", "response_format": "wav"})
assert response.status_code == 200
assert "audio/wav" in response.headers["content-type"]
tests/entrypoints/openai_api/test_video_server.py
+70 -68
View File
@@ -5,18 +5,20 @@ Unit tests for OpenAI-compatible video generation endpoints.
"""
import io
from unittest.mock import patch
import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient
from PIL import Image
from pydantic import ValidationError
from pytest_mock import MockerFixture
from vllm_omni.entrypoints.openai.api_server import router
from vllm_omni.entrypoints.openai.protocol.videos import VideoGenerationRequest, VideoResponseFormat
from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class MockVideoResult:
def __init__(self, videos):
@@ -55,27 +57,27 @@ def _make_test_image_bytes(size=(64, 64)) -> bytes:
return buf.getvalue()
def test_t2v_video_generation_form(test_client):
def test_t2v_video_generation_form(test_client, mocker: MockerFixture):
fps_values = []
def _fake_encode(video, fps):
fps_values.append(fps)
return "Zg=="
with patch(
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
side_effect=_fake_encode,
):
response = test_client.post(
"/v1/videos",
data={
"prompt": "A cat runs across the street.",
"size": "640x360",
"seconds": "2",
"fps": "12",
"n": "2",
},
)
)
response = test_client.post(
"/v1/videos",
data={
"prompt": "A cat runs across the street.",
"size": "640x360",
"seconds": "2",
"fps": "12",
"n": "2",
},
)
assert response.status_code == 200
data = response.json()
@@ -92,18 +94,18 @@ def test_t2v_video_generation_form(test_client):
assert fps_values == [12, 12]
def test_i2v_video_generation_form(test_client):
def test_i2v_video_generation_form(test_client, mocker: MockerFixture):
image_bytes = _make_test_image_bytes((48, 32))
with patch(
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
return_value="Zg==",
):
response = test_client.post(
"/v1/videos",
data={"prompt": "A bear playing with yarn."},
files={"input_reference": ("input.png", image_bytes, "image/png")},
)
)
response = test_client.post(
"/v1/videos",
data={"prompt": "A bear playing with yarn."},
files={"input_reference": ("input.png", image_bytes, "image/png")},
)
assert response.status_code == 200
data = response.json()
@@ -119,24 +121,24 @@ def test_i2v_video_generation_form(test_client):
assert input_image.size == (48, 32)
def test_seconds_defaults_fps_and_frames(test_client):
def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
fps_values = []
def _fake_encode(video, fps):
fps_values.append(fps)
return "Zg=="
with patch(
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
side_effect=_fake_encode,
):
response = test_client.post(
"/v1/videos",
data={
"prompt": "A bird flying.",
"seconds": "3",
},
)
)
response = test_client.post(
"/v1/videos",
data={
"prompt": "A bird flying.",
"seconds": "3",
},
)
assert response.status_code == 200
engine = test_client.app.state.openai_serving_video._engine_client
@@ -146,18 +148,18 @@ def test_seconds_defaults_fps_and_frames(test_client):
assert fps_values == [24]
def test_size_param_sets_width_height(test_client):
with patch(
def test_size_param_sets_width_height(test_client, mocker: MockerFixture):
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
return_value="Zg==",
):
response = test_client.post(
"/v1/videos",
data={
"prompt": "size test",
"size": "320x240",
},
)
)
response = test_client.post(
"/v1/videos",
data={
"prompt": "size test",
"size": "320x240",
},
)
assert response.status_code == 200
engine = test_client.app.state.openai_serving_video._engine_client
@@ -166,23 +168,23 @@ def test_size_param_sets_width_height(test_client):
assert captured.height == 240
def test_sampling_params_pass_through(test_client):
with patch(
def test_sampling_params_pass_through(test_client, mocker: MockerFixture):
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
return_value="Zg==",
):
response = test_client.post(
"/v1/videos",
data={
"prompt": "param pass",
"num_inference_steps": "30",
"guidance_scale": "6.5",
"guidance_scale_2": "8.0",
"true_cfg_scale": "4.0",
"boundary_ratio": "0.7",
"flow_shift": "0.25",
},
)
)
response = test_client.post(
"/v1/videos",
data={
"prompt": "param pass",
"num_inference_steps": "30",
"guidance_scale": "6.5",
"guidance_scale_2": "8.0",
"true_cfg_scale": "4.0",
"boundary_ratio": "0.7",
"flow_shift": "0.25",
},
)
assert response.status_code == 200
engine = test_client.app.state.openai_serving_video._engine_client
@@ -264,19 +266,19 @@ def test_invalid_n_raises_validation_error(test_client):
)
def test_negative_prompt_and_seed_pass_through(test_client):
with patch(
def test_negative_prompt_and_seed_pass_through(test_client, mocker: MockerFixture):
mocker.patch(
"vllm_omni.entrypoints.openai.serving_video.encode_video_base64",
return_value="Zg==",
):
response = test_client.post(
"/v1/videos",
data={
"prompt": "snowy mountain",
"negative_prompt": "blurry",
"seed": "123",
},
)
)
response = test_client.post(
"/v1/videos",
data={
"prompt": "snowy mountain",
"negative_prompt": "blurry",
"seed": "123",
},
)
assert response.status_code == 200
engine = test_client.app.state.openai_serving_video._engine_client
tests/entrypoints/test_omni_diffusion.py
+108 -94
View File
@@ -2,9 +2,9 @@ import uuid
import warnings
from queue import Empty, Queue
from typing import Any
from unittest.mock import MagicMock
import pytest
from pytest_mock import MockerFixture
from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -152,7 +152,7 @@ class _FakeZmqQueue:
class _FakeStage:
"""Lightweight Stage stub for multi-process pipeline version with queue support."""
def __init__(self, config, stage_init_timeout: int = 300):
def __init__(self, mocker: MockerFixture, config, stage_init_timeout: int = 300):
# Handle both dict and object configs
if isinstance(config, dict):
config = _FakeStageConfig(config)
@@ -178,6 +178,10 @@ class _FakeStage:
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."""
@@ -196,11 +200,11 @@ class _FakeStage:
):
"""Mock init_stage_worker: don't start real process, just send stage_ready message."""
# Create a mock process object
self._proc = MagicMock()
self._proc.start = MagicMock()
self._proc.join = MagicMock()
self._proc.is_alive = MagicMock(return_value=False)
self._proc.terminate = MagicMock()
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:
@@ -264,21 +268,23 @@ def fake_stage_config():
}
def _setup_engine_mocks(monkeypatch):
"""Helper function to set up common engine mocks."""
fake_engine = MagicMock()
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 = MagicMock()
fake_engine.tokenizer = mocker.MagicMock()
fake_engine.log_stats = False
fake_engine.vllm_config = MagicMock()
fake_engine.vllm_config.model_config = MagicMock()
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 = MagicMock(return_value=[])
fake_engine.model_config = MagicMock()
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 = MagicMock()
fake_registry.resolve_model_cls = MagicMock(return_value=(MagicMock(), "test_arch"))
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
@@ -327,21 +333,23 @@ def _setup_engine_mocks(monkeypatch):
)
def _setup_multiprocessing_mocks(monkeypatch):
"""Helper function to set up multiprocessing mocks."""
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 = MagicMock()
fake_process_instance = MagicMock()
fake_process_instance.start = MagicMock()
fake_process_instance.join = MagicMock()
fake_process_instance.is_alive = MagicMock(return_value=False)
fake_process_instance.terminate = MagicMock()
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=False)
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 = MagicMock()
fake_ctx = mocker.MagicMock()
fake_ctx.Queue = lambda maxsize=0: _FakeQueue(maxsize=maxsize)
fake_ctx.Process = fake_process_class
@@ -365,7 +373,7 @@ def _setup_multiprocessing_mocks(monkeypatch):
)
def _setup_ipc_mocks(monkeypatch):
def _setup_ipc_mocks(monkeypatch: pytest.MonkeyPatch):
"""Helper function to set up IPC function mocks."""
# Mock _encode: simple serialization
@@ -385,7 +393,7 @@ def _setup_ipc_mocks(monkeypatch):
monkeypatch.setattr("vllm_omni.entrypoints.omni._set", _fake_set, raising=False)
def _setup_log_mocks(monkeypatch):
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
@@ -416,7 +424,7 @@ def _setup_log_mocks(monkeypatch):
)
def _setup_connector_mocks(monkeypatch, omni_module=None):
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.
@@ -429,7 +437,7 @@ def _setup_connector_mocks(monkeypatch, omni_module=None):
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))] = MagicMock()
fake_connectors[(str(i), str(i + 1))] = mocker.MagicMock()
return None, fake_connectors
if omni_module is not None:
@@ -476,15 +484,15 @@ def _setup_connector_adapter_mock(monkeypatch, omni_module):
@pytest.fixture(autouse=True)
def mock_get_config(monkeypatch):
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 = MagicMock()
fake_tokenizer.encode = MagicMock(return_value=[1, 2, 3])
fake_tokenizer.decode = MagicMock(return_value="test")
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):
@@ -555,7 +563,7 @@ def mock_get_config(monkeypatch):
setattr(async_omni_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
# Now mock get_config and other functions
fake_hf_config = MagicMock()
fake_hf_config = mocker.MagicMock()
fake_hf_config.model_type = "qwen2_5_omni"
def _mock_get_config(model, **kwargs):
@@ -597,7 +605,9 @@ def mock_get_config(monkeypatch):
)
def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_config):
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(
@@ -627,11 +637,11 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
del sys.modules[module_name]
# Set up mocks
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch, mocker)
# Mock load_and_resolve_stage_configs
monkeypatch.setattr(
@@ -643,7 +653,7 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
# Replace OmniStage
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs),
raising=False,
)
@@ -651,7 +661,7 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
import vllm_omni.entrypoints.omni as omni_module
# Patch the imported class in the module
monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
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
@@ -672,7 +682,7 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
assert len(omni._stages_ready) == 2
def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
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(
@@ -693,11 +703,11 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch, mocker)
monkeypatch.setattr(
"vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
@@ -706,13 +716,13 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
)
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
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(cfg, **kwargs))
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
@@ -722,7 +732,7 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
omni.generate(prompts=["hi"], sampling_params_list=[])
def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
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
@@ -750,8 +760,8 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
@@ -762,16 +772,16 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
)
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
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(cfg, **kwargs))
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, 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
@@ -828,7 +838,7 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
assert omni.stage_list[1].process_engine_inputs([], []) is not None
def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config):
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
@@ -856,8 +866,8 @@ def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config):
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
@@ -868,16 +878,16 @@ def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config):
)
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
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(cfg, **kwargs))
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, 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
@@ -946,7 +956,9 @@ def test_generate_pipeline_with_batch_input(monkeypatch, fake_stage_config):
assert len(outputs) == 2
def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
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
@@ -975,8 +987,8 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
@@ -987,16 +999,16 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
)
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
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(cfg, **kwargs))
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, 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
@@ -1037,7 +1049,9 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
assert outputs == []
def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_config):
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
@@ -1066,8 +1080,8 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
@@ -1078,16 +1092,16 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
)
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
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(cfg, **kwargs))
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, 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
@@ -1121,7 +1135,7 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
omni.generate(prompts=["p"], sampling_params_list=None)
def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
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(
@@ -1142,11 +1156,11 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch, mocker)
monkeypatch.setattr(
"vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
@@ -1158,21 +1172,21 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
class _FakeStageNoReady(_FakeStage):
def init_stage_worker(self, *args, **kwargs):
# Don't send stage_ready message
self._proc = MagicMock()
self._proc.start = MagicMock()
self._proc.join = MagicMock()
self._proc.is_alive = MagicMock(return_value=False)
self._proc.terminate = MagicMock()
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(cfg, **kwargs),
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(cfg, **kwargs))
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
@@ -1183,7 +1197,7 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
assert len(omni._stages_ready) == 0
def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
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(
@@ -1204,11 +1218,11 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch, mocker)
monkeypatch.setattr(
"vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
@@ -1217,13 +1231,13 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
)
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
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(cfg, **kwargs))
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
@@ -1264,7 +1278,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
assert len(outputs) == 1
def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
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(
@@ -1285,11 +1299,11 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch, mocker)
monkeypatch.setattr(
"vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
@@ -1298,14 +1312,14 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
)
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
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(cfg, **kwargs))
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
tests/entrypoints/test_omni_llm.py
+98 -90
View File
@@ -3,9 +3,9 @@ import uuid
import warnings
from queue import Empty, Queue
from typing import Any
from unittest.mock import MagicMock
import pytest
from pytest_mock import MockerFixture
from vllm import SamplingParams
from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK
@@ -124,7 +124,7 @@ class _FakeZmqQueue:
class _FakeStage:
"""Lightweight Stage stub for multi-process pipeline version with queue support."""
def __init__(self, config, stage_init_timeout: int = 300):
def __init__(self, mocker: MockerFixture, config, stage_init_timeout: int = 300):
# Handle both dict and object configs
if isinstance(config, dict):
config = _FakeStageConfig(config)
@@ -150,6 +150,8 @@ class _FakeStage:
self._out_q = None
self._proc = None # Mock process reference
self._stage_init_timeout = max(0, int(stage_init_timeout))
# NOTE: mocker fixture object **must be provided**
self._mocker = mocker
def attach_queues(self, in_q, out_q):
"""Attach input and output queues."""
@@ -168,11 +170,11 @@ class _FakeStage:
):
"""Mock init_stage_worker: don't start real process, just send stage_ready message."""
# Create a mock process object
self._proc = MagicMock()
self._proc.start = MagicMock()
self._proc.join = MagicMock()
self._proc.is_alive = MagicMock(return_value=False)
self._proc.terminate = MagicMock()
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:
@@ -236,21 +238,21 @@ def fake_stage_config():
}
def _setup_engine_mocks(monkeypatch):
def _setup_engine_mocks(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture):
"""Helper function to set up common engine mocks."""
fake_engine = MagicMock()
fake_engine = mocker.MagicMock()
# Add necessary attributes to fake_engine
fake_engine.tokenizer = MagicMock()
fake_engine.tokenizer = mocker.MagicMock()
fake_engine.log_stats = False
fake_engine.vllm_config = MagicMock()
fake_engine.vllm_config.model_config = MagicMock()
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 = MagicMock(return_value=[])
fake_engine.model_config = MagicMock()
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 = MagicMock()
fake_registry.resolve_model_cls = MagicMock(return_value=(MagicMock(), "test_arch"))
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
@@ -299,21 +301,21 @@ def _setup_engine_mocks(monkeypatch):
)
def _setup_multiprocessing_mocks(monkeypatch):
def _setup_multiprocessing_mocks(monkeypatch: pytest.MonkeyPatch, mocker: MockerFixture):
"""Helper function to set up multiprocessing mocks."""
import multiprocessing as mp
# Mock Process
fake_process_class = MagicMock()
fake_process_instance = MagicMock()
fake_process_instance.start = MagicMock()
fake_process_instance.join = MagicMock()
fake_process_instance.is_alive = MagicMock(return_value=False)
fake_process_instance.terminate = MagicMock()
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=False)
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 = MagicMock()
fake_ctx = mocker.MagicMock()
fake_ctx.Queue = lambda maxsize=0: _FakeQueue(maxsize=maxsize)
fake_ctx.Process = fake_process_class
@@ -337,7 +339,7 @@ def _setup_multiprocessing_mocks(monkeypatch):
)
def _setup_ipc_mocks(monkeypatch):
def _setup_ipc_mocks(monkeypatch: pytest.MonkeyPatch):
"""Helper function to set up IPC function mocks."""
# Mock _encode: simple serialization
@@ -357,7 +359,7 @@ def _setup_ipc_mocks(monkeypatch):
monkeypatch.setattr("vllm_omni.entrypoints.omni._set", _fake_set, raising=False)
def _setup_log_mocks(monkeypatch):
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
@@ -388,7 +390,7 @@ def _setup_log_mocks(monkeypatch):
)
def _setup_connector_mocks(monkeypatch, omni_module=None):
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.
@@ -401,7 +403,7 @@ def _setup_connector_mocks(monkeypatch, omni_module=None):
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))] = MagicMock()
fake_connectors[(str(i), str(i + 1))] = mocker.MagicMock()
return None, fake_connectors
if omni_module is not None:
@@ -448,15 +450,15 @@ def _setup_connector_adapter_mock(monkeypatch, omni_module):
@pytest.fixture(autouse=True)
def mock_get_config(monkeypatch):
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 = MagicMock()
fake_tokenizer.encode = MagicMock(return_value=[1, 2, 3])
fake_tokenizer.decode = MagicMock(return_value="test")
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):
@@ -527,7 +529,7 @@ def mock_get_config(monkeypatch):
setattr(async_omni_module, "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
# Now mock get_config and other functions
fake_hf_config = MagicMock()
fake_hf_config = mocker.MagicMock()
fake_hf_config.model_type = "qwen2_5_omni"
def _mock_get_config(model, **kwargs):
@@ -569,7 +571,9 @@ def mock_get_config(monkeypatch):
)
def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_config):
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(
@@ -599,11 +603,11 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
del sys.modules[module_name]
# Set up mocks
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch, mocker)
# Mock load_and_resolve_stage_configs
monkeypatch.setattr(
@@ -615,7 +619,7 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
# Replace OmniStage
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs),
raising=False,
)
@@ -624,9 +628,9 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
# Patch the imported function and class in the module
monkeypatch.setattr(omni_module, "load_and_resolve_stage_configs", _fake_loader)
monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(cfg, **kwargs))
monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs))
# Apply connector and adapter mocks after importing omni module
_setup_connector_mocks(monkeypatch, omni_module)
_setup_connector_mocks(monkeypatch, mocker, omni_module)
from vllm_omni.entrypoints.omni import Omni
@@ -646,7 +650,7 @@ def test_initialize_stage_configs_called_when_none(monkeypatch, fake_stage_confi
assert len(omni._stages_ready) == 2
def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
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(
@@ -669,11 +673,11 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch, mocker)
monkeypatch.setattr(
"vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
@@ -682,14 +686,14 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
)
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
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(cfg, **kwargs))
monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs))
from vllm_omni.entrypoints.omni import Omni
@@ -698,7 +702,7 @@ def test_generate_raises_on_length_mismatch(monkeypatch, fake_stage_config):
omni.generate(prompts=["hi"], sampling_params_list=[])
def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
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
@@ -724,11 +728,11 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch, mocker)
monkeypatch.setattr(
"vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
@@ -737,16 +741,16 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
)
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
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(cfg, **kwargs))
monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs))
# Apply connector and adapter mocks after importing omni module
_setup_connector_mocks(monkeypatch, 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
@@ -803,7 +807,9 @@ def test_generate_pipeline_and_final_outputs(monkeypatch, fake_stage_config):
assert omni.stage_list[1].process_engine_inputs([], []) is not None
def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
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
@@ -830,11 +836,11 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch, mocker)
monkeypatch.setattr(
"vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
@@ -843,16 +849,16 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
)
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
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(cfg, **kwargs))
monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs))
# Apply connector and adapter mocks after importing omni module
_setup_connector_mocks(monkeypatch, 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
@@ -893,7 +899,9 @@ def test_generate_no_final_output_returns_empty(monkeypatch, fake_stage_config):
assert outputs == []
def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_config):
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
@@ -920,11 +928,11 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch, mocker)
monkeypatch.setattr(
"vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
@@ -933,16 +941,16 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
)
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
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(cfg, **kwargs))
monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs))
# Apply connector and adapter mocks after importing omni module
_setup_connector_mocks(monkeypatch, 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
@@ -976,7 +984,7 @@ def test_generate_sampling_params_none_use_default(monkeypatch, fake_stage_confi
omni.generate(prompts=["p"], sampling_params_list=None)
def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
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(
@@ -999,11 +1007,11 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch, mocker)
monkeypatch.setattr(
"vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
@@ -1015,22 +1023,22 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
class _FakeStageNoReady(_FakeStage):
def init_stage_worker(self, *args, **kwargs):
# Don't send stage_ready message
self._proc = MagicMock()
self._proc.start = MagicMock()
self._proc.join = MagicMock()
self._proc.is_alive = MagicMock(return_value=False)
self._proc.terminate = MagicMock()
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(cfg, **kwargs),
lambda cfg, **kwargs: _FakeStageNoReady(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: _FakeStageNoReady(cfg, **kwargs))
monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStageNoReady(mocker, cfg, **kwargs))
from vllm_omni.entrypoints.omni import Omni
@@ -1040,7 +1048,7 @@ def test_wait_for_stages_ready_timeout(monkeypatch, fake_stage_config):
assert len(omni._stages_ready) == 0
def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
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(
@@ -1063,11 +1071,11 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch, mocker)
monkeypatch.setattr(
"vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
@@ -1076,14 +1084,14 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
)
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
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(cfg, **kwargs))
monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs))
# Mock uuid.uuid4() to return a predictable value for request ID generation
test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000000")
@@ -1123,7 +1131,7 @@ def test_generate_handles_error_messages(monkeypatch, fake_stage_config):
assert len(outputs) == 1
def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
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(
@@ -1146,11 +1154,11 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
if module_name in sys.modules:
del sys.modules[module_name]
_setup_engine_mocks(monkeypatch)
_setup_multiprocessing_mocks(monkeypatch)
_setup_engine_mocks(monkeypatch, mocker)
_setup_multiprocessing_mocks(monkeypatch, mocker)
_setup_ipc_mocks(monkeypatch)
_setup_log_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch)
_setup_connector_mocks(monkeypatch, mocker)
monkeypatch.setattr(
"vllm_omni.entrypoints.utils.load_and_resolve_stage_configs",
@@ -1159,14 +1167,14 @@ def test_close_sends_shutdown_signal(monkeypatch, fake_stage_config):
)
monkeypatch.setattr(
"vllm_omni.entrypoints.omni_stage.OmniStage",
lambda cfg, **kwargs: _FakeStage(cfg, **kwargs),
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(cfg, **kwargs))
monkeypatch.setattr(omni_module, "OmniStage", lambda cfg, **kwargs: _FakeStage(mocker, cfg, **kwargs))
from vllm_omni.entrypoints.omni import Omni
tests/entrypoints/test_stage_utils.py
+10 -8
View File
@@ -1,8 +1,8 @@
import os
import sys
from unittest.mock import MagicMock
import pytest
from pytest_mock import MockerFixture
from vllm_omni.entrypoints.stage_utils import set_stage_devices
@@ -45,16 +45,18 @@ def _make_dummy_torch(call_log):
return _Torch
def _make_mock_platform(device_type: str = "cuda", env_var: str = "CUDA_VISIBLE_DEVICES"):
"""Create a mock platform for testing."""
mock_platform = MagicMock()
def _make_mock_platform(mocker, device_type: str = "cuda", env_var: str = "CUDA_VISIBLE_DEVICES"):
"""Create a mock platform for testing.
mocker object has to be passed in to utilize this helper function.
"""
mock_platform = mocker.MagicMock()
mock_platform.device_type = device_type
mock_platform.device_control_env_var = env_var
return mock_platform
@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
def test_set_stage_devices_respects_logical_ids(monkeypatch):
def test_set_stage_devices_respects_logical_ids(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
# Preserve an existing logical mapping and ensure devices "0,1" map through it.
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "6,7")
call_log: list[int] = []
@@ -62,7 +64,7 @@ def test_set_stage_devices_respects_logical_ids(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", dummy_torch)
# Mock the platform at the source module where it's defined
mock_platform = _make_mock_platform(device_type="cuda", env_var="CUDA_VISIBLE_DEVICES")
mock_platform = _make_mock_platform(mocker, device_type="cuda", env_var="CUDA_VISIBLE_DEVICES")
monkeypatch.setattr(
"vllm_omni.platforms.current_omni_platform",
mock_platform,
@@ -74,7 +76,7 @@ def test_set_stage_devices_respects_logical_ids(monkeypatch):
@pytest.mark.usefixtures("clean_gpu_memory_between_tests")
def test_set_stage_devices_npu_platform(monkeypatch):
def test_set_stage_devices_npu_platform(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
"""Test that set_stage_devices works correctly for NPU platform."""
monkeypatch.setenv("ASCEND_RT_VISIBLE_DEVICES", "4,5")
call_log: list[int] = []
@@ -99,7 +101,7 @@ def test_set_stage_devices_npu_platform(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", _NpuTorch)
# Mock NPU platform at the source module where it's defined
mock_platform = _make_mock_platform(device_type="npu", env_var="ASCEND_RT_VISIBLE_DEVICES")
mock_platform = _make_mock_platform(mocker, device_type="npu", env_var="ASCEND_RT_VISIBLE_DEVICES")
monkeypatch.setattr(
"vllm_omni.platforms.current_omni_platform",
mock_platform,
tests/entrypoints/test_utils.py
+22 -22
View File
@@ -2,9 +2,9 @@
from collections import Counter
from dataclasses import dataclass
from unittest.mock import patch
import pytest
from pytest_mock import MockerFixture
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.engine.arg_utils import OmniEngineArgs
@@ -140,7 +140,7 @@ class TestFilterDictLikeObject:
assert 42 in result
assert (1, 2) in result
def test_dict_with_recursive_structure(self):
def test_dict_with_recursive_structure(self, mocker: MockerFixture):
"""Test filtering dict with recursive/complex nested structure."""
input_dict = {
"level1": {
@@ -152,15 +152,15 @@ class TestFilterDictLikeObject:
"normal": "value",
}
with patch("vllm_omni.entrypoints.utils.logger"):
result = _filter_dict_like_object(input_dict)
mocker.patch("vllm_omni.entrypoints.utils.logger")
result = _filter_dict_like_object(input_dict)
# Normal key should exist
assert "normal" in result
# Level1 should exist
assert "level1" in result
# Normal key should exist
assert "normal" in result
# Level1 should exist
assert "level1" in result
def test_integration_with_convert_dataclasses(self):
def test_integration_with_convert_dataclasses(self, mocker: MockerFixture):
"""Test that _filter_dict_like_object integrates properly with _convert_dataclasses_to_dict."""
@dataclass
@@ -174,32 +174,32 @@ class TestFilterDictLikeObject:
"normal": "value",
}
with patch("vllm_omni.entrypoints.utils.logger"):
result = _filter_dict_like_object(input_dict)
mocker.patch("vllm_omni.entrypoints.utils.logger")
result = _filter_dict_like_object(input_dict)
# Callable should be filtered
assert "func" not in result
# Config should be converted to dict
assert "config" in result
assert "normal" in result
# Callable should be filtered
assert "func" not in result
# Config should be converted to dict
assert "config" in result
assert "normal" in result
class TestConvertDataclassesToDict:
"""Test suite for _convert_dataclasses_to_dict function."""
def test_uses_filter_dict_like_object(self):
def test_uses_filter_dict_like_object(self, mocker: MockerFixture):
"""Test that _convert_dataclasses_to_dict uses _filter_dict_like_object for dicts."""
input_dict = {
"normal": "value",
"callable": lambda x: x,
}
with patch("vllm_omni.entrypoints.utils.logger"):
result = _convert_dataclasses_to_dict(input_dict)
mocker.patch("vllm_omni.entrypoints.utils.logger")
result = _convert_dataclasses_to_dict(input_dict)
# Callable should be filtered out by _filter_dict_like_object
assert "normal" in result
assert "callable" not in result
# Callable should be filtered out by _filter_dict_like_object
assert "normal" in result
assert "callable" not in result
class TestFilterDataclassKwargs: