From 0c5caf6ab175e95bde23f233d71505de9b2ac11d Mon Sep 17 00:00:00 2001 From: hsliu Date: Tue, 16 Sep 2025 16:05:47 +0800 Subject: [PATCH] scheduler api design --- docs/api/API_DESIGN_TEMPLATE.md | 396 ++++++++++++++++++++ docs/api/README.md | 235 ++++++++++++ docs/api/core_module_api.md | 580 ++++++++++++++++++++++++++++++ vllm_omni/core/sched/__init__.py | 1 + vllm_omni/core/sched/scheduler.py | 41 +++ 5 files changed, 1253 insertions(+) create mode 100644 docs/api/API_DESIGN_TEMPLATE.md create mode 100644 docs/api/README.md create mode 100644 docs/api/core_module_api.md diff --git a/docs/api/API_DESIGN_TEMPLATE.md b/docs/api/API_DESIGN_TEMPLATE.md new file mode 100644 index 000000000..9c31eec34 --- /dev/null +++ b/docs/api/API_DESIGN_TEMPLATE.md @@ -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. diff --git a/docs/api/README.md b/docs/api/README.md new file mode 100644 index 000000000..681bbba81 --- /dev/null +++ b/docs/api/README.md @@ -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. diff --git a/docs/api/core_module_api.md b/docs/api/core_module_api.md new file mode 100644 index 000000000..b32bcd3a9 --- /dev/null +++ b/docs/api/core_module_api.md @@ -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 +``` diff --git a/vllm_omni/core/sched/__init__.py b/vllm_omni/core/sched/__init__.py index e69de29bb..0925dc9ca 100644 --- a/vllm_omni/core/sched/__init__.py +++ b/vllm_omni/core/sched/__init__.py @@ -0,0 +1 @@ +from vllm.v1.core.sched import SchedulerInterface \ No newline at end of file diff --git a/vllm_omni/core/sched/scheduler.py b/vllm_omni/core/sched/scheduler.py index e69de29bb..852cfde93 100644 --- a/vllm_omni/core/sched/scheduler.py +++ b/vllm_omni/core/sched/scheduler.py @@ -0,0 +1,41 @@ + +from vllm_omni.request import OmniRequest +from typing import List +from threading import Lock, Condition +from vllm_omni.config import OmniConfig + +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.core.sched import SchedulerInterface +from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry +from vllm.v1.structured_output import StructuredOutputManager + + +class OmniScheduler(SchedulerInterface): + """ + OmniScheduler: Scheduler for vLLM-omni multimodal processing. + + This scheduler extends vLLM's scheduler to support multimodal and non-autoregressive + processing with additional fields and methods specific to vLLM-omni. + """ + def __init__(self, + omni_config: OmniConfig, + kv_cache_config: KVCacheConfig, + structured_output_manager: StructuredOutputManager, + mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, + include_finished_set: bool = False, + log_stats: bool = False, + ): + super().__init__( + vllm_config=omni_config.vllm_config, + kv_cache_config=kv_cache_config, + multimodal_registry=mm_registry, + structured_output_manager=structured_output_manager, + include_finished_set=include_finished_set, + log_stats=log_stats, + ) + self.omni_config = omni_config + + def schedule(self, requests: List[OmniRequest]) -> List[OmniRequest]: + # TODO: Implement scheduling logic + return requests +