Coverage for src/lilbee/modelhub/model_manager/core.py: 100%
163 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""ModelManager: native and SDK-backed model lifecycle operations."""
3import logging
4import threading
5from collections.abc import Callable
6from dataclasses import replace
7from pathlib import Path
9from cachetools import TTLCache
11from lilbee.catalog.compat import UnsupportedQuantError
12from lilbee.catalog.models import CatalogModel
13from lilbee.catalog.types import ModelSource
14from lilbee.core.config import DEFAULT_HTTP_TIMEOUT
15from lilbee.core.security import validate_path_within
16from lilbee.modelhub.model_manager.types import ModelNotFoundError
17from lilbee.modelhub.registry import ModelRegistry
18from lilbee.providers.local_servers import LOCAL_SERVERS, local_server_for_key
19from lilbee.providers.model_ref import parse_model_ref
20from lilbee.runtime.cancellation import CancelSignal
22log = logging.getLogger(__name__)
24_INSTALLED_CACHE_TTL_SECONDS = 60.0
25_NATIVE_IDENTITIES_CACHE_KEY = "native-identities"
28def _prefixed_source(model: str) -> ModelSource | None:
29 """Map a provider-prefixed ref to its source, or ``None`` for a bare ref.
31 Local-server prefixes (``ollama/``, ``lm_studio/``) map to that server's
32 source; API-provider prefixes are FRONTIER. A bare name returns ``None``
33 so the caller falls back to backend membership.
34 """
35 for spec in LOCAL_SERVERS:
36 if model.startswith(spec.wire_prefix):
37 return ModelSource(spec.key)
38 try:
39 ref = parse_model_ref(model)
40 except ValueError:
41 return None
42 return ModelSource.FRONTIER if ref.is_api else None
45class ModelManager:
46 """Manages model lifecycle with distinct sources."""
48 def __init__(self, models_dir: Path) -> None:
49 self._models_dir = models_dir
50 self._registry = ModelRegistry(self._models_dir)
51 # Memoize list_installed results to avoid walking the registry
52 # filesystem and hitting the backend HTTP endpoint on every call.
53 # The catalog filter path fires this per request. Time-based TTL
54 # plus explicit invalidation on pull/remove keeps freshness.
55 self._installed_cache: TTLCache[ModelSource | None, list[str]] = TTLCache(
56 maxsize=8, ttl=_INSTALLED_CACHE_TTL_SECONDS
57 )
58 # Identity cache: refs + hf_repos of installed natives. The catalog
59 # screen reads this to mark rows as installed without re-walking
60 # the registry on every screen mount (~150-300 ms saved).
61 self._native_identities_cache: TTLCache[str, frozenset[str]] = TTLCache(
62 maxsize=1, ttl=_INSTALLED_CACHE_TTL_SECONDS
63 )
64 # TTLCache splices an internal link list without a lock; every access holds this one.
65 self._cache_lock = threading.Lock()
67 def list_installed(self, source: ModelSource | None = None) -> list[str]:
68 """List installed model names. ``source=None`` lists all sources.
70 Memoized with a ``_INSTALLED_CACHE_TTL_SECONDS`` TTL and
71 invalidated eagerly by ``pull``/``remove``.
72 """
73 with self._cache_lock:
74 cached = self._installed_cache.get(source)
75 if cached is not None:
76 return cached
78 if source is None:
79 native = set(self._list_native())
80 remote = set(self._list_remote())
81 result = sorted(native | remote)
82 elif source is ModelSource.NATIVE:
83 result = self._list_native()
84 else:
85 result = self._list_remote()
87 with self._cache_lock:
88 self._installed_cache[source] = result
89 return result
91 def list_native_identities(self) -> frozenset[str]:
92 """Return refs + hf_repos of installed native models.
94 Same TTL as ``list_installed``. The catalog screen reads this to
95 mark catalog rows as installed without re-walking the registry
96 on every screen mount.
98 Raises ``OSError`` when the registry cannot be read, so an
99 unreadable tree is never cached as nothing installed.
100 """
101 with self._cache_lock:
102 cached = self._native_identities_cache.get(_NATIVE_IDENTITIES_CACHE_KEY)
103 if cached is not None:
104 return cached
105 identities: set[str] = set()
106 for m in self._registry.list_installed():
107 identities.add(m.ref)
108 identities.add(m.hf_repo)
109 result = frozenset(identities)
110 with self._cache_lock:
111 self._native_identities_cache[_NATIVE_IDENTITIES_CACHE_KEY] = result
112 return result
114 def _invalidate_installed_cache(self) -> None:
115 """Drop all cached list_installed results and the route-layer cache."""
116 with self._cache_lock:
117 self._installed_cache.clear()
118 self._native_identities_cache.clear()
119 from lilbee.app.services import peek_services
121 # peek_services is None for a standalone ModelManager (test setup);
122 # the route isn't running so there's nothing to invalidate.
123 services = peek_services()
124 if services is not None:
125 services.known_models.invalidate()
127 def _list_native(self) -> list[str]:
128 """List native models from the registry only."""
129 return sorted(m.ref for m in self._registry.list_installed())
131 def _list_remote(self) -> list[str]:
132 """List model names across every configured local server (Ollama, LM Studio).
134 Reuses the discovery dispatch so each listing endpoint matches its
135 server (Ollama ``/api/tags`` vs LM Studio ``/v1/models``). Returns
136 ``[]`` when the backends are unreachable.
137 """
138 # circular: discovery -> app.services -> model_manager.__init__ -> core
139 from lilbee.modelhub.model_manager.discovery import classify_all_remote_models
141 models = classify_all_remote_models(timeout=DEFAULT_HTTP_TIMEOUT)
142 return [m.name for m in models]
144 def is_installed(self, model: str, source: ModelSource | None = None) -> bool:
145 """Check if model exists in specified source."""
146 if source is None:
147 return self._is_native(model) or self._is_remote(model)
148 if source is ModelSource.NATIVE:
149 return self._is_native(model)
150 return self._is_remote(model)
152 def _is_native(self, model: str) -> bool:
153 """True when the manifest registry holds *model*."""
154 return self._registry.is_installed(model)
156 def is_registered(self, model: str) -> bool:
157 """True when a manifest records *model*, so the model list shows it."""
158 return self._registry.has_manifest(model)
160 def _is_remote(self, model: str) -> bool:
161 return model in self.list_installed(ModelSource.REMOTE)
163 def get_source(self, model: str) -> ModelSource | None:
164 """Return the granular source a model lives in. Native takes precedence.
166 A provider-prefixed ref classifies without a network call; a bare name
167 that a backend reports installed is ``REMOTE`` (the prefix is what names
168 the specific server). ``None`` when the model is in no known source.
169 """
170 if self._is_native(model):
171 return ModelSource.NATIVE
172 prefixed = _prefixed_source(model)
173 if prefixed is not None:
174 return prefixed
175 if self._is_remote(model):
176 return ModelSource.REMOTE
177 return None
179 def pull(
180 self,
181 model: str,
182 source: ModelSource,
183 *,
184 on_bytes: Callable[[int, int], None] | None = None,
185 allow_unsupported: bool = False,
186 cancel: CancelSignal | None = None,
187 ) -> Path | None:
188 """Download a native GGUF model and return its path.
190 lilbee pulls native models only. Local servers (Ollama, LM Studio)
191 are read-only: their models are managed in their own app and surface
192 here once present, so a non-native *source* is refused.
194 Native pulls of architectures the bundled llama.cpp doesn't support
195 are refused with ``UnsupportedArchError``, and files whose tensors it
196 cannot decode with ``UnsupportedQuantError``, unless *allow_unsupported*
197 is True. *on_bytes* receives (downloaded_bytes, total_bytes) progress.
198 A *cancel* signal makes the download cancellable mid-transfer; see
199 :func:`~lilbee.catalog.download_model`.
200 """
201 if source is not ModelSource.NATIVE:
202 spec = local_server_for_key(source.value)
203 where = spec.display_name if spec is not None else "the configured server"
204 raise ValueError(
205 f"lilbee runs {where} models but doesn't download them. "
206 f"Add the model in {where}, then pick it here."
207 )
208 if not allow_unsupported:
209 self.enforce_arch_compat(model)
210 try:
211 return self._pull_native(
212 model, on_bytes=on_bytes, cancel=cancel, allow_unsupported=allow_unsupported
213 )
214 finally:
215 self._invalidate_installed_cache()
217 def enforce_arch_compat(self, ref: str) -> None:
218 """Raise UnsupportedArchError if *ref*'s architecture isn't in the supported set.
220 Public because the pull preflight on the HTTP surface runs the same
221 check before starting a download, so a caller learns the model is
222 unsupported before any bytes move.
223 """
224 from lilbee.app.services import get_services
225 from lilbee.catalog.compat import (
226 ModelCompat,
227 UnsupportedArchError,
228 classify,
229 resolve_arch_for_pull,
230 )
232 arch = resolve_arch_for_pull(ref, get_services().hf_client)
233 if classify(arch) is ModelCompat.UNSUPPORTED:
234 raise UnsupportedArchError(ref, arch)
236 def _pull_native(
237 self,
238 model: str,
239 *,
240 on_bytes: Callable[[int, int], None] | None = None,
241 cancel: CancelSignal | None = None,
242 allow_unsupported: bool = False,
243 ) -> Path:
244 """Download a featured or ad-hoc HuggingFace model to the native GGUF directory."""
245 # heavy: lilbee.catalog (>50ms; huggingface_hub fanout)
246 from lilbee.catalog import download_model, resolve_pull_target
247 from lilbee.modelhub.registry import register_downloaded_model
249 entry = resolve_pull_target(model)
250 if entry is None:
251 raise ModelNotFoundError(
252 f"Model '{model}' not recognized. "
253 "Pass a HuggingFace repo id (owner/name) or a featured model name."
254 )
255 entry = self._resolved_entry(entry, verify=not allow_unsupported)
256 path = download_model(
257 entry, on_progress=on_bytes, on_complete=register_downloaded_model, cancel=cancel
258 )
259 log.info("Downloaded %s to %s", model, path)
260 return path
262 @staticmethod
263 def _resolved_entry(entry: CatalogModel, *, verify: bool) -> CatalogModel:
264 """*entry* carrying the file the pull will fetch.
266 Resolution happens here because the engine's verdict is part of choosing
267 the file, not a check on the file already chosen. A repo can publish the
268 same weights in several packings, only some of which this build reads,
269 and the resolver walks its ranking until one of them answers. Passing the
270 resolved name down means the download does not work it out again.
272 A repo that will not resolve is not this check's verdict to deliver. The
273 download reports the gated repo or the missing file properly, so a
274 resolution failure here is left to it.
275 """
276 # heavy: lilbee.catalog (>50ms; huggingface_hub fanout)
277 from lilbee.catalog.download import resolve_filename
278 from lilbee.catalog.hf_client import hf_token
279 from lilbee.providers.fleet.loadability import assert_engine_can_load
281 token = hf_token()
283 def _can_load(hf_repo: str, filename: str) -> None:
284 assert_engine_can_load(hf_repo, filename, token)
286 try:
287 filename = resolve_filename(entry, can_load=_can_load if verify else None)
288 except (PermissionError, RuntimeError) as exc:
289 if isinstance(exc, UnsupportedQuantError):
290 raise
291 log.debug("Cannot name a file to fetch for %s: %s", entry.hf_repo, exc)
292 return entry
293 return replace(entry, gguf_filename=filename)
295 def remove(self, model: str, source: ModelSource | None = None) -> bool:
296 """Remove an installed native model. Returns True if removed.
298 lilbee removes only native GGUF models it downloaded. Local servers
299 (Ollama, LM Studio) are read-only: a model that lives on one is refused
300 (mirrors ``pull``), since its lifecycle is managed in that app. A bare
301 ``source`` is resolved so a local-server ref is caught either way.
302 """
303 effective = source if source is not None else self.get_source(model)
304 if effective is not None and effective is not ModelSource.NATIVE:
305 spec = local_server_for_key(effective.value)
306 where = spec.display_name if spec is not None else "the configured server"
307 raise ValueError(
308 f"lilbee runs {where} models but doesn't remove them. "
309 f"Manage them in {where} instead."
310 )
311 try:
312 return self._remove_native(model)
313 finally:
314 self._invalidate_installed_cache()
316 def _remove_native(self, model: str) -> bool:
317 if self._registry.remove(model):
318 log.info("Removed native model %s from registry", model)
319 return True
320 try:
321 path = validate_path_within(self._models_dir / model, self._models_dir)
322 except ValueError:
323 log.warning("Path traversal blocked: %s escapes %s", model, self._models_dir)
324 return False
325 if path.is_file():
326 path.unlink()
327 log.info("Removed native model %s", model)
328 return True
329 return False