# ------------------------------------------------------------------- # 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()` 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), ]