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