diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 7978a200c0..ae4b386589 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -364,6 +364,7 @@ jobs: tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py \ + tests/test_prefetch_snapshot_scope.py \ --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' # The deselected test monkeypatches flash_attn_varlen_func, which is # only bound on the module when `flash_attn` is importable. flash_attn diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py index 9e40fbf508..4d73213d15 100644 --- a/studio/backend/tests/test_hf_xet_fallback.py +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -1,18 +1,16 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Unit tests for utils.hf_xet_fallback: the no-progress watchdog, the Xet->HTTP -transport policy, and the HF_HUB_DISABLE_XET precondition the fallback rests on. -CPU-only, no network, no real subprocess (the per-attempt download seam is -monkeypatched). +"""Tests for the Studio shim over the shared unsloth_zoo Xet -> HTTP fallback. + +The transport-policy matrix is tested once in unsloth_zoo; here we assert only the +Studio seam: re-exporting the shared API and injecting the marker-aware +prepare_cache_for_transport on the HTTP retry. CPU-only, no network, no real subprocess. """ from __future__ import annotations -import subprocess import sys -import threading -import time import types as _types from pathlib import Path @@ -22,9 +20,8 @@ if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) -# Stub heavy/unavailable deps before importing the module under test. Use the -# real structlog when present; a bare stub left in sys.modules would break later -# modules that log at import time. +# Stub heavy/unavailable deps before importing the module under test. Use real structlog when present; +# a bare stub would break later modules that log at import time. _loggers_stub = _types.ModuleType("loggers") _loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) sys.modules.setdefault("loggers", _loggers_stub) @@ -34,171 +31,59 @@ sys.modules["structlog"] = _types.ModuleType("structlog") import huggingface_hub -from huggingface_hub import constants as hf_constants - -import utils.hf_xet_fallback as xf - - -# --------------------------------------------------------------------------- # -# Watchdog: fires only on a constant-size .incomplete, sparse-aware byte total. -# --------------------------------------------------------------------------- # -REPO = "ztest/xet-watchdog" - - -@pytest.fixture -def hf_cache(tmp_path, monkeypatch): - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) - return tmp_path - - -def _blobs_dir(root: Path, repo_id: str = REPO) -> Path: - d = root / f"models--{repo_id.replace('/', '--')}" / "blobs" - d.mkdir(parents = True, exist_ok = True) - return d - - -def _wait( - predicate, - timeout: float = 2.0, - step: float = 0.02, -) -> bool: - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if predicate(): - return True - time.sleep(step) - return predicate() - - -def test_constant_incomplete_fires_stall(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "deadbeef.incomplete").write_bytes(b"\0" * 1024) # never grows - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - assert _wait( - lambda: len(calls) >= 1, timeout = 3.0 - ), "watchdog never fired on a constant-size .incomplete" - finally: - stop.set() - assert "stalled" in calls[0].lower() - - -def test_growing_incomplete_never_stalls(hf_cache): - blobs = _blobs_dir(hf_cache) - part = blobs / "growing.incomplete" - part.write_bytes(b"\0" * 1024) - - grow_stop = threading.Event() - - def _grow(): - size = 1024 - while not grow_stop.wait(0.05): - size += 4096 - part.write_bytes(b"\0" * size) - - grower = threading.Thread(target = _grow, daemon = True) - grower.start() - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - time.sleep(1.0) # well past stall_timeout, but bytes keep growing - assert calls == [], "watchdog fired despite continuous progress" - finally: - stop.set() - grow_stop.set() - - -def test_no_incomplete_never_stalls(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "finalized_blob").write_bytes(b"\0" * 4096) # no .incomplete - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.3 - ) - try: - time.sleep(0.8) - assert calls == [], "watchdog fired with no active .incomplete" - finally: - stop.set() - - -def test_stall_fires_at_most_once(hf_cache): - blobs = _blobs_dir(hf_cache) - (blobs / "frozen.incomplete").write_bytes(b"\0" * 2048) - - calls: list[str] = [] - stop = xf.start_watchdog( - repo_ids = [REPO], on_stall = calls.append, interval = 0.05, stall_timeout = 0.2 - ) - try: - assert _wait(lambda: len(calls) >= 1, timeout = 3.0) - time.sleep(0.6) # keep ticking; must not fire again - assert len(calls) == 1, f"on_stall fired {len(calls)} times, expected exactly 1" - finally: - stop.set() +try: + import unsloth_zoo.hf_xet_fallback as _shared_mod + shared = _shared_mod +except Exception: # noqa: BLE001 - still collect degraded-path tests when unsloth_zoo is unavailable + shared = None -def test_get_state_empty_cache(hf_cache): - assert xf.get_hf_download_state([REPO]) == (0, False) +import utils.hf_xet_fallback as xf -def test_get_state_absent_cache_root(tmp_path, monkeypatch): - monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path / "no-such-cache")) - assert xf.get_hf_download_state([REPO]) == (0, False) +DL_REPO, FILE = "ztest/xet-dl", "model-Q4_K_XL.gguf" -def test_get_state_skips_local_paths(hf_cache): - # Filesystem paths are not HF repo IDs and must be ignored without error. - assert xf.get_hf_download_state(["/abs/path", "./rel", "~user", "c:\\x"]) == (0, False) +def _requires_shared(): + if shared is None: + pytest.skip("unsloth_zoo.hf_xet_fallback is not installed in this environment") -def test_get_state_sparse_aware(hf_cache): - blobs = _blobs_dir(hf_cache) - sparse = blobs / "sparse.incomplete" - with open(sparse, "wb") as f: - f.truncate(64 * 1024 * 1024) # large apparent size, few allocated blocks - st = sparse.stat() - if getattr(st, "st_blocks", 0) == 0: - pytest.skip("filesystem does not report st_blocks; sparse accounting unavailable") - total, has_incomplete = xf.get_hf_download_state([REPO]) - assert has_incomplete is True - assert total < st.st_size, "sparse partial counted at apparent size, not allocated blocks" +def test_shim_reexports_shared_api(): + _requires_shared() + assert xf.DownloadStallError is shared.DownloadStallError + for name in ( + "start_watchdog", + "get_hf_download_state", + "child_should_disable_xet", + "hf_hub_download_with_xet_fallback", + "snapshot_download_with_xet_fallback", + ): + assert hasattr(xf, name), f"shim missing {name}" -# --------------------------------------------------------------------------- # -# Transport policy: cached short-circuit, cancel, error propagation, and the -# single Xet->HTTP fallback. _run_download_attempt is faked, so no real spawn. -# --------------------------------------------------------------------------- # -DL_REPO, FILE = "ztest/xet-dl", "model-Q4_K_XL.gguf" +def test_child_should_disable_xet_truth_table(): + assert xf.child_should_disable_xet({"disable_xet": True}) is True + assert xf.child_should_disable_xet({"disable_xet": False}) is False + assert xf.child_should_disable_xet({}) is False -@pytest.fixture(autouse = True) -def _no_real_cache_hit(monkeypatch): - """Default: the cached probe misses; tests override it to force a hit.""" +def test_shim_injects_studio_prepare_on_http_retry(monkeypatch): + """A Xet stall retries over HTTP and the shim runs Studio's marker-aware + ``prepare_cache_for_transport(..., 'http')`` before the retry.""" + _requires_shared() + for var in ("UNSLOTH_DISABLE_XET", "UNSLOTH_STABLE_DOWNLOADS", "HF_HUB_DISABLE_XET"): + monkeypatch.delenv(var, raising = False) monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: None) + seen_disable_xet = [] -class _FakeAttempt: - """Records calls to the download seam and returns scripted results.""" - - def __init__(self, results): - self._results = list(results) - self.calls = [] - - def __call__( - self, + def fake_attempt( repo_id, - filename, - token, *, + kind, + params, + token, repo_type, disable_xet, cancel_event, @@ -208,146 +93,243 @@ def __call__( on_status, force_download = False, ): - self.calls.append( - _types.SimpleNamespace( - repo_id = repo_id, - filename = filename, - disable_xet = disable_xet, - repo_type = repo_type, - ) - ) - return self._results[len(self.calls) - 1] + seen_disable_xet.append(disable_xet) + return ("ok", "/cache/model.gguf") if disable_xet else ("stall", None) + monkeypatch.setattr(shared, "_run_download_attempt", fake_attempt) -def _install(monkeypatch, results): - fake = _FakeAttempt(results) - monkeypatch.setattr(xf, "_run_download_attempt", fake) - return fake - - -def test_cached_file_short_circuits(monkeypatch, tmp_path): - cached = tmp_path / "cached.gguf" - cached.write_bytes(b"\0" * 8) - monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: str(cached)) - fake = _install(monkeypatch, []) # must not be called - - out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert out == str(cached) - assert fake.calls == [], "spawned a download for an already-cached file" - - -def test_cancel_before_start_raises_no_attempt(monkeypatch): - fake = _install(monkeypatch, []) - ev = threading.Event() - ev.set() - with pytest.raises(RuntimeError, match = "Cancelled"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None, cancel_event = ev) - assert fake.calls == [] - - -def test_nonstall_error_propagates_without_fallback(monkeypatch): - fake = _install(monkeypatch, [("error", "RepositoryNotFoundError: 404 not found")]) - with pytest.raises(RuntimeError, match = "RepositoryNotFoundError"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 1, "deterministic error must not trigger an HTTP fallback" - assert fake.calls[0].disable_xet is False - - -def test_immediate_success_uses_xet_only(monkeypatch): - prepared = [] - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", - lambda *a, **k: prepared.append(a), - ) - fake = _install(monkeypatch, [("ok", "/cache/model.gguf")]) - out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert out == "/cache/model.gguf" - assert len(fake.calls) == 1 and fake.calls[0].disable_xet is False - assert prepared == [], "no cache prep should run when Xet succeeds first try" - - -def test_stall_then_http_fallback_succeeds(monkeypatch): prepared = [] monkeypatch.setattr( "hub.utils.download_registry.prepare_cache_for_transport", lambda repo_type, repo_id, mode, *a, **k: prepared.append((repo_type, repo_id, mode)), ) - fake = _install(monkeypatch, [("stall", None), ("ok", "/cache/model.gguf")]) out = xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) assert out == "/cache/model.gguf" - assert len(fake.calls) == 2 - assert fake.calls[0].disable_xet is False # Xet first - assert fake.calls[1].disable_xet is True # HTTP fallback - assert prepared == [("model", DL_REPO, "http")], "must prep cache for HTTP before the retry" - - -def test_second_stall_raises_download_stall_error(monkeypatch): - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None - ) - fake = _install(monkeypatch, [("stall", None), ("stall", None)]) - with pytest.raises(xf.DownloadStallError): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 2 - - -def test_cancelled_midattempt_raises_no_fallback(monkeypatch): - fake = _install(monkeypatch, [("cancelled", None)]) - with pytest.raises(RuntimeError, match = "Cancelled"): - xf.hf_hub_download_with_xet_fallback(DL_REPO, FILE, None) - assert len(fake.calls) == 1 - - -def test_per_file_independent_fallback(monkeypatch): - """A stalled shard falls back; a sibling shard that succeeds does not.""" - monkeypatch.setattr( - "hub.utils.download_registry.prepare_cache_for_transport", lambda *a, **k: None - ) - fake = _install(monkeypatch, [("ok", "/a"), ("stall", None), ("ok", "/b")]) - assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardA.gguf", None) == "/a" - assert xf.hf_hub_download_with_xet_fallback(DL_REPO, "shardB.gguf", None) == "/b" - assert [c.disable_xet for c in fake.calls] == [False, False, True] - - -# --------------------------------------------------------------------------- # -# Precondition: HF_HUB_DISABLE_XET is read at import time, so assert its effect -# in a FRESH interpreter (huggingface/huggingface_hub#3266 once ignored it). -# --------------------------------------------------------------------------- # -def _safe_path() -> str: + assert seen_disable_xet == [False, True] # Xet first, then HTTP + assert prepared == [("model", DL_REPO, "http")], "shim must run Studio's marker-aware prep" + + +def test_shim_snapshot_injects_studio_prepare(monkeypatch): + """The snapshot wrapper forwards Studio's marker-aware prep, like the file wrapper.""" + captured = {} + + def fake_snapshot(repo_id, **kwargs): + captured["repo_id"] = repo_id + captured["prepare_for_http_fn"] = kwargs.get("prepare_for_http_fn") + return "/tmp/snap-dir" + + monkeypatch.setattr(xf, "_shared_snapshot_download_with_xet_fallback", fake_snapshot) + out = xf.snapshot_download_with_xet_fallback("org/model") + assert out == "/tmp/snap-dir" + assert captured["repo_id"] == "org/model" + assert captured["prepare_for_http_fn"] is xf._studio_prepare_for_http + + +def test_degrades_gracefully_without_shared_helper(monkeypatch): + """On an older unsloth_zoo lacking the shared helper, the shim still imports (Studio + boots) and exposes stub API doing plain HF downloads with the watchdog disabled.""" + import importlib + + class _BlockShared: + def find_spec( + self, + name, + path = None, + target = None, + ): + if name == "unsloth_zoo.hf_xet_fallback": + raise ModuleNotFoundError(f"No module named '{name}'", name = name) + return None + + finder = _BlockShared() + saved_shared = sys.modules.pop("unsloth_zoo.hf_xet_fallback", None) + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + + # Boots without raising and mirrors the shared API surface. + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.child_should_disable_xet({"disable_xet": True}) is True + assert degraded.get_hf_download_state(["x"]) is None # unmeasurable + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() # never fires + + # Degraded mode still emits heartbeats so the inactivity deadline is not tripped. + import time as _time + + beats = [] + hb_stop = degraded.start_watchdog( + repo_ids = ["x"], + on_stall = lambda m: None, + on_heartbeat = beats.append, + interval = 0.02, + ) + try: + deadline = _time.monotonic() + 2.0 + while not beats and _time.monotonic() < deadline: + _time.sleep(0.02) + assert beats, "degraded watchdog emitted no heartbeat" + finally: + hb_stop.set() + + # Downloads fall back to plain huggingface_hub (no watchdog, no crash). + called = {} + + def _fake_snapshot(repo_id, **kwargs): + called["repo_id"] = repo_id + return "/snap-dir" + + monkeypatch.setattr(huggingface_hub, "snapshot_download", _fake_snapshot) + assert degraded.snapshot_download_with_xet_fallback("org/model") == "/snap-dir" + assert called["repo_id"] == "org/model" + + # Cancellation still holds: an already-set cancel_event aborts before the HF download. + import threading as _threading + + cancelled = _threading.Event() + cancelled.set() + called.clear() + with pytest.raises(RuntimeError, match = "Cancelled"): + degraded.snapshot_download_with_xet_fallback("org/model", cancel_event = cancelled) + assert "repo_id" not in called, "degraded download ran despite cancellation" + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + if saved_shared is not None: + sys.modules["unsloth_zoo.hf_xet_fallback"] = saved_shared + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim + + +def test_degrades_when_unsloth_zoo_entirely_absent(): + """When unsloth_zoo is absent entirely, the import raises + ModuleNotFoundError(name='unsloth_zoo') (top-level package). Guard that the shim still + degrades and does not re-raise, breaking every Studio import that pulls it in.""" + import importlib + + class _BlockZoo: + def find_spec( + self, + name, + path = None, + target = None, + ): + # Whole package absent, so ModuleNotFoundError.name is the top-level 'unsloth_zoo'. + if name == "unsloth_zoo" or name.startswith("unsloth_zoo."): + raise ModuleNotFoundError("No module named 'unsloth_zoo'", name = "unsloth_zoo") + return None + + finder = _BlockZoo() + saved = { + k: v + for k, v in list(sys.modules.items()) + if k == "unsloth_zoo" or k.startswith("unsloth_zoo.") + } + for k in saved: + del sys.modules[k] + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + # Boots without raising and exposes the stub API. + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.get_hf_download_state(["x"]) is None + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + sys.modules.update(saved) + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim + + +def test_degrades_when_shared_helper_import_raises_importerror(): + """unsloth_zoo can be installed yet fail to import when torch is missing (llama.cpp/GGUF-only + Studio), raising ImportError not ModuleNotFoundError. The shim must degrade for that too.""" + import importlib + + class _BlockWithImportError: + def find_spec( + self, + name, + path = None, + target = None, + ): + if name == "unsloth_zoo.hf_xet_fallback": + # Mirror a torch-less install: a plain ImportError with no .name. + raise ImportError("Unsloth: Pytorch is not installed.") + return None + + finder = _BlockWithImportError() + saved_shared = sys.modules.pop("unsloth_zoo.hf_xet_fallback", None) + saved_zoo = sys.modules.pop("unsloth_zoo", None) + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + assert issubclass(degraded.DownloadStallError, RuntimeError) + assert degraded.get_hf_download_state(["x"]) is None + event = degraded.start_watchdog(repo_ids = ["x"], on_stall = lambda m: None) + assert hasattr(event, "set") and not event.is_set() + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + if saved_shared is not None: + sys.modules["unsloth_zoo.hf_xet_fallback"] = saved_shared + if saved_zoo is not None: + sys.modules["unsloth_zoo"] = saved_zoo + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim + + +def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): + """GPU detection in unsloth_zoo's __init__ raises NotImplementedError on a GPU-less host. The shim + retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails.""" + import importlib import os - return os.environ.get("PATH", "") - - -def test_disable_xet_constant_set_in_fresh_interpreter(): - code = ( - "from huggingface_hub import constants as c; " - "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is True else 17)" - ) - proc = subprocess.run( - [sys.executable, "-c", code], - env = {"HF_HUB_DISABLE_XET": "1", "PATH": _safe_path()}, - capture_output = True, - text = True, - ) - assert proc.returncode == 0, ( - f"HF_HUB_DISABLE_XET=1 did not set constants.HF_HUB_DISABLE_XET=True " - f"(rc={proc.returncode}): {proc.stderr}" - ) - -def test_default_leaves_xet_enabled(): - code = ( - "from huggingface_hub import constants as c; " - "import sys; sys.exit(0 if c.HF_HUB_DISABLE_XET is False else 17)" - ) - proc = subprocess.run( - [sys.executable, "-c", code], - env = {"PATH": _safe_path()}, # no HF_HUB_DISABLE_XET - capture_output = True, - text = True, - ) - assert proc.returncode == 0, ( - f"without the env var, constants.HF_HUB_DISABLE_XET was not False " - f"(rc={proc.returncode}): {proc.stderr}" - ) + monkeypatch.delenv("UNSLOTH_ZOO_DISABLE_GPU_INIT", raising = False) + seen_env = [] + + class _GpuGatedBlocker: + def find_spec( + self, + name, + path = None, + target = None, + ): + # Crash is in unsloth_zoo's __init__, so intercept "unsloth_zoo" itself (the parent). + if name == "unsloth_zoo": + # Record the env each attempt sees; raise the no-GPU error both times so the shim + # degrades. + seen_env.append(os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT")) + raise NotImplementedError("Unsloth cannot find any torch accelerator") + return None + + finder = _GpuGatedBlocker() + saved = { + k: v + for k, v in list(sys.modules.items()) + if k == "unsloth_zoo" or k.startswith("unsloth_zoo.") + } + for k in saved: + del sys.modules[k] + saved_shim = sys.modules.pop("utils.hf_xet_fallback", None) + sys.meta_path.insert(0, finder) + try: + degraded = importlib.import_module("utils.hf_xet_fallback") + # First attempt without the light env, then a retry with it set. + assert seen_env == [None, "1"], seen_env + # Both attempts raised -> Studio still boots in degraded mode. + assert issubclass(degraded.DownloadStallError, RuntimeError) + # The env override must not leak past the import. + assert os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") is None + finally: + sys.meta_path.remove(finder) + sys.modules.pop("utils.hf_xet_fallback", None) + sys.modules.update(saved) + if saved_shim is not None: + sys.modules["utils.hf_xet_fallback"] = saved_shim diff --git a/studio/backend/tests/test_model_update_robustness.py b/studio/backend/tests/test_model_update_robustness.py index 9cf2a62c39..300eb587b3 100644 --- a/studio/backend/tests/test_model_update_robustness.py +++ b/studio/backend/tests/test_model_update_robustness.py @@ -5,8 +5,8 @@ Covers: * GGUF variant listing computes update_available from the already-fetched sibling metadata instead of a second Hub call. - * hf_hub_download_with_xet_fallback(force_download=True) bypasses the - try_to_load_from_cache cache-first early-return. + * hf_hub_download_with_xet_fallback forwards force_download through the shim to the + shared unsloth_zoo helper (which owns the cache-first early-return and its bypass). The cache "Update" action now runs through the download manager as a normal managed download (so it shows in the Downloads panel with progress + cancel), @@ -341,44 +341,26 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path): # ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ─── -def test_force_download_bypasses_cache_first_early_return(monkeypatch): - """force_download=True skips the try_to_load_from_cache early-return and - proceeds to the real download path; force_download=False returns the cached - path without ever attempting a download (X2/F2).""" - import huggingface_hub as hf +def test_force_download_is_forwarded_through_the_shim(monkeypatch): + """The shim's contract is to forward force_download unchanged to the shared helper (which owns the + cache-first early-return and bypass). Verify both False and True reach it (X2/F2).""" import utils.hf_xet_fallback as X - cached_path = "/cache/blob/cached.gguf" + seen = [] - # Pretend the blob IS cached on disk (try_to_load_from_cache is imported - # inside the function from huggingface_hub, and os.path.exists must agree). - monkeypatch.setattr(hf, "try_to_load_from_cache", lambda *a, **k: cached_path, raising = False) - monkeypatch.setattr(X.os.path, "exists", lambda p: True, raising = False) + def fake_shared(repo_id, filename, token, **kwargs): + seen.append(kwargs.get("force_download")) + return "/downloaded/path" - attempts = [] + monkeypatch.setattr(X, "_shared_hf_hub_download_with_xet_fallback", fake_shared, raising = True) - def fake_attempt(repo_id, filename, token, **kwargs): - attempts.append( - {"repo_id": repo_id, "filename": filename, "force": kwargs.get("force_download")} - ) - return ("ok", "/freshly/downloaded/path") - - monkeypatch.setattr(X, "_run_download_attempt", fake_attempt, raising = True) - - # force_download=False: cache-first early-return, no download attempt. - out = X.hf_hub_download_with_xet_fallback( + X.hf_hub_download_with_xet_fallback( "unsloth/repo", "model.gguf", token = None, force_download = False ) - assert out == cached_path - assert attempts == [] # never reached the real download - - # force_download=True: bypass the early-return, run the real download. - out2 = X.hf_hub_download_with_xet_fallback( + X.hf_hub_download_with_xet_fallback( "unsloth/repo", "model.gguf", token = None, force_download = True ) - assert out2 == "/freshly/downloaded/path" - assert len(attempts) == 1 - assert attempts[0]["force"] is True + assert seen == [False, True] # the shim forwards force_download to the shared helper unchanged # ── multi-revision GGUF blob comparison and update reclaim ── diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py index 15961ac03a..2dd2247396 100644 --- a/studio/backend/utils/hf_xet_fallback.py +++ b/studio/backend/utils/hf_xet_fallback.py @@ -1,341 +1,204 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""Xet-primary HF downloads with an automatic HTTP fallback on a no-progress stall. +"""Studio shim over the shared ``unsloth_zoo.hf_xet_fallback`` Xet -> HTTP stall fallback. -Xet (``hf_xet``) is the fast default but can hang with no progress and no -exception, and a blocked native thread cannot be killed. Keep Xet primary; fall -back to plain HTTP only when the parent observes a stall. ``HF_HUB_DISABLE_XET`` -is read at import time, so the fallback runs in a fresh ``spawn`` child (not a -thread) that sets the env before importing ``huggingface_hub``. Cached files -short-circuit with no child; deterministic errors (401/403/404/disk-full) and -cancellation propagate without a fallback. Mirrors the safetensors inference -recovery in core/inference/{orchestrator,worker}.py. +Re-exports the shared API and injects Studio's marker-aware cache purge +(``prepare_cache_for_transport``) so the download manager keeps its ``.transport`` +marker semantics on the HTTP retry. """ from __future__ import annotations -import multiprocessing as mp -import os -import queue -import signal -import sys import threading -import time from typing import Any, Callable, Optional -from loggers import get_logger - -logger = get_logger(__name__) - -_CTX = mp.get_context("spawn") - -# Defaults match the existing inference watchdog and hub shutdown deadline. -DEFAULT_HEARTBEAT_INTERVAL = 30.0 -DEFAULT_STALL_TIMEOUT = 180.0 -DEFAULT_GRACE_PERIOD = 10.0 -_POLL_INTERVAL = 0.5 - - -class DownloadStallError(RuntimeError): - """Raised when no download progress is observed for too long. - - Canonical home; orchestrator.py re-imports it so all paths share one type. - """ - - -def child_should_disable_xet(config: dict) -> bool: - """Single source of truth for the per-worker Xet env flip.""" - return bool(config.get("disable_xet")) - - -def get_hf_download_state( - repo_ids: Optional[list[str]] = None, *, repo_type: str = "model" -) -> Optional[tuple[int, bool]]: - """Return ``(total_on_disk_bytes, has_incomplete)`` for the active HF cache. - - Sparse-aware (st_blocks based) so a sparse Xet/``hf_transfer`` ``.incomplete`` - is not mistaken for full-size progress. ``None`` means the state could not be - measured, so callers skip stall logic for that tick. - """ +_shared_import_error = None +try: + import unsloth_zoo.hf_xet_fallback as _shared + _shared_available = True +except Exception as _exc: # noqa: BLE001 - any import failure must degrade, not crash + # unsloth_zoo's __init__ runs torch/GPU detection, which raises on a torch-less/GPU-less Studio + # host. The download helper needs none of it, so retry via the light UNSLOTH_ZOO_DISABLE_GPU_INIT + # path before giving up. + _shared_import_error = _exc + import os as _os + + _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" try: - from hub.utils.hf_cache_state import ( - blob_bytes_present, - has_active_incomplete_blobs, - hf_cache_root, - iter_active_repo_cache_dirs, - ) - - if hf_cache_root() is None: - return (0, False) - - total = 0 - has_incomplete = False - for repo_id in repo_ids or []: - # Skip local paths: HF IDs never start with / . ~ or contain "\". - if not repo_id or repo_id.startswith(("/", ".", "~")) or "\\" in repo_id: - continue - for entry in iter_active_repo_cache_dirs(repo_type, repo_id): - blobs_dir = entry / "blobs" - if not blobs_dir.is_dir(): - continue - for blob in blobs_dir.iterdir(): - try: - if blob.is_file(): - total += blob_bytes_present(blob) - except OSError: - pass - if has_active_incomplete_blobs(repo_type, repo_id): - has_incomplete = True - return (total, has_incomplete) - except Exception as e: - logger.debug("Failed to determine HF download state: %s", e) - return None - - -def start_watchdog( - *, - repo_ids: list[str], - on_stall: Callable[[str], None], - repo_type: str = "model", - interval: float = DEFAULT_HEARTBEAT_INTERVAL, - stall_timeout: float = DEFAULT_STALL_TIMEOUT, - xet_disabled: bool = False, - on_heartbeat: Optional[Callable[[str], None]] = None, -) -> threading.Event: - """Start a daemon thread that fires ``on_stall(message)`` exactly once iff a - ``*.incomplete`` is present AND the on-disk size is unchanged for - *stall_timeout* seconds. The timer resets while no ``*.incomplete`` exists, so - post-download init is never misread as a stall. Returns a stop event the - caller sets when the download phase ends. - """ - stop = threading.Event() - transport = "https" if xet_disabled else "xet" - fired = False - - def _beat() -> None: - nonlocal fired - state = get_hf_download_state(repo_ids, repo_type = repo_type) - last_size = state[0] if state is not None else 0 - last_change = time.monotonic() - - while not stop.wait(interval): - state = get_hf_download_state(repo_ids, repo_type = repo_type) - now = time.monotonic() + import unsloth_zoo.hf_xet_fallback as _shared + _shared_available = True + _shared_import_error = None + except Exception as _exc2: # noqa: BLE001 - degrade so Studio still boots with plain HF downloads + _shared_import_error = _exc2 + _shared_available = False + finally: + if _prev_gpu_init is None: + _os.environ.pop("UNSLOTH_ZOO_DISABLE_GPU_INIT", None) + else: + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = _prev_gpu_init + +if _shared_available: + # Bind by assignment so each public name shares one module-level binding with the degraded branch. + DEFAULT_GRACE_PERIOD = _shared.DEFAULT_GRACE_PERIOD + DEFAULT_HEARTBEAT_INTERVAL = _shared.DEFAULT_HEARTBEAT_INTERVAL + DEFAULT_STALL_TIMEOUT = _shared.DEFAULT_STALL_TIMEOUT + DownloadStallError = _shared.DownloadStallError + child_should_disable_xet = _shared.child_should_disable_xet + get_hf_download_state = _shared.get_hf_download_state + start_watchdog = _shared.start_watchdog + _shared_hf_hub_download_with_xet_fallback = _shared.hf_hub_download_with_xet_fallback + _shared_snapshot_download_with_xet_fallback = _shared.snapshot_download_with_xet_fallback +else: + # Degrade instead of crashing Studio: plain HF downloads, stall watchdog disabled. Thin stubs, + # not a second copy of the orchestration; recovery returns once unsloth_zoo is upgraded. + import logging as _logging + + _logging.getLogger(__name__).warning( + "unsloth_zoo.hf_xet_fallback unavailable (%s); the Xet stall watchdog is " + "disabled. Install/upgrade unsloth_zoo (and its torch dependency) to " + "re-enable automatic Xet -> HTTP download recovery.", + _shared_import_error, + ) - if state is None: - if on_heartbeat is not None: + DEFAULT_HEARTBEAT_INTERVAL = 30.0 + DEFAULT_STALL_TIMEOUT = 180.0 + DEFAULT_GRACE_PERIOD = 10.0 + + class DownloadStallError(RuntimeError): + """Stub mirror so callers' ``except`` clauses resolve; never raised in degraded mode.""" + + def child_should_disable_xet(config: dict) -> bool: + return bool(config.get("disable_xet")) + + def get_hf_download_state(*args: Any, **kwargs: Any) -> None: + return None # unmeasurable -> the (absent) watchdog never fires + + def start_watchdog( + *, + on_heartbeat: "Optional[Callable[[str], None]]" = None, + interval: float = DEFAULT_HEARTBEAT_INTERVAL, + xet_disabled: bool = False, + **kwargs: Any, + ) -> "threading.Event": + # No stall detection, but keep emitting heartbeats so the orchestrator's inactivity deadline + # is not tripped during a long download. + stop = threading.Event() + if on_heartbeat is None: + return stop + transport = "https" if xet_disabled else "xet" + + def _beat() -> None: + while not stop.wait(interval): + try: on_heartbeat(f"Downloading ({transport} transport)...") - continue - - current_size, has_incomplete = state - if current_size != last_size: - last_size = current_size - last_change = now - - # Reset unless .incomplete confirms an active download, so model init - # and lock waits are not counted as a stall. - if not has_incomplete: - last_change = now - elif now - last_change >= stall_timeout: - if not fired: - fired = True - on_stall( - f"Download appears stalled ({transport} transport) " - f"-- no progress for {int(now - last_change)}s" - ) - return - - if on_heartbeat is not None: - on_heartbeat(f"Downloading ({transport} transport)...") - - threading.Thread(target = _beat, daemon = True, name = "hf-xet-watchdog").start() - return stop - - -def _download_child_entry( - *, - repo_id: str, - filename: str, - token: Optional[str], - repo_type: str, - disable_xet: bool, - result_queue: Any, - force_download: bool = False, -) -> None: - """Spawn-child entrypoint: download one file and report the result. - - Top-level and picklable. Sets the Xet env BEFORE importing huggingface_hub, - forms its own process group so the parent can kill the whole transfer, and - never logs the token or signed URLs. - """ - # Die with Studio on Linux (this mp child gets no parent-set preexec_fn). - try: - from utils.process_lifetime import bind_current_process_to_parent_lifetime - bind_current_process_to_parent_lifetime() - except Exception: - pass - - if hasattr(os, "setsid"): - try: - os.setsid() - except OSError: - pass - - if disable_xet: - os.environ["HF_HUB_DISABLE_XET"] = "1" - # Keep the HTTP writer sequential and resumable (hf_transfer leaves sparse - # partials a sequential resume cannot safely continue). - os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" - os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") - - # Test-only fault injection (never set in production): stall the Xet attempt - # so the watchdog + HTTP fallback can be exercised against a real repo. - if not disable_xet and os.environ.get("UNSLOTH_HF_XET_FORCE_STALL") == "1": - import time as _t - try: - from huggingface_hub.constants import HF_HUB_CACHE - - blobs = os.path.join(HF_HUB_CACHE, "models--" + repo_id.replace("/", "--"), "blobs") - os.makedirs(blobs, exist_ok = True) - with open(os.path.join(blobs, "xet-force-stall.incomplete"), "wb") as fh: - fh.write(b"\0" * 4096) - except OSError: - pass - while True: - _t.sleep(3600) + except Exception: + pass + + threading.Thread( + target = _beat, + daemon = True, + name = "hf-xet-degraded-heartbeat", + ).start() + return stop + + def _degraded_cancelled(cancel_event: "Optional[threading.Event]") -> bool: + return cancel_event is not None and cancel_event.is_set() + + def _shared_hf_hub_download_with_xet_fallback( + repo_id: str, + filename: str, + token: Optional[str], + *, + repo_type: str = "model", + revision: Optional[str] = None, + cache_dir: Optional[str] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, + ) -> str: + # Keep the cancellation contract: do not start or return a download once cancelled. + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") - try: from huggingface_hub import hf_hub_download + path = hf_hub_download( repo_id = repo_id, filename = filename, - repo_type = repo_type, token = token, + repo_type = repo_type, + revision = revision, + cache_dir = cache_dir, force_download = force_download, ) - result_queue.put({"ok": True, "path": path}) - except BaseException as e: # noqa: BLE001 - report every failure to the parent - error = f"{type(e).__name__}: {e}" - try: - from hub.utils.download_registry import scrub_secrets - error = scrub_secrets(error, hf_token = token) - except Exception: - pass - result_queue.put({"ok": False, "error": error}) - - -def _terminate_process_group(proc: "mp.process.BaseProcess", grace_period: float) -> None: - """Kill *proc* and its whole process group (Xet may spawn helper procs). - - The child calls ``os.setsid()`` so its pgid equals its pid; signal via - ``os.killpg(pid, ...)`` -- NOT ``getpgid``, which before the child becomes a - group leader resolves to OUR group. SIGTERM, then SIGKILL after *grace_period*. - """ - pid = proc.pid - - def _signal_group(sig: int) -> None: - if pid is not None and hasattr(os, "killpg"): - try: - os.killpg(pid, sig) - return - except (ProcessLookupError, PermissionError, OSError): - pass - # Windows or pre-setsid: best effort on the single process. - try: - proc.terminate() if sig != getattr(signal, "SIGKILL", -9) else proc.kill() - except Exception: - pass - - _signal_group(getattr(signal, "SIGTERM", signal.SIGINT)) - proc.join(timeout = grace_period) - if proc.is_alive(): - _signal_group(getattr(signal, "SIGKILL", signal.SIGTERM)) - proc.join(timeout = 5.0) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + + def _shared_snapshot_download_with_xet_fallback( + repo_id: str, + *, + revision: Optional[str] = None, + token: Optional[str] = None, + repo_type: str = "model", + cache_dir: Optional[str] = None, + allow_patterns: Optional[Any] = None, + ignore_patterns: Optional[Any] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, + ) -> str: + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + from huggingface_hub import snapshot_download -def _run_download_attempt( - repo_id: str, - filename: str, - token: Optional[str], - *, - repo_type: str, - disable_xet: bool, - cancel_event: Optional[threading.Event], - stall_timeout: float, - interval: float, - grace_period: float, - on_status: Optional[Callable[[str], None]], - force_download: bool = False, -) -> tuple[str, Optional[str]]: - """Run one download in a spawn child supervised by the no-progress watchdog. - - Returns ``("ok", path)``, ``("stall", None)``, ``("cancelled", None)``, or - ``("error", message)``. This is the seam tests monkeypatch to avoid spawning. - """ - result_queue: Any = _CTX.Queue() - proc = _CTX.Process( - target = _download_child_entry, - kwargs = dict( + path = snapshot_download( repo_id = repo_id, - filename = filename, - token = token, repo_type = repo_type, - disable_xet = disable_xet, - result_queue = result_queue, + revision = revision, + token = token, + cache_dir = cache_dir, + allow_patterns = allow_patterns, + ignore_patterns = ignore_patterns, force_download = force_download, - ), - daemon = True, - ) - proc.start() - from utils.process_lifetime import adopt_pid - - adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep) - - stalled = threading.Event() - stop_watchdog = start_watchdog( - repo_ids = [repo_id], - on_stall = lambda msg: stalled.set(), - repo_type = repo_type, - interval = interval, - stall_timeout = stall_timeout, - xet_disabled = disable_xet, - on_heartbeat = on_status, - ) - - result: Optional[dict] = None - try: - while proc.is_alive(): - if cancel_event is not None and cancel_event.is_set(): - _terminate_process_group(proc, grace_period) - return ("cancelled", None) - if stalled.is_set(): - _terminate_process_group(proc, grace_period) - return ("stall", None) - try: - result = result_queue.get(timeout = _POLL_INTERVAL) - break - except queue.Empty: - continue - else: - # Process exited; drain any result it enqueued. - try: - result = result_queue.get_nowait() - except queue.Empty: - result = None - finally: - stop_watchdog.set() - proc.join(timeout = grace_period) - - if result is None: - return ( - "error", - f"download process for '{repo_id}/{filename}' exited " - f"(code={proc.exitcode}) without a result", ) - if result.get("ok"): - return ("ok", result["path"]) - return ("error", result.get("error") or "unknown download error") + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + + +__all__ = [ + "DEFAULT_GRACE_PERIOD", + "DEFAULT_HEARTBEAT_INTERVAL", + "DEFAULT_STALL_TIMEOUT", + "DownloadStallError", + "child_should_disable_xet", + "get_hf_download_state", + "start_watchdog", + "hf_hub_download_with_xet_fallback", + "snapshot_download_with_xet_fallback", +] + + +def _studio_prepare_for_http(repo_type: str, repo_id: str) -> None: + """Studio's marker-aware purge before an HTTP resume, keeping the download manager's ``.transport`` + accounting consistent (vs unsloth_zoo's generic default). Guarded: a purge failure is logged, + not fatal to the retry.""" + try: + from hub.utils.download_registry import prepare_cache_for_transport + prepare_cache_for_transport(repo_type, repo_id, "http") + except Exception as exc: + try: + from loggers import get_logger + get_logger(__name__).debug( + "Studio prepare_cache_for_transport failed for %s: %s", repo_id, exc + ) + except ModuleNotFoundError as logger_exc: + if logger_exc.name != "loggers": + raise def hf_hub_download_with_xet_fallback( @@ -345,83 +208,32 @@ def hf_hub_download_with_xet_fallback( *, cancel_event: Optional[threading.Event] = None, repo_type: str = "model", + revision: Optional[str] = None, stall_timeout: float = DEFAULT_STALL_TIMEOUT, interval: float = DEFAULT_HEARTBEAT_INTERVAL, grace_period: float = DEFAULT_GRACE_PERIOD, on_status: Optional[Callable[[str], None]] = None, force_download: bool = False, ) -> str: - """Download a single file with Xet primary and HTTP as a stall-only fallback. - - Returns the local cache path. Raises ``RuntimeError("Cancelled")`` if - *cancel_event* is set, re-raises a deterministic child error unchanged (no - fallback), and raises ``DownloadStallError`` only if BOTH transports stall. - - When *force_download* is True the cache-first early-return is skipped and the - flag is threaded to ``hf_hub_download`` so a newer remote blob is re-fetched - even if an older blob is already cached. - """ - # Finalized blob already cached: return it with no child and no network. - # Skipped when force_download is set so an update re-fetches a newer blob. - if not force_download: - try: - from huggingface_hub import try_to_load_from_cache - cached = try_to_load_from_cache(repo_id, filename, repo_type = repo_type) - if isinstance(cached, str) and os.path.exists(cached): - return cached - except Exception as e: - logger.debug("Cached probe failed for %s/%s: %s", repo_id, filename, e) - - if cancel_event is not None and cancel_event.is_set(): - raise RuntimeError("Cancelled") - - disable_xet = False - for attempt in range(2): - if disable_xet: - # Purge a non-HTTP partial before resuming over HTTP: an HTTP resume - # over a sparse Xet/hf_transfer partial silently corrupts the blob. - try: - from hub.utils.download_registry import prepare_cache_for_transport - prepare_cache_for_transport(repo_type, repo_id, "http") - except Exception as e: - logger.debug("prepare_cache_for_transport failed for %s: %s", repo_id, e) - - kind, payload = _run_download_attempt( - repo_id, - filename, - token, - repo_type = repo_type, - disable_xet = disable_xet, - cancel_event = cancel_event, - stall_timeout = stall_timeout, - interval = interval, - grace_period = grace_period, - on_status = on_status, - force_download = force_download, - ) + """Single-file download via the shared fallback with Studio's marker-aware HTTP-retry prep. + ``force_download`` re-fetches a newer blob over a cached one (Studio's model-update path).""" + return _shared_hf_hub_download_with_xet_fallback( + repo_id, + filename, + token, + cancel_event = cancel_event, + repo_type = repo_type, + revision = revision, + stall_timeout = stall_timeout, + interval = interval, + grace_period = grace_period, + on_status = on_status, + force_download = force_download, + prepare_for_http_fn = _studio_prepare_for_http, + ) - if kind == "ok": - return payload # type: ignore[return-value] - if kind == "cancelled": - raise RuntimeError("Cancelled") - if kind == "error": - # Deterministic failure: the other transport would fail identically. - raise RuntimeError(payload) - # kind == "stall" - if attempt == 0 and not disable_xet: - logger.warning( - "Download stalled for '%s/%s' -- retrying with HF_HUB_DISABLE_XET=1", - repo_id, - filename, - ) - if on_status is not None: - on_status(f"{repo_id}/{filename}: Xet stalled, retrying over HTTP") - disable_xet = True - continue - raise DownloadStallError( - f"Download stalled for '{repo_id}/{filename}' even with " - f"HF_HUB_DISABLE_XET=1 -- check your network connection" - ) - # Unreachable: the loop either returns or raises on each attempt. - raise DownloadStallError(f"Download failed for '{repo_id}/{filename}'") +def snapshot_download_with_xet_fallback(repo_id: str, **kwargs: Any) -> str: + """Whole-repo download via the shared fallback with Studio's marker-aware HTTP-retry prep.""" + kwargs.setdefault("prepare_for_http_fn", _studio_prepare_for_http) + return _shared_snapshot_download_with_xet_fallback(repo_id, **kwargs) diff --git a/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx b/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx index d00c812325..db7628777a 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/model-update-action.tsx @@ -42,10 +42,8 @@ export function ModelUpdateAction({ }: ModelUpdateActionProps) { const [open, setOpen] = useState(false); - // The update is a managed download (it surfaces in the global Downloads panel - // with progress + cancel). When this exact repo+variant finishes, refresh the - // caller so the "update available" cue clears once the new revision is on - // disk. A ref keeps the subscription stable across renders without resubscribing. + // Refresh the caller when this repo+variant's download finishes so the "update available" cue + // clears. A ref keeps the subscription stable across renders. const onUpdatedRef = useRef(onUpdated); onUpdatedRef.current = onUpdated; useEffect(() => { @@ -60,9 +58,8 @@ export function ModelUpdateAction({ }, [repoId, variant]); const handleConfirm = useCallback(() => { - // Start the background re-download and close the dialog immediately; the - // Downloads panel owns progress + cancel from here. Only a failure to START - // surfaces a toast — a failed download reports itself in the panel. + // Start the re-download and close the dialog; the Downloads panel owns progress + cancel. + // Only a failure to START toasts (a failed download shows in the panel). void Promise.resolve() .then(onConfirm) .catch((err) => { diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index e11a08f8ab..d3ae638e5c 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -1245,11 +1245,8 @@ export function HubModelPicker({ onEject?: () => void; }) { const gpu = useGpuInfo(); - // The currently-loaded/running model id. We read params.checkpoint from the - // runtime store (backend-mirrored from /api/inference/status.active_model, see - // chat-runtime-store) rather than the dropdown `isSelected` highlight (which is - // just `value === repo_id` and can reflect a staged, not-yet-loaded pick). Used - // to disable the cached-row update action for the model that's live in memory. + // Live model id from the runtime store (backend-mirrored active_model), not the dropdown + // highlight which can be a staged pick. Disables the update action for it. const loadedModelId = useChatRuntimeStore((s) => s.params.checkpoint); // Last-loaded timestamps power the "Recent" sort (vs "Downloaded" = file date). const loadTimes = useModelLoadTimes(value); @@ -1584,11 +1581,8 @@ export function HubModelPicker({ refreshLocalModelsList(); }, [hfToken, refreshLocalModelsList]); - // Updates run as MANAGED downloads (they show in the global Downloads panel - // with manifest-based progress + a working Cancel), instead of a blocking - // call. The worker re-resolves `main` and pulls only changed blobs, so the - // cached copy stays usable until the new revision lands. The row's - // ModelUpdateAction refreshes the list when this repo+variant completes. + // Updates run as managed downloads (Downloads panel: progress + Cancel), not a blocking + // call. The worker pulls only changed blobs, so the cached copy stays usable until done. const startManagedUpdate = useCallback((repoId: string, variant: string, expectedBytes: number) => { return downloadManager .requestStart({ diff --git a/tests/test_prefetch_snapshot_scope.py b/tests/test_prefetch_snapshot_scope.py new file mode 100644 index 0000000000..c7ec4f2c34 --- /dev/null +++ b/tests/test_prefetch_snapshot_scope.py @@ -0,0 +1,916 @@ +# Unsloth Zoo - Utilities for Unsloth +# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved. +# +# This program is free software: you can redistribute it and/or modify +# it under the terms of the GNU Affero General Public License as published +# by the Free Software Foundation, either version 3 of the License, or +# (at your option) any later version. +# +# This program is distributed in the hope that it will be useful, +# but WITHOUT ANY WARRANTY; without even the implied warranty of +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +# GNU Affero General Public License for more details. +# +# You should have received a copy of the GNU Affero General Public License +# along with this program. If not, see . + +"""Pure-CPU, no-network unit tests for prefetch snapshot scoping in unsloth/models/_utils.py. + +maybe_prefetch_hf_snapshot warms the HF cache before the in-process load. The warm must cover at +least what the load reads (else the missing file falls to an unprotected in-process Xet fetch) but +not pull weights the load never reads. These tests lock the allow/ignore patterns each mode hands +snapshot_download_with_xet_fallback. The zoo downloader is monkeypatched to capture its kwargs. +""" + +import fnmatch +import sys +import types + +import pytest + +from unsloth.models import _utils as U + + +def _filter(names, allow_patterns, ignore_patterns): + """Mirror HF filter_repo_objects: keep on allow match (or None), drop on ignore match.""" + kept = [] + for name in names: + if allow_patterns is not None and not any(fnmatch.fnmatch(name, p) for p in allow_patterns): + continue + if ignore_patterns and any(fnmatch.fnmatch(name, p) for p in ignore_patterns): + continue + kept.append(name) + return kept + + +@pytest.fixture +def capture(monkeypatch): + """Run maybe_prefetch_hf_snapshot with a fake repo, capturing the patterns forwarded to a + fake injected zoo downloader (independent of the installed unsloth_zoo). Offline env cleared.""" + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + + state = {} + + def fake_download(repo_id, **kw): + state["repo_id"] = repo_id + state["allow_patterns"] = kw.get("allow_patterns") + state["ignore_patterns"] = kw.get("ignore_patterns") + state["variant"] = kw.get("variant") + return "/tmp/fake-snapshot" + + fake_module = types.ModuleType("unsloth_zoo.hf_xet_fallback") + fake_module.snapshot_download_with_xet_fallback = fake_download + fake_module.DownloadStallError = type("DownloadStallError", (RuntimeError,), {}) + monkeypatch.setitem(sys.modules, "unsloth_zoo.hf_xet_fallback", fake_module) + + # Neutralize the model_info network call by default; tests exercising format selection + # install their own. + import huggingface_hub + + class _NoNetworkApi: + def model_info(self, *a, **k): + raise RuntimeError("no network in test") + + monkeypatch.setattr(huggingface_hub, "HfApi", _NoNetworkApi) + + def run(**call_kwargs): + state.clear() + ok = U.maybe_prefetch_hf_snapshot("some-org/some-repo", **call_kwargs) + return ok, state + + return run + + +# Representative repo listing: root weights + aux, subdir, adapter, checkpoint, merged weights. +_SAMPLE_FILES = [ + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors", + "model.safetensors.index.json", + "pytorch_model.bin", + "fp16/model.safetensors", + "experimental/model-00001-of-00002.safetensors", + "checkpoint-500/model.safetensors", + "adapter_config.json", + "adapter_model.safetensors", +] + + +def test_weights_at_root_excludes_subdir_weights(capture): + """A root load ignores subdir weights (fp16/, experimental/, checkpoint-500/) but keeps root weights.""" + ok, st = capture(weights_at_root = True, use_safetensors = True) + assert ok is True + assert st["allow_patterns"] is None + ig = st["ignore_patterns"] + assert "*/*.safetensors" in ig and "*/*.bin" in ig + kept = _filter(_SAMPLE_FILES, st["allow_patterns"], ig) + assert "model-00001-of-00002.safetensors" in kept + assert "model.safetensors.index.json" in kept + assert "config.json" in kept + assert "fp16/model.safetensors" not in kept + assert "experimental/model-00001-of-00002.safetensors" not in kept + assert "checkpoint-500/model.safetensors" not in kept + + +def test_adapter_only_excludes_merged_weights(capture): + """An adapter warm keeps adapter files + root aux, not merged full-model weights.""" + ok, st = capture(adapter_only = True) + assert ok is True + assert st["ignore_patterns"] is None + allow = st["allow_patterns"] + assert "adapter_config.json" in allow and "adapter_model*" in allow + kept = _filter(_SAMPLE_FILES, allow, st["ignore_patterns"]) + assert "adapter_config.json" in kept + assert "adapter_model.safetensors" in kept + assert "config.json" in kept and "tokenizer.json" in kept + assert "model-00001-of-00002.safetensors" not in kept + assert "pytorch_model.bin" not in kept + assert "fp16/model.safetensors" not in kept + + +def test_adapter_only_warms_sharded_adapter(capture): + """A sharded adapter is still covered by the adapter_model* glob.""" + _, st = capture(adapter_only = True) + sharded = [ + "adapter_config.json", + "adapter_model-00001-of-00002.safetensors", + "adapter_model-00002-of-00002.safetensors", + "adapter_model.safetensors.index.json", + ] + kept = _filter(sharded, st["allow_patterns"], st["ignore_patterns"]) + assert set(kept) == set(sharded) + + +def test_tokenizer_only_warms_only_aux_files(capture): + """A tokenizer-only repo warms tokenizer/config/vocab files, never weights.""" + _, st = capture(tokenizer_only = True) + assert st["ignore_patterns"] is None + assert st["allow_patterns"] == list(U._ROOT_AUX_PREFETCH_PATTERNS) + kept = _filter(_SAMPLE_FILES, st["allow_patterns"], st["ignore_patterns"]) + assert "tokenizer.json" in kept and "config.json" in kept + assert "model-00001-of-00002.safetensors" not in kept + assert "adapter_model.safetensors" not in kept + + +def test_aux_warm_covers_arbitrary_remote_code_modules(capture): + """The aux warm must cover any *.py, since trust_remote_code auto_map names modules freely.""" + _, st = capture(tokenizer_only = True) + allow = st["allow_patterns"] + assert "*.py" in allow + remote_code = [ + "config.json", + "modeling.py", + "tokenization.py", + "my_custom_code.py", + "configuration_foo.py", + ] + kept = _filter(remote_code, allow, st["ignore_patterns"]) + for name in ("modeling.py", "tokenization.py", "my_custom_code.py", "configuration_foo.py"): + assert name in kept, name + + +def test_subfolder_warms_subfolder_plus_root_aux(capture): + """A subfolder load warms that subfolder's weights plus root aux; other subdirs/root weights skipped.""" + _, st = capture(subfolder = "fp16") + allow = st["allow_patterns"] + assert "fp16/*" in allow + assert all(p in allow for p in U._ROOT_AUX_PREFETCH_PATTERNS) + kept = _filter(_SAMPLE_FILES, allow, st["ignore_patterns"]) + assert "fp16/model.safetensors" in kept + assert "config.json" in kept + assert "experimental/model-00001-of-00002.safetensors" not in kept + + +def test_subfolder_takes_precedence_over_weights_at_root(capture): + """When a subfolder is requested the subfolder branch wins over weights_at_root.""" + _, st = capture(subfolder = "fp16", weights_at_root = True) + assert "fp16/*" in st["allow_patterns"] + kept = _filter(_SAMPLE_FILES, st["allow_patterns"], st["ignore_patterns"]) + assert "fp16/model.safetensors" in kept + + +def test_local_dir_is_not_warmed(capture, tmp_path): + """A local directory path skips the warm (returns False).""" + d = tmp_path / "local-model" + d.mkdir() + ok = U.maybe_prefetch_hf_snapshot(str(d), weights_at_root = True) + assert ok is False + + +def _install_fake_model_info(monkeypatch, filenames): + """Make HfApi().model_info(...).siblings report filenames, with no network.""" + import huggingface_hub + + class _Sib: + def __init__(self, name): + self.rfilename = name + + class _Info: + def __init__(self, names): + self.siblings = [_Sib(n) for n in names] + + class _Api: + def model_info(self, *a, **k): + return _Info(filenames) + + monkeypatch.setattr(huggingface_hub, "HfApi", _Api) + + +# ----- Finding P: variant-aware weight-format selection ----- + + +def test_variant_keeps_bin_when_only_default_safetensors(monkeypatch): + """A default model.safetensors must not prove a variant .bin redundant; without a variant it does.""" + _install_fake_model_info(monkeypatch, ["model.safetensors", "pytorch_model.fp16.bin"]) + ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" not in ig + ig_default = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" in ig_default + + +def test_variant_drops_bin_when_variant_safetensors_present(monkeypatch): + """A variant-matching safetensors makes the variant .bin redundant, so .bin is dropped.""" + _install_fake_model_info(monkeypatch, ["model.fp16.safetensors", "pytorch_model.fp16.bin"]) + ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" in ig + + +def test_no_variant_keeps_bin_when_only_variant_safetensors(monkeypatch): + """For a no-variant load, only a canonical safetensors (not a lone variant) makes .bin redundant.""" + _install_fake_model_info(monkeypatch, ["model.fp16.safetensors", "pytorch_model.bin"]) + ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" not in ig + _install_fake_model_info(monkeypatch, ["model.safetensors", "pytorch_model.bin"]) + ig2 = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" in ig2 + + +def test_variant_keeps_bin_for_noncanonical_sidecar(monkeypatch): + """A non-canonical variant sidecar must not prove the variant .bin redundant; a canonical one does.""" + _install_fake_model_info( + monkeypatch, ["consolidated.fp16.safetensors", "pytorch_model.fp16.bin"] + ) + ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" not in ig + _install_fake_model_info(monkeypatch, ["model.fp16.safetensors", "pytorch_model.fp16.bin"]) + ig2 = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" in ig2 + + +def test_is_canonical_model_weight_safetensors(): + """The canonical detector matches only non-variant model-weight safetensors names.""" + assert U._is_canonical_model_weight_safetensors("model.safetensors") is True + assert U._is_canonical_model_weight_safetensors("model-00001-of-00002.safetensors") is True + assert U._is_canonical_model_weight_safetensors("model.safetensors.index.json") is True + assert U._is_canonical_model_weight_safetensors("model.fp16.safetensors") is False + assert ( + U._is_canonical_model_weight_safetensors("model.fp16-00001-of-00002.safetensors") is False + ) + assert U._is_canonical_model_weight_safetensors("adapter_model.safetensors") is False + + +def test_st_prefetch_resolves_env_cache_and_runs_after_validation(): + """The ST prefetch must resolve SENTENCE_TRANSFORMERS_HOME and run after load-mode validation.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + tree = ast.parse(src) + prefetch_calls = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "maybe_prefetch_hf_snapshot" + ] + assert len(prefetch_calls) == 1, "expected exactly one ST prefetch call" + call = prefetch_calls[0] + # cache_dir kwarg resolves SENTENCE_TRANSFORMERS_HOME. + cache_dir_kw = next((kw for kw in call.keywords if kw.arg == "cache_dir"), None) + assert cache_dir_kw is not None, "ST prefetch must pass cache_dir" + assert "SENTENCE_TRANSFORMERS_HOME" in ast.dump( + cache_dir_kw.value + ), "ST prefetch cache_dir must resolve SENTENCE_TRANSFORMERS_HOME" + # Load-mode validation runs before the prefetch (fewer source lines = earlier). + val_lineno = src[: src.index("Can only load in 4bit or 8bit or 16bit")].count("\n") + assert val_lineno < call.lineno, "load-mode validation must precede the ST prefetch" + + +def test_st_cache_resolutions_honor_explicit_hf_cache_dir(): + """Every ST cache resolution falling back to SENTENCE_TRANSFORMERS_HOME must first honor an explicit HF cache_dir.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + resolutions = [ + kw + for kw in ast.walk(tree) + if isinstance(kw, ast.keyword) + and kw.arg == "cache_dir" + and "SENTENCE_TRANSFORMERS_HOME" in ast.dump(kw.value) + ] + assert resolutions, "expected cache_dir resolutions referencing SENTENCE_TRANSFORMERS_HOME" + for kw in resolutions: + assert "'cache_dir'" in ast.dump( + kw.value + ), "an ST cache_dir resolution must read an explicit kwargs.get('cache_dir') first" + + +def test_st_native_loads_map_hf_cache_dir_to_cache_folder(): + """Native SentenceTransformer loads take cache_folder, so an explicit HF cache_dir must be mapped onto it.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + tree = ast.parse(src) + # Every native SentenceTransformer(...) forwarding cache_folder must read cache_dir. + st_calls = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "SentenceTransformer" + ] + cache_folder_kws = [kw for call in st_calls for kw in call.keywords if kw.arg == "cache_folder"] + assert cache_folder_kws, "expected a native SentenceTransformer call forwarding cache_folder" + for kw in cache_folder_kws: + assert "'cache_dir'" in ast.dump( + kw.value + ), "a native SentenceTransformer cache_folder must map the explicit HF cache_dir first" + # for_inference feeds cache_folder via st_kwargs; both native branches map cache_dir -> cache_folder. + normalized = "".join(src.split()) + assert ( + 'st_kwargs["cache_folder"]=' in normalized + ), "for_inference must set st_kwargs cache_folder" + assert ( + normalized.count('kwargs.get("cache_dir")orkwargs.get("cache_folder")') >= 2 + ), "both native ST branches (for_inference, fast-encoder) must map cache_dir -> cache_folder" + + +def test_vision_warms_vllm_tokenizer_after_remap(): + """On the vLLM path the tokenizer warm is deferred until after the fast_inference_setup remap.""" + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "vision.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + guard = "if _vllm_owns_weights and isinstance(tokenizer_name" + assert guard in src, "expected a vLLM-gated tokenizer warm" + assert src.index(guard) > src.index( + "fast_inference_setup(" + ), "the vLLM tokenizer warm must run after the fast_inference_setup remap" + + +def test_diffusion_forwards_variant_to_real_load(): + """FastDiffusionModel must forward variant to the real model_cls.from_pretrained load, not just the prefetch.""" + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "diffusion.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + assert ( + 'load_kwargs["variant"] = kwargs["variant"]' in src + ), "the diffusion load must forward variant to model_cls.from_pretrained" + + +def test_vision_prefetch_runs_after_load_mode_validation(): + """The FastBaseModel (vision) prefetch must run after the load-mode validation.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "vision.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + tree = ast.parse(src) + prefetch_calls = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "maybe_prefetch_hf_snapshot" + ] + assert prefetch_calls, "expected a vision prefetch call" + first_prefetch = min(call.lineno for call in prefetch_calls) + val_lineno = src[: src.index("Can only load in 4bit or 8bit or 16bit")].count("\n") + assert val_lineno < first_prefetch, "load-mode validation must precede the vision prefetch" + + +def test_llama_prefetch_skips_only_real_vllm_loads(): + """The llama prefetch's fast_inference skip must be gated on num_labels is None (a classification load still downloads).""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "llama.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + gated = False + for n in ast.walk(tree): + if not ( + isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "maybe_prefetch_hf_snapshot" + ): + continue + fi_kw = next((kw for kw in n.keywords if kw.arg == "fast_inference"), None) + if fi_kw is None: + continue + dumped = ast.dump(fi_kw.value) + if "fast_inference" in dumped and "num_labels" in dumped: + gated = True + assert gated, "llama prefetch fast_inference must be gated on num_labels is None" + + +def test_st_fallback_module_loads_resolve_env_cache(): + """Fallback module loads deriving cache_dir from cache_folder must also fall back to SENTENCE_TRANSFORMERS_HOME.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + src = f.read() + tree = ast.parse(src) + + # Fallback sites (cache_dir derived from cache_folder) must resolve SENTENCE_TRANSFORMERS_HOME. + checked = 0 + for node in ast.walk(tree): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)): + continue + if node.func.attr not in ("_module_path", "_load_modules"): + continue + cache_dir_kw = next((kw for kw in node.keywords if kw.arg == "cache_dir"), None) + if cache_dir_kw is None: + continue + dumped = ast.dump(cache_dir_kw.value) + if "cache_folder" not in dumped: + continue # internal pass-through, not a resolution site + checked += 1 + assert ( + "SENTENCE_TRANSFORMERS_HOME" in dumped + ), f"{node.func.attr} cache_dir resolves cache_folder but not SENTENCE_TRANSFORMERS_HOME" + assert ( + checked >= 2 + ), "expected the fallback _module_path and _load_modules calls to resolve the env cache" + + +def test_st_fallback_module_loads_forward_revision(): + """The fallback module loads must forward revision so module files match the revision-pinned weights. + Guards: (a) helpers accept revision, (b) every download primitive forwards it, (c) _load_modules + threads it into internal calls, (d) the from_pretrained fallback sites forward it.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + + funcs = { + n.name: n + for n in ast.walk(tree) + if isinstance(n, ast.FunctionDef) + and n.name in ("_module_path", "_read_pooling_mode", "_load_modules") + } + assert set(funcs) == {"_module_path", "_read_pooling_mode", "_load_modules"} + + # (a) each helper takes a revision parameter. + for name, fn in funcs.items(): + arg_names = {a.arg for a in fn.args.args + fn.args.kwonlyargs} + assert "revision" in arg_names, f"{name} must accept a revision argument" + + # (b) every download primitive inside the helpers forwards revision. + downloads = 0 + for name, fn in funcs.items(): + for node in ast.walk(fn): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Name)): + continue + if node.func.id not in ("hf_hub_download", "load_dir_path"): + continue + downloads += 1 + assert any( + kw.arg == "revision" for kw in node.keywords + ), f"{node.func.id} in {name} must forward revision" + assert downloads >= 3, "expected the module-download primitives to be revision-guarded" + + # (c) _load_modules threads revision into its internal _module_path / _read_pooling_mode calls. + internal = 0 + for node in ast.walk(funcs["_load_modules"]): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)): + continue + if node.func.attr not in ("_module_path", "_read_pooling_mode"): + continue + internal += 1 + assert any( + kw.arg == "revision" for kw in node.keywords + ), f"_load_modules must forward revision to {node.func.attr}" + assert internal >= 2, "expected _load_modules to call _module_path and _read_pooling_mode" + + # (d) the from_pretrained fallback _module_path / _load_modules sites forward revision. + checked = 0 + for node in ast.walk(tree): + if not (isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute)): + continue + if node.func.attr not in ("_module_path", "_load_modules"): + continue + cache_dir_kw = next((kw for kw in node.keywords if kw.arg == "cache_dir"), None) + if cache_dir_kw is None or "cache_folder" not in ast.dump(cache_dir_kw.value): + continue # internal pass-through, not a fallback site + checked += 1 + rev_kw = next((kw for kw in node.keywords if kw.arg == "revision"), None) + assert rev_kw is not None and "revision" in ast.dump( + rev_kw.value + ), f"{node.func.attr} fallback call must forward revision" + assert ( + checked >= 2 + ), "expected the fallback _module_path and _load_modules calls to forward revision" + + +def test_st_fallback_model_load_resolves_env_cache(): + """from_pretrained must resolve the warmed ST cache into kwargs['cache_dir'] before the FastModel weight load.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + + def _resolves_st_cache(value_node): + # Resolution may be inline or in the assignment to an intermediate variable the value references. + dumped = ast.dump(value_node) + if "cache_folder" in dumped and "SENTENCE_TRANSFORMERS_HOME" in dumped: + return True + if isinstance(value_node, ast.Name): + for n in ast.walk(tree): + if isinstance(n, ast.Assign) and any( + isinstance(t, ast.Name) and t.id == value_node.id for t in n.targets + ): + d = ast.dump(n.value) + if "cache_folder" in d and "SENTENCE_TRANSFORMERS_HOME" in d: + return True + return False + + resolved_lines = [] + for node in ast.walk(tree): + if not isinstance(node, ast.Assign): + continue + for tgt in node.targets: + if ( + isinstance(tgt, ast.Subscript) + and isinstance(tgt.value, ast.Name) + and tgt.value.id == "kwargs" + and isinstance(tgt.slice, ast.Constant) + and tgt.slice.value == "cache_dir" + and _resolves_st_cache(node.value) + ): + resolved_lines.append(node.lineno) + assert resolved_lines, "from_pretrained must resolve the ST cache into kwargs['cache_dir']" + + fastmodel_calls = [ + n.lineno + for n in ast.walk(tree) + if isinstance(n, ast.Call) + and isinstance(n.func, ast.Attribute) + and n.func.attr == "from_pretrained" + and isinstance(n.func.value, ast.Name) + and n.func.value.id == "FastModel" + ] + assert fastmodel_calls, "expected a FastModel.from_pretrained call" + assert min(resolved_lines) < min( + fastmodel_calls + ), "kwargs['cache_dir'] must be resolved before the fallback FastModel weight load" + + +def test_canonical_variant_model_weight_matches_transformers_names(): + """The variant safetensors detector matches only canonical variant names, rejecting sidecars and wrong variants.""" + f = U._is_canonical_variant_model_weight_safetensors + assert f("model.fp16.safetensors", "fp16") is True + assert f("model.fp16-00001-of-00002.safetensors", "fp16") is True + assert f("model-00001-of-00002.fp16.safetensors", "fp16") is True + assert f("model.safetensors.index.fp16.json", "fp16") is True + assert f("consolidated.fp16.safetensors", "fp16") is False + assert f("model.safetensors", "fp16") is False + assert f("model-00001-of-00002.safetensors", "fp16") is False + assert f("model.bf16.safetensors", "fp16") is False + + +def test_variant_is_forwarded_to_downloader(capture): + """maybe_prefetch_hf_snapshot must forward variant to the downloader (absent a variant, nothing is forwarded).""" + _, st = capture(weights_at_root = True, use_safetensors = True, variant = "fp16") + assert st["variant"] == "fp16" + _, st = capture(weights_at_root = True, use_safetensors = True) + assert st["variant"] is None + + +def test_variant_drops_bin_for_sharded_variant_safetensors(monkeypatch): + """A sharded variant safetensors is recognized, so its redundant variant .bin is dropped.""" + _install_fake_model_info( + monkeypatch, + [ + "model.fp16-00001-of-00002.safetensors", + "model.fp16-00002-of-00002.safetensors", + "pytorch_model.fp16-00001-of-00002.bin", + ], + ) + ig = U._prefetch_ignore_patterns("org/repo", variant = "fp16", weights_at_root = True) + assert "*.bin" in ig + + +def test_tokenizer_only_warms_extra_vocab_files(capture): + """tokenizer_only must warm SentencePiece / vocab / processor files, including a named jinja template.""" + _, st = capture(tokenizer_only = True) + allow = st["allow_patterns"] + for name in ( + "spm.model", + "normalizer.json", + "video_preprocessor_config.json", + "tokenizer.model.v3", + ): + assert name in allow, name + sample = [ + "spm.model", + "normalizer.json", + "video_preprocessor_config.json", + "tokenizer.model.v3", + "additional_chat_templates/custom.jinja", + ] + kept = _filter(sample, allow, st["ignore_patterns"]) + assert set(kept) == set(sample) + + +def test_format_probe_runs_even_when_config_cached(capture, monkeypatch): + """A cached config.json must not skip the weight-format probe; model_info still drops the redundant .bin.""" + import huggingface_hub + + # Pretend config.json is cached (the AutoConfig side effect); this must not gate the probe. + monkeypatch.setattr( + huggingface_hub, "try_to_load_from_cache", lambda *a, **k: "/cache/config.json" + ) + _install_fake_model_info(monkeypatch, ["model.safetensors", "pytorch_model.bin"]) + _, st = capture(weights_at_root = True) + ig = st["ignore_patterns"] or [] + assert "*.bin" in ig + + +def test_optimizer_safetensors_does_not_drop_bin(monkeypatch): + """An optimizer.safetensors sidecar must not count as model safetensors, so the real .bin weights are kept.""" + _install_fake_model_info(monkeypatch, ["pytorch_model.bin", "optimizer.safetensors"]) + ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" not in ig + + +def test_model_safetensors_still_drops_bin(monkeypatch): + """Control for the optimizer case: a real model.safetensors next to pytorch_model.bin still drops the .bin.""" + _install_fake_model_info( + monkeypatch, ["model.safetensors", "pytorch_model.bin", "optimizer.safetensors"] + ) + ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" in ig + + +def test_whole_multi_component_snapshot_keeps_subdir_bin(monkeypatch): + """A whole multi-component snapshot must not drop *.bin (it would strip a subdir module's weight); a root load still does.""" + _install_fake_model_info(monkeypatch, ["model.safetensors", "1_Dense/pytorch_model.bin"]) + ig = U._prefetch_ignore_patterns("org/repo", weights_at_root = False) + assert "*.bin" not in ig + ig_root = U._prefetch_ignore_patterns("org/repo", weights_at_root = True) + assert "*.bin" in ig_root + + +def test_is_model_weight_safetensors_classification(): + """Real model weights count; adapter / trainer-state sidecars do not.""" + assert U._is_model_weight_safetensors("model.safetensors") is True + assert U._is_model_weight_safetensors("model-00001-of-00002.safetensors") is True + assert U._is_model_weight_safetensors("model.safetensors.index.json") is True + assert U._is_model_weight_safetensors("consolidated.safetensors") is True + assert U._is_model_weight_safetensors("adapter_model.safetensors") is False + assert U._is_model_weight_safetensors("optimizer.safetensors") is False + assert U._is_model_weight_safetensors("scheduler.safetensors") is False + assert U._is_model_weight_safetensors("rng_state_0.safetensors") is False + + +def test_tokenizer_only_warms_slow_sentencepiece_vocab(capture): + """tokenizer_only must warm the slow-tokenizer SentencePiece / BPE vocab files AutoTokenizer fetches first.""" + _, st = capture(tokenizer_only = True) + allow = st["allow_patterns"] + for name in ( + "sentencepiece.bpe.model", + "source.spm", + "target.spm", + "bpe.codes", + "vocab.bpe", + "sentencepiece.model", + "vocab-src.json", + "vocab-tgt.json", + ): + assert name in allow, name + + +def test_adapter_safetensors_check_scoped_to_root(monkeypatch): + """_adapter_repo_has_safetensors must only count a root adapter_model*.safetensors, not a subdir one.""" + import huggingface_hub + + class _Sib: + def __init__(self, name): + self.rfilename = name + + class _Api: + def __init__(self, names): + self._names = names + + def model_info(self, *a, **k): + return type("MI", (), {"siblings": [_Sib(n) for n in self._names]})() + + # Subdir safetensors only -> not reported present. + monkeypatch.setattr( + huggingface_hub, + "HfApi", + lambda: _Api( + ["adapter_config.json", "adapter_model.bin", "checkpoint-5/adapter_model.safetensors"] + ), + ) + assert U._adapter_repo_has_safetensors("org/repo") is False + # Root safetensors -> reported present. + monkeypatch.setattr( + huggingface_hub, + "HfApi", + lambda: _Api(["adapter_config.json", "adapter_model.safetensors"]), + ) + assert U._adapter_repo_has_safetensors("org/repo") is True + + +def test_gguf_file_warm_keeps_gguf(capture): + """A gguf_file load allow-lists that GGUF while not pulling other quants the repo publishes.""" + _, st = capture(weights_at_root = True, gguf_file = "model-Q4_K_M.gguf") + allow = st["allow_patterns"] + ig = st["ignore_patterns"] + assert allow is not None and "model-Q4_K_M.gguf" in allow + sample = [ + "model-Q4_K_M.gguf", + "model-Q8_0.gguf", + "config.json", + "tokenizer.json", + ] + kept = _filter(sample, allow, ig) + assert "model-Q4_K_M.gguf" in kept + assert "config.json" in kept + assert "model-Q8_0.gguf" not in kept + + +# ----- Finding Q: adapter weight-format selection ----- + + +def test_adapter_only_prefers_safetensors_over_bin(capture, monkeypatch): + """A mixed-format adapter repo warms only the safetensors PeftModel reads, not both formats.""" + _install_fake_model_info( + monkeypatch, ["adapter_config.json", "adapter_model.safetensors", "adapter_model.bin"] + ) + _, st = capture(adapter_only = True) + ig = st["ignore_patterns"] + assert ig is not None and "adapter_model*.bin" in ig + kept = _filter( + ["adapter_config.json", "adapter_model.safetensors", "adapter_model.bin"], + st["allow_patterns"], + ig, + ) + assert "adapter_model.safetensors" in kept + assert "adapter_model.bin" not in kept + + +def test_adapter_only_bin_only_keeps_bin(capture, monkeypatch): + """A .bin-only adapter repo must keep adapter_model.bin (no safetensors found -> both formats eligible).""" + _install_fake_model_info(monkeypatch, ["adapter_config.json", "adapter_model.bin"]) + _, st = capture(adapter_only = True) + kept = _filter( + ["adapter_config.json", "adapter_model.bin"], st["allow_patterns"], st["ignore_patterns"] + ) + assert "adapter_model.bin" in kept + + +def test_adapter_only_explicit_use_safetensors_false_keeps_bin(capture): + """An explicit use_safetensors=False forces the .bin form without a model_info call.""" + _, st = capture(adapter_only = True, use_safetensors = False) + ig = st["ignore_patterns"] + assert ig is not None and "adapter_model*.safetensors" in ig + kept = _filter( + ["adapter_config.json", "adapter_model.safetensors", "adapter_model.bin"], + st["allow_patterns"], + ig, + ) + assert "adapter_model.bin" in kept + assert "adapter_model.safetensors" not in kept + + +def test_gguf_file_with_subfolder_warms_subfolder_path(capture): + """gguf_file + subfolder: the warm allow-lists /, not the bare root name.""" + _, st = capture(weights_at_root = True, gguf_file = "model-Q4_K_M.gguf", subfolder = "gguf") + allow = st["allow_patterns"] + assert "gguf/model-Q4_K_M.gguf" in allow + kept = _filter(["gguf/model-Q4_K_M.gguf", "config.json"], allow, st["ignore_patterns"]) + assert "gguf/model-Q4_K_M.gguf" in kept and "config.json" in kept + + +def test_from_tf_root_load_ignores_nested_h5(capture): + """A from_tf root load keeps the root .h5 but drops nested .h5 / .msgpack checkpoints.""" + _, st = capture(weights_at_root = True, from_tf = True) + ig = st["ignore_patterns"] + assert "*/*.h5" in ig and "*/*.msgpack" in ig + kept = _filter(["model.h5", "checkpoint-1/model.h5", "config.json"], st["allow_patterns"], ig) + assert "model.h5" in kept + assert "checkpoint-1/model.h5" not in kept + + +def test_sentence_transformer_from_pretrained_is_prefetch_wired(): + """from_pretrained must call maybe_prefetch_hf_snapshot as an unconditional top-level statement before any return.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + cls = next( + n for n in tree.body if isinstance(n, ast.ClassDef) and n.name == "FastSentenceTransformer" + ) + fp = next(n for n in cls.body if isinstance(n, ast.FunctionDef) and n.name == "from_pretrained") + + def _prefetch_call(node): + # a bare call statement, or one whose return is captured (e.g. _st_prefetched = ...) + value = node.value if isinstance(node, (ast.Expr, ast.Assign)) else None + if ( + isinstance(value, ast.Call) + and isinstance(value.func, ast.Name) + and value.func.id == "maybe_prefetch_hf_snapshot" + ): + return value + return None + + prefetch_pos = next((i for i, n in enumerate(fp.body) if _prefetch_call(n)), None) + return_pos = next((i for i, n in enumerate(fp.body) if isinstance(n, ast.Return)), len(fp.body)) + assert ( + prefetch_pos is not None + ), "from_pretrained must call maybe_prefetch_hf_snapshot at top level" + assert prefetch_pos < return_pos, "prefetch must run before any top-level return" + # local_files_only must be forwarded so an offline load does not start a Hub download. + prefetch_call = _prefetch_call(fp.body[prefetch_pos]) + assert "local_files_only" in { + kw.arg for kw in prefetch_call.keywords + }, "prefetch must forward local_files_only" + + +def test_st_module_download_forwards_cache_folder(): + """_load_modules must forward the custom cache_folder into load_dir_path so per-module subdirs read the warmed cache.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + calls = [ + n + for n in ast.walk(tree) + if isinstance(n, ast.Call) and isinstance(n.func, ast.Name) and n.func.id == "load_dir_path" + ] + assert calls, "expected a load_dir_path call in sentence_transformer.py" + assert all( + "cache_folder" in {kw.arg for kw in c.keywords} for c in calls + ), "every load_dir_path call must forward cache_folder" + + +def test_st_native_sentence_transformer_calls_forward_cache_folder(): + """Every native SentenceTransformer(model_name, ...) load must forward cache_folder; a modules-based build needs none.""" + import ast + import os + + src_path = os.path.join(os.path.dirname(U.__file__), "sentence_transformer.py") + with open(src_path, "r", encoding = "utf-8") as f: + tree = ast.parse(f.read()) + weight_loading_calls = [] + for n in ast.walk(tree): + if not ( + isinstance(n, ast.Call) + and isinstance(n.func, ast.Name) + and n.func.id == "SentenceTransformer" + ): + continue + kw_names = {kw.arg for kw in n.keywords} + # A modules-based build downloads nothing; only a repo-name load reads the cache. + if "modules" in kw_names: + continue + weight_loading_calls.append(n) + assert ( + weight_loading_calls + ), "expected a repo-name SentenceTransformer load in sentence_transformer.py" + # cache_folder is forwarded explicitly or via a **kwargs unpacking (kw.arg == None). + for c in weight_loading_calls: + kw_names = {kw.arg for kw in c.keywords} + forwards = "cache_folder" in kw_names or None in kw_names + assert forwards, ( + "a repo-name SentenceTransformer load must forward cache_folder " + f"(explicitly or via **kwargs) at line {c.lineno}" + ) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 599a5c0262..a52d559aad 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -83,6 +83,7 @@ "verify_fp8_support_if_applicable", "_get_inference_mode_context_manager", "hf_login", + "maybe_prefetch_hf_snapshot", "is_moe_model", "get_moe_target_parameters", "make_fast_generate_wrapper", @@ -905,6 +906,411 @@ def _run_temporary_patches(phase): TORCHAO_MSG = "Error: torchao not found, please install with `pip install torchao`" +# Artifacts a Transformers/PEFT load never reads (ONNX/TF/Flax/CoreML/GGUF/training state), skipped +# when prewarming so a mixed-format repo is not pulled in full. +_PREFETCH_IGNORE_PATTERNS = ( + "*.onnx", + "onnx/*", + "*.h5", + "*.msgpack", + "*.tflite", + "coreml/*", + "*.mlpackage/*", + "*.mlmodel", + "*.gguf", + # Training / checkpoint formats from_pretrained never reads. + "*.pt", + "*.pth", + "*.ckpt", + "optimizer.*", + "scheduler.*", + "rng_state*", + "trainer_state.json", + "events.out.tfevents*", + "checkpoint-*/*", +) + + +# Repo-root tokenizer / config / processor files from_pretrained reads from root even when weights +# load from a subfolder. Exact names (no wildcard) so they match only root-level files. +_ROOT_AUX_PREFETCH_PATTERNS = ( + "config.json", + "generation_config.json", + "tokenizer_config.json", + "tokenizer.json", + "tokenizer.model", + "special_tokens_map.json", + "added_tokens.json", + "vocab.json", + "vocab.txt", + "merges.txt", + "spiece.model", + # More VOCAB_FILES_NAMES the slow tokenizer may fetch (DeBERTa-v2, Whisper, Mistral, XLM-R/mBART, Marian, FSMT/XLM, GPT-2). + "spm.model", + "normalizer.json", + "tokenizer.model.v3", + "sentencepiece.bpe.model", + "source.spm", + "target.spm", + "bpe.codes", + "vocab.bpe", + # More VOCAB_FILES_NAMES (RemBERT, FSMT) a distinct-tokenizer-repo warm must cache too. + "sentencepiece.model", + "vocab-src.json", + "vocab-tgt.json", + "chat_template.jinja", + "chat_template.json", + # chat_template="" fetches additional_chat_templates/.jinja. + "additional_chat_templates/*.jinja", + "preprocessor_config.json", + "processor_config.json", + "video_preprocessor_config.json", # Qwen2.5-VL-style video processors + # trust_remote_code auto_map can name any module, so warm every *.py (tiny; none in a non-remote repo). + "*.py", + "*.tiktoken", # tiktoken vocab (e.g. Qwen's qwen.tiktoken) +) + + +# Files a PEFT adapter load reads: config + weights (glob covers sharded adapters). Any merged +# full-model weights the repo also ships match none of these. +_ADAPTER_PREFETCH_PATTERNS = ( + "adapter_config.json", + "adapter_model*", +) + + +# Weight files in a SUBDIRECTORY. A bare root load reads only root weights, so ignoring these drops +# alternate-precision/experimental dirs (fp16/, experimental/). "*/*" spans "/" (HF fnmatch), so nested +# weights match while root "model.safetensors" is kept. Only applied when weights_at_root (diffusion +# keeps weights in subfolders). +_SUBDIR_WEIGHT_IGNORE_PATTERNS = ( + "*/*.safetensors", + "*/*.bin", + "*/*.h5", + "*/*.msgpack", + "*/*.pt", + "*/*.pth", +) + + +def _in_requested_load_scope(filename, subfolder): + """True if *filename* is in the location being loaded (*subfolder*, else root). Scopes the ".bin is + redundant when safetensors exist" test so a .bin-only subfolder keeps its .bin.""" + filename = filename.replace("\\", "/") + if isinstance(subfolder, str) and subfolder.strip("/"): + return filename.startswith(subfolder.strip("/") + "/") + return "/" not in filename # root load: no directory component + + +# .safetensors training-state files that are NOT model weights (e.g. optimizer.safetensors next to a +# real pytorch_model.bin); counting them as "model safetensors present" would drop the needed .bin. +_NON_MODEL_WEIGHT_STEMS = frozenset( + { + "optimizer", + "scheduler", + "scaler", + "rng_state", + "training_args", + } +) + + +def _is_model_weight_safetensors(filename): + """True if *filename* is a model-weights safetensors, not a PEFT adapter/sidecar + (adapter_model.safetensors) or trainer-state (optimizer.safetensors). Only a real one proves the + .bin redundant; counting a sidecar would wrongly drop the needed .bin (fetched then without Xet fallback).""" + name = filename.replace("\\", "/").rsplit("/", 1)[-1] + if not name.endswith((".safetensors", ".safetensors.index.json")): + return False + if name.startswith("adapter_"): + return False + # Stem before first dot: "optimizer.safetensors" -> "optimizer" (real shards kept); rng_state via prefix. + stem = name.split(".", 1)[0].lower() + if stem in _NON_MODEL_WEIGHT_STEMS or stem.startswith("rng_state"): + return False + return True + + +def _is_canonical_variant_model_weight_safetensors(filename, variant): + """True for a canonical model-weights safetensors carrying the requested *variant*, in the forms + transformers reads (single, either numbered-shard layout, or the index). Strict (base must be + "model"): a sidecar like consolidated..safetensors does not prove the variant .bin redundant.""" + base = filename.replace("\\", "/").rsplit("/", 1)[-1] + v = re.escape(variant) + return bool( + re.match( + rf"^(?:model\.{v}\.safetensors" + rf"|model\.{v}-\d{{5}}-of-\d{{5}}\.safetensors" + rf"|model-\d{{5}}-of-\d{{5}}\.{v}\.safetensors" + rf"|model\.safetensors\.index\.{v}\.json)$", + base, + ) + ) + + +_CANONICAL_MODEL_WEIGHT_SAFETENSORS_RE = re.compile( + r"^(?:model\.safetensors|model-\d{5}-of-\d{5}\.safetensors|model\.safetensors\.index\.json)$" +) + + +def _is_canonical_model_weight_safetensors(filename): + """True for a canonical (non-variant) model-weights safetensors a default load reads (model.safetensors, + a numbered shard, or the index). Strict: an unrecognized name keeps both formats, so a variant-only + safetensors + pytorch_model.bin repo never has its .bin dropped for a no-variant load.""" + name = filename.replace("\\", "/").rsplit("/", 1)[-1] + return bool(_CANONICAL_MODEL_WEIGHT_SAFETENSORS_RE.match(name)) + + +def _adapter_repo_has_safetensors( + model_name, + *, + token = None, + revision = None, +): + """Best-effort: does the adapter repo ship a root safetensors adapter weight (making the .bin + redundant)? Scoped to root adapter_model* files; any failure returns False.""" + try: + from huggingface_hub import HfApi + siblings = HfApi().model_info(model_name, revision = revision, token = token).siblings or [] + return any( + "/" not in sibling.rfilename.replace("\\", "/") # root only + and sibling.rfilename.startswith("adapter_model") + and sibling.rfilename.endswith(".safetensors") + for sibling in siblings + ) + except Exception: + return False + + +def _prefetch_ignore_patterns( + model_name, + *, + token = None, + revision = None, + subfolder = None, + use_safetensors = None, + from_tf = False, + from_flax = False, + variant = None, + weights_at_root = False, +): + """ignore_patterns for the prewarm snapshot: the static skip list, minus the checkpoint guard when + loading from a checkpoint-* subfolder, minus the weight format the load will not read. use_safetensors + is a format allowlist (True -> skip *.bin, False -> skip *.safetensors); auto (None) skips *.bin only + when in-scope safetensors are shipped. from_tf/from_flax keep *.h5/*.msgpack. + + Suppressed for a whole multi-component snapshot (weights_at_root=False, no subfolder: ST/diffusers + repos with per-subfolder weights, each in its own format), since "*" spans "/" so dropping "*.bin" + would strip a module's only weight.""" + # Keep checkpoint-*/* under a checkpoint-* subfolder; keep *.h5 / *.msgpack under from_tf/flax. + ignore_patterns = [ + pattern + for pattern in _PREFETCH_IGNORE_PATTERNS + if not ( + ( + pattern == "checkpoint-*/*" + and isinstance(subfolder, str) + and subfolder.startswith("checkpoint-") + ) + or (from_tf and pattern == "*.h5") + or (from_flax and pattern == "*.msgpack") + ) + ] + # Drop the format the load will not read (the other doubles the download); skipped for a whole + # multi-component snapshot (see docstring). + whole_multi_component = not weights_at_root and not ( + isinstance(subfolder, str) and subfolder.strip("/") + ) + if whole_multi_component: + pass + elif from_tf or from_flax: + # TF / Flax loads never read the PyTorch formats; drop safetensors and .bin. + ignore_patterns.extend( + ( + "*.safetensors", + "*.safetensors.index.json", + "*.bin", + "*.bin.index.json", + ) + ) + elif use_safetensors is True: + # Explicit safetensors: load never reads .bin (no model_info call needed). + ignore_patterns.extend(("*.bin", "*.bin.index.json")) + elif use_safetensors is False: + # Explicit .bin: load never reads safetensors. + ignore_patterns.extend(("*.safetensors", "*.safetensors.index.json")) + else: + # Auto: skip .bin only once in-scope safetensors are confirmed (best-effort; any failure keeps both). + try: + from huggingface_hub import HfApi + + siblings = ( + HfApi() + .model_info( + model_name, + revision = revision, + token = token, + ) + .siblings + or [] + ) + # Count only in-scope model-weights safetensors (not adapters/sidecars): variant-matching if + # a variant is requested, else canonical, proving the .bin redundant. + has_safetensors = any( + _is_model_weight_safetensors(sibling.rfilename) + and _in_requested_load_scope(sibling.rfilename, subfolder) + and ( + _is_canonical_variant_model_weight_safetensors(sibling.rfilename, variant) + if variant + else _is_canonical_model_weight_safetensors(sibling.rfilename) + ) + for sibling in siblings + ) + if has_safetensors: + ignore_patterns.extend(("*.bin", "*.bin.index.json")) + except Exception: + pass + return ignore_patterns + + +def maybe_prefetch_hf_snapshot( + model_name, + token = None, + *, + revision = None, + cache_dir = None, + local_files_only = False, + fast_inference = False, + subfolder = None, + force_download = False, + use_safetensors = None, + from_tf = False, + from_flax = False, + tokenizer_only = False, + adapter_only = False, + weights_at_root = False, + variant = None, + gguf_file = None, +): + """Warm the HF cache for a remote repo before the in-process load. + + Xet can hang on a blob with no progress or exception, and a blocked native Xet thread cannot be + killed in-process. So pull the snapshot first in a killable subprocess that falls back Xet -> HTTP + on a stall (unsloth_zoo.hf_xet_fallback), making from_pretrained a cache hit. + + Returns True iff warmed (caller can clear force_download), else False (skipped: local/offline/ + local_files_only/fast_inference/old unsloth_zoo, or failed). Only a both-transports-stalled + DownloadStallError is raised; other failures are left for from_pretrained to surface. + """ + try: + from unsloth_zoo.hf_xet_fallback import ( + snapshot_download_with_xet_fallback, + DownloadStallError, + ) + except Exception: + return False # older unsloth_zoo without the helper: load normally + + if not isinstance(model_name, str) or not model_name: + return False + # Local path: nothing to download. Expand ~ first (os.path.exists does not). + model_path = os.path.expanduser(model_name) + if os.path.isdir(model_path) or os.path.exists(model_path): + return False + # Looks local but not yet on disk (e.g. an uncreated output dir): not a Hub repo id, so leave it + # for from_pretrained rather than download it. + if ( + os.path.isabs(model_path) + or model_name.startswith(("~", "./", "../", ".\\", "..\\")) + or "\\" in model_name + ): + return False + if local_files_only: # cache-only: never reach out + return False + if any( + os.environ.get(flag, "0").lower() in ("1", "true", "yes", "on") + for flag in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE") + ): + return False + if fast_inference: # vLLM has its own download path + return False + + # tokenizer-only / adapter-only warms allow-list exact files below, so the weight-format ignore + # list (and its auto-branch model_info call) is skipped. + ignore_patterns = ( + None + if tokenizer_only or adapter_only or gguf_file + else _prefetch_ignore_patterns( + model_name, + token = token, + revision = revision, + subfolder = subfolder, + use_safetensors = use_safetensors, + from_tf = from_tf, + from_flax = from_flax, + variant = variant, + weights_at_root = weights_at_root, + ) + ) + # Narrow the warm to what the load reads (skip extra checkpoints/precisions); every branch still warms + # root tokenizer/config/custom-code so those never fall in-process. + allow_patterns = None + if gguf_file: + # gguf_file=NAME reads exactly that GGUF, but the static ignore list drops *.gguf; so warm just + # that file (plus root aux), under / if set. + _gguf_path = ( + f"{subfolder.strip('/')}/{gguf_file}" + if isinstance(subfolder, str) and subfolder.strip("/") + else gguf_file + ) + allow_patterns = [_gguf_path, *_ROOT_AUX_PREFETCH_PATTERNS] + elif tokenizer_only: + # A distinct tokenizer repo: warm only tokenizer / config / vocab files, never its weights. + allow_patterns = list(_ROOT_AUX_PREFETCH_PATTERNS) + elif adapter_only: + # A PEFT adapter load reads only adapter_config.json + adapter_model.* (plus root aux), not any + # merged weights the repo may also publish. + allow_patterns = [*_ADAPTER_PREFETCH_PATTERNS, *_ROOT_AUX_PREFETCH_PATTERNS] + # PeftModel reads one format (safetensors when present): explicit use_safetensors wins, else + # prefer safetensors when shipped (best-effort; any failure keeps both). + if use_safetensors is False: + ignore_patterns = [ + "adapter_model*.safetensors", + "adapter_model*.safetensors.index.json", + ] + elif use_safetensors is True or _adapter_repo_has_safetensors( + model_name, token = token, revision = revision + ): + ignore_patterns = ["adapter_model*.bin", "adapter_model*.bin.index.json"] + elif isinstance(subfolder, str) and subfolder.strip("/"): + # subfolder=X: load resolves every weight under X/, so warm that subfolder (plus root aux). + allow_patterns = [f"{subfolder.strip('/')}/*", *_ROOT_AUX_PREFETCH_PATTERNS] + elif weights_at_root: + # A bare load reads only root weights: drop subdir weights (fp16/, checkpoint dirs) while keeping + # subdir configs. Diffusion leaves weights_at_root False. + ignore_patterns = [*(ignore_patterns or []), *_SUBDIR_WEIGHT_IGNORE_PATTERNS] + try: + snapshot_download_with_xet_fallback( + model_name, + token = token, + revision = revision, + cache_dir = cache_dir, + allow_patterns = allow_patterns, + ignore_patterns = ignore_patterns, + force_download = force_download, + variant = variant, + ) + return True + except DownloadStallError: + # Both transports stalled: surface a clear network error, not a silent in-process hang. + raise + except Exception as exception: + logger.warning_once( + f"Unsloth: Could not pre-download {model_name} " + f"({type(exception).__name__}: {exception}); continuing with the normal load." + ) + return False + + # Ignore logging messages class HideLoggingMessage(logging.Filter): __slots__ = ("text",) diff --git a/unsloth/models/diffusion.py b/unsloth/models/diffusion.py index 12596b432e..955bf55987 100644 --- a/unsloth/models/diffusion.py +++ b/unsloth/models/diffusion.py @@ -24,7 +24,7 @@ import torch from transformers import AutoConfig, AutoProcessor, AutoTokenizer -from ._utils import is_bfloat16_supported +from ._utils import is_bfloat16_supported, maybe_prefetch_hf_snapshot from .llama import logger __all__ = ["FastDiffusionModel", "DIFFUSION_MODEL_TYPES", "is_diffusion_model_type"] @@ -79,7 +79,14 @@ def _resolve_diffusion_model_class(config): ) -def _load_diffusion_config(model_name, token, trust_remote_code, revision, local_files_only): +def _load_diffusion_config( + model_name, + token, + trust_remote_code, + revision, + local_files_only, + cache_dir = None, +): """Load the config, aliasing the legacy ``diffusion_gemma`` model_type to the ``diffusion_gemma4`` classes current transformers ships. AutoConfig raises on the legacy type; catch that, rewrite the type/arch names in-memory, and rebuild.""" @@ -90,6 +97,7 @@ def _load_diffusion_config(model_name, token, trust_remote_code, revision, local trust_remote_code = trust_remote_code, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) except ValueError as e: if "diffusion_gemma" not in str(e): @@ -103,6 +111,7 @@ def _load_diffusion_config(model_name, token, trust_remote_code, revision, local token = token, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) with open(cfg_path, encoding = "utf-8") as f: cd = json.load(f) @@ -152,12 +161,16 @@ def from_pretrained( os.environ.get("HF_HUB_OFFLINE", "0") == "1" or os.environ.get("TRANSFORMERS_OFFLINE", "0") == "1" ) + + cache_dir = kwargs.get("cache_dir") + config = _load_diffusion_config( model_name, token, trust_remote_code, revision, local_files_only, + cache_dir = cache_dir, ) model_type = getattr(config, "model_type", None) if not is_diffusion_model_type(model_type): @@ -168,6 +181,21 @@ def from_pretrained( model_cls = _resolve_diffusion_model_class(config) + # Prefetch the whole repo root so the weight load is a cache hit. No subfolder: the pipeline + # loads every component subfolder, so narrowing would leave unet/vae/text_encoder to Xet. + maybe_prefetch_hf_snapshot( + model_name, + token = token, + revision = revision, + cache_dir = cache_dir, + local_files_only = local_files_only, + fast_inference = False, + force_download = kwargs.get("force_download", False), + use_safetensors = kwargs.get("use_safetensors"), + # Forward variant (e.g. "fp16") so the warm keeps variant weights. + variant = kwargs.get("variant"), + ) + load_kwargs = dict( dtype = dtype, device_map = device_map, @@ -176,7 +204,14 @@ def from_pretrained( attn_implementation = attn_implementation, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) + # Match the load's weight format to the warm (None/auto already matches). + if kwargs.get("use_safetensors") is not None: + load_kwargs["use_safetensors"] = kwargs["use_safetensors"] + # Forward variant to the real load so it reads the warmed variant weights. + if kwargs.get("variant") is not None: + load_kwargs["variant"] = kwargs["variant"] # Optional bitsandbytes quant. The MoE experts (3D Parameters) are not nn.Linear so bnb skips # them; only attention + dense MLP Linears quantize, lm_head/embeddings stay full precision. @@ -222,6 +257,7 @@ def from_pretrained( trust_remote_code = trust_remote_code, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) except Exception: tokenizer = AutoTokenizer.from_pretrained( @@ -230,6 +266,7 @@ def from_pretrained( trust_remote_code = trust_remote_code, revision = revision, local_files_only = local_files_only, + cache_dir = cache_dir, ) return model, tokenizer diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 14ee5ee24e..1c95c0719c 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -2420,6 +2420,73 @@ def from_pretrained( preferred_attn_impl = resolve_attention_implementation(model_function, model_config) + # Prefetch the repo (killable child) so the weight load is a cache hit. Runs after the + # AutoConfig/model-class check so an unsupported repo fails on its small config fetch. No + # revision: the load resolves model_name (maybe a remapped prequant repo) on its default branch. + _prefetched = maybe_prefetch_hf_snapshot( + model_name, + token = token, + cache_dir = kwargs.get("cache_dir"), + local_files_only = kwargs.get("local_files_only", False), + # Skip the warm only for a real vLLM load; a num_labels classification load still goes + # in-process below, so it must be warmed even under fast_inference. + fast_inference = fast_inference and num_labels is None, + subfolder = kwargs.get("subfolder"), + force_download = kwargs.get("force_download", False), + use_safetensors = kwargs.get("use_safetensors"), + from_tf = kwargs.get("from_tf", False), + from_flax = kwargs.get("from_flax", False), + # Bare load reads only ROOT weights; skip subdir weights. Ignored when a subfolder is set. + weights_at_root = True, + variant = kwargs.get("variant"), # forward so the warm keeps the variant .bin + gguf_file = kwargs.get( + "gguf_file" + ), # forward so the warm fetches the GGUF (else ignored) + ) + # Child did the forced download; clear the flag so the load reuses the warm cache. + if _prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + + # Tokenizer always loads in-process. Resolve the cache_dir the tokenizer load will actually + # use, mirroring load_correct_tokenizer: without an explicit cache_dir, Colab/Kaggle route to + # a special tokenizer cache (huggingface_tokenizers_cache / Kaggle tmp), NOT the HF-default + # cache the base snapshot warmed. So the base warm does not cover the tokenizer there. + from ..tokenizer_utils import ( + IS_COLAB_ENVIRONMENT, + IS_KAGGLE_ENVIRONMENT, + KAGGLE_TMP, + ) + + _tokenizer_repo = ( + tokenizer_name if (isinstance(tokenizer_name, str) and tokenizer_name) else model_name + ) + _tokenizer_cache_dir = kwargs.get("cache_dir") + if _tokenizer_cache_dir is None: + if IS_COLAB_ENVIRONMENT: + _tokenizer_cache_dir = "huggingface_tokenizers_cache" + elif IS_KAGGLE_ENVIRONMENT: + _tokenizer_cache_dir = os.path.join(KAGGLE_TMP, "huggingface_tokenizers_cache") + # Warm the tokenizer repo into the cache the load will use whenever the base warm did not + # cover it: a distinct tokenizer repo, fast_inference (base warm skipped), or a tokenizer + # cache_dir that differs from the base-warm cache_dir (Colab/Kaggle special cache). + _warm_tokenizer_repo = ( + isinstance(_tokenizer_repo, str) + and bool(_tokenizer_repo) + and ( + _tokenizer_repo != model_name + or fast_inference + or _tokenizer_cache_dir != kwargs.get("cache_dir") + ) + ) + if _warm_tokenizer_repo: + maybe_prefetch_hf_snapshot( + _tokenizer_repo, + token = token, + cache_dir = _tokenizer_cache_dir, + local_files_only = kwargs.get("local_files_only", False), + tokenizer_only = True, + ) + has_rope_scaling = False try: with open(inspect.getfile(model_function), "r", encoding = "utf-8") as file: @@ -2672,6 +2739,10 @@ def from_pretrained( # Counteract saved tokenizers tokenizer_name = model_name if tokenizer_name is None else tokenizer_name + # Route the tokenizer load to the custom cache_dir the prefetch warmed. + _tokenizer_cache_kwargs = {} + if kwargs.get("cache_dir") is not None: + _tokenizer_cache_kwargs["cache_dir"] = kwargs["cache_dir"] tokenizer = load_correct_tokenizer( tokenizer_name = tokenizer_name, model_max_length = max_position_embeddings, @@ -2679,6 +2750,7 @@ def from_pretrained( token = token, trust_remote_code = trust_remote_code, fix_tokenizer = fix_tokenizer, + **_tokenizer_cache_kwargs, ) model, tokenizer = patch_tokenizer(model, tokenizer) @@ -2805,6 +2877,7 @@ def from_pretrained( model_max_length = max_position_embeddings, padding_side = "right", token = token, + cache_dir = kwargs.get("cache_dir"), ) patch_saving_functions(tokenizer) diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 562afdd645..84f808d2b5 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -106,6 +106,7 @@ _is_family_text_decoder, _apply_text_only_key_mapping, set_task_config_attr, + maybe_prefetch_hf_snapshot, ) # Single source of truth is unsloth_zoo.model_lists. Re-exported so callers @@ -865,6 +866,28 @@ def from_pretrained( if is_peft: # From https://github.com/huggingface/peft/issues/184 # Now add PEFT adapters + # Warm the adapter repo: PeftModel downloads it in-process and can hang on Xet. + _prefetched = maybe_prefetch_hf_snapshot( + old_model_name, + token = token, + revision = revision, + cache_dir = kwargs.get("cache_dir"), + local_files_only = local_files_only, + # Adapter always loads in-process via PeftModel, so warm it even under fast_inference. + fast_inference = False, + force_download = kwargs.get("force_download", False), + # Leave use_safetensors auto (inheriting base format could skip a safetensors-only + # adapter). adapter_only restricts the warm to the adapter files + root aux. + adapter_only = True, + ) + # Child did the forced download; clear the flag so the load reuses the warm cache. + if _prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + # Forward cache_dir so the load reads the warmed adapter. No subfolder (that targets the + # base checkpoint; adapters live at the root). + peft_load_kwargs = {} + if kwargs.get("cache_dir") is not None: + peft_load_kwargs["cache_dir"] = kwargs["cache_dir"] model = PeftModel.from_pretrained( model, old_model_name, @@ -873,6 +896,7 @@ def from_pretrained( local_files_only = local_files_only, is_trainable = True, trust_remote_code = trust_remote_code, + **peft_load_kwargs, ) # Patch it as well! model = dispatch_model.patch_peft_model(model, use_gradient_checkpointing) @@ -1790,6 +1814,28 @@ def _patched_car( _LoraModel._create_and_replace = _patched_car + # Warm the adapter repo: PeftModel downloads it in-process and can hang on Xet. + _prefetched = maybe_prefetch_hf_snapshot( + old_model_name, + token = token, + revision = revision, + cache_dir = kwargs.get("cache_dir"), + local_files_only = local_files_only, + # Adapter always loads in-process via PeftModel, so warm it even under fast_inference. + fast_inference = False, + force_download = kwargs.get("force_download", False), + # Leave use_safetensors auto (inheriting base format could skip a safetensors-only + # adapter). adapter_only restricts the warm to the adapter files + root aux. + adapter_only = True, + ) + # Child did the forced download; clear the flag so the load reuses the warm cache. + if _prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + # Forward cache_dir so the load reads the warmed adapter. No subfolder (that targets the + # base checkpoint; adapters live at the root). + peft_load_kwargs = {} + if kwargs.get("cache_dir") is not None: + peft_load_kwargs["cache_dir"] = kwargs["cache_dir"] try: model = PeftModel.from_pretrained( model, @@ -1799,6 +1845,7 @@ def _patched_car( local_files_only = local_files_only, is_trainable = True, trust_remote_code = trust_remote_code, + **peft_load_kwargs, ) finally: # Always restore original PEFT method, even if loading fails diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index 7e43442bfd..c1172faa94 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -19,6 +19,7 @@ SUPPORTS_BFLOAT16, resolve_model_class, resolve_encoder_attention_implementation, + maybe_prefetch_hf_snapshot, ) import inspect import json @@ -541,7 +542,12 @@ def _save_with_base_config(self, output_path, *args, **kwargs): return transformer_module @staticmethod - def _read_pooling_mode(model_name, token): + def _read_pooling_mode( + model_name, + token, + cache_dir = None, + revision = None, + ): """Read the pooling mode from modules.json, else return "mean".""" try: if os.path.exists(model_name) and os.path.exists( @@ -549,7 +555,13 @@ def _read_pooling_mode(model_name, token): ): modules_json_path = os.path.join(model_name, "modules.json") else: - modules_json_path = hf_hub_download(model_name, "modules.json", token = token) + modules_json_path = hf_hub_download( + model_name, + "modules.json", + token = token, + cache_dir = cache_dir, + revision = revision, + ) with open(modules_json_path, "r", encoding = "utf-8") as f: modules_config = json.load(f) @@ -571,6 +583,8 @@ def _read_pooling_mode(model_name, token): model_name, os.path.join(pooling_path, "config.json"), token = token, + cache_dir = cache_dir, + revision = revision, ) break @@ -950,7 +964,12 @@ def _add_unsloth_branding(save_directory): f.write(content) @staticmethod - def _module_path(model_name, token = None): + def _module_path( + model_name, + token = None, + cache_dir = None, + revision = None, + ): """Return the path to the modules.json file, or None.""" try: if os.path.exists(model_name) and os.path.isdir(model_name): @@ -958,7 +977,13 @@ def _module_path(model_name, token = None): return path if os.path.exists(path) else None else: try: - return hf_hub_download(model_name, "modules.json", token = token) + return hf_hub_download( + model_name, + "modules.json", + token = token, + cache_dir = cache_dir, + revision = revision, + ) except: return None except: @@ -1135,6 +1160,8 @@ def _load_modules( max_seq_length, pooling_mode, trust_remote_code = False, + cache_dir = None, + revision = None, ) -> tuple[OrderedDict, bool]: """Load modules from modules.json, else fall back to hard-coded modules. @@ -1145,7 +1172,9 @@ def _load_modules( from sentence_transformers.models import Pooling, Normalize modules = OrderedDict() - modules_json_path = FastSentenceTransformer._module_path(model_name, token) + modules_json_path = FastSentenceTransformer._module_path( + model_name, token, cache_dir = cache_dir, revision = revision + ) if modules_json_path: with open(modules_json_path, encoding = "utf8") as f: @@ -1171,7 +1200,13 @@ def _load_modules( load_path = os.path.join(model_name, module_path) else: try: - load_path = load_dir_path(model_name, module_path, token = token) + load_path = load_dir_path( + model_name, + module_path, + token = token, + cache_folder = cache_dir, + revision = revision, + ) except Exception as e: print(f"Unsloth Warning: Could not download module {module_path}: {e}") continue @@ -1198,7 +1233,9 @@ def _load_modules( hidden_size = getattr(model.config, "hidden_size", 768) if pooling_mode == "mean": - pooling_mode = FastSentenceTransformer._read_pooling_mode(model_name, token) + pooling_mode = FastSentenceTransformer._read_pooling_mode( + model_name, token, cache_dir = cache_dir, revision = revision + ) modules["1"] = Pooling(word_embedding_dimension = hidden_size, pooling_mode = pooling_mode) modules["2"] = Normalize() @@ -1386,6 +1423,45 @@ def from_pretrained( "Run `pip install sentence-transformers` to install it." ) + # Validate the load modes BEFORE the prefetch so a bad config fails without downloading weights. + # Guard on not for_inference: that branch below never used these flags. + if not for_inference: + # sanity check, thanks Etherl: + if full_finetuning and (load_in_4bit or load_in_8bit): + print( + "Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA." + ) + load_in_4bit = False + load_in_8bit = False + load_in_fp8 = False + load_in_16bit = False + + if int(load_in_4bit) + int(load_in_8bit) + int(load_in_16bit) >= 2: + raise RuntimeError( + "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!\n" + "Also, we by default set `load_in_16bit = True`.\n" + "If you want 4bit LoRA finetuning, set `load_in_16bit = False` and `load_in_4bit = True`\n" + "If you want 8bit finetuning, set both `load_in_16bit = False` and `load_in_8bit = True`" + ) + + # Prefetch so the ST load below is a cache hit. weights_at_root stays False (ST component + # weights live in per-module subfolders). Resolve the same cache the load uses: HF cache_dir, + # else cache_folder, else SENTENCE_TRANSFORMERS_HOME, else default -- a wrong cache misses the warm. + _st_prefetched = maybe_prefetch_hf_snapshot( + model_name, + token = token, + revision = revision, + cache_dir = kwargs.get("cache_dir") + or kwargs.get("cache_folder") + or os.environ.get("SENTENCE_TRANSFORMERS_HOME"), + local_files_only = kwargs.get("local_files_only", False), + # Forward force_download so the refresh happens in the killable child, then clear it so the + # in-process ST load reuses the warm cache instead of re-downloading over unguarded Xet. + force_download = kwargs.get("force_download", False), + ) + if _st_prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + # if for_inference == True, skip Unsloth optimizations to avoid torch compile issues if for_inference: st_device = device_map @@ -1416,27 +1492,16 @@ def from_pretrained( if k in kwargs: st_kwargs[k] = kwargs[k] + # ST takes cache_folder, not cache_dir: map cache_dir onto it so this load hits the warm + # (None lets ST honor SENTENCE_TRANSFORMERS_HOME, matching the prefetch). + _st_cache = kwargs.get("cache_dir") or kwargs.get("cache_folder") + if _st_cache is not None: + st_kwargs["cache_folder"] = _st_cache + st_model = SentenceTransformer(model_name, **st_kwargs) return st_model - # sanity check, thanks Etherl: - if full_finetuning and (load_in_4bit or load_in_8bit): - print( - "Unsloth: You selected full finetuning support, but 4bit / 8bit is enabled - disabling LoRA / QLoRA." - ) - load_in_4bit = False - load_in_8bit = False - load_in_fp8 = False - load_in_16bit = False - - if int(load_in_4bit) + int(load_in_8bit) + int(load_in_16bit) >= 2: - raise RuntimeError( - "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!\n" - "Also, we by default set `load_in_16bit = True`.\n" - "If you want 4bit LoRA finetuning, set `load_in_16bit = False` and `load_in_4bit = True`\n" - "If you want 8bit finetuning, set both `load_in_16bit = False` and `load_in_8bit = True`" - ) - + # Load-mode validation already ran before the prefetch above. if "auto_model" not in kwargs: kwargs["auto_model"] = AutoModel @@ -1533,7 +1598,8 @@ def from_pretrained( elif is_mpnet: FastSentenceTransformer._patch_mpnet_v5() - # Load via native SentenceTransformer (bypasses Unsloth patching) + # ST takes cache_folder, not cache_dir: map cache_dir onto it so this load hits the warm + # (None lets ST honor SENTENCE_TRANSFORMERS_HOME, matching the prefetch). st_model = SentenceTransformer( model_name, device = st_device, @@ -1541,6 +1607,7 @@ def from_pretrained( token = token, revision = revision, model_kwargs = model_kwargs, + cache_folder = kwargs.get("cache_dir") or kwargs.get("cache_folder"), ) # Store metadata for get_peft_model @@ -1646,7 +1713,18 @@ def _push_to_hub_merged(self, repo_id, **push_kwargs): # No modules.json -> force 16-bit: saving is custom for these models and # 4-bit would need dequant in save_pretrained_merged, not worth it. - has_modules_json = FastSentenceTransformer._module_path(model_name, token) is not None + # Resolve the warmed cache: hf_hub_download ignores SENTENCE_TRANSFORMERS_HOME, so pass it as cache_dir. + has_modules_json = ( + FastSentenceTransformer._module_path( + model_name, + token, + cache_dir = kwargs.get("cache_dir") + or kwargs.get("cache_folder") + or os.environ.get("SENTENCE_TRANSFORMERS_HOME"), + revision = revision, + ) + is not None + ) if not has_modules_json and load_in_4bit: print( @@ -1656,6 +1734,12 @@ def _push_to_hub_merged(self, repo_id, **push_kwargs): load_in_4bit = False load_in_16bit = True + # The fallback FastModel load reads HF cache_dir, not ST's cache_folder/SENTENCE_TRANSFORMERS_HOME. + # Point it at the warmed cache, but only when no explicit cache_dir was passed (which wins). + _st_cache_dir = kwargs.get("cache_folder") or os.environ.get("SENTENCE_TRANSFORMERS_HOME") + if _st_cache_dir is not None and "cache_dir" not in kwargs: + kwargs["cache_dir"] = _st_cache_dir + try: model, tokenizer = FastModel.from_pretrained( model_name = model_name, @@ -1697,6 +1781,12 @@ def _push_to_hub_merged(self, repo_id, **push_kwargs): max_seq_length, pooling_mode, trust_remote_code = trust_remote_code, + # Same resolved cache as above so the fallback module loads hit the warm, not Xet. + cache_dir = kwargs.get("cache_dir") + or kwargs.get("cache_folder") + or os.environ.get("SENTENCE_TRANSFORMERS_HOME"), + # Same revision as the weight load so modules hit the warm (None = default branch). + revision = revision, ) st_device = device_map diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index bdc2bd9ef6..891bed20c8 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -479,6 +479,7 @@ def _construct_vlm_processor_fallback( model_type, token, trust_remote_code, + cache_dir = None, local_files_only = False, ): """Build a VLM processor manually when AutoProcessor.from_pretrained fails (some VLMs @@ -496,6 +497,7 @@ def _construct_vlm_processor_fallback( tokenizer_name, token = token, trust_remote_code = trust_remote_code, + cache_dir = cache_dir, local_files_only = local_files_only, ) # Load tokenizer via PreTrainedTokenizerFast (bypasses tokenizer_class check) @@ -504,6 +506,7 @@ def _construct_vlm_processor_fallback( padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = cache_dir, local_files_only = local_files_only, ) # Read tokenizer_config.json for special tokens: prefer the local file (offline @@ -529,6 +532,7 @@ def _construct_vlm_processor_fallback( tokenizer_name, "tokenizer_config.json", token = token, + cache_dir = cache_dir, local_files_only = local_files_only, ) with open(config_path, "r", encoding = "utf-8") as f: @@ -560,6 +564,7 @@ def _construct_vlm_processor_fallback( tokenizer_name, token = token, trust_remote_code = trust_remote_code, + cache_dir = cache_dir, local_files_only = local_files_only, ) proc_class_name = PROCESSOR_MAPPING_NAMES.get(config.model_type) @@ -800,6 +805,9 @@ def from_pretrained( # For debugging - we use a download counter to see if environments are not breaking or if HF is down get_statistics(kwargs.get("local_files_only", False)) + # The base + tokenizer prefetch runs AFTER the load-mode validation below, so an invalid + # load_in_* combination fails without first downloading a snapshot. + if dtype is None: dtype = torch.float16 if not SUPPORTS_BFLOAT16 else torch.bfloat16 elif os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": @@ -896,6 +904,53 @@ def from_pretrained( raise RuntimeError( "Unsloth: Can only load in 4bit or 8bit or 16bit, not a combination!" ) + + # Prefetch the repo (killable child) so the in-process load below is a cache hit. vLLM owns the + # weight download only when actually available; if fast_inference was requested but vLLM is + # missing, the load falls through in-process, so weights must still be warmed here. + _vllm_owns_weights = fast_inference and is_vLLM_available() + _prefetched = maybe_prefetch_hf_snapshot( + model_name, + token = token, + revision = kwargs.get("revision"), + cache_dir = kwargs.get("cache_dir"), + local_files_only = kwargs.get("local_files_only", False), + fast_inference = _vllm_owns_weights, + subfolder = kwargs.get("subfolder"), + force_download = kwargs.get("force_download", False), + use_safetensors = kwargs.get("use_safetensors"), + from_tf = kwargs.get("from_tf", False), + from_flax = kwargs.get("from_flax", False), + # Bare load reads only ROOT weights; skip subdir weights. Ignored when a subfolder is set. + weights_at_root = True, + variant = kwargs.get("variant"), # forward so the warm keeps the variant .bin + gguf_file = kwargs.get( + "gguf_file" + ), # forward so the warm fetches the GGUF (else ignored) + ) + # Child did the forced download; clear the flag so the load reuses the warm cache. + if _prefetched and kwargs.get("force_download", False): + kwargs["force_download"] = False + + # Warm a SEPARATE tokenizer repo only (model_name is covered above). Not model_name here: this + # runs before fast_inference_setup may remap the repo, so it would warm the wrong one. + _tokenizer_repo = ( + tokenizer_name if (isinstance(tokenizer_name, str) and tokenizer_name) else model_name + ) + _warm_tokenizer_repo = ( + isinstance(_tokenizer_repo, str) + and bool(_tokenizer_repo) + and _tokenizer_repo != model_name + ) + if _warm_tokenizer_repo: + maybe_prefetch_hf_snapshot( + _tokenizer_repo, + token = token, + cache_dir = kwargs.get("cache_dir"), + local_files_only = kwargs.get("local_files_only", False), + tokenizer_only = True, + ) + _skip_modules = SKIP_QUANTIZATION_MODULES.copy() # Nemotron-H uses 'mixer' (not 'mamba') for Mamba layers. # Mamba fused kernels pass out_proj.weight directly to F.linear, @@ -1195,6 +1250,18 @@ def from_pretrained( # Counteract saved tokenizers tokenizer_name = model_name if tokenizer_name is None else tokenizer_name + # On the vLLM path the tokenizer warm was deferred (fast_inference_setup may remap model_name). + # Warm the now-final tokenizer repo so the load below hits the cache (a cached/local repo is a no-op). + if _vllm_owns_weights and isinstance(tokenizer_name, str) and tokenizer_name: + maybe_prefetch_hf_snapshot( + tokenizer_name, + token = token, + revision = kwargs.get("revision"), + cache_dir = kwargs.get("cache_dir"), + local_files_only = kwargs.get("local_files_only", False), + tokenizer_only = True, + ) + # Fix _Unsloth_Patched_ prefix in local config files from old saves (issue #4085) if os.path.isdir(tokenizer_name): import json as _json @@ -1232,6 +1299,7 @@ def _acquire_processor(lfo): language = whisper_language, task = whisper_task, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception as _e: @@ -1244,6 +1312,7 @@ def _acquire_processor(lfo): padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception as _e: @@ -1254,6 +1323,7 @@ def _acquire_processor(lfo): padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception: @@ -1272,6 +1342,7 @@ def _acquire_processor(lfo): model_type_arch, token, trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception as _fe: @@ -1357,6 +1428,7 @@ def _is_degraded_vlm(_t): padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = local_files_only, ) model, _fallback_tok = patch_tokenizer(model, _fallback_tok) @@ -1386,6 +1458,7 @@ def _last_resort_tokenizer(lfo): padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) except Exception: @@ -1395,6 +1468,7 @@ def _last_resort_tokenizer(lfo): padding_side = "left", token = token, trust_remote_code = trust_remote_code, + cache_dir = kwargs.get("cache_dir"), local_files_only = lfo, ) diff --git a/unsloth/tokenizer_utils.py b/unsloth/tokenizer_utils.py index 93dfa9b2ad..3a91ef188d 100644 --- a/unsloth/tokenizer_utils.py +++ b/unsloth/tokenizer_utils.py @@ -563,8 +563,11 @@ def _load_correct_tokenizer( # /tmp of Kaggle seems has a 80GB limit! # Let's utilize them cache_dir = os.path.join(KAGGLE_TMP, cache_dir) - else: + elif cache_dir == "huggingface_tokenizers_cache": + # This default name is Colab/Kaggle-only; elsewhere use the HF default cache. cache_dir = None + # else: keep a caller-supplied cache_dir so the tokenizer loads from the prefetch-warmed dir instead + # of risking an in-process Hub/Xet transfer. # Try loading the slow tokenizer. If it fails, then try Fast only # Mainly to solve Deepseek models with no tokenizer.model file @@ -1323,6 +1326,7 @@ def check_tokenizer( padding_side = "right", token = None, _reload = True, + cache_dir = None, ): # Checks tokenizer for out of bounds ids. # Mainly a fix for https://huggingface.co/berkeley-nest/Starling-LM-7B-alpha @@ -1413,10 +1417,11 @@ def check_tokenizer( f"Fix your tokenizer since it'll perform out of bounds memory accesses." ) - if IS_COLAB_ENVIRONMENT or IS_KAGGLE_ENVIRONMENT: - cache_dir = "huggingface_tokenizers_cache" - else: - cache_dir = None + # Reuse a caller-supplied cache_dir (warmed cache) for the repair reload; else the + # Colab/Kaggle sentinel (HF default elsewhere), as load_correct_tokenizer does. + reload_cache_dir = cache_dir + if reload_cache_dir is None and (IS_COLAB_ENVIRONMENT or IS_KAGGLE_ENVIRONMENT): + reload_cache_dir = "huggingface_tokenizers_cache" # Sometimes slow tokenizer does not work like Deepseek try: @@ -1430,7 +1435,7 @@ def check_tokenizer( use_fast = False, legacy = False, from_slow = True, - cache_dir = cache_dir, + cache_dir = reload_cache_dir, ) return check_tokenizer( model = model, @@ -1440,6 +1445,7 @@ def check_tokenizer( padding_side = padding_side, token = token, _reload = False, + cache_dir = cache_dir, ) break except: