File size: 5,200 Bytes
ed962ac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | # -------------------------------------------------------------------
# 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),
]
|