257 lines
7.0 KiB
Python
257 lines
7.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""
|
|
Unified cache statistics for oMLX.
|
|
|
|
This module provides base classes and utilities for tracking cache performance
|
|
metrics across different cache implementations (prefix, paged, VLM, paged SSD).
|
|
"""
|
|
|
|
from dataclasses import dataclass, asdict, field
|
|
from typing import Any, Dict
|
|
|
|
|
|
@dataclass
|
|
class BaseCacheStats:
|
|
"""
|
|
Base statistics for all cache implementations.
|
|
|
|
This class provides common metrics shared across different cache types.
|
|
Subclasses can extend with additional type-specific metrics.
|
|
"""
|
|
|
|
hits: int = 0
|
|
misses: int = 0
|
|
evictions: int = 0
|
|
|
|
@property
|
|
def total_queries(self) -> int:
|
|
"""Get total number of cache queries."""
|
|
return self.hits + self.misses
|
|
|
|
@property
|
|
def hit_rate(self) -> float:
|
|
"""
|
|
Calculate cache hit rate.
|
|
|
|
Returns:
|
|
Hit rate as a float between 0.0 and 1.0.
|
|
"""
|
|
total = self.total_queries
|
|
if total == 0:
|
|
return 0.0
|
|
return self.hits / total
|
|
|
|
def record_hit(self) -> None:
|
|
"""Record a cache hit."""
|
|
self.hits += 1
|
|
|
|
def record_miss(self) -> None:
|
|
"""Record a cache miss."""
|
|
self.misses += 1
|
|
|
|
def record_eviction(self) -> None:
|
|
"""Record a cache eviction."""
|
|
self.evictions += 1
|
|
|
|
def reset(self) -> None:
|
|
"""Reset all statistics to zero."""
|
|
self.hits = 0
|
|
self.misses = 0
|
|
self.evictions = 0
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""
|
|
Convert stats to dictionary.
|
|
|
|
Returns:
|
|
Dictionary with all stats fields.
|
|
"""
|
|
d = asdict(self)
|
|
# Add computed properties
|
|
d["total_queries"] = self.total_queries
|
|
d["hit_rate"] = self.hit_rate
|
|
return d
|
|
|
|
|
|
@dataclass
|
|
class PrefixCacheStats(BaseCacheStats):
|
|
"""
|
|
Statistics for prefix cache performance.
|
|
|
|
Extends base stats with tokens_saved to track efficiency and
|
|
partial-block skip metrics for observability.
|
|
"""
|
|
|
|
tokens_saved: int = 0
|
|
partial_block_skips: int = 0
|
|
partial_tokens_skipped: int = 0
|
|
block_size: int = 0
|
|
last_partial_tokens_skipped: int = 0
|
|
last_tokens_to_next_block: int = 0
|
|
tokens_matched_total: int = 0
|
|
tokens_requested_total: int = 0
|
|
_total_queries: int = field(default=0, repr=False)
|
|
|
|
@property
|
|
def total_queries(self) -> int:
|
|
"""Get total number of cache queries."""
|
|
# Use explicit counter if set, otherwise compute from hits + misses
|
|
if self._total_queries > 0:
|
|
return self._total_queries
|
|
return self.hits + self.misses
|
|
|
|
@total_queries.setter
|
|
def total_queries(self, value: int) -> None:
|
|
"""Set total queries counter (for legacy compatibility)."""
|
|
self._total_queries = value
|
|
|
|
def reset(self) -> None:
|
|
"""Reset all statistics to zero."""
|
|
super().reset()
|
|
self.tokens_saved = 0
|
|
self.partial_block_skips = 0
|
|
self.partial_tokens_skipped = 0
|
|
self.last_partial_tokens_skipped = 0
|
|
self.last_tokens_to_next_block = 0
|
|
self.tokens_matched_total = 0
|
|
self.tokens_requested_total = 0
|
|
self._total_queries = 0
|
|
|
|
|
|
@dataclass
|
|
class PagedCacheStats(BaseCacheStats):
|
|
"""
|
|
Statistics for paged KV cache performance.
|
|
|
|
Extends base stats with block-level metrics.
|
|
"""
|
|
|
|
total_blocks: int = 0
|
|
allocated_blocks: int = 0
|
|
free_blocks: int = 0
|
|
shared_blocks: int = 0 # Blocks with ref_count > 1
|
|
total_tokens_cached: int = 0
|
|
cow_copies: int = 0 # Copy-on-write operations
|
|
|
|
@property
|
|
def utilization(self) -> float:
|
|
"""
|
|
Calculate block utilization rate.
|
|
|
|
Returns:
|
|
Utilization as a float between 0.0 and 1.0.
|
|
"""
|
|
if self.total_blocks == 0:
|
|
return 0.0
|
|
return self.allocated_blocks / self.total_blocks
|
|
|
|
def reset(self) -> None:
|
|
"""Reset runtime statistics (not capacity metrics)."""
|
|
super().reset()
|
|
self.cow_copies = 0
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert stats to dictionary."""
|
|
d = super().to_dict()
|
|
d["utilization"] = self.utilization
|
|
return d
|
|
|
|
|
|
@dataclass
|
|
class VLMCacheStats(BaseCacheStats):
|
|
"""
|
|
Statistics for VLM (Vision Language Model) cache performance.
|
|
|
|
Extends base stats with VLM-specific metrics.
|
|
"""
|
|
|
|
tokens_saved: int = 0
|
|
image_cache_hits: int = 0
|
|
|
|
def record_image_hit(self) -> None:
|
|
"""Record an image cache hit."""
|
|
self.image_cache_hits += 1
|
|
|
|
def reset(self) -> None:
|
|
"""Reset all statistics to zero."""
|
|
super().reset()
|
|
self.tokens_saved = 0
|
|
self.image_cache_hits = 0
|
|
|
|
|
|
@dataclass
|
|
class PagedSSDCacheStats(BaseCacheStats):
|
|
"""
|
|
Statistics for paged SSD cache performance.
|
|
|
|
Extends base stats with storage-specific and hot cache metrics.
|
|
"""
|
|
|
|
# ``saves`` counts blocks that PASSED the quota gate and were enqueued
|
|
# for the writer thread — it is incremented on the inference thread in
|
|
# ``save_block`` BEFORE the writer fsyncs the file, so a block whose
|
|
# background write later fails (ENOSPC, EDQUOT, OSError) still
|
|
# contributes 1 to ``saves`` and 1 to ``errors``. Use
|
|
# ``saves_persisted`` for the count of writes that actually landed on
|
|
# disk (incremented after the atomic rename in ``_writer_loop``);
|
|
# ``saves - saves_persisted`` is the steady-state in-flight depth.
|
|
saves: int = 0
|
|
saves_persisted: int = 0
|
|
loads: int = 0
|
|
errors: int = 0
|
|
ssd_write_drops: int = 0
|
|
ssd_inline_write_fallbacks: int = 0
|
|
evict_unlink_failures: int = 0
|
|
|
|
# Storage capacity
|
|
total_size_bytes: int = 0
|
|
max_size_bytes: int = 0
|
|
configured_max_size_bytes: int = 0
|
|
num_files: int = 0
|
|
|
|
# Hot cache (in-memory tier) metrics
|
|
hot_cache_entries: int = 0
|
|
hot_cache_size_bytes: int = 0
|
|
hot_cache_max_bytes: int = 0
|
|
hot_cache_hits: int = 0
|
|
hot_cache_evictions: int = 0
|
|
hot_cache_promotions: int = 0
|
|
|
|
@property
|
|
def save_rate(self) -> float:
|
|
"""Calculate successful save rate."""
|
|
total = self.saves + self.errors
|
|
if total == 0:
|
|
return 0.0
|
|
return self.saves / total
|
|
|
|
def record_save(self) -> None:
|
|
"""Record a successful save operation."""
|
|
self.saves += 1
|
|
|
|
def record_load(self) -> None:
|
|
"""Record a successful load operation."""
|
|
self.loads += 1
|
|
self.hits += 1
|
|
|
|
def record_error(self) -> None:
|
|
"""Record an error."""
|
|
self.errors += 1
|
|
|
|
def reset(self) -> None:
|
|
"""Reset runtime statistics."""
|
|
super().reset()
|
|
self.saves = 0
|
|
self.saves_persisted = 0
|
|
self.loads = 0
|
|
self.errors = 0
|
|
self.ssd_write_drops = 0
|
|
self.ssd_inline_write_fallbacks = 0
|
|
self.evict_unlink_failures = 0
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert stats to dictionary."""
|
|
d = super().to_dict()
|
|
d["save_rate"] = self.save_rate
|
|
return d
|