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

fix: auto-retry stalled HF downloads with HF_HUB_DISABLE_XET=1 (#4712)

Browse source 
* 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>
This commit is contained in:
Roland Tannous 2026-03-31 14:00:46 +04:00 • committed by GitHub
commit cc5e4fbf17
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 217 additions and 33 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -17,6 +17,7 @@ Pattern follows core/training/training.py.
import atexit
import base64
import os
import structlog
from loggers import get_logger
import multiprocessing as mp
@ -33,6 +34,11 @@ 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
@ -302,6 +308,11 @@ class InferenceOrchestrator:
deadline = time.monotonic() + timeout
continue
if rtype == "stall":
msg = resp.get("message", "Download stalled")
logger.warning("Subprocess reported stall: %s", msg)
raise DownloadStallError(msg)
# Other response types during wait — skip
logger.debug(
"Skipping response type '%s' while waiting for '%s'",
@ -627,36 +638,66 @@ class InferenceOrchestrator:
# Dead subprocess — clean up
self._shutdown_subprocess(timeout = 2)
logger.info(
"Spawning fresh inference subprocess for '%s' (transformers %s.x)",
model_name,
needed_major,
disable_xet = sub_config.get("disable_xet", False) or (
os.environ.get("HF_HUB_DISABLE_XET") == "1"
)
self._spawn_subprocess(sub_config)
resp = self._wait_response("loaded")
# Update local state from response
if resp.get("success"):
self._current_transformers_major = needed_major
model_info = resp.get("model_info", {})
self.active_model_name = model_info.get("identifier", model_name)
self.models[self.active_model_name] = {
"is_vision": model_info.get("is_vision", False),
"is_lora": model_info.get("is_lora", False),
"display_name": model_info.get("display_name", model_name),
"is_audio": model_info.get("is_audio", False),
"audio_type": model_info.get("audio_type"),
"has_audio_input": model_info.get("has_audio_input", False),
}
self.loading_models.discard(model_name)
logger.info("Model '%s' loaded successfully in subprocess", model_name)
return True
else:
error = resp.get("error", "Failed to load model")
self.loading_models.discard(model_name)
self.active_model_name = None
self.models.clear()
raise Exception(error)
for attempt in range(2):
logger.info(
"Spawning fresh inference subprocess for '%s' "
"(transformers %s.x, attempt %d/2%s)",
model_name,
needed_major,
attempt + 1,
", xet disabled" if disable_xet else "",
)
sub_config["disable_xet"] = disable_xet
self._spawn_subprocess(sub_config)
try:
resp = self._wait_response("loaded")
except DownloadStallError:
# First stall and Xet was enabled -> retry with Xet disabled
if attempt == 0 and not disable_xet:
logger.warning(
"Download stalled for '%s' -- retrying with "
"HF_HUB_DISABLE_XET=1",
model_name,
)
self._shutdown_subprocess(timeout = 5)
disable_xet = True
continue
# Second stall (or already had xet disabled) -> give up
self._shutdown_subprocess(timeout = 5)
raise RuntimeError(
f"Download stalled for '{model_name}' even with "
f"HF_HUB_DISABLE_XET=1 -- check your network connection"
)
# Got a response — check success
if resp.get("success"):
self._current_transformers_major = needed_major
model_info = resp.get("model_info", {})
self.active_model_name = model_info.get("identifier", model_name)
self.models[self.active_model_name] = {
"is_vision": model_info.get("is_vision", False),
"is_lora": model_info.get("is_lora", False),
"display_name": model_info.get("display_name", model_name),
"is_audio": model_info.get("is_audio", False),
"audio_type": model_info.get("audio_type"),
"has_audio_input": model_info.get("has_audio_input", False),
}
self.loading_models.discard(model_name)
logger.info(
"Model '%s' loaded successfully in subprocess", model_name
)
return True
else:
error = resp.get("error", "Failed to load model")
self.loading_models.discard(model_name)
self.active_model_name = None
self.models.clear()
raise Exception(error)
except Exception:
self.loading_models.discard(model_name)

151
studio/backend/core/inference/worker.py
View file

@ -115,20 +115,145 @@ def _build_model_config(config: dict):
return mc
def _start_heartbeat(resp_queue: Any, interval: float = 30.0) -> threading.Event:
def _get_hf_download_state(
model_names: list[str] | None = None,
) -> tuple[int, bool] | None:
"""Return (total_bytes, has_incomplete) for the HF Hub cache, or None on error.
When *model_names* is provided, only those models' ``blobs/``
directories are checked instead of scanning every cached model --
much faster on systems with many models. Accepts multiple names so
that LoRA loads can watch both the adapter repo and the base model
repo simultaneously.
*has_incomplete* is True when any ``*.incomplete`` files exist in the
watched blobs directories, indicating that ``huggingface_hub`` is
actively downloading.
Returns None if the state cannot be determined (import error,
permission error, etc.) so callers can skip stall logic.
"""
try:
from huggingface_hub.constants import HF_HUB_CACHE
cache = Path(HF_HUB_CACHE)
if not cache.exists():
return (0, False)
total = 0
has_incomplete = False
blobs_dirs: list[Path] = []
if model_names:
for name in model_names:
if not name:
continue
# Skip local filesystem paths -- HF model IDs use forward
# slashes (org/model) but never start with / . ~ or contain
# backslashes. This distinguishes them from absolute paths,
# relative paths, and Windows paths.
if name.startswith(("/", ".", "~")) or "\\" in name:
continue
# HF cache dir format: models--org--name (slashes -> --)
cache_dir_name = "models--" + name.replace("/", "--")
blobs_dir = cache / cache_dir_name / "blobs"
if blobs_dir.exists():
blobs_dirs.append(blobs_dir)
else:
blobs_dirs = list(cache.glob("models--*/blobs"))
for bdir in blobs_dirs:
for f in bdir.iterdir():
try:
if f.is_file():
total += f.stat().st_size
if f.name.endswith(".incomplete"):
has_incomplete = True
except OSError:
pass
return (total, has_incomplete)
except Exception as e:
logger.debug("Failed to determine HF download state: %s", e)
return None
def _start_heartbeat(
resp_queue: Any,
interval: float = 30.0,
stall_timeout: float = 180.0,
xet_disabled: bool = False,
model_names: list[str] | None = None,
) -> threading.Event:
"""Start a daemon thread that sends periodic status heartbeats.
Returns a stop event — set it to terminate the heartbeat thread.
Monitors the HF Hub cache directory for download activity. A stall
is only reported when ``*.incomplete`` files are present (indicating
``huggingface_hub`` is actively downloading) **and** the total cache
size has not changed for *stall_timeout* seconds.
Once the download finishes (no more ``.incomplete`` files), the stall
timer resets, so post-download initialization (quantization, GPU
weight loading) is never misclassified as a stalled download.
Returns a stop event -- set it to terminate the heartbeat thread.
"""
stop = threading.Event()
transport = "https" if xet_disabled else "xet"
def _beat():
state = _get_hf_download_state(model_names)
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(model_names)
now = time.monotonic()
# Skip stall logic if we cannot measure the cache
if state is None:
_send_response(
resp_queue,
{
"type": "status",
"message": f"Loading model ({transport} transport)...",
"ts": time.time(),
},
)
continue
current_size, has_incomplete = state
if current_size != last_size:
last_size = current_size
last_change = now
# Only fire stall when .incomplete files are present,
# confirming a download is actively in progress.
# Once downloads finish (no .incomplete), reset the timer
# so model init time is not counted as a stall.
if not has_incomplete:
last_change = now
elif now - last_change >= stall_timeout:
_send_response(
resp_queue,
{
"type": "stall",
"message": (
f"Download appears stalled ({transport} transport) "
f"-- no progress for {int(now - last_change)}s"
),
"ts": time.time(),
},
)
# Only fire once -- the orchestrator will kill us
return
_send_response(
resp_queue,
{
"type": "status",
"message": "Still loading model...",
"message": f"Loading model ({transport} transport)...",
"ts": time.time(),
},
)
@ -199,7 +324,21 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
# Send heartbeats every 30s so the orchestrator knows we're still alive
# (download / weight loading can take a long time on slow connections)
heartbeat_stop = _start_heartbeat(resp_queue, interval = 30.0)
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1"
# Watch both the model repo and base model repo (for LoRA loads
# where the base model download is the actual bottleneck)
watch_repos = [mc.identifier]
base = getattr(mc, "base_model", None)
if base and str(base) != mc.identifier:
watch_repos.append(str(base))
heartbeat_stop = _start_heartbeat(
resp_queue,
interval = 30.0,
xet_disabled = xet_disabled,
model_names = watch_repos,
)
try:
success = backend.load_model(
config = mc,
@ -522,6 +661,10 @@ def run_inference_process(
"ignore" # Suppress warnings at C-level before imports
)
if config.get("disable_xet"):
os.environ["HF_HUB_DISABLE_XET"] = "1"
logger.info("Xet transport disabled (HF_HUB_DISABLE_XET=1)")
import warnings
from loggers.config import LogConfig