From cc5e4fbf17d05066dd95d8c23c048f38c55af998 Mon Sep 17 00:00:00 2001 From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com> Date: Tue, 31 Mar 2026 14:00:46 +0400 Subject: [PATCH] fix: auto-retry stalled HF downloads with HF_HUB_DISABLE_XET=1 (#4712) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 --- studio/backend/core/inference/orchestrator.py | 97 +++++++---- studio/backend/core/inference/worker.py | 151 +++++++++++++++++- 2 files changed, 216 insertions(+), 32 deletions(-) diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 0366524078..f293a0dcd8 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -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) diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 9948142afc..e2513f43de 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -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