diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index 9ecfa73eee..709de16da1 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -172,6 +172,68 @@ def _activate_transformers_version(model_name: str, hf_token: str | None = None) activate_transformers_for_subprocess(model_name, hf_token) +def _reset_hf_sessions() -> None: + try: + from huggingface_hub.utils._http import reset_sessions + except Exception: + try: + from huggingface_hub.utils import reset_sessions + except Exception: + return + try: + reset_sessions() + except Exception: + pass + + +@contextlib.contextmanager +def _force_hf_offline_window(): + saved_env = {key: os.environ.get(key) for key in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE")} + saved_attrs = [] + try: + import huggingface_hub.constants as hub_constants + if hasattr(hub_constants, "HF_HUB_OFFLINE"): + saved_attrs.append( + ( + hub_constants, + "HF_HUB_OFFLINE", + hub_constants.HF_HUB_OFFLINE, + ) + ) + except Exception: + pass + try: + import transformers.utils.hub as transformers_hub + for attr in ("_is_offline_mode", "OFFLINE"): + if hasattr(transformers_hub, attr): + saved_attrs.append((transformers_hub, attr, getattr(transformers_hub, attr))) + except Exception: + pass + + try: + os.environ["HF_HUB_OFFLINE"] = "1" + os.environ["TRANSFORMERS_OFFLINE"] = "1" + for obj, attr, _ in saved_attrs: + try: + setattr(obj, attr, True) + except Exception: + pass + _reset_hf_sessions() + yield + finally: + for obj, attr, value in saved_attrs: + try: + setattr(obj, attr, value) + except Exception: + pass + for key, value in saved_env.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + _reset_hf_sessions() + + @contextlib.contextmanager def _offline_window_if_unreachable(step = "loading"): """Force HF offline for a network-touching step (transformers version activation, or the @@ -187,18 +249,21 @@ def _offline_window_if_unreachable(step = "loading"): force_ctx = None try: from utils.transformers_version import _env_offline, hf_endpoint_unreachable + probe_enabled = os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() not in ( "0", "false", "no", "off", ) - if not _env_offline() and probe_enabled and hf_endpoint_unreachable(): + should_force = _env_offline() + if not should_force and probe_enabled and hf_endpoint_unreachable(): + should_force = True logger.warning("Hugging Face endpoint unreachable; %s offline", step) - if "huggingface_hub" in sys.modules: + if should_force: + if "huggingface_hub" in sys.modules or "transformers" in sys.modules: try: - from unsloth.models.loader_utils import _force_hf_offline - force_ctx = _force_hf_offline() + force_ctx = _force_hf_offline_window() force_ctx.__enter__() # sets env + in-process flags + resets sessions except Exception: force_ctx = None @@ -659,7 +724,9 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None _handle_load(backend, cmd, resp_queue) elif cmd_type == "export": - _handle_export(backend, cmd, resp_queue) + # Export can trigger hidden Hub metadata calls from tokenizer save paths. + with _offline_window_if_unreachable(step = "exporting"): + _handle_export(backend, cmd, resp_queue) elif cmd_type == "cleanup": _handle_cleanup(backend, resp_queue) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 6e587c18e8..495d6dc544 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -129,6 +129,7 @@ try: from utils.models.model_config import ( _pick_best_gguf, _extract_quant_label, + _hf_metadata_unavailable, _is_big_endian_gguf_path, _is_mtp_drafter, is_audio_input_type, @@ -162,6 +163,7 @@ except ImportError: from utils.models.model_config import ( _pick_best_gguf, _extract_quant_label, + _hf_metadata_unavailable, _is_big_endian_gguf_path, _is_mtp_drafter, is_audio_input_type, @@ -1825,6 +1827,9 @@ def _get_max_position_embeddings(config) -> Optional[int]: def _get_model_size_bytes(model_name: str, hf_token: Optional[str] = None) -> Optional[int]: """Total size of model weight files from HF Hub.""" + if _hf_metadata_unavailable(): + return None + try: from huggingface_hub import HfApi diff --git a/studio/backend/tests/test_models_get_model_config_case_resolution.py b/studio/backend/tests/test_models_get_model_config_case_resolution.py index a50765898b..ae98841976 100644 --- a/studio/backend/tests/test_models_get_model_config_case_resolution.py +++ b/studio/backend/tests/test_models_get_model_config_case_resolution.py @@ -7,7 +7,7 @@ import types # Keep this test runnable in lightweight environments where optional logging # deps are not installed. -if "structlog" not in sys.modules: +if "structlog" not in sys.modules or not hasattr(sys.modules["structlog"], "get_logger"): class _DummyLogger: def __getattr__(self, _name): @@ -106,3 +106,14 @@ def test_repo_in_any_hf_cache_matches_case_variant_in_legacy_cache(tmp_path, mon assert models_route._repo_in_any_hf_cache("unsloth/foo") is True # Absent from every cache -> reported absent. assert models_route._repo_in_any_hf_cache("unsloth/not-cached") is False + + +def test_get_model_size_skips_hub_when_metadata_unavailable(monkeypatch): + monkeypatch.setattr(models_route, "_hf_metadata_unavailable", lambda: True) + + class _BoomApi: + def __init__(self, *args, **kwargs): + raise AssertionError("HfApi must not be constructed when metadata is unavailable") + + monkeypatch.setattr("huggingface_hub.HfApi", _BoomApi) + assert models_route._get_model_size_bytes("unsloth/a") is None diff --git a/studio/backend/tests/test_offline_embedding_minimal.py b/studio/backend/tests/test_offline_embedding_minimal.py index ccc6b5f76a..b39e9cf93b 100644 --- a/studio/backend/tests/test_offline_embedding_minimal.py +++ b/studio/backend/tests/test_offline_embedding_minimal.py @@ -127,6 +127,7 @@ def _clean_env(monkeypatch): """Start each test online with an empty detection cache; offline tests opt in.""" monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + monkeypatch.setenv("UNSLOTH_OFFLINE_PROBE", "0") from utils.models import model_config as mc mc._embedding_detection_cache.clear() @@ -322,6 +323,20 @@ def test_offline_ignores_stale_online_memo(hf_cache, monkeypatch): assert _is_embedding_model("org/uncached-emb") is False # recomputed from empty cache +def test_offline_keeps_verified_embedding_memo_for_loadable_cache(hf_cache, monkeypatch): + from utils.models import model_config + + _make_cache( + hf_cache, + "org/transformers-emb", + {"config.json": "{}", "model.safetensors": "x"}, + ) + model_config._embedding_detection_cache[("org/transformers-emb", None)] = True + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + with _no_network(): + assert _is_embedding_model("org/transformers-emb") is True + + def test_offline_recomputes_after_cache_materializes(hf_cache, monkeypatch): # Because the offline branch never records a memo, once an uncached repo's snapshot # materializes (another process populates the cache) the next call re-reports True. @@ -335,6 +350,72 @@ def test_offline_recomputes_after_cache_materializes(hf_cache, monkeypatch): # ── is_embedding_model: online (bounded + fallback) ────────────── +def test_unreachable_endpoint_uses_cached_embedding_marker(hf_cache, monkeypatch): + _make_cache(hf_cache, "org/emb", {"modules.json": MODULES_JSON}) + monkeypatch.setenv("UNSLOTH_OFFLINE_PROBE", "1") + monkeypatch.setattr( + "utils.transformers_version.hf_endpoint_unreachable", + lambda timeout = 3: True, + ) + with _no_network(): + assert _is_embedding_model("org/emb") is True + + +def test_unreachable_endpoint_ignores_stale_online_embedding_memo(hf_cache, monkeypatch): + from utils.models import model_config + + model_config._embedding_detection_cache[("org/uncached-emb", None)] = True + monkeypatch.setenv("UNSLOTH_OFFLINE_PROBE", "1") + monkeypatch.setattr(model_config, "_hf_metadata_probe_state", None) + monkeypatch.setattr( + "utils.transformers_version.hf_endpoint_unreachable", + lambda timeout = 3: True, + ) + with _no_network(): + assert _is_embedding_model("org/uncached-emb") is False + + +def test_unreachable_endpoint_keeps_verified_memo_for_loadable_cache(hf_cache, monkeypatch): + from utils.models import model_config + + _make_cache( + hf_cache, + "org/transformers-emb", + {"config.json": "{}", "model.safetensors": "x"}, + ) + model_config._embedding_detection_cache[("org/transformers-emb", None)] = True + monkeypatch.setenv("UNSLOTH_OFFLINE_PROBE", "1") + monkeypatch.setattr(model_config, "_hf_metadata_probe_state", None) + monkeypatch.setattr( + "utils.transformers_version.hf_endpoint_unreachable", + lambda timeout = 3: True, + ) + with _no_network(): + assert _is_embedding_model("org/transformers-emb") is True + + +def test_metadata_reachability_probe_is_coalesced(monkeypatch): + from utils.models import model_config + + calls = 0 + + def reachable(*, timeout): + nonlocal calls + calls += 1 + return False + + monkeypatch.setenv("UNSLOTH_OFFLINE_PROBE", "1") + monkeypatch.setattr(model_config, "_hf_metadata_probe_state", None) + monkeypatch.setattr( + "utils.transformers_version.hf_endpoint_unreachable", + reachable, + ) + + assert model_config._hf_metadata_unavailable() is False + assert model_config._hf_metadata_unavailable() is False + assert calls == 1 + + def test_online_passes_bounded_timeout(hf_cache): seen = {} diff --git a/studio/backend/tests/test_offline_export_worker.py b/studio/backend/tests/test_offline_export_worker.py new file mode 100644 index 0000000000..eea98f4c49 --- /dev/null +++ b/studio/backend/tests/test_offline_export_worker.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import pytest + + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + + +def test_offline_window_forces_in_process_hub_flags(monkeypatch): + pytest.importorskip("huggingface_hub") + pytest.importorskip("transformers") + + import huggingface_hub.constants as hub_constants + import transformers.utils.hub as transformers_hub + import utils.transformers_version as transformers_version + from core.export.worker import _offline_window_if_unreachable + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + monkeypatch.setattr(transformers_version, "_env_offline", lambda: False) + monkeypatch.setattr( + transformers_version, + "hf_endpoint_unreachable", + lambda timeout = 3: True, + ) + monkeypatch.setattr( + hub_constants, + "HF_HUB_OFFLINE", + False, + raising = False, + ) + if hasattr(transformers_hub, "_is_offline_mode"): + monkeypatch.setattr( + transformers_hub, + "_is_offline_mode", + False, + raising = False, + ) + if hasattr(transformers_hub, "OFFLINE"): + monkeypatch.setattr( + transformers_hub, + "OFFLINE", + False, + raising = False, + ) + + with _offline_window_if_unreachable(step = "test"): + assert os.environ.get("HF_HUB_OFFLINE") == "1" + assert os.environ.get("TRANSFORMERS_OFFLINE") == "1" + assert hub_constants.HF_HUB_OFFLINE is True + if hasattr(transformers_hub, "_is_offline_mode"): + assert transformers_hub._is_offline_mode is True + if hasattr(transformers_hub, "OFFLINE"): + assert transformers_hub.OFFLINE is True + + assert os.environ.get("HF_HUB_OFFLINE") is None + assert os.environ.get("TRANSFORMERS_OFFLINE") is None + + +def test_offline_window_falls_back_when_force_context_enter_fails(monkeypatch): + pytest.importorskip("transformers") + import utils.transformers_version as transformers_version + import core.export.worker as worker + + class _BrokenContext: + def __enter__(self): + raise RuntimeError("offline force unavailable") + + def __exit__(self, *_args): + raise AssertionError("__exit__ should not run when __enter__ failed") + + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + monkeypatch.setattr(transformers_version, "_env_offline", lambda: False) + monkeypatch.setattr( + transformers_version, + "hf_endpoint_unreachable", + lambda timeout = 3: True, + ) + monkeypatch.setattr( + worker, + "_force_hf_offline_window", + lambda: _BrokenContext(), + ) + + with worker._offline_window_if_unreachable(step = "test"): + assert os.environ.get("HF_HUB_OFFLINE") == "1" + assert os.environ.get("TRANSFORMERS_OFFLINE") == "1" + + assert os.environ.get("HF_HUB_OFFLINE") is None + assert os.environ.get("TRANSFORMERS_OFFLINE") is None diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index d1e61d0546..357f14a949 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -142,6 +142,7 @@ def clean_offline_env(monkeypatch): """Strip ``HF_HUB_OFFLINE`` / ``TRANSFORMERS_OFFLINE`` for the test.""" monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False) + monkeypatch.setenv("UNSLOTH_OFFLINE_PROBE", "0") class TestGgufVariantFileResolution: @@ -819,6 +820,40 @@ class TestDetectGgufFromCache: class TestDetectGgufModelRemoteOffline: + def test_unreachable_endpoint_short_circuits_retries( + self, hf_cache, clean_offline_env, monkeypatch + ): + _build_cache(hf_cache, "unsloth/a", {"a-Q4_K_M.gguf": 1}) + monkeypatch.setenv("UNSLOTH_OFFLINE_PROBE", "1") + monkeypatch.setattr( + "utils.transformers_version.hf_endpoint_unreachable", + lambda timeout = 3: True, + ) + + def boom(*args, **kwargs): + raise AssertionError("API must not be called when endpoint is unreachable") + + with patch("huggingface_hub.model_info", boom): + assert detect_gguf_model_remote("unsloth/a") == "a-Q4_K_M.gguf" + + def test_unreachable_endpoint_lists_cached_variants_without_api( + self, hf_cache, clean_offline_env, monkeypatch + ): + _build_cache(hf_cache, "unsloth/a", {"a-Q4_K_M.gguf": 1}) + monkeypatch.setenv("UNSLOTH_OFFLINE_PROBE", "1") + monkeypatch.setattr( + "utils.transformers_version.hf_endpoint_unreachable", + lambda timeout = 3: True, + ) + + def boom(*args, **kwargs): + raise AssertionError("API must not be called when endpoint is unreachable") + + with patch("huggingface_hub.model_info", boom): + variants, has_vision = list_gguf_variants("unsloth/a") + assert [variant.filename for variant in variants] == ["a-Q4_K_M.gguf"] + assert has_vision is False + def test_offline_env_short_circuits_retries(self, hf_cache, clean_offline_env, monkeypatch): _build_cache(hf_cache, "unsloth/a", {"a-Q4_K_M.gguf": 1}) monkeypatch.setenv("HF_HUB_OFFLINE", "1") @@ -843,6 +878,28 @@ class TestDetectGgufModelRemoteOffline: out = detect_gguf_model_remote("unsloth/a") assert out == "a-Q4_K_M.gguf" + def test_transient_timeout_retries_before_success(self, clean_offline_env): + calls = 0 + + class ReadTimeout(Exception): + pass + + def flaky(*_args, **_kwargs): + nonlocal calls + calls += 1 + if calls < 3: + raise ReadTimeout("temporary timeout") + return _types.SimpleNamespace( + siblings = [_types.SimpleNamespace(rfilename = "a-Q4_K_M.gguf")] + ) + + with ( + patch("huggingface_hub.model_info", flaky), + patch("time.sleep", lambda *_: None), + ): + assert detect_gguf_model_remote("unsloth/a") == "a-Q4_K_M.gguf" + assert calls == 3 + def test_remote_big_endian_only_repo_is_not_detected(self, clean_offline_env, monkeypatch): siblings = [ _types.SimpleNamespace(rfilename = "model-Q4_K_M-be.gguf"), diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 6270d9e03f..b1fdf62a63 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -33,6 +33,7 @@ from typing import List, Tuple import hashlib import json import threading +import time import yaml @@ -47,6 +48,9 @@ logger = get_logger(__name__) _OFFLINE_TRUE_VALUES = {"1", "true", "yes", "on"} +_HF_METADATA_PROBE_TTL_SECONDS = 5.0 +_hf_metadata_probe_lock = threading.Lock() +_hf_metadata_probe_state: tuple[str, float, bool] | None = None def _env_offline() -> bool: @@ -57,6 +61,42 @@ def _env_offline() -> bool: ) +def _hf_metadata_unavailable() -> bool: + """Return whether Hub metadata should be served from cache only.""" + if _env_offline(): + return True + probe_enabled = os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() not in ( + "0", + "false", + "no", + "off", + ) + if not probe_enabled: + return False + + endpoint = os.environ.get("HF_ENDPOINT", "https://huggingface.co") + now = time.monotonic() + global _hf_metadata_probe_state + with _hf_metadata_probe_lock: + cached = _hf_metadata_probe_state + if cached is not None: + cached_endpoint, expires_at, unavailable = cached + if cached_endpoint == endpoint and now < expires_at: + return unavailable + + try: + from utils.transformers_version import hf_endpoint_unreachable + unavailable = hf_endpoint_unreachable(timeout = 1) + except Exception: + unavailable = False + _hf_metadata_probe_state = ( + endpoint, + time.monotonic() + _HF_METADATA_PROBE_TTL_SECONDS, + unavailable, + ) + return unavailable + + # ── Model size extraction ──────────────────────────────────── import re as _re @@ -1002,7 +1042,7 @@ _AUDIO_TOKEN_PATTERNS = { and "<|text_start|>" in tokens and "<|text_end|>" in tokens ), - "snac": lambda tokens: (sum(1 for t in tokens if t.startswith(" 10000), + "snac": lambda tokens: sum(1 for t in tokens if t.startswith(" 10000, } @@ -1809,11 +1849,16 @@ def list_gguf_variants( """ from huggingface_hub import model_info as hf_model_info - # Offline: skip the API and serve from cache - if _env_offline(): + # Offline / unreachable: skip the API and serve from cache. + if _hf_metadata_unavailable(): cached = _list_gguf_variants_from_hf_cache(repo_id) if cached is not None: return cached + logger.debug( + "Offline/unreachable HF Hub -- skipping GGUF variant listing for '%s'", + repo_id, + ) + return [], False try: info = hf_model_info(repo_id, token = hf_token, files_metadata = True) @@ -2015,10 +2060,15 @@ def detect_gguf_model_remote(repo_id: str, hf_token: Optional[str] = None) -> Op import time from huggingface_hub import model_info as hf_model_info - if _env_offline(): + if _hf_metadata_unavailable(): cached = _detect_gguf_from_hf_cache(repo_id) if cached is not None: return cached + logger.debug( + "Offline/unreachable HF Hub -- skipping remote GGUF detection for '%s'", + repo_id, + ) + return None last_err: Optional[Exception] = None for attempt in range(3): @@ -2120,18 +2170,23 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: Returns: True if embedding model, else False (default for local paths or errors). """ - from utils.utils import hf_env_offline + from utils.utils import hf_cache_snapshot_is_loadable, hf_env_offline + + cache_key = (model_name, hf_token) + + def cached_embedding_classification() -> bool: + if _embedding_marker_in_hf_cache(model_name): + return True + return _embedding_detection_cache.get(cache_key) is True and hf_cache_snapshot_is_loadable( + model_name + ) # Offline (remote repo): reclassify from the local cache on every call, before/without the # memo. An online lookup can memoize True from tags with no weights cached, so trusting it once # the session goes offline would accept a repo _get() cannot load; a cached negative can also be # invalidated by later cache materialization. The cache probe is local-only, so it's cheap. if not is_local_path(model_name) and hf_env_offline(): - return _embedding_marker_in_hf_cache(model_name) - - cache_key = (model_name, hf_token) - if cache_key in _embedding_detection_cache: - return _embedding_detection_cache[cache_key] + return cached_embedding_classification() # Local paths: check for sentence-transformer marker (modules.json) if is_local_path(model_name): @@ -2140,6 +2195,16 @@ def is_embedding_model(model_name: str, hf_token: Optional[str] = None) -> bool: _embedding_detection_cache[cache_key] = is_emb return is_emb + if _hf_metadata_unavailable(): + logger.debug( + "Offline/unreachable HF Hub -- using cached embedding marker for %s", + model_name, + ) + return cached_embedding_classification() + + if cache_key in _embedding_detection_cache: + return _embedding_detection_cache[cache_key] + try: from huggingface_hub import model_info as hf_model_info