Coverage for src/local_deep_research/embeddings/model_cache.py: 97%
185 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""One resident local embedding backend; callers own weight-free handles.
3The cache contains model weights/tokenizer only, never user settings, documents,
4query results, or authorization decisions. Provider admission runs before this
5module is called. Inference is serial: one call at a time uses the resident
6model, and model replacement waits until every admitted call finishes.
7"""
9from __future__ import annotations
11import atexit
12from copy import deepcopy
13import gc
14import hashlib
15from pathlib import Path
16import re
17from stat import S_ISDIR
18from threading import Lock, local
19from typing import Any, Callable, TypeVar
21from langchain_core.embeddings import Embeddings
23from .model_loading import MODEL_LOAD_LOCK, load_sentence_transformer
24from ..security.secure_logging import logger
26_Result = TypeVar("_Result")
27_Factory = Callable[..., Embeddings]
28_CacheKey = tuple[tuple[_Factory, Any], str]
31def _freeze(value: Any) -> Any:
32 if isinstance(value, dict):
33 return tuple(
34 sorted((key, _freeze(item)) for key, item in value.items())
35 )
36 if isinstance(value, (list, tuple)): 36 ↛ 37line 36 didn't jump to line 37 because the condition on line 36 was never true
37 return tuple(_freeze(item) for item in value)
38 return value
41def _artifact_stamp(kwargs: dict) -> str:
42 """Notice local replacements and floating Hub ref changes without I/O
43 of model contents or contacting the network.
45 Hub snapshots consist of links to immutable blobs. Stat follows those links
46 so an operator replacing a blob also changes the key. Local module directory
47 links are followed too, visiting each directory once to stop cycles. A handle
48 cannot reload changed files after eviction and silently use different weights.
49 """
50 name = kwargs["model_name"]
51 # nosemgrep: semgrep.rules.path-traversal-risk, reason: Provider admission confines absolute model paths to the app models directory; the stamp only stats and lists entries
52 path = Path(name)
53 references = []
54 if path.is_absolute():
55 roots = [path]
56 else:
57 # nosemgrep: semgrep.rules.path-traversal-risk, reason: Operator model cache root; admitted repository ids have slashes replaced and revisions are catalog or cached commit ids
58 cache = Path(kwargs["cache_folder"])
59 repositories = [name]
60 if "/" not in name:
61 repositories.append(f"sentence-transformers/{name}")
62 roots = []
63 revision = kwargs.get("model_kwargs", {}).get("revision")
64 for repository in repositories:
65 cached_repo = cache / ("models--" + repository.replace("/", "--"))
66 snapshot = revision
67 if not revision or not re.fullmatch(r"[0-9a-f]{40}", revision):
68 reference = cached_repo / "refs" / (revision or "main")
69 references.append(reference)
70 try:
71 # nosemgrep: semgrep.rules.path-traversal-risk, reason: Reads refs/main of an admitted repository inside the operator model cache; repository ids have slashes replaced and provider revisions are commit ids
72 snapshot = reference.read_text(encoding="ascii").strip()
73 except FileNotFoundError:
74 snapshot = None
75 if snapshot and re.fullmatch(r"[0-9a-f]{40}", snapshot):
76 roots.append(cached_repo / "snapshots" / snapshot)
77 # Do not scan all snapshots/blobs or .no_exist markers: warming an
78 # unrelated revision must not invalidate handles pinned to old bytes.
79 roots.extend(references)
80 digest = hashlib.sha256()
81 for root in roots:
82 # Missing files are part of the identity too (initial download/repair).
83 digest.update(str(root).encode())
84 paths = [root]
85 seen_directories = set()
86 while paths:
87 path = paths.pop()
88 try:
89 stat = path.stat()
90 except FileNotFoundError:
91 continue
92 digest.update(str(path).encode())
93 digest.update(
94 repr(
95 (
96 stat.st_dev,
97 stat.st_ino,
98 stat.st_size,
99 stat.st_mtime_ns,
100 stat.st_ctime_ns,
101 )
102 ).encode()
103 )
104 if S_ISDIR(stat.st_mode):
105 # rglob does not descend into directory symlinks. Stamp the
106 # target and its files, while preventing alias/cycle expansion.
107 directory = (stat.st_dev, stat.st_ino)
108 if directory in seen_directories:
109 continue
110 seen_directories.add(directory)
111 # nosemgrep: semgrep.rules.path-traversal-risk, reason: Lists an admitted model root and its linked module directories only to stat entries; no file contents are read or written
112 paths.extend(sorted(path.iterdir(), reverse=True))
113 return digest.hexdigest()
116class _ModelCache:
117 def __init__(self) -> None:
118 self.key: _CacheKey | None = None
119 self.backend: Embeddings | None = None
120 # Tokenization and forward passes mutate model and tokenizer state,
121 # so admitted calls take turns on the one resident backend.
122 self._inference = Lock()
123 self._inference_state = local()
125 def _clear_locked(self) -> None:
126 had_model = self.backend is not None
127 self.backend = None
128 self.key = None
129 if had_model:
130 # SentenceTransformer -> model_card_data -> SentenceTransformer
131 # is a strong cycle. Dropping the wrapper alone retains hundreds
132 # of MiB until a generation-2 collection eventually happens.
133 gc.collect()
135 def clear(self, *, wait: bool = True) -> bool:
136 if MODEL_LOAD_LOCK.acquire(blocking=wait):
137 try:
138 self._clear_locked()
139 return True
140 finally:
141 MODEL_LOAD_LOCK.release()
142 return False
144 def prepare(
145 self, factory: _Factory, kwargs: dict[str, Any]
146 ) -> SharedSentenceTransformerEmbeddings:
147 # A warm handle for an identical configuration is created without
148 # waiting for active calls; its own calls then wait for their turn.
149 with MODEL_LOAD_LOCK.inference():
150 options, identity = self._options(factory, kwargs)
151 key = (identity, _artifact_stamp(kwargs))
152 if self.key == key:
153 return SharedSentenceTransformerEmbeddings(
154 self, factory, options, key
155 )
156 with MODEL_LOAD_LOCK:
157 return self._prepare_locked(factory, kwargs)
159 @staticmethod
160 def _options(factory: _Factory, kwargs: dict[str, Any]):
161 local_options = deepcopy(kwargs)
162 local_options.setdefault("model_kwargs", {})["local_files_only"] = True
163 return local_options, (factory, _freeze(local_options))
165 def _load_locked(self, factory: _Factory, kwargs: dict[str, Any]) -> None:
166 self.backend = load_sentence_transformer(factory, **kwargs)
168 def _prepare_locked(
169 self, factory: _Factory, kwargs: dict[str, Any]
170 ) -> SharedSentenceTransformerEmbeddings:
171 options = deepcopy(kwargs)
172 # Network permission is checked by the provider on EVERY request.
173 # It is not a property of the resident weights. After initial
174 # construction, an evicted handle may only reload local files.
175 local_options, identity = self._options(factory, options)
176 stamp = _artifact_stamp(options)
177 key = (identity, stamp)
178 if self.key != key:
179 self._clear_locked()
180 logger.info("Loading resident local embedding model")
181 try:
182 self._load_locked(factory, options)
183 loaded = _artifact_stamp(options)
184 if loaded != stamp:
185 # An authorized download/repair, or a replacement, changed
186 # the files while the constructor read them, so neither
187 # stamp need describe the loaded bytes. Load once more
188 # from local files and keep it only if they stay intact.
189 self._clear_locked()
190 logger.info(
191 "Reloading local embedding model after its files changed"
192 )
193 self._load_locked(factory, local_options)
194 stamp, loaded = loaded, _artifact_stamp(options)
195 except BaseException:
196 self._clear_locked()
197 raise
198 if loaded != stamp:
199 self._clear_locked()
200 raise ValueError(
201 "Local embedding model files changed; create a new embedding client"
202 )
203 key = (identity, stamp)
204 self.key = key
205 return SharedSentenceTransformerEmbeddings(
206 self, factory, local_options, key
207 )
209 def embed(
210 self,
211 handle: SharedSentenceTransformerEmbeddings,
212 operation: Callable[[Embeddings], _Result],
213 ) -> _Result:
214 while True:
215 with MODEL_LOAD_LOCK.inference():
216 if handle._closed:
217 raise RuntimeError("Embedding client is closed")
218 if self.key == handle._key:
219 return self._run_serially(handle, operation)
220 # Never upgrade shared access in place: another reader may also
221 # need a different model. Recheck after obtaining exclusive access.
222 with MODEL_LOAD_LOCK:
223 self._reload_locked(handle)
225 def _reload_locked(
226 self, handle: SharedSentenceTransformerEmbeddings
227 ) -> None:
228 if handle._closed: 228 ↛ 229line 228 didn't jump to line 229 because the condition on line 228 was never true
229 raise RuntimeError("Embedding client is closed")
230 if self.key != handle._key: 230 ↛ exitline 230 didn't return from function '_reload_locked' because the condition on line 230 was always true
231 if _artifact_stamp(handle._kwargs) != handle._key[1]:
232 raise ValueError(
233 "Local embedding model files changed; create a new embedding client"
234 )
235 self._clear_locked()
236 logger.info("Reloading evicted local embedding model")
237 try:
238 self._load_locked(handle._factory, handle._kwargs)
239 # An operator or another process can replace artifacts while
240 # the constructor reads them. Never give an old handle the
241 # resulting weights under its original configuration identity.
242 stamp = _artifact_stamp(handle._kwargs)
243 except BaseException:
244 self._clear_locked()
245 raise
246 if stamp != handle._key[1]:
247 self._clear_locked()
248 raise ValueError(
249 "Local embedding model files changed; create a new embedding client"
250 )
251 self.key = handle._key
253 def _run_serially(
254 self,
255 handle: SharedSentenceTransformerEmbeddings,
256 operation: Callable[[Embeddings], _Result],
257 ) -> _Result:
258 # The caller holds shared access, so no writer can replace the backend
259 # while this call waits for its turn or runs. The returned vectors are
260 # ordinary lists; no caller keeps a model reference that could defeat
261 # eviction or the one-model cap.
262 backend = self.backend
263 if backend is None: 263 ↛ 264line 263 didn't jump to line 264 because the condition on line 263 was never true
264 raise RuntimeError("Local embedding model was not initialized")
265 if getattr(self._inference_state, "active", False):
266 # The inference lock is not reentrant; waiting would deadlock.
267 raise RuntimeError(
268 "Recursive local embedding inference is not supported"
269 )
270 with self._inference:
271 try:
272 self._inference_state.active = True
273 if handle._closed:
274 raise RuntimeError("Embedding client is closed")
275 return operation(backend)
276 finally:
277 self._inference_state.active = False
280class SharedSentenceTransformerEmbeddings(Embeddings):
281 """Per-caller lifecycle/configuration, with no ownership of model weights."""
283 def __init__(
284 self,
285 cache: _ModelCache,
286 factory: _Factory,
287 kwargs: dict[str, Any],
288 key: _CacheKey,
289 ) -> None:
290 self._cache = cache
291 self._factory = factory
292 self._kwargs = kwargs
293 self._key = key
294 self._closed = False
296 def embed_documents(self, texts: list[str]) -> list[list[float]]:
297 return self._cache.embed(
298 self, lambda backend: backend.embed_documents(texts)
299 )
301 def embed_query(self, text: str) -> list[float]:
302 return self._cache.embed(
303 self, lambda backend: backend.embed_query(text)
304 )
306 def close(self) -> None:
307 # Closing one search must not invalidate another search's handle.
308 # Already-started inference may finish; future calls are refused.
309 self._closed = True
312_MODEL_CACHE = _ModelCache()
315def shared_sentence_transformer_embeddings(
316 factory: Callable[..., Embeddings], **kwargs: Any
317) -> Embeddings:
318 """Prepare a checked provider configuration and return a lightweight handle."""
319 return _MODEL_CACHE.prepare(factory, kwargs)
322def clear_sentence_transformer_cache(*, wait: bool = True) -> bool:
323 """Release weights, returning False if shutdown finds inference busy.
325 Shutdown must not wait indefinitely for a daemon research worker. In that
326 case the process exit reclaims its active model along with the worker.
327 """
328 return _MODEL_CACHE.clear(wait=wait)
331atexit.register(clear_sentence_transformer_cache, wait=False)