@tuned-tensor/cli 0.20.0 → 0.22.0

package/training/adapter/src/evaluate.py

@@ -10,12 +10,12 @@ from tempfile import TemporaryDirectory
 from typing import Any
 
 from model_contract import (
-    CERTIFIED_BASE_MODELS,
     CERTIFIED_BASE_MODEL,
     assert_certified_base_model,
     assert_certified_base_model_revision,
     assert_certified_model_config,
     chat_template_kwargs,
+    uncertified_load_kwargs,
 )
 
 MAX_ARCHIVE_MEMBERS = 20_000
@@ -145,10 +145,12 @@ def resolve_device(device: str) -> str:
     return device
 
 
-def _assert_certified_model_source(base_model: str) -> None:
-    if not Path(base_model).exists() and base_model not in CERTIFIED_BASE_MODELS:
+def _assert_certified_model_source(base_model: str, model_source: str) -> None:
+    # A hub source must be the requested model; certified or not, it is
+    # checked by assert_certified_base_model.
+    if not Path(model_source).exists() and model_source != base_model:
         raise ValueError(
-            f"The bundled evaluator does not certify {base_model!r}"
+            f"The bundled evaluator does not certify {model_source!r}"
         )
 
 
@@ -160,9 +162,14 @@ def load_text_model(
     base_model = str(payload["base_model"])
     model_source = str(payload.get("model_source", base_model))
     assert_certified_base_model(base_model, "Evaluation base model")
-    _assert_certified_model_source(model_source)
+    _assert_certified_model_source(base_model, model_source)
     revision = payload.get("base_model_revision")
-    assert_certified_base_model_revision(base_model, revision, "Evaluation base model revision")
+    assert_certified_base_model_revision(
+        base_model,
+        revision,
+        "Evaluation base model revision",
+        local_source=Path(model_source).is_dir(),
+    )
     device = resolve_device(str(payload.get("device", "cuda")))
     token = os.getenv("HF_TOKEN")
     source_kwargs: dict[str, Any] = {
@@ -184,7 +191,7 @@ def load_text_model(
     if device == "cuda":
         dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
 
-    model_kwargs: dict[str, Any] = dict(source_kwargs)
+    model_kwargs: dict[str, Any] = {**source_kwargs, **uncertified_load_kwargs(base_model)}
     if dtype is not None:
         model_kwargs["dtype"] = dtype
     if device == "cuda":

package/training/adapter/src/hub_download.py

@@ -0,0 +1,384 @@
+"""Resumable, observable Hugging Face Hub snapshot downloads.
+
+Model and dataset prefetch share this module so both get the same behavior:
+
+* the requested revision is resolved once to an immutable commit, so retries
+  can never mix files from two commits when a branch moves mid-download;
+* transient network failures are retried with exponential backoff. Completed
+  files stay in the Hugging Face cache, so each retry only fetches what is
+  still missing;
+* progress is emitted as ``@@tt-progress {json}`` lines on stdout, which the
+  Node runner renders as a progress bar. huggingface_hub's own tqdm bars are
+  disabled because redrawn carriage-return output is unreadable in logs.
+
+huggingface_hub is imported lazily so the pure helpers stay testable with the
+system Python.
+"""
+from __future__ import annotations
+
+import errno
+import hashlib
+import json
+import os
+import re
+import sys
+import threading
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Iterable
+
+PROGRESS_PREFIX = "@@tt-progress "
+COMMIT = re.compile(r"[0-9a-fA-F]{40}")
+DEFAULT_MAX_ATTEMPTS = 10
+MAX_BACKOFF_SECONDS = 60.0
+# A per-request read timeout that tolerates slow or briefly stalled Wi-Fi.
+DEFAULT_DOWNLOAD_TIMEOUT_SECONDS = "60"
+
+
+@dataclass(frozen=True)
+class PlannedFile:
+    """One repository file expected in the snapshot, keyed by its cache blob name."""
+
+    path: str
+    size: int
+    etag: str
+
+
+def configure_download_environment() -> None:
+    """Must run before huggingface_hub is imported; it reads these at import."""
+    os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
+    os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", DEFAULT_DOWNLOAD_TIMEOUT_SECONDS)
+
+
+def repo_folder_name(repo_id: str, repo_type: str) -> str:
+    return f"{repo_type}s--{repo_id.replace('/', '--')}"
+
+
+def max_attempts_from_env(environ: dict[str, str] | None = None) -> int:
+    raw = (environ if environ is not None else os.environ).get("TT_HF_DOWNLOAD_MAX_ATTEMPTS")
+    if raw is None or not raw.strip():
+        return DEFAULT_MAX_ATTEMPTS
+    try:
+        value = int(raw)
+    except ValueError as exc:
+        raise ValueError("TT_HF_DOWNLOAD_MAX_ATTEMPTS must be an integer from 1 to 100") from exc
+    if not 1 <= value <= 100:
+        raise ValueError("TT_HF_DOWNLOAD_MAX_ATTEMPTS must be an integer from 1 to 100")
+    return value
+
+
+def backoff_seconds(attempt: int) -> float:
+    """Delay before retry ``attempt`` (1-based): 2, 4, 8, ... capped at 60s."""
+    return min(MAX_BACKOFF_SECONDS, float(2 ** attempt))
+
+
+class IncompleteDownloadError(RuntimeError):
+    """snapshot_download returned a cached folder that still lacks planned files."""
+
+
+def is_retryable(error: BaseException) -> bool:
+    """Retry network and server failures; fail fast on caller or auth errors.
+
+    huggingface_hub reports a connection failure on a file that is not cached
+    yet as LocalEntryNotFoundError (a FileNotFoundError), so that and
+    IncompleteSnapshotError are retried during downloads.
+    """
+    if isinstance(error, (KeyboardInterrupt, SystemExit)):
+        return False
+    if isinstance(error, OSError) and error.errno in (errno.ENOSPC, errno.EACCES, errno.EROFS, errno.EDQUOT):
+        return False
+    name = type(error).__name__
+    if name in {
+        "RepositoryNotFoundError",
+        "GatedRepoError",
+        "RevisionNotFoundError",
+        "RemoteEntryNotFoundError",
+        "DisabledRepoError",
+        "OfflineModeIsEnabled",
+    }:
+        return False
+    if isinstance(error, (ValueError, TypeError)):
+        return False
+    response = getattr(error, "response", None)
+    status = getattr(response, "status_code", None)
+    if isinstance(status, int) and 400 <= status < 500 and status not in (408, 429):
+        return False
+    return True
+
+
+def plan_files(siblings: Iterable[Any]) -> list[PlannedFile]:
+    """Build the expected cache inventory from repo_info(files_metadata=True)."""
+    planned: list[PlannedFile] = []
+    for sibling in siblings:
+        lfs = getattr(sibling, "lfs", None)
+        etag = getattr(lfs, "sha256", None) if lfs is not None else None
+        etag = etag or getattr(sibling, "blob_id", None)
+        size = getattr(sibling, "size", None)
+        if not etag or not isinstance(size, int):
+            # Without a blob name and size the file cannot be tracked; the
+            # download still fetches and verifies it, only progress is coarser.
+            continue
+        planned.append(PlannedFile(path=sibling.rfilename, size=size, etag=str(etag)))
+    return planned
+
+
+def downloaded_bytes(blobs: Path, planned: PlannedFile) -> tuple[int, bool]:
+    """Bytes present for one planned file and whether its blob is complete."""
+    if (blobs / planned.etag).is_file():
+        return planned.size, True
+    partial = 0
+    try:
+        for candidate in blobs.glob(f"{planned.etag}*.incomplete"):
+            try:
+                partial = max(partial, candidate.stat().st_size)
+            except OSError:
+                continue
+    except OSError:
+        return 0, False
+    return min(partial, planned.size), False
+
+
+def measure(blobs: Path, planned: list[PlannedFile]) -> dict[str, int]:
+    completed_bytes = 0
+    completed_files = 0
+    for item in planned:
+        size, done = downloaded_bytes(blobs, item)
+        completed_bytes += size
+        completed_files += int(done)
+    return {
+        "completed_bytes": completed_bytes,
+        "total_bytes": sum(item.size for item in planned),
+        "files_completed": completed_files,
+        "files_total": len(planned),
+    }
+
+
+_STDOUT_LOCK = threading.Lock()
+
+
+def write_line(text: str, stream: Any = None) -> None:
+    """Write one whole line; the progress thread and log messages share stdout."""
+    target = stream if stream is not None else sys.stdout
+    with _STDOUT_LOCK:
+        target.write(text + "\n")
+        target.flush()
+
+
+def emit_progress(event: dict[str, Any], stream: Any = None) -> None:
+    write_line(PROGRESS_PREFIX + json.dumps(event, separators=(",", ":")), stream)
+
+
+class ProgressMonitor:
+    """Polls the cache blob directory and emits progress while a download runs."""
+
+    def __init__(
+        self,
+        blobs: Path,
+        planned: list[PlannedFile],
+        label: str,
+        emit: Callable[[dict[str, Any]], None] = emit_progress,
+        interval: float = 0.5,
+        heartbeat: float = 5.0,
+    ) -> None:
+        self.blobs = blobs
+        self.planned = planned
+        self.label = label
+        self.emit = emit
+        self.interval = interval
+        self.heartbeat = heartbeat
+        self._stop = threading.Event()
+        self._thread: threading.Thread | None = None
+        self._last: dict[str, int] | None = None
+        self._last_emit = 0.0
+
+    def snapshot(self, force: bool = False, phase: str = "download") -> None:
+        current = measure(self.blobs, self.planned)
+        now = time.monotonic()
+        if not force and current == self._last and now - self._last_emit < self.heartbeat:
+            return
+        self._last = current
+        self._last_emit = now
+        self.emit({"label": self.label, "phase": phase, **current})
+
+    def _run(self) -> None:
+        while not self._stop.wait(self.interval):
+            try:
+                self.snapshot()
+            except Exception:  # Progress is best effort; never break the download.
+                pass
+
+    def __enter__(self) -> "ProgressMonitor":
+        self.snapshot(force=True)
+        self._thread = threading.Thread(target=self._run, name="tt-progress", daemon=True)
+        self._thread.start()
+        return self
+
+    def __exit__(self, *_: object) -> None:
+        self._stop.set()
+        if self._thread:
+            self._thread.join(timeout=2)
+
+
+def with_retries(
+    action: Callable[[], Any],
+    *,
+    description: str,
+    max_attempts: int,
+    sleep: Callable[[float], None] | None = None,
+    log: Callable[[str], None] = write_line,
+) -> Any:
+    for attempt in range(1, max_attempts + 1):
+        try:
+            return action()
+        except Exception as error:  # noqa: BLE001 - classified below
+            if attempt >= max_attempts or not is_retryable(error):
+                raise
+            delay = backoff_seconds(attempt)
+            log(
+                f"{description} interrupted ({type(error).__name__}: {error}). "
+                f"Retrying in {delay:.0f}s (attempt {attempt + 1}/{max_attempts}); "
+                "completed files are kept in the cache."
+            )
+            (sleep or time.sleep)(delay)
+    raise AssertionError("unreachable")
+
+
+def verify_blob(path: Path) -> bool:
+    """Check a cached file against its content-addressed blob name.
+
+    Returns True when the blob name is a recognized digest and it matches;
+    False when the name carries no digest (for example a cache without
+    symlinks). Raises ValueError on a mismatch.
+    """
+    return verify_blob_digest(path)[0]
+
+
+def verify_blob_digest(path: Path) -> tuple[bool, str]:
+    """Like verify_blob, also returning the file's SHA-256 from the same read."""
+    blob = path.resolve()
+    expected = blob.name.lower()
+    sha256 = hashlib.sha256()
+    git_sha1 = hashlib.sha1(f"blob {blob.stat().st_size}\0".encode())
+    with blob.open("rb") as source:
+        for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""):
+            sha256.update(chunk)
+            git_sha1.update(chunk)
+    digests = {64: sha256.hexdigest(), 40: git_sha1.hexdigest()}
+    if len(expected) in digests and all(character in "0123456789abcdef" for character in expected):
+        if digests[len(expected)] != expected:
+            raise ValueError(f"Cached Hugging Face blob checksum mismatch: {path}")
+        return True, digests[64]
+    return False, digests[64]
+
+
+def download_snapshot(
+    *,
+    repo_id: str,
+    repo_type: str,
+    revision: str | None,
+    allow_patterns: list[str] | None = None,
+    ignore_patterns: list[str] | None = None,
+    local_files_only: bool = False,
+    label: str,
+) -> tuple[str, str]:
+    """Download (or verify locally) one snapshot. Returns (path, commit sha)."""
+    from huggingface_hub import constants, snapshot_download
+
+    token = os.getenv("HF_TOKEN")
+    if local_files_only:
+        path = snapshot_download(
+            repo_id=repo_id,
+            repo_type=repo_type,
+            revision=revision,
+            token=token,
+            allow_patterns=allow_patterns,
+            ignore_patterns=ignore_patterns,
+            local_files_only=True,
+        )
+        return path, Path(path).name
+
+    if revision and COMMIT.fullmatch(revision):
+        # A pinned commit that is already fully cached needs no network, so
+        # prefetch keeps working offline. huggingface_hub checks completeness
+        # against its cached tree listing and raises if files are missing.
+        try:
+            path = snapshot_download(
+                repo_id=repo_id,
+                repo_type=repo_type,
+                revision=revision,
+                token=token,
+                allow_patterns=allow_patterns,
+                ignore_patterns=ignore_patterns,
+                local_files_only=True,
+            )
+            write_line(f"{label}: {repo_id}@{revision} is already cached; no download needed.")
+            return path, revision
+        except Exception:  # noqa: BLE001 - not cached or incomplete: download below
+            pass
+
+    from huggingface_hub import HfApi
+    from huggingface_hub.utils import filter_repo_objects
+
+    max_attempts = max_attempts_from_env()
+    api = HfApi(token=token)
+    info = with_retries(
+        lambda: api.repo_info(repo_id, repo_type=repo_type, revision=revision, files_metadata=True),
+        description=f"Resolving {repo_id}",
+        max_attempts=max_attempts,
+    )
+    commit = str(info.sha)
+    siblings = list(filter_repo_objects(
+        info.siblings or [],
+        allow_patterns=allow_patterns,
+        ignore_patterns=ignore_patterns,
+        key=lambda sibling: sibling.rfilename,
+    ))
+    planned = plan_files(siblings)
+    blobs = Path(constants.HF_HUB_CACHE) / repo_folder_name(repo_id, repo_type) / "blobs"
+    blobs.mkdir(parents=True, exist_ok=True)
+    total = sum(item.size for item in planned)
+    write_line(
+        f"{label}: {repo_id}@{commit} — {len(planned)} file(s), {total} bytes. "
+        "Interrupted downloads can be restarted; completed files are reused."
+    )
+    monitor = ProgressMonitor(blobs, planned, label)
+    with monitor:
+        def attempt() -> str:
+            path = snapshot_download(
+                repo_id=repo_id,
+                repo_type=repo_type,
+                revision=commit,
+                token=token,
+                allow_patterns=allow_patterns,
+                ignore_patterns=ignore_patterns,
+            )
+            # If the Hub becomes unreachable mid-way, huggingface_hub may hand
+            # back the partial cached folder; treat that as a retryable failure.
+            missing = [sibling.rfilename for sibling in siblings if not (Path(path) / sibling.rfilename).exists()]
+            if missing:
+                raise IncompleteDownloadError(f"{len(missing)} file(s) still missing, e.g. {missing[0]}")
+            return path
+
+        path = with_retries(attempt, description=f"Downloading {repo_id}", max_attempts=max_attempts)
+    monitor.snapshot(force=True, phase="downloaded")
+    if revision and revision != commit:
+        # Downloading by commit skips huggingface_hub's ref bookkeeping. Record
+        # the branch or tag -> commit mapping it would have written, so later
+        # offline resolution of an unpinned revision finds this snapshot.
+        write_ref(Path(constants.HF_HUB_CACHE) / repo_folder_name(repo_id, repo_type), revision, commit)
+    elif not revision:
+        write_ref(Path(constants.HF_HUB_CACHE) / repo_folder_name(repo_id, repo_type), "main", commit)
+    return path, commit
+
+
+def write_ref(repository: Path, revision: str, commit: str) -> None:
+    if revision.lower() == commit.lower():
+        return
+    ref = repository / "refs" / revision
+    if repository.resolve() not in ref.resolve().parents:
+        raise ValueError(f"Invalid Hugging Face revision name: {revision}")
+    ref.parent.mkdir(parents=True, exist_ok=True)
+    temporary = ref.with_name(f".{ref.name}.{os.getpid()}.tmp")
+    temporary.write_text(commit, encoding="utf-8")
+    os.replace(temporary, ref)

package/training/adapter/src/model_contract.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import re
 from typing import Any
 
 
@@ -15,6 +16,10 @@ NEMOTRON_BASE_MODEL_REVISION = "ce38b6ab8b252b4b8ee7165b4605e93191cafd73"
 # Immutable Hugging Face revision reviewed and certified for Muse Glimmer.
 MUSE_GLIMMER_BASE_MODEL_REVISION = "a4e59da52a7bc87ae7251dd5545c0dd437c44b68"
 CERTIFIED_BASE_MODELS = (CERTIFIED_BASE_MODEL, NEMOTRON_BASE_MODEL, MUSE_GLIMMER_BASE_MODEL)
+# Any other Hugging Face repo may be trained as an uncertified causal LM when
+# the run pins an immutable commit. Node enforces the same contract.
+HUGGING_FACE_REPO_ID = re.compile(r"[A-Za-z0-9][A-Za-z0-9._-]{0,95}/[A-Za-z0-9][A-Za-z0-9._-]{0,95}")
+COMMIT_SHA = re.compile(r"[0-9a-fA-F]{40}")
 CERTIFIED_QWEN_TEXT_CONFIG = {
     "model_type": "qwen3_5_text",
     "hidden_size": 2048,
@@ -58,6 +63,9 @@ def assert_certified_model_config(
     label: str = "base-model config",
     expected_model_id: str | None = None,
 ) -> None:
+    if expected_model_id and not is_certified_base_model(expected_model_id):
+        # No reviewed architecture exists for an uncertified model.
+        return
     model_type = _field(value, "model_type")
     architectures = _field(value, "architectures")
 
@@ -126,17 +134,30 @@ def assert_certified_model_config(
         )
 
 
+def is_certified_base_model(model_id: str) -> bool:
+    return model_id in CERTIFIED_BASE_MODELS
+
+
 def assert_certified_base_model(model_id: str, label: str = "base model") -> None:
-    if model_id not in CERTIFIED_BASE_MODELS:
-        supported = ", ".join(CERTIFIED_BASE_MODELS)
-        raise ValueError(f"{label} must be one of {supported}; got {model_id!r}")
+    """Accept a certified model or an uncertified Hugging Face repo id."""
+    if is_certified_base_model(model_id) or HUGGING_FACE_REPO_ID.fullmatch(model_id):
+        return
+    supported = ", ".join(CERTIFIED_BASE_MODELS)
+    raise ValueError(f"{label} must be one of {supported} or a Hugging Face repo id; got {model_id!r}")
 
 
-def assert_certified_base_model_revision(base_model: str, revision: str | None, label: str = "base model revision") -> None:
+def assert_certified_base_model_revision(
+    base_model: str,
+    revision: str | None,
+    label: str = "base model revision",
+    *,
+    local_source: bool = False,
+) -> None:
     """Require the certified immutable revision for pinned training models.
 
     Qwen, Nemotron, and Muse Glimmer loads are bound here; local snapshot
-    contents are additionally verified by the Node runtime.
+    contents are additionally verified by the Node runtime. An uncertified
+    model must pin a commit SHA unless it loads from a local directory.
     """
     expected = {
         CERTIFIED_BASE_MODEL: QWEN_BASE_MODEL_REVISION,
@@ -144,6 +165,13 @@ def assert_certified_base_model_revision(base_model: str, revision: str | None,
         MUSE_GLIMMER_BASE_MODEL: MUSE_GLIMMER_BASE_MODEL_REVISION,
     }.get(base_model)
     if expected is None:
+        if revision is None and local_source:
+            return
+        if revision is None or not COMMIT_SHA.fullmatch(str(revision)):
+            raise ValueError(
+                f"{label} must be a 40-character commit SHA for uncertified base model {base_model}; "
+                f"got {revision!r}"
+            )
         return
     if revision is None:
         raise ValueError(f"{label} is required for {base_model}")
@@ -167,3 +195,8 @@ def chat_template_kwargs(model_id: str) -> dict[str, Any]:
     if model_id == NEMOTRON_BASE_MODEL:
         return {"enable_thinking": False}
     return {}
+
+
+def uncertified_load_kwargs(model_id: str) -> dict[str, Any]:
+    """Refuse pickle weights for models TT has not reviewed."""
+    return {} if is_certified_base_model(model_id) else {"use_safetensors": True}

package/training/adapter/src/prefetch.py

@@ -1,15 +1,16 @@
 import argparse
-import hashlib
 import json
 import os
 import re
 from pathlib import Path
 from typing import Any
 
+from hub_download import configure_download_environment, download_snapshot, verify_blob
 from model_contract import (
     assert_certified_base_model,
     assert_certified_base_model_revision,
     assert_certified_model_config,
+    is_certified_base_model,
 )
 
 ALLOW_PATTERNS = [
@@ -76,15 +77,6 @@ def configure_hugging_face_cache(cache_home: str | None) -> None:
         os.environ.pop(deprecated, None)
 
 
-def hash_file(path: Path, algorithm: str, prefix: bytes = b"") -> str:
-    digest = hashlib.new(algorithm)
-    digest.update(prefix)
-    with path.open("rb") as source:
-        for chunk in iter(lambda: source.read(8 * 1024 * 1024), b""):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
 def verify_snapshot(snapshot: Path, expected_model_id: str | None = None) -> tuple[list[Path], int]:
     config = snapshot / "config.json"
     if not config.is_file() or config.stat().st_size == 0:
@@ -99,8 +91,15 @@ def verify_snapshot(snapshot: Path, expected_model_id: str | None = None) -> tup
     except (OSError, json.JSONDecodeError) as exc:
         raise ValueError(f"Cached snapshot has an invalid config.json: {snapshot}") from exc
 
+    if expected_model_id and not is_certified_base_model(expected_model_id):
+        if not any(path.stat().st_size > 0 for path in snapshot.glob("*.safetensors")):
+            raise ValueError("Uncertified base models require non-empty safetensors weights; pickle-only checkpoints are unsupported")
+
     weight_files = sorted(snapshot.glob("*.safetensors")) + sorted(snapshot.glob("pytorch_model*.bin"))
     index_files = sorted(snapshot.glob("*.safetensors.index.json")) + sorted(snapshot.glob("pytorch_model*.bin.index.json"))
+    # Keep lexical snapshot names: normal HF file symlinks target sibling blobs.
+    files = [path for path in snapshot.rglob("*") if path.is_file()]
+    inventory = set(files)
     for index in index_files:
         try:
             weight_map = json.loads(index.read_text(encoding="utf-8")).get("weight_map", {})
@@ -108,8 +107,13 @@ def verify_snapshot(snapshot: Path, expected_model_id: str | None = None) -> tup
             raise ValueError(f"Cached snapshot has an invalid weight index: {index}") from exc
         if not isinstance(weight_map, dict) or not weight_map:
             raise ValueError(f"Cached snapshot has an empty weight index: {index}")
-        for name in set(weight_map.values()):
-            if not isinstance(name, str) or not (snapshot / name).is_file():
+        for name in weight_map.values():
+            if (not isinstance(name, str) or not name
+                    or re.match(r"^(?:[A-Za-z]:|[/\\])", name)
+                    or ".." in re.split(r"[/\\]", name)):
+                raise ValueError("Cached snapshot has an invalid indexed weight shard path")
+            shard = snapshot / name
+            if shard not in inventory or shard.stat().st_size == 0:
                 raise ValueError(f"Cached snapshot is missing an indexed weight shard: {name}")
     if not weight_files:
         raise ValueError(f"Cached snapshot contains no Transformers model weights: {snapshot}")
@@ -129,20 +133,7 @@ def verify_snapshot(snapshot: Path, expected_model_id: str | None = None) -> tup
     if not tokenizer_config.is_file() or not any(path.is_file() and path.stat().st_size > 0 for path in vocabulary):
         raise ValueError(f"Cached snapshot is missing tokenizer metadata or vocabulary files: {snapshot}")
 
-    files = [path for path in snapshot.rglob("*") if path.is_file()]
-    verified_blobs = 0
-    for path in files:
-        blob = path.resolve()
-        expected = blob.name.lower()
-        if len(expected) == 64 and all(character in "0123456789abcdef" for character in expected):
-            actual = hash_file(blob, "sha256")
-        elif len(expected) == 40 and all(character in "0123456789abcdef" for character in expected):
-            actual = hash_file(blob, "sha1", f"blob {blob.stat().st_size}\0".encode())
-        else:
-            continue
-        if actual != expected:
-            raise ValueError(f"Cached Hugging Face blob checksum mismatch: {path}")
-        verified_blobs += 1
+    verified_blobs = sum(1 for path in files if verify_blob(path))
     return files, verified_blobs
 
 
@@ -166,17 +157,20 @@ def main() -> None:
 
     # Import only after the cache environment is final. huggingface_hub reads
     # these variables while importing its constants module.
-    from huggingface_hub import constants, snapshot_download
+    configure_download_environment()
+    from huggingface_hub import constants
 
-    action = "Verifying cached" if payload.get("local_files_only") else "Prefetching"
+    local_only = bool(payload.get("local_files_only", False))
+    action = "Verifying cached" if local_only else "Prefetching"
     print(f"{action} {base_model} in {constants.HF_HUB_CACHE}...", flush=True)
-    snapshot_path = snapshot_download(
+    snapshot_path, _ = download_snapshot(
         repo_id=base_model,
+        repo_type="model",
         revision=payload.get("revision"),
-        token=os.getenv("HF_TOKEN"),
         allow_patterns=ALLOW_PATTERNS,
-        ignore_patterns=IGNORE_PATTERNS,
-        local_files_only=bool(payload.get("local_files_only", False)),
+        ignore_patterns=IGNORE_PATTERNS + ([] if is_certified_base_model(base_model) else ["pytorch_model*.bin"]),
+        local_files_only=local_only,
+        label="model_prefetch",
     )
     snapshot = Path(snapshot_path)
     snapshot_revision = snapshot.name