wehub-resource-sync
302 lines
9.4 KiB
Python
302 lines
9.4 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Centralized configuration for oMLX.
|
|
|
|
This module provides unified configuration management with:
|
|
- Pydantic validation
|
|
- Environment variable support
|
|
- CLI argument mapping
|
|
- Default values with sensible defaults
|
|
"""
|
|
|
|
import os
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional, Union
|
|
|
|
|
|
def parse_size(size_str: str) -> int:
|
|
"""
|
|
Parse a human-readable size string to bytes.
|
|
|
|
Args:
|
|
size_str: Size string like "100GB", "50MB", "1TB".
|
|
|
|
Returns:
|
|
Size in bytes.
|
|
"""
|
|
size_str = size_str.strip().upper()
|
|
|
|
units = {
|
|
"B": 1,
|
|
"KB": 1024,
|
|
"MB": 1024**2,
|
|
"GB": 1024**3,
|
|
"TB": 1024**4,
|
|
}
|
|
|
|
for unit, multiplier in units.items():
|
|
if size_str.endswith(unit):
|
|
try:
|
|
value = float(size_str[: -len(unit)])
|
|
return int(value * multiplier)
|
|
except ValueError:
|
|
pass
|
|
|
|
# Try parsing as plain number (bytes)
|
|
try:
|
|
return int(size_str)
|
|
except ValueError:
|
|
raise ValueError(f"Invalid size string: {size_str}")
|
|
|
|
|
|
@dataclass
|
|
class ServerConfig:
|
|
"""Server configuration."""
|
|
|
|
host: str = "0.0.0.0"
|
|
port: int = 8000
|
|
log_level: str = "info"
|
|
cors_origins: List[str] = field(default_factory=lambda: ["*"])
|
|
|
|
|
|
@dataclass
|
|
class ModelConfig:
|
|
"""Model configuration."""
|
|
|
|
model_name: str = ""
|
|
# Security: default off. HuggingFace repos can ship arbitrary modeling_*.py
|
|
# that gets executed at load time when this is True. Issue #926.
|
|
trust_remote_code: bool = False
|
|
model_path: Optional[str] = None
|
|
|
|
|
|
@dataclass
|
|
class GenerationConfig:
|
|
"""Generation parameters configuration."""
|
|
|
|
max_tokens: int = 32768
|
|
temperature: float = 1.0
|
|
top_p: float = 0.95
|
|
top_k: int = 0
|
|
force_sampling: bool = False
|
|
|
|
|
|
@dataclass
|
|
class SchedulerConfig:
|
|
"""Scheduler configuration."""
|
|
|
|
max_num_seqs: int = 8
|
|
completion_batch_size: int = 8
|
|
embedding_batch_size: int = 32
|
|
stream_interval: int = 1
|
|
enable_thinking: Optional[bool] = None
|
|
|
|
|
|
@dataclass
|
|
class CacheConfig:
|
|
"""Cache configuration (deprecated, kept for compatibility)."""
|
|
|
|
# All cache options moved to PagedSSDCacheConfig
|
|
pass
|
|
|
|
|
|
@dataclass
|
|
class PagedSSDCacheConfig:
|
|
"""Paged SSD cache configuration. oMLX only supports paged SSD-based caching."""
|
|
|
|
enabled: bool = False
|
|
hot_cache_only: bool = False
|
|
cache_dir: Optional[Path] = None
|
|
max_size: str = "100GB"
|
|
hot_cache_max_size: str = "0" # "0" = disabled, e.g. "8GB"
|
|
|
|
@property
|
|
def max_size_bytes(self) -> int:
|
|
"""Get max size in bytes."""
|
|
return parse_size(self.max_size)
|
|
|
|
@property
|
|
def hot_cache_max_size_bytes(self) -> int:
|
|
"""Get hot cache max size in bytes. 0 means disabled."""
|
|
return parse_size(self.hot_cache_max_size)
|
|
|
|
|
|
@dataclass
|
|
class MCPConfig:
|
|
"""MCP (Model Context Protocol) configuration."""
|
|
|
|
config_path: Optional[str] = None
|
|
enabled: bool = False
|
|
|
|
|
|
@dataclass
|
|
class OMLXConfig:
|
|
"""
|
|
Centralized configuration for oMLX.
|
|
|
|
This class combines all configuration sections and provides
|
|
environment variable overrides and CLI argument mapping.
|
|
"""
|
|
|
|
server: ServerConfig = field(default_factory=ServerConfig)
|
|
model: ModelConfig = field(default_factory=ModelConfig)
|
|
generation: GenerationConfig = field(default_factory=GenerationConfig)
|
|
scheduler: SchedulerConfig = field(default_factory=SchedulerConfig)
|
|
cache: CacheConfig = field(default_factory=CacheConfig)
|
|
paged_ssd_cache: PagedSSDCacheConfig = field(default_factory=PagedSSDCacheConfig)
|
|
mcp: MCPConfig = field(default_factory=MCPConfig)
|
|
|
|
# Feature flags
|
|
continuous_batching: bool = False
|
|
|
|
@classmethod
|
|
def from_env(cls) -> "OMLXConfig":
|
|
"""
|
|
Create config from environment variables.
|
|
|
|
Environment variables are prefixed with OMLX_.
|
|
"""
|
|
config = cls()
|
|
|
|
# Server settings
|
|
config.server.host = os.getenv("OMLX_HOST", config.server.host)
|
|
config.server.port = int(os.getenv("OMLX_PORT", str(config.server.port)))
|
|
config.server.log_level = os.getenv("OMLX_LOG_LEVEL", config.server.log_level)
|
|
|
|
# Model settings
|
|
config.model.model_name = os.getenv("OMLX_MODEL", config.model.model_name)
|
|
config.model.trust_remote_code = os.getenv(
|
|
"OMLX_TRUST_REMOTE_CODE", "false"
|
|
).lower() == "true"
|
|
|
|
# Generation settings
|
|
config.generation.max_tokens = int(
|
|
os.getenv("OMLX_MAX_TOKENS", str(config.generation.max_tokens))
|
|
)
|
|
config.generation.temperature = float(
|
|
os.getenv("OMLX_TEMPERATURE", str(config.generation.temperature))
|
|
)
|
|
|
|
# Paged SSD cache settings
|
|
config.paged_ssd_cache.hot_cache_only = os.getenv("OMLX_HOT_CACHE_ONLY", "false").lower() == "true"
|
|
paged_ssd_dir = os.getenv("OMLX_PAGED_SSD_CACHE_DIR")
|
|
if paged_ssd_dir:
|
|
config.paged_ssd_cache.enabled = True
|
|
config.paged_ssd_cache.cache_dir = Path(paged_ssd_dir)
|
|
config.paged_ssd_cache.max_size = os.getenv(
|
|
"OMLX_PAGED_SSD_CACHE_MAX_SIZE", config.paged_ssd_cache.max_size
|
|
)
|
|
|
|
# MCP settings
|
|
mcp_config = os.getenv("OMLX_MCP_CONFIG")
|
|
if mcp_config:
|
|
config.mcp.enabled = True
|
|
config.mcp.config_path = mcp_config
|
|
|
|
# Feature flags
|
|
config.continuous_batching = os.getenv(
|
|
"OMLX_CONTINUOUS_BATCHING", "false"
|
|
).lower() == "true"
|
|
|
|
return config
|
|
|
|
@classmethod
|
|
def from_cli_args(cls, args: Any) -> "OMLXConfig":
|
|
"""
|
|
Create config from argparse namespace.
|
|
|
|
Args:
|
|
args: Argparse namespace with CLI arguments.
|
|
|
|
Returns:
|
|
OMLXConfig instance.
|
|
"""
|
|
config = cls.from_env() # Start with env vars
|
|
|
|
# Override with CLI args if provided
|
|
if hasattr(args, "host") and args.host:
|
|
config.server.host = args.host
|
|
if hasattr(args, "port") and args.port:
|
|
config.server.port = args.port
|
|
if hasattr(args, "log_level") and args.log_level:
|
|
config.server.log_level = args.log_level
|
|
|
|
if hasattr(args, "model") and args.model:
|
|
config.model.model_name = args.model
|
|
if hasattr(args, "trust_remote_code"):
|
|
config.model.trust_remote_code = args.trust_remote_code
|
|
|
|
if hasattr(args, "max_tokens") and args.max_tokens:
|
|
config.generation.max_tokens = args.max_tokens
|
|
if hasattr(args, "temperature") and args.temperature is not None:
|
|
config.generation.temperature = args.temperature
|
|
if hasattr(args, "top_p") and args.top_p is not None:
|
|
config.generation.top_p = args.top_p
|
|
if hasattr(args, "top_k") and args.top_k is not None:
|
|
config.generation.top_k = args.top_k
|
|
|
|
if hasattr(args, "continuous_batching"):
|
|
config.continuous_batching = args.continuous_batching
|
|
|
|
# Paged SSD cache settings
|
|
if hasattr(args, "hot_cache_only") and args.hot_cache_only is not None:
|
|
config.paged_ssd_cache.hot_cache_only = args.hot_cache_only
|
|
if hasattr(args, "paged_ssd_cache_dir") and args.paged_ssd_cache_dir:
|
|
config.paged_ssd_cache.enabled = True
|
|
config.paged_ssd_cache.cache_dir = Path(args.paged_ssd_cache_dir)
|
|
if hasattr(args, "paged_ssd_cache_max_size") and args.paged_ssd_cache_max_size:
|
|
config.paged_ssd_cache.max_size = args.paged_ssd_cache_max_size
|
|
|
|
if hasattr(args, "mcp_config") and args.mcp_config:
|
|
config.mcp.enabled = True
|
|
config.mcp.config_path = args.mcp_config
|
|
|
|
return config
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert config to dictionary."""
|
|
from dataclasses import asdict
|
|
|
|
return {
|
|
"server": asdict(self.server),
|
|
"model": asdict(self.model),
|
|
"generation": asdict(self.generation),
|
|
"scheduler": asdict(self.scheduler),
|
|
"cache": asdict(self.cache),
|
|
"paged_ssd_cache": {
|
|
**asdict(self.paged_ssd_cache),
|
|
"cache_dir": str(self.paged_ssd_cache.cache_dir) if self.paged_ssd_cache.cache_dir else None,
|
|
},
|
|
"mcp": asdict(self.mcp),
|
|
"continuous_batching": self.continuous_batching,
|
|
}
|
|
|
|
def validate(self) -> List[str]:
|
|
"""
|
|
Validate configuration.
|
|
|
|
Returns:
|
|
List of validation errors (empty if valid).
|
|
"""
|
|
errors = []
|
|
|
|
# Server validation
|
|
if not 0 < self.server.port < 65536:
|
|
errors.append(f"Invalid port: {self.server.port}")
|
|
|
|
# Generation validation
|
|
if self.generation.max_tokens <= 0:
|
|
errors.append(f"max_tokens must be positive: {self.generation.max_tokens}")
|
|
if not 0.0 <= self.generation.temperature <= 2.0:
|
|
errors.append(f"temperature must be 0.0-2.0: {self.generation.temperature}")
|
|
if not 0.0 <= self.generation.top_p <= 1.0:
|
|
errors.append(f"top_p must be 0.0-1.0: {self.generation.top_p}")
|
|
|
|
# Paged SSD cache validation
|
|
if self.paged_ssd_cache.enabled:
|
|
if not self.paged_ssd_cache.cache_dir:
|
|
errors.append("Paged SSD cache enabled but no cache_dir specified")
|
|
|
|
return errors
|