core/model_cache.py

"""
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()