cache.py

"""Thread-safe in-memory cache for lazily loaded models."""

from __future__ import annotations

from collections import OrderedDict
from collections.abc import Callable
import threading

from config import MAX_CACHED_MODELS, MODELS, ModelSpec
from loader import LoadedModelBundle, load_model_bundle, release_model_bundle


class ModelCache:
    """Caches loaded models in memory and avoids duplicate loads."""

    def __init__(
        self,
        specs: tuple[ModelSpec, ...],
        loader: Callable[[ModelSpec], LoadedModelBundle],
        *,
        max_items: int,
    ) -> None:
        self._specs = {spec.name: spec for spec in specs}
        self._loader = loader
        self._max_items = max(1, max_items)
        self._items: OrderedDict[str, LoadedModelBundle] = OrderedDict()
        self._state_lock = threading.RLock()
        self._load_locks = {name: threading.Lock() for name in self._specs}

    def get(self, model_name: str) -> LoadedModelBundle:
        """Return a loaded model, loading and caching it if needed."""

        with self._state_lock:
            cached = self._items.get(model_name)
            if cached is not None:
                self._items.move_to_end(model_name)
        if cached is not None:
            return cached

        load_lock = self._load_locks[model_name]
        with load_lock:
            with self._state_lock:
                cached = self._items.get(model_name)
                if cached is not None:
                    self._items.move_to_end(model_name)
            if cached is not None:
                return cached

            spec = self._specs[model_name]
            loaded = self._loader(spec)
            evicted: LoadedModelBundle | None = None
            with self._state_lock:
                self._items[model_name] = loaded
                self._items.move_to_end(model_name)
                if len(self._items) > self._max_items:
                    _, evicted = self._items.popitem(last=False)
            if evicted is not None:
                release_model_bundle(evicted)
            with self._state_lock:
                self._items.move_to_end(model_name)
                return loaded

    def status_snapshot(self) -> dict[str, bool]:
        """Return whether each configured model has been loaded."""

        with self._state_lock:
            loaded_names = set(self._items)
        return {name: name in loaded_names for name in self._specs}

    def loaded_count(self) -> int:
        """Return the number of loaded models."""

        with self._state_lock:
            return len(self._items)


MODEL_CACHE = ModelCache(MODELS, load_model_bundle, max_items=MAX_CACHED_MODELS)