From 00566662c681ac9ef17bcfb27ca5f3f5ed642262 Mon Sep 17 00:00:00 2001 From: Unsloth Date: Mon, 27 Jul 2026 00:24:24 -0700 Subject: [PATCH] Bound scan gates, keep generation fallbacks local, and align inactive-cache load targets Inactive-cache rows now emit the newest snapshot holding config.json plus safetensors weights as their load_id, since the load consumes that path directly and a metadata-only newest revision would fail an eligible model. Ordinary offline guards no longer join a forced window (their block runs correctly under either env state), so concurrent online work sees the narrowest possible override; forced guards still share windows. The generation-time native template reload honors the load's local_files_only flag and resolved path, so a tool-calling turn on an auto-loaded model cannot download tokenizer files mid-Send. Every GGUF variant scan behind the model-kind gate now runs with an abortable 30s timeout (matching the inventory calls' own bound) so one hung request cannot gate Send forever. An HF cache snapshot registered as a custom scan folder dedupes against its cached row by expanding snapshot paths to their cache root, so shared files consume one load attempt. --- .../core/inference/chat_template_helpers.py | 9 ++++ studio/backend/core/inference/inference.py | 3 ++ studio/backend/core/inference/llama_cpp.py | 13 +++-- .../backend/core/inference/mlx_inference.py | 4 ++ .../hub/services/models/cache_inventory.py | 33 ++++++++++++ .../tests/test_offline_guard_refcount.py | 28 ++++++++-- .../src/features/chat/api/chat-adapter.ts | 52 +++++++++++++++++-- .../src/features/chat/api/chat-api.ts | 4 ++ tests/studio/test_model_picker_contracts.py | 51 +++++++++++++++++- 9 files changed, 181 insertions(+), 16 deletions(-) diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index 528c059fbc..8163caa639 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -432,6 +432,14 @@ def render_native_template( if native_tpl is None: # A LoRA adapter's native template lives on the base model, not the adapter id. template_source = model_info.get("base_model") or active_model_name + # A local-only (background) load resolved its weights to a local + # snapshot; the template reload must read the SAME files. For a + # non-LoRA model prefer the stored load path over the repo id, and + # either way pass local_files_only so a cache miss fails the fallback + # instead of downloading tokenizer files mid-generation. + local_files_only = bool(model_info.get("local_files_only", False)) + if local_files_only and not model_info.get("base_model"): + template_source = model_info.get("model_path") or template_source # Re-use the load-time trust_remote_code so a custom-code tokenizer repo can # instantiate its class (the stored flag already covers template_source). trust_remote_code = bool(model_info.get("trust_remote_code", False)) @@ -441,6 +449,7 @@ def render_native_template( template_source, token = hf_token if hf_token and hf_token.strip() else None, trust_remote_code = trust_remote_code, + local_files_only = local_files_only, ) native_tpl = nt.chat_template or False except Exception as exc: diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 6f5efd1642..170dfc6e46 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -365,6 +365,9 @@ class InferenceBackend: "trust_remote_code": trust_remote_code, "is_vision": config.is_vision, "is_lora": config.is_lora, + # Local-only loads: generation-time repo fallbacks (native + # template reload) must also resolve from cache, not the Hub. + "local_files_only": local_files_only, "is_audio": config.is_audio, "audio_type": config.audio_type, "has_audio_input": config.has_audio_input, diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 43a2bd20cf..79eff0a7aa 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -534,10 +534,15 @@ def _hf_offline_if_dns_dead(force: bool = False): owner = False # this guard created it (first in) with _OFFLINE_GUARD_LOCK: if _OFFLINE_GUARD_STATE["count"] > 0: - # Join the active override so the env survives until every guard - # exits, whichever request finishes first. - _OFFLINE_GUARD_STATE["count"] += 1 - entered = True + if force: + # Join the active override so the env survives until every + # local-only guard exits, whichever request finishes first. + _OFFLINE_GUARD_STATE["count"] += 1 + entered = True + # An ordinary guard never joins: its block tolerates either env + # state (it would have run online), so extending the forced + # window would only widen the exposure of concurrent online + # work to the process-global override. It no-ops instead. elif _hf_env_offline() and (not force or _hub_offline_env_truthy()): # User-set truthy env (count is 0): already offline, nothing to # arrange or restore. A forced guard only trusts HF_HUB_OFFLINE diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 293675e329..b3372026b7 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -635,6 +635,10 @@ class MLXInferenceBackend: "hf_token": hf_token, # Per-model trust_remote_code reused by the native-template reload (matches transformers). "trust_remote_code": trust_remote_code, + # Local-only loads: generation-time repo fallbacks (native + # template reload) must also resolve from cache, not the Hub. + "local_files_only": local_files_only, + "model_path": load_source, "model": self._model, "tokenizer": self._tokenizer, "processor": self._processor, diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py index 0a17790bf7..eebe242a9b 100644 --- a/studio/backend/hub/services/models/cache_inventory.py +++ b/studio/backend/hub/services/models/cache_inventory.py @@ -542,7 +542,40 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: ) +def _weightful_snapshot_path(repo_path: Path) -> Optional[Path]: + """Newest snapshot dir holding config.json plus a safetensors weight. + + Inactive-cache rows emit this snapshot path as their load_id, which the + load consumes directly (it bypasses the repo-id resolver), so a newest + metadata-only revision must not be emitted while an older revision holds + the weights the row was classified from. + """ + snapshots = repo_path / "snapshots" + try: + revisions = [entry for entry in snapshots.iterdir() if entry.is_dir()] + except OSError: + return None + + def _loadable(rev: Path) -> bool: + try: + names = [entry.name for entry in rev.iterdir()] + except OSError: + return False + return "config.json" in names and any(n.endswith(".safetensors") for n in names) + + candidates = [rev for rev in revisions if _loadable(rev)] + if not candidates: + return None + try: + return max(candidates, key = lambda rev: rev.stat().st_mtime).resolve() + except OSError: + return None + + def _cached_model_snapshot_path(repo_path: Path) -> Optional[Path]: + weightful = _weightful_snapshot_path(repo_path) + if weightful is not None: + return weightful resolved = hf_cache_scan.resolve_hf_cache_realpath(repo_path) if not resolved: return None diff --git a/studio/backend/tests/test_offline_guard_refcount.py b/studio/backend/tests/test_offline_guard_refcount.py index e9d7b285d9..c4c07c0ea0 100644 --- a/studio/backend/tests/test_offline_guard_refcount.py +++ b/studio/backend/tests/test_offline_guard_refcount.py @@ -114,15 +114,33 @@ def test_transformers_only_env_does_not_satisfy_forced_guard(clean_env): assert os.environ.get("TRANSFORMERS_OFFLINE") == "1" -def test_nonforce_joins_active_override(clean_env): - """A DNS-alive non-force guard entering while a forced guard is active must - JOIN the refcount (deferring the restore) rather than no-op.""" +def test_nonforce_never_extends_forced_window(clean_env): + """An ordinary (non-force) guard entering while a forced override is + active must no-op, not join: its block tolerates either env state, and + joining would only widen the exposure of concurrent online work to the + process-global override.""" guard = _load_guard() forced = guard(force = True) plain = guard(force = False) assert forced.__enter__() is True - assert plain.__enter__() is True + assert plain.__enter__() is False forced.__exit__(None, None, None) - assert os.environ.get("HF_HUB_OFFLINE") == "1" + assert "HF_HUB_OFFLINE" not in os.environ, ( + "the forced owner's exit restores; the ordinary no-op guard holds nothing" + ) plain.__exit__(None, None, None) assert "HF_HUB_OFFLINE" not in os.environ + + +def test_forced_guards_share_windows_with_dns_dead_owners(clean_env): + """A forced guard joining a DNS-dead override keeps the env until the + forced guard itself exits, even when the owner exits first.""" + guard = _load_guard(dns_dead = True) + plain = guard(force = False) + forced = guard(force = True) + assert plain.__enter__() is True + assert forced.__enter__() is True + plain.__exit__(None, None, None) + assert os.environ.get("HF_HUB_OFFLINE") == "1" + forced.__exit__(None, None, None) + assert "HF_HUB_OFFLINE" not in os.environ diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 7fc86ccd6a..1ad9a35d76 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1626,6 +1626,44 @@ const BEST_EFFORT_PREFETCH_GRACE_MS = 1_000; // path after this window and join the pool in size order as they resolve. const AUTO_LOAD_RESOLVE_GRACE_MS = 2500; +// Upper bound on ONE GGUF variant scan. Pending scans hold the model-kind +// gate, and the underlying fetch has no timeout of its own, so a single hung +// request would otherwise gate Send forever. Matches the inventory calls' +// own 30s bound; on expiry the scan job settles as failed and a ready +// safetensors candidate can proceed. +const AUTO_LOAD_VARIANT_SCAN_TIMEOUT_MS = 30_000; + +// An HF cache snapshot dir: captures the cache repo ROOT before /snapshots/. +const HF_SNAPSHOT_PATH_RE = /^(.*)[\\/]snapshots[\\/][^\\/]+[\\/]?$/; + +// A custom scan folder registered at an HF cache snapshot aliases the cached +// row for the same repo, but the cached row contributes the cache repo ROOT +// (cache_path) while the custom row contributes its snapshots/ dir. +// Expanding a snapshot path with its cache root lets the two rows collide on +// one seen-set key, so shared files consume one attempt. +function expandSeenValues(value: string): string[] { + const snapshotMatch = HF_SNAPSHOT_PATH_RE.exec(value); + return snapshotMatch ? [value, snapshotMatch[1]] : [value]; +} + +async function listGgufVariantsBounded( + repoId: string, + options: { preferLocalCache?: boolean; localPath?: string | null }, +) { + const controller = new AbortController(); + const timer = setTimeout(() => { + controller.abort(); + }, AUTO_LOAD_VARIANT_SCAN_TIMEOUT_MS); + try { + return await listGgufVariants(repoId, undefined, { + ...options, + signal: controller.signal, + }); + } finally { + clearTimeout(timer); + } +} + type ResolvedLocalCandidate = { candidate: AutoLoadCandidate; /** Size of what would actually load: the resolved quant's own size for a @@ -1649,7 +1687,7 @@ async function resolveLocalRowCandidate( // cache but missing at row's own path. A local-path repo id routes // the backend straight to the filesystem scan of that folder. const variantScanTarget = isLocalModelPath(row.id) ? row.id : row.path; - const variants = await listGgufVariants(variantScanTarget, undefined, { + const variants = await listGgufVariantsBounded(variantScanTarget, { preferLocalCache: true, localPath: row.path, }); @@ -2128,7 +2166,9 @@ export async function autoLoadOnDeviceModel(): Promise<{ ): void => { for (const value of values) { if (value) { - seenLoadTargets.add(`${kind}:${normalizeLoadTargetKey(value)}`); + for (const alias of expandSeenValues(value)) { + seenLoadTargets.add(`${kind}:${normalizeLoadTargetKey(alias)}`); + } } } }; @@ -2139,7 +2179,9 @@ export async function autoLoadOnDeviceModel(): Promise<{ values.some( (value) => !!value && - seenLoadTargets.has(`${kind}:${normalizeLoadTargetKey(value)}`), + expandSeenValues(value).some((alias) => + seenLoadTargets.has(`${kind}:${normalizeLoadTargetKey(alias)}`), + ), ); try { @@ -2188,7 +2230,7 @@ export async function autoLoadOnDeviceModel(): Promise<{ // may still pick another complete quant from this repo (only the // failed candidate key below is excluded). try { - const variants = await listGgufVariants(repo.repo_id, undefined, { + const variants = await listGgufVariantsBounded(repo.repo_id, { preferLocalCache: true, localPath: repo.cache_path, }); @@ -2346,7 +2388,7 @@ export async function autoLoadOnDeviceModel(): Promise<{ const resolveCachedGgufEntry = async ( repo: CachedGgufRepo, ): Promise | null> => { - const variants = await listGgufVariants(repo.repo_id, undefined, { + const variants = await listGgufVariantsBounded(repo.repo_id, { preferLocalCache: true, localPath: repo.cache_path, }); diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 4e6feb7a70..ee05ad8b0d 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -952,6 +952,9 @@ export async function listGgufVariants( options?: { preferLocalCache?: boolean; localPath?: string | null; + /** Aborts the underlying fetch; background scans pass a timeout signal so + * one hung request cannot gate Send forever. */ + signal?: AbortSignal; }, ): Promise { const params = new URLSearchParams({ repo_id: repoId }); @@ -964,6 +967,7 @@ export async function listGgufVariants( } const response = await authFetch(`/api/models/gguf-variants?${params}`, { headers: hubTokenHeader(hfToken), + signal: options?.signal, }); return parseJsonOrThrow(response); } diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index 609aba9858..340f5f44df 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -714,7 +714,7 @@ def test_autoload_deduplicates_cached_and_local_candidates(): assert "const seenLoadTargets = new Set()" in auto_load # Keys carry the model kind: a folder emitting both GGUF and safetensors # rows shares a path while holding two different models. - assert "seenLoadTargets.add(`${kind}:${normalizeLoadTargetKey(value)}`)" in auto_load + assert "seenLoadTargets.add(`${kind}:${normalizeLoadTargetKey(alias)}`)" in auto_load assert 'markSeen("gguf", repo.load_id || repo.repo_id, repo.cache_path)' in auto_load assert 'markSeen("model", repo.load_id || repo.repo_id, repo.cache_path)' in auto_load assert "isSeen(localCandidate.kind, row.load_id, row.id, row.path)" in auto_load @@ -809,7 +809,7 @@ def test_directory_gguf_rows_resolve_variant_like_picker(): # Quants must be resolved from the folder the row will load from, not # from a same-id HF cache repo whose quants may be absent locally. assert "const variantScanTarget = isLocalModelPath(row.id) ? row.id : row.path;" in resolve_fn - assert "listGgufVariants(variantScanTarget" in resolve_fn + assert "listGgufVariantsBounded(variantScanTarget" in resolve_fn assert "localPath: row.path" in resolve_fn assert "entry.downloaded && !entry.partial && isAutoLoadableGgufVariant(entry)" in resolve_fn # The cascade must keep directory GGUF rows as candidates. @@ -1244,6 +1244,53 @@ def test_forced_offline_is_hub_specific_and_covers_parent_preflight(): assert "prepare_gpu_selection(" in tier_probe.split("_spawn_subprocess", 1)[0] +def test_generation_and_scan_paths_stay_bounded_and_local(): + """Round-17 gates. Inactive-cache rows emit a weight-bearing snapshot as + their load_id (the load consumes it directly, bypassing the repo-id + resolver). Ordinary offline guards never extend a forced window. The + generation-time native template reload honors the load's local-only flag + and resolved path instead of refetching the repo id online. Every GGUF + variant scan behind the model-kind gate is bounded by an abortable + timeout, and an HF cache snapshot registered as a custom folder dedupes + against its cached row through the shared cache root.""" + inventory = _read_backend("hub/services/models/cache_inventory.py") + assert "def _weightful_snapshot_path(" in inventory + resolver = inventory.split("def _cached_model_snapshot_path(", 1)[1] + resolver = resolver.split("def ", 1)[0] + assert "_weightful_snapshot_path(repo_path)" in resolver + + llama = _read_backend("core/inference/llama_cpp.py") + join_branch = llama.split('if _OFFLINE_GUARD_STATE["count"] > 0:', 1)[1] + join_branch = join_branch.split("elif", 1)[0] + assert "if force:" in join_branch + + helpers = _read_backend("core/inference/chat_template_helpers.py") + reload_block = helpers.split("native_chat_template", 1)[1] + reload_block = reload_block.split("model_info[", 1)[0] + assert 'local_files_only = bool(model_info.get("local_files_only", False))' in reload_block + assert 'template_source = model_info.get("model_path") or template_source' in reload_block + assert "local_files_only = local_files_only," in reload_block + inference = _read_backend("core/inference/inference.py") + assert '"local_files_only": local_files_only,' in inference + mlx = _read_backend("core/inference/mlx_inference.py") + assert '"local_files_only": local_files_only,' in mlx + assert '"model_path": load_source,' in mlx + + adapter = _read("features/chat/api/chat-adapter.ts") + assert "AUTO_LOAD_VARIANT_SCAN_TIMEOUT_MS" in adapter + assert "async function listGgufVariantsBounded(" in adapter + assert "controller.abort()" in adapter + # No unbounded scan calls remain in the adapter: every call site routes + # through the bounded wrapper (the wrapper itself holds the one direct + # call, with the timeout signal attached). + assert adapter.count("await listGgufVariants(") == 1 + chat_api = _read("features/chat/api/chat-api.ts") + assert "signal: options?.signal," in chat_api + + assert "function expandSeenValues(" in adapter + assert adapter.count("expandSeenValues(value)") == 2 + + def test_gguf_background_loads_never_download_companions(): """A cached GGUF load can still fetch from the Hub through its optional companions (mmproj, MTP drafter) or a cache-miss main quant. Background