Align background picks with inventory selection and drop remaining install/network paths
Cached checkpoint repos (pickle weights) are excluded from background picks like local checkpoint rows: forced-offline validation cannot consult the Hub security scan and pickles can execute code on load. The worker keeps its entire bootstrap offline under local-only (base resolution, transformers activation, security gates, kernel probes and the initial load) and never pip-installs SSM kernels for a background load; a missing fatal kernel fails into candidate failover. The offline guard is refcounted so overlapping local-only requests share one env override restored only when the last exits, closing the race where one request finishing re-enabled network for another still running. Snapshot resolution now prefers the newest snapshot dir, the same selection the inventory scanner surfaces, before consulting refs/main. MLX loads read config.path so the live-cache rewrite is honored. Cached non-GGUF rows carry snapshot_size_bytes (the newest snapshot's weight bytes) and the cascade orders on it instead of the all-revisions blob total, so a small current revision no longer sinks behind larger candidates.
This commit is contained in:
Unsloth
parent
1dce5373aa
commit
db64cd2ab4
12 changed files with 405 additions and 60 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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": (
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
114
studio/backend/tests/test_offline_guard_refcount.py
Normal file
114
studio/backend/tests/test_offline_guard_refcount.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue