File size: 7,062 Bytes
4723b99 faef572 4723b99 faef572 4723b99 faef572 4723b99 faef572 4723b99 faef572 4723b99 faef572 4723b99 faef572 4723b99 faef572 4723b99 faef572 4723b99 faef572 4723b99 faef572 4723b99 faef572 4723b99 faef572 4723b99 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 | """Resolve and validate Hugging Face kernels before importing their binaries."""
from __future__ import annotations
import json
import os
from pathlib import Path
from typing import Any
def require_kernels_package() -> None:
"""Fail early when the precompiled-kernel runtime is not installed."""
try:
import kernels # noqa: F401
except ImportError as error:
raise RuntimeError(
"Precompiled FlashAttention requires requirements/features/flash.in."
) from error
def _kernel_lock_path() -> Path:
"""Return the kernel lock from a Hub artifact or source checkout."""
source_path = Path(__file__).resolve()
for candidate in (
source_path.parents[1] / "kernels.lock",
source_path.parents[3] / "kernels.lock",
):
if candidate.is_file():
return candidate
raise RuntimeError(
"kernels.lock is missing from the Hugging Face artifact or source checkout."
)
def _locked_entry(lock_path: Path, repository: str) -> dict[str, Any]:
try:
lock_entries = json.loads(lock_path.read_text(encoding="utf-8"))
except (OSError, json.JSONDecodeError) as error:
raise RuntimeError(f"Unable to read the kernel lock: {lock_path}") from error
if not isinstance(lock_entries, list):
raise RuntimeError("kernels.lock must contain a JSON list.")
if any(not isinstance(entry, dict) for entry in lock_entries):
raise RuntimeError("Every kernels.lock entry must be a JSON object.")
matches = [entry for entry in lock_entries if entry.get("repo_id") == repository]
if len(matches) != 1:
raise RuntimeError(
f"kernels.lock must contain exactly one entry for {repository!r}; found {len(matches)}."
)
return matches[0]
def _offline_mode() -> bool:
"""Return whether Hub access was explicitly disabled for this process."""
enabled_values = {"1", "on", "true", "yes"}
return any(
os.environ.get(name, "").strip().lower() in enabled_values
for name in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE")
)
def _offline_snapshot_path(repository: str, revision: str) -> Path:
"""Locate one exact, possibly sparse, kernel snapshot without using Hub APIs."""
try:
from huggingface_hub import constants
from huggingface_hub.file_download import repo_folder_name
except ImportError as error:
raise RuntimeError("Offline kernel loading requires huggingface-hub.") from error
cache_root = Path(os.environ.get("KERNELS_CACHE") or constants.HF_HUB_CACHE).resolve()
repository_root = (
cache_root / repo_folder_name(repo_id=repository, repo_type="kernel")
).resolve()
snapshot = repository_root / "snapshots" / revision
if not snapshot.is_dir():
raise RuntimeError(
f"The exact offline kernel snapshot {repository}@{revision} is not cached under "
f"{cache_root}. Run `kernels download` before enabling offline mode."
)
if repository_root not in snapshot.resolve().parents:
raise RuntimeError(f"Refusing kernel snapshot outside its cache repository: {snapshot}")
return snapshot
def _load_offline_locked_kernel(
repository: str,
revision: str,
variant_locks: dict[str, object],
) -> object:
"""Validate and import the one compatible variant from a sparse Hub snapshot."""
snapshot = _offline_snapshot_path(repository, revision)
build_root = snapshot / "build"
if not build_root.is_dir():
raise RuntimeError(f"The cached kernel snapshot has no build directory: {snapshot}")
cached_names = sorted(entry.name for entry in build_root.iterdir() if entry.is_dir())
unexpected = sorted(set(cached_names).difference(variant_locks))
if unexpected:
raise RuntimeError(
f"The cached {repository}@{revision} snapshot contains unlocked variants: "
f"{', '.join(unexpected)}"
)
try:
from kernels import get_local_kernel
from kernels.utils import validate_kernel
from kernels.variants import get_variants_local, resolve_variants
except ImportError as error:
raise RuntimeError(
"Precompiled FlashAttention requires requirements/features/flash.in."
) from error
cached_variants = get_variants_local(build_root)
parsed_names = {variant.variant_str for variant in cached_variants}
invalid = sorted(set(cached_names).difference(parsed_names))
if invalid:
raise RuntimeError(
f"The cached {repository}@{revision} snapshot contains invalid variants: "
f"{', '.join(invalid)}"
)
compatible, _ = resolve_variants(cached_variants)
if len(compatible) != 1:
names = ", ".join(variant.variant_str for variant in compatible) or "none"
raise RuntimeError(
f"Expected exactly one compatible cached variant for {repository}@{revision}; "
f"found {names}."
)
variant_name = compatible[0].variant_str
variant_lock = variant_locks.get(variant_name)
expected_hash = getattr(variant_lock, "hash", None)
if not isinstance(expected_hash, str) or not expected_hash.startswith("sha256-"):
raise RuntimeError(f"The kernel lock for {variant_name} has no valid SHA-256 digest.")
# Hash validation deliberately happens before import. This operates on the
# sparse snapshot produced by `kernels download` and avoids Hub 1.23's
# full-snapshot completeness check in offline mode.
validate_kernel(repo_path=snapshot, variant=variant_name, hash=expected_hash)
return get_local_kernel(build_root / variant_name)
def load_locked_kernel(repository: str, revision: str) -> object:
"""Download, hash-validate, then import one immutable precompiled kernel."""
require_kernels_package()
try:
from kernels import get_local_kernel, install_kernel
from kernels.lockfile import KernelLock
except ImportError as error:
raise RuntimeError(
"Precompiled FlashAttention requires requirements/features/flash.in."
) from error
lock_path = _kernel_lock_path()
kernel_lock = KernelLock.from_json(_locked_entry(lock_path, repository))
if kernel_lock.sha != revision:
raise RuntimeError(
f"The typed manifest pins {repository}@{revision}, but kernels.lock pins "
f"{kernel_lock.sha}."
)
if _offline_mode():
return _load_offline_locked_kernel(repository, revision, kernel_lock.variants)
# `install_kernel` downloads data without importing it and validates the
# selected build against the tracked variant hash. Only then is the exact
# validated path imported directly. Offline mode uses the sparse-cache
# resolver above because Hub 1.23 rejects partial snapshots as incomplete.
validated_path = install_kernel(
repository,
revision=kernel_lock.sha,
variant_locks=kernel_lock.variants,
)
return get_local_kernel(validated_path)
|