diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 3e93a267878..fec5214cf92 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -9,6 +9,7 @@ import atexit import contextlib +import functools import json import os import re @@ -999,6 +1000,283 @@ def _cached_colocated_split_main( return None +def _cached_variant_resolution(repo_id: str, hf_variant: str) -> tuple[Optional[str], list[str]]: + """Find a cached main GGUF and its shards for a variant.""" + candidate = next(_cached_variant_candidates(repo_id, hf_variant), None) + if candidate is None: + return None, [] + _, main, shards, _ = candidate + return main, shards + + +def _cached_variant_candidates( + repo_id: str, + hf_variant: str, + *, + require_mmproj: bool = False, +) -> Generator[tuple[str, str, list[str], Path], None, None]: + """Yield complete cached variant copies in snapshot preference order.""" + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + cached_files = _gguf_snapshot_files(snap) + matches = _gguf_files_for_variant(cached_files, hf_variant) + if not matches: + continue + main = matches[0] + shards = _gguf_extra_shards(matches, main) + split = _SHARD_FULL_RE.match(main) + if split: + numbers = { + int(match.group(2)) + for path in [main, *shards] + if (match := _SHARD_FULL_RE.match(path)) + } + if numbers != set(range(1, int(split.group(3)) + 1)): + continue + main_path = snap.joinpath(*main.replace("\\", "/").split("/")) + if not main_path.is_file() or not _snapshot_has_all_shards( + str(main_path), main, shards, {} + ): + continue + if require_mmproj and not _pick_mmproj(cached_files): + continue + yield str(main_path), main, shards, snap + except Exception as e: + logger.debug(f"Cache lookup for variant failed: {e}") + + +def _cached_candidate_matches_revision_size( + repo_id: str, candidate: tuple[str, str, list[str], Path], hf_token: Optional[str] +) -> bool: + """Check cached byte sizes against the snapshot's own Hub revision. + + A snapshot pointer is normally published only after its blob is complete. + When the old revision is still queryable, also compare every weight file's + size so a manually truncated cache entry is not treated as reusable. If + metadata cannot be reached, retain the cache's normal offline semantics. + """ + main_path, main, shards, snap = candidate + paths = [main, *shards] + try: + from huggingface_hub import get_paths_info + infos = list( + get_paths_info( + repo_id, + paths, + revision = snap.name, + token = hf_token, + ) + ) + except Exception as e: + logger.debug( + "Could not size-check cached GGUF %s at revision %s: %s", + repo_id, + snap.name, + e, + ) + return True + + if not infos: + # The Hub answers an unknown (e.g. force-pushed away) revision with an + # empty result, not an error; treat it like unreachable metadata. + return True + expected_sizes = {info.path: info.size for info in infos if info.size is not None} + if any(path not in expected_sizes for path in paths): + return False + try: + if os.path.getsize(main_path) < expected_sizes[main]: + return False + except OSError: + return False + return _snapshot_has_all_shards(main_path, main, shards, expected_sizes) + + +def _cached_complete_candidate( + repo_id: str, gguf_filename: Optional[str], shards: list[str] +) -> Optional[tuple[str, str, list[str], Path]]: + """Return one complete exact-filename cache candidate with snapshot context.""" + if not gguf_filename: + return None + if shards: + main_path = _cached_colocated_split_main(repo_id, gguf_filename, shards, {}) + else: + m = _SHARD_FULL_RE.match(gguf_filename) + if m and int(m.group(3)) > 1: + return None + main_path = _cached_hf_snapshot_file(repo_id, gguf_filename) + if main_path is None: + return None + snap = _snapshot_dir_of(main_path) + if snap is None: + return None + return main_path, gguf_filename, shards, snap + + +def cached_gguf_for_load( + hf_repo: str, + hf_variant: Optional[str], + *, + require_mmproj: bool = False, + verify_sizes: bool = False, + hf_token: Optional[str] = None, +) -> Optional[str]: + """Return a cached GGUF that can be loaded without downloading.""" + if not hf_variant: + return None + hf_repo = _resolve_repo_id_casing(hf_repo) + for candidate in _cached_variant_candidates( + hf_repo, + hf_variant, + require_mmproj = require_mmproj, + ): + if verify_sizes and not _cached_candidate_matches_revision_size( + hf_repo, candidate, hf_token + ): + continue + return candidate[0] + return None + + +def _snapshot_dir_of(path: str) -> Optional[Path]: + """Return the HF cache snapshot containing path, if any.""" + try: + p = Path(os.path.abspath(path)) + except OSError: + return None + for ancestor in p.parents: + if ancestor.parent.name == "snapshots": + return ancestor + return None + + +def _companion_snapshot_sibling( + near_path: str, pick: Callable[[list[str]], Optional[str]] +) -> Optional[str]: + """Find a companion in the same snapshot as near_path.""" + snap = _snapshot_dir_of(near_path) + if snap is None: + return None + try: + sibling = pick(_gguf_snapshot_files(snap)) + except Exception: + return None + if not sibling: + return None + candidate = snap / sibling + return str(candidate) if candidate.is_file() else None + + +def _pick_mmproj(candidates: list[str]) -> Optional[str]: + mmproj_files = sorted( + f for f in candidates if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower() + ) + if not mmproj_files: + return None + return next((f for f in mmproj_files if f.lower().endswith("-f16.gguf")), mmproj_files[0]) + + +def _hub_download_in_flight(hf_repo: str) -> bool: + try: + from hub.utils.download_registry import get_models_registry + return bool(get_models_registry().active_job_refs(hf_repo)) + except Exception: + return False + + +def _hub_download_blocks_gguf_load( + hf_repo: str, + hf_variant: Optional[str], + *, + require_mmproj: bool = False, + hf_token: Optional[str] = None, +) -> bool: + """Whether an active Hub job makes this GGUF load unsafe. + + Same-variant jobs can reclaim the stale snapshot a load would reuse, so + they always block. Other jobs block only when this load lacks a complete + cached copy and would write to the shared cache itself. + """ + try: + from hub.utils.download_registry import get_models_registry + + registry = get_models_registry() + if not registry.active_job_refs(hf_repo): + return False + if registry.has_active_variant(hf_repo, hf_variant): + return True + except Exception: + return False + return ( + cached_gguf_for_load( + hf_repo, + hf_variant, + require_mmproj = require_mmproj, + verify_sizes = True, + hf_token = hf_token, + ) + is None + ) + + +# Active GGUF loads by normalized repo ID. +_LOADS_IN_FLIGHT: dict[str, int] = {} +_LOADS_IN_FLIGHT_LOCK = threading.Lock() + + +@contextlib.contextmanager +def gguf_load_in_flight(hf_repo: Optional[str]): + """Track an HF GGUF load until the context exits.""" + key = (hf_repo or "").strip().lower() + if not key: + yield + return + with _LOADS_IN_FLIGHT_LOCK: + _LOADS_IN_FLIGHT[key] = _LOADS_IN_FLIGHT.get(key, 0) + 1 + try: + yield + finally: + with _LOADS_IN_FLIGHT_LOCK: + remaining = _LOADS_IN_FLIGHT.get(key, 1) - 1 + if remaining <= 0: + _LOADS_IN_FLIGHT.pop(key, None) + else: + _LOADS_IN_FLIGHT[key] = remaining + + +def hf_gguf_load_in_flight(hf_repo: str) -> bool: + """Return whether a GGUF load is active for hf_repo.""" + key = (hf_repo or "").strip().lower() + if not key: + return False + with _LOADS_IN_FLIGHT_LOCK: + return _LOADS_IN_FLIGHT.get(key, 0) > 0 + + +def _with_gguf_load_marker(load: Callable): + """Keep an HF repo marked for the full synchronous load call.""" + + @functools.wraps(load) + def wrapped(self, *args, **kwargs): + hf_repo = kwargs.get("hf_repo") + with gguf_load_in_flight(hf_repo): + if hf_repo and _hub_download_blocks_gguf_load( + hf_repo, + kwargs.get("hf_variant"), + require_mmproj = bool( + kwargs.get("is_vision") + and not extra_args_disable_mmproj(kwargs.get("extra_args")) + ), + hf_token = kwargs.get("hf_token"), + ): + raise RuntimeError( + f"'{hf_repo}' is currently being downloaded by the download manager" + ) + return load(self, *args, **kwargs) + + return wrapped + + def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]: m = _SHARD_FULL_RE.match(first_shard) if not m: @@ -3806,13 +4084,13 @@ def _find_smallest_fitting_variant( hf_repo: str, free_bytes: int, hf_token: Optional[str] = None, - ) -> Optional[tuple[str, int]]: + ) -> Optional[tuple[str, int, list[str]]]: """Find the smallest GGUF variant (including all shards) that fits. Groups split shards by variant prefix and sums their sizes (e.g. UD-Q4_K_XL with 9 shards of 50 GB each = 450 GB total). - Returns (first_shard_filename, total_size_bytes) or None. + Returns (first_shard_filename, total_size_bytes, extra_shards) or None. """ try: from huggingface_hub import get_paths_info, list_repo_files @@ -3848,9 +4126,13 @@ def _find_smallest_fitting_variant( # Smallest that fits variant_sizes.sort(key = lambda x: x[1]) - for first_file, total_size, _ in variant_sizes: + for first_file, total_size, shard_files in variant_sizes: if total_size > 0 and total_size <= free_bytes: - return first_file, total_size + return ( + first_file, + total_size, + [path for path in sorted(shard_files) if path != first_file], + ) return None except Exception: @@ -4446,42 +4728,52 @@ def _download_gguf( except Exception as e: logger.warning(f"Could not list repo files: {e}") - # Offline: resolve variant -> filename from the local HF cache. - # The heuristic below assumes filenames echo the repo name, which - # breaks for e.g. Qwen3.6-27B-MTP-GGUF (no "MTP" in file). Match - # against the rel path (not just basename) so subdir layouts like - # ``BF16/foo.gguf`` are findable. + # Fall back to the local cache when the repo listing is unavailable. if not gguf_filename: - try: - from utils.models.model_config import _iter_hf_cache_snapshots - for snap in _iter_hf_cache_snapshots(hf_repo): - cached_files = _gguf_snapshot_files(snap) - matches = _gguf_files_for_variant(cached_files, hf_variant) - if not matches: - continue - gguf_filename = matches[0] - gguf_extra_shards = _gguf_extra_shards(matches, gguf_filename) - logger.info( - "Resolved variant %s -> %s from local HF cache", - hf_variant, - gguf_filename, - ) - break - except Exception as e: - logger.debug(f"Offline cache lookup for variant failed: {e}") + cached_name, cached_shards = _cached_variant_resolution(hf_repo, hf_variant) + if cached_name: + gguf_filename = cached_name + gguf_extra_shards = cached_shards + logger.info( + "Resolved variant %s -> %s from local HF cache", + hf_variant, + gguf_filename, + ) if not gguf_filename: repo_name = hf_repo.split("/")[-1].replace("-GGUF", "") gguf_filename = f"{repo_name}-{hf_variant}.gguf" + # Prefer the existing model. Updates use force=True to fetch a new revision. + if not force: + if hf_variant: + # Resolve by variant so a newer revision's filename does not hide + # the complete older copy. Size-check against that older snapshot's + # own revision when its metadata remains available. + cached_main = cached_gguf_for_load( + hf_repo, + hf_variant, + verify_sizes = True, + hf_token = hf_token, + ) + else: + candidate = _cached_complete_candidate(hf_repo, gguf_filename, gguf_extra_shards) + cached_main = ( + candidate[0] + if candidate is not None + and _cached_candidate_matches_revision_size(hf_repo, candidate, hf_token) + else None + ) + if cached_main is not None: + logger.info(f"Reusing cached GGUF: {cached_main}") + return cached_main + # Check disk space; fall back to a smaller variant if needed all_gguf_files = [gguf_filename] + gguf_extra_shards - expected_sizes: dict[str, int] = {} try: from huggingface_hub import get_paths_info, try_to_load_from_cache path_infos = list(get_paths_info(hf_repo, all_gguf_files, token = hf_token)) - expected_sizes = {p.path: p.size for p in path_infos if p.size} total_bytes = sum((p.size or 0) for p in path_infos) # Subtract bytes already in the HF cache so we only preflight @@ -4490,25 +4782,10 @@ def _download_gguf( # cold whenever free disk is below the full weight footprint, # even though nothing needs downloading. already_cached_bytes = 0 - # Cross-snapshot / case-variant cache reuse is offline-only (see the download - # path below); online, hf_hub_download fetches the current revision and - # resumes partials, so an old snapshot must not be counted as cached here or - # the preflight would under-count the download and skip the disk fallback. + # Count only files that can resume this download. offline = _hf_env_offline() - # A split GGUF whose shards are not co-located in a single snapshot is - # refetched as a whole set later, so it must not be counted as cached here. - split_needs_refetch = False - if offline and not force and gguf_extra_shards: - # Scan all snapshots for one that holds the whole set co-located, so a - # newer snapshot with only the first shard does not mask an older - # complete one and needlessly trip the disk fallback. - if ( - _cached_colocated_split_main( - hf_repo, gguf_filename, gguf_extra_shards, expected_sizes - ) - is None - ): - split_needs_refetch = True + # Offline split sets are reusable only when every shard shares a snapshot. + split_needs_refetch = bool(offline and not force and gguf_extra_shards) if not force and not split_needs_refetch: for p in path_infos: if not p.size: @@ -4569,32 +4846,26 @@ def _download_gguf( hf_token, ) if smaller: - fallback_file, fallback_size = smaller + fallback_file, fallback_size, fallback_shards = smaller logger.info( f"Selected variant too large ({total_gb:.1f} GB), " f"falling back to {fallback_file} ({fallback_size / (1024**3):.1f} GB)" ) gguf_filename = fallback_file - _m = _SHARD_RE.match(gguf_filename) - _prefix = _m.group(1) if _m else None - if _prefix: - prefix_lower = _prefix.lower() - gguf_extra_shards = sorted( - f - for f in all_gguf_files - if f.lower().startswith(prefix_lower) - and f != gguf_filename - and not _is_companion_gguf_path(f) + gguf_extra_shards = fallback_shards + + # The selected fallback is a new load target. Apply the + # same any-revision reuse policy before starting a fetch. + fallback_candidate = _cached_complete_candidate( + hf_repo, gguf_filename, gguf_extra_shards + ) + if fallback_candidate is not None and ( + _cached_candidate_matches_revision_size( + hf_repo, fallback_candidate, hf_token ) - else: - gguf_extra_shards = [] - # Record the fallback's size so the later cache-reuse probe can - # size-verify it; only for a single-file fallback, since - # _find_smallest_fitting_variant returns the whole-variant size - # and using that as the first shard's expected size would reject - # a valid cached first shard of a split fallback. - if not gguf_extra_shards: - expected_sizes[fallback_file] = fallback_size + ): + logger.info(f"Reusing cached fallback GGUF: {fallback_candidate[0]}") + return fallback_candidate[0] else: raise RuntimeError( f"Not enough disk space to download any variant. " @@ -4614,45 +4885,25 @@ def _download_gguf( raise RuntimeError("Cancelled") dl_start = time.monotonic() # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached. - local_path = None - # Reuse a cached copy from another snapshot / case-variant repo dir only when - # offline. Online, fall through to hf_hub_download so its revision/etag check - # fetches the current file (and resumes a partial) instead of serving a stale - # same-name blob from an older revision. - if not force and _hf_env_offline(): - if gguf_extra_shards: - # A split GGUF must load every shard from one snapshot; reuse only a - # snapshot that holds the whole set co-located, scanning past a newer - # snapshot that has just the first shard while an older one is complete. - local_path = _cached_colocated_split_main( - hf_repo, gguf_filename, gguf_extra_shards, expected_sizes - ) - else: - local_path = _cached_hf_snapshot_file( - hf_repo, - gguf_filename, - expected_size = expected_sizes.get(gguf_filename), - ) - if local_path is None: - local_path = hf_hub_download_with_xet_fallback( + local_path = hf_hub_download_with_xet_fallback( + hf_repo, + gguf_filename, + hf_token, + cancel_event = cancel_event, + on_status = lambda m: logger.info(m), + force_download = force, + ) + for shard in gguf_extra_shards: + if cancel_event.is_set(): + raise RuntimeError("Cancelled") + logger.info(f"Resolving GGUF shard: {shard}") + hf_hub_download_with_xet_fallback( hf_repo, - gguf_filename, + shard, hf_token, cancel_event = cancel_event, - on_status = lambda m: logger.info(m), force_download = force, ) - for shard in gguf_extra_shards: - if cancel_event.is_set(): - raise RuntimeError("Cancelled") - logger.info(f"Resolving GGUF shard: {shard}") - hf_hub_download_with_xet_fallback( - hf_repo, - shard, - hf_token, - cancel_event = cancel_event, - force_download = force, - ) except Exception as e: if isinstance(e, RuntimeError) and "Cancelled" in str(e): raise @@ -4675,10 +4926,12 @@ def _download_companion_gguf( pick: Callable[[list[str]], Optional[str]], label: str, cancel_event: Optional[threading.Event] = None, + near_path: Optional[str] = None, ) -> Optional[str]: """Resolve and fetch a companion GGUF (mmproj / MTP drafter) by name. - Tries the live repo file list, then the local HF cache snapshots + Prefers a companion co-located with ``near_path``'s cache snapshot, + then tries the live repo file list, then the local HF cache snapshots (offline, same fallback as _download_gguf), then hf_hub_download. Runs WITHOUT self._lock (like _download_gguf); honors _cancel_event so an /unload between the main download and here skips the fetch. @@ -4688,6 +4941,17 @@ def _download_companion_gguf( if cancel_event.is_set(): return None + # Keep companion files in the main GGUF's snapshot. + if near_path: + cached = _companion_snapshot_sibling(near_path, pick) + if cached: + logger.info("Reusing cached %s: %s", label, cached) + return cached + + if _hub_download_in_flight(hf_repo): + logger.info("Skipping %s download while a hub download is active", label) + return None + target: Optional[str] = None from huggingface_hub import list_repo_files @@ -4760,33 +5024,23 @@ def _download_mmproj( hf_repo: str, hf_token: Optional[str] = None, cancel_event: Optional[threading.Event] = None, + near_path: Optional[str] = None, ) -> Optional[str]: """Download the mmproj (vision projection) file from a GGUF repo. Prefers mmproj-F16.gguf, else any mmproj*.gguf. Returns the local path, or None if none exists. ``cancel_event`` overrides - ``self._cancel_event`` (defaults to it). + ``self._cancel_event`` (defaults to it). ``near_path`` prefers a + copy co-located with the main GGUF's cache snapshot. """ - def _pick_mmproj(candidates: list[str]) -> Optional[str]: - mmproj_files = sorted( - f - for f in candidates - if f.lower().endswith(".gguf") and "mmproj" in Path(f).name.lower() - ) - if not mmproj_files: - return None - for f in mmproj_files: - if f.lower().endswith("-f16.gguf"): - return f - return mmproj_files[0] - return self._download_companion_gguf( hf_repo = hf_repo, hf_token = hf_token, pick = _pick_mmproj, label = "mmproj", cancel_event = cancel_event, + near_path = near_path, ) def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]: @@ -4817,6 +5071,7 @@ def _download_mtp( *, hf_repo: str, hf_token: Optional[str] = None, + near_path: Optional[str] = None, ) -> Optional[str]: """Download the separate MTP drafter (speculative head) from a GGUF repo. @@ -4828,16 +5083,6 @@ def _download_mtp( are intentionally skipped. Returns the local path, or None. """ - # Offline, reuse any drafter already on disk (a fresh copy can't be - # fetched). Online, _download_companion_gguf/hf_hub_download reuse the - # current cached file and refetch a changed one, so skip the probe here - # rather than pair new weights with a stale draft. - if _hf_env_offline(): - cached = self._cached_repo_mtp_drafter(hf_repo) - if cached: - logger.info(f"Reusing cached MTP drafter (offline): {cached}") - return cached - def _pick_mtp(candidates: list[str]) -> Optional[str]: # Root-level only: MTP/ subdir copies now share the mtp- prefix but # are explicit-selection, not auto-fetch (they'd sort ahead of root). @@ -4850,11 +5095,28 @@ def _pick_mtp(candidates: list[str]) -> Optional[str]: ) return mtp_files[0] if mtp_files else None + if near_path: + cached = _companion_snapshot_sibling(near_path, _pick_mtp) + if cached: + logger.info("Reusing cached MTP drafter: %s", cached) + return cached + + # Offline, reuse any drafter already on disk (a fresh copy can't be + # fetched). Online, _download_companion_gguf/hf_hub_download reuse the + # current cached file and refetch a changed one, so skip the probe here + # rather than pair new weights with a stale draft. + if _hf_env_offline(): + cached = self._cached_repo_mtp_drafter(hf_repo) + if cached: + logger.info(f"Reusing cached MTP drafter (offline): {cached}") + return cached + return self._download_companion_gguf( hf_repo = hf_repo, hf_token = hf_token, pick = _pick_mtp, label = "MTP drafter", + near_path = near_path, ) def _resolve_launch_mmproj_path( @@ -5436,6 +5698,7 @@ def _start_llama_process(self, cmd: list[str], env: dict) -> None: ) self._stdout_thread.start() + @_with_gguf_load_marker def load_model( self, *, @@ -5581,6 +5844,7 @@ def load_model( mmproj_path = self._download_mmproj( hf_repo = hf_repo, hf_token = hf_token, + near_path = model_path, ) # Auto-download the separate MTP drafter (e.g. Gemma) when # the requested spec mode can use it. Repos with the head @@ -5598,6 +5862,7 @@ def load_model( mtp_draft_path = self._download_mtp( hf_repo = hf_repo, hf_token = hf_token, + near_path = model_path, ) elif gguf_path: if not Path(gguf_path).is_file(): diff --git a/studio/backend/hub/services/models/downloads.py b/studio/backend/hub/services/models/downloads.py index c2ffe7bffc2..862af0141ab 100644 --- a/studio/backend/hub/services/models/downloads.py +++ b/studio/backend/hub/services/models/downloads.py @@ -60,6 +60,30 @@ def _job_status( return DownloadJobStatus(state = state, error = error, generation = generation) +def _load_in_flight(repo_id: str) -> bool: + try: + from core.inference.llama_cpp import hf_gguf_load_in_flight + return hf_gguf_load_in_flight(repo_id) + except Exception: + return False + + +def _load_in_flight_error(repo_id: str) -> HTTPException: + return HTTPException( + status_code = 409, + detail = ( + f"A model load for '{repo_id}' is in progress and may be " + "downloading it. Wait for the load to finish (or cancel it), " + "then start the download." + ), + ) + + +def _reject_if_load_in_flight(repo_id: str) -> None: + if _load_in_flight(repo_id): + raise _load_in_flight_error(repo_id) + + def _spawn_download_worker( repo_id: str, variant: Optional[str], @@ -89,6 +113,9 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional # Canonicalize so two different-cased paste-ins share one job + cache dir. repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model") + # Avoid concurrent writers to the same HF cache files. + _reject_if_load_in_flight(repo_id) + variant = (body.gguf_variant or "").strip() or None if variant is not None and not _is_valid_gguf_variant(variant): raise HTTPException( @@ -147,9 +174,12 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional blob_hashes = variant_blob_hashes, progress_blob_hashes = variant_progress_blob_hashes, completed_baseline_bytes = completed_baseline_bytes, + admission_check = lambda: not _load_in_flight(repo_id), ) generation = _registry.current_generation(key) if not claimed: + if claim_state == "admission_blocked": + raise _load_in_flight_error(repo_id) # claim_state is the blocking job's state. The client can attach only # when the blocker is this key's own in-flight job (adoptable); a # cross-variant conflict or in-progress delete is not accepted. diff --git a/studio/backend/hub/utils/download_registry.py b/studio/backend/hub/utils/download_registry.py index 274038b2926..b6bdee3bcef 100644 --- a/studio/backend/hub/utils/download_registry.py +++ b/studio/backend/hub/utils/download_registry.py @@ -47,7 +47,7 @@ import weakref from dataclasses import dataclass, field, replace from pathlib import Path -from typing import Iterator, Literal, Optional +from typing import Callable, Iterator, Literal, Optional from loggers import get_logger @@ -1028,6 +1028,7 @@ def claim( blob_hashes: Optional[frozenset[str]] = None, progress_blob_hashes: Optional[frozenset[str]] = None, completed_baseline_bytes: int = 0, + admission_check: Optional[Callable[[], bool]] = None, generation: Optional[int] = None, replace_active: bool = False, metadata_transport: Optional[str] = None, @@ -1038,6 +1039,13 @@ def claim( requested_hashes = blob_hashes or frozenset() requested_progress_hashes = progress_blob_hashes or frozenset() with self._lock: + # Run the final external admission check while the registry lock is + # held, immediately before inspecting and publishing active state. + # The GGUF load path establishes its marker before calling + # its active-job probe, so either this claim observes that marker + # or the load's later probe observes this claim. + if admission_check is not None and not admission_check(): + return False, "admission_blocked" deleting_scopes = self._deleting.get(repo) if deleting_scopes is not None and ( None in deleting_scopes or variant_from_key(key) in deleting_scopes @@ -1222,6 +1230,23 @@ def active_job_refs(self, repo_id: Optional[str] = None) -> list[ActiveDownloadR ) return refs + def has_active_variant(self, repo_id: str, variant: Optional[str]) -> bool: + """Whether an active model job targets this exact GGUF variant. + + Scans the job table rather than only ``_repo_active`` so an XET-to-HTTP + retry handoff remains visible while it has temporarily released its + active slot. + """ + repo_key = normalize_repo_key(repo_id) + target = (variant or "").strip().lower() or None + with self._lock: + for key, job in self._jobs.items(): + if _repo_of_key(key) != repo_key or job.state not in _ACTIVE_STATES: + continue + if self._active_job_variant_locked(key) == target: + return True + return False + def begin_delete( self, repo_id: str, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 0541a81ff20..81dea96274e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -20,6 +20,7 @@ import asyncio import threading import weakref +from contextlib import ExitStack import re as _re @@ -998,6 +999,7 @@ def _classify_llama_generation_error(exc: Exception) -> Optional[bool]: from core.inference.llama_server_args import ( _effective_tensor_parallel, _tensor_parallel_matches_loaded, + extra_args_disable_mmproj, parse_split_mode_override, resolve_tensor_parallel, strip_shadowing_flags, @@ -1035,6 +1037,7 @@ def _classify_llama_generation_error(exc: Exception) -> Optional[bool]: from core.inference.llama_server_args import ( _effective_tensor_parallel, _tensor_parallel_matches_loaded, + extra_args_disable_mmproj, parse_split_mode_override, resolve_tensor_parallel, strip_shadowing_flags, @@ -3968,6 +3971,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre native_grant_backed = False model_log_label = request.model_path + gguf_load_stack = ExitStack() try: # Validate user pass-through args up front so a managed-flag collision # returns 400 before any model work. @@ -4187,15 +4191,9 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre llama_backend = get_llama_cpp_backend() unsloth_backend = get_inference_backend() - # Unload any active Unsloth model to free VRAM (off the event loop: - # unload takes _gen_lock and can wait on an in-flight stream). - if unsloth_backend.active_model_name: - logger.info( - f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF" - ) - await asyncio.to_thread( - unsloth_backend.unload_model, unsloth_backend.active_model_name - ) + if config.gguf_hf_repo: + from core.inference.llama_cpp import gguf_load_in_flight + gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo)) # Inherit llama_extra_args from the previous load when the request # omits the field (the chat-settings Apply path doesn't round-trip @@ -4275,6 +4273,38 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre extra_llama_args, ) + # Block cache writes that would race the download manager. This runs + # after pass-through argument inheritance so a carried --no-mmproj + # changes the companion requirement exactly as it does for the load. + if config.gguf_hf_repo: + from core.inference.llama_cpp import _hub_download_blocks_gguf_load + if await asyncio.to_thread( + _hub_download_blocks_gguf_load, + config.gguf_hf_repo, + config.gguf_variant, + require_mmproj = bool( + config.is_vision and not extra_args_disable_mmproj(extra_llama_args) + ), + hf_token = request.hf_token, + ): + raise HTTPException( + status_code = 409, + detail = ( + f"'{model_log_label}' is currently being downloaded " + "by the download manager. Wait for the download to " + "finish (or cancel it), then load the model." + ), + ) + + # Unload any active Unsloth model only after every hub conflict check. + if unsloth_backend.active_model_name: + logger.info( + f"Unloading Unsloth model '{unsloth_backend.active_model_name}' before loading GGUF" + ) + await asyncio.to_thread( + unsloth_backend.unload_model, unsloth_backend.active_model_name + ) + # Route to HF or local mode based on config. Run in a thread so the # event loop stays free for progress polling and other requests # during the (potentially long) GGUF download + llama-server start. @@ -4639,6 +4669,8 @@ async def _attempt_gguf_load( logger.error(f"Error loading model: {e}", exc_info = True) msg = _maybe_unsupported_message(redacted_msg) raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}") + finally: + gguf_load_stack.close() def _requires_trust_remote_code_for_model( diff --git a/studio/backend/tests/test_gguf_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py new file mode 100644 index 00000000000..15d91cd3243 --- /dev/null +++ b/studio/backend/tests/test_gguf_load_cache_reuse.py @@ -0,0 +1,740 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for cached GGUF reuse and load/download exclusion. + +No GPU, network, or subprocesses are required. +""" + +from __future__ import annotations + +import asyncio +import sys +import threading +import types as _types +from pathlib import Path +from unittest.mock import patch + +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Stub optional dependencies before importing the modules under test. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) + +_structlog_stub = _types.ModuleType("structlog") +sys.modules.setdefault("structlog", _structlog_stub) + +try: + import httpx # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc_name in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "HTTPError", + "RequestError", + "HTTPStatusError", + ): + setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {})) + _httpx_stub.Response = type("Response", (), {}) + _httpx_stub.Request = type("Request", (), {}) + + class _FakeTimeout: + def __init__(self, *a, **kw): + pass + + _httpx_stub.Timeout = _FakeTimeout + _httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **kw: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, + ) + sys.modules.setdefault("httpx", _httpx_stub) + + +from huggingface_hub import constants as hf_constants + +from core.inference.llama_cpp import ( + LlamaCppBackend, + cached_gguf_for_load, + gguf_load_in_flight, + hf_gguf_load_in_flight, +) + + +REPO = "unsloth/gemma-test-GGUF" +VARIANT = "UD-Q4_K_XL" +MAIN = f"gemma-test-{VARIANT}.gguf" + + +def _build_cache( + root: Path, + repo_id: str, + files: dict[str, int], + *, + snapshot_sha: str = "a" * 40, +) -> Path: + """Create ``$root/models--/snapshots//`` for each entry.""" + repo_dir = root / f"models--{repo_id.replace('/', '--')}" + (repo_dir / "blobs").mkdir(parents = True, exist_ok = True) + snap = repo_dir / "snapshots" / snapshot_sha + snap.mkdir(parents = True, exist_ok = True) + for rel, size in files.items(): + full = snap / rel + full.parent.mkdir(parents = True, exist_ok = True) + full.write_bytes(b"\0" * size) + return snap + + +@pytest.fixture +def hf_cache(tmp_path, monkeypatch): + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + return tmp_path + + +def _fail_download(*_args, **_kwargs): + raise AssertionError("must reuse the cached GGUF instead of downloading") + + +def _fail_get_paths_info(*_args, **_kwargs): + raise AssertionError("cached reuse must return before the sizing preflight") + + +class TestLoadReusesCachedCopy: + def test_online_reuse_after_revision_bump(self, hf_cache): + """A new repo revision does not replace a complete cached model.""" + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", _fail_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / MAIN) + + def test_reuse_size_check_uses_cached_snapshot_revision(self, hf_cache): + """Current-revision size changes do not invalidate an older complete copy.""" + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + revisions: list[str | None] = [] + + def fake_get_paths_info( + _repo, + paths, + *, + revision = None, + token = None, + ): + revisions.append(revision) + size = 4 if revision == snap.name else 8 + return [_types.SimpleNamespace(path = path, size = size) for path in paths] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / MAIN) + assert revisions == [snap.name] + + def test_reuse_when_cached_revision_vanished_from_hub(self, hf_cache): + """The Hub answers an unknown revision with an empty result, not an error.""" + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", lambda *_a, **_k: []), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / MAIN) + + def test_truncated_cached_file_is_not_reused(self, hf_cache): + backend = LlamaCppBackend() + _build_cache(hf_cache, REPO, {MAIN: 4}) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo, + paths, + *, + revision = None, + token = None, + ): + return [_types.SimpleNamespace(path = path, size = 8) for path in paths] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert downloaded == [MAIN] + assert out == f"/fake/{REPO}/{MAIN}" + + def test_truncated_cached_split_shard_is_not_reused(self, hf_cache): + backend = LlamaCppBackend() + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf" + _build_cache(hf_cache, REPO, {shard1: 8, shard2: 4}) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo, + paths, + *, + revision = None, + token = None, + ): + return [_types.SimpleNamespace(path = path, size = 8) for path in paths] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert downloaded == [shard1, shard2] + assert out == f"/fake/{REPO}/{shard1}" + + def test_online_reuse_when_reupload_renamed_the_file(self, hf_cache): + """A renamed variant still reuses its cached file.""" + backend = LlamaCppBackend() + old_name = f"gemma-test-old-{VARIANT}.gguf" + snap = _build_cache(hf_cache, REPO, {old_name: 4}) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", _fail_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / old_name) + + def test_downloads_when_nothing_cached(self, hf_cache): + backend = LlamaCppBackend() + downloaded: list[str] = [] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 1) for p in paths if p is not None] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert downloaded == [MAIN] + assert out == f"/fake/{REPO}/{MAIN}" + + def test_force_redownloads_despite_cache(self, hf_cache): + """A forced download ignores a complete cached copy.""" + backend = LlamaCppBackend() + _build_cache(hf_cache, REPO, {MAIN: 4}) + downloaded: list[str] = [] + + def fake_download( + repo_id, + filename, + token = None, + **kwargs, + ): + assert kwargs.get("force_download") is True + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 1) for p in paths if p is not None] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT, force = True) + + assert downloaded == [MAIN] + assert out == f"/fake/{REPO}/{MAIN}" + + def test_split_reused_only_when_colocated(self, hf_cache): + backend = LlamaCppBackend() + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf" + snap = _build_cache(hf_cache, REPO, {shard1: 4, shard2: 4}) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]), + patch("huggingface_hub.get_paths_info", _fail_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / shard1) + + def test_partial_split_set_downloads(self, hf_cache): + """A partial split set is not reused.""" + backend = LlamaCppBackend() + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf" + _build_cache(hf_cache, REPO, {shard1: 4}) + downloaded: list[str] = [] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p is not None] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert downloaded == [shard1, shard2] + assert out == f"/fake/{REPO}/{shard1}" + + def test_reuse_prefers_newest_snapshot_after_update(self, hf_cache): + """Loads prefer the newest complete snapshot.""" + import os + + backend = LlamaCppBackend() + old_snap = _build_cache(hf_cache, REPO, {MAIN: 4}, snapshot_sha = "a" * 40) + new_snap = _build_cache(hf_cache, REPO, {MAIN: 6}, snapshot_sha = "b" * 40) + os.utime(old_snap, (1_000_000, 1_000_000)) + os.utime(new_snap, (2_000_000, 2_000_000)) + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", _fail_get_paths_info), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(new_snap / MAIN) + + def test_low_disk_fallback_reuses_cached_copy(self, hf_cache): + backend = LlamaCppBackend() + fallback = "gemma-test-Q2_K.gguf" + snap = _build_cache(hf_cache, REPO, {fallback: 4}) + + def fake_get_paths_info( + _repo, + paths, + *, + revision = None, + token = None, + ): + size = 4 if revision == snap.name else 100 + return [_types.SimpleNamespace(path = path, size = size) for path in paths] + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("shutil.disk_usage", lambda *_a, **_k: _types.SimpleNamespace(free = 10)), + patch.object( + backend, + "_find_smallest_fitting_variant", + lambda *_a, **_k: (fallback, 4, []), + ), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT) + + assert out == str(snap / fallback) + + def test_companion_prefers_main_snapshot_sibling(self, hf_cache): + """A cached mmproj is reused from the main model's snapshot.""" + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4, "mmproj-F16.gguf": 2}) + + def _fail_list(*_args, **_kwargs): + raise AssertionError("snapshot sibling must resolve without a repo listing") + + with patch("huggingface_hub.list_repo_files", _fail_list): + out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN)) + + assert out == str(snap / "mmproj-F16.gguf") + + def test_companion_finds_snapshot_through_hf_symlink(self, hf_cache): + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {}) + blobs = snap.parent.parent / "blobs" + main_blob = blobs / "main" + mmproj_blob = blobs / "mmproj" + main_blob.write_bytes(b"main") + mmproj_blob.write_bytes(b"mmproj") + try: + (snap / MAIN).symlink_to(main_blob) + (snap / "mmproj-F16.gguf").symlink_to(mmproj_blob) + except OSError as exc: + pytest.skip(f"symlinks unavailable: {exc}") + + with patch("huggingface_hub.list_repo_files", _fail_download): + out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN)) + + assert out == str(snap / "mmproj-F16.gguf") + + def test_companion_does_not_download_during_hub_job(self, hf_cache): + backend = LlamaCppBackend() + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + registry = _types.SimpleNamespace(active_job_refs = lambda _repo: [object()]) + + with ( + patch("huggingface_hub.list_repo_files", _fail_download), + patch("hub.utils.download_registry.get_models_registry", lambda: registry), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download), + ): + out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN)) + + assert out is None + + +class TestCachedGgufForLoadProbe: + def test_complete_copy_found(self, hf_cache): + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + assert cached_gguf_for_load(REPO, VARIANT) == str(snap / MAIN) + + def test_absent_copy_is_none(self, hf_cache): + assert cached_gguf_for_load(REPO, VARIANT) is None + + def test_partial_split_is_none(self, hf_cache): + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + _build_cache(hf_cache, REPO, {shard1: 4}) + assert cached_gguf_for_load(REPO, VARIANT) is None + + def test_partial_new_snapshot_does_not_hide_complete_split(self, hf_cache): + import os + + shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf" + old = _build_cache( + hf_cache, + REPO, + {shard1: 4, shard2: 4}, + snapshot_sha = "a" * 40, + ) + new = _build_cache(hf_cache, REPO, {shard1: 4}, snapshot_sha = "b" * 40) + os.utime(old, (1_000_000, 1_000_000)) + os.utime(new, (2_000_000, 2_000_000)) + + assert cached_gguf_for_load(REPO, VARIANT) == str(old / shard1) + + def test_split_requires_every_declared_shard(self, hf_cache): + shard1 = f"gemma-test-{VARIANT}-00001-of-00003.gguf" + shard2 = f"gemma-test-{VARIANT}-00002-of-00003.gguf" + _build_cache(hf_cache, REPO, {shard1: 4, shard2: 4}) + + assert cached_gguf_for_load(REPO, VARIANT) is None + + def test_required_mmproj_must_share_main_snapshot(self, hf_cache): + snap = _build_cache(hf_cache, REPO, {MAIN: 4}) + assert cached_gguf_for_load(REPO, VARIANT) == str(snap / MAIN) + assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) is None + + (snap / "mmproj-F16.gguf").write_bytes(b"mmproj") + assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) == str(snap / MAIN) + + def test_required_mmproj_scans_past_newer_main_only_snapshot(self, hf_cache): + import os + + old = _build_cache( + hf_cache, + REPO, + {MAIN: 4, "mmproj-F16.gguf": 2}, + snapshot_sha = "a" * 40, + ) + new = _build_cache(hf_cache, REPO, {MAIN: 4}, snapshot_sha = "b" * 40) + os.utime(old, (1_000_000, 1_000_000)) + os.utime(new, (2_000_000, 2_000_000)) + + assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) == str(old / MAIN) + + +class TestLoadHubDownloadExclusion: + def test_in_flight_marker_counts_and_normalizes_case(self): + assert not hf_gguf_load_in_flight(REPO) + with gguf_load_in_flight(REPO): + assert hf_gguf_load_in_flight(REPO.upper()) + with gguf_load_in_flight(REPO.lower()): + assert hf_gguf_load_in_flight(REPO) + assert hf_gguf_load_in_flight(REPO) + assert not hf_gguf_load_in_flight(REPO) + + def test_marker_noops_for_local_loads(self): + with gguf_load_in_flight(None): + assert not hf_gguf_load_in_flight("") + + def test_marker_cleared_on_exception(self): + with pytest.raises(RuntimeError): + with gguf_load_in_flight(REPO): + raise RuntimeError("boom") + assert not hf_gguf_load_in_flight(REPO) + + def test_hub_download_refused_while_load_in_flight(self): + from fastapi import HTTPException + + from hub.schemas.downloads import DownloadModelRequest + from hub.services.models import downloads as dl + + body = DownloadModelRequest(repo_id = REPO, gguf_variant = VARIANT) + with ( + patch.object(dl, "resolve_cached_repo_id_case", lambda repo_id, repo_type: repo_id), + gguf_load_in_flight(REPO), + ): + with pytest.raises(HTTPException) as exc_info: + asyncio.run(dl.download_model_response(body)) + + assert exc_info.value.status_code == 409 + assert "load" in exc_info.value.detail.lower() + + def test_hub_download_rechecks_marker_before_claim(self): + from fastapi import HTTPException + + from hub.schemas.downloads import DownloadModelRequest + from hub.services.models import downloads as dl + + scope = None + + def mark_load(*_args, **_kwargs): + nonlocal scope + if scope is None: + scope = gguf_load_in_flight(REPO) + scope.__enter__() + return frozenset() + + class _Registry: + def claim(self, *_args, admission_check, **_kwargs): + assert admission_check() is False + return False, "admission_blocked" + + def current_generation(self, _key): + return 0 + + registry = _Registry() + body = DownloadModelRequest(repo_id = REPO, gguf_variant = VARIANT) + try: + with ( + patch.object(dl, "resolve_cached_repo_id_case", lambda repo_id, repo_type: repo_id), + patch.object(dl.gguf_variants, "gguf_variant_blob_hashes", mark_load), + patch.object(dl, "_registry", registry), + ): + with pytest.raises(HTTPException) as exc_info: + asyncio.run(dl.download_model_response(body)) + finally: + if scope is not None: + scope.__exit__(None, None, None) + + assert exc_info.value.status_code == 409 + + def test_registry_admission_check_prevents_claim(self): + from hub.utils.download_registry import DownloadRegistry, TRANSPORT_HTTP + + registry = DownloadRegistry() + claimed, state = registry.claim( + f"{REPO}::{VARIANT}", + TRANSPORT_HTTP, + repo_type = "model", + repo_id = REPO, + variant = VARIANT, + admission_check = lambda: False, + ) + + assert claimed is False + assert state == "admission_blocked" + assert registry.active_jobs(REPO) == {} + + def test_same_variant_job_stays_visible_during_retry_handoff(self): + from hub.utils.download_registry import DownloadRegistry, TRANSPORT_XET + from core.inference.llama_cpp import _hub_download_blocks_gguf_load + + registry = DownloadRegistry() + key = f"{REPO}::{VARIANT}" + claimed, _ = registry.claim( + key, + TRANSPORT_XET, + repo_type = "model", + repo_id = REPO, + variant = VARIANT, + ) + assert claimed is True + assert registry.has_active_variant(REPO, VARIANT.lower()) is True + + registry.release_active_slot(key) + + assert registry.active_jobs(REPO) == {} + assert registry.active_job_refs(REPO) + assert registry.has_active_variant(REPO, VARIANT) is True + with ( + patch("hub.utils.download_registry.get_models_registry", lambda: registry), + patch( + "core.inference.llama_cpp.cached_gguf_for_load", + side_effect = AssertionError("same-variant jobs must block before cache reuse"), + ), + ): + assert _hub_download_blocks_gguf_load(REPO, VARIANT) is True + + registry.set_job(key, "complete") + assert registry.has_active_variant(REPO, VARIANT) is False + + def test_other_variant_job_still_allows_complete_cached_load(self): + from core.inference.llama_cpp import _hub_download_blocks_gguf_load + from hub.utils.download_registry import DownloadRegistry, TRANSPORT_HTTP + + registry = DownloadRegistry() + registry.claim( + f"{REPO}::Q8_0", + TRANSPORT_HTTP, + repo_type = "model", + repo_id = REPO, + variant = "Q8_0", + ) + with ( + patch("hub.utils.download_registry.get_models_registry", lambda: registry), + patch( + "core.inference.llama_cpp.cached_gguf_for_load", + return_value = "/cached/model.gguf", + ) as cached_probe, + ): + assert _hub_download_blocks_gguf_load(REPO, VARIANT) is False + + cached_probe.assert_called_once_with( + REPO, + VARIANT, + require_mmproj = False, + verify_sizes = True, + hf_token = None, + ) + + def test_cancelled_request_keeps_marker_until_load_thread_finishes(self): + from core.inference.llama_cpp import _with_gguf_load_marker + + started = threading.Event() + release = threading.Event() + finished = threading.Event() + + class FakeBackend: + @_with_gguf_load_marker + def load_model(self, *, hf_repo): + started.set() + release.wait(timeout = 2) + finished.set() + return True + + async def scenario(): + with patch( + "core.inference.llama_cpp._hub_download_blocks_gguf_load", + return_value = False, + ): + task = asyncio.create_task( + asyncio.to_thread(FakeBackend().load_model, hf_repo = REPO) + ) + assert await asyncio.to_thread(started.wait, 1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert hf_gguf_load_in_flight(REPO) + + release.set() + assert await asyncio.to_thread(finished.wait, 1) + for _ in range(100): + if not hf_gguf_load_in_flight(REPO): + break + await asyncio.sleep(0.001) + assert not hf_gguf_load_in_flight(REPO) + + asyncio.run(scenario()) + + def test_load_marker_precedes_hub_guard_and_unload(self): + source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text() + gguf_branch = source[source.index("if config.is_gguf:") :] + + assert ( + gguf_branch.index("enter_context(gguf_load_in_flight") + < gguf_branch.index("if request.llama_extra_args is None") + < gguf_branch.index("_hub_download_blocks_gguf_load") + < gguf_branch.index("unsloth_backend.unload_model") + ) + llama_source = ( + Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py" + ).read_text() + assert "@_with_gguf_load_marker\n def load_model(" in llama_source diff --git a/studio/backend/tests/test_mtp_drafter_companion.py b/studio/backend/tests/test_mtp_drafter_companion.py index 86e528ae678..02230632b65 100644 --- a/studio/backend/tests/test_mtp_drafter_companion.py +++ b/studio/backend/tests/test_mtp_drafter_companion.py @@ -328,6 +328,7 @@ def _fake_companion( pick, label, cancel_event = None, + near_path = None, ): captured["pick"] = pick return None @@ -428,6 +429,32 @@ def test_download_mtp_reuse_follows_snapshot_order_offline(tmp_path, monkeypatch assert got is not None and Path(got).parent.parent.name == "newest" +def test_download_mtp_prefers_main_snapshot_offline(tmp_path, monkeypatch): + import utils.models.model_config as mc + from core.inference.llama_cpp import LlamaCppBackend + + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + snapshots = tmp_path / "models--unsloth--gemma" / "snapshots" + old = snapshots / "old" + new = snapshots / "new" + old.mkdir(parents = True) + new.mkdir(parents = True) + main = old / "gemma-UD-Q4_K_XL.gguf" + old_drafter = old / "mtp-gemma.gguf" + new_drafter = new / "mtp-gemma.gguf" + main.write_bytes(b"main") + old_drafter.write_bytes(b"old") + new_drafter.write_bytes(b"new") + monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda _repo: [new, old]) + + got = LlamaCppBackend()._download_mtp( + hf_repo = "unsloth/gemma-GGUF", + near_path = str(main), + ) + + assert got == str(old_drafter) + + def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch): # Online, do not reuse a cached copy: go to the download path so a changed # drafter is refetched (hf_hub_download checks the current revision). @@ -447,6 +474,7 @@ def _fake_companion( pick, label, cancel_event = None, + near_path = None, ): reached["hit"] = True return None diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index e24e2ca4511..0b9a1e704f3 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -244,9 +244,7 @@ def fake_download( def test_download_reuses_older_snapshot_when_current_ref_snapshot_is_partial( self, monkeypatch, hf_cache ): - # Cross-snapshot reuse is an offline-resilience path: online, hf_hub_download - # resumes the partial current-ref download and revalidates the revision instead - # of serving an older snapshot's same-name blob. + # Keep coverage for offline reuse; online reuse is tested separately. monkeypatch.setenv("HF_HUB_OFFLINE", "1") backend = LlamaCppBackend() repo = "unsloth/vision-GGUF" @@ -292,8 +290,7 @@ def fail_download(*_args, **_kwargs): def test_download_reuses_cached_gguf_when_lowercase_partial_cache_shadows_it( self, monkeypatch, hf_cache ): - # Case-variant cross-dir reuse is offline-only; online the canonical repo id - # resolves up front and hf_hub_download fetches the current revision. + # Keep coverage for case-insensitive offline cache lookup. monkeypatch.setenv("HF_HUB_OFFLINE", "1") backend = LlamaCppBackend() canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" @@ -348,45 +345,26 @@ def fail_download(*_args, **_kwargs): assert out == str(snap / gguf_file) assert seen_repos - def test_download_online_does_not_reuse_old_snapshot(self, monkeypatch, hf_cache): - # Online, an older same-name snapshot must not be served (it may be a stale - # revision); hf_hub_download is called so the current revision is fetched and - # its etag revalidated. + def test_download_online_reuses_complete_cached_snapshot(self, monkeypatch, hf_cache): + # Loads reuse complete cached models across repo revisions. monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) backend = LlamaCppBackend() repo = "unsloth/vision-GGUF" - _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) - downloaded: list[str] = [] - - def fake_get_paths_info( - _repo_id, - paths, - token = None, - ): - return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + snap = _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) - def fake_download( - repo_id, - filename, - token = None, - **kwargs, - ): - downloaded.append(filename) - return f"/fresh/{filename}" + def fail_download(*_args, **_kwargs): + raise AssertionError("must reuse the cached GGUF instead of downloading") with ( patch( "huggingface_hub.list_repo_files", lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"], ), - patch("huggingface_hub.get_paths_info", fake_get_paths_info), - patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), - patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), ): out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") - assert downloaded == ["model-UD-Q4_K_XL.gguf"] - assert out == "/fresh/model-UD-Q4_K_XL.gguf" + assert out == str(snap / "model-UD-Q4_K_XL.gguf") def test_download_reuses_older_snapshot_when_offline_env_is_true(self, monkeypatch, hf_cache): # HF_HUB_OFFLINE accepts truthy spellings beyond "1" (true/yes/on); the offline