diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 0f4abe04f5..7b6c0b0668 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -51,6 +51,7 @@ from core.tool_healing import ( strip_tool_call_markup, ) from utils.native_path_leases import child_env_without_native_path_secret +from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) @@ -2996,7 +2997,7 @@ class LlamaCppBackend: any time; checks it between each shard download. """ try: - from huggingface_hub import hf_hub_download + import huggingface_hub # noqa: F401 -- presence check only except ImportError: raise RuntimeError( "huggingface_hub is required for HF model loading. " @@ -3140,19 +3141,23 @@ class LlamaCppBackend: if self._cancel_event.is_set(): raise RuntimeError("Cancelled") dl_start = time.monotonic() - local_path = hf_hub_download( - repo_id = hf_repo, - filename = gguf_filename, - token = hf_token, + # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached. + local_path = hf_hub_download_with_xet_fallback( + hf_repo, + gguf_filename, + hf_token, + cancel_event = self._cancel_event, + on_status = lambda m: logger.info(m), ) for shard in gguf_extra_shards: if self._cancel_event.is_set(): raise RuntimeError("Cancelled") logger.info(f"Resolving GGUF shard: {shard}") - hf_hub_download( - repo_id = hf_repo, - filename = shard, - token = hf_token, + hf_hub_download_with_xet_fallback( + hf_repo, + shard, + hf_token, + cancel_event = self._cancel_event, ) except RuntimeError as e: if "Cancelled" in str(e): @@ -3213,12 +3218,13 @@ class LlamaCppBackend: return None try: - from huggingface_hub import hf_hub_download logger.info(f"Downloading {label}: {hf_repo}/{target}") - return hf_hub_download( - repo_id = hf_repo, - filename = target, - token = hf_token, + # Same policy; companions are best-effort (caller below swallows failures to None). + return hf_hub_download_with_xet_fallback( + hf_repo, + target, + hf_token, + cancel_event = self._cancel_event, ) except Exception as e: logger.warning(f"Could not download {label}: {e}") diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 4fa48a2d88..6b6deb1265 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -30,15 +30,15 @@ from pathlib import Path from typing import Any, Generator, Optional, Tuple, Union from utils.hardware import prepare_gpu_selection +# Re-exported from the shared helper so GGUF, training, and inference share one +# type; kept importable here for backwards compatibility. +from utils.hf_xet_fallback import DownloadStallError + logger = get_logger(__name__) _CTX = mp.get_context("spawn") -class DownloadStallError(RuntimeError): - """Raised when the worker reports no download progress for too long.""" - - # Dispatcher timeout constants (seconds) _DISPATCH_READ_TIMEOUT = 30.0 _DISPATCH_POLL_INTERVAL = 0.5 diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index d69246f584..0c3d42a2fe 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -217,6 +217,12 @@ class TrainingBackend: self._db_config: Optional[dict] = None self._db_started_at: Optional[str] = None + # Xet -> HTTP model-load fallback state (config kept for the respawn). + self._last_full_config: Optional[dict] = None + self._in_model_load: bool = False + self._xet_fallback_used: bool = False + self._needs_xet_respawn: bool = False + logger.info("TrainingBackend initialized (subprocess mode)") # ------------------------------------------------------------------ @@ -311,6 +317,8 @@ class TrainingBackend: "trust_remote_code": kwargs.get("trust_remote_code", False), "gpu_ids": kwargs.get("gpu_ids"), "s3_config": kwargs.get("s3_config"), + # Flipped to True only by the HTTP-fallback respawn after a stall. + "disable_xet": kwargs.get("disable_xet", False), } # Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request. @@ -386,6 +394,11 @@ class TrainingBackend: self._db_total_steps_set = False self._db_config = _sanitize_db_config(config) self._db_started_at = datetime.now(timezone.utc).isoformat() + # Start each job Xet-first; keep config so a stall can respawn over HTTP. + self._last_full_config = config + self._in_model_load = False + self._xet_fallback_used = False + self._needs_xet_respawn = False # Assign subprocess handles after state reset. self._event_queue = event_queue @@ -446,6 +459,97 @@ class TrainingBackend: output_dir, ) + def _handle_stall_event(self, event: dict) -> None: + """A worker reported a no-progress download stall. + + On the first model-load, terminate the worker so the pump loop respawns it + over HTTP. A later stall (already on HTTP, or outside model-load) surfaces + as an error instead. + """ + msg = event.get("message", "Download stalled") + with self._lock: + recover = self._in_model_load and not self._xet_fallback_used + proc = self._proc + if recover: + self._xet_fallback_used = True + self._needs_xet_respawn = True + self._progress.status_message = ( + "Model download stalled on Xet; retrying over HTTP..." + ) + else: + self._progress.error = self._progress.error or ( + "Model download stalled even over HTTP -- check your network connection" + ) + if recover: + logger.warning("Training model-load stalled on Xet; respawning over HTTP: %s", msg) + else: + logger.error("Training download stalled with no further fallback: %s", msg) + # Terminate either way so the pump loop proceeds (respawn or finalize). + if proc is not None and proc.is_alive(): + proc.terminate() + + def _respawn_worker_disable_xet(self) -> None: + """Respawn the worker once with HF_HUB_DISABLE_XET=1 after a model-load + stall. Runs on the exiting pump thread, reaps the terminated worker, and + starts a fresh worker + pump. DB/progress run-state is preserved so the + history row is not duplicated; the new worker re-formats and loads over HTTP. + """ + config = self._last_full_config + if config is None: + logger.error("Cannot respawn training worker: no stored config") + return + + with self._lock: + old_proc = self._proc + if old_proc is not None: + old_proc.join(timeout = 5.0) + if old_proc.is_alive(): + old_proc.kill() + old_proc.join(timeout = 2.0) + + config = {**config, "disable_xet": True} + self._last_full_config = config + logger.warning("Respawning training worker with HF_HUB_DISABLE_XET=1 after Xet stall") + + from .worker import run_training_process + + try: + with native_path_secret_removed_for_child_start(): + event_queue = _CTX.Queue() + stop_queue = _CTX.Queue() + new_proc = _CTX.Process( + target = run_without_native_path_secret, + args = (run_training_process,), + kwargs = { + "event_queue": event_queue, + "stop_queue": stop_queue, + "config": config, + }, + daemon = True, + ) + new_proc.start() + except Exception: + logger.error("Failed to respawn training subprocess", exc_info = True) + with self._lock: + self._progress.is_training = False + self._progress.error = "Failed to recover stalled model download" + self._ensure_db_run_created() + self._finalize_run_in_db( + status = "error", + error_message = "Failed to recover stalled model download", + ) + return + + logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid) + new_pump = threading.Thread(target = self._pump_loop, daemon = True) + with self._lock: + self._in_model_load = False + self._event_queue = event_queue + self._stop_queue = stop_queue + self._proc = new_proc + self._pump_thread = new_pump + new_pump.start() + def is_training_active(self) -> bool: """Check if training is currently active.""" with self._lock: @@ -566,6 +670,14 @@ class TrainingBackend: for e in self._drain_queue(self._event_queue): self._handle_event(e) + # Model-load stall: respawn over HTTP instead of finalizing as failure. + # Runs on THIS exiting pump thread and starts a fresh pump (never joins + # the current thread); DB run-state is preserved. + if self._needs_xet_respawn: + self._needs_xet_respawn = False + self._respawn_worker_disable_xet() + return + # Mark done if no explicit complete/error was received. with self._lock: if self._progress.is_training: @@ -597,6 +709,19 @@ class TrainingBackend: db_action: Optional[str] = None db_action_kwargs: dict = {} + # Model-load lifecycle + stall recovery (no DB metrics); handled first. + if etype == "model_load_started": + with self._lock: + self._in_model_load = True + return + if etype == "model_load_completed": + with self._lock: + self._in_model_load = False + return + if etype == "stall": + self._handle_stall_event(event) + return + with self._lock: if etype == "progress": self._progress.step = event.get("step", self._progress.step) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index b106aa1298..65533b6946 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1987,6 +1987,19 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> os.environ["TOKENIZERS_PARALLELISM"] = "false" os.environ["PYTHONWARNINGS"] = "ignore" # before imports + # HTTP-fallback respawn: disable Xet before any huggingface_hub import (the + # var is read at import time). Mirrors core/inference/worker.py. + from utils.hf_xet_fallback import child_should_disable_xet + + if child_should_disable_xet(config): + os.environ["HF_HUB_DISABLE_XET"] = "1" + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" + print( + "Xet transport disabled for this training worker (HF_HUB_DISABLE_XET=1).", + file = sys.stderr, + flush = True, + ) + # Offline auto-detect: skip ~25s of HF retries per call when DNS is dead. if "HF_HUB_OFFLINE" not in os.environ: import socket as _socket @@ -2706,18 +2719,33 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> cpt_trains_embeddings = False # ── 4c. Load training model (uses VRAM — dataset already formatted) ── + # Watchdog lets the parent recover a stalled Xet download via respawn. _send_status(event_queue, "Loading model...") - success = trainer.load_model( - model_name = model_name, - max_seq_length = config["max_seq_length"], - load_in_4bit = config["load_in_4bit"], - full_finetuning = not use_lora, - hf_token = hf_token, - is_dataset_image = config.get("is_dataset_image", False), - is_dataset_audio = config.get("is_dataset_audio", False), - trust_remote_code = config.get("trust_remote_code", False), - gpu_ids = config.get("resolved_gpu_ids"), + from utils.hf_xet_fallback import start_watchdog + + event_queue.put({"type": "model_load_started", "ts": time.time()}) + _load_watchdog_stop = start_watchdog( + repo_ids = [model_name], + on_stall = lambda msg: event_queue.put( + {"type": "stall", "message": msg, "ts": time.time()} + ), + xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1", ) + try: + success = trainer.load_model( + model_name = model_name, + max_seq_length = config["max_seq_length"], + load_in_4bit = config["load_in_4bit"], + full_finetuning = not use_lora, + hf_token = hf_token, + is_dataset_image = config.get("is_dataset_image", False), + is_dataset_audio = config.get("is_dataset_audio", False), + trust_remote_code = config.get("trust_remote_code", False), + gpu_ids = config.get("resolved_gpu_ids"), + ) + finally: + _load_watchdog_stop.set() + event_queue.put({"type": "model_load_completed", "ts": time.time()}) if not success or trainer.should_stop: if trainer.should_stop: event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) diff --git a/studio/backend/tests/test_gguf_xet_fallback_integration.py b/studio/backend/tests/test_gguf_xet_fallback_integration.py new file mode 100644 index 0000000000..cbcc73847e --- /dev/null +++ b/studio/backend/tests/test_gguf_xet_fallback_integration.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Integration: GGUF Chat-Mode downloads route through the Xet->HTTP helper, +preserving cancellation and the best-effort companion contract. No GPU, no +network, no real subprocess (the helper is patched). +""" + +from __future__ import annotations + +import sys +import threading +import types as _types +from pathlib import Path +from unittest.mock import patch + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +# Heavy-dep stubbing; prefer the real structlog so a bare stub never leaks to +# 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) +try: + import structlog # noqa: F401 +except ImportError: + sys.modules["structlog"] = _types.ModuleType("structlog") +try: + import httpx # noqa: F401 +except ImportError: + _httpx_stub = _types.ModuleType("httpx") + for _exc in ( + "ConnectError", + "TimeoutException", + "ReadTimeout", + "ReadError", + "RemoteProtocolError", + "CloseError", + "HTTPError", + "RequestError", + "HTTPStatusError", + ): + setattr(_httpx_stub, _exc, type(_exc, (Exception,), {})) + _httpx_stub.Response = type("Response", (), {}) + _httpx_stub.Request = type("Request", (), {}) + _httpx_stub.Timeout = type("Timeout", (), {"__init__": lambda self, *a, **k: None}) + _httpx_stub.Client = type( + "Client", + (), + { + "__init__": lambda self, **k: None, + "__enter__": lambda self: self, + "__exit__": lambda self, *a: None, + }, + ) + sys.modules.setdefault("httpx", _httpx_stub) + +from huggingface_hub import constants as hf_constants + +from core.inference.llama_cpp import LlamaCppBackend +from utils.hf_xet_fallback import DownloadStallError + +REPO = "unsloth/vision-GGUF" + + +@pytest.fixture +def hf_cache(tmp_path, monkeypatch): + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) + return tmp_path + + +def _build_cache( + root: Path, + repo_id: str, + files: dict[str, int], + sha: str = "a" * 40, +) -> Path: + repo_dir = root / f"models--{repo_id.replace('/', '--')}" + (repo_dir / "blobs").mkdir(parents = True, exist_ok = True) + snap = repo_dir / "snapshots" / sha + snap.mkdir(parents = True, exist_ok = True) + for rel, size in files.items(): + (snap / rel).write_bytes(b"\0" * size) + return snap + + +def test_companion_routes_through_helper(hf_cache): + _build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1}) + backend = LlamaCppBackend() + captured = {} + + def fake_helper( + repo_id, + filename, + token = None, + **kwargs, + ): + captured["filename"] = filename + captured["cancel_event"] = kwargs.get("cancel_event") + return f"/fake/{filename}" + + with ( + patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_helper), + ): + out = backend._download_mmproj(hf_repo = REPO, hf_token = None) + + assert out == "/fake/mmproj-vision-F16.gguf" + # _cancel_event must be threaded through so /unload can abort the download. + assert captured["cancel_event"] is backend._cancel_event + + +def test_companion_swallows_terminal_stall_to_none(hf_cache): + _build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1}) + backend = LlamaCppBackend() + + def stalling_helper( + repo_id, + filename, + token = None, + **kwargs, + ): + raise DownloadStallError("both transports stalled") + + with ( + patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", stalling_helper), + ): + out = backend._download_mmproj(hf_repo = REPO, hf_token = None) + + assert out is None, "a companion download is best-effort; a terminal stall must not raise" + + +def test_companion_cancelled_skips_download(hf_cache): + _build_cache(hf_cache, REPO, {"mmproj-vision-F16.gguf": 1}) + backend = LlamaCppBackend() + backend._cancel_event.set() + called = {"n": 0} + + def helper( + repo_id, + filename, + token = None, + **kwargs, + ): + called["n"] += 1 + return "/should-not-happen" + + with ( + patch("huggingface_hub.list_repo_files", lambda *a, **k: ["mmproj-vision-F16.gguf"]), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", helper), + ): + out = backend._download_mmproj(hf_repo = REPO, hf_token = None) + + assert out is None + assert called["n"] == 0, "a cancelled load must not start a companion download" diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py new file mode 100644 index 0000000000..39ecebd328 --- /dev/null +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -0,0 +1,352 @@ +# 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). +""" + +from __future__ import annotations + +import subprocess +import sys +import threading +import time +import types as _types +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) + +# 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. +_loggers_stub = _types.ModuleType("loggers") +_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name) +sys.modules.setdefault("loggers", _loggers_stub) +try: + import structlog # noqa: F401 +except ImportError: + 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() + + +def test_get_state_empty_cache(hf_cache): + assert xf.get_hf_download_state([REPO]) == (0, False) + + +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) + + +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 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" + + +# --------------------------------------------------------------------------- # +# 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" + + +@pytest.fixture(autouse = True) +def _no_real_cache_hit(monkeypatch): + """Default: the cached probe misses; tests override it to force a hit.""" + monkeypatch.setattr(huggingface_hub, "try_to_load_from_cache", lambda *a, **k: None) + + +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, + repo_id, + filename, + token, + *, + repo_type, + disable_xet, + cancel_event, + stall_timeout, + interval, + grace_period, + on_status, + ): + 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] + + +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: + 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}" + ) diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index a9e0b98247..c00e98ad7e 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -639,17 +639,20 @@ class TestDownloadMmprojOfflineCacheFallback: raise OSError("offline") def fake_download( - *, repo_id, filename, token = None, + **kwargs, ): # Echo back so the test can verify the cache-resolved filename return f"/fake/cache/{repo_id}/{filename}" with ( patch("huggingface_hub.list_repo_files", boom_list), - patch("huggingface_hub.hf_hub_download", fake_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fake_download, + ), ): out = backend._download_mmproj( hf_repo = "unsloth/vision-GGUF", @@ -675,17 +678,20 @@ class TestDownloadMmprojOfflineCacheFallback: captured = {} def fake_download( - *, repo_id, filename, token = None, + **kwargs, ): captured["filename"] = filename return f"/fake/{filename}" with ( patch("huggingface_hub.list_repo_files", boom_list), - patch("huggingface_hub.hf_hub_download", fake_download), + patch( + "core.inference.llama_cpp.hf_hub_download_with_xet_fallback", + fake_download, + ), ): backend._download_mmproj( hf_repo = "unsloth/vision-GGUF", diff --git a/studio/backend/tests/test_training_xet_fallback.py b/studio/backend/tests/test_training_xet_fallback.py new file mode 100644 index 0000000000..b4a3864334 --- /dev/null +++ b/studio/backend/tests/test_training_xet_fallback.py @@ -0,0 +1,209 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Parent-side training Xet->HTTP fallback: a model-load stall respawns the +worker once with Xet disabled, preserving the DB run row. Driven via +_handle_event with a fake spawn context; no GPU, no network, no real subprocess. +""" + +from __future__ import annotations + +import contextlib +import logging +import queue +import sys +import threading +import types as _types +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) + +# Stub the heavy module-level imports of core/training/training.py so it imports +# under CPU-only/no-network, then restore them (see the restore loop below). +_SAVED: dict = {} + + +def _stub(name, mod): + _SAVED[name] = sys.modules.get(name) + sys.modules[name] = mod + + +_lg = _types.ModuleType("loggers") +_lg.get_logger = lambda name: logging.getLogger(name) +_stub("loggers", _lg) +_stub("structlog", _types.ModuleType("structlog")) +_mpl = _types.ModuleType("matplotlib") +_plt = _types.ModuleType("matplotlib.pyplot") +_plt.Figure = type("Figure", (), {}) # referenced in a class-def annotation +_mpl.pyplot = _plt +_stub("matplotlib", _mpl) +_stub("matplotlib.pyplot", _plt) +_hw = _types.ModuleType("utils.hardware") +_hw.prepare_gpu_selection = lambda *a, **k: (None, None) +_stub("utils.hardware", _hw) +_npl = _types.ModuleType("utils.native_path_leases") +_npl.native_path_secret_removed_for_child_start = lambda: contextlib.nullcontext() +_npl.run_without_native_path_secret = lambda fn: fn +_stub("utils.native_path_leases", _npl) +_pth = _types.ModuleType("utils.paths") +_pth.outputs_root = lambda *a, **k: "/tmp/outputs" +_stub("utils.paths", _pth) + +import core.training.training as training_mod +from core.training.training import TrainingBackend + +# Restore every stubbed module so this file never pollutes the shared session: a +# leaked bare ``structlog`` (no ``get_logger``) would break every later module +# that logs at import. training_mod already bound the stubs it needs at runtime. +for _name in ( + "loggers", + "structlog", + "matplotlib", + "matplotlib.pyplot", + "utils.hardware", + "utils.native_path_leases", + "utils.paths", +): + _prev = _SAVED.get(_name) + if _prev is None: + sys.modules.pop(_name, None) + else: + sys.modules[_name] = _prev + + +@pytest.fixture(autouse = True) +def _stub_worker_module(): + """Stub ``core.training.worker`` so the respawn's lazy import of the + torch-heavy worker is never required.""" + prev = sys.modules.get("core.training.worker") + stub = _types.ModuleType("core.training.worker") + stub.run_training_process = lambda **kwargs: None + sys.modules["core.training.worker"] = stub + yield + if prev is None: + sys.modules.pop("core.training.worker", None) + else: + sys.modules["core.training.worker"] = prev + + +class _FakeProc: + def __init__(self, **kwargs): + self._alive = True + self.pid = 4321 + self.kwargs = kwargs + + def start(self): + pass + + def is_alive(self): + return self._alive + + def terminate(self): + self._alive = False + + def kill(self): + self._alive = False + + def join(self, timeout = None): + self._alive = False + + +class _FakeQueue: + def put(self, *a, **k): + pass + + def get(self, *a, **k): + raise queue.Empty + + +class _FakeCtx: + def __init__(self): + self.spawned: list = [] + + def Queue(self): + return _FakeQueue() + + def Process(self, **kwargs): + self.spawned.append(kwargs) + return _FakeProc(**kwargs) + + +def _backend_mid_load(): + b = TrainingBackend() + b._last_full_config = {"model_name": "org/model", "disable_xet": False, "hf_token": "tok"} + b._in_model_load = True + b._xet_fallback_used = False + proc = _FakeProc() + b._proc = proc + return b, proc + + +def test_stall_during_load_arms_respawn_and_terminates_worker(): + b, proc = _backend_mid_load() + b._handle_event({"type": "stall", "message": "no progress for 180s"}) + assert b._needs_xet_respawn is True + assert b._xet_fallback_used is True + assert proc.is_alive() is False, "stalled worker must be terminated" + + +def test_respawn_uses_disable_xet_and_preserves_run_row(monkeypatch): + b, _ = _backend_mid_load() + b._handle_event({"type": "stall", "message": "x"}) + + fake_ctx = _FakeCtx() + monkeypatch.setattr(training_mod, "_CTX", fake_ctx) + monkeypatch.setattr(b, "_pump_loop", lambda: None) # neutralize the new pump + created = {"n": 0} + finalized = {"n": 0} + monkeypatch.setattr( + b, "_ensure_db_run_created", lambda: created.__setitem__("n", created["n"] + 1) + ) + monkeypatch.setattr( + b, "_finalize_run_in_db", lambda **k: finalized.__setitem__("n", finalized["n"] + 1) + ) + + b._respawn_worker_disable_xet() + + assert len(fake_ctx.spawned) == 1, "respawn must start exactly one worker" + cfg = fake_ctx.spawned[0]["kwargs"]["config"] + assert cfg["disable_xet"] is True, "respawned worker must run with Xet disabled" + assert cfg["model_name"] == "org/model" + assert created["n"] == 0, "respawn must not recreate the DB run row" + assert finalized["n"] == 0, "a successful respawn must not finalize the run as error" + + +def test_second_stall_surfaces_error_without_respawn(): + b, proc = _backend_mid_load() + b._xet_fallback_used = True # HTTP fallback already spent + b._handle_event({"type": "stall", "message": "stalled again over http"}) + assert b._needs_xet_respawn is False + assert b._progress.error and "stalled" in b._progress.error.lower() + assert proc.is_alive() is False + + +def test_model_load_completed_disarms_recovery(): + b, _ = _backend_mid_load() + b._handle_event({"type": "model_load_completed"}) + assert b._in_model_load is False + # A stall after the load finished is not a transport stall to recover from. + b._handle_event({"type": "stall", "message": "post-load"}) + assert b._needs_xet_respawn is False + + +def test_model_load_started_arms_recovery_window(): + b = TrainingBackend() + assert b._in_model_load is False + b._handle_event({"type": "model_load_started"}) + assert b._in_model_load is True + + +def test_child_should_disable_xet_truth_table(): + from utils.hf_xet_fallback import child_should_disable_xet + + assert child_should_disable_xet({"disable_xet": True}) is True + assert child_should_disable_xet({"disable_xet": False}) is False + assert child_should_disable_xet({}) is False diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py new file mode 100644 index 0000000000..80f42f92d1 --- /dev/null +++ b/studio/backend/utils/hf_xet_fallback.py @@ -0,0 +1,405 @@ +# 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. + +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. +""" + +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. + """ + 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() + + if state is None: + if on_heartbeat is not None: + 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, +) -> 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. + """ + 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) + + try: + from huggingface_hub import hf_hub_download + path = hf_hub_download( + repo_id = repo_id, + filename = filename, + repo_type = repo_type, + token = token, + ) + 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) + + +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]], +) -> 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( + repo_id = repo_id, + filename = filename, + token = token, + repo_type = repo_type, + disable_xet = disable_xet, + result_queue = result_queue, + ), + daemon = True, + ) + proc.start() + + 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") + + +def hf_hub_download_with_xet_fallback( + repo_id: str, + filename: str, + token: Optional[str], + *, + cancel_event: Optional[threading.Event] = None, + repo_type: str = "model", + stall_timeout: float = DEFAULT_STALL_TIMEOUT, + interval: float = DEFAULT_HEARTBEAT_INTERVAL, + grace_period: float = DEFAULT_GRACE_PERIOD, + on_status: Optional[Callable[[str], None]] = None, +) -> 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. + """ + # Finalized blob already cached: return it with no child and no network. + 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, + ) + + 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}'")