fix(studio): recover stalled Hub downloads over HTTP

studio/backend/hub/services/download_lifecycle.py

@@ -206,7 +206,8 @@ def finalize_worker_exit(
     repo_type: Optional[RepoType] = None,
     repo_id: Optional[str] = None,
     transport: Optional[str] = None,
-) -> None:
+    defer_error: bool = False,
+) -> str:
     """Block until *proc* exits, then record the job's terminal state in
     *registry*. Drains and scrubs stderr first, then classifies the exit code.
     A no-op when the process was already dropped (e.g. superseded).
@@ -218,7 +219,7 @@ def finalize_worker_exit(
     rc = proc.wait()
     cancel_requested = registry.cancel_requested(key)
     if not registry.drop_process(key, proc):
-        return
+        return "idle"
     stderr_text = download_registry.scrub_secrets(
         (stderr_data or b"").decode("utf-8", "replace").strip(),
         hf_token = hf_token,
@@ -262,14 +263,201 @@ def finalize_worker_exit(
             logger = logger,
         )
     else:
-        registry.set_job(
+        if not defer_error:
+            registry.set_job(
+                key,
+                "error",
+                stderr_text or f"worker exited with code {rc}",
+            )
+        logger.error(
+            f"{log_prefix} failed for {label} (rc={rc}): {stderr_text}",
+        )
+    return state
+
+
+def _set_job_transport(
+    registry: download_registry.DownloadRegistry, key: str, transport: str
+) -> None:
+    registry.update_job_transport(key, transport)
+
+
+def _set_retry_failure_state(
+    registry: download_registry.DownloadRegistry,
+    key: str,
+    error: str,
+    *,
+    repo_type: RepoType,
+    repo_id: str,
+    fallback_variant: Optional[str],
+    fallback_transport: Optional[str],
+    logger,
+) -> str:
+    state, metadata = registry.set_error_unless_cancelled(key, error)
+    if state == "cancelled":
+        download_registry.persist_cancel_marker(
+            repo_type,
+            repo_id,
+            metadata.variant if metadata is not None and metadata.variant else fallback_variant,
+            metadata.transport
+            if metadata is not None and metadata.transport
+            else fallback_transport,
+            logger = logger,
+        )
+    return state
+
+
+def _try_http_retry(
+    registry: download_registry.DownloadRegistry,
+    key: str,
+    *,
+    hf_token: Optional[str],
+    label: str,
+    log_prefix: str,
+    logger,
+    repo_type: RepoType,
+    repo_id: str,
+    watch_name: str,
+) -> bool:
+    """Reclaim *key* with HTTP transport and spawn a recovery worker.
+
+    Returns ``True`` when the HTTP worker was successfully registered.
+    Caller is responsible for ensuring this is only called when: the job is
+    in ``"error"`` state, the original transport was XET, and HTTP is available.
+
+    Derives variant and blob-hash metadata from the registry entry written by
+    the original XET claim so callers do not re-construct worker arguments.
+    Re-queries peer protection hashes at spawn time to reflect any concurrent
+    sibling changes between the XET failure and this call.
+    """
+    original_metadata = registry.get_job_metadata(key)
+    if original_metadata is None:
+        logger.debug("%s XET retry skipped for %s; metadata unavailable", log_prefix, label)
+        _set_retry_failure_state(
+            registry,
             key,
-            "error",
-            stderr_text or f"worker exited with code {rc}",
+            "XET retry skipped: metadata unavailable",
+            repo_type = repo_type,
+            repo_id = repo_id,
+            fallback_variant = download_registry.variant_from_key(key),
+            fallback_transport = download_registry.TRANSPORT_XET,
+            logger = logger,
+        )
+        return False
+    if original_metadata.transport != download_registry.TRANSPORT_XET:
+        logger.debug(
+            "%s XET retry skipped for %s; original transport was %s",
+            log_prefix,
+            label,
+            original_metadata.transport,
         )
+        _set_retry_failure_state(
+            registry,
+            key,
+            f"XET retry skipped: original transport was {original_metadata.transport}",
+            repo_type = repo_type,
+            repo_id = repo_id,
+            fallback_variant = original_metadata.variant,
+            fallback_transport = original_metadata.transport,
+            logger = logger,
+        )
+        return False
+    variant = original_metadata.variant
+    blob_hashes = original_metadata.blob_hashes
+    progress_blob_hashes = original_metadata.progress_blob_hashes
+    completed_baseline_bytes = (
+        download_registry.completed_blob_bytes(
+            repo_type,
+            repo_id,
+            progress_blob_hashes,
+        )
+        if progress_blob_hashes
+        else 0
+    )
+    generation = registry.current_generation(key)
+
+    claimed, _ = registry.claim(
+        key,
+        download_registry.TRANSPORT_HTTP,
+        repo_type = repo_type,
+        repo_id = repo_id,
+        variant = variant,
+        blob_hashes = blob_hashes,
+        progress_blob_hashes = progress_blob_hashes,
+        completed_baseline_bytes = completed_baseline_bytes,
+        generation = generation,
+        replace_active = True,
+    )
+    if not claimed:
+        logger.debug(
+            "%s XET retry claim rejected for %s; another job took the slot", log_prefix, label
+        )
+        _set_retry_failure_state(
+            registry,
+            key,
+            "HTTP retry could not reclaim the download slot",
+            repo_type = repo_type,
+            repo_id = repo_id,
+            fallback_variant = variant,
+            fallback_transport = original_metadata.transport,
+            logger = logger,
+        )
+        return False
+
+    args: list[str] = ["--repo-id", repo_id]
+    if repo_type == "dataset":
+        args.append("--dataset")
+    elif variant:
+        args.extend(["--variant", variant])
+
+    # Re-query at spawn time: sibling state may have changed since XET failed.
+    peer_hashes = registry.peer_blob_hashes(key) if variant else frozenset()
+
+    logger.warning(
+        "%s XET worker failed for %s; retrying over HTTP",
+        log_prefix,
+        label,
+    )
+    try:
+        proc = spawn_worker(
+            args,
+            hf_token,
+            use_xet = False,
+            protected_blob_hashes = peer_hashes or None,
+        )
+    except Exception as exc:
+        scrubbed = download_registry.scrub_secrets(str(exc), hf_token = hf_token)
         logger.error(
-            f"{log_prefix} failed for {label} (rc={rc}): {stderr_text}",
+            "%s HTTP retry spawn failed for %s: %s",
+            log_prefix,
+            label,
+            scrubbed,
+        )
+        _set_job_transport(registry, key, download_registry.TRANSPORT_XET)
+        _set_retry_failure_state(
+            registry,
+            key,
+            scrubbed,
+            repo_type = repo_type,
+            repo_id = repo_id,
+            fallback_variant = variant,
+            fallback_transport = original_metadata.transport,
+            logger = logger,
         )
+        return False
+
+    return register_worker(
+        registry,
+        key,
+        proc,
+        hf_token = hf_token,
+        label = label,
+        log_prefix = log_prefix,
+        logger = logger,
+        repo_type = repo_type,
+        repo_id = repo_id,
+        transport = download_registry.TRANSPORT_HTTP,
+        watch_name = watch_name,
+    )
 
 
 def kill_and_reap_process(
@@ -315,7 +503,14 @@ def register_worker(
 
     def _watch() -> None:
         try:
-            finalize_worker_exit(
+            can_retry_http = (
+                transport == download_registry.TRANSPORT_XET
+                and download_registry.download_transport_unavailable_reason(
+                    download_registry.TRANSPORT_HTTP
+                )
+                is None
+            )
+            state = finalize_worker_exit(
                 registry,
                 key,
                 proc,
@@ -326,7 +521,24 @@ def _watch() -> None:
                 repo_type = repo_type,
                 repo_id = repo_id,
                 transport = transport,
+                defer_error = can_retry_http,
             )
+            # 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
+            # never satisfies `transport == TRANSPORT_XET`, so it stays terminal.
+            if can_retry_http and state == "error":
+                _try_http_retry(
+                    registry,
+                    key,
+                    hf_token = worker_token,
+                    label = label,
+                    log_prefix = log_prefix,
+                    logger = logger,
+                    repo_type = repo_type,
+                    repo_id = repo_id,
+                    watch_name = watch_name,
+                )
         except Exception:
             # 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.