unslothai--unsloth/studio/backend/hub/workers/hf_download.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""HuggingFace Hub download worker, spawned as a subprocess so SIGKILL stops all chunk threads.

Resume safety
-------------
Downloads here MUST be single-stream sequential writers so the parent's
SIGKILL → restart loop can rely on ``os.path.getsize(.incomplete)`` to
compute the correct resume offset.

Enforced by:
- Setting ``HF_HUB_DISABLE_XET=1`` and ``HF_HUB_ENABLE_HF_TRANSFER=0`` on
  the spawning side (see :mod:`hub.utils.download_registry`) for transport=http.
- Passing ``max_workers=1`` to ``snapshot_download`` so files download
  serially, making the at-most-one-active-`.incomplete` invariant hold
  globally and simplifying reasoning about partial state during a SIGKILL.
- Letting ``prepare_cache_for_transport`` purge any pre-existing
  ``.incomplete`` blobs not provably from the same sequential writer.

If the final byte count doesn't match what HF declared, huggingface_hub
raises ``EnvironmentError`` ("Consistency check failed: …"); we surface
that on stderr so the watcher can show the exact message to the user.
"""

from __future__ import annotations

import argparse
import os
import signal
import sys
import threading
import time
from pathlib import Path

_HERE = Path(__file__).resolve().parent
_BACKEND = _HERE.parent.parent
if str(_BACKEND) not in sys.path:
    sys.path.insert(0, str(_BACKEND))

from hub.utils.snapshot_filters import (
    SNAPSHOT_IGNORE_PATTERNS,
)
from hub.utils.gguf_plan import (
    GgufVariantPlan,
    build_gguf_variant_plans,
    plan_from_expected_files,
    sibling_sha256,
)
from hub.utils.state_dir import RepoType

HfTokenArg = str | bool | None


# Bound the metadata fetch so a stalled connection fails the worker (exit 1)
# instead of hanging at 0%. The file download itself is governed separately by
# huggingface_hub's own timeout.
_METADATA_REQUEST_TIMEOUT = 10.0
_METADATA_RETRY_TIMEOUT = 30.0
_METADATA_RETRY_DELAY = 1.0


def _on_signal(signum, frame):
    # 130 is what `classify_exit` maps to the "cancelled" job state.
    sys.exit(130)


def _install_signal_handlers() -> None:
    signal.signal(signal.SIGTERM, _on_signal)
    signal.signal(signal.SIGINT, _on_signal)
    sigpipe = getattr(signal, "SIGPIPE", None)
    if sigpipe is not None:
        signal.signal(sigpipe, _on_signal)


def _parent_poll_seconds() -> float:
    raw = os.environ.get("UNSLOTH_HF_WORKER_PARENT_POLL_SECONDS")
    if raw:
        try:
            value = float(raw)
            if value > 0:
                return value
        except ValueError:
            pass
    return 2.0


def _protected_blob_hashes() -> frozenset[str]:
    """Blob hashes a concurrent same-repo peer is writing (passed by the backend
    as a plain env list). Excluded from this worker's purge so a shared
    ``.incomplete`` (e.g. a bundled mmproj) is never deleted under the peer."""
    raw = os.environ.get("UNSLOTH_PROTECTED_BLOB_HASHES", "")
    return frozenset(h for h in raw.split(",") if h)


def _parent_is_alive(parent_pid: int) -> bool:
    """Whether the recorded parent (the backend) is still running.

    Liveness ONLY: ``os.kill(pid, 0)`` on POSIX, an ``OpenProcess`` handle on
    Windows, against the *recorded* PID (never os.getppid(), so POSIX
    reparenting to init after the backend dies still resolves as dead). Probe
    ambiguity is treated as alive so a transient error never kills a healthy
    download.

    We deliberately do NOT compare psutil ``create_time()`` for PID-reuse
    detection: it isn't stable across reads on some platforms, so an exact match
    can spuriously kill a live download. PID-reuse after parent death is covered
    by the boot-time orphan reaper.
    """
    if sys.platform == "win32":
        import ctypes
        from ctypes import wintypes

        SYNCHRONIZE = 0x00100000
        WAIT_OBJECT_0 = 0x0
        ERROR_INVALID_PARAMETER = 87
        kernel32 = ctypes.WinDLL("kernel32", use_last_error = True)
        kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD]
        kernel32.OpenProcess.restype = wintypes.HANDLE
        kernel32.WaitForSingleObject.argtypes = [wintypes.HANDLE, wintypes.DWORD]
        kernel32.WaitForSingleObject.restype = wintypes.DWORD
        kernel32.CloseHandle.argtypes = [wintypes.HANDLE]
        kernel32.CloseHandle.restype = wintypes.BOOL
        ctypes.set_last_error(0)
        handle = kernel32.OpenProcess(SYNCHRONIZE, False, parent_pid)
        if not handle:
            return ctypes.get_last_error() != ERROR_INVALID_PARAMETER
        try:
            return kernel32.WaitForSingleObject(handle, 0) != WAIT_OBJECT_0
        finally:
            kernel32.CloseHandle(handle)
    try:
        os.kill(parent_pid, 0)
    except ProcessLookupError:
        return False
    except OSError:
        return True
    return True


def _terminate_orphaned_self() -> None:
    # Hard exit from the watchdog thread: a self-SIGTERM would be deferred while
    # the main thread is GIL-blocked in a C socket read. The partial .incomplete
    # resumes byte-exact and marker/manifest writes are atomic, so cancelled code
    # 130 is safe. The diagnostic is best-effort: a dead parent's closed stderr
    # pipe can raise BrokenPipeError, which must never preempt the exit.
    try:
        print(
            "Parent process exited; stopping orphaned download worker.",
            file = sys.stderr,
        )
        sys.stderr.flush()
    except Exception:
        pass
    os._exit(130)


def _install_parent_death_watchdog(parent_pid: int | None) -> None:
    if not parent_pid or parent_pid <= 0:
        return
    interval = _parent_poll_seconds()

    def _watch() -> None:
        while True:
            try:
                alive = _parent_is_alive(parent_pid)
            except Exception:
                alive = True
            if not alive:
                _terminate_orphaned_self()
                return
            time.sleep(interval)

    threading.Thread(
        target = _watch,
        name = "parent-death-watchdog",
        daemon = True,
    ).start()


def _hf_token_arg(hf_token: str | None) -> HfTokenArg:
    return hf_token if hf_token else False


def _retry_metadata_fetch(repo_id: str, fetch, *, label: str):
    for attempt, timeout in enumerate((_METADATA_REQUEST_TIMEOUT, _METADATA_RETRY_TIMEOUT)):
        try:
            return fetch(timeout)
        except Exception as e:
            if attempt == 1:
                raise
            print(
                f"{label} request failed for {repo_id} " f"({type(e).__name__}: {e}); retrying.",
                file = sys.stderr,
            )
            time.sleep(_METADATA_RETRY_DELAY)
    raise RuntimeError(f"{label} unavailable for {repo_id}")


def _model_info_with_retry(repo_id: str, hf_token: str | None):
    from huggingface_hub import model_info as hf_model_info
    return _retry_metadata_fetch(
        repo_id,
        lambda timeout: hf_model_info(
            repo_id,
            token = _hf_token_arg(hf_token),
            timeout = timeout,
            files_metadata = True,
        ),
        label = "Metadata",
    )


def _dataset_info_with_retry(repo_id: str, hf_token: str | None):
    from huggingface_hub import HfApi
    api = HfApi(token = _hf_token_arg(hf_token))
    return _retry_metadata_fetch(
        repo_id,
        lambda timeout: api.dataset_info(
            repo_id,
            timeout = timeout,
            files_metadata = True,
        ),
        label = "Dataset metadata",
    )


# Tied to drain_stderr_excerpt's 500-byte head/tail window in the parent (see
# hub/utils/download_registry.py): listing every expected file would blow past
# it and lose the diagnostic. Cap the preview so the summary line survives.
_VERIFY_PATH_LIST_CAP = 10


def _format_path_list(paths: tuple[str, ...], cap: int = _VERIFY_PATH_LIST_CAP) -> str:
    if len(paths) <= cap:
        return ", ".join(paths)
    head = ", ".join(paths[:cap])
    return f"{head}, ... and {len(paths) - cap} more"


def _verify_completed_download(
    repo_type: RepoType,
    repo_id: str,
    variant: str | None,
    snapshot_path: str,
    *,
    metadata_unavailable: bool = False,
) -> None:
    """Verify every manifest file is on disk at its declared size; exit nonzero
    with a diagnostic if not.

    No-op when no manifest exists: the manifest write is best-effort, so absence
    means "verification unavailable, trust snapshot_download's exit code".
    """
    from hub.utils import download_manifest

    manifest = download_manifest.read_manifest(repo_type, repo_id, variant)
    if manifest is None:
        return
    result = download_manifest.verify_against_disk(
        manifest,
        Path(snapshot_path),
    )
    if result.ok:
        return
    label = f"{repo_id}{f' [{variant}]' if variant else ''}"
    if metadata_unavailable:
        print(
            f"Could not reach Hugging Face for {label} and the copy on disk is "
            f"incomplete ({len(result.missing)} file(s) missing, "
            f"{len(result.size_mismatched)} the wrong size). Access to a private "
            "or restricted repo may have been lost (HF token removed or "
            "changed), the connection dropped, or Hugging Face is temporarily "
            "unavailable. Set a valid HF token or reconnect, then resume the "
            "download.",
            file = sys.stderr,
        )
    else:
        print(
            f"Verification failed for {label}: snapshot_download completed but "
            f"{len(result.missing)} expected file(s) are missing and "
            f"{len(result.size_mismatched)} have incorrect size on disk.",
            file = sys.stderr,
        )
    if result.missing:
        print(
            f"Missing: {_format_path_list(result.missing)}",
            file = sys.stderr,
        )
    if result.size_mismatched:
        print(
            f"Size mismatched: {_format_path_list(result.size_mismatched)}",
            file = sys.stderr,
        )
    sys.exit(1)


def _preflight_disk_space(repo_type: str, repo_id: str, expected_files: list) -> None:
    """Fail fast when the active HF cache filesystem can't hold what's left to
    download. Fail-open: any inability to size the work or read free space skips
    the check, so a real download is never blocked by an estimation gap."""
    import shutil

    from hub.utils.download_registry import existing_blob_bytes
    from hub.utils.hf_cache_state import hf_cache_root

    try:
        size_by_hash: dict[str, int] = {}
        unhashed_bytes = 0
        for expected in expected_files:
            size = int(getattr(expected, "size", 0) or 0)
            if size <= 0:
                continue
            blob_hash = getattr(expected, "sha256", None)
            if blob_hash:
                # Dedup by content hash: a blob listed under two filenames is
                # written once, so it must be counted once.
                size_by_hash[blob_hash] = size
            else:
                unhashed_bytes += size
        total_expected = sum(size_by_hash.values()) + unhashed_bytes
        if total_expected <= 0:
            return
        already_have = existing_blob_bytes(
            repo_type,
            repo_id,
            frozenset(size_by_hash),
        )
        remaining = max(0, total_expected - already_have)
        if remaining <= 0:
            return
        root = hf_cache_root(create = True)
        if root is None:
            return
        free = shutil.disk_usage(root).free
    except Exception:
        return

    if free < remaining:
        print(
            f"Not enough disk space to download {repo_id}: need about "
            f"{remaining / (1024 ** 3):.1f} GB free in {root}, but only "
            f"{free / (1024 ** 3):.1f} GB is available. Free up space and "
            "try again.",
            file = sys.stderr,
        )
        sys.exit(1)


def _snapshot_download_plan(info) -> tuple[list[str], list]:
    from hub.utils.download_manifest import ExpectedFile
    from hub.utils.snapshot_filters import (
        resolve_snapshot_ignore_patterns_for_files,
        snapshot_download_siblings,
    )

    filenames = [s.rfilename for s in info.siblings if isinstance(s.rfilename, str)]
    filtered = snapshot_download_siblings(info.siblings)
    expected_files = [
        ExpectedFile(
            path = s.rfilename,
            size = int(getattr(s, "size", 0) or 0),
            sha256 = sibling_sha256(s),
        )
        for s in filtered
        if isinstance(s.rfilename, str)
    ]
    return resolve_snapshot_ignore_patterns_for_files(filenames), expected_files


def _dataset_expected_files(info) -> list:
    from hub.utils.download_manifest import ExpectedFile
    return [
        ExpectedFile(
            path = s.rfilename,
            size = int(getattr(s, "size", 0) or 0),
            sha256 = sibling_sha256(s),
        )
        for s in info.siblings
        if isinstance(s.rfilename, str)
    ]


def _recover_manifest_after_download(
    repo_type: RepoType,
    repo_id: str,
    snapshot_path: str,
    mode: str,
    *,
    fetch_info,
    expected_files_from_info,
    label: str = "",
) -> None:
    """Best-effort manifest write for a download whose metadata was unavailable
    at start: re-fetch and record the expected files, else fall back to the
    on-disk file list. Shared by the model and dataset workers.

    A pre-existing manifest is authoritative and is preserved untouched. This is
    load-bearing: when access is lost on resume (token revoked/changed on a
    gated/private repo), snapshot_download returns the cached partial snapshot
    WITHOUT downloading, so rebuilding the manifest from on-disk files would record
    the partial set as expected and let _verify_completed_download certify a
    half-finished download as complete.

    The same hazard exists with NO prior manifest. When metadata is still
    unavailable here, leftover ``.incomplete`` blobs prove a cached partial was
    returned without downloading, so we fail (exit 1) instead of deriving a
    self-certifying manifest, leaving the partial intact for a later resume. That
    signal misses a file that never started (no ``.incomplete``), so a kill
    between files is accepted optimistically from the on-disk subset; a later
    metadata-bearing attempt writes the true manifest and catches any shortfall."""
    from hub.utils import download_manifest
    from hub.utils.hf_cache_state import has_active_incomplete_blobs

    if download_manifest.read_manifest(repo_type, repo_id, None) is not None:
        return

    try:
        if download_manifest.write_manifest(
            repo_type,
            repo_id,
            None,
            expected_files_from_info(fetch_info()),
            mode,
        ):
            return
        reason = "manifest write failed"
    except Exception as e:
        reason = f"{type(e).__name__}: {e}"

    if has_active_incomplete_blobs(repo_type, repo_id):
        print(
            f"{label}could not reach Hugging Face for {repo_id} and the copy on "
            "disk is still incomplete. Access to a private or restricted repo may "
            "have been lost (HF token removed or changed), the connection dropped, "
            "or Hugging Face is temporarily unavailable. Set a valid HF token or "
            "reconnect, then resume the download.",
            file = sys.stderr,
        )
        sys.exit(1)

    fallback_files = download_manifest.expected_files_from_snapshot_dir(Path(snapshot_path))
    if fallback_files and download_manifest.write_manifest(
        repo_type,
        repo_id,
        None,
        fallback_files,
        mode,
    ):
        print(
            f"{label}could not record the metadata manifest for {repo_id}, "
            "recorded one from the downloaded files so completion "
            f"is tracked ({reason})",
            file = sys.stderr,
        )
    else:
        print(
            f"{label}could not record the metadata manifest for {repo_id}, "
            f"{download_manifest.MANIFEST_DEGRADED_MARKER} ({reason})",
            file = sys.stderr,
        )


def _download_snapshot(repo_id: str, hf_token: str | None, mode: str) -> None:
    from huggingface_hub import snapshot_download
    from hub.utils.download_registry import prepare_cache_for_transport
    from hub.utils import download_manifest

    # One metadata fetch powers both the ignore-pattern decision (drop
    # consolidated.* when transformers weights exist) and the manifest's
    # expected_files. A failure is non-fatal: fall back to the legacy
    # ignore-pattern set (keeping consolidated) and skip the manifest, so
    # download proceeds without the verification + partial detection it enables.
    try:
        info = _model_info_with_retry(repo_id, hf_token)
    except Exception as e:
        print(
            f"metadata unavailable, downloading full snapshot for {repo_id} "
            f"({type(e).__name__}: {e})",
            file = sys.stderr,
        )
        info = None

    download_manifest.clear_cancel_marker("model", repo_id, None)
    if info is not None:
        ignore_patterns, expected_files = _snapshot_download_plan(info)
        # Written for every transport. The manifest verifies the finalized
        # files under snapshots/, which both transports produce identically
        # (XET also renames a full, correctly-sized blob into place). XET's
        # block-level dedup lives only in the chunk-cache the manifest never
        # inspects, so per-file size verification is valid regardless of transport.
        download_manifest.write_manifest("model", repo_id, None, expected_files, mode)
    else:
        ignore_patterns = list(SNAPSHOT_IGNORE_PATTERNS)
        expected_files = []

    purged = prepare_cache_for_transport("model", repo_id, mode)
    if purged:
        print(
            f"Purged {purged} untrusted partial blob(s) for {repo_id} "
            f"before starting {mode} download.",
            file = sys.stderr,
        )
    _preflight_disk_space("model", repo_id, expected_files)
    snapshot_path = snapshot_download(
        repo_id = repo_id,
        token = _hf_token_arg(hf_token),
        ignore_patterns = ignore_patterns,
        max_workers = 1,
    )
    if info is None:
        _recover_manifest_after_download(
            "model",
            repo_id,
            snapshot_path,
            mode,
            fetch_info = lambda: _model_info_with_retry(repo_id, hf_token),
            expected_files_from_info = lambda recovered: _snapshot_download_plan(recovered)[1],
        )
    _verify_completed_download(
        "model",
        repo_id,
        None,
        snapshot_path,
        metadata_unavailable = info is None,
    )


def _gguf_variant_target_plan(
    repo_id: str, variant: str, hf_token: str | None
) -> GgufVariantPlan | None:
    try:
        info = _model_info_with_retry(repo_id, hf_token)
    except Exception as e:
        print(
            f"metadata unavailable, cannot resolve GGUF variant '{variant}' "
            f"for {repo_id} ({type(e).__name__}: {e})",
            file = sys.stderr,
        )
        raise RuntimeError(
            f"Metadata unavailable while resolving GGUF variant '{variant}' " f"for {repo_id}"
        ) from e
    return build_gguf_variant_plans(list(info.siblings)).get(variant.lower())


def _download_gguf_variant(repo_id: str, variant: str, hf_token: str | None, mode: str) -> None:
    from huggingface_hub import snapshot_download
    from hub.utils.download_registry import prepare_cache_for_transport
    from hub.utils.hf_cache_state import has_active_incomplete_blobs
    from hub.utils import download_manifest

    metadata_unavailable = False
    try:
        plan = _gguf_variant_target_plan(repo_id, variant, hf_token)
    except RuntimeError:
        plan = None
        metadata_unavailable = True

    if not metadata_unavailable:
        if plan is None:
            print(
                f"No GGUF shards matching variant '{variant}' in {repo_id}",
                file = sys.stderr,
            )
            sys.exit(1)
        targets = list(plan.target_filenames)
        expected_files = list(plan.expected_files)
        main_blob_hashes = plan.main_hashes
        companion_blob_hashes = plan.companion_hashes
        download_manifest.write_manifest(
            "model",
            repo_id,
            variant,
            expected_files,
            mode,
        )
    else:
        # Metadata unreachable (offline / gated / private). Resume the exact
        # shards the original attempt recorded so snapshot_download can range
        # over the surviving .incomplete blobs without a model_info call.
        manifest = download_manifest.read_manifest("model", repo_id, variant)
        if manifest is None or not manifest.expected_files:
            print(
                f"Metadata unavailable and no manifest to resume GGUF "
                f"variant '{variant}' for {repo_id}",
                file = sys.stderr,
            )
            sys.exit(1)
        plan = plan_from_expected_files(variant, manifest.expected_files)
        targets = list(plan.target_filenames)
        expected_files = list(plan.expected_files)
        download_manifest.write_manifest(
            "model",
            repo_id,
            variant,
            expected_files,
            mode,
        )
        main_blob_hashes = plan.main_hashes
        companion_blob_hashes = plan.companion_hashes
        print(
            f"Metadata unavailable; resuming GGUF variant '{variant}' for "
            f"{repo_id} from the existing manifest.",
            file = sys.stderr,
        )

    download_manifest.clear_cancel_marker("model", repo_id, variant)
    purge_blob_hashes = main_blob_hashes
    if not main_blob_hashes:
        if has_active_incomplete_blobs("model", repo_id):
            print(
                f"GGUF variant '{variant}' for {repo_id} has partial cache state "
                "but no resolvable blob hashes; delete the partial download or "
                "retry when metadata is available.",
                file = sys.stderr,
            )
            sys.exit(1)
        purge_blob_hashes = frozenset()
        print(
            f"GGUF variant '{variant}' for {repo_id} has no resolvable blob "
            "hashes; starting without partial cache reuse.",
            file = sys.stderr,
        )
    # Main quant blobs are owned by this variant (variant-scoped marker). The
    # shared vision companion (mmproj) is judged by a separate companion marker
    # and never purged while a concurrent peer is writing it.
    purged = prepare_cache_for_transport(
        "model",
        repo_id,
        mode,
        variant,
        only_blob_hashes = purge_blob_hashes,
        companion_blob_hashes = companion_blob_hashes,
        protected_blob_hashes = _protected_blob_hashes(),
    )
    if purged:
        print(
            f"Purged {purged} untrusted partial blob(s) for {repo_id} "
            f"before starting {mode} download.",
            file = sys.stderr,
        )
    _preflight_disk_space("model", repo_id, expected_files)
    snapshot_path = snapshot_download(
        repo_id = repo_id,
        token = _hf_token_arg(hf_token),
        allow_patterns = targets,
        max_workers = 1,
    )
    _verify_completed_download(
        "model",
        repo_id,
        variant,
        snapshot_path,
        metadata_unavailable = metadata_unavailable,
    )
    if plan is not None:
        try:
            from hub.services.models.deletion import reclaim_replaced_gguf_variant
            reclaim_replaced_gguf_variant(
                repo_id,
                variant,
                plan.main_hashes,
                hf_token,
            )
        except Exception as e:
            print(
                f"Verified GGUF update for {repo_id} [{variant}], but stale-cache "
                f"reclaim failed ({type(e).__name__}: {e})",
                file = sys.stderr,
            )


def _download_dataset(repo_id: str, hf_token: str | None, mode: str) -> None:
    from huggingface_hub import snapshot_download
    from hub.utils.download_registry import prepare_cache_for_transport
    from hub.utils import download_manifest

    try:
        info = _dataset_info_with_retry(repo_id, hf_token)
    except Exception as e:
        print(
            f"dataset metadata unavailable, downloading full dataset for {repo_id} "
            f"({type(e).__name__}: {e})",
            file = sys.stderr,
        )
        info = None
    # Cancel-marker clear and manifest write run on every transport. See
    # _download_snapshot for why per-file size verification is valid under XET.
    download_manifest.clear_cancel_marker("dataset", repo_id, None)
    if info is not None:
        expected_files = _dataset_expected_files(info)
        download_manifest.write_manifest(
            "dataset",
            repo_id,
            None,
            expected_files,
            mode,
        )
    else:
        expected_files = []
    purged = prepare_cache_for_transport("dataset", repo_id, mode)
    if purged:
        print(
            f"Purged {purged} untrusted partial blob(s) for {repo_id} "
            f"before starting {mode} download.",
            file = sys.stderr,
        )
    _preflight_disk_space("dataset", repo_id, expected_files)
    snapshot_path = snapshot_download(
        repo_id = repo_id,
        token = _hf_token_arg(hf_token),
        repo_type = "dataset",
        max_workers = 1,
    )
    if info is None:
        _recover_manifest_after_download(
            "dataset",
            repo_id,
            snapshot_path,
            mode,
            fetch_info = lambda: _dataset_info_with_retry(repo_id, hf_token),
            expected_files_from_info = _dataset_expected_files,
            label = "dataset ",
        )
    _verify_completed_download(
        "dataset",
        repo_id,
        None,
        snapshot_path,
        metadata_unavailable = info is None,
    )


def main() -> None:
    parser = argparse.ArgumentParser(description = "HuggingFace Hub download worker")
    parser.add_argument("--repo-id", required = True)
    parser.add_argument("--variant", default = None)
    parser.add_argument("--dataset", action = "store_true")
    parser.add_argument("--transport", choices = ("http", "xet"), default = "http")
    parser.add_argument("--parent-pid", type = int, default = None)
    args = parser.parse_args()

    _install_signal_handlers()
    _install_parent_death_watchdog(args.parent_pid)

    hf_token = os.environ.get("HF_TOKEN") or None

    try:
        if args.dataset:
            _download_dataset(args.repo_id, hf_token, args.transport)
        elif args.variant:
            _download_gguf_variant(args.repo_id, args.variant, hf_token, args.transport)
        else:
            _download_snapshot(args.repo_id, hf_token, args.transport)
        sys.exit(0)
    except SystemExit:
        raise
    except Exception as e:
        # Surface a precise message so the UI doesn't show a generic "worker
        # exited with code 1". huggingface_hub's consistency check recommends
        # force_download=True to recover, which our "Restart" UI maps to a fresh
        # start by purging the partial via prepare_cache_for_transport.
        print(f"{type(e).__name__}: {e}", file = sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()