ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: don't re-download updated GGUFs on load (#7209)

Browse source
This commit is contained in:
oobabooga 2026-07-17 19:05:46 -03:00 • committed by GitHub
commit 2139200b3f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 1264 additions and 166 deletions
Download patch file Download diff file Expand all files Collapse all files

515
studio/backend/core/inference/llama_cpp.py
View file

@ -9,6 +9,7 @@ OpenAI-compatible /v1/chat/completions endpoint.
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 @@ class LlamaCppBackend:
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 @@ class LlamaCppBackend:
# 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 @@ class LlamaCppBackend:
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 @@ class LlamaCppBackend:
# 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 @@ class LlamaCppBackend:
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 @@ class LlamaCppBackend:
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 @@ class LlamaCppBackend:
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 @@ class LlamaCppBackend:
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 @@ class LlamaCppBackend:
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 @@ class LlamaCppBackend:
*,
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 @@ class LlamaCppBackend:
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 @@ class LlamaCppBackend:
)
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 @@ class LlamaCppBackend:
)
self._stdout_thread.start()
@_with_gguf_load_marker
def load_model(
self,
*,
@ -5581,6 +5844,7 @@ class LlamaCppBackend:
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 @@ class LlamaCppBackend:
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():

30
studio/backend/hub/services/models/downloads.py
View file

@ -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.

27
studio/backend/hub/utils/download_registry.py
View file

@ -47,7 +47,7 @@ import time
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 @@ class DownloadRegistry:
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 @@ class DownloadRegistry:
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 @@ class DownloadRegistry:
)
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,

50
studio/backend/routes/inference.py
View file

@ -20,6 +20,7 @@ from loggers import get_logger
import asyncio
import threading
import weakref
from contextlib import ExitStack
import re as _re
@ -998,6 +999,7 @@ try:
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 @@ except ImportError:
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 _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
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(

740
studio/backend/tests/test_gguf_load_cache_reuse.py Normal file
View file

@ -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--<repo>/snapshots/<sha>/<rel>`` 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

28
studio/backend/tests/test_mtp_drafter_companion.py
View file

@ -328,6 +328,7 @@ def test_download_mtp_prefers_root_over_new_scheme_copies(monkeypatch):
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 test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch):
pick,
label,
cancel_event = None,
near_path = None,
):
reached["hit"] = True
return None

40
studio/backend/tests/test_offline_gguf_cache_fallback.py
View file

@ -244,9 +244,7 @@ class TestGgufVariantFileResolution:
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 @@ class TestGgufVariantFileResolution:
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 @@ class TestGgufVariantFileResolution:
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] = []
snap = _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40)
def fake_get_paths_info(
_repo_id,
paths,
token = None,
):
return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p]
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