Studio: Auto download transport, Xet memory caps, and stall detection on the hub path

studio/backend/hub/schemas/downloads.py

@@ -26,7 +26,14 @@ class DownloadModelRequest(BaseModel):
     )
     use_xet: bool = Field(
         True,
-        description = "Use Xet parallel chunked transport. Default True; set False for HTTP Range-resume.",
+        description = "Legacy transport flag, superseded by transport_mode. Kept so an older "
+        "frontend or a scripted caller keeps working.",
+    )
+    transport_mode: Optional[Literal["auto", "xet", "http"]] = Field(
+        None,
+        description = "Transport preference. 'auto' (the default in the UI) lets the backend pick "
+        "per machine: it knows this host's RAM, its hf_xet build, and whether Xet has "
+        "been failing here. 'xet'/'http' force one. Omitted -> use_xet decides.",
     )
 
 
@@ -93,6 +100,8 @@ class TransportCapability(BaseModel):
 class TransportCapabilities(BaseModel):
     http: TransportCapability
     xet: TransportCapability
+    auto_resolves_to: Literal["xet", "http"] = "xet"
+    auto_reason: Optional[str] = None
 
 
 class TransportStatusResponse(BaseModel):
@@ -126,7 +135,14 @@ class DownloadDatasetRequest(BaseModel):
     repo_id: str = Field(..., description = "HuggingFace dataset repo ID")
     use_xet: bool = Field(
         True,
-        description = "Use Xet parallel chunked transport. Default True; set False for HTTP Range-resume.",
+        description = "Legacy transport flag, superseded by transport_mode. Kept so an older "
+        "frontend or a scripted caller keeps working.",
+    )
+    transport_mode: Optional[Literal["auto", "xet", "http"]] = Field(
+        None,
+        description = "Transport preference. 'auto' (the default in the UI) lets the backend pick "
+        "per machine: it knows this host's RAM, its hf_xet build, and whether Xet has "
+        "been failing here. 'xet'/'http' force one. Omitted -> use_xet decides.",
     )
 
 

studio/backend/hub/services/datasets/downloads.py

@@ -157,8 +157,15 @@ async def download_dataset_response(
     repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "dataset")
     key = _download_job_key(repo_id)
 
-    use_xet = download_lifecycle.resolve_effective_use_xet(body.use_xet)
+    # Off the event loop: resolving "auto" can run the Xet reachability probe, and a blackholed
+    # DNS makes that outlast its 3s budget while every other Studio request waits behind it.
+    use_xet, transport_reason = await asyncio.to_thread(
+        download_lifecycle.resolve_requested_use_xet,
+        getattr(body, "transport_mode", None),
+        body.use_xet,
+    )
     transport = download_lifecycle.resolve_transport(use_xet)
+    logger.info("Download transport for %s: %s (%s)", repo_id, transport, transport_reason)
     from utils.hf_cache_settings import get_hf_cache_paths
 
     cache_paths = get_hf_cache_paths()

studio/backend/hub/services/download_lifecycle.py

@@ -43,6 +43,49 @@ def resolve_effective_use_xet(use_xet: bool) -> bool:
     return False
 
 
+def resolve_requested_use_xet(transport_mode: Optional[str], use_xet: bool) -> tuple[bool, str]:
+    """Turn a download request's transport preference into ``(use_xet, reason)``.
+
+    ``transport_mode`` is the current field ("auto" / "xet" / "http"); ``use_xet`` is the older
+    boolean, still honoured so an older frontend (or a scripted API caller) keeps working. An
+    explicit "xet" is respected even on a machine the health check dislikes -- the user asked -- but
+    it still gets the memory caps and the stall fallback.
+    """
+    mode = (transport_mode or "").strip().lower()
+    if mode == download_registry.TRANSPORT_HTTP:
+        return (False, "HTTP (requested)")
+    if mode == download_registry.TRANSPORT_XET:
+        return (resolve_effective_use_xet(True), "Xet (requested)")
+    if mode == download_registry.TRANSPORT_AUTO:
+        return resolve_auto_use_xet()
+    resolved = resolve_effective_use_xet(use_xet)
+    return (resolved, "Xet" if resolved else "HTTP")
+
+
+def resolve_auto_use_xet() -> tuple[bool, str]:
+    """Pick a transport for a download the user left on "Auto". Returns ``(use_xet, reason)``.
+
+    Server-side on purpose: only the backend can see this machine's RAM, its hf_xet build, and
+    whether Xet has been failing here, and the browser must not have to guess any of it.
+
+    Probing IS allowed here (unlike on the per-download path) because Auto resolution happens once
+    per download request and the answer is cached for the whole machine -- a few hundred ms buys a
+    verdict that saves every subsequent download a stalled attempt.
+    """
+    if not resolve_effective_use_xet(True):
+        return (False, "hf_xet is not installed")
+    try:
+        from utils.hf_xet_fallback import xet_health
+        health = xet_health()
+    except Exception as exc:  # noqa: BLE001 - never let a health probe block a download
+        logger.debug("Xet health probe failed, defaulting to Xet: %s", exc)
+        return (True, "Xet (health check unavailable)")
+    if health is None:
+        # Older unsloth_zoo without the health module: no opinion, keep the existing default.
+        return (True, "Xet")
+    return (bool(health.use_xet), str(health.reason))
+
+
 def resolve_transport(use_xet: bool) -> str:
     transport = download_registry.TRANSPORT_XET if use_xet else download_registry.TRANSPORT_HTTP
     unavailable_reason = download_registry.download_transport_unavailable_reason(transport)
