src/model_cache.py

# -------------------------------------------------------------------
# This source file is available under the terms of the
# Pimcore Open Core License (POCL)
# Full copyright and license information is available in
# LICENSE.md which is distributed with this source code.
#
#  @copyright  Copyright (c) Pimcore GmbH (https://www.pimcore.com)
#  @license    Pimcore Open Core License (POCL)
# -------------------------------------------------------------------

"""A bounded, least-recently-used cache for loaded models and processors.

Callers name the model in the request, so the set of models one running instance may be asked for
is open-ended. An unbounded cache therefore grows until the container is killed — no attacker
needed, just a handful of callers each bringing their own model. Bounding it trades a reload for
staying alive, which is the right trade for a service whose working set is realistically a couple
of models used constantly plus a long tail used once.

The bound is a *count*, not a byte budget, because it is predictable and cheap to reason about.
That is deliberately crude: ten small models and ten multi-GB ones both count as ten, so an
operator running large vision models should lower `MAX_CACHED_MODELS` accordingly.

Eviction only drops the cache's reference. A model still in use by an in-flight request stays
alive through that request's own reference and is collected afterwards, so evicting under load is
safe.
"""

import logging
from collections import OrderedDict
from typing import Any, Optional


class LruCache:
    """Dict-like, bounded by entry count, evicting least-recently-used first.

    Deliberately dict-shaped (`in`, `[]`, `len`) so it drops into the existing
    `if key not in cache: load` call sites without restructuring them. Reads count as uses:
    `__getitem__` refreshes recency, otherwise the model every request hits would be the one
    evicted.

    Not thread-safe, and does nothing to prevent two concurrent misses loading the same model
    twice — that race predates this class and is unchanged by it.
    """

    def __init__(self, max_entries: int, name: str, logger: Optional[logging.Logger] = None) -> None:
        # A limit below one would evict the entry a request just loaded before it could be used,
        # making every request a cold load. Clamp rather than raise: this comes from an env var,
        # and a typo should degrade, not take the service down at import.
        self._max_entries = max(1, max_entries)
        self._name = name
        self._logger = logger or logging.getLogger(__name__)
        self._entries: OrderedDict[str, Any] = OrderedDict()

    def __contains__(self, key: str) -> bool:
        return key in self._entries

    def __getitem__(self, key: str) -> Any:
        value = self._entries[key]
        self._entries.move_to_end(key)
        return value

    def __setitem__(self, key: str, value: Any) -> None:
        if key in self._entries:
            # Overwriting an existing key must not evict anything: the size is unchanged.
            self._entries[key] = value
            self._entries.move_to_end(key)
            return

        self._entries[key] = value
        while len(self._entries) > self._max_entries:
            evicted_key, _ = self._entries.popitem(last=False)
            # Logged because thrashing is otherwise invisible — the only symptom is slowness, and
            # a repeated evict/reload of the same id is the signal that the limit is too low.
            self._logger.info(
                f"Evicted {self._name} '{evicted_key}' (cache limit {self._max_entries} reached); "
                f"it will be reloaded if requested again"
            )

    def __len__(self) -> int:
        return len(self._entries)

    def keys(self) -> Any:
        return self._entries.keys()