unsloth/studio
Daniel Han 58c2ec1ebd
Studio: Xet-primary model downloads with automatic HTTP fallback on stall (#6372)
* 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>
2026-06-16 06:17:54 -07:00
..
backend Studio: Xet-primary model downloads with automatic HTTP fallback on stall (#6372) 2026-06-16 06:17:54 -07:00
frontend Fix GGUF variant file selection (#6342) 2026-06-16 12:42:58 +02:00
src-tauri feat(studio): Hub + Download Manager (#5916) 2026-06-09 04:11:24 -07:00
__init__.py Final cleanup 2026-03-12 18:28:04 +00:00
install_llama_prebuilt.py Studio: don't silently fall back to a CPU prebuilt on NVIDIA Linux GPU hosts (#6310) 2026-06-14 02:06:31 -03:00
install_python_stack.py fix: use partial hipinfo output on crash to avoid CPU fallback (RDNA 4 / gfx1200) (#6292) 2026-06-15 13:26:04 +02:00
LICENSE.AGPL-3.0 Add AGPL-3.0 license to studio folder 2026-03-09 19:36:25 +00:00
package-lock.json ci: advisory lockfile supply-chain audit (no install-script changes) (#5604) 2026-05-19 05:56:56 -07:00
package.json ci: advisory lockfile supply-chain audit (no install-script changes) (#5604) 2026-05-19 05:56:56 -07:00
setup.bat Final cleanup 2026-03-12 18:28:04 +00:00
setup.ps1 Make the Studio installers (sh + ps1) resilient to transient uv download failures (#6281) 2026-06-16 03:48:01 -07:00
setup.sh Make the Studio installers (sh + ps1) resilient to transient uv download failures (#6281) 2026-06-16 03:48:01 -07:00
Unsloth_Studio_Colab.ipynb Fix/studio colab proxy and iframe - Unsloth Studio not loading in Colab (iframe "refused to connect" and wrong URL) (#5844) 2026-05-28 23:54:48 -07:00