ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: Xet-primary model downloads with automatic HTTP fallback on stall (#6372)

Browse source 
* Studio: add shared Xet-primary download helper with HTTP stall fallback

Xet is the fast default transport in huggingface_hub, but a stalled Xet
transfer hangs with no progress and no exception, and a blocked native thread
cannot be killed. The safetensors inference path already recovers (subprocess
watchdog + respawn with HF_HUB_DISABLE_XET=1); the GGUF and training paths do
not. Add a reusable helper that the in-process paths can adopt.

utils/hf_xet_fallback.py:
- DownloadStallError (moved here from core/inference/orchestrator.py, which now
  imports it; behavior unchanged, still a RuntimeError subclass).
- get_hf_download_state / start_watchdog: a no-progress watchdog built on the
  sparse-aware hub.utils.hf_cache_state helpers; fires only while a .incomplete
  is present and the on-disk byte total is unchanged for stall_timeout.
- hf_hub_download_with_xet_fallback: cached files short-circuit; otherwise the
  download runs in a spawn child (own process group) supervised by the watchdog.
  On a stall it kills the child, makes the partial safe for HTTP via
  prepare_cache_for_transport, and respawns once with HF_HUB_DISABLE_XET=1. Cancel
  and deterministic errors (auth/missing/disk) propagate without a fallback.

Tests cover the watchdog state machine, the transport decision logic, and a
regression lock that HF_HUB_DISABLE_XET is honored in a fresh interpreter.

* Studio: route GGUF Chat-Mode downloads through the Xet->HTTP fallback

The GGUF load path (_download_gguf main+shards, _download_companion_gguf for
mmproj/MTP) called a bare blocking hf_hub_download with no recovery, so a Xet
stall hung the Chat-Mode load with no fallback. Route those three calls through
hf_hub_download_with_xet_fallback: Xet stays primary, HTTP is used only if Xet
stalls, per-file so finished shards stay cached. The existing _cancel_event is
threaded through, the Cancelled sentinel is preserved, and companions stay
best-effort (a terminal stall is swallowed to None). Cached files short-circuit
in the helper with no subprocess, so the fast path is unchanged.

The two offline mmproj tests are repointed from huggingface_hub.hf_hub_download
to the new call boundary (the helper) since the download now goes through it.

* Studio: recover a stalled training model-load via Xet->HTTP respawn

Training runs in a spawn subprocess and FastModel.from_pretrained downloads
internally, so the download cannot be wrapped per-file like GGUF. Instead the
worker now watches the HF cache during the model-load phase (emitting
model_load_started / model_load_completed and a stall event), and the parent
recovers a stall by terminating the worker and respawning it once with
HF_HUB_DISABLE_XET=1.

worker.py: set HF_HUB_DISABLE_XET=1 before any HF import when the parent passes
disable_xet (respawn), and wrap trainer.load_model with start_watchdog.

training.py: plumb disable_xet through the config; track the model-load window;
on a first-load stall arm a one-shot respawn (handled on the exiting pump thread,
so no pump self-join) that preserves the DB run row (history is not duplicated)
and re-runs the load over HTTP. A second stall, or a stall outside model-load,
surfaces as a normal error. W&B init happens after model-load, so a pre-load
respawn cannot duplicate it; the dataset is re-formatted in the new worker.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: add gated test-only fault-injection hook for the Xet stall path

UNSLOTH_HF_XET_FORCE_STALL=1 makes the Xet download attempt write a partial
blob and hang, so the no-progress watchdog and the HTTP fallback can be
exercised end to end against a real repo (never set in production). Used to
verify recovery on real models: a forced Xet stall on a 5.37GB Qwen3.5-35B-A3B
shard triggered the watchdog and the HTTP retry downloaded the correct file
(sha256 verified).

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: tighten Xet-fallback comments and consolidate its tests

Trim docstrings and inline comments across the Xet->HTTP fallback code to
the non-obvious why (spawn-not-thread, killpg-not-getpgid, the sparse-partial
HTTP-resume hazard); drop comments that merely restate the code. Verified
comment-only with an AST signature check.

Merge the three helper-level test files (watchdog, transport policy, and the
HF_HUB_DISABLE_XET regression lock) into tests/test_hf_xet_fallback.py, and
prefer the real structlog over a bare stub so test collection order cannot
leak an incomplete module to others that log at import.

Full backend suite: 3455 passed, 14 pre-existing flash-attn failures only.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Daniel Han 2026-06-16 06:17:54 -07:00 • committed by GitHub
commit 58c2ec1ebd
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 1323 additions and 32 deletions
Download patch file Download diff file Expand all files Collapse all files

34
studio/backend/core/inference/llama_cpp.py
View file

@ -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}")

8
studio/backend/core/inference/orchestrator.py
View file

@ -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

125
studio/backend/core/training/training.py
View file

@ -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)

48
studio/backend/core/training/worker.py
View file

@ -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()})

160
studio/backend/tests/test_gguf_xet_fallback_integration.py Normal file
View file

@ -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"

352
studio/backend/tests/test_hf_xet_fallback.py Normal file
View file

@ -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}"
)

14
studio/backend/tests/test_offline_gguf_cache_fallback.py
View file

@ -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",

209
studio/backend/tests/test_training_xet_fallback.py Normal file
View file

@ -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

405
studio/backend/utils/hf_xet_fallback.py Normal file
View file

@ -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}'")