unsloth/studio/backend
Roland Tannous cc5e4fbf17
fix: auto-retry stalled HF downloads with HF_HUB_DISABLE_XET=1 (#4712)
* fix: auto-retry stalled HF downloads with HF_HUB_DISABLE_XET=1

The heartbeat thread now monitors the HF Hub cache directory for
file-size growth. If no bytes are written for 3 minutes, it sends a
"stall" message to the orchestrator, which kills the subprocess and
retries with HF_HUB_DISABLE_XET=1 (falling back from Xet to standard
HTTPS). If the retry also stalls, it errors out with a clear message.

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

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

* fix: include transport type (xet/https) in heartbeat and stall log messages

Makes it clear in backend logs whether the download is using xet or
https transport, and which transport stalled — helpful for debugging.

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

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

* fix: monitor HF Hub .tmp dir to avoid false stall detections

huggingface_hub downloads into .tmp/ before atomically moving to
blobs/. Without monitoring .tmp, a large shard actively downloading
for several minutes would show zero blob growth and trigger a false
stall.

* fix: scope HF cache size check to specific model being loaded

Instead of scanning every models--*/blobs directory (O(N) with cached
models), only check the specific model's blobs dir plus the global
.tmp dir. Much faster on systems with many cached models.

* Fix false stall detection on cached/local models and cleanup issues

- Only fire stall if download activity was observed (cache size changed
  at least once). Previously, any model load taking >180s would trigger
  a false stall, even for already-cached or local models where no
  download is happening.
- Return -1 from _get_hf_cache_size on exception to distinguish
  "unable to measure" from "genuinely zero bytes". Skip stall logic
  when measurement fails.
- Add _shutdown_subprocess before raising on terminal stall path to
  prevent leaking a stuck subprocess.
- Detect pre-existing HF_HUB_DISABLE_XET=1 in the parent environment
  to avoid a redundant retry cycle when Xet is already disabled.
- Remove global .tmp directory scanning (not used by modern
  huggingface_hub; in-progress downloads use .incomplete files in
  blobs/ which are already captured by iterdir).
- Add f.is_file() guard in cache size calculation.
- Replace em dashes with ASCII dashes for Windows terminal compat.

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

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

* Harden stall detection edge cases

- Guard -1 to valid value transition: when initial _get_hf_cache_size
  returns -1 (error) and later recovers to a real value, do not count
  that as download activity. Only set saw_download_activity when the
  previous measurement was also valid (>= 0).
- Move os import to top-level in orchestrator.py instead of inline
  import os as _os.
- Fix misleading comment about post-download protection.

* Use .incomplete files to detect active downloads for stall detection

Replace the saw_download_activity heuristic with direct .incomplete file
detection. huggingface_hub creates *.incomplete files in blobs/ during
active downloads and removes them on completion. This gives a reliable
signal for whether a download is actually in progress.

Benefits:
- Cached models: no .incomplete files -> no stall fired even after 180s
- Post-download init (quantization, GPU loading): .incomplete files gone
  so stall timer resets, long init phases are not killed
- Pre-download hangs (XET handshake stall): .incomplete files are
  created at download start, so zero-byte stalls are now detected
- No more false positives from -1 to valid measurement transitions

The _get_hf_download_state function now returns (total_bytes,
has_incomplete) tuple or None on error, replacing _get_hf_cache_size.

* Add debug logging to download state exception handler

Log the exception at debug level when _get_hf_download_state fails,
instead of silently returning None. Helps with troubleshooting cache
measurement issues.

* Watch both adapter and base model repos for LoRA stall detection

When loading a LoRA adapter, the actual download bottleneck is often
the base model, not the adapter itself. Update the heartbeat to watch
both mc.identifier and mc.base_model cache directories so stall
detection works for LoRA loads where the base model stalls on Xet.

Also update _get_hf_download_state to accept multiple model names and
skip names without "/" (local paths) since those do not have HF cache
directories.

* Fix model name filtering for official HF models without org prefix

Models like gpt2 and bert-base-uncased do not contain a slash but are
still valid HF Hub models with cache directories. Replace the "/" check
with a proper local-path detection that checks for path separators and
path-like prefixes instead.

Also fix the base_model watch list to not require "/" in the base model
name, so official models used as LoRA bases are also monitored.

* Fix local path detection that broke all org/model names on Linux

The os.path.sep check matched "/" in HF model IDs like "org/model" on
Linux, causing the stall detector to skip ALL standard HF models.

Replace with a check that only skips names starting with "/" (absolute
paths), "." (relative paths), "~" (home-relative), or containing "\"
(Windows paths). HF model IDs like "org/model" or "gpt2" pass through
correctly on all platforms.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-03-31 03:00:46 -07:00
..
assets fix(studio): correct default weight_decay and learning rate (#4695) 2026-03-31 13:50:25 +04:00
auth fix: remove old comments (#4292) 2026-03-14 16:50:13 +04:00
core fix: auto-retry stalled HF downloads with HF_HUB_DISABLE_XET=1 (#4712) 2026-03-31 03:00:46 -07:00
loggers Final cleanup 2026-03-12 18:28:04 +00:00
models fix(studio): correct default weight_decay and learning rate (#4695) 2026-03-31 13:50:25 +04:00
plugins Bump Data Designer to 0.5.4 (removes litellm dependency) (#4569) 2026-03-25 02:01:43 -07:00
requirements fix: no-torch install deps without pulling torch transitively (#4650) 2026-03-27 05:19:26 -07:00
routes [Studio] multi gpu finetuning/inference via "balanced_low0/sequential" device_map (#4602) 2026-03-30 02:33:15 -07:00
state Final cleanup 2026-03-12 18:28:04 +00:00
storage feat(studio): training history persistence and past runs viewer (#4501) 2026-03-25 00:58:55 -07:00
tests [studio] multi gpu: revert to balanced for inference. (#4698) 2026-03-31 01:24:41 -07:00
utils [studio] multi gpu: revert to balanced for inference. (#4698) 2026-03-31 01:24:41 -07:00
__init__.py Final cleanup 2026-03-12 18:28:04 +00:00
_platform_compat.py Fix Studio crash on Anaconda/conda-forge Python (#4484) 2026-03-22 05:36:55 -07:00
colab.py Allow install_python_stack to run on Colab (#4633) 2026-03-27 00:29:27 +04:00
main.py [Studio] multi gpu finetuning/inference via "balanced_low0/sequential" device_map (#4602) 2026-03-30 02:33:15 -07:00
run.py fix(studio): avoid UnicodeEncodeError on Windows cp1252 consoles (#4699) 2026-03-30 06:40:47 -07:00
startup_banner.py studio: unify Windows installer/setup logging style, verbosity controls, and startup messaging (#4651) 2026-03-30 00:53:23 -07:00