tests/test_model_cache_sharing.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.
# -------------------------------------------------------------------
"""One loaded model per id, not one per (id, use-site).

The text service, the image service and the vector-size lookup all want the same
`AutoModel.from_pretrained(<id>)` object. When they each kept their own cache and appended an
image/text suffix to the key, a single multi-GB model was resident three times.
"""

import logging

import pytest
from src import embeddings
from src.embeddings import ImageEmbeddingTaskService, TextEmbeddingTaskService
from src.model_cache import LruCache

_TEST_LIMIT = 3


@pytest.fixture(autouse=True)
def _isolate_caches(monkeypatch):
    """The caches are module-level, so give each test its own pair rather than leaking across.

    Real `LruCache` instances, not plain dicts — patching in a dict would quietly stop these tests
    from exercising the type the service actually uses.
    """
    logger = logging.getLogger("test")
    monkeypatch.setattr(embeddings, "_MODEL_CACHE", LruCache(_TEST_LIMIT, "model", logger))
    monkeypatch.setattr(embeddings, "_PROCESSOR_CACHE", LruCache(_TEST_LIMIT, "processor", logger))


class _FakeModel:
    def to(self, device):
        return self


def _count_loads(monkeypatch) -> list:
    """Record every from_pretrained call so we can assert how many actually happened."""
    calls: list = []

    def fake_from_pretrained(model_name, **kwargs):
        calls.append((model_name, kwargs.get("trust_remote_code")))
        return _FakeModel()

    monkeypatch.setattr(embeddings.AutoModel, "from_pretrained", fake_from_pretrained)
    return calls


def _services():
    """Fresh instances bound to the current (test-isolated) module-level caches."""
    logger = logging.getLogger("test")
    return ImageEmbeddingTaskService(logger), TextEmbeddingTaskService(logger)


def test_image_and_text_services_share_one_loaded_model(monkeypatch):
    calls = _count_loads(monkeypatch)
    image_service, text_service = _services()

    first = image_service._load_model("google/siglip-so400m-patch14-384")
    second = text_service._load_model("google/siglip-so400m-patch14-384")

    assert len(calls) == 1, "the text service reloaded a model the image service had already loaded"
    assert first is second


def test_vector_size_lookup_reuses_the_same_model(monkeypatch):
    """The third copy: get_embedding_vector_size used to key the model without any suffix."""
    calls = _count_loads(monkeypatch)
    image_service, text_service = _services()

    image_service._load_model("openai/clip-vit-large-patch14")
    text_service._load_model("openai/clip-vit-large-patch14")
    image_service._load_model("openai/clip-vit-large-patch14")

    assert len(calls) == 1
    assert len(embeddings._MODEL_CACHE) == 1, f"expected one cache entry, got {list(embeddings._MODEL_CACHE)}"


def test_distinct_models_are_cached_separately(monkeypatch):
    """Guard against over-sharing: deduplicating must not collapse different models together."""
    calls = _count_loads(monkeypatch)
    image_service, _ = _services()

    image_service._load_model("openai/clip-vit-large-patch14")
    image_service._load_model("google/siglip-so400m-patch14-384")

    assert len(calls) == 2
    assert len(embeddings._MODEL_CACHE) == 2


def test_loader_evicts_once_past_the_limit(monkeypatch):
    """End to end through `_load_model`: a caller naming more models than fit must not grow the
    cache without bound — the coldest is dropped instead."""
    calls = _count_loads(monkeypatch)
    image_service, _ = _services()

    for index in range(_TEST_LIMIT + 1):
        image_service._load_model(f"org/model-{index}")

    assert len(calls) == _TEST_LIMIT + 1
    assert len(embeddings._MODEL_CACHE) == _TEST_LIMIT
    assert "org/model-0" not in embeddings._MODEL_CACHE, "the coldest model should have been evicted"
    assert f"org/model-{_TEST_LIMIT}" in embeddings._MODEL_CACHE


def test_an_evicted_model_is_reloaded_on_next_use(monkeypatch):
    """Eviction costs a reload rather than an error — the caller must not see a failure."""
    calls = _count_loads(monkeypatch)
    image_service, _ = _services()

    for index in range(_TEST_LIMIT + 1):
        image_service._load_model(f"org/model-{index}")
    image_service._load_model("org/model-0")

    assert len(calls) == _TEST_LIMIT + 2
    assert "org/model-0" in embeddings._MODEL_CACHE


def test_remote_code_flag_follows_the_trusted_org_list(monkeypatch):
    """The loader must pass the per-model decision through, not a hardcoded True."""
    monkeypatch.setenv("TRUSTED_CODE_ORGS", "nomic-ai")
    calls = _count_loads(monkeypatch)
    image_service, _ = _services()

    image_service._load_model("nomic-ai/nomic-embed-text-v1.5")
    image_service._load_model("someone-else/custom-model")

    assert calls == [
        ("nomic-ai/nomic-embed-text-v1.5", True),
        ("someone-else/custom-model", False),
    ]