Studio autoload: local-only cached loads, bounded platform fetch, file-scoped BE marker
Background loads now send local_files_only and the load route rewrites a cached repo id to its locally resolved snapshot directory before the weight load, keeping the registry identity intact. A cache populated outside Studio passes the partial check while missing shards because a missing download manifest reads as non-partial, and from_pretrained on a repo id would download the gaps; resolution never touches the network and an uncached repo returns 409, so incomplete caches fail over to the next candidate instead of fetching on Send. The platform probe behind the picker's format gates joins the bounded best-effort prefetch (its raw fetch has no timeout), falling back to the boot-detected platform past the grace. The name-based big-endian marker now applies only to direct .gguf files: a directory named foo-be says nothing about the files inside it, which the per-file variant filter already covers.
This commit is contained in:
parent
184075a383
commit
d57ca52b83
7 changed files with 257 additions and 14 deletions
33
studio/backend/hub/utils/local_snapshot.py
Normal file
33
studio/backend/hub/utils/local_snapshot.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
102
studio/backend/tests/test_local_snapshot_resolution.py
Normal file
102
studio/backend/tests/test_local_snapshot_resolution.py
Normal file
|
|
@ -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
|
||||
)
|
||||
|
|
@ -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<void> = 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<void>((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(),
|
||||
|
|
|
|||
|
|
@ -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. */
|
||||
|
|
|
|||
|
|
@ -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<void> = 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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue