diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 30d251727b..0c562708fe 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -506,6 +506,13 @@ def _hf_env_offline() -> bool: return os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in {"1", "true", "yes", "on"} +# Overlapping offline guards share one env override: the count tracks active +# guards and the saved values are restored only when the LAST guard exits, so +# a request finishing early cannot re-enable network for one still running. +_OFFLINE_GUARD_LOCK = threading.Lock() +_OFFLINE_GUARD_STATE: dict = {"count": 0, "hub_prev": None, "transformers_prev": None} + + @contextlib.contextmanager def _hf_offline_if_dns_dead(force: bool = False): """Set HF_HUB_OFFLINE for this block only when DNS to huggingface.co fails; @@ -513,41 +520,56 @@ def _hf_offline_if_dns_dead(force: bool = False): No-op when the user already set it to a truthy value. ``force`` skips the DNS probe and goes offline unconditionally (local-only background loads resolve metadata from the cache without any network), overriding even an - explicitly falsy HF_HUB_OFFLINE=0 for the block and restoring it after.""" - if _hf_env_offline(): + explicitly falsy HF_HUB_OFFLINE=0 for the block and restoring it after. + Guards are refcounted so overlapping loads/validations each keep offline + until the last one exits.""" + entered = False # this guard joined or created the env override + 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 + elif _hf_env_offline(): + # User-set truthy env (count is 0): already offline, nothing to + # arrange or restore. + pass + elif not force and "HF_HUB_OFFLINE" in os.environ: + # A user-pinned falsy value stays authoritative for ordinary loads. + pass + elif force or _probe_dns_dead(): + _OFFLINE_GUARD_STATE["hub_prev"] = os.environ.get("HF_HUB_OFFLINE") + _OFFLINE_GUARD_STATE["transformers_prev"] = os.environ.get("TRANSFORMERS_OFFLINE") + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + _OFFLINE_GUARD_STATE["count"] = 1 + entered = True + owner = True + if not entered: yield False return - if not force: - if "HF_HUB_OFFLINE" in os.environ: - # A user-pinned falsy value stays authoritative for ordinary loads. - yield False - return - if not _probe_dns_dead(): - yield False - return - - hub_prev = os.environ.get("HF_HUB_OFFLINE") - transformers_prev = os.environ.get("TRANSFORMERS_OFFLINE") - wrote_transformers = transformers_prev is None or force - os.environ["HF_HUB_OFFLINE"] = "1" - if wrote_transformers: - os.environ["TRANSFORMERS_OFFLINE"] = "1" - if force: - logger.info("Local-only load: forcing HF offline for this block.") - else: - logger.warning("huggingface.co unreachable; using local HF cache for this load.") + if owner: + if force: + logger.info("Local-only load: forcing HF offline for this block.") + else: + logger.warning("huggingface.co unreachable; using local HF cache for this load.") try: yield True finally: - if hub_prev is None: - os.environ.pop("HF_HUB_OFFLINE", None) - else: - os.environ["HF_HUB_OFFLINE"] = hub_prev - if wrote_transformers: - if transformers_prev is None: - os.environ.pop("TRANSFORMERS_OFFLINE", None) - else: - os.environ["TRANSFORMERS_OFFLINE"] = transformers_prev + with _OFFLINE_GUARD_LOCK: + _OFFLINE_GUARD_STATE["count"] -= 1 + if _OFFLINE_GUARD_STATE["count"] == 0: + for key, prev in ( + ("HF_HUB_OFFLINE", _OFFLINE_GUARD_STATE["hub_prev"]), + ("TRANSFORMERS_OFFLINE", _OFFLINE_GUARD_STATE["transformers_prev"]), + ): + if prev is None: + os.environ.pop(key, None) + else: + os.environ[key] = prev + _OFFLINE_GUARD_STATE["hub_prev"] = None + _OFFLINE_GUARD_STATE["transformers_prev"] = None try: diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 03a98c18dc..293675e329 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -606,8 +606,13 @@ class MLXInferenceBackend: else: load_kwargs["tensor_group"] = distributed_group + # Registry identity stays the repo id (model_name); the LOAD source + # honors config.path so a route-resolved local snapshot (local-only + # loads against a moved live cache) is read instead of re-resolving + # the id through the import-time cache location. + load_source = getattr(config, "path", None) or model_name model, tokenizer_or_processor = FastMLXModel.from_pretrained( - model_name, + load_source, **load_kwargs, ) diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index e7e9e7336b..4e01ece482 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -204,19 +204,45 @@ def _resolve_lora_4bit(mc, load_in_4bit: bool) -> bool: return load_in_4bit -def _ensure_ssm_kernels(targets: list, resp_queue: Any) -> bool: +def _ensure_ssm_kernels(targets: list, resp_queue: Any, local_files_only: bool = False) -> bool: """Install the SSM kernels the given model(s) lazy-import in from_pretrained; no-op for non-SSM models, idempotent. Returns True on success; on a fatal mamba-ssm failure sends a 'loaded' failure response and returns False. Call BEFORE importing transformers, which snapshots its optional-backend gates at import (a later install may not be picked up). + Under ``local_files_only`` nothing is ever installed: kernels already present are used, + a missing fatal kernel fails the load into candidate failover, and the optional + causal-conv1d fast path is skipped (its torch fallback covers it). """ try: - from utils.ssm_runtime import ensure_ssm_runtime + from utils.ssm_runtime import ensure_ssm_runtime, model_is_ssm except Exception as exc: logger.debug("ssm_runtime unavailable (%s); skipping SSM kernel pre-install", exc) return True _ssm_status = lambda m: _send_response(resp_queue, {"type": "status", "message": m}) + if local_files_only: + import importlib.util + + for ssm_target in dict.fromkeys(t for t in targets if t): + try: + needs_mamba = model_is_ssm(ssm_target) + except Exception: + needs_mamba = False + if needs_mamba and importlib.util.find_spec("mamba_ssm") is None: + _send_response( + resp_queue, + { + "type": "loaded", + "success": False, + "message": ( + "This model needs the mamba-ssm kernel, which is not " + "installed; select the model explicitly to install it." + ), + "error_kind": "ssm_runtime_install_failed", + }, + ) + return False + return True try: for ssm_target in dict.fromkeys(t for t in targets if t): ensure_ssm_runtime(ssm_target, status_cb = _ssm_status) @@ -370,7 +396,11 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: str(mc.base_model) if (mc.is_lora and getattr(mc, "base_model", None)) else None ) ssm_targets = [ssm_probe_identifier(config["model_name"], _ssm_base)] - if not _ensure_ssm_kernels(ssm_targets, resp_queue): + if not _ensure_ssm_kernels( + ssm_targets, + resp_queue, + local_files_only = bool(config.get("local_files_only", False)), + ): return # Heartbeat keeps the orchestrator's inactivity deadline alive during slow @@ -834,6 +864,16 @@ def run_inference_process( apply_gpu_ids(config.get("resolved_gpu_ids"), backend = config.get("device_backend")) + # Local-only background loads keep the ENTIRE bootstrap offline: base + # resolution, transformers activation, security gates, kernel probes and + # the initial load can all reach the Hub otherwise. Closed before the + # command loop so generation-time fetches (e.g. the chat template + # fallback) still work; error-path returns end the process anyway. + _bootstrap_offline = contextlib.ExitStack() + _bootstrap_offline.enter_context( + _local_only_offline_env(bool(config.get("local_files_only", False))) + ) + model_name = config["model_name"] # ── 0. MLX fast-path — skip torch/transformers ── @@ -895,6 +935,7 @@ def run_inference_process( return # Enter the same command loop as the GPU path. + _bootstrap_offline.close() logger.info("MLX inference subprocess ready, entering command loop") while True: try: @@ -1050,7 +1091,11 @@ def run_inference_process( from utils.ssm_runtime import ssm_probe_identifier _ssm_targets = [ssm_probe_identifier(model_name, _base)] - if not _ensure_ssm_kernels(_ssm_targets, resp_queue): + if not _ensure_ssm_kernels( + _ssm_targets, + resp_queue, + local_files_only = bool(config.get("local_files_only", False)), + ): return # ── 2. Import ML libraries (fresh in this clean process) ── @@ -1112,6 +1157,8 @@ def run_inference_process( ) return + _bootstrap_offline.close() + # ── 4. Command loop — process commands until shutdown ── # cancel_event is an mp.Event the parent can set anytime to cancel # generation instantly (no queue polling needed). diff --git a/studio/backend/hub/schemas/inventory.py b/studio/backend/hub/schemas/inventory.py index ca0f4658a3..581d6dc90d 100644 --- a/studio/backend/hub/schemas/inventory.py +++ b/studio/backend/hub/schemas/inventory.py @@ -188,6 +188,13 @@ class CachedModelRepo(CachedRepoBase): pipeline_tag: Optional[str] = None library_name: Optional[str] = None tags: Optional[List[str]] = None + snapshot_size_bytes: Optional[int] = Field( + None, + description = ( + "Weight bytes of the newest cached snapshot only (what a load " + "resolves); size_bytes sums blobs across every cached revision." + ), + ) class CachedModelsResponse(BaseModel): diff --git a/studio/backend/hub/services/models/cache_inventory.py b/studio/backend/hub/services/models/cache_inventory.py index 807ec70991..9cc835819f 100644 --- a/studio/backend/hub/services/models/cache_inventory.py +++ b/studio/backend/hub/services/models/cache_inventory.py @@ -415,16 +415,34 @@ async def list_cached_gguf_response(hf_token: Optional[str] = None): class _CachedNonGgufPayload(NamedTuple): size_bytes: int + snapshot_size_bytes: int has_runnable_weights: bool model_format: ModelFormat last_modified: float +def _snapshot_dir_mtime(revision) -> float: + """mtime of a revision's snapshot dir; the same signal latest_snapshot_dir + (and therefore the load-side snapshot resolution) selects by.""" + snapshot_path = getattr(revision, "snapshot_path", None) + if not snapshot_path: + return 0.0 + try: + return float(Path(snapshot_path).stat().st_mtime) + except OSError: + return 0.0 + + def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: all_weight_blobs: dict[str, tuple[int, float]] = {} adapter_blobs: dict[str, tuple[int, float]] = {} safetensors_blobs: dict[str, tuple[int, float]] = {} checkpoint_blobs: dict[str, tuple[int, float]] = {} + # Per-revision selected-format byte sums, so the row can also report the + # size of the ONE snapshot a load resolves (newest by snapshot-dir mtime) + # rather than only the all-revisions total. + rev_category_sizes: dict[str, dict[str, int]] = {} + rev_snapshot_mtimes: dict[str, float] = {} has_config = False has_adapter_config = False has_adapter_weights = False @@ -433,7 +451,11 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: has_checkpoint = False def _record_blob( - target: dict[str, tuple[int, float]], file_obj, rev_id: str, file_name: str + target: dict[str, tuple[int, float]], + file_obj, + rev_id: str, + file_name: str, + category: str, ) -> None: blob_path = getattr(file_obj, "blob_path", None) size = int(file_obj.size_on_disk or 0) @@ -441,9 +463,13 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: value = (size, _blob_mtime(file_obj)) target[key] = value all_weight_blobs[key] = value + per_rev = rev_category_sizes.setdefault(rev_id, {}) + per_rev[category] = per_rev.get(category, 0) + size + per_rev["all"] = per_rev.get("all", 0) + size for revision in repo_info.revisions: rev_id = getattr(revision, "commit_hash", None) or str(id(revision)) + rev_snapshot_mtimes[rev_id] = _snapshot_dir_mtime(revision) for f in revision.files: file_name = str(f.file_name) lower = file_name.lower() @@ -461,15 +487,15 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: is_checkpoint = _is_checkpoint_weight_name(name) if is_adapter: has_adapter_weights = True - _record_blob(adapter_blobs, f, rev_id, file_name) + _record_blob(adapter_blobs, f, rev_id, file_name, "adapter") if is_safetensors: has_safetensors = True if _is_transformers_safetensors_weight_name(name): has_transformers_safetensors = True - _record_blob(safetensors_blobs, f, rev_id, file_name) + _record_blob(safetensors_blobs, f, rev_id, file_name, "safetensors") if is_checkpoint: has_checkpoint = True - _record_blob(checkpoint_blobs, f, rev_id, file_name) + _record_blob(checkpoint_blobs, f, rev_id, file_name, "checkpoint") model_format = ( _classify_non_gguf_model_format( @@ -485,15 +511,35 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload: ) if model_format == "adapter": selected_blobs = adapter_blobs + selected_category = "adapter" elif model_format == "safetensors": selected_blobs = safetensors_blobs + selected_category = "safetensors" elif model_format == "checkpoint": selected_blobs = checkpoint_blobs + selected_category = "checkpoint" else: selected_blobs = all_weight_blobs + selected_category = "all" + + size_bytes = sum(size for size, _mtime in selected_blobs.values()) + # The one snapshot a load resolves (newest snapshot-dir mtime) among the + # revisions actually holding selected-format weights; falls back to the + # all-revisions total when no revision reports one. + weight_revs = [ + rev_id + for rev_id, sizes in rev_category_sizes.items() + if sizes.get(selected_category, 0) > 0 + ] + if weight_revs: + newest_rev = max(weight_revs, key = lambda rev_id: rev_snapshot_mtimes.get(rev_id, 0.0)) + snapshot_size_bytes = rev_category_sizes[newest_rev].get(selected_category, 0) + else: + snapshot_size_bytes = size_bytes return _CachedNonGgufPayload( - size_bytes = sum(size for size, _mtime in selected_blobs.values()), + size_bytes = size_bytes, + snapshot_size_bytes = snapshot_size_bytes, has_runnable_weights = model_format != "unknown", model_format = model_format, last_modified = max((mtime for _size, mtime in selected_blobs.values()), default = 0.0), @@ -636,6 +682,7 @@ def _scan_cached_models() -> list[dict]: row = { "repo_id": repo_id, "size_bytes": payload.size_bytes, + "snapshot_size_bytes": payload.snapshot_size_bytes, "cache_path": str(repo_info.repo_path), "partial": snapshot_partial, "partial_transport": ( diff --git a/studio/backend/hub/utils/local_snapshot.py b/studio/backend/hub/utils/local_snapshot.py index da37bc8ab2..23d3f0a243 100644 --- a/studio/backend/hub/utils/local_snapshot.py +++ b/studio/backend/hub/utils/local_snapshot.py @@ -8,10 +8,12 @@ from typing import Optional def _snapshot_dir_fallback(repo_id: str, cache_dir: Optional[str]) -> Optional[str]: - """Newest snapshots/* dir holding a config.json for a cache entry whose - refs/ are missing or pruned. snapshot_download(local_files_only = True) - needs refs/main to map the ref to a revision, but the inventory scanner - accepts revision-only layouts, so background loads must resolve them too. + """Newest snapshots/* dir (by mtime) holding a config.json. + + This mirrors the inventory scanner's latest_snapshot_dir selection, so the + load targets the snapshot that made the row eligible. It also covers + revision-only layouts (refs/ missing or pruned) that + snapshot_download(local_files_only = True) cannot map through refs/main. """ if cache_dir is None: try: @@ -42,12 +44,20 @@ def resolve_local_snapshot_path( """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. + The newest snapshot dir (by mtime) holding a config.json is preferred: + that is the same selection the inventory scanner surfaces, so the load + targets the revision that made the row eligible. refs/main can lag it + when a non-main commit was downloaded later, so consulting refs first + could load an older revision or fail on its missing files. + ``snapshot_download(local_files_only = True)`` stays as the fallback for + layouts the directory scan cannot interpret. Either way resolution never + touches the network, and a missing-file snapshot still resolves so the + subsequent weight load fails instead of downloading the gaps, which is + the fail-closed behavior background loads need. """ + resolved = _snapshot_dir_fallback(repo_id, cache_dir) + if resolved is not None: + return resolved try: from huggingface_hub import snapshot_download return snapshot_download( @@ -57,4 +67,4 @@ def resolve_local_snapshot_path( cache_dir = cache_dir or None, ) except Exception: - return _snapshot_dir_fallback(repo_id, cache_dir) + return None diff --git a/studio/backend/tests/test_local_snapshot_resolution.py b/studio/backend/tests/test_local_snapshot_resolution.py index 3d9f0197cc..037d26f68e 100644 --- a/studio/backend/tests/test_local_snapshot_resolution.py +++ b/studio/backend/tests/test_local_snapshot_resolution.py @@ -84,6 +84,34 @@ 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_newest_snapshot_preferred_over_refs_main(tmp_path): + """A newer snapshot downloaded at an explicit revision outranks the older + refs/main target: the inventory surfaces the newest snapshot by mtime, so + the load must resolve the same one instead of an older (possibly + incomplete) main revision.""" + import os + import time + + old_main = _build_cached_repo( + tmp_path, + "org/newer-rev", + {"config.json": "{}"}, + rev = "a" * 40, + ) + stale = time.time() - 1000 + os.utime(old_main, (stale, stale)) + newer = _build_cached_repo( + tmp_path, + "org/newer-rev", + {"config.json": "{}", "model.safetensors": "weights"}, + with_refs = False, + rev = "b" * 40, + ) + resolved = resolve_local_snapshot_path("org/newer-rev", cache_dir = str(tmp_path)) + assert resolved is not None + assert Path(resolved).resolve() == newer.resolve() + + def test_revision_only_snapshot_resolves_without_refs(tmp_path): """snapshot_download(local_files_only = True) needs refs/main, but the inventory scanner accepts revision-only layouts (pruned refs), so the diff --git a/studio/backend/tests/test_offline_guard_refcount.py b/studio/backend/tests/test_offline_guard_refcount.py new file mode 100644 index 0000000000..1be7e1e324 --- /dev/null +++ b/studio/backend/tests/test_offline_guard_refcount.py @@ -0,0 +1,114 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Refcounted HF offline guard semantics. + +Overlapping local-only validations/loads share one process-global env +override. The refcount means a request finishing first cannot restore the +environment while another local-only request still runs (which would let its +remaining metadata checks reach the Hub), and forced mode overrides an +explicitly falsy HF_HUB_OFFLINE=0 then restores it. The guard is extracted +from source and exercised with stubbed logging/DNS so no ML dependencies are +needed. +""" + +from __future__ import annotations + +import contextlib +import os +import threading +from pathlib import Path + +import pytest + +_LLAMA_CPP = Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py" + + +class _NullLogger: + def info(self, *args, **kwargs): + pass + + def warning(self, *args, **kwargs): + pass + + +def _load_guard(dns_dead: bool = False): + src = _LLAMA_CPP.read_text() + start = src.index("# Overlapping offline guards") + end = src.index("_SLOT_SAVE_MAX_BYTES") + end = src.rindex("try:", start, end) + block = src[start:end] + ns = { + "threading": threading, + "contextlib": contextlib, + "os": os, + "logger": _NullLogger(), + "_hf_env_offline": lambda: os.environ.get("HF_HUB_OFFLINE", "").strip().lower() + in {"1", "true", "yes", "on"}, + "_probe_dns_dead": lambda: dns_dead, + } + exec(block, ns) + return ns["_hf_offline_if_dns_dead"] + + +@pytest.fixture +def clean_env(monkeypatch): + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + + +def test_overlapping_guards_restore_only_after_last_exit(clean_env): + guard = _load_guard() + a = guard(force = True) + b = guard(force = True) + assert a.__enter__() is True + assert os.environ.get("HF_HUB_OFFLINE") == "1" + assert b.__enter__() is True + a.__exit__(None, None, None) + assert os.environ.get("HF_HUB_OFFLINE") == "1", ( + "first exit must not restore while another guard is active" + ) + b.__exit__(None, None, None) + assert "HF_HUB_OFFLINE" not in os.environ + + +def test_force_overrides_and_restores_falsy_env(clean_env): + guard = _load_guard() + os.environ["HF_HUB_OFFLINE"] = "0" + g = guard(force = True) + assert g.__enter__() is True + assert os.environ["HF_HUB_OFFLINE"] == "1" + g.__exit__(None, None, None) + assert os.environ["HF_HUB_OFFLINE"] == "0" + + +def test_falsy_env_stays_authoritative_for_ordinary_loads(clean_env): + guard = _load_guard(dns_dead = True) + os.environ["HF_HUB_OFFLINE"] = "0" + g = guard(force = False) + assert g.__enter__() is False + assert os.environ["HF_HUB_OFFLINE"] == "0" + g.__exit__(None, None, None) + + +def test_truthy_user_env_is_a_noop(clean_env): + guard = _load_guard() + os.environ["HF_HUB_OFFLINE"] = "1" + g = guard(force = True) + assert g.__enter__() is False + g.__exit__(None, None, None) + assert os.environ["HF_HUB_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.""" + guard = _load_guard() + forced = guard(force = True) + plain = guard(force = False) + assert forced.__enter__() is True + assert plain.__enter__() is True + forced.__exit__(None, None, None) + assert os.environ.get("HF_HUB_OFFLINE") == "1" + plain.__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 2e4cff7a8a..7fc86ccd6a 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -1468,6 +1468,12 @@ function isAutoLoadableCachedRepo(repo: { }): boolean { if (repo.partial) return false; if (repo.model_format === "adapter") return false; + // Cached checkpoint repos (pickle .bin/.pt weights) stay interactive-only, + // like local checkpoint rows: forced-offline validation cannot consult the + // Hub security scan, and deserializing a pickle can execute code. + if (repo.model_format === "checkpoint") { + return false; + } if (repo.capabilities?.can_chat === false) return false; return !isHiddenModelId(repo.repo_id); } @@ -2326,7 +2332,13 @@ export async function autoLoadOnDeviceModel(): Promise<{ insertReady({ type: "cached-model", repo, - sizeBytes: sizeOrUnknownBytes(repo.size_bytes), + // Order by the snapshot the load will actually resolve: the row's + // size_bytes sums weight blobs across EVERY cached revision, so a + // small current revision beside a huge stale one would otherwise be + // ranked as their total and sink behind genuinely larger candidates. + sizeBytes: sizeOrUnknownBytes( + repo.snapshot_size_bytes ?? repo.size_bytes, + ), }); } // Smallest complete, auto-loadable, not-yet-skipped quant of a managed diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 8690671380..4e6feb7a70 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/studio/frontend/src/features/chat/api/chat-api.ts @@ -383,6 +383,9 @@ export interface CachedModelRepo { repo_id: string; load_id?: string | null; size_bytes: number; + /** Weight bytes of the newest cached snapshot only (what a load resolves); + * size_bytes sums selected-format blobs across every cached revision. */ + snapshot_size_bytes?: number | null; /** Epoch seconds of the newest downloaded weight file; sorts Downloaded * newest-first. Optional for older-backend compatibility. */ last_modified?: number; diff --git a/studio/frontend/src/features/hub/inventory/api.ts b/studio/frontend/src/features/hub/inventory/api.ts index d3a1abb540..42f3cb4d64 100644 --- a/studio/frontend/src/features/hub/inventory/api.ts +++ b/studio/frontend/src/features/hub/inventory/api.ts @@ -65,6 +65,9 @@ export interface CachedModelRepo { format_variant?: string | null; capabilities?: BackendModelCapabilities | null; size_bytes: number; + /** Weight bytes of the newest cached snapshot only (what a load resolves); + * size_bytes sums selected-format blobs across every cached revision. */ + snapshot_size_bytes?: number | null; cache_path?: string; last_modified?: number | null; partial?: boolean; diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index cb21e8eb15..aba294af69 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -848,12 +848,11 @@ def test_fallback_orders_by_resolved_quant_size(): # Cached GGUF repos order on the resolved quant size too. assert "const resolveCachedGgufEntry" in auto_load assert "sizeBytes: sizeOrUnknownBytes(variant.size_bytes)" in auto_load - # The all-variant row sum only orders non-GGUF cached repos, whose - # snapshot loads whole. + # Non-GGUF cached repos order on the SELECTED snapshot's size (falling + # back to the all-revisions row sum for older backends). seed_block = auto_load.split("for (const repo of platform.chatOnly ? [] : modelRepos)", 1)[1] seed_block = seed_block.split("const resolveCachedGgufEntry", 1)[0] - assert "sizeOrUnknownBytes(repo.size_bytes)" in seed_block - assert auto_load.count("sizeOrUnknownBytes(repo.size_bytes)") == 1 + assert "repo.snapshot_size_bytes ?? repo.size_bytes" in seed_block def test_cascade_retries_next_quant_after_load_failure(): @@ -1112,11 +1111,13 @@ def test_background_candidate_filters_have_no_side_effects(): llama = _read_backend("core/inference/llama_cpp.py") assert "def _hf_offline_if_dns_dead(force: bool = False):" in llama - # force must also override an explicitly falsy HF_HUB_OFFLINE=0: only a - # TRUTHY env value short-circuits, and prior values are restored on exit. - assert "if _hf_env_offline():" in llama - assert 'hub_prev = os.environ.get("HF_HUB_OFFLINE")' in llama - assert 'os.environ["HF_HUB_OFFLINE"] = hub_prev' in llama + # force must also override an explicitly falsy HF_HUB_OFFLINE=0 (only a + # TRUTHY env value short-circuits), and overlapping guards are refcounted + # so the env is restored only when the LAST one exits; behavior is + # exercised directly in test_offline_guard_refcount.py. + assert "elif _hf_env_offline():" in llama + assert "_OFFLINE_GUARD_LOCK" in llama + assert '_OFFLINE_GUARD_STATE["count"] += 1' in llama helper = _read_backend("hub/utils/local_snapshot.py") assert "def _snapshot_dir_fallback(" in helper @@ -1170,6 +1171,52 @@ def test_local_only_covers_every_load_and_validate_network_path(): assert "local_files_only = local_files_only," in processor_call +def test_background_picks_mirror_inventory_and_skip_installers(): + """Round-15 gates. Cached checkpoint repos (pickle weights) are excluded + from background picks like local checkpoint rows. The worker keeps its + ENTIRE bootstrap offline under local-only and never pip-installs SSM + kernels (missing fatal kernels fail into candidate failover). Snapshot + resolution prefers the newest snapshot dir, matching the inventory + scanner's latest_snapshot_dir selection, before consulting refs/main. + MLX loads read config.path so the live-cache rewrite is honored. Cached + non-GGUF ordering uses the selected snapshot's size, not the + all-revisions blob total.""" + adapter = _read("features/chat/api/chat-adapter.ts") + cached_filter = adapter.split("function isAutoLoadableCachedRepo", 1)[1] + cached_filter = cached_filter.split("AUTO_LOAD_LOCAL_SOURCES", 1)[0] + assert 'if (repo.model_format === "checkpoint") {' in cached_filter + assert "repo.snapshot_size_bytes ?? repo.size_bytes" in adapter + + worker = _read_backend("core/inference/worker.py") + assert "_bootstrap_offline = contextlib.ExitStack()" in worker + # Entered before base resolution / gates / kernels, closed before BOTH + # command loops (MLX and GPU paths). + bootstrap = worker.split("_bootstrap_offline = contextlib.ExitStack()", 1)[1] + assert bootstrap.count("_bootstrap_offline.close()") == 2 + assert "def _ensure_ssm_kernels(targets: list, resp_queue: Any, local_files_only: bool = False) -> bool:" in worker + ssm = worker.split("def _ensure_ssm_kernels", 1)[1] + ssm = ssm.split("def _run_security_gates", 1)[0] + assert "if local_files_only:" in ssm + assert 'importlib.util.find_spec("mamba_ssm") is None' in ssm + + helper = _read_backend("hub/utils/local_snapshot.py") + resolve = helper.split("def resolve_local_snapshot_path", 1)[1] + # Newest-snapshot scan runs BEFORE the refs/main-based resolver (compare + # the actual calls, not docstring mentions). + assert resolve.index("resolved = _snapshot_dir_fallback(") < resolve.index( + "return snapshot_download(" + ) + + mlx = _read_backend("core/inference/mlx_inference.py") + assert 'load_source = getattr(config, "path", None) or model_name' in mlx + + inventory = _read_backend("hub/services/models/cache_inventory.py") + assert "snapshot_size_bytes" in inventory + assert "def _snapshot_dir_mtime(" in inventory + schema = _read_backend("hub/schemas/inventory.py") + assert "snapshot_size_bytes" in schema + + 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