"""Configuration loader using KDL format."""
import logging
from pathlib import Path
from typing import Any
import kdl
from pydantic import BaseModel, Field
logger = logging.getLogger(__name__)
class LimitsConfig(BaseModel):
"""Resource limits configuration."""
max_array_size: int = Field(default=100_000_000, description="Maximum array elements")
max_particles: int = Field(default=10_000_000, description="Maximum particles in MD")
max_time_steps: int = Field(default=1_000_000, description="Maximum simulation timesteps")
max_epochs: int = Field(default=1000, description="Maximum training epochs")
class GPUConfig(BaseModel):
"""GPU configuration."""
memory_fraction: float = Field(default=0.8, ge=0.0, le=1.0, description="GPU memory fraction")
backend: str = Field(default="cuda", description="GPU backend (cuda/rocm)")
class MCPConfig(BaseModel):
"""MCP server configuration."""
server_name: str
limits: LimitsConfig = Field(default_factory=LimitsConfig)
gpu: GPUConfig = Field(default_factory=GPUConfig)
extra: dict[str, Any] = Field(default_factory=dict)
def load_config(config_path: Path) -> MCPConfig:
"""Load configuration from KDL file."""
if not config_path.exists():
logger.warning(f"Config file not found: {config_path}, using defaults")
return MCPConfig(server_name="unknown")
try:
with open(config_path) as f:
doc = kdl.parse(f.read())
# Extract server name
server_name = "unknown"
for node in doc:
if node.name == "server":
server_name = str(node.args[0]) if node.args else "unknown"
break
# Extract limits
limits_dict: dict[str, Any] = {}
for node in doc:
if node.name == "limits":
for prop in node.props:
# Convert kebab-case to snake_case
key = prop.replace("-", "_")
limits_dict[key] = node.props[prop]
# Extract GPU config
gpu_dict: dict[str, Any] = {}
for node in doc:
if node.name == "gpu":
for prop in node.props:
key = prop.replace("-", "_")
gpu_dict[key] = node.props[prop]
# Build config
config = MCPConfig(
server_name=server_name,
limits=LimitsConfig(**limits_dict) if limits_dict else LimitsConfig(),
gpu=GPUConfig(**gpu_dict) if gpu_dict else GPUConfig(),
)
logger.info(f"Loaded config from {config_path}: {config.server_name}")
return config
except Exception as e:
logger.error(f"Failed to load config from {config_path}: {e}")
raise
