diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py index 4d73213d1..2fff744b6 100644 --- a/studio/backend/tests/test_hf_xet_fallback.py +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -287,7 +287,9 @@ def test_degrades_when_shared_helper_import_raises_importerror(): 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.""" + retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails. + The backend loads lazily (first use of a heavy helper), so this triggers the load explicitly + before asserting the retry/degrade behavior.""" import importlib import os @@ -321,11 +323,15 @@ def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): 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. + # Import is light (lazy backend); unsloth_zoo not loaded yet. + assert seen_env == [], seen_env + # First use of a heavy helper triggers the load (attempt without the light env, then a retry + # with it set); accessing DownloadStallError drives it via __getattr__. + stall_error = degraded.DownloadStallError 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 issubclass(stall_error, RuntimeError) + # The env override must not leak past the load. assert os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") is None finally: sys.meta_path.remove(finder) @@ -333,3 +339,31 @@ def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): sys.modules.update(saved) if saved_shim is not None: sys.modules["utils.hf_xet_fallback"] = saved_shim + + +def test_importing_child_should_disable_xet_stays_light(monkeypatch): + """Regression guard for the stale-transformers-sidecar bug: importing the shim (and + ``child_should_disable_xet``) must NOT pull in ``transformers``/``unsloth_zoo``. The worker calls + this at startup to decide the Xet env flip BEFORE activating the sidecar; an eager import here + would cache the default transformers 4.57.x in sys.modules, defeating the sidecar sys.path prepend + and breaking 5.x models (Qwen3.5/GLM/gemma-4).""" + import importlib + + for name in [ + m + for m in list(sys.modules) + if m == "transformers" + or m.startswith("transformers.") + or m == "unsloth_zoo" + or m.startswith("unsloth_zoo.") + or m == "utils.hf_xet_fallback" + ]: + monkeypatch.delitem(sys.modules, name, raising = False) + + mod = importlib.import_module("utils.hf_xet_fallback") + # The lightweight decision works without the heavy backend. + assert mod.child_should_disable_xet({"disable_xet": True}) is True + assert mod.child_should_disable_xet({}) is False + # And nothing heavy was imported as a side effect. + assert "transformers" not in sys.modules, "importing the shim must not import transformers" + assert "unsloth_zoo" not in sys.modules, "importing the shim must not import unsloth_zoo" diff --git a/studio/backend/tests/test_training_worker_import_discipline.py b/studio/backend/tests/test_training_worker_import_discipline.py new file mode 100644 index 000000000..a047c9170 --- /dev/null +++ b/studio/backend/tests/test_training_worker_import_discipline.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Invariant: the training worker must not import ``transformers`` before it activates the +transformers sidecar. + +``core/training/worker.py:run_training_process`` runs a preflight (Xet decision, logging, hardware +detection) and only THEN calls ``_activate_transformers_version`` -> ``activate_transformers_for_subprocess``, +which prepends the correct ``.venv_t5_*`` (5.x) sidecar to ``sys.path``. Because activation only edits +``sys.path``, it is a no-op for any module already cached in ``sys.modules``. So if the preflight imports +``transformers`` (directly or transitively via ``unsloth_zoo``), the default 4.57.x gets pinned before +the sidecar is on the path -- and 5.x models (Qwen3.5, GLM-4.7, gemma-4) then fail to load their +tokenizer/config ("Tokenizer class TokenizersBackend does not exist"). + +This regression shipped once when ``utils/hf_xet_fallback.py`` eagerly imported ``unsloth_zoo`` (which +imports ``transformers``) at module load; the worker imports that shim during preflight to decide the +Xet env flip (see issue #6951). This test locks the invariant in a fresh interpreter. It is CPU-only, +needs no network/GPU/weights/sidecars, so it runs in the standard ``studio-backend-ci`` matrix. +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend + +# Mirrors run_training_process's imports that run BEFORE _activate_transformers_version (worker.py); +# keep in sync. torch-dependent imports are optional (a no-torch CI shard skips them) but must still +# not drag in transformers. +_PREFLIGHT_SNIPPET = r""" +import sys + +# worker.py: from utils.hf_xet_fallback import child_should_disable_xet (+ call it) +from utils.hf_xet_fallback import child_should_disable_xet +child_should_disable_xet({}) + +# worker.py: from loggers.config import LogConfig +from loggers.config import LogConfig # noqa: F401 + +# worker.py: from utils.hardware import hardware (imports torch, not transformers) +try: + from utils.hardware import hardware as _hw # noqa: F401 +except Exception: + pass # torch may be absent in a no-torch shard; the invariant below still applies + +# worker.py: from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend +# (the MLX-dispatch preflight; must also stay clear of transformers). Guarded because it may pull +# unsloth/trl, absent in a minimal shard -- but a partial import that leaked transformers would still +# be caught by the assertion below. +try: + from core.training.training import ( # noqa: F401 + is_apple_silicon_training_platform as _is_apple, + should_use_mlx_training_backend as _use_mlx, + ) +except Exception: + pass + +leaked_tf = sorted(m for m in sys.modules if m == "transformers" or m.startswith("transformers.")) +leaked_zoo = sorted(m for m in sys.modules if m == "unsloth_zoo" or m.startswith("unsloth_zoo.")) +assert not leaked_tf, f"transformers imported during worker preflight (before sidecar activation): {leaked_tf}" +assert not leaked_zoo, f"unsloth_zoo imported during worker preflight (before sidecar activation): {leaked_zoo}" +print("PREFLIGHT_CLEAN") +""" + + +def test_worker_preflight_does_not_import_transformers(): + """A fresh interpreter running the worker's pre-activation imports must leave ``transformers`` + (and ``unsloth_zoo``) unimported, so the 5.x sidecar prepend is not defeated by a stale module.""" + result = subprocess.run( + [sys.executable, "-c", _PREFLIGHT_SNIPPET], + cwd = str(_BACKEND_DIR), + capture_output = True, + text = True, + ) + assert result.returncode == 0, ( + "Worker preflight imported transformers before sidecar activation.\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + assert "PREFLIGHT_CLEAN" in result.stdout, result.stdout diff --git a/studio/backend/tests/test_worker_activates_correct_transformers.py b/studio/backend/tests/test_worker_activates_correct_transformers.py new file mode 100644 index 000000000..fe7b8dd25 --- /dev/null +++ b/studio/backend/tests/test_worker_activates_correct_transformers.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Invariant: after the training worker runs its preflight and then activates the transformers +sidecar, the in-process ``transformers`` must be the sidecar version the model requires -- not the +default 4.57.x that the base environment ships. + +The CPU-only "does it choose the correct transformers version" guard, stronger than the pure +import-order check in ``test_training_worker_import_discipline.py``: it runs the REAL tier detection +(``get_transformers_tier``) and REAL activation (``activate_transformers_for_subprocess``) for a +transformers-5.x model (Qwen3.5, tier 530) and asserts the version actually switched. It catches the +whole failure family at once: + + * a stale pre-activation ``transformers`` import (the #6951 / ``TokenizersBackend`` regression: an + already-cached 4.57.x defeats the sidecar's ``sys.path`` prepend), + * a wrong tier selected for a 5.x model, and + * activation not actually swapping the resident module. + +Why the CUDA spoof matters (verified): ``unsloth_zoo``'s eager ``import transformers`` only happens on +its full, GPU-present init path. On a GPU-less runner it silently degrades and never preloads +transformers -- which would MASK the stale-import bug (the check would falsely pass). Spoofing +``torch.cuda`` so ``unsloth_zoo`` believes a GPU is present forces the real init path, exposing the +regression on CPU CI. The spoof mirrors ``tests/_zoo_aggressive_cuda_spoof.py`` but is inlined so the +test is self-contained in the ``studio-backend-ci`` matrix (whose conftest does not apply the shared +spoof). No GPU/network/weights/real sidecar needed: a one-line stub sidecar stands in for the 5.x venv, +so we only assert activation lands on it. + +Proven: passes on the fixed tree (active == 5.3.0) and fails on the buggy tree (active == 4.57.x) on +a simulated GPU-less runner. +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend +# Canonical CUDA spoof at the repo root (studio/backend -> studio -> repo root). Loaded by the +# subprocess when present (matches the consolidated CI); absent in a standalone studio checkout, where +# the subprocess falls back to a minimal inline spoof. +_SPOOF_PATH = _BACKEND_DIR.parent.parent / "tests" / "_zoo_aggressive_cuda_spoof.py" + +# Runs in a fresh interpreter with cwd == studio/backend so ``utils.*`` resolves like the worker. +# STUB_HOME (a pytest tmp dir) holds a throwaway ``.venv_t5_530`` sidecar exporting transformers 5.3.0. +_SNIPPET = r""" +import os, sys +sys.path.insert(0, os.getcwd()) + +# CUDA spoof so unsloth_zoo takes its full, transformers-importing init path on a GPU-less runner. +# Without it unsloth_zoo degrades and never preloads transformers, which would MASK the stale-import +# regression under test (verified). Prefer the repo's canonical spoof (single source of truth, and the +# one the consolidated CI already relies on); fall back to a minimal inline spoof so this also works in +# a standalone studio checkout. If torch is absent the fixed tree still passes below; the bug just +# would not be exposable in that shard. +try: + import torch # noqa: F401 + _sp = os.environ.get("SPOOF_PATH") + if _sp and os.path.exists(_sp): + import importlib.util + _spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", _sp) + _mod = importlib.util.module_from_spec(_spec) + _spec.loader.exec_module(_mod) + _mod.apply() + else: + torch.cuda.is_available = lambda: True + torch.cuda.device_count = lambda: 1 + torch.cuda.current_device = lambda: 0 + torch.cuda.get_device_capability = lambda *a, **k: (8, 0) + torch.cuda.get_device_name = lambda *a, **k: "NVIDIA A100-SPOOFED" + torch.cuda.is_bf16_supported = lambda *a, **k: True + class _Props: + name = "NVIDIA A100-SPOOFED" + major = 8 + minor = 0 + total_memory = 80 * 1024**3 + multi_processor_count = 108 + torch.cuda.get_device_properties = lambda *a, **k: _Props() + torch.cuda.mem_get_info = lambda *a, **k: (0, 80 * 1024**3) +except Exception: + pass +os.environ["UNSLOTH_IS_PRESENT"] = "1" + +# Stub 5.x sidecar: activation only edits sys.path, so a package that merely exports __version__ is +# enough to prove the resident transformers switched to it. +home = os.environ["STUB_HOME"] +pkg = os.path.join(home, ".venv_t5_530", "transformers") +os.makedirs(pkg, exist_ok = True) +with open(os.path.join(pkg, "__init__.py"), "w") as f: + f.write('__version__ = "5.3.0"\n') +os.environ["UNSLOTH_STUDIO_HOME"] = home + +# Faithful worker preflight (worker.py: from utils.hf_xet_fallback import child_should_disable_xet). +# This is the exact stale-import trigger: on the buggy tree it pulls unsloth_zoo -> transformers 4.57.x +# into sys.modules BEFORE activation. +from utils.hf_xet_fallback import child_should_disable_xet +child_should_disable_xet({}) +_tf = sys.modules.get("transformers") +preload = _tf.__version__ if _tf is not None else None + +# Real tier detection + real activation, with the 530 sidecar pointed at the stub above. +import utils.transformers_version as tv +tv._VENV_T5_530_DIR = os.path.join(home, ".venv_t5_530") +tv._ensure_venv_t5_530_exists = lambda: True +tier = tv.get_transformers_tier("Qwen/Qwen3.5-9B", None) +tv.activate_transformers_for_subprocess("Qwen/Qwen3.5-9B", None) + +import transformers +print(f"RESULT tier={tier} preload={preload} active={transformers.__version__}") +""" + + +def _parse(stdout: str) -> dict[str, str]: + for line in stdout.splitlines(): + if line.startswith("RESULT "): + return dict(kv.split("=", 1) for kv in line.split()[1:]) + return {} + + +def test_worker_activates_correct_transformers_version(tmp_path): + """The worker's real preflight + activation for a transformers-5.x model (Qwen3.5, tier 530) must + leave the in-process ``transformers`` on the 5.x sidecar. A stale pre-activation import leaves the + default 4.57.x pinned and fails this assertion -- exactly the #6951 ``TokenizersBackend`` regression.""" + result = subprocess.run( + [sys.executable, "-c", _SNIPPET], + cwd = str(_BACKEND_DIR), + env = { + **__import__("os").environ, + "STUB_HOME": str(tmp_path), + **({"SPOOF_PATH": str(_SPOOF_PATH)} if _SPOOF_PATH.exists() else {}), + }, + capture_output = True, + text = True, + ) + assert result.returncode == 0, ( + "Worker preflight + activation harness crashed.\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + parsed = _parse(result.stdout) + assert parsed, f"No RESULT line.\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + + # Correct tier chosen for a transformers-5.x model (pure, deterministic; no network/GPU). + assert parsed["tier"] == "530", ( + f"Wrong transformers tier for Qwen3.5 (expected 530, got {parsed['tier']}). " + "Tier detection regressed." + ) + + # Activation must actually swap the resident transformers to the sidecar version. If a preflight + # import cached 4.57.x first, the sidecar prepend is a no-op and this stays 4.57.x -- the bug. + assert parsed["active"] == "5.3.0", ( + "Sidecar activation did NOT switch the in-process transformers to the model's 5.x version " + f"(active={parsed['active']}, preloaded-before-activation={parsed['preload']}). A pre-activation " + "transformers import (directly or via unsloth_zoo) defeated the sidecar; 5.x models (Qwen3.5, " + "GLM-4.7, gemma-4) then fail with 'Tokenizer class TokenizersBackend does not exist'. See #6951." + ) diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py index 2dd224739..9bc4a60fa 100644 --- a/studio/backend/utils/hf_xet_fallback.py +++ b/studio/backend/utils/hf_xet_fallback.py @@ -6,6 +6,16 @@ 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. + +Import discipline: ``unsloth_zoo``'s ``__init__`` eagerly imports ``transformers``. The workers +import this shim at startup (to decide the per-worker Xet env flip) *before* activating the model's +``transformers`` sidecar. Activation only prepends the sidecar to ``sys.path``, so a ``transformers`` +already cached in ``sys.modules`` (via an eager ``unsloth_zoo`` import here) wins -- pinning the +default 4.57.x and regressing Qwen3.5 / GLM-4.7 / gemma-4 training with +``Tokenizer class TokenizersBackend does not exist``. So the shared backend is loaded **lazily** +(``_load_shared``), only on first use of a heavy download helper, i.e. after the sidecar is active. +``child_should_disable_xet`` and the ``DEFAULT_*`` constants are defined locally so importing them +never triggers the heavy load. """ from __future__ import annotations @@ -13,161 +23,230 @@ from __future__ import annotations import threading from typing import Any, Callable, Optional -_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 +# Defaults mirror unsloth_zoo.hf_xet_fallback; plain literals so they resolve (including as +# default args below) without importing unsloth_zoo/transformers. +DEFAULT_GRACE_PERIOD = 10.0 +DEFAULT_HEARTBEAT_INTERVAL = 30.0 +DEFAULT_STALL_TIMEOUT = 180.0 - _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") - _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" - try: - 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 +# --- lazy shared-backend loader ---------------------------------------------------------------- +_shared: Any = None +_shared_available: Optional[bool] = None # None = not yet attempted +_shared_import_error: Optional[BaseException] = None +_load_lock = threading.Lock() -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, - ) +def _load_shared() -> bool: + """Import ``unsloth_zoo.hf_xet_fallback`` on demand; return True if available. Deferred so + importing this module at worker startup does not pull transformers in before the sidecar is + activated. Degrades (returns False) rather than crashing when unsloth_zoo is unavailable.""" + global _shared, _shared_available, _shared_import_error + if _shared_available is not None: + return _shared_available + with _load_lock: + if _shared_available is not None: + return _shared_available + try: + import unsloth_zoo.hf_xet_fallback as shared - DEFAULT_HEARTBEAT_INTERVAL = 30.0 - DEFAULT_STALL_TIMEOUT = 180.0 - DEFAULT_GRACE_PERIOD = 10.0 + _shared = shared + _shared_available = True + _shared_import_error = None + return 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 + # host. The download helper needs none of it, so retry via UNSLOTH_ZOO_DISABLE_GPU_INIT. + _shared_import_error = exc + import os as _os - class DownloadStallError(RuntimeError): - """Stub mirror so callers' ``except`` clauses resolve; never raised in degraded mode.""" + _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" + try: + import unsloth_zoo.hf_xet_fallback as shared - def child_should_disable_xet(config: dict) -> bool: - return bool(config.get("disable_xet")) + _shared = shared + _shared_available = True + _shared_import_error = None + return True + except Exception as exc2: # noqa: BLE001 - degrade so Studio still boots with plain HF + _shared_import_error = exc2 + _shared_available = False + import logging as _logging - def get_hf_download_state(*args: Any, **kwargs: Any) -> None: - return None # unmeasurable -> the (absent) watchdog never fires + _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, + ) + return 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 - 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)...") - except Exception: - pass +def child_should_disable_xet(config: dict) -> bool: + """Single source of truth for the per-worker Xet env flip (mirrors + ``unsloth_zoo.hf_xet_fallback.child_should_disable_xet``). Deliberately lightweight: importing or + calling it must NOT pull in unsloth_zoo/transformers, so the worker can decide before activating + the transformers sidecar (see the module docstring).""" + return bool(config.get("disable_xet")) - threading.Thread( - target = _beat, - daemon = True, - name = "hf-xet-degraded-heartbeat", - ).start() + +# --- degraded stubs (used only when unsloth_zoo is unavailable) ------------------------------- +class _DegradedDownloadStallError(RuntimeError): + """Stub mirror so callers' ``except`` clauses resolve; never raised in degraded mode.""" + + +def _degraded_get_hf_download_state(*args: Any, **kwargs: Any) -> None: + return None # unmeasurable -> the (absent) watchdog never fires + + +def _degraded_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 _degraded_cancelled(cancel_event: "Optional[threading.Event]") -> bool: - return cancel_event is not None and cancel_event.is_set() + def _beat() -> None: + while not stop.wait(interval): + try: + on_heartbeat(f"Downloading ({transport} transport)...") + except Exception: + pass - 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") + threading.Thread( + target = _beat, + daemon = True, + name = "hf-xet-degraded-heartbeat", + ).start() + return stop - from huggingface_hub import hf_hub_download - path = hf_hub_download( - repo_id = repo_id, - filename = filename, - token = token, - repo_type = repo_type, - revision = revision, - cache_dir = cache_dir, - force_download = force_download, - ) - if _degraded_cancelled(cancel_event): - raise RuntimeError("Cancelled") - return path +def _degraded_cancelled(cancel_event: "Optional[threading.Event]") -> bool: + return cancel_event is not None and cancel_event.is_set() - 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 _degraded_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") - path = snapshot_download( - repo_id = repo_id, - repo_type = repo_type, - revision = revision, - token = token, - cache_dir = cache_dir, - allow_patterns = allow_patterns, - ignore_patterns = ignore_patterns, - force_download = force_download, - ) - if _degraded_cancelled(cancel_event): - raise RuntimeError("Cancelled") - return path + from huggingface_hub import hf_hub_download + + path = hf_hub_download( + repo_id = repo_id, + filename = filename, + token = token, + repo_type = repo_type, + revision = revision, + cache_dir = cache_dir, + force_download = force_download, + ) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + + +def _degraded_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 + + path = snapshot_download( + repo_id = repo_id, + repo_type = repo_type, + revision = revision, + token = token, + cache_dir = cache_dir, + allow_patterns = allow_patterns, + ignore_patterns = ignore_patterns, + force_download = force_download, + ) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + + +# --- lazy attribute access for the heavy shared API ------------------------------------------- +# ``DownloadStallError`` (class identity matters for ``except``), ``start_watchdog`` and +# ``get_hf_download_state`` come from the shared backend when available, else the degraded stubs. +# Resolved via PEP 562 ``__getattr__`` so ``from utils.hf_xet_fallback import X`` triggers the load +# only for these heavy names, not for ``child_should_disable_xet`` / ``DEFAULT_*``. +_DEGRADED_ATTRS = { + "DownloadStallError": _DegradedDownloadStallError, + "start_watchdog": _degraded_start_watchdog, + "get_hf_download_state": _degraded_get_hf_download_state, +} + +# Annotation-only declarations for the three names above: they bind NO value, so lookup still misses +# and PEP 562 ``__getattr__`` resolves them lazily -- but ruff/pyflakes see them as defined, so listing +# them in ``__all__`` does not trip F822 (while F822 still catches a real typo elsewhere in the list). +DownloadStallError: type +start_watchdog: Any +get_hf_download_state: Any + + +def __getattr__(name: str) -> Any: + if name in _DEGRADED_ATTRS: + if _load_shared(): + return getattr(_shared, name) + return _DEGRADED_ATTRS[name] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +# Indirection seam the public wrappers call (and tests monkeypatch): lazy-load the shared backend, +# then dispatch to it or the degraded stub. The ``_shared_*`` names preserve the pre-refactor contract. +def _shared_hf_hub_download_with_xet_fallback(*args: Any, **kwargs: Any) -> str: + impl = ( + _shared.hf_hub_download_with_xet_fallback + if _load_shared() + else _degraded_hf_hub_download_with_xet_fallback + ) + return impl(*args, **kwargs) + + +def _shared_snapshot_download_with_xet_fallback(*args: Any, **kwargs: Any) -> str: + impl = ( + _shared.snapshot_download_with_xet_fallback + if _load_shared() + else _degraded_snapshot_download_with_xet_fallback + ) + return impl(*args, **kwargs) __all__ = [