Coverage for src/lilbee/catalog/download.py: 100%
216 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""GGUF download, mmproj resolution, post-download hooks."""
3import fnmatch
4import logging
5import os
6import re
7import shutil
8from collections.abc import Callable
9from http import HTTPStatus
10from pathlib import Path
11from typing import Any
13import httpx
14from pydantic import BaseModel
16from lilbee.catalog.download_progress import ProgressCallback, _ProgressTracker
17from lilbee.catalog.hf_client import (
18 DEFAULT_TIMEOUT,
19 HF_API_URL,
20 hf_headers,
21 hf_token,
22 repo_has_mmproj,
23)
24from lilbee.catalog.models import CatalogModel
25from lilbee.catalog.refs import DEFAULT_MMPROJ_PATTERN, pick_best_gguf
26from lilbee.catalog.types import ModelTask
27from lilbee.runtime.cancellation import TaskCancelledError
29CompleteCallback = Callable[[CatalogModel, Path], None]
31log = logging.getLogger(__name__)
34def _models_dir() -> Path:
35 """Deferred cfg read: a module-level cfg import is circular via Config()'s
36 model-ref validator (config -> model_ref -> catalog -> here -> config)."""
37 from lilbee.core.config.model import cfg
39 return cfg.models_dir
42class DownloadConfig(BaseModel):
43 model_config = {"arbitrary_types_allowed": True}
45 repo_id: str
46 filename: str
47 token: str | None
48 force_download: bool = False
49 cache_dir: str | None = None
50 tqdm_class: Any = None
53_BYTES_PER_GB = 1024**3
56def _repo_partial_bytes(models_dir: Path, hf_repo: str) -> int:
57 """Bytes an interrupted attempt at *hf_repo* already holds on disk.
59 A resume needs only the remainder, so these count toward available space.
60 """
61 from huggingface_hub.file_download import repo_folder_name
63 repo_dir = models_dir / repo_folder_name(repo_id=hf_repo, repo_type="model")
64 if not repo_dir.is_dir():
65 return 0
66 return sum(f.stat().st_size for f in repo_dir.glob("blobs/*.incomplete") if f.is_file())
69def _free_bytes(path: Path) -> int | None:
70 """Free space on the volume that will hold *path*, which need not exist yet.
72 Measured at the nearest existing ancestor, since shutil.disk_usage raises on
73 a missing path.
74 """
75 probe = path.resolve()
76 while True:
77 try:
78 return shutil.disk_usage(probe).free
79 except OSError:
80 if probe.parent == probe:
81 return None
82 probe = probe.parent
85def disk_shortfall(models_dir: Path, hf_repo: str, needed: int) -> str | None:
86 """Describe why *needed* bytes will not fit, or None when they will."""
87 if needed == _SIZE_UNKNOWN:
88 return None # offline or unresolvable; nothing to compare against
89 free = _free_bytes(models_dir)
90 if free is None:
91 return None # unmeasurable volume; let the download report the truth
92 available = free + _repo_partial_bytes(models_dir, hf_repo)
93 if needed <= available:
94 return None
95 return (
96 f"Not enough disk space for {hf_repo}: needs "
97 f"{needed / _BYTES_PER_GB:.1f} GB, {available / _BYTES_PER_GB:.1f} GB free."
98 )
101def _require_disk_space(entry: CatalogModel, models_dir: Path, needed: int) -> None:
102 """Refuse a download the disk cannot hold, naming the shortfall.
104 huggingface_hub only warns, and the xet path reports a full disk as a
105 reconstruction error naming neither the disk nor the file.
106 """
107 message = disk_shortfall(models_dir, entry.hf_repo, needed)
108 if message is not None:
109 raise RuntimeError(message)
112_LOW_DISK_FLOOR = 512 * 1024**2
113"""Free bytes below which a failed download is reported as a full disk.
115Catches a volume that filled mid-transfer, which the pre-flight cannot see."""
118def _raise_if_disk_exhausted(
119 entry: CatalogModel, config: DownloadConfig, cause: BaseException
120) -> None:
121 """Re-raise a failed download as a disk problem when the volume is full.
123 Low free space is a heuristic, not a diagnosis, so *cause* stays in the
124 message.
125 """
126 if config.cache_dir is None:
127 return
128 try:
129 free = shutil.disk_usage(config.cache_dir).free
130 except OSError:
131 return # the path went away with the failure; leave the original error
132 if free >= _LOW_DISK_FLOOR:
133 return
134 raise RuntimeError(
135 f"Ran out of disk space downloading {entry.hf_repo}: "
136 f"{free / _BYTES_PER_GB:.1f} GB free. {type(cause).__name__}: {cause}"
137 ) from None
140_XET_CANCELLED_MARKER = "Operation cancelled"
143def abort_active_download() -> None:
144 """Stop the xet transfer running in this process.
146 Aborts at session granularity; hf_xet exposes nothing finer.
147 """
148 from huggingface_hub.utils._xet import abort_xet_session
150 abort_xet_session()
153_XET_HIGH_PERFORMANCE_ENV = "HF_XET_HIGH_PERFORMANCE"
156def _apply_fast_download_mode() -> None:
157 """Publish the high-performance setting to xet before it builds a session.
159 hf_xet reads it from the environment in Rust and caches it when the session
160 is built, so a change lands on restart.
161 """
162 # circular: catalog.download -> core.config via cfg, the same cycle
163 # _models_dir documents (config -> model_ref -> catalog -> here).
164 from lilbee.core.config.model import cfg
166 if cfg.fast_model_downloads:
167 os.environ[_XET_HIGH_PERFORMANCE_ENV] = "1"
168 else:
169 os.environ.pop(_XET_HIGH_PERFORMANCE_ENV, None)
172def _hf_download_or_translate(entry: CatalogModel, config: DownloadConfig) -> Path:
173 """Run the HF download and translate every error class into a clean exception."""
174 from huggingface_hub import hf_hub_download
175 from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError
177 _apply_fast_download_mode()
178 try:
179 return Path(hf_hub_download(**config.model_dump(exclude_none=True)))
180 except TaskCancelledError:
181 raise
182 except GatedRepoError:
183 raise PermissionError(
184 f"{entry.hf_repo} requires HuggingFace authentication. "
185 "Set HF_TOKEN env var or visit the repo page to request access."
186 ) from None
187 except RepositoryNotFoundError:
188 raise RuntimeError(f"Repository {entry.hf_repo!r} not found on HuggingFace.") from None
189 except (httpx.TimeoutException, httpx.ConnectError) as exc:
190 raise RuntimeError(f"Network error downloading {entry.hf_repo}: {exc}") from None
191 except OSError as exc:
192 raise RuntimeError(f"I/O error downloading {entry.hf_repo}: {exc}") from None
193 except Exception as exc:
194 if _XET_CANCELLED_MARKER in str(exc):
195 # An aborted session surfaces as a bare RuntimeError.
196 raise TaskCancelledError(str(exc)) from None
197 _raise_if_disk_exhausted(entry, config, exc)
198 raise RuntimeError(
199 f"Failed to download {entry.hf_repo}: {type(exc).__name__}: {exc}"
200 ) from None
203_SPLIT_SHARD_RE = re.compile(r"^(?P<base>.+)-(?P<idx>\d{5})-of-(?P<total>\d{5})\.gguf$")
206def split_shard_filenames(filename: str) -> list[str]:
207 """Return every shard of a split GGUF in order, or ``[filename]`` if it isn't split.
209 A split GGUF names its parts ``<base>-00001-of-0000N.gguf`` through
210 ``<base>-0000N-of-0000N.gguf``. llama.cpp loads the whole set from the first
211 shard but needs every part on disk, so the catalog must fetch all of them and
212 only consider the model installed once the full set is present.
213 """
214 match = _SPLIT_SHARD_RE.match(filename)
215 if match is None:
216 return [filename]
217 base = match.group("base")
218 total = int(match.group("total"))
219 return [f"{base}-{index:05d}-of-{total:05d}.gguf" for index in range(1, total + 1)]
222def download_model(
223 entry: CatalogModel,
224 *,
225 on_progress: ProgressCallback | None = None,
226 on_complete: CompleteCallback | None = None,
227) -> Path:
228 """Download a GGUF model from HuggingFace to the models dir.
229 Uses huggingface_hub for resumable downloads, caching, and auth.
230 The optional *on_progress(downloaded, total)* callback receives byte counts.
231 The optional *on_complete(entry, file_path)* callback runs after the file
232 is on disk; modelhub uses it to write a registry manifest. For vision
233 models, also downloads the mmproj (CLIP projection) file.
235 A split GGUF has every shard fetched before the model is finalized, so the
236 registry manifest (and thus "installed") only lands once the full set is on
237 disk; an interrupted multi-part pull leaves the model not-installed and
238 re-pullable rather than registered-but-unloadable.
240 Raises:
241 PermissionError: gated repo requiring authentication
242 RuntimeError: repo not found or download failure with details
243 """
244 models_dir = _models_dir()
245 models_dir.mkdir(parents=True, exist_ok=True)
247 filename = resolve_filename(entry)
248 shards = split_shard_filenames(filename)
249 dest = models_dir / shards[0]
250 if all(
251 (models_dir / shard).exists()
252 and _cached_file_is_complete(entry.hf_repo, shard, models_dir / shard)
253 for shard in shards
254 ):
255 log.info("Model already downloaded: %s", dest)
256 if on_progress is not None:
257 size = sum((models_dir / shard).stat().st_size for shard in shards)
258 on_progress(size, size) # Report 100% immediately (every shard)
259 return _finalize_download(entry, dest, on_progress=on_progress, on_complete=on_complete)
261 shard_sizes = [fetch_expected_file_size(entry.hf_repo, shard) for shard in shards]
262 sizes_known = all(size != _SIZE_UNKNOWN for size in shard_sizes)
263 _require_disk_space(entry, models_dir, sum(shard_sizes) if sizes_known else 0)
265 # Sum the shard sizes up front so a multi-shard pull reports one monotonic
266 # 0->100% against the real total, not N separate per-shard cycles. Only use
267 # the sum when every shard size is known (0 = unresolved/offline); a partial
268 # sum would undercount the total and let progress run past 100%.
269 grand_total = sum(shard_sizes) if len(shards) > 1 and sizes_known else 0
270 tracker = _ProgressTracker(on_progress, grand_total=grand_total) if on_progress else None
271 shard_paths: list[Path] = []
272 for shard in shards:
273 log.info("Downloading %s/%s → %s", entry.hf_repo, shard, models_dir)
274 config = DownloadConfig(
275 repo_id=entry.hf_repo,
276 filename=shard,
277 token=hf_token(),
278 cache_dir=str(models_dir),
279 tqdm_class=tracker.make_tqdm_class() if tracker else None,
280 )
281 shard_path = _hf_download_or_translate(entry, config)
282 shard_paths.append(shard_path)
283 if tracker is not None:
284 tracker.shard_done(shard_path.stat().st_size)
285 first_shard_path = shard_paths[0] # the 00001-of-N shard llama.cpp loads from
287 if on_progress:
288 total_size = sum(path.stat().st_size for path in shard_paths)
289 if not tracker or not tracker.was_used:
290 log.info("Model found in HuggingFace cache: %s", first_shard_path)
291 on_progress(total_size, total_size)
292 return _finalize_download(
293 entry, first_shard_path, on_progress=on_progress, on_complete=on_complete
294 )
297def _finalize_download(
298 entry: CatalogModel,
299 dest: Path,
300 *,
301 on_progress: ProgressCallback | None = None,
302 on_complete: CompleteCallback | None = None,
303) -> Path:
304 """Run post-download hooks: registry write (via on_complete) + mmproj fetch.
306 The mmproj is fetched whenever the repo ships one, not only for VISION-task
307 entries: dual-use VL repos (Qwen-VL, InternVL, SmolVLM, gemma-3) classify as
308 chat by name and arch, and without their projector the vision role dies at
309 plan time with a missing-mmproj warning a re-pull cannot cure.
310 """
311 if on_complete is not None:
312 on_complete(entry, dest)
313 if entry.task == ModelTask.VISION or repo_has_mmproj(entry.hf_repo):
314 download_mmproj(entry, on_progress=on_progress)
315 return dest
318def download_mmproj(
319 entry: CatalogModel,
320 *,
321 on_progress: ProgressCallback | None = None,
322) -> Path | None:
323 """Download the mmproj (CLIP projection) file for a vision model.
324 Returns the path to the downloaded file, or None if no mmproj is configured.
325 The optional ``on_progress`` callback receives ``(downloaded, total)`` byte
326 counts and is wired through the same tqdm hook used by the main download.
327 """
328 mmproj_filename = _resolve_mmproj_filename(entry.hf_repo, DEFAULT_MMPROJ_PATTERN)
329 if not mmproj_filename:
330 log.warning("Could not resolve mmproj file for %s", entry.hf_repo)
331 return None
333 models_dir = _models_dir()
334 tracker = _ProgressTracker(on_progress) if on_progress else None
335 log.info("Downloading mmproj %s/%s → %s", entry.hf_repo, mmproj_filename, models_dir)
336 _require_disk_space(entry, models_dir, fetch_expected_file_size(entry.hf_repo, mmproj_filename))
337 # The projector gets the same error translation as the GGUF.
338 path = _hf_download_or_translate(
339 entry,
340 DownloadConfig(
341 repo_id=entry.hf_repo,
342 filename=mmproj_filename,
343 token=hf_token(),
344 cache_dir=str(models_dir),
345 tqdm_class=tracker.make_tqdm_class() if tracker else None,
346 ),
347 )
348 if on_progress is not None and (not tracker or not tracker.was_used):
349 # Cache hit: HF returned the cached path without invoking tqdm.
350 size = path.stat().st_size
351 on_progress(size, size)
352 return path
355def _resolve_mmproj_filename(hf_repo: str, pattern: str) -> str | None:
356 """Resolve an mmproj filename pattern to a concrete filename via the HF API."""
357 if "*" not in pattern:
358 return pattern
360 try:
361 resp = httpx.get(
362 f"{HF_API_URL}/{hf_repo}",
363 timeout=DEFAULT_TIMEOUT,
364 headers=hf_headers(),
365 )
366 resp.raise_for_status()
367 siblings = resp.json().get("siblings", [])
368 except Exception as exc:
369 log.warning("Cannot query mmproj files for %s: %s", hf_repo, exc)
370 return None
372 mmproj_files: list[str] = [
373 s.get("rfilename", "") for s in siblings if fnmatch.fnmatch(s.get("rfilename", ""), pattern)
374 ]
375 if not mmproj_files:
376 return None
378 # Prefer an F16 mmproj when one is offered; otherwise take the first match.
379 for preference in ("f16", "F16"):
380 for f in mmproj_files:
381 if preference in f:
382 return f
383 return mmproj_files[0]
386def resolve_filename(entry: CatalogModel) -> str:
387 """Resolve a GGUF filename pattern to the best concrete filename.
388 For exact filenames, return as-is. For wildcards, query the HF API
389 and pick the best quantization (prefer Q4_K_M for balance of size/quality).
390 """
391 if "*" not in entry.gguf_filename:
392 return entry.gguf_filename
394 try:
395 resp = httpx.get(
396 f"{HF_API_URL}/{entry.hf_repo}",
397 timeout=DEFAULT_TIMEOUT,
398 headers=hf_headers(),
399 )
400 if resp.status_code == HTTPStatus.UNAUTHORIZED:
401 raise PermissionError(
402 f"{entry.hf_repo} requires HuggingFace authentication. "
403 "Set HF_TOKEN env var or visit the repo page to request access."
404 )
405 resp.raise_for_status()
406 siblings = resp.json().get("siblings", [])
407 except PermissionError:
408 raise
409 except Exception as exc:
410 raise RuntimeError(f"Cannot query files for {entry.hf_repo}: {exc}") from exc
412 gguf_files = [
413 s.get("rfilename", "") for s in siblings if s.get("rfilename", "").endswith(".gguf")
414 ]
415 if not gguf_files:
416 raise RuntimeError(f"No GGUF files found in {entry.hf_repo}")
418 return pick_best_gguf(gguf_files)
421_SIZE_UNKNOWN = 0
424def _cached_file_is_complete(hf_repo: str, filename: str, dest: Path) -> bool:
425 """Decide whether an existing cached file may be accepted as complete.
427 Verifies the on-disk byte size against the size HuggingFace reports for
428 *filename*. A mismatch means a truncated / corrupt download, so the file
429 is rejected and re-fetched. When the size can't be fetched (offline, API
430 error) it stays unknown and the cached file is accepted: there's nothing
431 to verify against and refusing would block all offline reuse.
432 """
433 expected = fetch_expected_file_size(hf_repo, filename)
434 if expected == _SIZE_UNKNOWN:
435 return True
436 actual = dest.stat().st_size
437 if actual == expected:
438 return True
439 log.warning(
440 "Cached %s is %d bytes but HuggingFace reports %d; re-downloading",
441 dest,
442 actual,
443 expected,
444 )
445 return False
448def _hf_file_size(hf_repo: str, filename: str) -> int | None:
449 """Byte size huggingface_hub resolves for *filename* (None if unreported)."""
450 from huggingface_hub import get_hf_file_metadata, hf_hub_url
452 return get_hf_file_metadata(hf_hub_url(hf_repo, filename), token=hf_token()).size
455def fetch_expected_file_size(hf_repo: str, filename: str) -> int:
456 """Return the byte size huggingface_hub reports for *filename*, or _SIZE_UNKNOWN.
458 Resolves via hf_hub's own file metadata (correct revision, redirects, and
459 LFS/Xet handled uniformly) instead of scraping the repo tree. Returns 0 when
460 offline or unresolvable, in which case the caller keeps the cached file.
461 """
462 try:
463 return _hf_file_size(hf_repo, filename) or _SIZE_UNKNOWN
464 except Exception:
465 return _SIZE_UNKNOWN