scheduler api design
This commit is contained in:
hsliu committed
1 parent
8dbe6db1be
commit
0c5caf6ab1
5 files changed
+1253
No files matched your search
@@ -0,0 +1,396 @@
|
||||
# API Design Template for vLLM-omni Modules
|
||||
|
||||
This template provides a standardized structure for designing APIs for core, engine, executor, and worker modules in vLLM-omni. Use this template to ensure consistency and completeness across all modules.
|
||||
|
||||
## Template Structure
|
||||
|
||||
### 1. Module Overview
|
||||
- **Purpose**: What this module does
|
||||
- **Responsibilities**: Key responsibilities of the module
|
||||
- **Dependencies**: What other modules this depends on
|
||||
- **Integration Points**: How it integrates with other modules
|
||||
|
||||
### 2. Core Classes/Interfaces
|
||||
- **Base Classes**: Abstract base classes and interfaces
|
||||
- **Implementation Classes**: Concrete implementations
|
||||
- **Data Structures**: Key data structures and models
|
||||
|
||||
### 3. Public API Methods
|
||||
- **Initialization**: Constructor and setup methods
|
||||
- **Core Operations**: Main functionality methods
|
||||
- **Configuration**: Configuration and parameter methods
|
||||
- **Lifecycle Management**: Start, stop, cleanup methods
|
||||
- **Monitoring**: Status, metrics, and debugging methods
|
||||
|
||||
### 4. Configuration
|
||||
- **Configuration Classes**: Dataclasses or config objects
|
||||
- **Required Parameters**: Must-have configuration
|
||||
- **Optional Parameters**: Optional configuration with defaults
|
||||
- **Validation**: Parameter validation rules
|
||||
|
||||
### 5. Error Handling
|
||||
- **Custom Exceptions**: Module-specific exceptions
|
||||
- **Error Codes**: Standardized error codes
|
||||
- **Recovery Strategies**: How to handle and recover from errors
|
||||
|
||||
### 6. Examples
|
||||
- **Basic Usage**: Simple usage examples
|
||||
- **Advanced Usage**: Complex scenarios
|
||||
- **Integration Examples**: How to use with other modules
|
||||
|
||||
---
|
||||
|
||||
## Example: Core Module Template
|
||||
|
||||
### 1. Module Overview
|
||||
|
||||
**Purpose**: The core module provides fundamental scheduling and caching functionality for vLLM-omni.
|
||||
|
||||
**Responsibilities**:
|
||||
- Request scheduling and prioritization
|
||||
- DiT cache management
|
||||
- Resource allocation and coordination
|
||||
- Inter-module communication
|
||||
|
||||
**Dependencies**:
|
||||
- `vllm_omni.request` - Request handling
|
||||
- `vllm_omni.config` - Configuration management
|
||||
- `vllm_omni.utils` - Utility functions
|
||||
|
||||
**Integration Points**:
|
||||
- Receives requests from entrypoints
|
||||
- Coordinates with engine modules
|
||||
- Manages worker allocation
|
||||
- Provides status to monitoring systems
|
||||
|
||||
### 2. Core Classes/Interfaces
|
||||
|
||||
```python
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Optional, Dict, Any
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
|
||||
class SchedulerType(Enum):
|
||||
"""Types of schedulers available."""
|
||||
FIFO = "fifo"
|
||||
PRIORITY = "priority"
|
||||
MULTIMODAL = "multimodal"
|
||||
DIFFUSION = "diffusion"
|
||||
|
||||
@dataclass
|
||||
class SchedulerConfig:
|
||||
"""Configuration for scheduler."""
|
||||
scheduler_type: SchedulerType = SchedulerType.FIFO
|
||||
max_queue_size: int = 1000
|
||||
priority_weights: Dict[str, float] = None
|
||||
timeout_seconds: int = 300
|
||||
|
||||
class BaseScheduler(ABC):
|
||||
"""Abstract base class for all schedulers."""
|
||||
|
||||
@abstractmethod
|
||||
async def schedule_request(self, request: OmniRequest) -> bool:
|
||||
"""Schedule a request for processing."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_next_request(self) -> Optional[OmniRequest]:
|
||||
"""Get the next request to process."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def remove_request(self, request_id: str) -> bool:
|
||||
"""Remove a request from the queue."""
|
||||
pass
|
||||
|
||||
class DiTCacheManager:
|
||||
"""Manages DiT cache for diffusion models."""
|
||||
|
||||
def __init__(self, config: DiTCacheConfig):
|
||||
self.config = config
|
||||
self.cache = {}
|
||||
self.cache_stats = {}
|
||||
|
||||
async def get_cache(self, cache_key: str) -> Optional[Any]:
|
||||
"""Retrieve cached data."""
|
||||
pass
|
||||
|
||||
async def set_cache(self, cache_key: str, data: Any) -> None:
|
||||
"""Store data in cache."""
|
||||
pass
|
||||
|
||||
async def invalidate_cache(self, cache_key: str) -> None:
|
||||
"""Invalidate cached data."""
|
||||
pass
|
||||
```
|
||||
|
||||
### 3. Public API Methods
|
||||
|
||||
#### Initialization
|
||||
```python
|
||||
class CoreScheduler:
|
||||
def __init__(self, config: SchedulerConfig):
|
||||
"""Initialize the core scheduler."""
|
||||
self.config = config
|
||||
self.scheduler = self._create_scheduler()
|
||||
self.cache_manager = DiTCacheManager(config.dit_cache_config)
|
||||
self._running = False
|
||||
|
||||
def _create_scheduler(self) -> BaseScheduler:
|
||||
"""Factory method to create appropriate scheduler."""
|
||||
if self.config.scheduler_type == SchedulerType.FIFO:
|
||||
return FIFOScheduler(self.config)
|
||||
elif self.config.scheduler_type == SchedulerType.PRIORITY:
|
||||
return PriorityScheduler(self.config)
|
||||
# ... other scheduler types
|
||||
```
|
||||
|
||||
#### Core Operations
|
||||
```python
|
||||
async def schedule_request(self, request: OmniRequest) -> bool:
|
||||
"""
|
||||
Schedule a request for processing.
|
||||
|
||||
Args:
|
||||
request: The request to schedule
|
||||
|
||||
Returns:
|
||||
bool: True if successfully scheduled, False otherwise
|
||||
|
||||
Raises:
|
||||
SchedulerError: If scheduling fails
|
||||
QueueFullError: If queue is at capacity
|
||||
"""
|
||||
pass
|
||||
|
||||
async def get_next_request(self) -> Optional[OmniRequest]:
|
||||
"""
|
||||
Get the next request to process.
|
||||
|
||||
Returns:
|
||||
OmniRequest or None: Next request to process
|
||||
|
||||
Raises:
|
||||
SchedulerError: If retrieval fails
|
||||
"""
|
||||
pass
|
||||
|
||||
async def process_request(self, request: OmniRequest) -> Any:
|
||||
"""
|
||||
Process a request through the appropriate pipeline.
|
||||
|
||||
Args:
|
||||
request: The request to process
|
||||
|
||||
Returns:
|
||||
Any: Processing result
|
||||
|
||||
Raises:
|
||||
ProcessingError: If processing fails
|
||||
"""
|
||||
pass
|
||||
```
|
||||
|
||||
#### Configuration
|
||||
```python
|
||||
def update_config(self, new_config: SchedulerConfig) -> None:
|
||||
"""Update scheduler configuration."""
|
||||
pass
|
||||
|
||||
def get_config(self) -> SchedulerConfig:
|
||||
"""Get current configuration."""
|
||||
pass
|
||||
```
|
||||
|
||||
#### Lifecycle Management
|
||||
```python
|
||||
async def start(self) -> None:
|
||||
"""Start the scheduler."""
|
||||
self._running = True
|
||||
# Start background tasks
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the scheduler gracefully."""
|
||||
self._running = False
|
||||
# Cleanup resources
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
"""Force shutdown the scheduler."""
|
||||
# Immediate cleanup
|
||||
```
|
||||
|
||||
#### Monitoring
|
||||
```python
|
||||
def get_status(self) -> Dict[str, Any]:
|
||||
"""Get current scheduler status."""
|
||||
return {
|
||||
"running": self._running,
|
||||
"queue_size": self.scheduler.queue_size(),
|
||||
"processed_requests": self.scheduler.processed_count(),
|
||||
"cache_hit_rate": self.cache_manager.get_hit_rate()
|
||||
}
|
||||
|
||||
def get_metrics(self) -> Dict[str, float]:
|
||||
"""Get performance metrics."""
|
||||
pass
|
||||
```
|
||||
|
||||
### 4. Configuration
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class CoreModuleConfig:
|
||||
"""Configuration for the core module."""
|
||||
|
||||
# Scheduler configuration
|
||||
scheduler: SchedulerConfig = field(default_factory=SchedulerConfig)
|
||||
|
||||
# Cache configuration
|
||||
dit_cache: DiTCacheConfig = field(default_factory=DiTCacheConfig)
|
||||
|
||||
# Resource limits
|
||||
max_memory_gb: float = 16.0
|
||||
max_gpu_utilization: float = 0.8
|
||||
|
||||
# Timeouts
|
||||
request_timeout: int = 300
|
||||
worker_timeout: int = 60
|
||||
|
||||
def validate(self) -> None:
|
||||
"""Validate configuration parameters."""
|
||||
if self.max_memory_gb <= 0:
|
||||
raise ValueError("max_memory_gb must be positive")
|
||||
if not 0 < self.max_gpu_utilization <= 1:
|
||||
raise ValueError("max_gpu_utilization must be between 0 and 1")
|
||||
```
|
||||
|
||||
### 5. Error Handling
|
||||
|
||||
```python
|
||||
class CoreModuleError(Exception):
|
||||
"""Base exception for core module errors."""
|
||||
pass
|
||||
|
||||
class SchedulerError(CoreModuleError):
|
||||
"""Scheduler-related errors."""
|
||||
pass
|
||||
|
||||
class QueueFullError(SchedulerError):
|
||||
"""Queue is at capacity."""
|
||||
pass
|
||||
|
||||
class CacheError(CoreModuleError):
|
||||
"""Cache-related errors."""
|
||||
pass
|
||||
|
||||
class ResourceError(CoreModuleError):
|
||||
"""Resource allocation errors."""
|
||||
pass
|
||||
```
|
||||
|
||||
### 6. Examples
|
||||
|
||||
#### Basic Usage
|
||||
```python
|
||||
from vllm_omni.core import CoreScheduler, CoreModuleConfig
|
||||
from vllm_omni.request import create_text_request
|
||||
|
||||
# Create configuration
|
||||
config = CoreModuleConfig(
|
||||
scheduler=SchedulerConfig(scheduler_type=SchedulerType.FIFO),
|
||||
max_memory_gb=8.0
|
||||
)
|
||||
|
||||
# Initialize scheduler
|
||||
scheduler = CoreScheduler(config)
|
||||
await scheduler.start()
|
||||
|
||||
# Create and schedule a request
|
||||
request = create_text_request(
|
||||
request_id="req_001",
|
||||
prompt="Hello, world!",
|
||||
sampling_params=sampling_params
|
||||
)
|
||||
|
||||
success = await scheduler.schedule_request(request)
|
||||
if success:
|
||||
result = await scheduler.process_request(request)
|
||||
print(f"Result: {result}")
|
||||
```
|
||||
|
||||
#### Advanced Usage
|
||||
```python
|
||||
# Custom scheduler with priority handling
|
||||
config = CoreModuleConfig(
|
||||
scheduler=SchedulerConfig(
|
||||
scheduler_type=SchedulerType.PRIORITY,
|
||||
priority_weights={"high": 2.0, "normal": 1.0, "low": 0.5}
|
||||
)
|
||||
)
|
||||
|
||||
scheduler = CoreScheduler(config)
|
||||
await scheduler.start()
|
||||
|
||||
# Monitor scheduler status
|
||||
status = scheduler.get_status()
|
||||
print(f"Queue size: {status['queue_size']}")
|
||||
print(f"Cache hit rate: {status['cache_hit_rate']:.2%}")
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Module-Specific Guidelines
|
||||
|
||||
### Core Module
|
||||
- Focus on scheduling, caching, and resource management
|
||||
- Provide clear interfaces for other modules
|
||||
- Handle concurrency and thread safety
|
||||
- Implement comprehensive monitoring
|
||||
|
||||
### Engine Module
|
||||
- Handle model loading and inference
|
||||
- Support both AR and diffusion models
|
||||
- Provide unified interface for different model types
|
||||
- Implement efficient memory management
|
||||
|
||||
### Executor Module
|
||||
- Coordinate between different engines
|
||||
- Handle request routing and load balancing
|
||||
- Manage execution pipelines
|
||||
- Provide error recovery mechanisms
|
||||
|
||||
### Worker Module
|
||||
- Handle actual model execution
|
||||
- Manage GPU resources
|
||||
- Implement batching strategies
|
||||
- Provide performance optimization
|
||||
|
||||
---
|
||||
|
||||
## Checklist for API Design
|
||||
|
||||
- [ ] **Clear Purpose**: Module purpose is well-defined
|
||||
- [ ] **Complete Interface**: All public methods are documented
|
||||
- [ ] **Error Handling**: Comprehensive error handling strategy
|
||||
- [ ] **Configuration**: Flexible configuration system
|
||||
- [ ] **Examples**: Basic and advanced usage examples
|
||||
- [ ] **Type Hints**: All methods have proper type hints
|
||||
- [ ] **Documentation**: Comprehensive docstrings
|
||||
- [ ] **Testing**: Unit tests for all public methods
|
||||
- [ ] **Integration**: Clear integration points with other modules
|
||||
- [ ] **Performance**: Performance considerations documented
|
||||
|
||||
---
|
||||
|
||||
## Submission Guidelines
|
||||
|
||||
1. Create a new file: `docs/api/[module_name]_api.md`
|
||||
2. Follow the template structure exactly
|
||||
3. Include all required sections
|
||||
4. Provide working code examples
|
||||
5. Ensure all methods have proper docstrings
|
||||
6. Include error handling strategies
|
||||
7. Submit for review before implementation
|
||||
|
||||
This template ensures consistency and completeness across all vLLM-omni modules.
|
||||
@@ -0,0 +1,235 @@
|
||||
# vLLM-omni API Documentation
|
||||
|
||||
This directory contains comprehensive API design documentation for all core modules in vLLM-omni. These templates provide a standardized structure for designing and implementing the core, engine, executor, and worker modules.
|
||||
|
||||
## 📋 API Design Templates
|
||||
|
||||
### 1. [API Design Template](API_DESIGN_TEMPLATE.md)
|
||||
**Master template** for creating API documentation for any module in vLLM-omni. Use this as a starting point for new modules.
|
||||
|
||||
### 2. [Core Module API](core_module_api.md)
|
||||
**Core module** provides fundamental scheduling, caching, and resource management functionality.
|
||||
|
||||
**Key Components:**
|
||||
- Request scheduling and prioritization
|
||||
- DiT cache management for diffusion models
|
||||
- Resource allocation and coordination
|
||||
- Inter-module communication
|
||||
|
||||
### 3. [Engine Module API](engine_module_api.md)
|
||||
**Engine module** handles model loading, inference execution, and output processing.
|
||||
|
||||
**Key Components:**
|
||||
- Model loading and initialization
|
||||
- Inference execution for AR and diffusion models
|
||||
- Input preprocessing for various modalities
|
||||
- Output postprocessing and formatting
|
||||
|
||||
### 4. [Executor Module API](executor_module_api.md)
|
||||
**Executor module** coordinates and manages request execution across different engines and workers.
|
||||
|
||||
**Key Components:**
|
||||
- Request routing and load balancing
|
||||
- Execution pipeline coordination
|
||||
- Worker management and task distribution
|
||||
- Error handling and recovery
|
||||
|
||||
### 5. [Worker Module API](worker_module_api.md)
|
||||
**Worker module** provides the actual execution environment for model inference.
|
||||
|
||||
**Key Components:**
|
||||
- Model execution and inference
|
||||
- GPU resource management
|
||||
- Request batching and processing
|
||||
- Performance optimization
|
||||
|
||||
## 🎯 How to Use These Templates
|
||||
|
||||
### For New Module Development
|
||||
|
||||
1. **Start with the Master Template**
|
||||
- Copy `API_DESIGN_TEMPLATE.md`
|
||||
- Rename to `[module_name]_api.md`
|
||||
- Follow the structure exactly
|
||||
|
||||
2. **Fill in Module-Specific Details**
|
||||
- Update the module overview
|
||||
- Define core classes and interfaces
|
||||
- Specify public API methods
|
||||
- Add configuration options
|
||||
- Define error handling strategy
|
||||
- Provide usage examples
|
||||
|
||||
3. **Review and Validate**
|
||||
- Ensure all sections are complete
|
||||
- Verify code examples work
|
||||
- Check for consistency with other modules
|
||||
- Validate configuration options
|
||||
|
||||
### For Existing Module Updates
|
||||
|
||||
1. **Update API Documentation**
|
||||
- Modify existing API files
|
||||
- Add new methods and classes
|
||||
- Update examples and configuration
|
||||
- Maintain backward compatibility notes
|
||||
|
||||
2. **Version Control**
|
||||
- Track changes in git
|
||||
- Use clear commit messages
|
||||
- Tag major API changes
|
||||
|
||||
## 📚 Template Structure
|
||||
|
||||
Each API template follows this standardized structure:
|
||||
|
||||
### 1. Module Overview
|
||||
- **Purpose**: What the module does
|
||||
- **Responsibilities**: Key responsibilities
|
||||
- **Dependencies**: Required modules
|
||||
- **Integration Points**: How it connects with other modules
|
||||
|
||||
### 2. Core Classes/Interfaces
|
||||
- **Base Classes**: Abstract base classes
|
||||
- **Implementation Classes**: Concrete implementations
|
||||
- **Data Structures**: Key data models
|
||||
|
||||
### 3. Public API Methods
|
||||
- **Initialization**: Constructor and setup
|
||||
- **Core Operations**: Main functionality
|
||||
- **Configuration**: Settings management
|
||||
- **Lifecycle Management**: Start/stop/cleanup
|
||||
- **Monitoring**: Status and metrics
|
||||
|
||||
### 4. Configuration
|
||||
- **Configuration Classes**: Dataclasses for settings
|
||||
- **Required Parameters**: Must-have settings
|
||||
- **Optional Parameters**: Optional settings with defaults
|
||||
- **Validation**: Parameter validation rules
|
||||
|
||||
### 5. Error Handling
|
||||
- **Custom Exceptions**: Module-specific errors
|
||||
- **Error Codes**: Standardized error codes
|
||||
- **Recovery Strategies**: Error recovery approaches
|
||||
|
||||
### 6. Examples
|
||||
- **Basic Usage**: Simple usage examples
|
||||
- **Advanced Usage**: Complex scenarios
|
||||
- **Integration Examples**: Multi-module usage
|
||||
|
||||
## 🔧 Implementation Guidelines
|
||||
|
||||
### Code Standards
|
||||
- Use type hints for all methods
|
||||
- Include comprehensive docstrings
|
||||
- Follow PEP 8 style guidelines
|
||||
- Use async/await for I/O operations
|
||||
|
||||
### Error Handling
|
||||
- Define specific exception types
|
||||
- Provide meaningful error messages
|
||||
- Include error recovery strategies
|
||||
- Log errors appropriately
|
||||
|
||||
### Configuration
|
||||
- Use dataclasses for configuration
|
||||
- Provide sensible defaults
|
||||
- Include validation methods
|
||||
- Support environment variables
|
||||
|
||||
### Testing
|
||||
- Write unit tests for all public methods
|
||||
- Include integration tests
|
||||
- Test error conditions
|
||||
- Validate configuration options
|
||||
|
||||
## 🚀 Getting Started
|
||||
|
||||
### For Developers
|
||||
|
||||
1. **Choose a Module**
|
||||
- Pick the module you want to work on
|
||||
- Read the corresponding API documentation
|
||||
- Understand the responsibilities and interfaces
|
||||
|
||||
2. **Set Up Development Environment**
|
||||
```bash
|
||||
# Clone the repository
|
||||
git clone https://github.com/hsliuustc0106/vllm-omni.git
|
||||
cd vllm-omni
|
||||
|
||||
# Install dependencies
|
||||
pip install -r requirements-dev.txt
|
||||
|
||||
# Set up pre-commit hooks
|
||||
pre-commit install
|
||||
```
|
||||
|
||||
3. **Start Implementation**
|
||||
- Create the module directory
|
||||
- Implement base classes first
|
||||
- Add concrete implementations
|
||||
- Write tests as you go
|
||||
|
||||
### For Contributors
|
||||
|
||||
1. **Review API Documentation**
|
||||
- Read through all module APIs
|
||||
- Understand the overall architecture
|
||||
- Identify areas for improvement
|
||||
|
||||
2. **Propose Changes**
|
||||
- Create issues for API changes
|
||||
- Discuss in pull requests
|
||||
- Update documentation accordingly
|
||||
|
||||
3. **Submit Contributions**
|
||||
- Follow the coding standards
|
||||
- Include tests for new features
|
||||
- Update documentation
|
||||
- Submit pull requests
|
||||
|
||||
## 📖 Additional Resources
|
||||
|
||||
### Architecture Overview
|
||||
- [vLLM-omni Architecture Design](../architecture/vLLM-omni%20arch%20design%20doc.md)
|
||||
- [Component Categorization](../../ARCHITECTURE_CATEGORIZATION.md)
|
||||
|
||||
### Development Guides
|
||||
- [Development Setup](../../README.md#development)
|
||||
- [Testing Guidelines](../../tests/README.md)
|
||||
- [Contributing Guidelines](../../CONTRIBUTING.md)
|
||||
|
||||
### Examples
|
||||
- [Basic Examples](../../examples/basic/)
|
||||
- [Advanced Examples](../../examples/advanced/)
|
||||
- [Multimodal Examples](../../examples/multimodal/)
|
||||
|
||||
## 🤝 Contributing
|
||||
|
||||
We welcome contributions to improve the API documentation and implementation. Please:
|
||||
|
||||
1. **Follow the Template Structure**
|
||||
- Use the master template as a guide
|
||||
- Maintain consistency across modules
|
||||
- Include all required sections
|
||||
|
||||
2. **Provide Working Examples**
|
||||
- Test all code examples
|
||||
- Include both basic and advanced usage
|
||||
- Show integration patterns
|
||||
|
||||
3. **Keep Documentation Updated**
|
||||
- Update docs when APIs change
|
||||
- Version control documentation changes
|
||||
- Maintain backward compatibility notes
|
||||
|
||||
## 📝 Notes
|
||||
|
||||
- All API documentation should be kept up-to-date with implementation
|
||||
- Code examples should be tested and working
|
||||
- Configuration options should be validated
|
||||
- Error handling should be comprehensive
|
||||
- Performance considerations should be documented
|
||||
|
||||
For questions or suggestions about the API documentation, please open an issue or start a discussion in the repository.
|
||||
@@ -0,0 +1,580 @@
|
||||
# Core Module API Design
|
||||
|
||||
## 1. Module Overview
|
||||
|
||||
**Purpose**: The core module provides fundamental scheduling, caching, and resource management functionality for vLLM-omni.
|
||||
|
||||
**Responsibilities**:
|
||||
- Request scheduling and prioritization across different model types
|
||||
- DiT cache management for diffusion models
|
||||
- Resource allocation and coordination between workers
|
||||
- Inter-module communication and coordination
|
||||
- Queue management and load balancing
|
||||
|
||||
**Dependencies**:
|
||||
- `vllm_omni.request` - Request handling and types
|
||||
- `vllm_omni.config` - Configuration management
|
||||
- `vllm_omni.utils` - Utility functions and helpers
|
||||
- `vllm` - vLLM core components for integration
|
||||
|
||||
**Integration Points**:
|
||||
- Receives requests from entrypoints
|
||||
- Coordinates with engine modules for model execution
|
||||
- Manages worker allocation and task distribution
|
||||
- Provides status and metrics to monitoring systems
|
||||
- Handles cache coordination between different model types
|
||||
|
||||
## 2. Core Classes/Interfaces
|
||||
|
||||
```python
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Optional, Dict, Any, Union
|
||||
from dataclasses import dataclass, field
|
||||
from enum import Enum
|
||||
import asyncio
|
||||
from vllm_omni.request import OmniRequest, RequestType, ProcessingStage
|
||||
|
||||
class SchedulerType(Enum):
|
||||
"""Types of schedulers available."""
|
||||
FIFO = "fifo"
|
||||
PRIORITY = "priority"
|
||||
MULTIMODAL = "multimodal"
|
||||
DIFFUSION = "diffusion"
|
||||
HYBRID = "hybrid"
|
||||
|
||||
class CacheStrategy(Enum):
|
||||
"""Cache strategies for different model types."""
|
||||
LRU = "lru"
|
||||
LFU = "lfu"
|
||||
TTL = "ttl"
|
||||
ADAPTIVE = "adaptive"
|
||||
|
||||
@dataclass
|
||||
class SchedulerConfig:
|
||||
"""Configuration for scheduler."""
|
||||
scheduler_type: SchedulerType = SchedulerType.FIFO
|
||||
max_queue_size: int = 1000
|
||||
priority_weights: Dict[str, float] = field(default_factory=lambda: {
|
||||
"high": 2.0, "normal": 1.0, "low": 0.5
|
||||
})
|
||||
timeout_seconds: int = 300
|
||||
batch_size: int = 32
|
||||
enable_preemption: bool = True
|
||||
|
||||
@dataclass
|
||||
class CacheConfig:
|
||||
"""Configuration for cache management."""
|
||||
max_memory_gb: float = 8.0
|
||||
strategy: CacheStrategy = CacheStrategy.LRU
|
||||
ttl_seconds: int = 3600
|
||||
enable_compression: bool = True
|
||||
compression_ratio: float = 0.7
|
||||
|
||||
class BaseScheduler(ABC):
|
||||
"""Abstract base class for all schedulers."""
|
||||
|
||||
@abstractmethod
|
||||
async def schedule_request(self, request: OmniRequest) -> bool:
|
||||
"""Schedule a request for processing."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def get_next_request(self) -> Optional[OmniRequest]:
|
||||
"""Get the next request to process."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def remove_request(self, request_id: str) -> bool:
|
||||
"""Remove a request from the queue."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_queue_size(self) -> int:
|
||||
"""Get current queue size."""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_processed_count(self) -> int:
|
||||
"""Get number of processed requests."""
|
||||
pass
|
||||
|
||||
class DiTCacheManager:
|
||||
"""Manages DiT cache for diffusion models."""
|
||||
|
||||
def __init__(self, config: CacheConfig):
|
||||
self.config = config
|
||||
self.cache = {}
|
||||
self.cache_stats = {"hits": 0, "misses": 0}
|
||||
self.memory_usage = 0.0
|
||||
|
||||
async def get_cache(self, cache_key: str) -> Optional[Any]:
|
||||
"""Retrieve cached data."""
|
||||
pass
|
||||
|
||||
async def set_cache(self, cache_key: str, data: Any) -> None:
|
||||
"""Store data in cache."""
|
||||
pass
|
||||
|
||||
async def invalidate_cache(self, cache_key: str) -> None:
|
||||
"""Invalidate cached data."""
|
||||
pass
|
||||
|
||||
def get_hit_rate(self) -> float:
|
||||
"""Get cache hit rate."""
|
||||
pass
|
||||
|
||||
def get_memory_usage(self) -> float:
|
||||
"""Get current memory usage in GB."""
|
||||
pass
|
||||
|
||||
class ResourceManager:
|
||||
"""Manages system resources and allocation."""
|
||||
|
||||
def __init__(self, config: Dict[str, Any]):
|
||||
self.config = config
|
||||
self.allocated_resources = {}
|
||||
self.available_resources = {}
|
||||
|
||||
async def allocate_resources(self, request: OmniRequest) -> Dict[str, Any]:
|
||||
"""Allocate resources for a request."""
|
||||
pass
|
||||
|
||||
async def deallocate_resources(self, request_id: str) -> None:
|
||||
"""Deallocate resources for a request."""
|
||||
pass
|
||||
|
||||
def get_resource_status(self) -> Dict[str, Any]:
|
||||
"""Get current resource allocation status."""
|
||||
pass
|
||||
```
|
||||
|
||||
## 3. Public API Methods
|
||||
|
||||
### CoreScheduler Class
|
||||
|
||||
```python
|
||||
class OmniScheduler(SchedulerInterface):
|
||||
"""Main scheduler for vLLM-omni core module."""
|
||||
|
||||
def __init__(self, config: SchedulerConfig, cache_config: CacheConfig):
|
||||
"""
|
||||
Initialize the core scheduler.
|
||||
|
||||
Args:
|
||||
config: Scheduler configuration
|
||||
cache_config: Cache configuration
|
||||
"""
|
||||
self.config = config
|
||||
self.cache_manager = DiTCacheManager(cache_config)
|
||||
self.resource_manager = ResourceManager({})
|
||||
self.scheduler = self._create_scheduler()
|
||||
self._running = False
|
||||
self._stats = {
|
||||
"total_requests": 0,
|
||||
"processed_requests": 0,
|
||||
"failed_requests": 0,
|
||||
"average_processing_time": 0.0
|
||||
}
|
||||
|
||||
def _create_scheduler(self) -> BaseScheduler:
|
||||
"""Factory method to create appropriate scheduler."""
|
||||
if self.config.scheduler_type == SchedulerType.FIFO:
|
||||
return FIFOScheduler(self.config)
|
||||
elif self.config.scheduler_type == SchedulerType.PRIORITY:
|
||||
return PriorityScheduler(self.config)
|
||||
elif self.config.scheduler_type == SchedulerType.MULTIMODAL:
|
||||
return MultimodalScheduler(self.config)
|
||||
elif self.config.scheduler_type == SchedulerType.DIFFUSION:
|
||||
return DiffusionScheduler(self.config)
|
||||
elif self.config.scheduler_type == SchedulerType.HYBRID:
|
||||
return HybridScheduler(self.config)
|
||||
else:
|
||||
raise ValueError(f"Unknown scheduler type: {self.config.scheduler_type}")
|
||||
```
|
||||
|
||||
#### Core Operations
|
||||
|
||||
```python
|
||||
async def schedule_request(self, request: OmniRequest) -> bool:
|
||||
"""
|
||||
Schedule a request for processing.
|
||||
|
||||
Args:
|
||||
request: The request to schedule
|
||||
|
||||
Returns:
|
||||
bool: True if successfully scheduled, False otherwise
|
||||
|
||||
Raises:
|
||||
SchedulerError: If scheduling fails
|
||||
QueueFullError: If queue is at capacity
|
||||
ResourceError: If insufficient resources
|
||||
"""
|
||||
try:
|
||||
# Validate request
|
||||
if not self._validate_request(request):
|
||||
return False
|
||||
|
||||
# Check resource availability
|
||||
if not await self.resource_manager.can_allocate(request):
|
||||
raise ResourceError("Insufficient resources")
|
||||
|
||||
# Schedule the request
|
||||
success = await self.scheduler.schedule_request(request)
|
||||
if success:
|
||||
self._stats["total_requests"] += 1
|
||||
await self.resource_manager.allocate_resources(request)
|
||||
|
||||
return success
|
||||
|
||||
except Exception as e:
|
||||
self._stats["failed_requests"] += 1
|
||||
raise SchedulerError(f"Failed to schedule request: {e}")
|
||||
|
||||
async def get_next_request(self) -> Optional[OmniRequest]:
|
||||
"""
|
||||
Get the next request to process.
|
||||
|
||||
Returns:
|
||||
OmniRequest or None: Next request to process
|
||||
|
||||
Raises:
|
||||
SchedulerError: If retrieval fails
|
||||
"""
|
||||
try:
|
||||
request = await self.scheduler.get_next_request()
|
||||
if request:
|
||||
request.update_processing_stage(ProcessingStage.PREPROCESSING)
|
||||
return request
|
||||
except Exception as e:
|
||||
raise SchedulerError(f"Failed to get next request: {e}")
|
||||
|
||||
async def process_request(self, request: OmniRequest) -> Any:
|
||||
"""
|
||||
Process a request through the appropriate pipeline.
|
||||
|
||||
Args:
|
||||
request: The request to process
|
||||
|
||||
Returns:
|
||||
Any: Processing result
|
||||
|
||||
Raises:
|
||||
ProcessingError: If processing fails
|
||||
"""
|
||||
start_time = time.time()
|
||||
try:
|
||||
# Update processing stage
|
||||
request.update_processing_stage(ProcessingStage.AR_GENERATION)
|
||||
|
||||
# Check cache first
|
||||
cache_key = request.generate_cache_key()
|
||||
cached_result = await self.cache_manager.get_cache(cache_key)
|
||||
if cached_result:
|
||||
return cached_result
|
||||
|
||||
# Process based on request type
|
||||
if request.request_type == RequestType.TEXT:
|
||||
result = await self._process_text_request(request)
|
||||
elif request.request_type == RequestType.IMAGE:
|
||||
result = await self._process_image_request(request)
|
||||
elif request.request_type == RequestType.MULTIMODAL:
|
||||
result = await self._process_multimodal_request(request)
|
||||
elif request.request_type == RequestType.DIFFUSION:
|
||||
result = await self._process_diffusion_request(request)
|
||||
else:
|
||||
raise ProcessingError(f"Unknown request type: {request.request_type}")
|
||||
|
||||
# Cache the result
|
||||
await self.cache_manager.set_cache(cache_key, result)
|
||||
|
||||
# Update statistics
|
||||
processing_time = time.time() - start_time
|
||||
self._update_stats(processing_time)
|
||||
|
||||
request.update_processing_stage(ProcessingStage.COMPLETED)
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
request.add_error(str(e))
|
||||
self._stats["failed_requests"] += 1
|
||||
raise ProcessingError(f"Failed to process request: {e}")
|
||||
finally:
|
||||
# Cleanup resources
|
||||
await self.resource_manager.deallocate_resources(request.request_id)
|
||||
```
|
||||
|
||||
#### Configuration Methods
|
||||
|
||||
```python
|
||||
def update_config(self, new_config: SchedulerConfig) -> None:
|
||||
"""
|
||||
Update scheduler configuration.
|
||||
|
||||
Args:
|
||||
new_config: New configuration to apply
|
||||
|
||||
Raises:
|
||||
ConfigError: If configuration is invalid
|
||||
"""
|
||||
try:
|
||||
new_config.validate()
|
||||
self.config = new_config
|
||||
# Recreate scheduler with new config
|
||||
self.scheduler = self._create_scheduler()
|
||||
except Exception as e:
|
||||
raise ConfigError(f"Failed to update config: {e}")
|
||||
|
||||
def get_config(self) -> SchedulerConfig:
|
||||
"""Get current configuration."""
|
||||
return self.config
|
||||
|
||||
def update_cache_config(self, new_config: CacheConfig) -> None:
|
||||
"""Update cache configuration."""
|
||||
self.cache_manager.config = new_config
|
||||
```
|
||||
|
||||
#### Lifecycle Management
|
||||
|
||||
```python
|
||||
async def start(self) -> None:
|
||||
"""Start the scheduler."""
|
||||
if self._running:
|
||||
raise SchedulerError("Scheduler is already running")
|
||||
|
||||
self._running = True
|
||||
# Start background tasks
|
||||
asyncio.create_task(self._background_cleanup())
|
||||
asyncio.create_task(self._monitor_resources())
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the scheduler gracefully."""
|
||||
if not self._running:
|
||||
return
|
||||
|
||||
self._running = False
|
||||
# Wait for current requests to complete
|
||||
await self._wait_for_completion()
|
||||
# Cleanup resources
|
||||
await self._cleanup_resources()
|
||||
|
||||
async def shutdown(self) -> None:
|
||||
"""Force shutdown the scheduler."""
|
||||
self._running = False
|
||||
# Immediate cleanup
|
||||
await self._force_cleanup()
|
||||
```
|
||||
|
||||
#### Monitoring Methods
|
||||
|
||||
```python
|
||||
def get_status(self) -> Dict[str, Any]:
|
||||
"""Get current scheduler status."""
|
||||
return {
|
||||
"running": self._running,
|
||||
"queue_size": self.scheduler.get_queue_size(),
|
||||
"processed_requests": self.scheduler.get_processed_count(),
|
||||
"cache_hit_rate": self.cache_manager.get_hit_rate(),
|
||||
"memory_usage": self.cache_manager.get_memory_usage(),
|
||||
"resource_status": self.resource_manager.get_resource_status(),
|
||||
"stats": self._stats.copy()
|
||||
}
|
||||
|
||||
def get_metrics(self) -> Dict[str, float]:
|
||||
"""Get performance metrics."""
|
||||
return {
|
||||
"throughput": self._calculate_throughput(),
|
||||
"average_latency": self._stats["average_processing_time"],
|
||||
"cache_hit_rate": self.cache_manager.get_hit_rate(),
|
||||
"memory_utilization": self.cache_manager.get_memory_usage() / self.cache_manager.config.max_memory_gb,
|
||||
"queue_utilization": self.scheduler.get_queue_size() / self.config.max_queue_size
|
||||
}
|
||||
|
||||
def get_health_status(self) -> Dict[str, Any]:
|
||||
"""Get health status of the scheduler."""
|
||||
return {
|
||||
"healthy": self._is_healthy(),
|
||||
"issues": self._get_health_issues(),
|
||||
"recommendations": self._get_recommendations()
|
||||
}
|
||||
```
|
||||
|
||||
## 4. Configuration
|
||||
|
||||
```python
|
||||
@dataclass
|
||||
class CoreModuleConfig:
|
||||
"""Complete configuration for the core module."""
|
||||
|
||||
# Scheduler configuration
|
||||
scheduler: SchedulerConfig = field(default_factory=SchedulerConfig)
|
||||
|
||||
# Cache configuration
|
||||
cache: CacheConfig = field(default_factory=CacheConfig)
|
||||
|
||||
# Resource limits
|
||||
max_memory_gb: float = 16.0
|
||||
max_gpu_utilization: float = 0.8
|
||||
max_cpu_utilization: float = 0.9
|
||||
|
||||
# Timeouts
|
||||
request_timeout: int = 300
|
||||
worker_timeout: int = 60
|
||||
cleanup_interval: int = 30
|
||||
|
||||
# Monitoring
|
||||
enable_metrics: bool = True
|
||||
metrics_interval: int = 10
|
||||
enable_health_checks: bool = True
|
||||
|
||||
def validate(self) -> None:
|
||||
"""Validate configuration parameters."""
|
||||
if self.max_memory_gb <= 0:
|
||||
raise ValueError("max_memory_gb must be positive")
|
||||
if not 0 < self.max_gpu_utilization <= 1:
|
||||
raise ValueError("max_gpu_utilization must be between 0 and 1")
|
||||
if not 0 < self.max_cpu_utilization <= 1:
|
||||
raise ValueError("max_cpu_utilization must be between 0 and 1")
|
||||
if self.request_timeout <= 0:
|
||||
raise ValueError("request_timeout must be positive")
|
||||
```
|
||||
|
||||
## 5. Error Handling
|
||||
|
||||
```python
|
||||
class CoreModuleError(Exception):
|
||||
"""Base exception for core module errors."""
|
||||
pass
|
||||
|
||||
class SchedulerError(CoreModuleError):
|
||||
"""Scheduler-related errors."""
|
||||
pass
|
||||
|
||||
class QueueFullError(SchedulerError):
|
||||
"""Queue is at capacity."""
|
||||
pass
|
||||
|
||||
class CacheError(CoreModuleError):
|
||||
"""Cache-related errors."""
|
||||
pass
|
||||
|
||||
class ResourceError(CoreModuleError):
|
||||
"""Resource allocation errors."""
|
||||
pass
|
||||
|
||||
class ProcessingError(CoreModuleError):
|
||||
"""Request processing errors."""
|
||||
pass
|
||||
|
||||
class ConfigError(CoreModuleError):
|
||||
"""Configuration errors."""
|
||||
pass
|
||||
```
|
||||
|
||||
## 6. Examples
|
||||
|
||||
### Basic Usage
|
||||
|
||||
```python
|
||||
from vllm_omni.core import CoreScheduler, CoreModuleConfig
|
||||
from vllm_omni.request import create_text_request
|
||||
|
||||
# Create configuration
|
||||
config = CoreModuleConfig(
|
||||
scheduler=SchedulerConfig(scheduler_type=SchedulerType.FIFO),
|
||||
cache=CacheConfig(max_memory_gb=8.0)
|
||||
)
|
||||
|
||||
# Initialize scheduler
|
||||
scheduler = CoreScheduler(config.scheduler, config.cache)
|
||||
await scheduler.start()
|
||||
|
||||
# Create and schedule a request
|
||||
request = create_text_request(
|
||||
request_id="req_001",
|
||||
prompt="Hello, world!",
|
||||
sampling_params=sampling_params
|
||||
)
|
||||
|
||||
success = await scheduler.schedule_request(request)
|
||||
if success:
|
||||
result = await scheduler.process_request(request)
|
||||
print(f"Result: {result}")
|
||||
|
||||
# Get status
|
||||
status = scheduler.get_status()
|
||||
print(f"Queue size: {status['queue_size']}")
|
||||
print(f"Cache hit rate: {status['cache_hit_rate']:.2%}")
|
||||
|
||||
await scheduler.stop()
|
||||
```
|
||||
|
||||
### Advanced Usage
|
||||
|
||||
```python
|
||||
# Custom scheduler with priority handling and monitoring
|
||||
config = CoreModuleConfig(
|
||||
scheduler=SchedulerConfig(
|
||||
scheduler_type=SchedulerType.PRIORITY,
|
||||
priority_weights={"high": 3.0, "normal": 1.0, "low": 0.3},
|
||||
enable_preemption=True
|
||||
),
|
||||
cache=CacheConfig(
|
||||
strategy=CacheStrategy.ADAPTIVE,
|
||||
max_memory_gb=16.0,
|
||||
enable_compression=True
|
||||
),
|
||||
enable_metrics=True,
|
||||
enable_health_checks=True
|
||||
)
|
||||
|
||||
scheduler = CoreScheduler(config.scheduler, config.cache)
|
||||
await scheduler.start()
|
||||
|
||||
# Monitor scheduler
|
||||
async def monitor_scheduler():
|
||||
while True:
|
||||
status = scheduler.get_status()
|
||||
metrics = scheduler.get_metrics()
|
||||
health = scheduler.get_health_status()
|
||||
|
||||
print(f"Status: {status}")
|
||||
print(f"Metrics: {metrics}")
|
||||
print(f"Health: {health}")
|
||||
|
||||
await asyncio.sleep(10)
|
||||
|
||||
# Start monitoring
|
||||
asyncio.create_task(monitor_scheduler())
|
||||
```
|
||||
|
||||
### Integration Example
|
||||
|
||||
```python
|
||||
# Integration with other modules
|
||||
from vllm_omni.engine import EngineManager
|
||||
from vllm_omni.executor import ExecutorManager
|
||||
|
||||
class OmniSystem:
|
||||
def __init__(self, config: CoreModuleConfig):
|
||||
self.scheduler = CoreScheduler(config.scheduler, config.cache)
|
||||
self.engine_manager = EngineManager(config.engine)
|
||||
self.executor_manager = ExecutorManager(config.executor)
|
||||
|
||||
async def start(self):
|
||||
await self.scheduler.start()
|
||||
await self.engine_manager.start()
|
||||
await self.executor_manager.start()
|
||||
|
||||
async def process_request(self, request: OmniRequest):
|
||||
# Schedule request
|
||||
await self.scheduler.schedule_request(request)
|
||||
|
||||
# Get next request
|
||||
next_request = await self.scheduler.get_next_request()
|
||||
if next_request:
|
||||
# Route to appropriate executor
|
||||
executor = self.executor_manager.get_executor(next_request)
|
||||
result = await executor.execute(next_request)
|
||||
return result
|
||||
```
|
||||
Reference in new issue
Block a user