diff --git a/studio/backend/hub/utils/local_snapshot.py b/studio/backend/hub/utils/local_snapshot.py new file mode 100644 index 0000000000..d63334ff12 --- /dev/null +++ b/studio/backend/hub/utils/local_snapshot.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Local-only snapshot resolution for background model loads.""" + +from typing import Optional + + +def resolve_local_snapshot_path( + repo_id: str, + hf_token: Optional[str] = None, + cache_dir: Optional[str] = None, +) -> Optional[str]: + """Resolve a Hub repo id to its snapshot directory in the local HF cache + without any network access; None when the repo is not cached. + + ``snapshot_download(local_files_only = True)`` reads only the on-disk + refs/snapshots, so a cache populated outside Studio that is missing files + still resolves to its snapshot directory; the subsequent weight load on + that local path then fails instead of downloading the gaps, which is the + fail-closed behavior background loads need. + """ + try: + from huggingface_hub import snapshot_download + + return snapshot_download( + repo_id, + local_files_only = True, + token = hf_token or None, + cache_dir = cache_dir or None, + ) + except Exception: + return None diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index add3228a28..cba1b0e7f1 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -40,6 +40,13 @@ class LoadRequest(BaseModel): gguf_variant: Optional[str] = Field( None, description = "GGUF quantization variant (e.g. 'Q4_K_M')" ) + local_files_only: bool = Field( + False, + description = ( + "Resolve a cached repo id against the local snapshot only and " + "never download missing files (background auto-loads)." + ), + ) trust_remote_code: bool = Field( False, description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.", diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 7197483841..bdba631a9f 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -4497,6 +4497,35 @@ async def _load_model_impl( detail = f"Invalid model identifier: {model_log_label}", ) + # A background auto-load promises on-device-only resolution, but a + # cache populated outside Studio can pass the partial check while + # missing shards, and from_pretrained on a repo id would download the + # gaps. Rewriting the load path to the locally resolved snapshot keeps + # the registry identity (config.identifier) intact while forcing the + # weight load to the files actually on disk, so an incomplete cache + # fails over to the next candidate instead of fetching. + from utils.paths import is_local_path + + if ( + request.local_files_only + and not config.is_gguf + and not is_local_path(config.path) + ): + from hub.utils.local_snapshot import resolve_local_snapshot_path + + local_snapshot = await asyncio.to_thread( + resolve_local_snapshot_path, config.path, request.hf_token + ) + if local_snapshot is None: + raise HTTPException( + status_code = 409, + detail = ( + f"Model '{model_log_label}' is not available on device; " + "select it explicitly to download it." + ), + ) + config.path = local_snapshot + # Normalize gpu_ids: empty list means auto-selection, same as None effective_gpu_ids = request.gpu_ids if request.gpu_ids else None diff --git a/studio/backend/tests/test_local_snapshot_resolution.py b/studio/backend/tests/test_local_snapshot_resolution.py new file mode 100644 index 0000000000..f4e2561905 --- /dev/null +++ b/studio/backend/tests/test_local_snapshot_resolution.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Local-only snapshot resolution for background auto-loads. + +A cache populated outside Studio (no download manifest) passes the partial +check while missing shard files, and ``from_pretrained`` on a repo id would +download the gaps. Background loads therefore rewrite the load path to the +LOCALLY resolved snapshot: resolution never touches the network, an uncached +repo resolves to None (409 upstream), and an incomplete snapshot still +resolves so the weight load fails on the missing files instead of fetching +them. No GPU or network required. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +pytest.importorskip("huggingface_hub") + +from hub.utils.local_snapshot import resolve_local_snapshot_path + + +_REV = "0123456789abcdef0123456789abcdef01234567" + + +def _build_cached_repo( + cache_dir: Path, + repo_id: str, + files: dict[str, str], +) -> Path: + """Lay out a minimal HF hub cache entry the way huggingface_hub expects: + ``models--org--name/refs/main`` pointing at a snapshot directory.""" + repo_dir = cache_dir / f"models--{repo_id.replace('/', '--')}" + snapshot = repo_dir / "snapshots" / _REV + snapshot.mkdir(parents = True) + (repo_dir / "refs").mkdir() + (repo_dir / "refs" / "main").write_text(_REV) + for name, content in files.items(): + (snapshot / name).write_text(content) + return snapshot + + +def test_cached_repo_resolves_to_its_snapshot_dir(tmp_path): + snapshot = _build_cached_repo( + tmp_path, + "org/tiny-model", + {"config.json": "{}", "model.safetensors": "weights"}, + ) + resolved = resolve_local_snapshot_path( + "org/tiny-model", cache_dir = str(tmp_path) + ) + assert resolved is not None + assert Path(resolved).resolve() == snapshot.resolve() + + +def test_incomplete_snapshot_still_resolves_locally(tmp_path): + """Missing shards must not block resolution: the local path is what makes + the subsequent weight load fail closed instead of downloading.""" + snapshot = _build_cached_repo( + tmp_path, + "org/half-downloaded", + { + "config.json": "{}", + "model-00001-of-00002.safetensors": "first shard only", + }, + ) + resolved = resolve_local_snapshot_path( + "org/half-downloaded", cache_dir = str(tmp_path) + ) + assert resolved is not None + assert Path(resolved).resolve() == snapshot.resolve() + + +def test_uncached_repo_resolves_to_none(tmp_path): + assert ( + resolve_local_snapshot_path("org/never-downloaded", cache_dir = str(tmp_path)) + is None + ) + + +def test_resolution_never_uses_the_network(tmp_path, monkeypatch): + """local_files_only resolution must not open any connection even when the + repo is absent (the tempting fallback would be a Hub metadata call).""" + import socket + + def _no_network(*_args, **_kwargs): + raise AssertionError("network access attempted during local resolution") + + monkeypatch.setattr(socket.socket, "connect", _no_network) + _build_cached_repo(tmp_path, "org/offline-ok", {"config.json": "{}"}) + assert resolve_local_snapshot_path("org/offline-ok", cache_dir = str(tmp_path)) + assert ( + resolve_local_snapshot_path("org/absent", cache_dir = str(tmp_path)) is None + ) diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index f9466f895a..9136b3db5c 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1531,8 +1531,13 @@ function isAutoLoadableLocalRow( // auto-load must never start. Adapters stay interactive-only. if (row.model_format === "adapter") return false; if (isHiddenModelId(row.model_id, row.id, row.path)) return false; + // The name-based big-endian marker only applies to direct .gguf files: a + // DIRECTORY named e.g. /models/foo-be says nothing about the files inside + // it, and those are already filtered per-file by isAutoLoadableGgufVariant + // when the folder's quants are resolved. if ( row.model_format === "gguf" && + row.path.toLowerCase().endsWith(".gguf") && hasBigEndianGgufMarker(row.path, row.format_variant) ) { return false; @@ -1596,10 +1601,11 @@ function normalizeLoadTargetKey(value: string): string { const AUTO_LOAD_VARIANT_SCAN_CONCURRENCY = 4; // How long after the bounded inventory calls resolve the best-effort -// hidden-model matcher fetch may still be waited on. It usually settles -// while the inventory requests run; past this bound the static needles -// filter alone rather than letting an unbounded request stall Send. -const HIDDEN_MATCHERS_GRACE_MS = 1_000; +// prefetches (hidden-model matchers, platform snapshot) may still be +// waited on. Both use unbounded fetches of their own, so past this bound +// their client-side fallbacks apply (static needles; boot-detected +// platform) rather than letting a stalled request hang Send. +const BEST_EFFORT_PREFETCH_GRACE_MS = 1_000; // Settle window before the fallback starts consuming resolved candidates: // when scans finish quickly (the common case) the pool is complete first and @@ -1884,6 +1890,11 @@ export async function autoLoadOnDeviceModel(): Promise<{ load_in_4bit: true, is_lora: false, gguf_variant: candidate.ggufVariant, + // A background load never downloads: the backend resolves a cached + // repo id against the local snapshot only, so a cache populated + // outside Studio with missing shards fails over to the next + // candidate instead of silently fetching the gaps. + local_files_only: true, trust_remote_code: trustRemoteCode, chat_template_override: effectiveChatTemplateOverride, cache_type_kv: config.kvCacheDtype, @@ -2032,14 +2043,25 @@ export async function autoLoadOnDeviceModel(): Promise<{ const hiddenMatchersReady = ensureHiddenModelMatchers().catch( () => undefined, ); + // The platform fetch behind the picker's format gates is unbounded + // too (raw fetch, no signal); run it alongside as well. Past the + // grace, the store's boot-detected client-side platform applies. + const platformReady: Promise = fetchDeviceType().then( + () => undefined, + () => undefined, + ); const [cachedGguf, cachedModels, localList] = await Promise.all([ listCachedGguf(hfToken), listCachedModels(hfToken), listLocalModels(), ]); + const bestEffortPrefetches = Promise.all([ + hiddenMatchersReady, + platformReady, + ]); await new Promise((resolve) => { - const matcherTimer = setTimeout(resolve, HIDDEN_MATCHERS_GRACE_MS); - hiddenMatchersReady.then(() => { + const matcherTimer = setTimeout(resolve, BEST_EFFORT_PREFETCH_GRACE_MS); + bestEffortPrefetches.then(() => { clearTimeout(matcherTimer); resolve(); }); @@ -2063,10 +2085,9 @@ export async function autoLoadOnDeviceModel(): Promise<{ }; } - // Resolve the platform snapshot the picker's format gates key on. Cached - // after the app's boot fetch and never throws (falls back to client-side - // detection when the backend is not ready). - await fetchDeviceType(); + // The platform snapshot the picker's format gates key on: hydrated by the + // bounded best-effort prefetch above, else the store's boot-detected + // client-side values (the same defaults the whole UI uses pre-hydration). const platformState = usePlatformStore.getState(); const platform: AutoLoadPlatform = { chatOnly: platformState.isChatOnly(), diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index e6d3b79015..a0663044d0 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -41,6 +41,9 @@ export interface LoadModelRequest { load_in_4bit: boolean; is_lora: boolean; gguf_variant?: string | null; + /** Resolve a cached repo id against the local snapshot only and never + * download missing files (background auto-loads). */ + local_files_only?: boolean; /** Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust. */ trust_remote_code?: boolean; /** sha256 fingerprint pinning user approval of this exact custom-code version. */ diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index fd32b1ac38..d707c9d345 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -660,6 +660,10 @@ def test_autoload_filters_match_picker_policy(): # the implicit remote fetch a background auto-load must never trigger. assert 'row.model_format === "adapter"' in local_fn assert "isHiddenModelId(row.model_id, row.id, row.path)" in local_fn + # The name-based marker only applies to direct .gguf files; a directory + # named e.g. /models/foo-be says nothing about the files inside it, which + # are filtered per-file during variant resolution. + assert 'row.path.toLowerCase().endsWith(".gguf") &&' in local_fn assert "hasBigEndianGgufMarker(row.path, row.format_variant)" in local_fn cached_fn = src.split("function isAutoLoadableCachedRepo", 1)[1] cached_fn = cached_fn.split("\nconst ", 1)[0] @@ -952,8 +956,12 @@ def test_hidden_matcher_fetch_never_blocks_send_unbounded(): auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] assert "await ensureHiddenModelMatchers()" not in auto_load assert "const hiddenMatchersReady = ensureHiddenModelMatchers().catch(" in auto_load - assert "const HIDDEN_MATCHERS_GRACE_MS" in src - assert "setTimeout(resolve, HIDDEN_MATCHERS_GRACE_MS)" in auto_load + # The platform probe backing the picker's format gates is unbounded too + # (raw fetch, no signal), so it joins the same bounded prefetch wait. + assert "await fetchDeviceType();" not in auto_load + assert "const platformReady: Promise = fetchDeviceType().then(" in auto_load + assert "const BEST_EFFORT_PREFETCH_GRACE_MS" in src + assert "setTimeout(resolve, BEST_EFFORT_PREFETCH_GRACE_MS)" in auto_load assert "clearTimeout(matcherTimer)" in auto_load @@ -983,8 +991,9 @@ def test_local_rows_apply_picker_platform_gate(): assert "localRowIsGgufLike(row)" in local_fn assert "platform.isMac && localRowIsMlxNamed(row)" in local_fn auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] - # The platform snapshot hydrates through the cached, non-throwing fetch. - assert "await fetchDeviceType();" in auto_load + # The platform snapshot hydrates through the bounded best-effort prefetch + # (see test_hidden_matcher_fetch_never_blocks_send_unbounded). + assert "fetchDeviceType().then(" in auto_load # Cascade seeding of cached non-GGUF repos mirrors the picker's # chat-only exclusion; the remembered lookup above it stays unfiltered. assert "platform.chatOnly ? [] : modelRepos" in auto_load @@ -1024,3 +1033,42 @@ def test_local_variant_scans_bounded_concurrency(): auto_load = src.split("async function autoLoadOnDeviceModel", 1)[1] assert "Math.min(AUTO_LOAD_VARIANT_SCAN_CONCURRENCY, resolutionJobs.length)" in auto_load assert "await Promise.all(\n cascadeLocalRows.map(" not in auto_load + + +BACKEND = WORKDIR / "studio" / "backend" + + +def _read_backend(rel: str) -> str: + path = BACKEND / rel + assert path.exists(), f"missing backend source file: {path}" + return path.read_text() + + +def test_background_loads_resolve_local_files_only(): + """A cache populated outside Studio can pass the partial check while + missing shard files, and from_pretrained on a repo id downloads the gaps. + Background auto-loads therefore send local_files_only and the load route + rewrites the path to the locally resolved snapshot (identity intact), so + an incomplete cache fails over to the next candidate instead of + downloading on Send.""" + src = _read("features/chat/api/chat-adapter.ts") + load_fn = src.split("async function loadAutoLoadCandidate", 1)[1] + load_fn = load_fn.split("loadAttempts += 1;", 1)[1] + assert "local_files_only: true," in load_fn + types = _read("features/chat/types/api.ts") + assert "local_files_only?: boolean;" in types + + request_model = _read_backend("models/inference.py") + assert "local_files_only: bool = Field(" in request_model + + route = _read_backend("routes/inference.py") + assert "request.local_files_only" in route + assert "resolve_local_snapshot_path" in route + assert "config.path = local_snapshot" in route + # Uncached repos fail closed with a conflict, never a download. + rewrite = route.split("request.local_files_only", 1)[1] + rewrite = rewrite.split("config.path = local_snapshot", 1)[0] + assert "status_code = 409" in rewrite + + helper = _read_backend("hub/utils/local_snapshot.py") + assert "local_files_only = True" in helper