@@ -82,6 +125,36 @@ def spawn_worker(
     env["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
     env["HF_HUB_DISABLE_TELEMETRY"] = "1"
     env["HF_HUB_DISABLE_XET"] = "0" if use_xet else "1"
+    if use_xet:
+        # hf_xet sizes its reconstruction buffers from constants (up to 8GB stock, 64GB under
+        # high-performance mode), not from the machine. Cap them from this host's RAM and cores
+        # BEFORE the worker starts: hf_xet reads its config natively at import, so setting these
+        # inside the worker would be too late. Anything already in `env` was set deliberately by
+        # the caller and is left alone.
+        from utils.hf_xet_fallback import xet_env_overrides
+
+        allow_high_perf = os.environ.get(
+            "UNSLOTH_XET_ALLOW_HIGH_PERFORMANCE", ""
+        ).strip().lower() in (
+            "1",
+            "true",
+            "yes",
+            "on",
+        )
+        if not allow_high_perf:
+            # Unconditional, and deliberately not routed through xet_env_overrides(): an older
+            # unsloth_zoo has no tuning module, so the overrides come back empty -- and that same
+            # older zoo is the one that sets HF_XET_HIGH_PERFORMANCE=1 at import. `env` is seeded
+            # from the parent environment, so the inherited "1" would arrive here and raise the
+            # buffer ceiling to 64GB, voiding every cap below (xet-core applies the
+            # high-performance preset AFTER reading the environment). Overwrite, never setdefault.
+            for key in ("HF_XET_HIGH_PERFORMANCE", "HF_XET_HP"):
+                env[key] = "0"
+        for key, value in xet_env_overrides().items():
+            if key in ("HF_XET_HIGH_PERFORMANCE", "HF_XET_HP") and not allow_high_perf:
+                env[key] = value
+            else:
+                env.setdefault(key, value)
     # No token in Unsloth settings: fall back to the backend's own HF_TOKEN so
     # private repos stay downloadable (needed while inkling repos are private).
     # Not for a repo an API caller named: that would lend them the owner's identity.
@@ -554,6 +627,101 @@ def kill_and_reap_process(
         pass
 
 
+def _record_xet_failure(reason: str, logger) -> None:
+    """Tell the health tracker a Xet transfer failed here; best-effort, never fatal to a download."""
+    try:
+        from utils.hf_xet_fallback import record_xet_outcome
+        record_xet_outcome(False, reason)
+    except Exception as exc:  # noqa: BLE001
+        logger.debug("could not record Xet outcome: %s", exc)
+
+
+def _repo_bytes_on_disk(repo_type, repo_id: str, cache_dir) -> "Optional[int]":
+    """Bytes present for this repo, or None when unmeasurable.
+
+    Used only to tell an actual Xet transfer from a job that found everything already cached: the
+    worker reports nothing but an exit code, and the .transport marker is written before the
+    transfer starts, so there is no other signal.
+    """
+    try:
+        from utils.hf_xet_fallback import get_hf_download_state
+        state = get_hf_download_state([repo_id], repo_type = repo_type, cache_dir = cache_dir)
+    except Exception:  # noqa: BLE001 - a missing measurement must never fail a download
+        return None
+    return None if state is None else int(state[0])
+
+
+def _record_xet_success(logger) -> None:
+    """Tell the health tracker a Xet transfer completed here, which resets the failure streak."""
+    try:
+        from utils.hf_xet_fallback import record_xet_outcome
+        record_xet_outcome(True, "Xet download completed")
+    except Exception as exc:  # noqa: BLE001
+        logger.debug("could not record Xet outcome: %s", exc)
+
+
+def _start_stall_watchdog(
+    registry: download_registry.DownloadRegistry,
+    key: str,
+    proc: subprocess.Popen,
+    *,
+    repo_type: RepoType,
+    repo_id: str,
+    label: str,
+    log_prefix: str,
+    logger,
+    on_stall: Callable[[str], None],
+):
+    """Kill *proc* if its download stops making byte-level progress. Returns a stop event, or
+    ``None`` when no watchdog could be started.
+
+    SIGKILL rather than a polite signal: the worker traps SIGTERM and exits 130 ("cancelled"), which
+    would be recorded as a user cancel and skip the HTTP retry. An untrapped kill lands as "error",
+    which is the state that triggers the retry. The worker's writer is sequential and resumable, so
+    the partial it leaves behind is not lost work.
+    """
+    try:
+        from utils.hf_xet_fallback import start_watchdog
+    except Exception as exc:  # noqa: BLE001 - degraded unsloth_zoo: keep the old behaviour
+        logger.debug("%s stall watchdog unavailable for %s: %s", log_prefix, label, exc)
+        return None
+
+    metadata = registry.get_job_metadata(key)
+    cache_dir = getattr(metadata, "hub_cache", None) if metadata is not None else None
+
+    def _on_stall(message: str) -> None:
+        logger.warning(
+            "%s %s for %s; killing the worker to retry over HTTP", log_prefix, message, label
+        )
+        on_stall(message)
+        try:
+            # Kill only -- the _watch thread is already blocked reaping this process, and a second
+            # wait() here would just race it for the exit status.
+            proc.kill()
+        except ProcessLookupError:
+            pass  # already exited between the stall verdict and the kill
+        except Exception:
+            logger.exception("%s failed to kill stalled worker for %s", log_prefix, label)
+
+    try:
+        return start_watchdog(
+            repo_ids = [repo_id],
+            repo_type = repo_type,
+            cache_dir = cache_dir,
+            on_stall = _on_stall,
+            child_pid = proc.pid,
+            # Scope the measurement to partials this worker actually holds open. Without it the
+            # shared helper stays repo-wide and child_pid does nothing, so two same-transport GGUF
+            # variants of one repo (which the registry deliberately allows to run concurrently)
+            # reset each other's stall timer and a hung variant never falls back.
+            watch_new_partials_only = True,
+            xet_disabled = False,
+        )
+    except Exception as exc:  # noqa: BLE001
+        logger.debug("%s could not start stall watchdog for %s: %s", log_prefix, label, exc)
+        return None
+
+
 def register_worker(
     registry: download_registry.DownloadRegistry,
     key: str,
@@ -574,8 +742,17 @@ def register_worker(
         return False
 
     worker_token = hf_token
+    # getattr, not a direct call: test doubles and older registries do not all implement this.
+    _get_metadata = getattr(registry, "get_job_metadata", None)
+    _metadata = _get_metadata(key) if callable(_get_metadata) else None
+    _cache_dir = getattr(_metadata, "hub_cache", None) if _metadata is not None else None
+    # Sampled before the worker can write anything, so "did this job actually move bytes over Xet"
+    # is answerable when it exits.
+    _bytes_before = _repo_bytes_on_disk(repo_type, repo_id, _cache_dir)
 
     def _watch() -> None:
+        stalled: list[str] = []
+        watchdog_stop = None
         try:
             can_retry_http = (
                 transport == download_registry.TRANSPORT_XET
@@ -584,6 +761,23 @@ def _watch() -> None:
                 )
                 is None
             )
+            # Until now this path had NO stall detection at all: it relied on the worker's own exit
+            # code, and a Xet transfer that hangs with no progress and no error never produces one.
+            # That is the common failure the model-hub page shows as a frozen progress bar. Watch
+            # the cache for byte-level progress and kill a hung worker so the HTTP retry below can
+            # take over; the SIGKILL surfaces as "error", which is exactly what triggers it.
+            if can_retry_http:
+                watchdog_stop = _start_stall_watchdog(
+                    registry,
+                    key,
+                    proc,
+                    repo_type = repo_type,
+                    repo_id = repo_id,
+                    label = label,
+                    log_prefix = log_prefix,
+                    logger = logger,
+                    on_stall = stalled.append,
+                )
             state = finalize_worker_exit(
                 registry,
                 key,
@@ -598,6 +792,30 @@ def _watch() -> None:
                 cancel_marker_transport = cancel_marker_transport,
                 defer_error = can_retry_http,
             )
+            if watchdog_stop is not None:
+                # Stop measuring the moment the worker is reaped: post-download work (symlinking,
+                # verification) makes no byte-level progress and must not read as a stall.
+                watchdog_stop.set()
+            if stalled:
+                # A machine whose Xet transfers hang is one that should stop starting on Xet.
+                _record_xet_failure(stalled[0], logger)
+            elif transport == download_registry.TRANSPORT_XET and state == "complete":
+                # Clear the streak, so "two failures in a row" means in a row. Without this a
+                # stall today and another next week are counted as consecutive despite every
+                # download in between succeeding, pinning Auto to HTTP for no reason.
+                #
+                # Only a job that actually moved bytes says anything about Xet's health, though: a
+                # fully cached repo (the UI's re-download action on an up-to-date model) exits 0
+                # without touching the network, and clearing a correctly earned demotion on that
+                # would put a bad machine back on Xet. Unmeasurable means do not clear -- a missed
+                # clear costs one extra streak entry, a wrong clear undoes the demotion.
+                bytes_after = _repo_bytes_on_disk(repo_type, repo_id, _cache_dir)
+                if (
+                    _bytes_before is not None
+                    and bytes_after is not None
+                    and bytes_after > _bytes_before
+                ):
+                    _record_xet_success(logger)
             # XET-to-HTTP recovery: when a non-cancelled XET worker fails and
             # HTTP is available, attempt one automatic retry over HTTP.  The
             # transport check is the recursion guard: an HTTP worker that errors
@@ -615,6 +833,8 @@ def _watch() -> None:
                     watch_name = watch_name,
                 )
         except Exception:
+            if watchdog_stop is not None:
+                watchdog_stop.set()
             # finalize_worker_exit is the only thing that clears running/cancelling;
             # if it raises, force a terminal state so claim() isn't blocked until restart.
             logger.exception("download watcher crashed for %s", key)

studio/backend/hub/services/models/downloads.py

@@ -136,8 +136,15 @@ async def download_model_response(
             detail = f"Invalid gguf_variant: {variant!r}",
         )
     key = _download_job_key(repo_id, variant)
-    use_xet = download_lifecycle.resolve_effective_use_xet(body.use_xet)
+    # Off the event loop: resolving "auto" can run the Xet reachability probe, and a blackholed
+    # DNS makes that outlast its 3s budget while every other Studio request waits behind it.
+    use_xet, transport_reason = await asyncio.to_thread(
+        download_lifecycle.resolve_requested_use_xet,
+        getattr(body, "transport_mode", None),
+        body.use_xet,
+    )
     transport = download_lifecycle.resolve_transport(use_xet)
+    logger.info("Download transport for %s: %s (%s)", repo_id, transport, transport_reason)
     from utils.hf_cache_settings import get_hf_cache_paths
 
     cache_paths = get_hf_cache_paths()