Spaces:
Sleeping
Sleeping
File size: 3,348 Bytes
ea107b5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | """
LRU Model Cache Manager for AEGIS-AI.
Implements lazy loading and memory-aware eviction for deep learning models.
"""
import os
import logging
import threading
from typing import Any, Callable
from collections import OrderedDict
from typing import OrderedDict as OrderedDictType
try:
import psutil
HAS_PSUTIL = True
except ImportError:
HAS_PSUTIL = False
from config import MODEL_CACHE_MAX_MB
logger = logging.getLogger(__name__)
class ModelCache:
"""
Thread-safe, memory-aware LRU cache for models.
"""
def __init__(self, max_memory_mb: int = MODEL_CACHE_MAX_MB):
"""
Initialize the ModelCache.
Args:
max_memory_mb (int): Maximum memory in MB allowed before evicting.
"""
self._cache: OrderedDictType[str, Any] = OrderedDict()
self._max_memory = max_memory_mb
self._lock = threading.Lock()
def _estimate_memory(self) -> float:
"""
Estimate current memory usage of the process in MB.
Returns:
float: Memory usage in MB, or 0.0 if psutil is unavailable.
"""
if not HAS_PSUTIL:
return 0.0
process = psutil.Process(os.getpid())
return process.memory_info().rss / (1024 * 1024)
def get_or_load(self, key: str, loader_fn: Callable[[], Any]) -> Any:
"""
Retrieve a model from the cache or load it if not present.
Moves the accessed model to the end of the LRU cache.
Args:
key (str): Unique identifier for the model.
loader_fn (Callable): Function to load the model if not cached.
Returns:
Any: The loaded model.
"""
with self._lock:
if key in self._cache:
logger.debug(f"Cache hit for model: {key}")
self._cache.move_to_end(key)
return self._cache[key]
logger.info(f"Cache miss. Loading model: {key}")
# Execute loader_fn outside the lock to avoid blocking other cache reads
model = loader_fn()
with self._lock:
self._cache[key] = model
self._evict_if_needed()
return model
def _evict_if_needed(self) -> None:
"""
Evict least recently used models if memory exceeds max limit.
Leaves at least one model in cache.
"""
if not HAS_PSUTIL:
return
while self._estimate_memory() > self._max_memory and len(self._cache) > 1:
evicted_key, _ = self._cache.popitem(last=False)
logger.warning(f"Memory threshold exceeded. Evicted model: {evicted_key}")
def clear_group(self, group_name: str) -> None:
"""
Evict all models whose key starts with the group_name.
Groups e.g., 'image', 'audio', 'support'.
Args:
group_name (str): Group name of the models to clear.
"""
with self._lock:
keys_to_remove = [k for k in self._cache.keys() if k.startswith(group_name)]
for k in keys_to_remove:
del self._cache[k]
logger.info(f"Evicted model {k} from group {group_name}")
# Global instance
model_cache = ModelCache()
|