diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index b4f8fe1ca6..b07bd33076 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -377,6 +377,19 @@ def _probe_dns_dead(host: str = "huggingface.co", timeout: float = 2.0) -> bool: return True if result[0] is None else result[0] +def _hf_env_offline() -> bool: + """True when an HF offline env var is set to any truthy value (1/true/yes/on). + + Mirrors utils.models.model_config._env_offline so a user-set HF_HUB_OFFLINE=true + (not just "1") still routes through the local-cache reuse path below. + """ + try: + from utils.models.model_config import _env_offline + return _env_offline() + except Exception: + return os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in {"1", "true", "yes", "on"} + + @contextlib.contextmanager def _hf_offline_if_dns_dead(): """Set HF_HUB_OFFLINE for this block only when DNS to huggingface.co fails; @@ -837,6 +850,112 @@ def _gguf_snapshot_files(snapshot: Path) -> list[str]: ] +def _cached_hf_snapshot_file( + repo_id: str, + filename: str, + *, + expected_size: Optional[int] = None, +) -> Optional[str]: + """Return a cached snapshot file even when HF's current-ref probe misses it.""" + if not filename: + return None + parts = [part for part in filename.replace("\\", "/").split("/") if part] + if not parts or any(part in (".", "..") for part in parts): + return None + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + candidate = snap.joinpath(*parts) + if not candidate.is_file(): + continue + if expected_size: + try: + if candidate.stat().st_size < expected_size: + continue + except OSError: + continue + return str(candidate) + except Exception as e: + logger.debug("Snapshot cache lookup failed for %s/%s: %s", repo_id, filename, e) + return None + + +def _snapshot_has_all_shards( + main_path: str, main_filename: str, shards: Iterable[str], expected_sizes: dict[str, int] +) -> bool: + """True when every shard sits beside ``main_path`` in the same cache snapshot. + + llama.cpp loads a split GGUF by resolving its siblings from the main shard's + directory, so a cached main shard is only safe to reuse when the rest of the + set is co-located; otherwise the caller must fetch the whole set together. + """ + root = Path(main_path) + for _ in [part for part in main_filename.replace("\\", "/").split("/") if part]: + root = root.parent + for shard in shards: + parts = [part for part in shard.replace("\\", "/").split("/") if part] + if not parts or any(part in (".", "..") for part in parts): + return False + sibling = root.joinpath(*parts) + try: + if not sibling.is_file(): + return False + expected = expected_sizes.get(shard) + if expected and sibling.stat().st_size < expected: + return False + except OSError: + return False + return True + + +def _resolve_repo_id_casing(hf_repo: str) -> str: + """Map a requested repo id to its cached canonical casing, or return it unchanged. + + A case-variant request (for example a lowercased id) resolves to the + canonical-cased cache directory so the main GGUF and its companions + (mmproj / MTP drafter) all read the same cache entry. Returns ``hf_repo`` + unchanged when resolution is unavailable or errors. + """ + try: + from utils.paths import resolve_cached_repo_id_case + return resolve_cached_repo_id_case(hf_repo) + except Exception: + return hf_repo + + +def _cached_colocated_split_main( + repo_id: str, main_filename: str, shards: Iterable[str], expected_sizes: dict[str, int] +) -> Optional[str]: + """Main-shard path from a cache snapshot that also holds every sibling shard. + + A newer snapshot may hold only the first shard while an older snapshot has the + complete split set. ``_cached_hf_snapshot_file`` would return that newer partial + main and the co-location check would then force a refetch, so scan snapshots for + one where the whole set is present and return that main path instead. None when + no snapshot holds the full set. + """ + main_parts = [part for part in main_filename.replace("\\", "/").split("/") if part] + if not main_parts or any(part in (".", "..") for part in main_parts): + return None + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + main_path = snap.joinpath(*main_parts) + if not main_path.is_file(): + continue + expected_main = expected_sizes.get(main_filename) + try: + if expected_main and main_path.stat().st_size < expected_main: + continue + except OSError: + continue + if _snapshot_has_all_shards(str(main_path), main_filename, shards, expected_sizes): + return str(main_path) + except Exception as e: + logger.debug("Co-located split snapshot lookup failed for %s: %s", repo_id, e) + return None + + def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]: m = _SHARD_FULL_RE.match(first_shard) if not m: @@ -3992,6 +4111,15 @@ class LlamaCppBackend: "Install it with: pip install huggingface_hub" ) + resolved_hf_repo = _resolve_repo_id_casing(hf_repo) + if resolved_hf_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + resolved_hf_repo, + hf_repo, + ) + hf_repo = resolved_hf_repo + # Resolve the filename from the variant gguf_filename = None gguf_extra_shards: list[str] = [] @@ -4037,10 +4165,12 @@ class LlamaCppBackend: # 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 @@ -4049,7 +4179,26 @@ class LlamaCppBackend: # cold whenever free disk is below the full weight footprint, # even though nothing needs downloading. already_cached_bytes = 0 - if not force: + # 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. + 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 + if not force and not split_needs_refetch: for p in path_infos: if not p.size: continue @@ -4057,6 +4206,15 @@ class LlamaCppBackend: cached_path = try_to_load_from_cache(hf_repo, p.path) except Exception: cached_path = None + if ( + not (isinstance(cached_path, str) and os.path.exists(cached_path)) + and offline + ): + cached_path = _cached_hf_snapshot_file( + hf_repo, + p.path, + expected_size = p.size, + ) if isinstance(cached_path, str) and os.path.exists(cached_path): try: on_disk = os.path.getsize(cached_path) @@ -4119,6 +4277,13 @@ class LlamaCppBackend: ) 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 else: raise RuntimeError( f"Not enough disk space to download any variant. " @@ -4138,25 +4303,45 @@ class LlamaCppBackend: raise RuntimeError("Cancelled") dl_start = time.monotonic() # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached. - 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( + 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( hf_repo, - shard, + 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, + shard, + hf_token, + cancel_event = cancel_event, + force_download = force, + ) except Exception as e: if isinstance(e, RuntimeError) and "Cancelled" in str(e): raise @@ -4234,6 +4419,17 @@ class LlamaCppBackend: if target is None or cancel_event.is_set(): return None + # Offline, resolve the companion straight from the cache snapshot that + # holds it. resolve_cached_repo_id_case can return a partial lower-case + # spelling when any dir exists under the requested casing, so calling + # hf_hub_download with hf_repo would miss the canonical file and silently + # drop the companion. _cached_hf_snapshot_file scans every case variant. + if _hf_env_offline(): + cached = _cached_hf_snapshot_file(hf_repo, target) + if cached: + logger.info("Resolved %s from local HF cache: %s", label, cached) + return cached + try: logger.info(f"Downloading {label}: {hf_repo}/{target}") # Same policy; companions are best-effort (caller below swallows failures to None). @@ -5012,6 +5208,19 @@ class LlamaCppBackend: # dead; cleanup runs even on exception so a transient hiccup # can't quarantine future loads. if hf_repo: + # Resolve the requested repo id to its cached canonical casing once, + # up front, so the main GGUF and its companions (mmproj / MTP drafter) + # all resolve from the same cache entry. Otherwise a case-variant + # request resolves the main file from the canonical cache dir while the + # companions keep the requested casing and miss the cached files. + _resolved_repo = _resolve_repo_id_casing(hf_repo) + if _resolved_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + _resolved_repo, + hf_repo, + ) + hf_repo = _resolved_repo with _hf_offline_if_dns_dead(): model_path = self._download_gguf( hf_repo = hf_repo, diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index 4499881c4d..e24e2ca451 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -79,9 +79,11 @@ from huggingface_hub import constants as hf_constants from core.inference.llama_cpp import ( LlamaCppBackend, + _cached_colocated_split_main, _gguf_files_for_variant, _hf_offline_if_dns_dead, _probe_dns_dead, + _resolve_repo_id_casing, ) from utils.models.model_config import ( _detect_gguf_from_hf_cache, @@ -217,7 +219,7 @@ class TestGgufVariantFileResolution: downloaded.append(filename) return f"/fake/{repo_id}/{filename}" - monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) with ( patch( "huggingface_hub.list_repo_files", @@ -239,6 +241,214 @@ class TestGgufVariantFileResolution: assert downloaded == ["tinyllamas/stories260K.gguf"] assert out == "/fake/ggml-org/models/tinyllamas/stories260K.gguf" + 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. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + old = _build_cache( + hf_cache, + repo, + {"model-UD-Q4_K_XL.gguf": 4}, + snapshot_sha = "a" * 40, + ) + _build_cache( + hf_cache, + repo, + {"mtp-model.gguf": 1}, + snapshot_sha = "b" * 40, + ) + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = path, size = 4) for path in paths if path] + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch( + "huggingface_hub.list_repo_files", + lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf", "mtp-model.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", fail_download), + ): + out = backend._download_gguf( + hf_repo = repo, + hf_variant = "UD-Q4_K_XL", + ) + + assert out == str(old / "model-UD-Q4_K_XL.gguf") + + 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. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" + requested_repo = "unsloth/gemma-4-e2b-it-gguf" + gguf_file = "gemma-4-E2B-it-UD-Q4_K_XL.gguf" + snap = _build_cache( + hf_cache, + canonical_repo, + {gguf_file: 4}, + snapshot_sha = "a" * 40, + ) + lower_snap = _build_cache( + hf_cache, + requested_repo, + {"mtp-gemma-4-E2B-it.gguf": 1}, + snapshot_sha = "b" * 40, + ) + os.utime(lower_snap, (2000, 2000)) + os.utime(snap, (1000, 1000)) + seen_repos: list[str] = [] + + def fake_list_repo_files(repo_id, token = None): + seen_repos.append(repo_id) + return [gguf_file] + + def fake_get_paths_info( + repo_id, + paths, + token = None, + ): + seen_repos.append(repo_id) + return [_types.SimpleNamespace(path = path, size = 4) for path in paths if path] + + def fake_cache(repo_id, filename, *args, **kwargs): + seen_repos.append(repo_id) + return str(snap / filename) if repo_id == canonical_repo else None + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch("huggingface_hub.list_repo_files", fake_list_repo_files), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", fake_cache), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf( + hf_repo = requested_repo, + hf_variant = "UD-Q4_K_XL", + ) + + 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. + 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] + + def fake_download( + repo_id, + filename, + token = None, + **kwargs, + ): + downloaded.append(filename) + return f"/fresh/{filename}" + + 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), + ): + 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" + + 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 + # cache reuse must trigger for those too, otherwise the earlier Hub calls run + # offline while this branch still attempts hf_hub_download and the cached GGUF + # cannot load. + monkeypatch.setenv("HF_HUB_OFFLINE", "true") + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + old = _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 fail_download(*_args, **_kwargs): + raise AssertionError("should 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", fail_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") + + assert out == str(old / "model-UD-Q4_K_XL.gguf") + + def test_download_companion_resolves_from_case_variant_snapshot_offline( + self, monkeypatch, hf_cache + ): + # Offline, resolve_cached_repo_id_case can keep a partial lower-case spelling, + # so the companion (mmproj) must resolve from whichever case-variant snapshot + # actually holds it rather than being dropped by an hf_hub_download on the + # wrong casing. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" + requested_repo = "unsloth/gemma-4-e2b-it-gguf" + snap = _build_cache(hf_cache, canonical_repo, {"mmproj-F16.gguf": 4}, snapshot_sha = "a" * 40) + # A partial lower-case dir exists so casing resolution keeps the requested spelling. + _build_cache(hf_cache, requested_repo, {"config.json": 1}, snapshot_sha = "b" * 40) + + _offline_exc = type("OfflineModeIsEnabled", (Exception,), {}) + + def fake_list_repo_files(repo_id, token = None): + raise _offline_exc("offline") + + def fail_download(*_args, **_kwargs): + raise AssertionError("should resolve the companion from cache, not download") + + with ( + patch("huggingface_hub.list_repo_files", fake_list_repo_files), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_mmproj(hf_repo = requested_repo) + + assert out == str(snap / "mmproj-F16.gguf") + def test_download_includes_uppercase_split_gguf_shards(self, monkeypatch, tmp_path): backend = LlamaCppBackend() downloaded: list[str] = [] @@ -264,7 +474,7 @@ class TestGgufVariantFileResolution: downloaded.append(filename) return f"/fake/{repo_id}/{filename}" - monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) with ( patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), patch("huggingface_hub.get_paths_info", fake_get_paths_info), @@ -279,6 +489,48 @@ class TestGgufVariantFileResolution: assert downloaded == files assert out == "/fake/org/repo/model-Q4_K_M-00001-of-00002.GGUF" + def test_download_refetches_split_gguf_when_shards_span_snapshots(self, monkeypatch, hf_cache): + # The cached main shard lives in an older snapshot; its sibling shard is only + # in a newer, separate snapshot. Reusing the main shard alone would leave + # llama.cpp unable to resolve the sibling, so the whole set must be re-fetched + # together (co-located) rather than served split across snapshot dirs. + backend = LlamaCppBackend() + repo = "org/split" + files = [ + "model-Q4_K_M-00001-of-00002.gguf", + "model-Q4_K_M-00002-of-00002.gguf", + ] + _build_cache(hf_cache, repo, {files[0]: 4}, snapshot_sha = "a" * 40) + _build_cache(hf_cache, repo, {files[1]: 4}, snapshot_sha = "b" * 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] + + 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: files), + 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 = "Q4_K_M") + + assert downloaded == files + assert out == f"/fake/{repo}/{files[0]}" + def _siblings(items: dict[str, int]): """Mock ``hf_model_info(...).siblings`` payload.""" @@ -315,6 +567,21 @@ class TestIterHfCacheSnapshots: out = list(_iter_hf_cache_snapshots("unsloth/multi")) assert [p.name for p in out] == ["b" * 40, "a" * 40] + def test_skips_snapshot_when_mtime_is_unavailable(self, hf_cache, monkeypatch): + stale = _build_cache(hf_cache, "unsloth/multi", {"x.gguf": 1}, snapshot_sha = "a" * 40) + good = _build_cache(hf_cache, "unsloth/multi", {"y.gguf": 1}, snapshot_sha = "b" * 40) + original_stat = Path.stat + + def flaky_stat(self, *args, **kwargs): + if self == stale: + raise FileNotFoundError(str(self)) + return original_stat(self, *args, **kwargs) + + monkeypatch.setattr(Path, "stat", flaky_stat) + + out = list(_iter_hf_cache_snapshots("unsloth/multi")) + assert out == [good] + def test_repo_id_match_is_case_insensitive(self, hf_cache): _build_cache(hf_cache, "unsloth/Foo-GGUF", {"Foo-Q4_K_M.gguf": 1}) # Lookup with different org/name casing still resolves @@ -347,6 +614,87 @@ class TestListGgufVariantsFromCache: assert _list_gguf_variants_from_hf_cache("unsloth/absent") is None +class TestCachedColocatedSplitMain: + def test_prefers_older_complete_snapshot_over_newer_partial(self, hf_cache): + # Newer snapshot has only shard 1; older snapshot has the complete set. The + # complete older snapshot must win so the split GGUF can load co-located. + shard1 = "m-00001-of-00002.gguf" + shard2 = "m-00002-of-00002.gguf" + old = _build_cache( + hf_cache, "unsloth/split-GGUF", {shard1: 100, shard2: 100}, snapshot_sha = "a" * 40 + ) + new = _build_cache(hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "b" * 40) + os.utime(old, (1000, 1000)) + os.utime(new, (2000, 2000)) + + main = _cached_colocated_split_main("unsloth/split-GGUF", shard1, [shard2], {}) + assert main is not None + assert main.startswith(str(old)) + + def test_returns_none_when_shards_span_snapshots(self, hf_cache): + shard1 = "m-00001-of-00002.gguf" + shard2 = "m-00002-of-00002.gguf" + a = _build_cache(hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "a" * 40) + b = _build_cache(hf_cache, "unsloth/split-GGUF", {shard2: 100}, snapshot_sha = "b" * 40) + os.utime(a, (1000, 1000)) + os.utime(b, (2000, 2000)) + + assert _cached_colocated_split_main("unsloth/split-GGUF", shard1, [shard2], {}) is None + + +class TestResolveRepoIdCasing: + def test_maps_to_canonical_casing(self, monkeypatch): + monkeypatch.setattr( + "utils.paths.resolve_cached_repo_id_case", + lambda repo: "unsloth/Gemma-4-GGUF" if repo.lower() == "unsloth/gemma-4-gguf" else repo, + ) + # A companion download passed the resolved id reads the same cache entry + # as the main GGUF instead of missing it under the requested casing. + assert _resolve_repo_id_casing("unsloth/gemma-4-gguf") == "unsloth/Gemma-4-GGUF" + + def test_passthrough_on_resolver_error(self, monkeypatch): + def boom(_repo): + raise RuntimeError("resolver unavailable") + + monkeypatch.setattr("utils.paths.resolve_cached_repo_id_case", boom) + assert _resolve_repo_id_casing("unsloth/gemma-4-gguf") == "unsloth/gemma-4-gguf" + + def test_companion_only_newer_snapshot_does_not_shadow_real_variants(self, hf_cache): + # A newer snapshot holds only a vision projector fetched on demand, + # while the quant files live in an older snapshot. The newer snapshot + # must not shadow the real variants; the vision flag carries over. + old = _build_cache( + hf_cache, + "unsloth/vision-GGUF", + {"vision-Q4_K_M.gguf": 100}, + snapshot_sha = "a" * 40, + ) + new = _build_cache( + hf_cache, + "unsloth/vision-GGUF", + {"mmproj-vision-F16.gguf": 10}, + snapshot_sha = "b" * 40, + ) + os.utime(old, (1000, 1000)) + os.utime(new, (2000, 2000)) + + out = _list_gguf_variants_from_hf_cache("unsloth/vision-GGUF") + assert out is not None + variants, has_vision = out + assert [v.quant for v in variants] == ["Q4_K_M"] + assert has_vision is True + + def test_companion_only_cache_returns_empty_variants_with_vision(self, hf_cache): + # Only a vision projector is cached anywhere: report the vision flag + # with an empty variant list rather than None. + _build_cache(hf_cache, "unsloth/vision-GGUF", {"mmproj-vision-F16.gguf": 10}) + out = _list_gguf_variants_from_hf_cache("unsloth/vision-GGUF") + assert out is not None + variants, has_vision = out + assert variants == [] + assert has_vision is True + + class TestListGgufVariantsOffline: def test_offline_env_short_circuits_api(self, hf_cache, clean_offline_env, monkeypatch): _build_cache(hf_cache, "unsloth/a", {"a-UD-Q4_K_XL.gguf": 1}) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 5d8458e5f0..281ca24281 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1617,36 +1617,60 @@ def _iter_hf_cache_snapshots(repo_id: str): cache_dir = Path(hf_constants.HF_HUB_CACHE) target = f"models--{repo_id.replace('/', '--')}".lower() - repo_dir: Optional[Path] = None + repo_dirs: list[Path] = [] try: if not cache_dir.is_dir(): return for entry in cache_dir.iterdir(): if entry.is_dir() and entry.name.lower() == target: - repo_dir = entry - break + repo_dirs.append(entry) except OSError: return - if repo_dir is None: + if not repo_dirs: return - snapshots = repo_dir / "snapshots" - try: - if not snapshots.is_dir(): - return - snap_dirs = [s for s in snapshots.iterdir() if s.is_dir()] - except OSError: + snap_dirs: list[Path] = [] + for repo_dir in repo_dirs: + snapshots = repo_dir / "snapshots" + try: + if snapshots.is_dir(): + for snap_dir in snapshots.iterdir(): + try: + if snap_dir.is_dir(): + snap_dirs.append(snap_dir) + except OSError: + continue + except OSError: + continue + if not snap_dirs: return - snap_dirs.sort(key = lambda s: s.stat().st_mtime, reverse = True) - yield from snap_dirs + snap_dirs_with_mtime = [] + for snap_dir in snap_dirs: + try: + snap_dirs_with_mtime.append((snap_dir.stat().st_mtime, snap_dir)) + except OSError: + continue + snap_dirs_with_mtime.sort(key = lambda item: item[0], reverse = True) + yield from (snap_dir for _, snap_dir in snap_dirs_with_mtime) def _list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]: - """Variants from the local HF cache snapshot, or None if not cached.""" + """Variants from the local HF cache snapshot, or None if not cached. + + A newer snapshot can hold only a companion file (for example a vision + projector fetched on demand) while the quant files live in an older + snapshot. Returning the first snapshot that merely reports a vision flag + would shadow those real variants, so keep scanning older snapshots for + actual variants and carry the vision flag across snapshots. + """ + any_vision = False for snap in _iter_hf_cache_snapshots(repo_id): variants, has_vision = list_local_gguf_variants(str(snap)) - if variants or has_vision: - return variants, has_vision + any_vision = any_vision or has_vision + if variants: + return variants, any_vision + if any_vision: + return [], True return None diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 01014de4d3..477a47cc3d 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -597,6 +597,59 @@ def _loaded_models(base: str, key: str) -> list: return _http_json("GET", f"{base}/v1/models", key, error = "Couldn't list models").get("data", []) +_HF_REPO_ID_SEGMENT_RE = re.compile(r"^[A-Za-z0-9._-]+$") + + +def _is_hub_model_id(value: object) -> bool: + if not isinstance(value, str): + return False + text = value.strip() + if "\\" in text: + return False + if text.startswith(("/", "./", "../", "~")): + return False + if len(text) >= 2 and text[1] == ":" and text[0].isalpha(): + return False + # A hub id is exactly "namespace/name" over a restricted charset. Anything with + # extra path segments (e.g. a server-side relative path such as + # models/Llama/Foo.gguf on a remote Studio) is not a hub id and must not be + # casefold-matched against a differently cased path on a case-sensitive + # filesystem. This is host independent, unlike the existence probe below which + # cannot see a path that only exists on the server. + parts = text.split("/") + if len(parts) != 2: + return False + if any(part in ("", ".", "..") or not _HF_REPO_ID_SEGMENT_RE.match(part) for part in parts): + return False + try: + if Path(os.path.expanduser(text)).exists(): + return False + except OSError: + return False + return True + + +def _model_id_matches( + actual: object, + requested: object, + *, + allow_casefold: bool = True, +) -> bool: + if actual == requested: + return True + # Case-insensitive matching is only safe when the local existence probe in + # _is_hub_model_id is authoritative, i.e. against a loopback Studio on this host. + # Against a remote Studio a two-segment string is indistinguishable from a + # server-side relative path (e.g. Models/Foo vs models/foo), so casefolding it + # could attach to the wrong model on a case-sensitive server; defer to an exact + # match there and let the load endpoint resolve the requested path. + if not allow_casefold: + return False + if not (_is_hub_model_id(actual) and _is_hub_model_id(requested)): + return False + return str(actual).casefold() == str(requested).casefold() + + def _resolve_model( base: str, key: str, @@ -604,6 +657,9 @@ def _resolve_model( load: LoadOptions = LoadOptions(), ) -> dict: models = _loaded_models(base, key) + # Only casefold-match ids against a loopback Studio, where _is_hub_model_id's + # local existence probe can actually reject a server-side path; see the note there. + allow_casefold = is_loopback_url(base) # /v1/models reports the model id but not the active GGUF variant or runtime load # settings, so an id match alone can hide the wrong quant (Q8_0 serving while the # user asked for UD-Q4_K_XL). When the user passed any explicit load knob, defer to @@ -613,10 +669,21 @@ def _resolve_model( load_has_overrides = bool( load.gguf_variant or load.max_seq_length or not load.load_in_4bit or load.tensor_parallel ) + # /v1/models also lists cached-but-unloaded catalog entries (loaded == False); + # matching one would skip /api/inference/load and leave the agent pointed at a + # model that is not resident, so only attach to an entry that is actually loaded. match = ( None if requested and load_has_overrides - else next((m for m in models if m["id"] == requested), None) + else next( + ( + m + for m in models + if _model_id_matches(m.get("id"), requested, allow_casefold = allow_casefold) + and m.get("loaded") is not False + ), + None, + ) ) if requested and match is None: typer.echo( @@ -651,7 +718,16 @@ def _resolve_model( if isinstance(loaded, dict): wanted |= {loaded.get("model"), loaded.get("display_name")} - {None} models = _loaded_models(base, key) - match = next((m for m in models if m["id"] in wanted), None) + match = next( + ( + m + for m in models + if any( + _model_id_matches(m.get("id"), w, allow_casefold = allow_casefold) for w in wanted + ) + ), + None, + ) if match is not None: return match if requested: diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index bd964d5e54..ee5b442c27 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -457,6 +457,189 @@ def test_connect_codex_no_launch(fake_studio, tmp_path): assert (home / "unsloth_api.config.toml").exists() +def test_connect_codex_matches_requested_model_case_insensitively(fake_studio, tmp_path): + result = CliRunner().invoke( + start.start_app, + [ + "codex", + "--no-launch", + "--model", + "unsloth/gemma-4-26b-a4b-it-gguf", + ], + ) + assert result.exit_code == 0, result.output + home = tmp_path / "agents" / "codex" + profile = _parse_toml((home / "unsloth_api.config.toml").read_text()) + assert profile["model"] == MODEL["id"] + + +def test_resolve_model_matches_loaded_canonical_case_after_load(monkeypatch): + calls = [] + state = {"loaded": False} + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url, payload)) + if url.endswith("/v1/models"): + return { + "data": [ + { + "id": "unsloth/gemma-4-E2B-it-GGUF" if state["loaded"] else "other/model", + "context_length": 131072, + } + ] + } + if url.endswith("/api/inference/load"): + state["loaded"] = True + return {"model": "unsloth/gemma-4-E2B-it-GGUF"} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "_http_json", http_json) + + entry = start._resolve_model( + BASE, + "sk-test", + "unsloth/gemma-4-e2b-it-gguf", + start.LoadOptions(gguf_variant = "UD-Q4_K_XL"), + ) + + assert entry["id"] == "unsloth/gemma-4-E2B-it-GGUF" + assert any(c[1].endswith("/api/inference/load") for c in calls) + + +def test_resolve_model_loads_when_catalog_hit_is_not_loaded(monkeypatch): + # A cached-but-unloaded catalog entry (loaded == False) that only case-differs must + # not be treated as ready; the load endpoint must still be called so the requested + # model becomes resident instead of the agent preflighting a different backend. + calls = [] + state = {"loaded": False} + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url)) + if url.endswith("/v1/models"): + return { + "data": [ + { + "id": "unsloth/Gemma-4-GGUF", + "loaded": state["loaded"], + "context_length": 131072, + } + ] + } + if url.endswith("/api/inference/load"): + state["loaded"] = True + return {"model": "unsloth/Gemma-4-GGUF"} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "_http_json", http_json) + + entry = start._resolve_model(BASE, "sk-test", "unsloth/gemma-4-gguf") + + assert entry["id"] == "unsloth/Gemma-4-GGUF" + assert any(u.endswith("/api/inference/load") for _, u in calls) + + +def test_resolve_model_attaches_to_loaded_catalog_hit_without_reload(monkeypatch): + # The mirror case: a loaded entry (loaded == True) that case-matches attaches with + # no /api/inference/load call. + calls = [] + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url)) + if url.endswith("/v1/models"): + return { + "data": [{"id": "unsloth/Gemma-4-GGUF", "loaded": True, "context_length": 131072}] + } + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "_http_json", http_json) + + entry = start._resolve_model(BASE, "sk-test", "unsloth/gemma-4-gguf") + + assert entry["id"] == "unsloth/Gemma-4-GGUF" + assert not any(u.endswith("/api/inference/load") for _, u in calls) + + +def test_resolve_model_remote_studio_does_not_casefold_attach(monkeypatch): + # Against a remote Studio the local existence probe cannot see server-side paths, + # so a case-variant loaded id must NOT attach without a load: it could be a distinct + # server-side path on a case-sensitive host. The load endpoint resolves the request. + calls = [] + state = {"loaded": False} + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url)) + if url.endswith("/v1/models"): + return { + "data": [{"id": "unsloth/Gemma-4-GGUF", "loaded": True, "context_length": 131072}] + } + if url.endswith("/api/inference/load"): + state["loaded"] = True + return {"model": "unsloth/Gemma-4-GGUF"} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "_http_json", http_json) + + entry = start._resolve_model("http://10.0.0.5:8888", "sk-test", "unsloth/gemma-4-gguf") + + # The load endpoint was consulted (no casefold shortcut), and we still attach to the + # server's canonical id it reports back. + assert entry["id"] == "unsloth/Gemma-4-GGUF" + assert any(u.endswith("/api/inference/load") for _, u in calls) + + +def test_model_id_matching_does_not_casefold_local_paths(tmp_path): + existing_local = tmp_path / "Org" / "Foo" + existing_local.mkdir(parents = True) + + assert start._model_id_matches("Org/Foo", "org/foo") + assert not start._model_id_matches(str(existing_local), str(existing_local).lower()) + assert not start._model_id_matches("./Models/Foo", "./models/foo") + assert not start._model_id_matches(r".\Models\Foo", r".\models\foo") + # A server-side relative path (extra path segments) is not a hub id even when it + # does not exist on the CLI host, so it must not casefold-match a differently + # cased path on a case-sensitive server filesystem. + assert not start._is_hub_model_id("models/Llama/Foo.gguf") + assert not start._model_id_matches("models/Llama/Foo.gguf", "models/llama/foo.gguf") + # A genuine two-segment hub id still matches case-insensitively. + assert start._is_hub_model_id("unsloth/Gemma-3-4b-it-GGUF") + assert start._model_id_matches("unsloth/Gemma-3-4b-it-GGUF", "unsloth/gemma-3-4b-it-gguf") + # Casefolding is gated to loopback studios (allow_casefold). With it disabled (a + # remote studio, where a two-segment string could be a server-side path), even a + # genuine hub-id case variant must not match, so the load endpoint resolves it. + assert not start._model_id_matches( + "unsloth/Gemma-3-4b-it-GGUF", "unsloth/gemma-3-4b-it-gguf", allow_casefold = False + ) + assert start._model_id_matches("unsloth/Foo", "unsloth/Foo", allow_casefold = False) + + def test_connect_codex_launch_uses_ephemeral_home(fake_studio, monkeypatch): # Launch mode writes config to a throwaway temp CODEX_HOME and removes it after # the agent exits; the user's real ~/.codex is never the target.