0xSerocleand98768b10 months agoHistory
121 lines · 3.1 KB · Python
RawDownload
"""Pydantic models for vLLM Studio.""" from enum import Enumfrom typing import Any, Dict, List, Optional from pydantic import BaseModel, Field, model_validator  class Backend(str, Enum):    """Supported inference backends."""     VLLM = "vllm"    SGLANG = "sglang"  class Recipe(BaseModel):    """Model launch configuration."""     id: str = Field(..., description="Unique identifier")    name: str = Field(..., description="Display name")    model_path: str = Field(..., description="Path to model")    backend: Backend = Field(default=Backend.VLLM)     # Parallelism    tensor_parallel_size: int = Field(default=1, alias="tp")    pipeline_parallel_size: int = Field(default=1, alias="pp")     # Memory    max_model_len: int = Field(default=32768)    gpu_memory_utilization: float = Field(default=0.9)    kv_cache_dtype: str = Field(default="auto")     # Batching    max_num_seqs: int = Field(default=256)     # Features    trust_remote_code: bool = Field(default=True)    tool_call_parser: Optional[str] = Field(default=None)     # Quantization    quantization: Optional[str] = Field(default=None)    dtype: Optional[str] = Field(default=None)     # Networking    host: str = Field(default="0.0.0.0")    port: int = Field(default=8000)    served_model_name: Optional[str] = Field(default=None)     # Custom Python    python_path: Optional[str] = Field(default=None)     # Extra CLI args    extra_args: Dict[str, Any] = Field(default_factory=dict)     @model_validator(mode="before")    @classmethod    def normalize(cls, data: Any) -> Any:        if not isinstance(data, dict):            return data        d = dict(data)        # Legacy: engine -> backend        if "backend" not in d and "engine" in d:            d["backend"] = d.pop("engine")        # Fold unknown keys into extra_args        extra = dict(d.get("extra_args") or {})        known = set(cls.model_fields.keys()) | {"tp", "pp", "engine", "extra_args"}        for k in list(d.keys()):            if k not in known:                extra[k] = d.pop(k)        d["extra_args"] = extra        return d     model_config = {"populate_by_name": True}  class ProcessInfo(BaseModel):    """Running inference process."""     pid: int    backend: str    model_path: Optional[str] = None    port: int    served_model_name: Optional[str] = None  class LaunchResult(BaseModel):    """Result of launching a model."""     success: bool    pid: Optional[int] = None    message: str    log_file: Optional[str] = None  class HealthResponse(BaseModel):    """Health check response."""     status: str    version: str    inference_ready: bool = False    # Frontend compatibility: UI expects `backend_reachable`.    backend_reachable: bool = False    running_model: Optional[str] = None  class OpenAIModelInfo(BaseModel):    """OpenAI-compatible model info."""     id: str    object: str = "model"    created: int    owned_by: str = "vllm-studio"    active: bool = False    max_model_len: Optional[int] = None  class OpenAIModelList(BaseModel):    """OpenAI-compatible model list response."""     object: str = "list"    data: List[OpenAIModelInfo]