Coverage for src/lilbee/modelhub/model_manager/core.py: 100%
161 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-28 17:20 +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_remote(self, model: str) -> bool:
157 return model in self.list_installed(ModelSource.REMOTE)
159 def get_source(self, model: str) -> ModelSource | None:
160 """Return the granular source a model lives in. Native takes precedence.
162 A provider-prefixed ref classifies without a network call; a bare name
163 that a backend reports installed is ``REMOTE`` (the prefix is what names
164 the specific server). ``None`` when the model is in no known source.
165 """
166 if self._is_native(model):
167 return ModelSource.NATIVE
168 prefixed = _prefixed_source(model)
169 if prefixed is not None:
170 return prefixed
171 if self._is_remote(model):
172 return ModelSource.REMOTE
173 return None
175 def pull(
176 self,
177 model: str,
178 source: ModelSource,
179 *,
180 on_bytes: Callable[[int, int], None] | None = None,
181 allow_unsupported: bool = False,
182 cancel: CancelSignal | None = None,
183 ) -> Path | None:
184 """Download a native GGUF model and return its path.
186 lilbee pulls native models only. Local servers (Ollama, LM Studio)
187 are read-only: their models are managed in their own app and surface
188 here once present, so a non-native *source* is refused.
190 Native pulls of architectures the bundled llama.cpp doesn't support
191 are refused with ``UnsupportedArchError``, and files whose tensors it
192 cannot decode with ``UnsupportedQuantError``, unless *allow_unsupported*
193 is True. *on_bytes* receives (downloaded_bytes, total_bytes) progress.
194 A *cancel* signal makes the download cancellable mid-transfer; see
195 :func:`~lilbee.catalog.download_model`.
196 """
197 if source is not ModelSource.NATIVE:
198 spec = local_server_for_key(source.value)
199 where = spec.display_name if spec is not None else "the configured server"
200 raise ValueError(
201 f"lilbee runs {where} models but doesn't download them. "
202 f"Add the model in {where}, then pick it here."
203 )
204 if not allow_unsupported:
205 self.enforce_arch_compat(model)
206 try:
207 return self._pull_native(
208 model, on_bytes=on_bytes, cancel=cancel, allow_unsupported=allow_unsupported
209 )
210 finally:
211 self._invalidate_installed_cache()
213 def enforce_arch_compat(self, ref: str) -> None:
214 """Raise UnsupportedArchError if *ref*'s architecture isn't in the supported set.
216 Public because the pull preflight on the HTTP surface runs the same
217 check before starting a download, so a caller learns the model is
218 unsupported before any bytes move.
219 """
220 from lilbee.app.services import get_services
221 from lilbee.catalog.compat import (
222 ModelCompat,
223 UnsupportedArchError,
224 classify,
225 resolve_arch_for_pull,
226 )
228 arch = resolve_arch_for_pull(ref, get_services().hf_client)
229 if classify(arch) is ModelCompat.UNSUPPORTED:
230 raise UnsupportedArchError(ref, arch)
232 def _pull_native(
233 self,
234 model: str,
235 *,
236 on_bytes: Callable[[int, int], None] | None = None,
237 cancel: CancelSignal | None = None,
238 allow_unsupported: bool = False,
239 ) -> Path:
240 """Download a featured or ad-hoc HuggingFace model to the native GGUF directory."""
241 # heavy: lilbee.catalog (>50ms; huggingface_hub fanout)
242 from lilbee.catalog import download_model, resolve_pull_target
243 from lilbee.modelhub.registry import register_downloaded_model
245 entry = resolve_pull_target(model)
246 if entry is None:
247 raise ModelNotFoundError(
248 f"Model '{model}' not recognized. "
249 "Pass a HuggingFace repo id (owner/name) or a featured model name."
250 )
251 entry = self._resolved_entry(entry, verify=not allow_unsupported)
252 path = download_model(
253 entry, on_progress=on_bytes, on_complete=register_downloaded_model, cancel=cancel
254 )
255 log.info("Downloaded %s to %s", model, path)
256 return path
258 @staticmethod
259 def _resolved_entry(entry: CatalogModel, *, verify: bool) -> CatalogModel:
260 """*entry* carrying the file the pull will fetch.
262 Resolution happens here because the engine's verdict is part of choosing
263 the file, not a check on the file already chosen. A repo can publish the
264 same weights in several packings, only some of which this build reads,
265 and the resolver walks its ranking until one of them answers. Passing the
266 resolved name down means the download does not work it out again.
268 A repo that will not resolve is not this check's verdict to deliver. The
269 download reports the gated repo or the missing file properly, so a
270 resolution failure here is left to it.
271 """
272 # heavy: lilbee.catalog (>50ms; huggingface_hub fanout)
273 from lilbee.catalog.download import resolve_filename
274 from lilbee.catalog.hf_client import hf_token
275 from lilbee.providers.fleet.loadability import assert_engine_can_load
277 token = hf_token()
279 def _can_load(hf_repo: str, filename: str) -> None:
280 assert_engine_can_load(hf_repo, filename, token)
282 try:
283 filename = resolve_filename(entry, can_load=_can_load if verify else None)
284 except (PermissionError, RuntimeError) as exc:
285 if isinstance(exc, UnsupportedQuantError):
286 raise
287 log.debug("Cannot name a file to fetch for %s: %s", entry.hf_repo, exc)
288 return entry
289 return replace(entry, gguf_filename=filename)
291 def remove(self, model: str, source: ModelSource | None = None) -> bool:
292 """Remove an installed native model. Returns True if removed.
294 lilbee removes only native GGUF models it downloaded. Local servers
295 (Ollama, LM Studio) are read-only: a model that lives on one is refused
296 (mirrors ``pull``), since its lifecycle is managed in that app. A bare
297 ``source`` is resolved so a local-server ref is caught either way.
298 """
299 effective = source if source is not None else self.get_source(model)
300 if effective is not None and effective is not ModelSource.NATIVE:
301 spec = local_server_for_key(effective.value)
302 where = spec.display_name if spec is not None else "the configured server"
303 raise ValueError(
304 f"lilbee runs {where} models but doesn't remove them. "
305 f"Manage them in {where} instead."
306 )
307 try:
308 return self._remove_native(model)
309 finally:
310 self._invalidate_installed_cache()
312 def _remove_native(self, model: str) -> bool:
313 if self._registry.remove(model):
314 log.info("Removed native model %s from registry", model)
315 return True
316 try:
317 path = validate_path_within(self._models_dir / model, self._models_dir)
318 except ValueError:
319 log.warning("Path traversal blocked: %s escapes %s", model, self._models_dir)
320 return False
321 if path.is_file():
322 path.unlink()
323 log.info("Removed native model %s", model)
324 return True
325 return False