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

[Misc] Migrate L1 tests to use pytest-mock (#1315)

Browse Source 
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
Signed-off-by: yuanheng <jonathan.zhaoyh@gmail.com>
This commit is contained in:
Yuanheng Zhao authored and GitHub committed 2026-02-26 17:08:39 +08:00
1 parent 9913c00144
commit ed08a216c1
20 files changed
+977 -954
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

pyproject.toml
+1
View File
@@ -41,6 +41,7 @@ dev = [
"pytest>=7.0.0",
"pytest-asyncio>=0.21.0",
"pytest-cov>=4.0.0",
"pytest-mock>=3.10.0",
"mypy==1.11.1",
"pre-commit==4.0.1",
"openai-whisper>=20250625",
tests/benchmarks/patch/test_patch.py
+34 -34
View File
@@ -7,9 +7,9 @@ Unit tests for patch.py
import asyncio
import json
from unittest.mock import AsyncMock, MagicMock
import pytest
from pytest_mock import MockerFixture
from vllm.benchmarks.lib.endpoint_request_func import RequestFuncInput
from vllm_omni.benchmarks.patch.patch import MixRequestFuncOutput, async_request_openai_chat_omni_completions
@@ -51,7 +51,7 @@ def create_sse_chunk(data_dict):
@pytest.mark.asyncio
async def test_output_tokens_assigned_with_metrics():
async def test_output_tokens_assigned_with_metrics(mocker: MockerFixture):
"""Test that output.output_tokens is assigned when metrics are present"""
# Arrange
request_input = RequestFuncInput(
@@ -82,8 +82,8 @@ async def test_output_tokens_assigned_with_metrics():
]
mock_response = MockResponse(200, chunks)
mock_session = AsyncMock()
mock_session.post = MagicMock(return_value=mock_response)
mock_session = mocker.AsyncMock()
mock_session.post = mocker.MagicMock(return_value=mock_response)
# Act
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -95,7 +95,7 @@ async def test_output_tokens_assigned_with_metrics():
@pytest.mark.asyncio
async def test_output_tokens_not_assigned_without_metrics():
async def test_output_tokens_not_assigned_without_metrics(mocker: MockerFixture):
"""Test that output.output_tokens defaults to 0 when no metrics present"""
# Arrange
request_input = RequestFuncInput(
@@ -125,8 +125,8 @@ async def test_output_tokens_not_assigned_without_metrics():
]
mock_response = MockResponse(200, chunks)
mock_session = AsyncMock()
mock_session.post = MagicMock(return_value=mock_response)
mock_session = mocker.AsyncMock()
mock_session.post = mocker.MagicMock(return_value=mock_response)
# Act
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -138,7 +138,7 @@ async def test_output_tokens_not_assigned_without_metrics():
@pytest.mark.asyncio
async def test_output_tokens_assigned_multiple_metrics():
async def test_output_tokens_assigned_multiple_metrics(mocker: MockerFixture):
"""Test that output.output_tokens is updated with the latest metrics value"""
# Arrange
request_input = RequestFuncInput(
@@ -177,8 +177,8 @@ async def test_output_tokens_assigned_multiple_metrics():
]
mock_response = MockResponse(200, chunks)
mock_session = AsyncMock()
mock_session.post = MagicMock(return_value=mock_response)
mock_session = mocker.AsyncMock()
mock_session.post = mocker.MagicMock(return_value=mock_response)
# Act
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -190,7 +190,7 @@ async def test_output_tokens_assigned_multiple_metrics():
@pytest.mark.asyncio
async def test_output_tokens_with_audio_and_text():
async def test_output_tokens_with_audio_and_text(mocker: MockerFixture):
"""Test output_tokens assignment in mixed audio and text response"""
# Arrange
request_input = RequestFuncInput(
@@ -226,8 +226,8 @@ async def test_output_tokens_with_audio_and_text():
]
mock_response = MockResponse(200, chunks)
mock_session = AsyncMock()
mock_session.post = MagicMock(return_value=mock_response)
mock_session = mocker.AsyncMock()
mock_session.post = mocker.MagicMock(return_value=mock_response)
# Act
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -238,7 +238,7 @@ async def test_output_tokens_with_audio_and_text():
@pytest.mark.asyncio
async def test_output_tokens_with_missing_num_tokens_out():
async def test_output_tokens_with_missing_num_tokens_out(mocker: MockerFixture):
"""Test that output_tokens defaults to 0 when num_tokens_out is missing in metrics"""
# Arrange
request_input = RequestFuncInput(
@@ -269,8 +269,8 @@ async def test_output_tokens_with_missing_num_tokens_out():
]
mock_response = MockResponse(200, chunks)
mock_session = AsyncMock()
mock_session.post = MagicMock(return_value=mock_response)
mock_session = mocker.AsyncMock()
mock_session.post = mocker.MagicMock(return_value=mock_response)
# Act
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -312,7 +312,7 @@ class TestTextLatencyAttribute:
assert text_latency is not None or hasattr(output, "text_latency"), "text_latency attribute should exist"
@pytest.mark.asyncio
async def test_text_latency_assigned_with_text_response(self):
async def test_text_latency_assigned_with_text_response(self, mocker: MockerFixture):
"""Test that text_latency is assigned when text response is received"""
request_input = RequestFuncInput(
model="test-model",
@@ -341,8 +341,8 @@ class TestTextLatencyAttribute:
]
mock_response = MockResponse(200, chunks)
mock_session = AsyncMock()
mock_session.post = MagicMock(return_value=mock_response)
mock_session = mocker.AsyncMock()
mock_session.post = mocker.MagicMock(return_value=mock_response)
# Act
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -353,7 +353,7 @@ class TestTextLatencyAttribute:
assert output.text_latency > 0, "text_latency should be greater than 0 for text response"
@pytest.mark.asyncio
async def test_text_latency_updated_with_multiple_text_chunks(self):
async def test_text_latency_updated_with_multiple_text_chunks(self, mocker: MockerFixture):
"""Test that text_latency is updated as more text chunks arrive"""
request_input = RequestFuncInput(
model="test-model",
@@ -388,8 +388,8 @@ class TestTextLatencyAttribute:
]
mock_response = MockResponse(200, chunks)
mock_session = AsyncMock()
mock_session.post = MagicMock(return_value=mock_response)
mock_session = mocker.AsyncMock()
mock_session.post = mocker.MagicMock(return_value=mock_response)
# Act
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -401,7 +401,7 @@ class TestTextLatencyAttribute:
assert output.generated_text == "First second third"
@pytest.mark.asyncio
async def test_text_latency_with_only_audio_response(self):
async def test_text_latency_with_only_audio_response(self, mocker: MockerFixture):
"""Test text_latency behavior when only audio is received"""
request_input = RequestFuncInput(
model="test-model",
@@ -424,8 +424,8 @@ class TestTextLatencyAttribute:
]
mock_response = MockResponse(200, chunks)
mock_session = AsyncMock()
mock_session.post = MagicMock(return_value=mock_response)
mock_session = mocker.AsyncMock()
mock_session.post = mocker.MagicMock(return_value=mock_response)
# Act
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -437,7 +437,7 @@ class TestTextLatencyAttribute:
assert output.text_latency >= 0, "text_latency should be non-negative"
@pytest.mark.asyncio
async def test_text_latency_not_affected_by_metrics(self):
async def test_text_latency_not_affected_by_metrics(self, mocker: MockerFixture):
"""Test that text_latency is independent of metrics data"""
request_input = RequestFuncInput(
model="test-model",
@@ -467,8 +467,8 @@ class TestTextLatencyAttribute:
]
mock_response = MockResponse(200, chunks)
mock_session = AsyncMock()
mock_session.post = MagicMock(return_value=mock_response)
mock_session = mocker.AsyncMock()
mock_session.post = mocker.MagicMock(return_value=mock_response)
# Act
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -481,7 +481,7 @@ class TestTextLatencyAttribute:
assert output.output_tokens == 100, "metrics should not affect text_latency"
@pytest.mark.asyncio
async def test_text_latency_mixed_modalities(self):
async def test_text_latency_mixed_modalities(self, mocker: MockerFixture):
"""Test text_latency with mixed text and audio modalities"""
request_input = RequestFuncInput(
model="test-model",
@@ -516,8 +516,8 @@ class TestTextLatencyAttribute:
]
mock_response = MockResponse(200, chunks)
mock_session = AsyncMock()
mock_session.post = MagicMock(return_value=mock_response)
mock_session = mocker.AsyncMock()
mock_session.post = mocker.MagicMock(return_value=mock_response)
# Act
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
@@ -529,7 +529,7 @@ class TestTextLatencyAttribute:
assert output.generated_text == "Text more text"
@pytest.mark.asyncio
async def test_text_latency_value_consistency(self):
async def test_text_latency_value_consistency(self, mocker: MockerFixture):
"""Test that text_latency matches latency minus ttft relationship"""
request_input = RequestFuncInput(
model="test-model",
@@ -557,8 +557,8 @@ class TestTextLatencyAttribute:
]
mock_response = MockResponse(200, chunks)
mock_session = AsyncMock()
mock_session.post = MagicMock(return_value=mock_response)
mock_session = mocker.AsyncMock()
mock_session.post = mocker.MagicMock(return_value=mock_response)
# Act
output = await async_request_openai_chat_omni_completions(request_input, mock_session)
tests/diffusion/test_diffusion_worker.py
+78 -83
View File
@@ -10,10 +10,9 @@ This module tests the DiffusionWorker implementation:
- wake_up: Waking worker from sleep mode
"""
from unittest.mock import Mock, patch
import pytest
import torch
from pytest_mock import MockerFixture
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
@@ -21,9 +20,9 @@ pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@pytest.fixture
def mock_od_config():
def mock_od_config(mocker: MockerFixture):
"""Create a mock OmniDiffusionConfig."""
config = Mock()
config = mocker.Mock()
config.num_gpus = 1
config.master_port = 12345
config.enable_sleep_mode = False
@@ -34,22 +33,23 @@ def mock_od_config():
@pytest.fixture
def mock_gpu_worker(mock_od_config):
def mock_gpu_worker(mocker: MockerFixture, mock_od_config):
"""Create a DiffusionWorker with mocked initialization."""
with patch.object(DiffusionWorker, "init_device"), patch.object(DiffusionWorker, "load_model"):
worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
# Mock the model_runner with pipeline
worker.model_runner = Mock()
worker.model_runner.pipeline = Mock()
worker.device = torch.device("cuda", 0)
worker._sleep_saved_buffers = {}
return worker
mocker.patch.object(DiffusionWorker, "init_device")
mocker.patch.object(DiffusionWorker, "load_model")
worker = DiffusionWorker(local_rank=0, rank=0, od_config=mock_od_config)
# Mock the model_runner with pipeline
worker.model_runner = mocker.Mock()
worker.model_runner.pipeline = mocker.Mock()
worker.device = torch.device("cuda", 0)
worker._sleep_saved_buffers = {}
return worker
class TestDiffusionWorkerLoadWeights:
"""Test DiffusionWorker.load_weights method."""
def test_load_weights_calls_pipeline(self, mock_gpu_worker):
def test_load_weights_calls_pipeline(self, mocker: MockerFixture, mock_gpu_worker):
"""Test that load_weights delegates to model_runner.load_weights."""
# Setup mock weights
mock_weights = [
@@ -59,7 +59,7 @@ class TestDiffusionWorkerLoadWeights:
expected_loaded = {"layer1.weight", "layer2.weight"}
# Configure model_runner mock
mock_gpu_worker.model_runner.load_weights = Mock(return_value=expected_loaded)
mock_gpu_worker.model_runner.load_weights = mocker.Mock(return_value=expected_loaded)
# Call load_weights
result = mock_gpu_worker.load_weights(mock_weights)
@@ -68,9 +68,9 @@ class TestDiffusionWorkerLoadWeights:
mock_gpu_worker.model_runner.load_weights.assert_called_once_with(mock_weights)
assert result == expected_loaded
def test_load_weights_empty_iterable(self, mock_gpu_worker):
def test_load_weights_empty_iterable(self, mocker: MockerFixture, mock_gpu_worker):
"""Test load_weights with empty weights iterable."""
mock_gpu_worker.model_runner.load_weights = Mock(return_value=set())
mock_gpu_worker.model_runner.load_weights = mocker.Mock(return_value=set())
result = mock_gpu_worker.load_weights([])
@@ -81,11 +81,12 @@ class TestDiffusionWorkerLoadWeights:
class TestDiffusionWorkerSleep:
"""Test DiffusionWorker.sleep method."""
@patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
@patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
@patch("vllm.device_allocator.cumem.CuMemAllocator")
def test_sleep_level_1(self, mock_allocator_class, mock_platform, mock_get_process_memory, mock_gpu_worker):
def test_sleep_level_1(self, mocker: MockerFixture, mock_gpu_worker):
"""Test sleep mode level 1 (offload weights only)."""
mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
# Setup process-scoped memory mocks
# Before sleep: 3GB used
# After sleep: 1GB used (freed 2GB)
@@ -95,9 +96,9 @@ class TestDiffusionWorkerSleep:
]
# Setup allocator mock
mock_allocator = Mock()
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
mock_allocator.sleep = Mock()
mock_allocator = mocker.Mock()
mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
mock_allocator.sleep = mocker.Mock()
# Call sleep with level 1
result = mock_gpu_worker.sleep(level=1)
@@ -108,11 +109,12 @@ class TestDiffusionWorkerSleep:
# Verify buffers were NOT saved (level 1 doesn't save buffers)
assert len(mock_gpu_worker._sleep_saved_buffers) == 0
@patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
@patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
@patch("vllm.device_allocator.cumem.CuMemAllocator")
def test_sleep_level_2(self, mock_allocator_class, mock_platform, mock_get_process_memory, mock_gpu_worker):
def test_sleep_level_2(self, mocker: MockerFixture, mock_gpu_worker):
"""Test sleep mode level 2 (offload all, save buffers)."""
mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
# Setup process-scoped memory mocks
mock_get_process_memory.side_effect = [
5 * 1024**3, # Before sleep
@@ -120,14 +122,14 @@ class TestDiffusionWorkerSleep:
]
# Setup allocator mock
mock_allocator = Mock()
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
mock_allocator.sleep = Mock()
mock_allocator = mocker.Mock()
mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
mock_allocator.sleep = mocker.Mock()
# Mock pipeline buffers
mock_buffer1 = torch.randn(10, 10)
mock_buffer2 = torch.randn(20, 20)
mock_gpu_worker.model_runner.pipeline.named_buffers = Mock(
mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(
return_value=[
("buffer1", mock_buffer1),
("buffer2", mock_buffer2),
@@ -146,42 +148,32 @@ class TestDiffusionWorkerSleep:
assert "buffer1" in mock_gpu_worker._sleep_saved_buffers
assert "buffer2" in mock_gpu_worker._sleep_saved_buffers
@patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
@patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
@patch("vllm.device_allocator.cumem.CuMemAllocator")
def test_sleep_memory_freed_validation(
self,
mock_allocator_class,
mock_platform,
mock_get_process_memory,
mock_gpu_worker,
):
def test_sleep_memory_freed_validation(self, mocker: MockerFixture, mock_gpu_worker):
"""Test that sleep validates memory was actually freed."""
mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
# Simulate process memory increase (should trigger assertion error)
mock_get_process_memory.side_effect = [
1 * 1024**3, # Before sleep: 1GB used
3 * 1024**3, # After sleep: 3GB used (negative freed)
]
mock_allocator = Mock()
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
mock_allocator.sleep = Mock()
mock_allocator = mocker.Mock()
mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
mock_allocator.sleep = mocker.Mock()
# This should raise an assertion error
with pytest.raises(AssertionError, match="Memory usage increased after sleeping"):
mock_gpu_worker.sleep(level=1)
@patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
@patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
@patch("vllm.device_allocator.cumem.CuMemAllocator")
def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(
self,
mock_allocator_class,
mock_platform,
mock_get_process_memory,
mock_gpu_worker,
):
def test_sleep_falls_back_to_device_memory_when_nvml_unavailable(self, mocker: MockerFixture, mock_gpu_worker):
"""Test sleep uses device-scoped fallback when NVML is unavailable."""
mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
mock_platform = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.current_omni_platform")
mock_get_process_memory = mocker.patch("vllm_omni.diffusion.worker.diffusion_worker.get_process_gpu_memory")
mock_get_process_memory.side_effect = [None, None]
mock_platform.get_free_memory.side_effect = [
1 * 1024**3, # Before sleep
@@ -189,9 +181,9 @@ class TestDiffusionWorkerSleep:
]
mock_platform.get_device_total_memory.return_value = 8 * 1024**3
mock_allocator = Mock()
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
mock_allocator.sleep = Mock()
mock_allocator = mocker.Mock()
mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
mock_allocator.sleep = mocker.Mock()
result = mock_gpu_worker.sleep(level=1)
@@ -202,13 +194,14 @@ class TestDiffusionWorkerSleep:
class TestDiffusionWorkerWakeUp:
"""Test DiffusionWorker.wake_up method."""
@patch("vllm.device_allocator.cumem.CuMemAllocator")
def test_wake_up_without_buffers(self, mock_allocator_class, mock_gpu_worker):
def test_wake_up_without_buffers(self, mocker: MockerFixture, mock_gpu_worker):
"""Test wake_up without saved buffers (level 1 sleep)."""
mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
# Setup allocator mock
mock_allocator = Mock()
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
mock_allocator.wake_up = Mock()
mock_allocator = mocker.Mock()
mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
mock_allocator.wake_up = mocker.Mock()
# Ensure no saved buffers
mock_gpu_worker._sleep_saved_buffers = {}
@@ -220,13 +213,14 @@ class TestDiffusionWorkerWakeUp:
mock_allocator.wake_up.assert_called_once_with(["weights"])
assert result is True
@patch("vllm.device_allocator.cumem.CuMemAllocator")
def test_wake_up_with_buffers(self, mock_allocator_class, mock_gpu_worker):
def test_wake_up_with_buffers(self, mocker: MockerFixture, mock_gpu_worker):
"""Test wake_up with saved buffers (level 2 sleep)."""
mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
# Setup allocator mock
mock_allocator = Mock()
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
mock_allocator.wake_up = Mock()
mock_allocator = mocker.Mock()
mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
mock_allocator.wake_up = mocker.Mock()
# Create saved buffers
saved_buffer1 = torch.randn(10, 10)
@@ -237,12 +231,12 @@ class TestDiffusionWorkerWakeUp:
}
# Mock pipeline buffers (these will be restored)
mock_buffer1 = Mock()
mock_buffer1.data = Mock()
mock_buffer2 = Mock()
mock_buffer2.data = Mock()
mock_buffer1 = mocker.Mock()
mock_buffer1.data = mocker.Mock()
mock_buffer2 = mocker.Mock()
mock_buffer2.data = mocker.Mock()
mock_gpu_worker.model_runner.pipeline.named_buffers = Mock(
mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(
return_value=[
("buffer1", mock_buffer1),
("buffer2", mock_buffer2),
@@ -263,13 +257,14 @@ class TestDiffusionWorkerWakeUp:
assert len(mock_gpu_worker._sleep_saved_buffers) == 0
assert result is True
@patch("vllm.device_allocator.cumem.CuMemAllocator")
def test_wake_up_partial_buffer_restore(self, mock_allocator_class, mock_gpu_worker):
def test_wake_up_partial_buffer_restore(self, mocker: MockerFixture, mock_gpu_worker):
"""Test wake_up only restores buffers that were saved."""
mock_allocator_class = mocker.patch("vllm.device_allocator.cumem.CuMemAllocator")
# Setup allocator mock
mock_allocator = Mock()
mock_allocator_class.get_instance = Mock(return_value=mock_allocator)
mock_allocator.wake_up = Mock()
mock_allocator = mocker.Mock()
mock_allocator_class.get_instance = mocker.Mock(return_value=mock_allocator)
mock_allocator.wake_up = mocker.Mock()
# Only save buffer1, not buffer2
saved_buffer1 = torch.randn(10, 10)
@@ -278,12 +273,12 @@ class TestDiffusionWorkerWakeUp:
}
# Mock pipeline has both buffers
mock_buffer1 = Mock()
mock_buffer1.data = Mock()
mock_buffer2 = Mock()
mock_buffer2.data = Mock()
mock_buffer1 = mocker.Mock()
mock_buffer1.data = mocker.Mock()
mock_buffer2 = mocker.Mock()
mock_buffer2.data = mocker.Mock()
mock_gpu_worker.model_runner.pipeline.named_buffers = Mock(
mock_gpu_worker.model_runner.pipeline.named_buffers = mocker.Mock(
return_value=[
("buffer1", mock_buffer1),
("buffer2", mock_buffer2),
tests/diffusion/test_worker_wrapper_base.py
+170 -165
View File
@@ -13,9 +13,9 @@ This module tests the WorkerWrapperBase implementation:
"""
from typing import Any
from unittest.mock import Mock, patch
import pytest
from pytest_mock import MockerFixture
from vllm_omni.diffusion.worker.diffusion_worker import (
CustomPipelineWorkerExtension,
@@ -23,15 +23,18 @@ from vllm_omni.diffusion.worker.diffusion_worker import (
WorkerWrapperBase,
)
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
# -------------------------------------------------------------------------
# Fixtures
# -------------------------------------------------------------------------
@pytest.fixture
def mock_od_config():
def mock_od_config(mocker: MockerFixture):
"""Create a mock OmniDiffusionConfig for use in tests."""
config = Mock()
config = mocker.Mock()
config.num_gpus = 1
config.master_port = 12345
config.enable_sleep_mode = False
@@ -72,9 +75,10 @@ class MockCustomPipeline:
class TestWorkerWrapperBaseInitialization:
"""Test WorkerWrapperBase initialization behavior."""
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_basic_initialization(self, mock_worker_init, mock_od_config):
def test_basic_initialization(self, mocker: MockerFixture, mock_od_config):
"""Test basic initialization without extensions."""
mock_worker_init = mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
@@ -103,49 +107,49 @@ class TestWorkerWrapperBaseInitialization:
class TestWorkerWrapperBaseExtension:
"""Test WorkerWrapperBase worker extension functionality."""
def test_prepare_worker_class_without_extension(self, mock_od_config):
def test_prepare_worker_class_without_extension(self, mocker: MockerFixture, mock_od_config):
"""Test _prepare_worker_class without a worker extension."""
with patch.object(DiffusionWorker, "__init__", return_value=None):
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
)
worker_class = wrapper._prepare_worker_class()
assert worker_class == DiffusionWorker
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
)
worker_class = wrapper._prepare_worker_class()
assert worker_class == DiffusionWorker
def test_prepare_worker_class_with_extension_class(self, mock_od_config):
def test_prepare_worker_class_with_extension_class(self, mocker: MockerFixture, mock_od_config):
"""Test _prepare_worker_class with an explicit extension class."""
class TestExtension:
def custom_method(self):
return "extension_method"
with patch.object(DiffusionWorker, "__init__", return_value=None):
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=TestExtension,
)
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=TestExtension,
)
assert hasattr(wrapper.worker.__class__, "custom_method")
assert TestExtension in wrapper.worker.__class__.__bases__
assert hasattr(wrapper.worker.__class__, "custom_method")
assert TestExtension in wrapper.worker.__class__.__bases__
@patch("vllm.utils.import_utils.resolve_obj_by_qualname")
def test_prepare_worker_class_with_extension_string(self, mock_resolve, mock_od_config):
def test_prepare_worker_class_with_extension_string(self, mocker: MockerFixture, mock_od_config):
"""Test _prepare_worker_class with worker extension as string."""
mock_resolve = mocker.patch("vllm.utils.import_utils.resolve_obj_by_qualname")
mock_resolve.return_value = TestExtension
with patch.object(DiffusionWorker, "__init__", return_value=None):
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls="tests.diffusion.test_worker_wrapper_base.TestExtension",
)
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls="tests.diffusion.test_worker_wrapper_base.TestExtension",
)
assert hasattr(wrapper.worker.__class__, "custom_method")
assert hasattr(wrapper.worker.__class__, "custom_method")
# -------------------------------------------------------------------------
@@ -156,70 +160,70 @@ class TestWorkerWrapperBaseExtension:
class TestWorkerWrapperBaseDelegation:
"""Test WorkerWrapperBase delegation to wrapped worker."""
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_generate_delegation(self, mock_worker_init, mock_od_config):
def test_generate_delegation(self, mocker: MockerFixture, mock_od_config):
"""Test that generate() delegates to worker.generate()."""
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
mock_output = Mock()
wrapper.worker.generate = Mock(return_value=mock_output)
mock_output = mocker.Mock()
wrapper.worker.generate = mocker.Mock(return_value=mock_output)
mock_requests = [Mock()]
mock_requests = [mocker.Mock()]
result = wrapper.generate(mock_requests)
wrapper.worker.generate.assert_called_once_with(mock_requests)
assert result == mock_output
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_execute_model_delegation(self, mock_worker_init, mock_od_config):
def test_execute_model_delegation(self, mocker: MockerFixture, mock_od_config):
"""Test that execute_model() delegates to worker.execute_model()."""
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
mock_output = Mock()
wrapper.worker.execute_model = Mock(return_value=mock_output)
mock_output = mocker.Mock()
wrapper.worker.execute_model = mocker.Mock(return_value=mock_output)
mock_reqs = [Mock()]
mock_reqs = [mocker.Mock()]
result = wrapper.execute_model(mock_reqs, mock_od_config)
wrapper.worker.execute_model.assert_called_once_with(mock_reqs, mock_od_config)
assert result == mock_output
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_load_weights_delegation(self, mock_worker_init, mock_od_config):
def test_load_weights_delegation(self, mocker: MockerFixture, mock_od_config):
"""Test that load_weights() delegates to worker.load_weights()."""
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
expected_result = {"weight1", "weight2"}
wrapper.worker.load_weights = Mock(return_value=expected_result)
wrapper.worker.load_weights = mocker.Mock(return_value=expected_result)
mock_weights = [("weight1", Mock()), ("weight2", Mock())]
mock_weights = [("weight1", mocker.Mock()), ("weight2", mocker.Mock())]
result = wrapper.load_weights(mock_weights)
wrapper.worker.load_weights.assert_called_once_with(mock_weights)
assert result == expected_result
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_sleep_delegation(self, mock_worker_init, mock_od_config):
def test_sleep_delegation(self, mocker: MockerFixture, mock_od_config):
"""Test that sleep() delegates to worker.sleep()."""
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.sleep = Mock(return_value=True)
wrapper.worker.sleep = mocker.Mock(return_value=True)
result = wrapper.sleep(level=1)
wrapper.worker.sleep.assert_called_once_with(1)
assert result is True
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_wake_up_delegation(self, mock_worker_init, mock_od_config):
def test_wake_up_delegation(self, mocker: MockerFixture, mock_od_config):
"""Test that wake_up() delegates to worker.wake_up()."""
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.wake_up = Mock(return_value=True)
wrapper.worker.wake_up = mocker.Mock(return_value=True)
result = wrapper.wake_up(tags=["weights"])
wrapper.worker.wake_up.assert_called_once_with(["weights"])
assert result is True
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_shutdown_delegation(self, mock_worker_init, mock_od_config):
def test_shutdown_delegation(self, mocker: MockerFixture, mock_od_config):
"""Test that shutdown() delegates to worker.shutdown()."""
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.shutdown = Mock(return_value=None)
wrapper.worker.shutdown = mocker.Mock(return_value=None)
result = wrapper.shutdown()
wrapper.worker.shutdown.assert_called_once()
@@ -234,39 +238,39 @@ class TestWorkerWrapperBaseDelegation:
class TestWorkerWrapperBaseExecuteMethod:
"""Test WorkerWrapperBase.execute_method functionality."""
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_execute_method_success(self, mock_worker_init, mock_od_config):
def test_execute_method_success(self, mocker: MockerFixture, mock_od_config):
"""Test execute_method successfully calls worker method."""
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.test_method = Mock(return_value="method_result")
wrapper.worker.test_method = mocker.Mock(return_value="method_result")
result = wrapper.execute_method("test_method", "arg1", kwarg1="value1")
wrapper.worker.test_method.assert_called_once_with("arg1", kwarg1="value1")
assert result == "method_result"
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_execute_method_with_no_args(self, mock_worker_init, mock_od_config):
def test_execute_method_with_no_args(self, mocker: MockerFixture, mock_od_config):
"""Test execute_method with no arguments."""
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.no_args_method = Mock(return_value="no_args_result")
wrapper.worker.no_args_method = mocker.Mock(return_value="no_args_result")
result = wrapper.execute_method("no_args_method")
wrapper.worker.no_args_method.assert_called_once_with()
assert result == "no_args_result"
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_execute_method_error(self, mock_worker_init, mock_od_config):
def test_execute_method_error(self, mocker: MockerFixture, mock_od_config):
"""Test execute_method raises exception on error."""
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.error_method = Mock(side_effect=RuntimeError("Test error"))
wrapper.worker.error_method = mocker.Mock(side_effect=RuntimeError("Test error"))
with pytest.raises(RuntimeError, match="Test error"):
wrapper.execute_method("error_method")
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_execute_method_invalid_type(self, mock_worker_init, mock_od_config):
def test_execute_method_invalid_type(self, mocker: MockerFixture, mock_od_config):
"""Test execute_method with invalid method type."""
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
with pytest.raises(AssertionError, match="Method must be str"):
@@ -281,26 +285,26 @@ class TestWorkerWrapperBaseExecuteMethod:
class TestWorkerWrapperBaseGetAttr:
"""Test WorkerWrapperBase.__getattr__ delegation."""
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_getattr_delegation(self, mock_worker_init, mock_od_config):
def test_getattr_delegation(self, mocker: MockerFixture, mock_od_config):
"""Test __getattr__ delegates to worker attributes."""
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.custom_attribute = "test_value"
assert wrapper.custom_attribute == "test_value"
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_getattr_method_access(self, mock_worker_init, mock_od_config):
def test_getattr_method_access(self, mocker: MockerFixture, mock_od_config):
"""Test __getattr__ delegates to worker methods."""
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
wrapper.worker.custom_method = Mock(return_value="method_result")
wrapper.worker.custom_method = mocker.Mock(return_value="method_result")
result = wrapper.custom_method()
wrapper.worker.custom_method.assert_called_once()
assert result == "method_result"
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_getattr_missing_attribute(self, mock_worker_init, mock_od_config):
def test_getattr_missing_attribute(self, mocker: MockerFixture, mock_od_config):
"""Test __getattr__ raises AttributeError for missing attributes."""
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
with pytest.raises(AttributeError):
_ = wrapper.nonexistent_attribute
@@ -314,47 +318,46 @@ class TestWorkerWrapperBaseGetAttr:
class TestWorkerWrapperBaseEdgeCases:
"""Test WorkerWrapperBase edge cases and special scenarios."""
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_extension_conflict_warning(self, mock_worker_init, mock_od_config, caplog):
def test_extension_conflict_warning(self, mocker: MockerFixture, mock_od_config, caplog):
"""Test a warning is logged when an extension conflicts with worker."""
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
class ConflictExtension:
def load_model(self):
return "extension_load_model"
with patch.object(DiffusionWorker, "load_model"):
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=ConflictExtension,
)
assert wrapper.worker is not None
mocker.patch.object(DiffusionWorker, "load_model")
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=ConflictExtension,
)
assert wrapper.worker is not None
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_multiple_extensions_same_class(self, mock_worker_init, mock_od_config):
def test_multiple_extensions_same_class(self, mocker: MockerFixture, mock_od_config):
"""Test that applying same extension twice doesn't duplicate it."""
class TestExtension:
def custom_method(self):
return "extension"
with patch.object(DiffusionWorker, "__init__", return_value=None):
wrapper1 = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=TestExtension,
)
wrapper2 = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=TestExtension,
)
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper1 = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=TestExtension,
)
wrapper2 = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=TestExtension,
)
assert hasattr(wrapper1.worker, "custom_method")
assert hasattr(wrapper2.worker, "custom_method")
assert hasattr(wrapper1.worker, "custom_method")
assert hasattr(wrapper2.worker, "custom_method")
# -------------------------------------------------------------------------
@@ -365,11 +368,12 @@ class TestWorkerWrapperBaseEdgeCases:
class TestCustomPipelineWorkerExtension:
"""Test CustomPipelineWorkerExtension functionality."""
@patch("torch.cuda.empty_cache")
@patch("gc.collect")
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_re_init_pipeline_basic(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config):
def test_re_init_pipeline_basic(self, mocker: MockerFixture, mock_od_config):
"""Test basic re_init_pipeline functionality."""
mocker.patch("torch.cuda.empty_cache")
mocker.patch("gc.collect")
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
@@ -378,12 +382,12 @@ class TestCustomPipelineWorkerExtension:
)
# Setup mock model_runner and pipeline
mock_model_runner = Mock()
mock_pipeline = Mock()
mock_model_runner = mocker.Mock()
mock_pipeline = mocker.Mock()
mock_model_runner.pipeline = mock_pipeline
wrapper.worker.model_runner = mock_model_runner
wrapper.worker.init_lora_manager = Mock()
wrapper.worker.load_model = Mock()
wrapper.worker.init_lora_manager = mocker.Mock()
wrapper.worker.load_model = mocker.Mock()
custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
@@ -397,11 +401,12 @@ class TestCustomPipelineWorkerExtension:
)
wrapper.worker.init_lora_manager.assert_called_once()
@patch("torch.cuda.empty_cache")
@patch("gc.collect")
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_re_init_pipeline_cleanup(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config):
def test_re_init_pipeline_cleanup(self, mocker: MockerFixture, mock_od_config):
"""Test that re_init_pipeline properly cleans up old pipeline."""
mock_gc_collect = mocker.patch("gc.collect")
mock_empty_cache = mocker.patch("torch.cuda.empty_cache")
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
@@ -410,12 +415,12 @@ class TestCustomPipelineWorkerExtension:
)
# Setup mock model_runner with pipeline
mock_model_runner = Mock()
mock_pipeline = Mock()
mock_model_runner = mocker.Mock()
mock_pipeline = mocker.Mock()
mock_model_runner.pipeline = mock_pipeline
wrapper.worker.model_runner = mock_model_runner
wrapper.worker.init_lora_manager = Mock()
wrapper.worker.load_model = Mock()
wrapper.worker.init_lora_manager = mocker.Mock()
wrapper.worker.load_model = mocker.Mock()
custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
@@ -426,11 +431,12 @@ class TestCustomPipelineWorkerExtension:
mock_gc_collect.assert_called_once()
mock_empty_cache.assert_called_once()
@patch("torch.cuda.empty_cache")
@patch("gc.collect")
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_re_init_pipeline_none_pipeline(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config):
def test_re_init_pipeline_none_pipeline(self, mocker: MockerFixture, mock_od_config):
"""Test re_init_pipeline when pipeline is None."""
mocker.patch("torch.cuda.empty_cache")
mocker.patch("gc.collect")
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
@@ -439,11 +445,11 @@ class TestCustomPipelineWorkerExtension:
)
# Setup mock model_runner with None pipeline
mock_model_runner = Mock()
mock_model_runner = mocker.Mock()
mock_model_runner.pipeline = None
wrapper.worker.model_runner = mock_model_runner
wrapper.worker.init_lora_manager = Mock()
wrapper.worker.load_model = Mock()
wrapper.worker.init_lora_manager = mocker.Mock()
wrapper.worker.load_model = mocker.Mock()
custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
@@ -453,32 +459,30 @@ class TestCustomPipelineWorkerExtension:
# Verify load_model was still called
wrapper.worker.load_model.assert_called_once()
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_custom_pipeline_args_initialization(self, mock_worker_init, mock_od_config):
def test_custom_pipeline_args_initialization(self, mocker: MockerFixture, mock_od_config):
"""Test initialization with custom_pipeline_args calls re_init_pipeline."""
custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
with patch.object(DiffusionWorker, "__init__", return_value=None):
with patch.object(WorkerWrapperBase, "_prepare_worker_class") as mock_prepare:
# Create a mock worker class with re_init_pipeline
mock_worker_class = Mock()
mock_worker_instance = Mock()
mock_worker_instance.re_init_pipeline = Mock()
mock_worker_class.return_value = mock_worker_instance
mock_prepare.return_value = mock_worker_class
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
mock_prepare = mocker.patch.object(WorkerWrapperBase, "_prepare_worker_class")
# Create a mock worker class with re_init_pipeline
mock_worker_class = mocker.Mock()
mock_worker_instance = mocker.Mock()
mock_worker_instance.re_init_pipeline = mocker.Mock()
mock_worker_class.return_value = mock_worker_instance
mock_prepare.return_value = mock_worker_class
_ = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
custom_pipeline_args=custom_args,
)
_ = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
custom_pipeline_args=custom_args,
)
# Verify re_init_pipeline was called with custom_pipeline_args
mock_worker_instance.re_init_pipeline.assert_called_once_with(custom_args)
# Verify re_init_pipeline was called with custom_pipeline_args
mock_worker_instance.re_init_pipeline.assert_called_once_with(custom_args)
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_custom_pipeline_with_explicit_extension(self, mock_worker_init, mock_od_config):
def test_custom_pipeline_with_explicit_extension(self, mocker: MockerFixture, mock_od_config):
"""Test that explicit worker_extension_cls is preserved when custom_pipeline_args is provided."""
class CustomExtension:
@@ -490,24 +494,25 @@ class TestCustomPipelineWorkerExtension:
custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
with patch.object(DiffusionWorker, "__init__", return_value=None):
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=CustomExtension,
custom_pipeline_args=custom_args,
)
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
base_worker_class=DiffusionWorker,
worker_extension_cls=CustomExtension,
custom_pipeline_args=custom_args,
)
# Should still have the explicitly provided extension
assert CustomExtension in wrapper.worker.__class__.__bases__
assert hasattr(wrapper.worker, "custom_extension_method")
# Should still have the explicitly provided extension
assert CustomExtension in wrapper.worker.__class__.__bases__
assert hasattr(wrapper.worker, "custom_extension_method")
@patch("torch.cuda.empty_cache")
@patch("gc.collect")
@patch.object(DiffusionWorker, "__init__", return_value=None)
def test_re_init_pipeline_multiple_calls(self, mock_worker_init, mock_gc_collect, mock_empty_cache, mock_od_config):
def test_re_init_pipeline_multiple_calls(self, mocker: MockerFixture, mock_od_config):
"""Test calling re_init_pipeline multiple times."""
mocker.patch("torch.cuda.empty_cache")
mocker.patch("gc.collect")
mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
wrapper = WorkerWrapperBase(
gpu_id=0,
od_config=mock_od_config,
@@ -516,13 +521,13 @@ class TestCustomPipelineWorkerExtension:
)
# Setup mock model_runner
mock_model_runner = Mock()
mock_pipeline1 = Mock()
mock_pipeline2 = Mock()
mock_model_runner = mocker.Mock()
mock_pipeline1 = mocker.Mock()
mock_pipeline2 = mocker.Mock()
mock_model_runner.pipeline = mock_pipeline1
wrapper.worker.model_runner = mock_model_runner
wrapper.worker.init_lora_manager = Mock()
wrapper.worker.load_model = Mock()
wrapper.worker.init_lora_manager = mocker.Mock()
wrapper.worker.load_model = mocker.Mock()
# First call
custom_args1 = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
tests/distributed/omni_connectors/test_adapter_and_flow.py
+9 -8
View File
@@ -1,19 +1,20 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import MagicMock
import pytest
from pytest_mock import MockerFixture
from vllm_omni.distributed.omni_connectors.adapter import try_recv_via_connector, try_send_via_connector
from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec, OmniTransferConfig
from vllm_omni.distributed.omni_connectors.utils.initialization import get_connectors_config_for_stage
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@pytest.fixture
def mock_objects():
return {"connector": MagicMock(), "metrics": MagicMock(), "queue_fn": MagicMock()}
def mock_objects(mocker: MockerFixture):
return {"connector": mocker.MagicMock(), "metrics": mocker.MagicMock(), "queue_fn": mocker.MagicMock()}
def test_send_success(mock_objects):
@@ -138,7 +139,7 @@ def test_recv_no_connector():
assert inputs is None
def test_shm_connector_flow():
def test_shm_connector_flow(mocker: MockerFixture):
"""
Verify the full flow: Send -> Adapter -> Connector -> Adapter -> Recv.
Using real SharedMemoryConnector (inline mode for simplicity).
@@ -161,7 +162,7 @@ def test_shm_connector_flow():
def mock_submit(payload):
queue_capture.append(payload)
mock_metrics = MagicMock()
mock_metrics = mocker.MagicMock()
# 3. Send
success = try_send_via_connector(
@@ -217,7 +218,7 @@ def test_get_connectors_for_stage():
assert stage_2_config["from_stage_1"]["spec"]["name"] == "C2"
def test_recv_with_missing_metadata():
def test_recv_with_missing_metadata(mocker: MockerFixture):
"""Test recv when queue payload is malformed (missing metadata)."""
# Connector expects metadata but task doesn't have it
task = {
@@ -226,7 +227,7 @@ def test_recv_with_missing_metadata():
"from_stage": "0",
# Missing "connector_metadata"
}
mock_conn = MagicMock()
mock_conn = mocker.MagicMock()
# If get is called with None metadata, connector usually handles it or adapter handles exception
mock_conn.get.side_effect = Exception("Get failed")
tests/distributed/omni_connectors/test_basic_connectors.py
+6 -4
View File
@@ -1,14 +1,16 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import MagicMock
import pytest
from pytest_mock import MockerFixture
from vllm_omni.distributed.omni_connectors.connectors.shm_connector import SharedMemoryConnector
from vllm_omni.distributed.omni_connectors.factory import OmniConnectorFactory
from vllm_omni.distributed.omni_connectors.utils.config import ConnectorSpec
from vllm_omni.distributed.omni_connectors.utils.serialization import OmniSerializer
# pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_basic_serialization():
"""Test basic msgpack serialization."""
@@ -79,19 +81,19 @@ def test_put_get_inline(shm_connector):
assert size == ret_size
def test_put_get_shm(shm_connector, monkeypatch):
def test_put_get_shm(mocker: MockerFixture, shm_connector, monkeypatch: pytest.MonkeyPatch):
"""Test SHM transfer logic for large data (Mocked)."""
# Create data larger than 100 bytes
data = {"large": "x" * 200}
# Mock SHM return values
mock_handle = {"name": "test_shm", "size": 200}
mock_write = MagicMock(return_value=mock_handle)
mock_write = mocker.MagicMock(return_value=mock_handle)
monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_write_bytes", mock_write)
# When reading, return the serialized bytes of the data
serialized_data = shm_connector.serialize_obj(data)
mock_read = MagicMock(return_value=serialized_data)
mock_read = mocker.MagicMock(return_value=serialized_data)
monkeypatch.setattr("vllm_omni.distributed.omni_connectors.connectors.shm_connector.shm_read_bytes", mock_read)
# Put
tests/distributed/omni_connectors/test_chunk_transfer_adapter.py
+3 -3
View File
@@ -4,10 +4,10 @@
import threading
from collections import deque
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
import torch
from pytest_mock import MockerFixture
from vllm.v1.request import RequestStatus
from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase
@@ -39,9 +39,9 @@ def _req(req_id: str, status: RequestStatus, external_req_id: str | None = None)
@pytest.fixture
def build_adapter(monkeypatch):
def build_adapter(monkeypatch, mocker: MockerFixture):
def _build(*, stage_id: int = 1, model_mode: str = "ar", max_num_seqs: int = 2):
connector = MagicMock()
connector = mocker.MagicMock()
connector.stage_id = stage_id
connector.get.return_value = None
connector.put.return_value = (True, 1, {})
tests/distributed/omni_connectors/test_mooncake_transfer_engine_buffer.py
+27 -30
View File
@@ -7,7 +7,6 @@ These tests do NOT require Mooncake or RDMA environment.
"""
import threading
import unittest
import pytest
import torch
@@ -22,7 +21,7 @@ pytestmark = [pytest.mark.cpu, pytest.mark.parallel]
@pytest.mark.core_model
class TestBufferAllocator(unittest.TestCase):
class TestBufferAllocator:
"""Unit tests for BufferAllocator."""
def test_basic_alloc_free(self):
@@ -30,15 +29,15 @@ class TestBufferAllocator(unittest.TestCase):
allocator = BufferAllocator(total_size=4096, alignment=64)
offset1 = allocator.alloc(512)
self.assertEqual(offset1, 0)
assert offset1 == 0
offset2 = allocator.alloc(512)
self.assertGreater(offset2, 0)
assert offset2 > 0
# Free first block, should be reusable
allocator.free(offset1, 512)
offset3 = allocator.alloc(512)
self.assertEqual(offset3, 0)
assert offset3 == 0
def test_alignment(self):
"""Verify allocation respects alignment."""
@@ -47,22 +46,22 @@ class TestBufferAllocator(unittest.TestCase):
_offset1 = allocator.alloc(100)
offset2 = allocator.alloc(100)
self.assertEqual(offset2 % 128, 0)
self.assertEqual(offset2, 128)
assert offset2 % 128 == 0
assert offset2 == 128
def test_exhaustion_and_recovery(self):
"""Test that full allocation fails, then succeeds after free."""
allocator = BufferAllocator(total_size=1024, alignment=64)
offset = allocator.alloc(1024)
self.assertEqual(offset, 0)
assert offset == 0
with self.assertRaises(MemoryError):
with pytest.raises(MemoryError):
allocator.alloc(64)
allocator.free(offset, 1024)
offset2 = allocator.alloc(1024)
self.assertEqual(offset2, 0)
assert offset2 == 0
def test_thread_safety(self):
"""Verify allocator is thread-safe under concurrent access."""
@@ -84,10 +83,10 @@ class TestBufferAllocator(unittest.TestCase):
for t in threads:
t.join()
self.assertEqual(len(errors), 0, f"Thread safety errors: {errors}")
assert len(errors) == 0, f"Thread safety errors: {errors}"
class TestAllocatorInvariants(unittest.TestCase):
class TestAllocatorInvariants:
"""
Defensive invariant tests for BufferAllocator: double-free, partial
overlap corruption, and adjacent-block merging.
@@ -106,7 +105,7 @@ class TestAllocatorInvariants(unittest.TestCase):
allocator.free(offset, 256)
# Pool should still be consistent: allocate full size back
offset2 = allocator.alloc(4096)
self.assertEqual(offset2, 0)
assert offset2 == 0
@pytest.mark.slow
def test_double_free_after_merge_is_safe(self):
@@ -123,7 +122,7 @@ class TestAllocatorInvariants(unittest.TestCase):
allocator.free(a, 64) # should not raise
# Pool should still be fully usable
offset = allocator.alloc(4096)
self.assertEqual(offset, 0)
assert offset == 0
@pytest.mark.slow
def test_partial_overlap_raises_corruption(self):
@@ -133,7 +132,7 @@ class TestAllocatorInvariants(unittest.TestCase):
b = allocator.alloc(128)
allocator.free(a, 128) # [0, 128) is now free
# Try to free a region that starts inside [0,128) but extends beyond
with self.assertRaises(RuntimeError):
with pytest.raises(RuntimeError):
allocator.free(64, 128) # [64, 192) overlaps with free [0, 128)
# b is still allocated; freeing b should be fine
@@ -154,7 +153,7 @@ class TestAllocatorInvariants(unittest.TestCase):
# covering at least 192 bytes (3 * 64).
# Verify by allocating a contiguous block of 192 bytes.
offset = allocator.alloc(192)
self.assertEqual(offset, 0, "Adjacent blocks were not merged properly")
assert offset == 0, "Adjacent blocks were not merged properly"
@pytest.mark.slow
def test_fragmentation_and_defrag(self):
@@ -174,7 +173,7 @@ class TestAllocatorInvariants(unittest.TestCase):
allocator.free(d, 64) # free blocks: [64, 128) and [192, 256)
# Pool has two 64-byte holes; contiguous 128 is unavailable
with self.assertRaises(MemoryError):
with pytest.raises(MemoryError):
allocator.alloc(128)
allocator.free(a, 64)
@@ -182,14 +181,16 @@ class TestAllocatorInvariants(unittest.TestCase):
# After freeing everything, full pool should be available
offset = allocator.alloc(256)
self.assertEqual(offset, 0)
assert offset == 0
@pytest.mark.core_model
class TestManagedBuffer(unittest.TestCase):
class TestManagedBuffer:
"""Unit tests for ManagedBuffer."""
def setUp(self):
@pytest.fixture(autouse=True)
def _setup(self):
# automatically invoked for every test method in the class
self.allocator = BufferAllocator(total_size=4096, alignment=64)
self.pool = torch.zeros(4096, dtype=torch.uint8)
@@ -203,12 +204,12 @@ class TestManagedBuffer(unittest.TestCase):
self.pool[offset : offset + 64] = src.view(torch.uint8)
# Raw uint8 view
self.assertEqual(buf.tensor.shape[0], 64)
assert buf.tensor.shape[0] == 64
# Typed view
typed = buf.as_tensor(dtype=torch.float32, shape=(4, 4))
self.assertEqual(typed.shape, (4, 4))
self.assertTrue(torch.equal(typed.flatten(), src))
assert typed.shape == (4, 4)
assert torch.equal(typed.flatten(), src)
buf.release()
@@ -217,14 +218,10 @@ class TestManagedBuffer(unittest.TestCase):
offset = self.allocator.alloc(128)
with ManagedBuffer(self.allocator, offset, 128, self.pool) as buf:
self.assertFalse(buf._released)
assert not buf._released
self.assertTrue(buf._released)
assert buf._released
# Space should be reusable
new_offset = self.allocator.alloc(128)
self.assertEqual(new_offset, offset)
if __name__ == "__main__":
unittest.main()
assert new_offset == offset
tests/distributed/omni_connectors/test_mooncake_transfer_engine_rdma.py
+83 -92
View File
@@ -13,7 +13,6 @@ import os
import subprocess
import threading
import time
import unittest
import pytest
import torch
@@ -195,43 +194,43 @@ def _md5(tensor: torch.Tensor) -> str:
# ---------------------------------------------------------------------------
@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
class TestBasicConnector(unittest.TestCase):
@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
class TestBasicConnector:
"""Verify connector initialization, put, cleanup, and health check."""
def test_initialization(self):
port = _free_port()
with MooncakeTransferEngineConnector(_connector_config(port, pool_size=1024 * 1024)) as c:
self.assertNotEqual(c.rpc_port, 0)
self.assertEqual(c.pool_size, 1024 * 1024)
self.assertTrue(c.pool.is_pinned())
assert c.rpc_port != 0
assert c.pool_size == 1024 * 1024
assert c.pool.is_pinned()
health = c.health()
self.assertEqual(health["status"], "healthy")
assert health["status"] == "healthy"
def test_put_tensor_bytes_object(self):
"""Put tensor / bytes / dict and verify metadata."""
port = _free_port()
with MooncakeTransferEngineConnector(_connector_config(port)) as c:
ok, sz, meta = c.put("s0", "s1", "t", torch.randn(100))
self.assertTrue(ok)
self.assertTrue(meta["is_fast_path"])
assert ok
assert meta["is_fast_path"]
ok, sz, meta = c.put("s0", "s1", "b", b"hello" * 100)
self.assertTrue(ok)
self.assertTrue(meta["is_fast_path"])
assert ok
assert meta["is_fast_path"]
ok, sz, meta = c.put("s0", "s1", "d", {"k": [1, 2, 3]})
self.assertTrue(ok)
self.assertFalse(meta["is_fast_path"])
assert ok
assert not meta["is_fast_path"]
def test_cleanup_releases_buffer(self):
port = _free_port()
with MooncakeTransferEngineConnector(_connector_config(port)) as c:
c.put("s0", "s1", "r1", torch.randn(100))
key = MooncakeTransferEngineConnector._make_key("r1", "s0", "s1")
self.assertIn(key, c._local_buffers)
assert key in c._local_buffers
c.cleanup("r1", from_stage="s0", to_stage="s1")
self.assertNotIn(key, c._local_buffers)
assert key not in c._local_buffers
def test_pool_exhaustion_and_recovery(self):
"""Fill pool, verify failure, free, verify recovery."""
@@ -247,7 +246,7 @@ class TestBasicConnector(unittest.TestCase):
for rid in ids:
c.cleanup(rid, from_stage="s", to_stage="s")
ok, _, _ = c.put("s", "s", "recovery", torch.randn(1000))
self.assertTrue(ok, "Pool recovery failed after cleanup")
assert ok, "Pool recovery failed after cleanup"
# ---------------------------------------------------------------------------
@@ -255,8 +254,8 @@ class TestBasicConnector(unittest.TestCase):
# ---------------------------------------------------------------------------
@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
class TestEndToEnd(unittest.TestCase):
@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
class TestEndToEnd:
"""E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload, mixed types."""
def _pair(self, pool_size=16 * 1024 * 1024):
@@ -269,14 +268,14 @@ class TestEndToEnd(unittest.TestCase):
try:
orig = torch.randn(1024, 1024, dtype=torch.float32)
ok, sz, meta = p.put("s0", "s1", "t1", orig)
self.assertTrue(ok)
assert ok
time.sleep(0.5)
result = c.get("s0", "s1", "t1", meta)
self.assertIsNotNone(result)
assert result is not None
buf, _ = result
self.assertIsInstance(buf, ManagedBuffer)
assert isinstance(buf, ManagedBuffer)
recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
self.assertTrue(torch.equal(recon, orig))
assert torch.equal(recon, orig)
buf.release()
finally:
p.close()
@@ -287,15 +286,15 @@ class TestEndToEnd(unittest.TestCase):
try:
orig = b"Hello RDMA! " * 1000
ok, _, meta = p.put("s0", "s1", "b1", orig)
self.assertTrue(ok)
assert ok
time.sleep(0.5)
result = c.get("s0", "s1", "b1", meta)
self.assertIsNotNone(result)
assert result is not None
buf, _ = result
recv = buf.to_bytes() if hasattr(buf, "to_bytes") else buf
if hasattr(buf, "release"):
buf.release()
self.assertEqual(recv, orig)
assert recv == orig
finally:
p.close()
c.close()
@@ -305,13 +304,13 @@ class TestEndToEnd(unittest.TestCase):
try:
orig = {"msg": "hello", "nums": [1, 2, 3], "nested": {"k": "v"}}
ok, _, meta = p.put("s0", "s1", "o1", orig)
self.assertTrue(ok)
self.assertFalse(meta["is_fast_path"])
assert ok
assert not meta["is_fast_path"]
time.sleep(1.0)
result = c.get("s0", "s1", "o1", meta)
self.assertIsNotNone(result)
assert result is not None
recv, _ = result
self.assertEqual(recv, orig)
assert recv == orig
finally:
p.close()
c.close()
@@ -327,15 +326,15 @@ class TestEndToEnd(unittest.TestCase):
sbuf.as_tensor(dtype=torch.float32, shape=shape).copy_(ref)
ok, sz, meta = p.put("s0", "s1", "zc", sbuf)
self.assertTrue(ok)
self.assertTrue(meta.get("is_fast_path"))
assert ok
assert meta.get("is_fast_path")
time.sleep(0.5)
result = c.get("s0", "s1", "zc", meta)
self.assertIsNotNone(result)
assert result is not None
rbuf, _ = result
recon = rbuf.as_tensor(dtype=torch.float32, shape=shape)
self.assertTrue(torch.equal(recon, ref))
assert torch.equal(recon, ref)
rbuf.release()
finally:
p.close()
@@ -348,13 +347,13 @@ class TestEndToEnd(unittest.TestCase):
orig = torch.randn(5000, 5000, dtype=torch.float32)
md5_send = _md5(orig)
ok, _, meta = p.put("s0", "s1", "lg", orig)
self.assertTrue(ok)
assert ok
time.sleep(1.0)
result = c.get("s0", "s1", "lg", meta)
self.assertIsNotNone(result)
assert result is not None
buf, _ = result
recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
self.assertEqual(_md5(recon), md5_send)
assert _md5(recon) == md5_send
buf.release()
finally:
p.close()
@@ -374,23 +373,23 @@ class TestEndToEnd(unittest.TestCase):
time.sleep(0.5)
for name, data in cases:
ok, _, meta = p.put("s0", "s1", name, data)
self.assertTrue(ok, f"put failed: {name}")
assert ok, f"put failed: {name}"
time.sleep(0.3)
result = c.get("s0", "s1", name, meta)
self.assertIsNotNone(result, f"get failed: {name}")
assert result is not None, f"get failed: {name}"
recv, _ = result
if isinstance(data, torch.Tensor):
self.assertIsInstance(recv, ManagedBuffer)
assert isinstance(recv, ManagedBuffer)
recon = recv.as_tensor(dtype=data.dtype, shape=data.shape)
self.assertTrue(torch.equal(recon, data), f"mismatch: {name}")
assert torch.equal(recon, data), f"mismatch: {name}"
recv.release()
elif isinstance(data, bytes):
got = recv.to_bytes() if hasattr(recv, "to_bytes") else recv
if hasattr(recv, "release"):
recv.release()
self.assertEqual(got, data)
assert got == data
else:
self.assertEqual(recv, data)
assert recv == data
finally:
p.close()
c.close()
@@ -420,24 +419,24 @@ class TestEndToEnd(unittest.TestCase):
t.start()
for t in threads:
t.join()
self.assertEqual(len(errors), 0, f"errors: {errors}")
assert len(errors) == 0, f"errors: {errors}"
def test_auto_cleanup(self):
"""Producer buffer should be released after consumer get."""
p, c = self._pair()
try:
ok, _, meta = p.put("s0", "s1", "ac", torch.randn(100))
self.assertTrue(ok)
assert ok
key = MooncakeTransferEngineConnector._make_key("ac", "s0", "s1")
self.assertIn(key, p._local_buffers)
assert key in p._local_buffers
time.sleep(0.5)
result = c.get("s0", "s1", "ac", meta)
self.assertIsNotNone(result)
assert result is not None
buf, _ = result
if isinstance(buf, ManagedBuffer):
buf.release()
time.sleep(0.3)
self.assertNotIn(key, p._local_buffers)
assert key not in p._local_buffers
finally:
p.close()
c.close()
@@ -448,22 +447,22 @@ class TestEndToEnd(unittest.TestCase):
# ---------------------------------------------------------------------------
@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
class TestLifecycle(unittest.TestCase):
@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
class TestLifecycle:
"""Close, context manager, double-close safety."""
def test_close_releases_resources(self):
c = MooncakeTransferEngineConnector(_connector_config(_free_port(), pool_size=1024 * 1024))
c.put("s0", "s1", "x", torch.randn(100))
c.close()
self.assertTrue(c._stop_event.is_set())
self.assertEqual(len(c._local_buffers), 0)
assert c._stop_event.is_set()
assert len(c._local_buffers) == 0
def test_context_manager(self):
with MooncakeTransferEngineConnector(_connector_config(_free_port())) as c:
ok, _, _ = c.put("s0", "s1", "ctx", torch.randn(50))
self.assertTrue(ok)
self.assertTrue(c._stop_event.is_set())
assert ok
assert c._stop_event.is_set()
def test_double_close_safe(self):
c = MooncakeTransferEngineConnector(_connector_config(_free_port()))
@@ -477,9 +476,9 @@ class TestLifecycle(unittest.TestCase):
@pytest.mark.cuda
@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
@unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
class TestGPUPool(unittest.TestCase):
@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
class TestGPUPool:
"""GPU memory pool: initialization, put (CPU/GPU tensor), E2E transfer."""
@staticmethod
@@ -488,18 +487,18 @@ class TestGPUPool(unittest.TestCase):
def test_gpu_pool_init(self):
with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c:
self.assertEqual(c.pool_device, "cuda:0")
self.assertTrue(c.pool.is_cuda)
assert c.pool_device == "cuda:0"
assert c.pool.is_cuda
def test_gpu_pool_put_cpu_and_gpu_tensor(self):
with MooncakeTransferEngineConnector(self._gpu_cfg(_free_port())) as c:
ok, _, meta = c.put("s0", "s1", "h2d", torch.randn(256, 256))
self.assertTrue(ok)
self.assertTrue(meta["is_fast_path"])
assert ok
assert meta["is_fast_path"]
ok, _, meta = c.put("s0", "s1", "d2d", torch.randn(256, 256, device="cuda:0"))
self.assertTrue(ok)
self.assertTrue(meta["is_fast_path"])
assert ok
assert meta["is_fast_path"]
def test_gpu_e2e_transfer(self):
p = MooncakeTransferEngineConnector(self._gpu_cfg(_free_port()))
@@ -507,14 +506,14 @@ class TestGPUPool(unittest.TestCase):
try:
orig = torch.randn(512, 512, dtype=torch.float32, device="cuda:0")
ok, _, meta = p.put("s0", "s1", "ge", orig)
self.assertTrue(ok)
assert ok
time.sleep(0.5)
result = c.get("s0", "s1", "ge", meta)
self.assertIsNotNone(result)
assert result is not None
buf, _ = result
recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
self.assertTrue(recon.is_cuda)
self.assertTrue(torch.equal(recon.cpu(), orig.cpu()))
assert recon.is_cuda
assert torch.equal(recon.cpu(), orig.cpu())
buf.release()
finally:
p.close()
@@ -527,8 +526,8 @@ class TestGPUPool(unittest.TestCase):
@pytest.mark.slow
@unittest.skipIf(TransferEngine is None, "Mooncake TransferEngine not available")
class TestStressCorrectness(unittest.TestCase):
@pytest.mark.skipif(TransferEngine is None, reason="Mooncake TransferEngine not available")
class TestStressCorrectness:
"""
Slow but high-value regression tests: concurrent put+get with data
integrity, large payload, edge cases, and sustained stress.
@@ -581,20 +580,20 @@ class TestStressCorrectness(unittest.TestCase):
for t in threads:
t.join()
self.assertEqual(len(errors), 0, f"Producer errors: {errors}")
self.assertEqual(len(meta_map), num_workers)
assert len(errors) == 0, f"Producer errors: {errors}"
assert len(meta_map) == num_workers
time.sleep(1.0)
# Consumer: get each tensor and verify integrity
for rid, meta in meta_map.items():
result = c.get("s0", "s1", rid, meta)
self.assertIsNotNone(result, f"get returned None for {rid}")
assert result is not None, f"get returned None for {rid}"
buf, _ = result
self.assertIsInstance(buf, ManagedBuffer, f"{rid} not ManagedBuffer")
assert isinstance(buf, ManagedBuffer), f"{rid} not ManagedBuffer"
recon = buf.as_tensor(dtype=torch.float32, shape=(512, 512))
recv_md5 = _md5(recon)
self.assertEqual(recv_md5, md5_sent[rid], f"MD5 mismatch for {rid}")
assert recv_md5 == md5_sent[rid], f"MD5 mismatch for {rid}"
buf.release()
finally:
p.close()
@@ -683,12 +682,8 @@ class TestStressCorrectness(unittest.TestCase):
for t in threads:
t.join()
ct.join(timeout=60)
self.assertEqual(len(errors), 0, f"errors: {errors}")
self.assertEqual(
consumed_count[0],
num_items,
f"Consumer only processed {consumed_count[0]}/{num_items} items",
)
assert len(errors) == 0, f"errors: {errors}"
assert consumed_count[0] == num_items, f"Consumer only processed {consumed_count[0]}/{num_items} items"
finally:
p.close()
c.close()
@@ -701,13 +696,13 @@ class TestStressCorrectness(unittest.TestCase):
try:
orig = torch.tensor([3.14], dtype=torch.float32)
ok, _, meta = p.put("s0", "s1", "tiny", orig)
self.assertTrue(ok)
assert ok
time.sleep(0.5)
result = c.get("s0", "s1", "tiny", meta)
self.assertIsNotNone(result)
assert result is not None
buf, _ = result
recon = buf.as_tensor(dtype=torch.float32, shape=(1,))
self.assertTrue(torch.equal(recon, orig))
assert torch.equal(recon, orig)
buf.release()
finally:
p.close()
@@ -718,7 +713,7 @@ class TestStressCorrectness(unittest.TestCase):
port = _free_port()
with MooncakeTransferEngineConnector(_connector_config(port, pool_size=8 * 1024 * 1024)) as c:
ok, sz, meta = c.put("s0", "s1", "empty_b", b"")
self.assertFalse(ok, "Empty bytes should be rejected by connector")
assert not ok, "Empty bytes should be rejected by connector"
# -- Large payload stress (500 MB) --
@@ -731,13 +726,13 @@ class TestStressCorrectness(unittest.TestCase):
orig = torch.randn(side, side, dtype=torch.float32)
md5_send = _md5(orig)
ok, _, meta = p.put("s0", "s1", "500mb", orig)
self.assertTrue(ok)
assert ok
time.sleep(2.0)
result = c.get("s0", "s1", "500mb", meta)
self.assertIsNotNone(result)
assert result is not None
buf, _ = result
recon = buf.as_tensor(dtype=orig.dtype, shape=orig.shape)
self.assertEqual(_md5(recon), md5_send)
assert _md5(recon) == md5_send
buf.release()
finally:
p.close()
@@ -752,12 +747,8 @@ class TestStressCorrectness(unittest.TestCase):
for i in range(50):
rid = f"cycle_{i}"
ok, _, _ = c.put("s0", "s1", rid, torch.randn(256, 256))
self.assertTrue(ok, f"Put failed at iteration {i}")
assert ok, f"Put failed at iteration {i}"
c.cleanup(rid, from_stage="s0", to_stage="s1")
# After all cycles, pool should be fully recovered
ok, _, _ = c.put("s0", "s1", "final", torch.randn(1024, 1024))
self.assertTrue(ok, "Pool not fully recovered after rapid cycles")
if __name__ == "__main__":
unittest.main()
assert ok, "Pool not fully recovered after rapid cycles"
tests/e2e/offline_inference/test_ovis_image.py
+51 -49
View File
@@ -12,11 +12,11 @@ Strategy:
and correctness of the model definition itself, independent of the pipeline mocks.
"""
from unittest.mock import MagicMock, patch
import pytest
import torch
from pytest_mock import MockerFixture
from tests.utils import hardware_test
from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
# Mock the OvisImageTransformer2DModel to avoid complex init if needed,
@@ -29,15 +29,15 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@pytest.fixture
def mock_dependencies(monkeypatch):
def mock_dependencies(mocker: MockerFixture, monkeypatch: pytest.MonkeyPatch):
"""
Mock external dependencies to avoid loading real models.
"""
device = get_local_device()
# Mock Tokenizer
mock_tokenizer = MagicMock()
mock_tokenizer.return_value = MagicMock(
mock_tokenizer = mocker.MagicMock()
mock_tokenizer.return_value = mocker.MagicMock(
input_ids=torch.zeros((1, 50), dtype=torch.long, device=device),
attention_mask=torch.ones((1, 50), dtype=torch.long, device=device),
)
@@ -45,13 +45,13 @@ def mock_dependencies(monkeypatch):
mock_tokenizer.model_max_length = 1024
# Mock Text Encoder
mock_text_encoder = MagicMock()
mock_text_encoder = mocker.MagicMock()
mock_text_encoder.dtype = torch.float32
# Output of text encoder must be on the same device as inputs (which are moved to execution_device)
mock_text_encoder.return_value.last_hidden_state = torch.randn(1, 50, 32, device=device)
# Mock VAE
mock_vae = MagicMock()
mock_vae = mocker.MagicMock()
mock_vae.config.block_out_channels = [128, 256, 512, 512] # Scale factor 8
mock_vae.config.scale_factor_temporal = 1
mock_vae.config.scale_factor_spatial = 8
@@ -63,8 +63,8 @@ def mock_dependencies(monkeypatch):
mock_vae.to.return_value = mock_vae
# Mock Scheduler
mock_scheduler = MagicMock()
mock_scheduler.config = MagicMock()
mock_scheduler = mocker.MagicMock()
mock_scheduler.config = mocker.MagicMock()
# Timesteps on device to match latents during denoising loop interaction if needed
mock_scheduler.timesteps = torch.tensor([1.0, 0.5, 0.0], device=device)
mock_scheduler.set_timesteps.return_value = None
@@ -95,7 +95,7 @@ def mock_dependencies(monkeypatch):
@pytest.fixture
def ovis_pipeline(mock_dependencies, monkeypatch):
def ovis_pipeline(mocker: MockerFixture, mock_dependencies, monkeypatch: pytest.MonkeyPatch):
"""
Creates an OvisImagePipeline instance with mocked components.
"""
@@ -122,8 +122,8 @@ def ovis_pipeline(mock_dependencies, monkeypatch):
# Mock Transformer Layer separately to avoid full init
# We patch OvisImageTransformer2DModel class in the module
mock_transformer_cls = MagicMock()
mock_transformer_instance = MagicMock()
mock_transformer_cls = mocker.MagicMock()
mock_transformer_instance = mocker.MagicMock()
mock_transformer_instance.dtype = torch.float32
mock_transformer_instance.in_channels = 16 # Must be 16 so num_channel_latents=4, packed=16
# Forward return: noise prediction
@@ -144,8 +144,8 @@ def ovis_pipeline(mock_dependencies, monkeypatch):
# Initialize pipeline
# We use a dummy model path check override
with patch("os.path.exists", return_value=True):
pipeline = OvisImagePipeline(od_config=od_config)
mocker.patch("os.path.exists", return_value=True)
pipeline = OvisImagePipeline(od_config=od_config)
return pipeline
@@ -223,9 +223,11 @@ def test_resolution_check(ovis_pipeline):
assert output is not None
def test_real_transformer_init_and_forward():
@pytest.mark.core_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4", "rocm": "MI325"})
def test_real_transformer_init_and_forward(mocker: MockerFixture):
"""Test the real OvisImageTransformer2DModel initialization and forward pass for coverage."""
from unittest.mock import patch
from vllm_omni.diffusion.models.ovis_image.ovis_image_transformer import OvisImageTransformer2DModel
@@ -249,42 +251,42 @@ def test_real_transformer_init_and_forward():
# Mock distributed state for QKVParallelLinear initialization
# We patch get_tp_group because get_tensor_model_parallel_rank calls it and asserts _TP is not None
mock_group = MagicMock()
mock_group = mocker.MagicMock()
mock_group.rank_in_group = 0
mock_group.world_size = 1
with patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group):
# Initialize real model
model = OvisImageTransformer2DModel(
od_config=od_config,
patch_size=1,
in_channels=16,
out_channels=16,
num_single_layers=1,
attention_head_dim=8,
num_attention_heads=2,
joint_attention_dim=32,
axes_dims_rope=(2, 2, 4),
).to(device)
mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
# Initialize real model
model = OvisImageTransformer2DModel(
od_config=od_config,
patch_size=1,
in_channels=16,
out_channels=16,
num_single_layers=1,
attention_head_dim=8,
num_attention_heads=2,
joint_attention_dim=32,
axes_dims_rope=(2, 2, 4),
).to(device)
# Create dummy inputs
B, Seq, C = 1, 16, 16
hidden_states = torch.randn(B, Seq, C, device=device)
encoder_hidden_states = torch.randn(B, 10, 32, device=device) # joint_attention_dim=32
timestep = torch.tensor([1], device=device)
img_ids = torch.zeros(Seq, 3, device=device)
txt_ids = torch.zeros(10, 3, device=device)
# Create dummy inputs
B, Seq, C = 1, 16, 16
hidden_states = torch.randn(B, Seq, C, device=device)
encoder_hidden_states = torch.randn(B, 10, 32, device=device) # joint_attention_dim=32
timestep = torch.tensor([1], device=device)
img_ids = torch.zeros(Seq, 3, device=device)
txt_ids = torch.zeros(10, 3, device=device)
# Run forward
output = model(
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
timestep=timestep,
img_ids=img_ids,
txt_ids=txt_ids,
return_dict=False,
)
# Run forward
output = model(
hidden_states=hidden_states,
encoder_hidden_states=encoder_hidden_states,
timestep=timestep,
img_ids=img_ids,
txt_ids=txt_ids,
return_dict=False,
)
assert output is not None
assert isinstance(output, tuple)
assert output[0].shape == hidden_states.shape
assert output is not None
assert isinstance(output, tuple)
assert output[0].shape == hidden_states.shape
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
+85 -84
View File
@@ -1,13 +1,13 @@
# 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 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 +26,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 +52,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 +65,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 +160,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 +185,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 +252,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 +288,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 +307,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 +380,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 +408,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 +443,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 +466,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
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
@@ -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 import SamplingParams
from vllm_omni.entrypoints.stage_utils import SHUTDOWN_TASK
@@ -123,7 +123,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)
@@ -149,6 +149,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."""
@@ -167,11 +169,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:
@@ -235,21 +237,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
@@ -298,21 +300,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
@@ -336,7 +338,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
@@ -356,7 +358,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
@@ -387,7 +389,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.
@@ -400,7 +402,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:
@@ -447,15 +449,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):
@@ -526,7 +528,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):
@@ -568,7 +570,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(
@@ -598,11 +602,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(
@@ -614,7 +618,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,
)
@@ -623,9 +627,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
@@ -645,7 +649,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(
@@ -668,11 +672,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",
@@ -681,14 +685,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
@@ -697,7 +701,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
@@ -723,11 +727,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",
@@ -736,16 +740,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
@@ -802,7 +806,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
@@ -829,11 +835,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",
@@ -842,16 +848,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
@@ -892,7 +898,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
@@ -919,11 +927,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",
@@ -932,16 +940,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
@@ -975,7 +983,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(
@@ -998,11 +1006,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",
@@ -1014,22 +1022,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
@@ -1039,7 +1047,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(
@@ -1062,11 +1070,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",
@@ -1075,14 +1083,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")
@@ -1122,7 +1130,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(
@@ -1145,11 +1153,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",
@@ -1158,14 +1166,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:
tests/test_outputs.py
+23 -20
View File
@@ -2,12 +2,14 @@
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for OmniRequestOutput class."""
from unittest.mock import Mock
import pytest
from PIL import Image
from pytest_mock import MockerFixture
from vllm_omni.outputs import OmniRequestOutput
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestOmniRequestOutput:
"""Tests for OmniRequestOutput class."""
@@ -28,17 +30,17 @@ class TestOmniRequestOutput:
assert output.is_diffusion_output
assert output.num_images == 1
def test_from_pipeline(self):
def test_from_pipeline(self, mocker: MockerFixture):
"""Test creating output from pipeline stage."""
mock_request_output = Mock()
mock_request_output = mocker.Mock()
mock_request_output.request_id = "pipeline-123"
mock_request_output.prompt_token_ids = [1, 2, 3]
mock_request_output.outputs = [Mock()]
mock_request_output.outputs = [mocker.Mock()]
mock_request_output.encoder_prompt_token_ids = None
mock_request_output.prompt_logprobs = None
mock_request_output.num_cached_tokens = 10
mock_request_output.kv_transfer_params = None
mock_request_output.multimodal_output = {"image": Mock()}
mock_request_output.multimodal_output = {"image": mocker.Mock()}
output = OmniRequestOutput.from_pipeline(
stage_id=0,
@@ -51,9 +53,9 @@ class TestOmniRequestOutput:
assert output.final_output_type == "text"
assert output.is_pipeline_output
def test_prompt_token_ids_property(self):
def test_prompt_token_ids_property(self, mocker: MockerFixture):
"""Test prompt_token_ids property for streaming compatibility."""
mock_request_output = Mock()
mock_request_output = mocker.Mock()
mock_request_output.prompt_token_ids = [1, 2, 3, 4, 5]
output = OmniRequestOutput.from_pipeline(
@@ -73,10 +75,10 @@ class TestOmniRequestOutput:
)
assert output.prompt_token_ids is None
def test_outputs_property(self):
def test_outputs_property(self, mocker: MockerFixture):
"""Test outputs property for chat completion compatibility."""
mock_output = Mock()
mock_request_output = Mock()
mock_output = mocker.Mock()
mock_request_output = mocker.Mock()
mock_request_output.outputs = [mock_output]
output = OmniRequestOutput.from_pipeline(
@@ -96,9 +98,9 @@ class TestOmniRequestOutput:
)
assert output.outputs == []
def test_encoder_prompt_token_ids_property(self):
def test_encoder_prompt_token_ids_property(self, mocker: MockerFixture):
"""Test encoder_prompt_token_ids property."""
mock_request_output = Mock()
mock_request_output = mocker.Mock()
mock_request_output.encoder_prompt_token_ids = [10, 20, 30]
output = OmniRequestOutput.from_pipeline(
@@ -109,9 +111,9 @@ class TestOmniRequestOutput:
assert output.encoder_prompt_token_ids == [10, 20, 30]
def test_num_cached_tokens_property(self):
def test_num_cached_tokens_property(self, mocker: MockerFixture):
"""Test num_cached_tokens property."""
mock_request_output = Mock()
mock_request_output = mocker.Mock()
mock_request_output.num_cached_tokens = 42
output = OmniRequestOutput.from_pipeline(
@@ -122,11 +124,12 @@ class TestOmniRequestOutput:
assert output.num_cached_tokens == 42
def test_multimodal_output_property(self):
def test_multimodal_output_property(self, mocker: MockerFixture):
"""Test multimodal_output property."""
mock_request_output = Mock()
mock_audio = Mock()
mock_request_output = mocker.Mock()
mock_audio = mocker.Mock()
expected_output = {"audio": mock_audio}
mock_request_output.outputs = []
mock_request_output.multimodal_output = expected_output
output = OmniRequestOutput.from_pipeline(
@@ -153,9 +156,9 @@ class TestOmniRequestOutput:
assert result["num_images"] == 1
assert result["prompt"] == "a cat"
def test_to_dict_pipeline(self):
def test_to_dict_pipeline(self, mocker: MockerFixture):
"""Test to_dict for pipeline output."""
mock_request_output = Mock()
mock_request_output = mocker.Mock()
mock_request_output.request_id = "pipeline-123"
output = OmniRequestOutput.from_pipeline(
tests/worker/test_process_gpu_memory.py
+71 -77
View File
@@ -1,52 +1,54 @@
"""Tests for process-scoped GPU memory accounting."""
import os
from unittest import mock
import pytest
from pytest_mock import MockerFixture
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestParseCudaVisibleDevices:
def test_empty(self):
def test_empty(self, mocker: MockerFixture):
from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
with mock.patch.dict(os.environ, {}, clear=True):
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
assert parse_cuda_visible_devices() == []
mocker.patch.dict(os.environ, {}, clear=True)
os.environ.pop("CUDA_VISIBLE_DEVICES", None)
assert parse_cuda_visible_devices() == []
with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""}):
assert parse_cuda_visible_devices() == []
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": ""})
assert parse_cuda_visible_devices() == []
def test_integer_indices(self):
def test_integer_indices(self, mocker: MockerFixture):
from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3,5"}):
assert parse_cuda_visible_devices() == [2, 3, 5]
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "2,3,5"})
assert parse_cuda_visible_devices() == [2, 3, 5]
with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0"}):
assert parse_cuda_visible_devices() == [0]
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "0"})
assert parse_cuda_visible_devices() == [0]
def test_uuids(self):
def test_uuids(self, mocker: MockerFixture):
from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
uuid1 = "GPU-12345678-1234-1234-1234-123456789abc"
uuid2 = "GPU-87654321-4321-4321-4321-cba987654321"
with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{uuid1},{uuid2}"}):
assert parse_cuda_visible_devices() == [uuid1, uuid2]
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{uuid1},{uuid2}"})
assert parse_cuda_visible_devices() == [uuid1, uuid2]
def test_mig_ids(self):
def test_mig_ids(self, mocker: MockerFixture):
from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
mig1 = "MIG-GPU-12345678-1234-1234-1234-123456789abc/0/0"
mig2 = "MIG-GPU-12345678-1234-1234-1234-123456789abc/1/0"
with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{mig1},{mig2}"}):
assert parse_cuda_visible_devices() == [mig1, mig2]
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": f"{mig1},{mig2}"})
assert parse_cuda_visible_devices() == [mig1, mig2]
def test_spaces(self):
def test_spaces(self, mocker: MockerFixture):
from vllm_omni.worker.gpu_memory_utils import parse_cuda_visible_devices
with mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": " 2 , 3 , 5 "}):
assert parse_cuda_visible_devices() == [2, 3, 5]
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": " 2 , 3 , 5 "})
assert parse_cuda_visible_devices() == [2, 3, 5]
class TestGetProcessGpuMemory:
@@ -68,87 +70,79 @@ class TestGetProcessGpuMemory:
del tensor
torch.cuda.empty_cache()
def test_raises_on_invalid_device(self):
def test_raises_on_invalid_device(self, mocker: MockerFixture):
from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
with (
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
mock.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=1),
):
with pytest.raises(RuntimeError, match="Invalid GPU device"):
get_process_gpu_memory(5)
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit")
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown")
mocker.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=1)
def test_returns_zero_when_process_not_found(self):
with pytest.raises(RuntimeError, match="Invalid GPU device"):
get_process_gpu_memory(5)
def test_returns_zero_when_process_not_found(self, mocker: MockerFixture):
from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
with (
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
mock.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=8),
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetHandleByIndex"),
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[]),
):
memory = get_process_gpu_memory(0)
assert memory == 0
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit")
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown")
mocker.patch("vllm.third_party.pynvml.nvmlDeviceGetCount", return_value=8)
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetHandleByIndex")
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[])
def test_uses_uuid_when_provided(self):
memory = get_process_gpu_memory(0)
assert memory == 0
def test_uses_uuid_when_provided(self, mocker: MockerFixture):
from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
uuid = "GPU-12345678-1234-1234-1234-123456789abc"
mock_handle = mock.MagicMock()
mock_handle = mocker.MagicMock()
with (
mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid}),
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
mock.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", return_value=mock_handle) as mock_by_uuid,
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[]),
):
memory = get_process_gpu_memory(0)
assert memory == 0
mock_by_uuid.assert_called_once_with(uuid)
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid})
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit")
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown")
mock_by_uuid = mocker.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", return_value=mock_handle)
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlDeviceGetComputeRunningProcesses", return_value=[])
def test_raises_on_invalid_uuid(self):
memory = get_process_gpu_memory(0)
assert memory == 0
mock_by_uuid.assert_called_once_with(uuid)
def test_raises_on_invalid_uuid(self, mocker: MockerFixture):
from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
uuid = "GPU-invalid-uuid"
with (
mock.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid}),
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
mock.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", side_effect=Exception("Invalid UUID")),
):
with pytest.raises(RuntimeError, match="Failed to get NVML handle"):
get_process_gpu_memory(0)
mocker.patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": uuid})
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit")
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown")
mocker.patch("vllm.third_party.pynvml.nvmlDeviceGetHandleByUUID", side_effect=Exception("Invalid UUID"))
def test_returns_none_on_nvml_init_failure(self):
with pytest.raises(RuntimeError, match="Failed to get NVML handle"):
get_process_gpu_memory(0)
def test_returns_none_on_nvml_init_failure(self, mocker: MockerFixture):
from vllm_omni.worker.gpu_memory_utils import get_process_gpu_memory
with (
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit", side_effect=Exception("NVML unavailable")),
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
):
result = get_process_gpu_memory(0)
assert result is None
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit", side_effect=Exception("NVML unavailable"))
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown")
result = get_process_gpu_memory(0)
assert result is None
class TestIsProcessScopedMemoryAvailable:
def test_returns_true_when_nvml_works(self):
def test_returns_true_when_nvml_works(self, mocker: MockerFixture):
from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available
with (
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit"),
mock.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown"),
):
assert is_process_scoped_memory_available() is True
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlInit")
mocker.patch("vllm_omni.worker.gpu_memory_utils.nvmlShutdown")
assert is_process_scoped_memory_available() is True
def test_returns_false_when_nvml_fails(self):
def test_returns_false_when_nvml_fails(self, mocker: MockerFixture):
from vllm_omni.worker.gpu_memory_utils import is_process_scoped_memory_available
with mock.patch(
mocker.patch(
"vllm_omni.worker.gpu_memory_utils.nvmlInit",
side_effect=Exception("NVML unavailable"),
):
assert is_process_scoped_memory_available() is False
)
assert is_process_scoped_memory_available() is False