1197 lines
49 KiB
Python
1197 lines
49 KiB
Python
"""Per-model settings management for oMLX.
|
||
|
||
This module provides dataclasses and a manager for storing and retrieving
|
||
per-model configuration settings, including sampling parameters, pinned/default
|
||
flags, and metadata.
|
||
"""
|
||
|
||
import copy
|
||
import json
|
||
import logging
|
||
import threading
|
||
from dataclasses import dataclass, field, fields
|
||
from pathlib import Path
|
||
from typing import Any, Callable, Dict, Optional
|
||
|
||
from .model_profiles import (
|
||
MODEL_SPECIFIC_PROFILE_FIELDS,
|
||
filter_profile_fields,
|
||
filter_universal_fields,
|
||
slugify_profile_api_name,
|
||
validate_profile_name,
|
||
utcnow,
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# Current settings file format version
|
||
SETTINGS_VERSION = 1
|
||
PROFILES_VERSION = 1
|
||
TEMPLATES_VERSION = 1
|
||
|
||
|
||
@dataclass
|
||
class ModelSettings:
|
||
"""Per-model configuration settings.
|
||
|
||
Attributes:
|
||
max_context_window: Maximum prompt token count before rejection (None = use global default).
|
||
max_tokens: Maximum number of tokens to generate (None = use global default).
|
||
temperature: Sampling temperature (None = use global default).
|
||
top_p: Nucleus sampling probability (None = use global default).
|
||
top_k: Top-k sampling parameter (None = use global default).
|
||
min_p: Minimum probability threshold (None = use global default).
|
||
repetition_penalty: Repetition penalty (None = use default 1.0, i.e. disabled).
|
||
presence_penalty: Presence penalty (None = use global default).
|
||
force_sampling: Force sampling even with temperature=0.
|
||
max_tool_result_tokens: Maximum tokens in tool result (None = use global default).
|
||
chat_template_kwargs: Extra chat template keyword arguments.
|
||
forced_ct_kwargs: Keys in chat_template_kwargs that cannot be overridden.
|
||
ttl_seconds: Auto-unload after idle seconds (None = no TTL).
|
||
model_type_override: "llm", "vlm", "embedding", "reranker", or None (auto-detect).
|
||
model_alias: API-visible alternative to the directory name.
|
||
index_cache_freq: IndexCache: every Nth layer keeps indexer (DeepSeek DSA
|
||
only; GLM-5.2 uses its native checkpoint schedule).
|
||
enable_thinking: Explicit toggle for thinking/reasoning mode (None = auto).
|
||
thinking_budget_enabled: Whether a thinking token budget is active.
|
||
thinking_budget_tokens: Max tokens for thinking/reasoning.
|
||
reasoning_parser: xgrammar builtin name: "qwen", "harmony", "llama", etc.
|
||
guided_grammar_enabled: Whether a default guided grammar is active.
|
||
guided_grammar: Default EBNF grammar for constrained decoding.
|
||
turboquant_kv_enabled: Enable TurboQuant KV cache compression.
|
||
turboquant_kv_bits: TurboQuant bit depth (2/2.5/3/3.5/4/6/8).
|
||
turboquant_skip_last: Skip last KVCache layer to prevent corruption.
|
||
specprefill_enabled: Enable SpecPrefill (experimental sparse prefill for MoE).
|
||
specprefill_draft_model: Path to draft model for SpecPrefill.
|
||
specprefill_keep_pct: Keep rate for SpecPrefill (0.1–0.5).
|
||
specprefill_threshold: Min tokens to trigger SpecPrefill.
|
||
dflash_enabled: Enable DFlash speculative decoding.
|
||
dflash_draft_model: Path/repo for DFlash draft checkpoint.
|
||
dflash_draft_quant_enabled: Enable draft model quantization.
|
||
dflash_draft_quant_weight_bits: Quantization weight bits (2, 4, 8).
|
||
dflash_draft_quant_activation_bits: Quantization activation bits (16, 32).
|
||
dflash_draft_quant_group_size: Quantization group size (32, 64, 128).
|
||
dflash_max_ctx: Token threshold to fall back to BatchedEngine (None = unlimited).
|
||
dflash_in_memory_cache: Enable DFlash L1 (RAM) prefix cache.
|
||
dflash_in_memory_cache_max_entries: L1 cache max entries (default 4, matches dflash balanced profile).
|
||
dflash_in_memory_cache_max_bytes: L1 cache byte budget.
|
||
dflash_ssd_cache: Enable DFlash L2 (SSD) prefix cache spill (uses omlx SSD cache dir).
|
||
dflash_ssd_cache_max_bytes: L2 (SSD) disk budget; dflash evicts oldest entries when exceeded.
|
||
dflash_draft_window_size: Draft model sliding-attention window (None = dflash default 1024).
|
||
Helps stabilise acceptance rate on long-context prompts.
|
||
dflash_draft_sink_size: Attention-sink tokens always kept regardless of window
|
||
(None = dflash default 64).
|
||
dflash_verify_mode: Verifier algorithm — "dflash", "adaptive", "ddtree", or "off"
|
||
(None = dflash default "adaptive"). "adaptive" can shrink block size when
|
||
acceptance drops.
|
||
mtp_enabled: Enable native multi-token prediction (mlx-lm PR 990 / PR 15 monkey-patch).
|
||
When True, BatchGenerator uses MTP draft+verify for singleton decode and
|
||
for multi-row decode batches whose cache positions are aligned. Unaligned
|
||
continuous batches fall back to standard decoding automatically. Compatible
|
||
model_types: qwen3_5*, qwen3_6*, deepseek_v4*. Mutually exclusive with
|
||
dflash_enabled.
|
||
vlm_mtp_enabled: Enable VLM MTP speculative decoding via an external assistant
|
||
drafter (mlx-vlm 191d7c8+). Target = Gemma4 VLM body, drafter must be a
|
||
"gemma4_assistant" model.
|
||
vlm_mtp_draft_model: Path/repo of the assistant drafter (e.g. "gemma-4-26B-A4B-it-assistant").
|
||
vlm_mtp_draft_block_size: Tokens drafted per round (None = mlx-vlm default).
|
||
is_pinned: Keep model loaded in memory.
|
||
is_default: Use this model when no model is specified.
|
||
display_name: Human-readable name for UI display.
|
||
description: Optional description of the model.
|
||
active_profile_name: Name of the currently-applied profile (None = no profile).
|
||
"""
|
||
|
||
# Sampling parameters (None means use global default)
|
||
max_context_window: Optional[int] = None
|
||
max_tokens: Optional[int] = None
|
||
temperature: Optional[float] = None
|
||
top_p: Optional[float] = None
|
||
top_k: Optional[int] = None
|
||
repetition_penalty: Optional[float] = None
|
||
min_p: Optional[float] = None
|
||
presence_penalty: Optional[float] = None
|
||
force_sampling: bool = False
|
||
max_tool_result_tokens: Optional[int] = None
|
||
chat_template_kwargs: Optional[Dict[str, Any]] = None
|
||
forced_ct_kwargs: Optional[list[str]] = (
|
||
None # Keys that cannot be overridden by API requests
|
||
)
|
||
ttl_seconds: Optional[int] = None # Auto-unload after idle seconds (None = no TTL)
|
||
model_type_override: Optional[str] = (
|
||
None # "llm", "vlm", "embedding", "reranker", or None (auto-detect)
|
||
)
|
||
model_alias: Optional[str] = (
|
||
None # API-visible name (alternative to directory name)
|
||
)
|
||
index_cache_freq: Optional[int] = (
|
||
None # IndexCache: every Nth layer keeps indexer (DeepSeek DSA only)
|
||
)
|
||
enable_thinking: Optional[bool] = (
|
||
None # Explicit toggle for thinking/reasoning mode (None = auto)
|
||
)
|
||
preserve_thinking: Optional[bool] = (
|
||
None # Keep <think> blocks in historical turns (None = auto, True when template supports it)
|
||
)
|
||
thinking_budget_enabled: bool = False
|
||
thinking_budget_tokens: Optional[int] = None
|
||
reasoning_parser: Optional[str] = (
|
||
None # xgrammar builtin name: "qwen", "harmony", "llama", etc.
|
||
)
|
||
guided_grammar_enabled: bool = False
|
||
guided_grammar: Optional[str] = None
|
||
|
||
# TurboQuant KV cache (mlx-vlm backend)
|
||
turboquant_kv_enabled: bool = False
|
||
turboquant_kv_bits: float = 4 # 2, 2.5, 3, 3.5, 4, 6, 8
|
||
turboquant_skip_last: bool = (
|
||
True # Skip last KVCache layer (prevents corruption on sensitive models)
|
||
)
|
||
|
||
# SpecPrefill (experimental: attention-based sparse prefill for MoE models)
|
||
specprefill_enabled: bool = False
|
||
specprefill_draft_model: Optional[str] = (
|
||
None # Path to draft model (must share tokenizer)
|
||
)
|
||
specprefill_keep_pct: Optional[float] = None # Keep rate (0.1-0.5, default 0.2)
|
||
specprefill_threshold: Optional[int] = None # Min tokens to trigger (default 8192)
|
||
|
||
# DFlash (block diffusion speculative decoding)
|
||
dflash_enabled: bool = False
|
||
dflash_draft_model: Optional[str] = None # Path/repo for DFlash draft checkpoint
|
||
dflash_draft_quant_enabled: Optional[bool] = None
|
||
dflash_draft_quant_weight_bits: Optional[int] = None # 2, 4, 8
|
||
dflash_draft_quant_activation_bits: Optional[int] = None # 16, 32
|
||
dflash_draft_quant_group_size: Optional[int] = None # 32, 64, 128
|
||
dflash_max_ctx: Optional[int] = (
|
||
None # None = unlimited; trigger BatchedEngine fallback when prompt_len >= this
|
||
)
|
||
# DFlash prefix cache (private to dflash; separate from omlx tiered cache because
|
||
# snapshots include draft model GDN state and target hidden chunks omlx never tracks)
|
||
dflash_in_memory_cache: bool = True
|
||
dflash_in_memory_cache_max_entries: int = (
|
||
4 # Matches dflash balanced profile default
|
||
)
|
||
dflash_in_memory_cache_max_bytes: int = (
|
||
8 * 1024 * 1024 * 1024
|
||
) # 8 GiB (balanced profile default)
|
||
dflash_ssd_cache: bool = (
|
||
False # Requires in-memory cache and an omlx paged SSD cache dir
|
||
)
|
||
dflash_ssd_cache_max_bytes: int = 20 * 1024 * 1024 * 1024 # 20 GiB L2 disk budget
|
||
# DFlash runtime tuning knobs. None = let dflash-mlx pick its own DEFAULT_RUNTIME_CONFIG
|
||
# value (currently window=1024, sink=64, verify_mode="adaptive"). Surfaced for long-context
|
||
# agentic workloads where acceptance drops on the default sliding window.
|
||
dflash_draft_window_size: Optional[int] = None
|
||
dflash_draft_sink_size: Optional[int] = None
|
||
dflash_verify_mode: Optional[str] = None # "dflash" | "adaptive" | "ddtree" | "off"
|
||
|
||
# Native MTP (mlx-lm PR 990 / PR 15 monkey-patch). When enabled, BatchGenerator
|
||
# uses MTP draft+verify for singleton decode and aligned multi-row decode batches.
|
||
# Compatible model_types: qwen3_5*, qwen3_6*, deepseek_v4*. Mutually exclusive
|
||
# with dflash.
|
||
mtp_enabled: bool = False
|
||
# Maximum chained MTP draft tokens per verify cycle (speculative depth).
|
||
# None = default (3). Effective for Qwen3.5/3.6 native MTP only;
|
||
# DeepSeek-V4 native MTP always runs depth 1. An adaptive controller
|
||
# picks 1..max per sequence from rolling acceptance/latency estimates;
|
||
# set to 1 for a fixed depth-1 cycle.
|
||
mtp_num_draft_tokens: Optional[int] = None
|
||
|
||
# VLM MTP speculative decoding via external MTP drafter (mlx-vlm f96138e+).
|
||
# Supported drafter types: gemma4_assistant (for Gemma 4 VLMs), qwen3_5_mtp
|
||
# (for Qwen 3.5/3.6). Both resolve to draft_kind="mtp" in mlx-vlm.
|
||
# Mutually exclusive with all other speculative paths because the wrapper
|
||
# bypasses mlx-lm BatchGenerator at decode time.
|
||
vlm_mtp_enabled: bool = False
|
||
vlm_mtp_draft_model: Optional[str] = (
|
||
None # Path / model id of the assistant drafter
|
||
)
|
||
vlm_mtp_draft_block_size: Optional[int] = (
|
||
None # Tokens per draft round (None = mlx-vlm default)
|
||
)
|
||
|
||
# Model management flags
|
||
is_pinned: bool = False
|
||
is_default: bool = False # Only one model can be default
|
||
is_hidden: bool = False # Hidden from /v1/models (still shown, badged, in admin)
|
||
is_favorite: bool = False # Listed first in /v1/models and admin lists
|
||
|
||
# Security: opt-in per model. When True, mlx-lm/mlx-vlm/mlx-embeddings/reranker
|
||
# loaders are allowed to execute custom Python from the model repository
|
||
# (modeling_*.py, tokenization_*.py). Off by default — see issue #926.
|
||
trust_remote_code: bool = False
|
||
|
||
# Metadata
|
||
display_name: Optional[str] = None
|
||
description: Optional[str] = None
|
||
active_profile_name: Optional[str] = None # Name of the currently-applied profile
|
||
|
||
def __post_init__(self) -> None:
|
||
# Native MTP is mutually exclusive with DFlash (also speculative).
|
||
# Reject the combo at construction time so the conflict surfaces in
|
||
# the admin UI / API rather than at model load. TurboQuant KV is
|
||
# compatible: its attention patch routes MTP's decode-shaped
|
||
# multi-row verify through the quantized decode kernels.
|
||
if self.mtp_enabled and self.dflash_enabled:
|
||
raise ValueError(
|
||
"mtp_enabled and dflash_enabled cannot both be True; choose one "
|
||
"speculative-decoding path per model"
|
||
)
|
||
# vlm_mtp wraps mlx-vlm's MTP loop and bypasses mlx-lm BatchGenerator
|
||
# at decode time, so it cannot coexist with any other speculative path
|
||
# or with TurboQuant (which mutates the same cache objects).
|
||
if self.vlm_mtp_enabled:
|
||
conflicts = [
|
||
("dflash_enabled", self.dflash_enabled),
|
||
("specprefill_enabled", self.specprefill_enabled),
|
||
("mtp_enabled", self.mtp_enabled),
|
||
("turboquant_kv_enabled", self.turboquant_kv_enabled),
|
||
]
|
||
for name, value in conflicts:
|
||
if value:
|
||
raise ValueError(
|
||
f"vlm_mtp_enabled and {name} cannot both be True; "
|
||
"choose one speculative path per model"
|
||
)
|
||
|
||
def to_dict(self) -> dict:
|
||
"""Convert to dictionary, excluding None values.
|
||
|
||
Returns:
|
||
Dictionary representation with None values filtered out.
|
||
"""
|
||
result = {}
|
||
for f in fields(self):
|
||
value = getattr(self, f.name)
|
||
if value is not None:
|
||
result[f.name] = value
|
||
return result
|
||
|
||
@classmethod
|
||
def from_dict(cls, data: dict) -> "ModelSettings":
|
||
"""Create ModelSettings from a dictionary.
|
||
|
||
Args:
|
||
data: Dictionary containing settings values.
|
||
|
||
Returns:
|
||
New ModelSettings instance with values from dict.
|
||
"""
|
||
# Get valid field names
|
||
valid_fields = {f.name for f in fields(cls)}
|
||
|
||
# Filter to only valid keys
|
||
filtered_data = {k: v for k, v in data.items() if k in valid_fields}
|
||
|
||
return cls(**filtered_data)
|
||
|
||
|
||
class ModelSettingsManager:
|
||
"""Manager for per-model settings with file persistence.
|
||
|
||
Handles loading, saving, and accessing model settings from a JSON file.
|
||
Thread-safe for concurrent access.
|
||
|
||
Attributes:
|
||
base_path: Base directory for settings storage.
|
||
settings_file: Path to the settings JSON file.
|
||
"""
|
||
|
||
def __init__(self, base_path: Path):
|
||
"""Initialize the settings manager.
|
||
|
||
Args:
|
||
base_path: Base directory for settings storage.
|
||
"""
|
||
self.base_path = Path(base_path)
|
||
self.settings_file = self.base_path / "model_settings.json"
|
||
self.profiles_file = self.base_path / "model_profiles.json"
|
||
self.templates_file = self.base_path / "global_templates.json"
|
||
self._lock = threading.Lock()
|
||
self._settings: Dict[str, ModelSettings] = {}
|
||
self._profiles: Dict[str, Dict[str, Dict[str, Any]]] = {}
|
||
self._templates: Dict[str, Dict[str, Any]] = {}
|
||
|
||
# Ensure base directory exists
|
||
self.base_path.mkdir(parents=True, exist_ok=True)
|
||
|
||
# Load existing settings
|
||
self._load()
|
||
self._load_profiles()
|
||
self._load_templates()
|
||
|
||
def _load(self) -> None:
|
||
"""Load settings from the JSON file.
|
||
|
||
If the file doesn't exist or is invalid, starts with empty settings.
|
||
"""
|
||
if not self.settings_file.exists():
|
||
logger.debug(f"Settings file not found: {self.settings_file}")
|
||
self._settings = {}
|
||
return
|
||
|
||
try:
|
||
with open(self.settings_file, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
|
||
# Check version
|
||
version = data.get("version", 1)
|
||
if version != SETTINGS_VERSION:
|
||
logger.warning(
|
||
f"Settings file version {version} differs from current {SETTINGS_VERSION}"
|
||
)
|
||
|
||
# Load model settings
|
||
models_data = data.get("models", {})
|
||
self._settings = {}
|
||
|
||
for model_id, model_data in models_data.items():
|
||
try:
|
||
self._settings[model_id] = ModelSettings.from_dict(model_data)
|
||
except Exception as e:
|
||
logger.warning(
|
||
f"Failed to load settings for model '{model_id}': {e}"
|
||
)
|
||
|
||
logger.info(f"Loaded settings for {len(self._settings)} models")
|
||
|
||
except json.JSONDecodeError as e:
|
||
logger.error(f"Invalid JSON in settings file: {e}")
|
||
self._settings = {}
|
||
except Exception as e:
|
||
logger.error(f"Failed to load settings file: {e}")
|
||
self._settings = {}
|
||
|
||
def _save(self) -> None:
|
||
"""Save settings to the JSON file.
|
||
|
||
Must be called while holding the lock.
|
||
"""
|
||
data = {
|
||
"version": SETTINGS_VERSION,
|
||
"models": {
|
||
model_id: settings.to_dict()
|
||
for model_id, settings in self._settings.items()
|
||
},
|
||
}
|
||
|
||
try:
|
||
# Write to temp file first, then rename for atomicity
|
||
temp_file = self.settings_file.with_suffix(".tmp")
|
||
with open(temp_file, "w", encoding="utf-8") as f:
|
||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||
|
||
temp_file.replace(self.settings_file)
|
||
logger.debug(f"Saved settings for {len(self._settings)} models")
|
||
|
||
except Exception as e:
|
||
logger.error(f"Failed to save settings file: {e}")
|
||
raise
|
||
|
||
def get_settings(self, model_id: str) -> ModelSettings:
|
||
"""Get settings for a specific model.
|
||
|
||
Args:
|
||
model_id: The model identifier.
|
||
|
||
Returns:
|
||
ModelSettings for the model, or default settings if not found.
|
||
"""
|
||
with self._lock:
|
||
if model_id in self._settings:
|
||
# Return a copy to prevent external modification
|
||
settings = self._settings[model_id]
|
||
return ModelSettings.from_dict(settings.to_dict())
|
||
|
||
return ModelSettings()
|
||
|
||
def get_settings_for_request(
|
||
self,
|
||
model_id: str,
|
||
resolved_model_id: Optional[str] = None,
|
||
) -> ModelSettings:
|
||
"""Get settings for an API-requested model name.
|
||
|
||
Exposed profile model IDs return the base model's settings merged
|
||
with the profile's universal (request-time) overrides. Engine-
|
||
construction fields are handled separately by
|
||
get_exposed_profile_runtime_settings_for_request(), which may trigger
|
||
a transient engine variant reload without mutating persisted settings.
|
||
Any other name falls back to the settings of the already-resolved
|
||
physical model.
|
||
"""
|
||
with self._lock:
|
||
candidates = [model_id]
|
||
if "/" in model_id:
|
||
candidates.append(model_id.split("/", 1)[1])
|
||
for candidate in candidates:
|
||
profile_match = self._find_exposed_profile_locked(candidate)
|
||
if profile_match is not None:
|
||
base_model_id, profile = profile_match
|
||
return self._settings_with_profile_locked(base_model_id, profile)
|
||
|
||
return self.get_settings(resolved_model_id or model_id)
|
||
|
||
def set_settings(self, model_id: str, settings: ModelSettings) -> None:
|
||
"""Set settings for a specific model.
|
||
|
||
If the new settings have is_default=True, clears is_default from all
|
||
other models to maintain the exclusive default constraint.
|
||
|
||
Args:
|
||
model_id: The model identifier.
|
||
settings: The settings to apply.
|
||
"""
|
||
with self._lock:
|
||
# Handle exclusive default constraint
|
||
if settings.is_default:
|
||
for mid, s in self._settings.items():
|
||
if mid != model_id and s.is_default:
|
||
s.is_default = False
|
||
logger.info(
|
||
f"Cleared is_default from model '{mid}' "
|
||
f"(new default: '{model_id}')"
|
||
)
|
||
|
||
# Store a copy of the settings
|
||
self._settings[model_id] = ModelSettings.from_dict(settings.to_dict())
|
||
logger.info(f"Updated settings for model '{model_id}'")
|
||
|
||
self._save()
|
||
|
||
def delete_settings(self, model_id: str) -> bool:
|
||
"""Remove all persisted state for a model (settings + profiles).
|
||
|
||
Called when a model is deleted so its alias and other settings are
|
||
released and can be reused by another model.
|
||
|
||
Args:
|
||
model_id: The model identifier.
|
||
|
||
Returns:
|
||
True if any state was removed, False if nothing was stored.
|
||
"""
|
||
with self._lock:
|
||
removed = False
|
||
if model_id in self._settings:
|
||
del self._settings[model_id]
|
||
self._save()
|
||
removed = True
|
||
if model_id in self._profiles:
|
||
del self._profiles[model_id]
|
||
self._save_profiles()
|
||
removed = True
|
||
if removed:
|
||
logger.info(f"Deleted settings for model '{model_id}'")
|
||
return removed
|
||
|
||
def get_default_model_id(self) -> Optional[str]:
|
||
"""Get the ID of the default model.
|
||
|
||
Returns:
|
||
The model ID marked as default, or None if no default is set.
|
||
"""
|
||
with self._lock:
|
||
for model_id, settings in self._settings.items():
|
||
if settings.is_default:
|
||
return model_id
|
||
return None
|
||
|
||
def get_pinned_model_ids(self) -> list[str]:
|
||
"""Get list of all pinned model IDs.
|
||
|
||
Returns:
|
||
List of model IDs that are marked as pinned.
|
||
"""
|
||
with self._lock:
|
||
return [
|
||
model_id
|
||
for model_id, settings in self._settings.items()
|
||
if settings.is_pinned
|
||
]
|
||
|
||
def get_all_settings(self) -> Dict[str, ModelSettings]:
|
||
"""Get a copy of all model settings.
|
||
|
||
Returns:
|
||
Dictionary mapping model IDs to their settings (deep copy).
|
||
"""
|
||
with self._lock:
|
||
return {
|
||
model_id: ModelSettings.from_dict(settings.to_dict())
|
||
for model_id, settings in self._settings.items()
|
||
}
|
||
|
||
# ==================== Profiles ====================
|
||
|
||
def _load_profiles(self) -> None:
|
||
if not self.profiles_file.exists():
|
||
self._profiles = {}
|
||
return
|
||
try:
|
||
with open(self.profiles_file, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
version = data.get("version", 1)
|
||
if version != PROFILES_VERSION:
|
||
logger.warning(
|
||
f"Profiles file version {version} differs from current {PROFILES_VERSION}"
|
||
)
|
||
self._profiles = data.get("profiles", {}) or {}
|
||
# Migration: strip ttl_seconds from existing profile settings and
|
||
# add api_name for exposed-model IDs without changing internal keys.
|
||
changed = False
|
||
for model_id, profiles in self._profiles.items():
|
||
model_changed = False
|
||
used_api_names: set[str] = set()
|
||
for name, profile in profiles.items():
|
||
settings = profile.get("settings")
|
||
if settings and "ttl_seconds" in settings:
|
||
del settings["ttl_seconds"]
|
||
changed = True
|
||
model_changed = True
|
||
current_api_name = profile.get("api_name")
|
||
if current_api_name:
|
||
try:
|
||
validate_profile_name(current_api_name)
|
||
base_api_name = current_api_name
|
||
except Exception:
|
||
base_api_name = slugify_profile_api_name(
|
||
profile.get("display_name") or name,
|
||
fallback="profile",
|
||
)
|
||
else:
|
||
base_api_name = slugify_profile_api_name(
|
||
profile.get("display_name") or name,
|
||
fallback="profile",
|
||
)
|
||
api_name = self._dedupe_profile_api_name(
|
||
base_api_name,
|
||
used_api_names,
|
||
)
|
||
if current_api_name != api_name:
|
||
profile["api_name"] = api_name
|
||
changed = True
|
||
model_changed = True
|
||
if model_changed:
|
||
logger.info(f"Migrated profile api names for model '{model_id}'")
|
||
if changed:
|
||
self._save_profiles()
|
||
except Exception as e:
|
||
logger.error(f"Failed to load profiles file: {e}")
|
||
self._profiles = {}
|
||
|
||
def _save_profiles(self) -> None:
|
||
"""Write profiles to disk atomically (temp file + rename)."""
|
||
data = {"version": PROFILES_VERSION, "profiles": self._profiles}
|
||
temp_file = self.profiles_file.with_suffix(".tmp")
|
||
try:
|
||
with open(temp_file, "w", encoding="utf-8") as f:
|
||
json.dump(data, f, indent=2, ensure_ascii=False, default=str)
|
||
temp_file.replace(self.profiles_file)
|
||
except Exception as e:
|
||
logger.error(f"Failed to save profiles file: {e}")
|
||
if temp_file.exists():
|
||
temp_file.unlink(missing_ok=True)
|
||
raise
|
||
|
||
@staticmethod
|
||
def _dedupe_profile_api_name(base: str, used: set[str]) -> str:
|
||
validate_profile_name(base)
|
||
candidate = base
|
||
index = 2
|
||
while candidate in used:
|
||
suffix = f"-{index}"
|
||
root = base[: 32 - len(suffix)].rstrip("-_") or "profile"
|
||
candidate = f"{root}{suffix}"
|
||
index += 1
|
||
used.add(candidate)
|
||
return candidate
|
||
|
||
@staticmethod
|
||
def _profile_api_name(profile: Dict[str, Any]) -> str:
|
||
api_name = profile.get("api_name") or profile["name"]
|
||
validate_profile_name(api_name)
|
||
return api_name
|
||
|
||
def _allocate_profile_api_name_locked(
|
||
self,
|
||
profiles: Dict[str, Dict[str, Any]],
|
||
value: str | None,
|
||
*,
|
||
display_name: str | None,
|
||
internal_name: str,
|
||
exclude_name: str | None = None,
|
||
) -> str:
|
||
if value:
|
||
validate_profile_name(value)
|
||
base = value
|
||
else:
|
||
base = slugify_profile_api_name(
|
||
display_name or internal_name,
|
||
fallback="profile",
|
||
)
|
||
used = {
|
||
self._profile_api_name(profile)
|
||
for name, profile in profiles.items()
|
||
if name != exclude_name
|
||
}
|
||
return self._dedupe_profile_api_name(base, used)
|
||
|
||
def _profile_model_id(self, model_id: str, api_name: str) -> str:
|
||
return f"{model_id}:{api_name}"
|
||
|
||
def _display_profile_model_id_locked(
|
||
self, model_id: str, profile: Dict[str, Any]
|
||
) -> str:
|
||
"""Advertised form of an exposed profile's model ID.
|
||
|
||
Uses the base model's alias when set, mirroring how /v1/models
|
||
lists the base model itself. The directory-name form remains
|
||
accepted for requests, exactly like the base model's directory
|
||
name.
|
||
"""
|
||
base = self._settings.get(model_id)
|
||
display_base = base.model_alias if base and base.model_alias else model_id
|
||
return self._profile_model_id(display_base, self._profile_api_name(profile))
|
||
|
||
def _find_exposed_profile_locked(
|
||
self, model_id: str
|
||
) -> Optional[tuple[str, Dict[str, Any]]]:
|
||
for base_model_id, profiles in self._profiles.items():
|
||
base = self._settings.get(base_model_id)
|
||
alias = base.model_alias if base else None
|
||
for profile in profiles.values():
|
||
if not profile.get("expose_as_model"):
|
||
continue
|
||
api_name = self._profile_api_name(profile)
|
||
if model_id == self._profile_model_id(base_model_id, api_name):
|
||
return base_model_id, profile
|
||
if alias and model_id == self._profile_model_id(alias, api_name):
|
||
return base_model_id, profile
|
||
return None
|
||
|
||
def _settings_with_profile_locked(
|
||
self, model_id: str, profile: Dict[str, Any]
|
||
) -> ModelSettings:
|
||
base = self._settings.get(model_id)
|
||
merged = base.to_dict() if base is not None else {}
|
||
# Request-time settings still use only universal fields. Engine-
|
||
# construction fields are handled separately by
|
||
# get_exposed_profile_runtime_settings_for_request(), which can
|
||
# trigger an engine variant reload without persisting base settings.
|
||
merged.update(filter_universal_fields(profile.get("settings", {}) or {}))
|
||
return ModelSettings.from_dict(merged)
|
||
|
||
def _runtime_settings_with_profile_locked(
|
||
self, model_id: str, profile: Dict[str, Any]
|
||
) -> ModelSettings:
|
||
base = self._settings.get(model_id)
|
||
merged = base.to_dict() if base is not None else {}
|
||
merged.update(filter_profile_fields(profile.get("settings", {}) or {}))
|
||
return ModelSettings.from_dict(merged)
|
||
|
||
def get_exposed_profile_source_model_id(self, model_id: str) -> Optional[str]:
|
||
"""Return the base model for an exposed profile model ID, if any."""
|
||
with self._lock:
|
||
candidates = [model_id]
|
||
if "/" in model_id:
|
||
candidates.append(model_id.split("/", 1)[1])
|
||
for candidate in candidates:
|
||
match = self._find_exposed_profile_locked(candidate)
|
||
if match is not None:
|
||
return match[0]
|
||
return None
|
||
|
||
def get_exposed_profile_runtime_settings_for_request(
|
||
self,
|
||
model_id: str,
|
||
) -> Optional[tuple[str, ModelSettings]]:
|
||
"""Return full runtime settings for an exposed profile request.
|
||
|
||
Unlike ``get_settings_for_request()``, this includes model-specific
|
||
engine-construction fields. It is used only for transient engine
|
||
loading and never mutates the base model's persisted settings.
|
||
"""
|
||
with self._lock:
|
||
candidates = [model_id]
|
||
if "/" in model_id:
|
||
candidates.append(model_id.split("/", 1)[1])
|
||
for candidate in candidates:
|
||
match = self._find_exposed_profile_locked(candidate)
|
||
if match is not None:
|
||
base_model_id, profile = match
|
||
return (
|
||
base_model_id,
|
||
self._runtime_settings_with_profile_locked(
|
||
base_model_id, profile
|
||
),
|
||
)
|
||
return None
|
||
|
||
def get_exposed_profile_model_ids(
|
||
self,
|
||
*,
|
||
exclude_model_id: str | None = None,
|
||
exclude_profile_name: str | None = None,
|
||
) -> set[str]:
|
||
"""Return every request ID accepted by exposed profiles."""
|
||
with self._lock:
|
||
model_ids: set[str] = set()
|
||
for base_model_id, profiles in self._profiles.items():
|
||
base = self._settings.get(base_model_id)
|
||
alias = base.model_alias if base else None
|
||
for profile in profiles.values():
|
||
if not profile.get("expose_as_model"):
|
||
continue
|
||
if (
|
||
exclude_model_id == base_model_id
|
||
and exclude_profile_name == profile["name"]
|
||
):
|
||
continue
|
||
api_name = self._profile_api_name(profile)
|
||
model_ids.add(self._profile_model_id(base_model_id, api_name))
|
||
if alias:
|
||
model_ids.add(self._profile_model_id(alias, api_name))
|
||
return model_ids
|
||
|
||
def _profile_request_ids_locked(
|
||
self,
|
||
model_id: str,
|
||
profile: Dict[str, Any],
|
||
) -> set[str]:
|
||
base = self._settings.get(model_id)
|
||
base_ids = {model_id}
|
||
if base and base.model_alias:
|
||
base_ids.add(base.model_alias)
|
||
api_name = self._profile_api_name(profile)
|
||
return {self._profile_model_id(base_id, api_name) for base_id in base_ids}
|
||
|
||
def _validate_exposed_profile_ids_available_locked(
|
||
self,
|
||
model_id: str,
|
||
profile: Dict[str, Any],
|
||
*,
|
||
exclude_profile_name: str | None = None,
|
||
reserved_model_ids: set[str] | None = None,
|
||
) -> None:
|
||
if not profile.get("expose_as_model"):
|
||
return
|
||
candidate_ids = self._profile_request_ids_locked(model_id, profile)
|
||
|
||
if reserved_model_ids:
|
||
for candidate_id in candidate_ids:
|
||
if candidate_id != model_id and candidate_id in reserved_model_ids:
|
||
raise ValueError(
|
||
f"Exposed profile model ID '{candidate_id}' conflicts "
|
||
"with a model directory name"
|
||
)
|
||
|
||
for mid, settings in self._settings.items():
|
||
if settings.model_alias and settings.model_alias in candidate_ids:
|
||
raise ValueError(
|
||
f"Exposed profile model ID '{settings.model_alias}' "
|
||
f"conflicts with model alias for '{mid}'"
|
||
)
|
||
|
||
existing_ids: set[str] = set()
|
||
for base_model_id, profiles in self._profiles.items():
|
||
for other in profiles.values():
|
||
if not other.get("expose_as_model"):
|
||
continue
|
||
if base_model_id == model_id and other["name"] == exclude_profile_name:
|
||
continue
|
||
existing_ids.update(
|
||
self._profile_request_ids_locked(base_model_id, other)
|
||
)
|
||
conflict = candidate_ids & existing_ids
|
||
if conflict:
|
||
conflict_id = sorted(conflict)[0]
|
||
raise ValueError(f"Exposed profile model ID '{conflict_id}' already exists")
|
||
|
||
def list_exposed_profile_models(self) -> list[dict]:
|
||
"""Return profile records promoted to independently visible model IDs."""
|
||
with self._lock:
|
||
exposed = []
|
||
for base_model_id, profiles in self._profiles.items():
|
||
for profile in profiles.values():
|
||
if not profile.get("expose_as_model"):
|
||
continue
|
||
item = dict(profile)
|
||
item["model_id"] = self._display_profile_model_id_locked(
|
||
base_model_id, profile
|
||
)
|
||
item["source_model_id"] = base_model_id
|
||
item["settings"] = self._settings_with_profile_locked(
|
||
base_model_id, item
|
||
).to_dict()
|
||
exposed.append(item)
|
||
return exposed
|
||
|
||
@staticmethod
|
||
def _has_engine_fields(profile: Dict[str, Any]) -> bool:
|
||
"""True when the profile overrides any engine-construction field.
|
||
|
||
Those fields are ignored by the request-time sampling overlay but are
|
||
used for transient engine variant reloads on exposed-profile requests.
|
||
UIs use this flag to indicate that a profile changes load-time state.
|
||
"""
|
||
settings = profile.get("settings", {}) or {}
|
||
return any(
|
||
k in MODEL_SPECIFIC_PROFILE_FIELDS and v is not None
|
||
for k, v in settings.items()
|
||
)
|
||
|
||
def list_profiles(self, model_id: str) -> list[dict]:
|
||
"""Return all profiles for ``model_id`` as serializable dicts."""
|
||
with self._lock:
|
||
per_model = self._profiles.get(model_id, {})
|
||
return [
|
||
{
|
||
**p,
|
||
"model_id": self._display_profile_model_id_locked(model_id, p),
|
||
"has_engine_fields": self._has_engine_fields(p),
|
||
}
|
||
for p in per_model.values()
|
||
]
|
||
|
||
def get_profile(self, model_id: str, name: str) -> Optional[dict]:
|
||
with self._lock:
|
||
return dict(self._profiles.get(model_id, {}).get(name, {})) or None
|
||
|
||
def save_profile(
|
||
self,
|
||
model_id: str,
|
||
name: str,
|
||
display_name: str,
|
||
description: Optional[str],
|
||
settings: Dict[str, Any],
|
||
source_template: Optional[str] = None,
|
||
expose_as_model: bool = False,
|
||
api_name: Optional[str] = None,
|
||
reserved_model_ids: Optional[set[str]] = None,
|
||
) -> dict:
|
||
"""Create a new profile. Raises if name is invalid or already exists."""
|
||
validate_profile_name(name)
|
||
filtered = filter_profile_fields(settings or {})
|
||
with self._lock:
|
||
per_model = self._profiles.setdefault(model_id, {})
|
||
if name in per_model:
|
||
raise ValueError(
|
||
f"Profile '{name}' already exists for model '{model_id}'"
|
||
)
|
||
now = utcnow().isoformat()
|
||
profile_api_name = self._allocate_profile_api_name_locked(
|
||
per_model,
|
||
api_name,
|
||
display_name=display_name,
|
||
internal_name=name,
|
||
)
|
||
profile_record = {
|
||
"name": name,
|
||
"display_name": display_name or name,
|
||
"api_name": profile_api_name,
|
||
"description": description,
|
||
"created_at": now,
|
||
"updated_at": now,
|
||
"settings": filtered,
|
||
"source_template": source_template,
|
||
"expose_as_model": bool(expose_as_model),
|
||
}
|
||
self._validate_exposed_profile_ids_available_locked(
|
||
model_id,
|
||
profile_record,
|
||
reserved_model_ids=reserved_model_ids,
|
||
)
|
||
per_model[name] = profile_record
|
||
self._save_profiles()
|
||
return dict(per_model[name])
|
||
|
||
def update_profile(
|
||
self,
|
||
model_id: str,
|
||
name: str,
|
||
*,
|
||
new_name: Optional[str] = None,
|
||
display_name: Optional[str] = None,
|
||
description: Optional[str] = None,
|
||
settings: Optional[Dict[str, Any]] = None,
|
||
source_template: Optional[str] = None,
|
||
expose_as_model: Optional[bool] = None,
|
||
api_name: Optional[str] = None,
|
||
reserved_model_ids: Optional[set[str]] = None,
|
||
) -> Optional[dict]:
|
||
"""Update a profile's metadata/settings. Returns updated dict or None if not found."""
|
||
with self._lock:
|
||
per_model = self._profiles.get(model_id, {})
|
||
if name not in per_model:
|
||
return None
|
||
profile = dict(per_model[name])
|
||
target_name = name
|
||
rename_mode = False
|
||
if new_name is not None and new_name != name:
|
||
validate_profile_name(new_name)
|
||
if new_name in per_model:
|
||
raise ValueError(
|
||
f"Profile '{new_name}' already exists for model '{model_id}'"
|
||
)
|
||
target_name = new_name
|
||
profile["name"] = new_name
|
||
rename_mode = True
|
||
if display_name is not None:
|
||
profile["display_name"] = display_name
|
||
if api_name is not None:
|
||
profile["api_name"] = self._allocate_profile_api_name_locked(
|
||
per_model,
|
||
api_name,
|
||
display_name=profile.get("display_name"),
|
||
internal_name=target_name,
|
||
exclude_name=name,
|
||
)
|
||
if description is not None:
|
||
profile["description"] = description
|
||
if settings is not None:
|
||
profile["settings"] = filter_profile_fields(settings)
|
||
if source_template is not None:
|
||
profile["source_template"] = source_template or None
|
||
if expose_as_model is not None:
|
||
profile["expose_as_model"] = bool(expose_as_model)
|
||
profile["updated_at"] = utcnow().isoformat()
|
||
self._validate_exposed_profile_ids_available_locked(
|
||
model_id,
|
||
profile,
|
||
exclude_profile_name=name,
|
||
reserved_model_ids=reserved_model_ids,
|
||
)
|
||
|
||
# Snapshot for rollback on write failure
|
||
profiles_snapshot = copy.deepcopy(self._profiles)
|
||
settings_snapshot = copy.deepcopy(self._settings)
|
||
|
||
# Also update ModelSettings.active_profile_name if renamed and it was active
|
||
old_active = None
|
||
if rename_mode:
|
||
old_active = self._settings.get(model_id)
|
||
if old_active is not None and old_active.active_profile_name == name:
|
||
old_active.active_profile_name = target_name
|
||
del per_model[name]
|
||
|
||
per_model[target_name] = profile
|
||
|
||
# Write profiles first; if this throws, rollback everything
|
||
try:
|
||
self._save_profiles()
|
||
if rename_mode and old_active is not None:
|
||
self._save()
|
||
except Exception:
|
||
self._profiles = profiles_snapshot
|
||
self._settings = settings_snapshot
|
||
raise
|
||
|
||
return dict(profile)
|
||
|
||
def delete_profile(self, model_id: str, name: str) -> bool:
|
||
with self._lock:
|
||
per_model = self._profiles.get(model_id, {})
|
||
if name not in per_model:
|
||
return False
|
||
|
||
profiles_snapshot = copy.deepcopy(self._profiles)
|
||
settings_snapshot = copy.deepcopy(self._settings)
|
||
|
||
del per_model[name]
|
||
if not per_model and model_id in self._profiles:
|
||
del self._profiles[model_id]
|
||
# Clear active_profile_name if it referenced this profile
|
||
old_active = self._settings.get(model_id)
|
||
if old_active is not None and old_active.active_profile_name == name:
|
||
old_active.active_profile_name = None
|
||
|
||
try:
|
||
self._save_profiles()
|
||
if old_active is not None and old_active.active_profile_name is None:
|
||
self._save()
|
||
except Exception:
|
||
self._profiles = profiles_snapshot
|
||
self._settings = settings_snapshot
|
||
raise
|
||
return True
|
||
|
||
def apply_profile(
|
||
self,
|
||
model_id: str,
|
||
name: str,
|
||
settings_sanitizer: Optional[Callable[[dict[str, Any]], None]] = None,
|
||
) -> Optional[ModelSettings]:
|
||
"""Merge profile settings into the model's live settings and persist."""
|
||
with self._lock:
|
||
per_model = self._profiles.get(model_id, {})
|
||
if name not in per_model:
|
||
return None
|
||
profile_settings = per_model[name].get("settings", {}) or {}
|
||
|
||
settings_snapshot = copy.deepcopy(self._settings)
|
||
|
||
current = self._settings.get(model_id)
|
||
if current is None:
|
||
current = ModelSettings()
|
||
merged = current.to_dict()
|
||
for k, v in profile_settings.items():
|
||
merged[k] = v
|
||
merged["active_profile_name"] = name
|
||
if settings_sanitizer is not None:
|
||
settings_sanitizer(merged)
|
||
new_settings = ModelSettings.from_dict(merged)
|
||
self._settings[model_id] = new_settings
|
||
try:
|
||
self._save()
|
||
except Exception:
|
||
self._settings = settings_snapshot
|
||
raise
|
||
return ModelSettings.from_dict(new_settings.to_dict())
|
||
|
||
# ==================== Templates ====================
|
||
|
||
def _load_templates(self) -> None:
|
||
# Built-in defaults ship inside the package (omlx/default_global_templates.json)
|
||
# and are merged in at read time — they are NEVER copied to disk and never
|
||
# appear in `self._templates`. The user file under <base_path> holds
|
||
# ONLY user-created templates; a missing/empty file is the legitimate
|
||
# initial state.
|
||
if not self.templates_file.exists():
|
||
self._templates = {}
|
||
return
|
||
try:
|
||
with open(self.templates_file, "r", encoding="utf-8") as f:
|
||
data = json.load(f)
|
||
version = data.get("version", 1)
|
||
if version != TEMPLATES_VERSION:
|
||
logger.warning(
|
||
f"Templates file version {version} differs from current {TEMPLATES_VERSION}"
|
||
)
|
||
self._templates = data.get("templates", {}) or {}
|
||
# Migration: strip ttl_seconds from existing template settings
|
||
for name, template in self._templates.items():
|
||
settings = template.get("settings")
|
||
if settings and "ttl_seconds" in settings:
|
||
del settings["ttl_seconds"]
|
||
except Exception as e:
|
||
logger.error(f"Failed to load templates file: {e}")
|
||
self._templates = {}
|
||
|
||
def _save_templates(self) -> None:
|
||
"""Must be called while holding the lock."""
|
||
data = {"version": TEMPLATES_VERSION, "templates": self._templates}
|
||
try:
|
||
temp_file = self.templates_file.with_suffix(".tmp")
|
||
with open(temp_file, "w", encoding="utf-8") as f:
|
||
json.dump(data, f, indent=2, ensure_ascii=False, default=str)
|
||
temp_file.replace(self.templates_file)
|
||
except Exception as e:
|
||
logger.error(f"Failed to save templates file: {e}")
|
||
raise
|
||
|
||
def list_templates(self) -> list[dict]:
|
||
# Shipped JSON seeds were retired in favor of the client-side preset
|
||
# bundle (`omlx/admin/static/omlx_preset.json`); every entry on this
|
||
# surface is user-created. Callers that distinguish presets from
|
||
# user templates do so via the preset bundle, not an `is_builtin`
|
||
# flag on this response.
|
||
with self._lock:
|
||
return [dict(t) for t in self._templates.values()]
|
||
|
||
def get_template(self, name: str) -> Optional[dict]:
|
||
with self._lock:
|
||
u = self._templates.get(name)
|
||
return dict(u) if u is not None else None
|
||
|
||
def save_template(
|
||
self,
|
||
name: str,
|
||
display_name: str,
|
||
description: Optional[str],
|
||
settings: Dict[str, Any],
|
||
) -> dict:
|
||
validate_profile_name(name)
|
||
filtered = filter_universal_fields(settings or {})
|
||
with self._lock:
|
||
if name in self._templates:
|
||
raise ValueError(f"Template '{name}' already exists")
|
||
now = utcnow().isoformat()
|
||
self._templates[name] = {
|
||
"name": name,
|
||
"display_name": display_name or name,
|
||
"description": description,
|
||
"created_at": now,
|
||
"updated_at": now,
|
||
"settings": filtered,
|
||
}
|
||
self._save_templates()
|
||
return dict(self._templates[name])
|
||
|
||
def upsert_template(
|
||
self,
|
||
name: str,
|
||
display_name: str,
|
||
description: Optional[str],
|
||
settings: Dict[str, Any],
|
||
) -> dict:
|
||
"""Create or replace a template with the given settings."""
|
||
validate_profile_name(name)
|
||
filtered = filter_universal_fields(settings or {})
|
||
with self._lock:
|
||
now = utcnow().isoformat()
|
||
existing = self._templates.get(name)
|
||
created_at = existing["created_at"] if existing else now
|
||
self._templates[name] = {
|
||
"name": name,
|
||
"display_name": display_name or name,
|
||
"description": description,
|
||
"created_at": created_at,
|
||
"updated_at": now,
|
||
"settings": filtered,
|
||
}
|
||
self._save_templates()
|
||
return dict(self._templates[name])
|
||
|
||
def update_template(
|
||
self,
|
||
name: str,
|
||
*,
|
||
new_name: Optional[str] = None,
|
||
display_name: Optional[str] = None,
|
||
description: Optional[str] = None,
|
||
settings: Optional[Dict[str, Any]] = None,
|
||
) -> Optional[dict]:
|
||
with self._lock:
|
||
if name not in self._templates:
|
||
return None
|
||
template = dict(self._templates[name])
|
||
target = name
|
||
if new_name is not None and new_name != name:
|
||
validate_profile_name(new_name)
|
||
if new_name in self._templates:
|
||
raise ValueError(f"Template '{new_name}' already exists")
|
||
target = new_name
|
||
template["name"] = new_name
|
||
if display_name is not None:
|
||
template["display_name"] = display_name
|
||
if description is not None:
|
||
template["description"] = description
|
||
if settings is not None:
|
||
template["settings"] = filter_universal_fields(settings)
|
||
template["updated_at"] = utcnow().isoformat()
|
||
if target != name:
|
||
del self._templates[name]
|
||
self._templates[target] = template
|
||
self._save_templates()
|
||
return dict(template)
|
||
|
||
def delete_template(self, name: str) -> bool:
|
||
with self._lock:
|
||
if name not in self._templates:
|
||
return False
|
||
del self._templates[name]
|
||
self._save_templates()
|
||
return True
|