init request.py
This commit is contained in:
hsliu committed
1 parent
dcb8868477
commit
8dbe6db1be
1 file changed
+294
-2
+294
-2
@@ -1,7 +1,299 @@
|
||||
"""
|
||||
OmniRequest: Extended request class for vLLM-omni multimodal processing.
|
||||
|
||||
This class extends vLLM's Request to support multimodal and non-autoregressive
|
||||
processing with additional fields and methods specific to vLLM-omni.
|
||||
"""
|
||||
|
||||
import enum
|
||||
import time
|
||||
from typing import Optional, Dict, Any
|
||||
from typing import Optional, Dict, Any, List, Union
|
||||
from dataclasses import dataclass, field
|
||||
from vllm.v1.request import Request as vLLMRequest
|
||||
|
||||
|
||||
class RequestType(enum.Enum):
|
||||
"""Types of requests supported by vLLM-omni."""
|
||||
TEXT = "text"
|
||||
IMAGE = "image"
|
||||
AUDIO = "audio"
|
||||
VIDEO = "video"
|
||||
MULTIMODAL = "multimodal"
|
||||
DIFFUSION = "diffusion"
|
||||
|
||||
|
||||
class ProcessingStage(enum.Enum):
|
||||
"""Processing stages for multi-stage models."""
|
||||
PREPROCESSING = "preprocessing"
|
||||
AR_GENERATION = "ar_generation"
|
||||
DIFFUSION_GENERATION = "diffusion_generation"
|
||||
POSTPROCESSING = "postprocessing"
|
||||
COMPLETED = "completed"
|
||||
|
||||
|
||||
@dataclass
|
||||
class MultimodalData:
|
||||
"""Container for multimodal input data."""
|
||||
data_type: str # "image", "audio", "video", etc.
|
||||
data: Any # The actual data (numpy array, bytes, etc.)
|
||||
metadata: Dict[str, Any] = field(default_factory=dict)
|
||||
processed: bool = False
|
||||
|
||||
|
||||
@dataclass
|
||||
class DiffusionParams:
|
||||
"""Parameters specific to diffusion model generation."""
|
||||
num_inference_steps: int = 50
|
||||
guidance_scale: float = 7.5
|
||||
seed: Optional[int] = None
|
||||
scheduler: str = "ddpm"
|
||||
strength: float = 1.0
|
||||
eta: float = 0.0
|
||||
|
||||
|
||||
class OmniRequest(vLLMRequest):
|
||||
# pass
|
||||
"""
|
||||
Extended request class for vLLM-omni multimodal and non-autoregressive processing.
|
||||
|
||||
This class extends vLLM's Request with additional fields and methods to support:
|
||||
- Multimodal input processing
|
||||
- Non-autoregressive generation (diffusion models)
|
||||
- Multi-stage processing pipelines
|
||||
- Enhanced caching for different model types
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
request_id: str,
|
||||
prompt: Optional[str] = None,
|
||||
prompt_token_ids: Optional[List[int]] = None,
|
||||
sampling_params: Optional[Any] = None,
|
||||
arrival_time: Optional[float] = None,
|
||||
lora_request: Optional[Any] = None,
|
||||
multi_modal_data: Optional[Dict[str, Any]] = None,
|
||||
multi_modal_placeholders: Optional[Dict[str, str]] = None,
|
||||
priority: int = 0,
|
||||
# vLLM-omni specific parameters
|
||||
request_type: RequestType = RequestType.TEXT,
|
||||
processing_stage: ProcessingStage = ProcessingStage.PREPROCESSING,
|
||||
multimodal_inputs: Optional[List[MultimodalData]] = None,
|
||||
diffusion_params: Optional[DiffusionParams] = None,
|
||||
output_format: str = "text",
|
||||
cache_key: Optional[str] = None,
|
||||
stage_configs: Optional[Dict[str, Any]] = None,
|
||||
**kwargs
|
||||
):
|
||||
# Initialize parent class
|
||||
super().__init__(
|
||||
request_id=request_id,
|
||||
prompt=prompt,
|
||||
prompt_token_ids=prompt_token_ids,
|
||||
sampling_params=sampling_params,
|
||||
arrival_time=arrival_time or time.time(),
|
||||
lora_request=lora_request,
|
||||
multi_modal_data=multi_modal_data,
|
||||
multi_modal_placeholders=multi_modal_placeholders,
|
||||
priority=priority,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
# vLLM-omni specific attributes
|
||||
self.request_type = request_type
|
||||
self.processing_stage = processing_stage
|
||||
self.multimodal_inputs = multimodal_inputs or []
|
||||
self.diffusion_params = diffusion_params or DiffusionParams()
|
||||
self.output_format = output_format
|
||||
self.cache_key = cache_key
|
||||
self.stage_configs = stage_configs or {}
|
||||
|
||||
# Processing state
|
||||
self.current_stage = 0
|
||||
self.stage_results = {}
|
||||
self.hidden_states = None
|
||||
self.intermediate_outputs = []
|
||||
|
||||
# Timing and metrics
|
||||
self.stage_timings = {}
|
||||
self.total_processing_time = 0.0
|
||||
|
||||
# Error handling
|
||||
self.errors = []
|
||||
self.retry_count = 0
|
||||
self.max_retries = 3
|
||||
|
||||
def add_multimodal_input(self, data: MultimodalData) -> None:
|
||||
"""Add a multimodal input to the request."""
|
||||
self.multimodal_inputs.append(data)
|
||||
|
||||
def get_multimodal_inputs_by_type(self, data_type: str) -> List[MultimodalData]:
|
||||
"""Get all multimodal inputs of a specific type."""
|
||||
return [inp for inp in self.multimodal_inputs if inp.data_type == data_type]
|
||||
|
||||
def update_processing_stage(self, stage: ProcessingStage) -> None:
|
||||
"""Update the current processing stage."""
|
||||
self.processing_stage = stage
|
||||
self.stage_timings[stage.value] = time.time()
|
||||
|
||||
def add_stage_result(self, stage: str, result: Any) -> None:
|
||||
"""Add a result from a processing stage."""
|
||||
self.stage_results[stage] = result
|
||||
self.intermediate_outputs.append({
|
||||
'stage': stage,
|
||||
'result': result,
|
||||
'timestamp': time.time()
|
||||
})
|
||||
|
||||
def get_stage_result(self, stage: str) -> Any:
|
||||
"""Get a result from a specific processing stage."""
|
||||
return self.stage_results.get(stage)
|
||||
|
||||
def set_hidden_states(self, hidden_states: Any) -> None:
|
||||
"""Set hidden states for the request."""
|
||||
self.hidden_states = hidden_states
|
||||
|
||||
def get_hidden_states(self) -> Any:
|
||||
"""Get hidden states for the request."""
|
||||
return self.hidden_states
|
||||
|
||||
def add_error(self, error: str) -> None:
|
||||
"""Add an error to the request."""
|
||||
self.errors.append({
|
||||
'error': error,
|
||||
'timestamp': time.time(),
|
||||
'stage': self.processing_stage.value
|
||||
})
|
||||
|
||||
def has_errors(self) -> bool:
|
||||
"""Check if the request has any errors."""
|
||||
return len(self.errors) > 0
|
||||
|
||||
def can_retry(self) -> bool:
|
||||
"""Check if the request can be retried."""
|
||||
return self.retry_count < self.max_retries
|
||||
|
||||
def increment_retry(self) -> None:
|
||||
"""Increment the retry count."""
|
||||
self.retry_count += 1
|
||||
|
||||
def generate_cache_key(self) -> str:
|
||||
"""Generate a cache key for the request."""
|
||||
if self.cache_key:
|
||||
return self.cache_key
|
||||
|
||||
# Generate cache key based on request content
|
||||
key_parts = [
|
||||
self.request_id,
|
||||
str(self.request_type.value),
|
||||
str(self.prompt_token_ids) if self.prompt_token_ids else str(self.prompt),
|
||||
str(self.diffusion_params.num_inference_steps) if self.diffusion_params else "0",
|
||||
str(len(self.multimodal_inputs))
|
||||
]
|
||||
return "_".join(key_parts)
|
||||
|
||||
def get_processing_time(self) -> float:
|
||||
"""Get the total processing time for the request."""
|
||||
if self.stage_timings:
|
||||
return time.time() - min(self.stage_timings.values())
|
||||
return 0.0
|
||||
|
||||
def is_completed(self) -> bool:
|
||||
"""Check if the request is completed."""
|
||||
return self.processing_stage == ProcessingStage.COMPLETED
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
"""Convert the request to a dictionary for serialization."""
|
||||
return {
|
||||
'request_id': self.request_id,
|
||||
'request_type': self.request_type.value,
|
||||
'processing_stage': self.processing_stage.value,
|
||||
'prompt': self.prompt,
|
||||
'prompt_token_ids': self.prompt_token_ids,
|
||||
'output_format': self.output_format,
|
||||
'multimodal_inputs_count': len(self.multimodal_inputs),
|
||||
'diffusion_params': self.diffusion_params.__dict__ if self.diffusion_params else None,
|
||||
'stage_results': list(self.stage_results.keys()),
|
||||
'has_errors': self.has_errors(),
|
||||
'retry_count': self.retry_count,
|
||||
'processing_time': self.get_processing_time()
|
||||
}
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""String representation of the request."""
|
||||
return (f"OmniRequest(id={self.request_id}, "
|
||||
f"type={self.request_type.value}, "
|
||||
f"stage={self.processing_stage.value}, "
|
||||
f"multimodal_inputs={len(self.multimodal_inputs)})")
|
||||
|
||||
|
||||
# Factory functions for creating different types of requests
|
||||
def create_text_request(
|
||||
request_id: str,
|
||||
prompt: str,
|
||||
sampling_params: Optional[Any] = None,
|
||||
**kwargs
|
||||
) -> OmniRequest:
|
||||
"""Create a text-only request."""
|
||||
return OmniRequest(
|
||||
request_id=request_id,
|
||||
prompt=prompt,
|
||||
request_type=RequestType.TEXT,
|
||||
sampling_params=sampling_params,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
|
||||
def create_image_request(
|
||||
request_id: str,
|
||||
prompt: str,
|
||||
image_data: Any,
|
||||
diffusion_params: Optional[DiffusionParams] = None,
|
||||
**kwargs
|
||||
) -> OmniRequest:
|
||||
"""Create an image generation request."""
|
||||
multimodal_input = MultimodalData(
|
||||
data_type="image",
|
||||
data=image_data,
|
||||
metadata={"is_input": True}
|
||||
)
|
||||
|
||||
return OmniRequest(
|
||||
request_id=request_id,
|
||||
prompt=prompt,
|
||||
request_type=RequestType.IMAGE,
|
||||
multimodal_inputs=[multimodal_input],
|
||||
diffusion_params=diffusion_params,
|
||||
output_format="image",
|
||||
**kwargs
|
||||
)
|
||||
|
||||
|
||||
def create_multimodal_request(
|
||||
request_id: str,
|
||||
prompt: str,
|
||||
multimodal_inputs: List[MultimodalData],
|
||||
**kwargs
|
||||
) -> OmniRequest:
|
||||
"""Create a multimodal request."""
|
||||
return OmniRequest(
|
||||
request_id=request_id,
|
||||
prompt=prompt,
|
||||
request_type=RequestType.MULTIMODAL,
|
||||
multimodal_inputs=multimodal_inputs,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
|
||||
def create_diffusion_request(
|
||||
request_id: str,
|
||||
prompt: str,
|
||||
diffusion_params: DiffusionParams,
|
||||
**kwargs
|
||||
) -> OmniRequest:
|
||||
"""Create a diffusion model request."""
|
||||
return OmniRequest(
|
||||
request_id=request_id,
|
||||
prompt=prompt,
|
||||
request_type=RequestType.DIFFUSION,
|
||||
diffusion_params=diffusion_params,
|
||||
**kwargs
|
||||
)
|
||||
Reference in new issue
Block a user