# SPDX-License-Identifier: Apache-2.0 """ Server-level metrics for the oMLX admin dashboard. Provides a thread-safe singleton that aggregates serving metrics across all engines/models. Session metrics reset on server start, while all-time metrics persist across restarts via JSON file. """ import json import logging import threading import time from pathlib import Path from typing import Any, Dict, Optional logger = logging.getLogger(__name__) # Interval between periodic saves of all-time stats (seconds) _SAVE_INTERVAL = 300 class ServerMetrics: """ Global server-level metrics for the Status dashboard. Thread-safe: uses threading.Lock since scheduler runs in ThreadPoolExecutor. Tracks cumulative totals and average speeds across all requests, with optional per-model breakdown. Supports two scopes: - "session": resets on server restart (default, backward-compatible) - "alltime": persisted across restarts via stats_path JSON file """ def __init__(self, stats_path: Optional[Path] = None): self._lock = threading.Lock() self._stats_path = stats_path # Session totals (reset on server restart or clear) self.total_prompt_tokens: int = 0 self.total_completion_tokens: int = 0 self.total_cached_tokens: int = 0 self.total_requests: int = 0 self.total_prefill_duration: float = 0.0 self.total_generation_duration: float = 0.0 self._per_model: Dict[str, Dict[str, Any]] = {} # Preflight rejections — split by reason so operators can tell # "user fed an oversize prompt" (hard_limit) apart from "system # under memory pressure, admission paused" (admission_paused). # This is the only operator-visible signal for the prefill-peak # memory guard; the guard was previously dead and shipping it # without telemetry repeats that mistake. self.preflight_rejections: Dict[str, int] = { "hard_limit": 0, "admission_paused": 0, } # All-time totals (persisted across restarts) self._alltime_prompt_tokens: int = 0 self._alltime_completion_tokens: int = 0 self._alltime_cached_tokens: int = 0 self._alltime_requests: int = 0 self._alltime_prefill_duration: float = 0.0 self._alltime_generation_duration: float = 0.0 self._alltime_per_model: Dict[str, Dict[str, Any]] = {} self._start_time = time.time() self._last_save_time = time.time() # Load persisted all-time stats if stats_path: self._load_alltime() @staticmethod def _new_model_counters() -> Dict[str, Any]: return { "prompt_tokens": 0, "completion_tokens": 0, "cached_tokens": 0, "requests": 0, "prefill_duration": 0.0, "generation_duration": 0.0, } def _load_alltime(self) -> None: """Load all-time stats from disk. Called once during __init__.""" if not self._stats_path or not self._stats_path.exists(): return try: with open(self._stats_path) as f: data = json.load(f) self._alltime_prompt_tokens = int(data.get("total_prompt_tokens", 0)) self._alltime_completion_tokens = int( data.get("total_completion_tokens", 0) ) self._alltime_cached_tokens = int(data.get("total_cached_tokens", 0)) self._alltime_requests = int(data.get("total_requests", 0)) self._alltime_prefill_duration = float( data.get("total_prefill_duration", 0.0) ) self._alltime_generation_duration = float( data.get("total_generation_duration", 0.0) ) per_model = data.get("per_model", {}) for model_id, counters in per_model.items(): self._alltime_per_model[model_id] = { "prompt_tokens": int(counters.get("prompt_tokens", 0)), "completion_tokens": int(counters.get("completion_tokens", 0)), "cached_tokens": int(counters.get("cached_tokens", 0)), "requests": int(counters.get("requests", 0)), "prefill_duration": float(counters.get("prefill_duration", 0.0)), "generation_duration": float( counters.get("generation_duration", 0.0) ), } logger.info("Loaded all-time stats from %s", self._stats_path) except (json.JSONDecodeError, TypeError, KeyError, ValueError, OSError) as e: logger.warning("Failed to load all-time stats from %s: %s", self._stats_path, e) def save_alltime(self) -> None: """Save all-time stats to disk. Thread-safe.""" if not self._stats_path: return with self._lock: data = { "total_prompt_tokens": self._alltime_prompt_tokens, "total_completion_tokens": self._alltime_completion_tokens, "total_cached_tokens": self._alltime_cached_tokens, "total_requests": self._alltime_requests, "total_prefill_duration": self._alltime_prefill_duration, "total_generation_duration": self._alltime_generation_duration, "per_model": dict(self._alltime_per_model), } self._last_save_time = time.time() try: self._stats_path.parent.mkdir(parents=True, exist_ok=True) tmp_path = self._stats_path.with_suffix(".json.tmp") with open(tmp_path, "w") as f: json.dump(data, f, indent=2) tmp_path.replace(self._stats_path) except OSError as e: logger.warning("Failed to save all-time stats to %s: %s", self._stats_path, e) def _maybe_save_alltime(self) -> None: """Save all-time stats if enough time has passed. Called within lock.""" if not self._stats_path: return now = time.time() if now - self._last_save_time >= _SAVE_INTERVAL: # Release lock before I/O self._lock.release() try: self.save_alltime() finally: self._lock.acquire() def record_request_complete( self, prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0, prefill_duration: float = 0.0, generation_duration: float = 0.0, model_id: str = "", ) -> None: """Record a completed request. Thread-safe.""" with self._lock: # Session counters self.total_prompt_tokens += prompt_tokens self.total_completion_tokens += completion_tokens self.total_cached_tokens += cached_tokens self.total_requests += 1 self.total_prefill_duration += prefill_duration self.total_generation_duration += generation_duration # All-time counters self._alltime_prompt_tokens += prompt_tokens self._alltime_completion_tokens += completion_tokens self._alltime_cached_tokens += cached_tokens self._alltime_requests += 1 self._alltime_prefill_duration += prefill_duration self._alltime_generation_duration += generation_duration # Per-model counters (session) if model_id: if model_id not in self._per_model: self._per_model[model_id] = self._new_model_counters() m = self._per_model[model_id] m["prompt_tokens"] += prompt_tokens m["completion_tokens"] += completion_tokens m["cached_tokens"] += cached_tokens m["requests"] += 1 m["prefill_duration"] += prefill_duration m["generation_duration"] += generation_duration # Per-model counters (all-time) if model_id not in self._alltime_per_model: self._alltime_per_model[model_id] = self._new_model_counters() am = self._alltime_per_model[model_id] am["prompt_tokens"] += prompt_tokens am["completion_tokens"] += completion_tokens am["cached_tokens"] += cached_tokens am["requests"] += 1 am["prefill_duration"] += prefill_duration am["generation_duration"] += generation_duration # Periodic save self._maybe_save_alltime() def record_preflight_rejection(self, reason: str) -> None: """Increment the preflight-rejection counter for ``reason``. Unknown reasons are bucketed under ``"other"`` rather than silently dropped so an operator notices when a new reject path is added without updating the metric. """ with self._lock: if reason not in self.preflight_rejections: reason = "other" self.preflight_rejections.setdefault("other", 0) self.preflight_rejections[reason] += 1 def _build_snapshot( self, prompt: int, completion: int, cached: int, requests: int, prefill_dur: float, gen_dur: float, uptime: float, ) -> Dict[str, Any]: """Build a metrics snapshot dict from raw values.""" actual_processed = prompt - cached avg_prefill_tps = ( actual_processed / prefill_dur if prefill_dur > 0 else 0.0 ) avg_generation_tps = completion / gen_dur if gen_dur > 0 else 0.0 cache_efficiency = (cached / prompt * 100) if prompt > 0 else 0.0 return { "total_tokens_served": prompt + completion, "total_cached_tokens": cached, "cache_efficiency": round(cache_efficiency, 1), "total_prompt_tokens": prompt, "total_completion_tokens": completion, "total_requests": requests, "avg_prefill_tps": round(avg_prefill_tps, 1), "avg_generation_tps": round(avg_generation_tps, 1), "uptime_seconds": round(uptime, 1), } def get_snapshot( self, model_id: str = "", scope: str = "session" ) -> Dict[str, Any]: """Get current metrics snapshot. Thread-safe. Args: model_id: If provided and tracked, return per-model metrics. Otherwise return global aggregate. scope: "session" for current session, "alltime" for persisted totals. """ with self._lock: now = time.time() uptime = now - self._start_time if scope == "alltime": if model_id: m = self._alltime_per_model.get(model_id) if m: return self._build_snapshot( m["prompt_tokens"], m["completion_tokens"], m["cached_tokens"], m["requests"], m["prefill_duration"], m["generation_duration"], uptime, ) return self._build_snapshot(0, 0, 0, 0, 0.0, 0.0, uptime) return self._build_snapshot( self._alltime_prompt_tokens, self._alltime_completion_tokens, self._alltime_cached_tokens, self._alltime_requests, self._alltime_prefill_duration, self._alltime_generation_duration, uptime, ) # scope == "session" (default) if model_id: m = self._per_model.get(model_id) if m: return self._build_snapshot( m["prompt_tokens"], m["completion_tokens"], m["cached_tokens"], m["requests"], m["prefill_duration"], m["generation_duration"], uptime, ) return self._build_snapshot(0, 0, 0, 0, 0.0, 0.0, uptime) return self._build_snapshot( self.total_prompt_tokens, self.total_completion_tokens, self.total_cached_tokens, self.total_requests, self.total_prefill_duration, self.total_generation_duration, uptime, ) def clear_metrics(self) -> None: """Clear session metrics. Thread-safe.""" with self._lock: self.total_prompt_tokens = 0 self.total_completion_tokens = 0 self.total_cached_tokens = 0 self.total_requests = 0 self.total_prefill_duration = 0.0 self.total_generation_duration = 0.0 self._per_model.clear() def clear_alltime_metrics(self) -> None: """Clear all-time metrics and delete the persisted file. Thread-safe.""" with self._lock: self._alltime_prompt_tokens = 0 self._alltime_completion_tokens = 0 self._alltime_cached_tokens = 0 self._alltime_requests = 0 self._alltime_prefill_duration = 0.0 self._alltime_generation_duration = 0.0 self._alltime_per_model.clear() if self._stats_path and self._stats_path.exists(): try: self._stats_path.unlink() except OSError as e: logger.warning( "Failed to delete stats file %s: %s", self._stats_path, e ) # Global singleton _server_metrics: Optional[ServerMetrics] = None def get_server_metrics() -> ServerMetrics: """Get the global ServerMetrics singleton.""" global _server_metrics if _server_metrics is None: _server_metrics = ServerMetrics() return _server_metrics def reset_server_metrics(stats_path: Optional[Path] = None) -> None: """Reset metrics (called on server start). If a previous instance exists and has a stats_path, save before resetting. """ global _server_metrics if _server_metrics is not None: _server_metrics.save_alltime() _server_metrics = ServerMetrics(stats_path=stats_path)