anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
e58e62fcbef60de8bd275542d8a481c2320bd8a5
vllm-omni/docs/testing/test_design.md
T
hsliu afe5c2692c feat: enhance serving functionality and add comprehensive examples 
- Improve API server with better error handling and response formatting
- Enhance CLI with additional options for DiT stages and configuration
- Add comprehensive examples in examples/basic/ including:
  - API client with health checks and text generation
  - Docker setup and usage examples
  - Simple usage patterns for different scenarios
- Add utility scripts for model downloading and Docker setup
- Update documentation with implementation details and testing guidelines
- Fix configuration validation issues in OmniLLM
- Improve stage configuration handling for AR and DiT stages
- Add proper error handling and fallback mechanisms

Tested with Qwen3-0.6B model:
- Server starts successfully on port 8000
- Health and info endpoints working correctly
- Text generation with various parameters functioning
- API client examples working as expected
- CLI help and configuration options working properly
2025-09-26 15:47:55 +08:00

760 lines
22 KiB
Markdown
Raw Blame History

# vLLM-omni Test Design Document
## 1. Testing Strategy Overview
### 1.1 Testing Pyramid
```
E2E Tests (10%)
/ \
Integration Tests (30%)
/ \
Unit Tests (60%)
```
### 1.2 Test Categories
- **Unit Tests**: Individual component testing with mocked dependencies
- **Integration Tests**: Component interaction testing with real vLLM integration
- **End-to-End Tests**: Full pipeline testing with real models
- **Performance Tests**: Benchmarking and profiling
- **Compatibility Tests**: vLLM version compatibility validation
## 2. Test Structure
```
tests/
├── __init__.py
├── conftest.py # Shared fixtures and configuration
├── unit/ # Unit tests (60%)
│ ├── __init__.py
│ ├── test_config/ # Configuration testing
│ │ ├── test_stage_config.py
│ │ ├── test_dit_config.py
│ │ └── test_cache_config.py
│ ├── test_core/ # Core component testing
│ │ ├── test_omni_llm.py
│ │ ├── test_async_omni_llm.py
│ │ ├── test_stage_manager.py
│ │ └── test_dit_cache_manager.py
│ ├── test_scheduler/ # Scheduler testing
│ │ ├── test_diffusion_scheduler.py
│ │ └── test_scheduler_interface.py
│ ├── test_executor/ # Executor testing
│ │ ├── test_base_executor.py
│ │ └── test_diffusers_executor.py
│ ├── test_model_executor/ # Model runner testing
│ │ ├── test_ar_model_runner.py
│ │ └── test_dit_model_runner.py
│ ├── test_engine/ # Engine testing
│ │ ├── test_output_processor.py
│ │ └── test_multimodal_processor.py
│ ├── test_worker/ # Worker testing
│ │ └── test_omni_worker.py
│ └── test_utils/ # Utility testing
│ ├── test_multimodal.py
│ └── test_vae.py
├── integration/ # Integration tests (30%)
│ ├── __init__.py
│ ├── test_vllm_integration.py
│ ├── test_stage_processing.py
│ ├── test_cli_integration.py
│ └── test_api_compatibility.py
├── e2e/ # End-to-end tests (10%)
│ ├── __init__.py
│ ├── test_full_pipeline.py
│ ├── test_multimodal_generation.py
│ └── test_performance.py
├── benchmarks/ # Performance benchmarks
│ ├── __init__.py
│ ├── test_memory_usage.py
│ ├── test_latency.py
│ └── test_throughput.py
└── fixtures/ # Test data and fixtures
├── __init__.py
├── sample_models/
├── test_images/
└── test_configs/
```
## 3. Unit Test Specifications
### 3.1 Configuration Tests
#### test_stage_config.py
```python
import pytest
from vllm_omni.config.stage_config import OmniStageConfig, DiTConfig, DiTCacheConfig
class TestOmniStageConfig:
def test_ar_stage_config_creation(self):
"""Test AR stage configuration creation"""
config = OmniStageConfig(
stage_id=0,
engine_type="AR",
model_path="test-model",
input_modalities=["text"],
output_modalities=["text"]
)
assert config.engine_type == "AR"
assert config.stage_id == 0
def test_dit_stage_config_creation(self):
"""Test DiT stage configuration creation"""
dit_config = DiTConfig(
model_type="dit",
scheduler_type="ddpm",
num_inference_steps=50
)
config = OmniStageConfig(
stage_id=1,
engine_type="DiT",
model_path="test-dit-model",
input_modalities=["text"],
output_modalities=["image"],
dit_config=dit_config
)
assert config.engine_type == "DiT"
assert config.dit_config is not None
def test_invalid_engine_type(self):
"""Test validation of engine type"""
with pytest.raises(ValueError):
OmniStageConfig(
stage_id=0,
engine_type="INVALID",
model_path="test-model",
input_modalities=["text"],
output_modalities=["text"]
)
```
#### test_dit_config.py
```python
import pytest
from vllm_omni.config.dit_config import DiTConfig
class TestDiTConfig:
def test_dit_config_defaults(self):
"""Test DiT configuration with defaults"""
config = DiTConfig(
model_type="dit",
scheduler_type="ddpm",
num_inference_steps=50
)
assert config.use_diffusers is False
assert config.diffusers_pipeline is None
def test_diffusers_config(self):
"""Test DiT configuration with diffusers"""
config = DiTConfig(
model_type="dit",
scheduler_type="ddpm",
num_inference_steps=50,
use_diffusers=True,
diffusers_pipeline="stable-diffusion"
)
assert config.use_diffusers is True
assert config.diffusers_pipeline == "stable-diffusion"
```
### 3.2 Core Component Tests
#### test_omni_llm.py
```python
import pytest
from unittest.mock import Mock, patch
from vllm_omni.core.omni_llm import OmniLLM
from vllm_omni.config.stage_config import OmniStageConfig
class TestOmniLLM:
@pytest.fixture
def mock_stage_configs(self):
"""Create mock stage configurations"""
return [
OmniStageConfig(
stage_id=0,
engine_type="AR",
model_path="test-ar-model",
input_modalities=["text"],
output_modalities=["text"]
),
OmniStageConfig(
stage_id=1,
engine_type="DiT",
model_path="test-dit-model",
input_modalities=["text"],
output_modalities=["image"]
)
]
@patch('vllm_omni.core.omni_llm.LLMEngine')
def test_omni_llm_initialization(self, mock_llm_engine, mock_stage_configs):
"""Test OmniLLM initialization"""
omni_llm = OmniLLM(mock_stage_configs)
assert len(omni_llm.engine_list) == 2
assert omni_llm.stage_configs == mock_stage_configs
@patch('vllm_omni.core.omni_llm.LLMEngine')
def test_stage_engine_creation(self, mock_llm_engine, mock_stage_configs):
"""Test stage engine creation"""
omni_llm = OmniLLM(mock_stage_configs)
# Verify engines are created for each stage
assert mock_llm_engine.from_vllm_config.call_count == 2
def test_process_stage_inputs_ar(self, mock_stage_configs):
"""Test AR stage input processing"""
omni_llm = OmniLLM(mock_stage_configs)
ar_config = mock_stage_configs[0]
stage_args = {"prompt": "test prompt"}
result = omni_llm._process_stage_inputs(ar_config, stage_args, None)
assert "prompt" in result
def test_process_stage_inputs_dit(self, mock_stage_configs):
"""Test DiT stage input processing"""
omni_llm = OmniLLM(mock_stage_configs)
dit_config = mock_stage_configs[1]
stage_args = {"image": "test_image.jpg"}
with patch.object(omni_llm, 'vae') as mock_vae:
mock_vae.encode.return_value = "encoded_image"
result = omni_llm._process_stage_inputs(dit_config, stage_args, None)
assert "encoded_image" in result
```
#### test_async_omni_llm.py
```python
import pytest
import asyncio
from unittest.mock import Mock, patch, AsyncMock
from vllm_omni.core.omni_llm import AsyncOmniLLM
class TestAsyncOmniLLM:
@pytest.fixture
def mock_stage_configs(self):
"""Create mock stage configurations"""
return [
OmniStageConfig(
stage_id=0,
engine_type="AR",
model_path="test-ar-model",
input_modalities=["text"],
output_modalities=["text"]
)
]
@pytest.mark.asyncio
async def test_async_generation(self, mock_stage_configs):
"""Test async generation"""
with patch('vllm_omni.core.omni_llm.AsyncLLM') as mock_async_llm:
async_omni_llm = AsyncOmniLLM(mock_stage_configs)
# Mock the generate_async method
mock_async_llm.return_value.generate_async = AsyncMock(return_value=[])
result = await async_omni_llm.generate_async([{"prompt": "test"}])
assert result == []
```
### 3.3 Scheduler Tests
#### test_diffusion_scheduler.py
```python
import pytest
from unittest.mock import Mock, patch
from vllm_omni.core.sched.diffusion_scheduler import OmniDiffusionScheduler
class TestOmniDiffusionScheduler:
@pytest.fixture
def mock_configs(self):
"""Create mock configurations"""
return {
'vllm_config': Mock(),
'kv_cache_config': Mock(),
'dit_cache_config': Mock(),
'structured_output_manager': Mock(),
'mm_registry': Mock(),
'include_finished_set': False,
'log_stats': False
}
def test_scheduler_initialization(self, mock_configs):
"""Test scheduler initialization"""
scheduler = OmniDiffusionScheduler(**mock_configs)
assert scheduler.dit_cache_manager is not None
def test_schedule_method(self, mock_configs):
"""Test scheduling method"""
scheduler = OmniDiffusionScheduler(**mock_configs)
with patch.object(scheduler, 'dit_cache_manager') as mock_cache:
mock_cache.allocate_cache.return_value = Mock()
result = scheduler.schedule()
# Verify scheduling logic
assert result is not None
```
### 3.4 Model Runner Tests
#### test_ar_model_runner.py
```python
import pytest
import torch
from unittest.mock import Mock, patch
from vllm_omni.model_executor.ar_model_runner import OmniARModelRunner
class TestOmniARModelRunner:
@pytest.fixture
def mock_runner(self):
"""Create mock model runner"""
with patch('vllm_omni.model_executor.ar_model_runner.GPUModelRunner.__init__'):
runner = OmniARModelRunner(Mock(), Mock(), Mock())
return runner
def test_execute_model_ar(self, mock_runner):
"""Test AR model execution"""
mock_scheduler_output = Mock()
mock_scheduler_output.req_ids = ["req1"]
mock_scheduler_output.req_id_to_index = {"req1": 0}
with patch.object(mock_runner, 'model') as mock_model:
mock_model.forward.return_value = (torch.randn(1, 10, 768), None)
result = mock_runner.execute_model(mock_scheduler_output)
assert result.req_ids == ["req1"]
assert len(result.pooler_output) > 0 # Hidden states
```
#### test_dit_model_runner.py
```python
import pytest
import torch
from unittest.mock import Mock, patch
from vllm_omni.model_executor.dit_model_runner import OmniDiffusionModelRunner
class TestOmniDiffusionModelRunner:
@pytest.fixture
def mock_runner(self):
"""Create mock DiT model runner"""
with patch('vllm_omni.model_executor.dit_model_runner.GPUModelRunner.__init__'):
runner = OmniDiffusionModelRunner(Mock(), Mock(), Mock())
return runner
def test_execute_model_dit(self, mock_runner):
"""Test DiT model execution"""
mock_scheduler_output = Mock()
mock_scheduler_output.req_ids = ["req1"]
mock_scheduler_output.req_id_to_index = {"req1": 0}
with patch.object(mock_runner, 'model') as mock_model:
mock_model.forward.return_value = torch.randn(1, 3, 512, 512)
result = mock_runner.execute_model(mock_scheduler_output)
assert result.req_ids == ["req1"]
assert len(result.pooler_output) > 0 # DiT output tensors
```
## 4. Integration Tests
### 4.1 vLLM Integration Tests
#### test_vllm_integration.py
```python
import pytest
from vllm_omni.core.omni_llm import OmniLLM
from vllm_omni.config.stage_config import OmniStageConfig
class TestVLLMIntegration:
@pytest.mark.integration
def test_vllm_engine_creation(self):
"""Test creation of vLLM engines"""
stage_configs = [
OmniStageConfig(
stage_id=0,
engine_type="AR",
model_path="Qwen/Qwen3-0.6B", # Small test model
input_modalities=["text"],
output_modalities=["text"]
)
]
omni_llm = OmniLLM(stage_configs)
assert len(omni_llm.engine_list) == 1
assert omni_llm.engine_list[0] is not None
@pytest.mark.integration
def test_stage_processing_integration(self):
"""Test integration between stages"""
# Test with real vLLM components
pass
```
### 4.2 CLI Integration Tests
#### test_cli_integration.py
```python
import pytest
import subprocess
import sys
from vllm_omni.entrypoints.cli.main import main
class TestCLIIntegration:
def test_omni_flag_detection(self):
"""Test --omni flag detection"""
original_argv = sys.argv
try:
sys.argv = ["vllm", "serve", "test-model", "--omni"]
# Test that omni command is triggered
with patch('vllm_omni.entrypoints.omni.OmniServeCommand') as mock_serve:
main()
mock_serve.assert_called_once()
finally:
sys.argv = original_argv
def test_forward_to_vllm(self):
"""Test forwarding to vLLM when --omni not present"""
original_argv = sys.argv
try:
sys.argv = ["vllm", "serve", "test-model"]
with patch('vllm.entrypoints.cli.main.main') as mock_vllm_main:
main()
mock_vllm_main.assert_called_once()
finally:
sys.argv = original_argv
```
## 5. End-to-End Tests
### 5.1 Full Pipeline Tests
#### test_full_pipeline.py
```python
import pytest
from vllm_omni.core.omni_llm import OmniLLM
from vllm_omni.config.stage_config import OmniStageConfig, DiTConfig
class TestFullPipeline:
@pytest.mark.e2e
@pytest.mark.slow
def test_ar_to_dit_pipeline(self):
"""Test complete AR to DiT pipeline"""
stage_configs = [
OmniStageConfig(
stage_id=0,
engine_type="AR",
model_path="Qwen/Qwen3-0.6B",
input_modalities=["text"],
output_modalities=["text"]
),
OmniStageConfig(
stage_id=1,
engine_type="DiT",
model_path="stabilityai/stable-diffusion-2-1",
input_modalities=["text"],
output_modalities=["image"],
dit_config=DiTConfig(
model_type="dit",
scheduler_type="ddpm",
num_inference_steps=10 # Reduced for testing
)
)
]
omni_llm = OmniLLM(stage_configs)
stage_args = [
{"prompt": "A beautiful landscape"},
{"prompt": "A beautiful landscape"}
]
result = omni_llm.generate(stage_args)
assert result is not None
assert len(result) > 0
```
### 5.2 Multimodal Generation Tests
#### test_multimodal_generation.py
```python
import pytest
from vllm_omni.core.omni_llm import OmniLLM
class TestMultimodalGeneration:
@pytest.mark.e2e
def test_text_to_image_generation(self):
"""Test text to image generation"""
# Test implementation
pass
@pytest.mark.e2e
def test_image_to_text_generation(self):
"""Test image to text generation"""
# Test implementation
pass
@pytest.mark.e2e
def test_text_and_image_generation(self):
"""Test combined text and image generation"""
# Test implementation
pass
```
## 6. Performance Tests
### 6.1 Memory Usage Tests
#### test_memory_usage.py
```python
import pytest
import psutil
import torch
from vllm_omni.core.omni_llm import OmniLLM
class TestMemoryUsage:
@pytest.mark.benchmark
def test_memory_usage_ar_stage(self):
"""Test memory usage for AR stage"""
process = psutil.Process()
initial_memory = process.memory_info().rss
# Run AR stage
# ... test implementation
final_memory = process.memory_info().rss
memory_increase = final_memory - initial_memory
# Assert memory usage is reasonable
assert memory_increase < 1024 * 1024 * 1024 # Less than 1GB
@pytest.mark.benchmark
def test_memory_usage_dit_stage(self):
"""Test memory usage for DiT stage"""
# Test implementation
pass
@pytest.mark.benchmark
def test_cache_memory_management(self):
"""Test DiT cache memory management"""
# Test implementation
pass
```
### 6.2 Latency Tests
#### test_latency.py
```python
import pytest
import time
from vllm_omni.core.omni_llm import OmniLLM
class TestLatency:
@pytest.mark.benchmark
def test_ar_stage_latency(self):
"""Test AR stage latency"""
start_time = time.time()
# Run AR stage
# ... test implementation
end_time = time.time()
latency = end_time - start_time
# Assert latency is reasonable
assert latency < 5.0 # Less than 5 seconds
@pytest.mark.benchmark
def test_dit_stage_latency(self):
"""Test DiT stage latency"""
# Test implementation
pass
@pytest.mark.benchmark
def test_end_to_end_latency(self):
"""Test end-to-end pipeline latency"""
# Test implementation
pass
```
## 7. Test Configuration
### 7.1 pytest.ini
```ini
[tool:pytest]
testpaths = tests
python_files = test_*.py
python_classes = Test*
python_functions = test_*
addopts =
--strict-markers
--strict-config
--cov=vllm_omni
--cov-report=term-missing
--cov-report=html
--cov-report=xml
--tb=short
markers =
unit: Unit tests
integration: Integration tests
e2e: End-to-end tests
benchmark: Performance benchmark tests
slow: Slow running tests
```
### 7.2 conftest.py
```python
import pytest
import torch
from unittest.mock import Mock
from vllm_omni.config.stage_config import OmniStageConfig, DiTConfig
@pytest.fixture(scope="session")
def device():
"""Get available device for testing"""
return "cuda" if torch.cuda.is_available() else "cpu"
@pytest.fixture
def sample_stage_configs():
"""Sample stage configurations for testing"""
return [
OmniStageConfig(
stage_id=0,
engine_type="AR",
model_path="test-ar-model",
input_modalities=["text"],
output_modalities=["text"]
),
OmniStageConfig(
stage_id=1,
engine_type="DiT",
model_path="test-dit-model",
input_modalities=["text"],
output_modalities=["image"],
dit_config=DiTConfig(
model_type="dit",
scheduler_type="ddpm",
num_inference_steps=10
)
)
]
@pytest.fixture
def mock_vllm_config():
"""Mock vLLM configuration"""
config = Mock()
config.model = "test-model"
config.tensor_parallel_size = 1
config.pipeline_parallel_size = 1
return config
```
## 8. Test Execution
### 8.1 Running Tests
```bash
# Run all tests
pytest
# Run specific test categories
pytest -m unit
pytest -m integration
pytest -m e2e
pytest -m benchmark
# Run with coverage
pytest --cov=vllm_omni --cov-report=html
# Run specific test file
pytest tests/unit/test_omni_llm.py
# Run with verbose output
pytest -v
# Run in parallel
pytest -n auto
```
### 8.2 Continuous Integration
```yaml
# .github/workflows/test.yml
name: Tests
on: [push, pull_request]
jobs:
test:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: [3.8, 3.9, 3.10, 3.11]
steps:
- uses: actions/checkout@v3
- name: Set up Python
uses: actions/setup-python@v4
with:
python-version: ${{ matrix.python-version }}
- name: Install dependencies
run: |
pip install -e ".[dev]"
- name: Run tests
run: |
pytest --cov=vllm_omni --cov-report=xml
- name: Upload coverage
uses: codecov/codecov-action@v3
```
## 9. Test Data Management
### 9.1 Fixtures Directory
```
tests/fixtures/
├── sample_models/
│ ├── ar_model/
│ └── dit_model/
├── test_images/
│ ├── landscape.jpg
│ └── portrait.png
├── test_configs/
│ ├── ar_stage_config.json
│ └── dit_stage_config.json
└── expected_outputs/
├── ar_output.json
└── dit_output.json
```
### 9.2 Test Data Loading
```python
import json
import os
from pathlib import Path
def load_test_config(config_name: str) -> dict:
"""Load test configuration from fixtures"""
config_path = Path(__file__).parent / "fixtures" / "test_configs" / f"{config_name}.json"
with open(config_path) as f:
return json.load(f)
def load_test_image(image_name: str) -> str:
"""Load test image path from fixtures"""
return str(Path(__file__).parent / "fixtures" / "test_images" / image_name)
```
## 10. Test Quality Metrics
### 10.1 Coverage Targets
- **Overall Coverage**: > 90%
- **Unit Tests**: > 95%
- **Integration Tests**: > 80%
- **Critical Paths**: 100%
### 10.2 Performance Targets
- **Unit Test Execution**: < 30 seconds
- **Integration Test Execution**: < 5 minutes
- **E2E Test Execution**: < 15 minutes
- **Memory Usage**: < 8GB for full test suite
### 10.3 Quality Gates
- All unit tests must pass
- Integration tests must pass
- Coverage threshold must be met
- No critical security vulnerabilities
- Performance regression tests must pass
Reference in new issue View Git Blame Copy Permalink