diff --git a/kernels/src/kernels/cli/__init__.py b/kernels/src/kernels/cli/__init__.py index ed7390bfc..c8bb7cff4 100644 --- a/kernels/src/kernels/cli/__init__.py +++ b/kernels/src/kernels/cli/__init__.py @@ -8,7 +8,7 @@ from kernels.cli.verify_signature import verify_signature from kernels.cli.versions import print_kernel_versions from kernels.compat import tomllib -from kernels.lockfile import KernelLock, get_kernel_locks +from kernels.locking import KernelLock, get_kernel_locks from kernels.utils import ( install_kernel, install_kernel_all_variants, @@ -153,14 +153,12 @@ def download_kernels(args): install_kernel_all_variants( kernel_lock.repo_id, revision=kernel_lock.sha, - variant_locks=kernel_lock.variants, ) else: try: install_kernel( kernel_lock.repo_id, revision=kernel_lock.sha, - variant_locks=kernel_lock.variants, ) except FileNotFoundError as e: print(e, file=sys.stderr) diff --git a/kernels/src/kernels/layer/func.py b/kernels/src/kernels/layer/func.py index 4994dd658..c4a3af8e3 100644 --- a/kernels/src/kernels/layer/func.py +++ b/kernels/src/kernels/layer/func.py @@ -5,9 +5,11 @@ from typing import TYPE_CHECKING, Protocol, Type from .._versions import select_revision_or_version +from ..locking import ( + get_caller_locked_kernel_revision, + get_locked_kernel_revision, +) from ..utils import ( - _get_caller_locked_kernel, - _get_locked_kernel, get_kernel, get_local_kernel, ) @@ -298,10 +300,10 @@ def __init__( def _resolve_revision(self) -> str: if self._lockfile is None: - locked_sha = _get_caller_locked_kernel(self._repo_id) + locked_sha = get_caller_locked_kernel_revision(self._repo_id) else: with open(self._lockfile, "r") as f: - locked_sha = _get_locked_kernel(self._repo_id, f.read()) + locked_sha = get_locked_kernel_revision(self._repo_id, f.read()) if locked_sha is None: raise ValueError(f"Kernel `{self._repo_id}` is not locked") diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py index 3c9c41cee..903ec14ea 100644 --- a/kernels/src/kernels/layer/layer.py +++ b/kernels/src/kernels/layer/layer.py @@ -10,9 +10,11 @@ from typing import TYPE_CHECKING, Callable, Protocol, Type from .._versions import select_revision_or_version +from ..locking import ( + get_caller_locked_kernel_revision, + get_locked_kernel_revision, +) from ..utils import ( - _get_caller_locked_kernel, - _get_locked_kernel, get_kernel, get_local_kernel, ) @@ -207,10 +209,10 @@ def __init__( def _resolve_revision(self) -> str: if self._lockfile is None: - locked_sha = _get_caller_locked_kernel(self._repo_id) + locked_sha = get_caller_locked_kernel_revision(self._repo_id) else: with open(self._lockfile, "r") as f: - locked_sha = _get_locked_kernel(self._repo_id, f.read()) + locked_sha = get_locked_kernel_revision(self._repo_id, f.read()) if locked_sha is None: raise ValueError(f"Kernel `{self._repo_id}` is not locked") diff --git a/kernels/src/kernels/lockfile.py b/kernels/src/kernels/locking.py similarity index 71% rename from kernels/src/kernels/lockfile.py rename to kernels/src/kernels/locking.py index 749088325..1cafa4e7b 100644 --- a/kernels/src/kernels/lockfile.py +++ b/kernels/src/kernels/locking.py @@ -1,6 +1,11 @@ import hashlib +import importlib.metadata +import inspect +import json from dataclasses import dataclass +from importlib.metadata import Distribution from pathlib import Path +from types import ModuleType from huggingface_hub.dataclasses import strict from huggingface_hub.hf_api import RepoFile @@ -129,3 +134,47 @@ def write_egg_lockfile(cmd, basename, filename): data = open(lock_path, "r").read() cmd.write_or_delete_file(basename, filename, data) + + +def get_locked_kernel_revision(repo_id: str, lock_json: str) -> str | None: + for kernel_lock_json in json.loads(lock_json): + kernel_lock = KernelLock.from_json(kernel_lock_json) + if kernel_lock.repo_id == repo_id: + return kernel_lock.sha + return None + + +def get_caller_locked_kernel_revision(repo_id: str) -> str | None: + for dist in _get_caller_distributions(): + lock_json = dist.read_text("kernels.lock") + if lock_json is None: + continue + locked_sha = get_locked_kernel_revision(repo_id, lock_json) + if locked_sha is not None: + return locked_sha + return None + + +def _get_caller_distributions() -> list[Distribution]: + module = _get_caller_module() + if module is None: + return [] + + # Look up all possible distributions that this module could be from. + package = module.__name__.split(".")[0] + dist_names = importlib.metadata.packages_distributions().get(package) + if dist_names is None: + return [] + + return [importlib.metadata.distribution(dist_name) for dist_name in dist_names] + + +def _get_caller_module() -> ModuleType | None: + stack = inspect.stack() + # Get first module in the stack that is not the current module. + first_module = inspect.getmodule(stack[0][0]) + for frame in stack[1:]: + module = inspect.getmodule(frame[0]) + if module is not None and module != first_module: + return module + return first_module diff --git a/kernels/src/kernels/utils.py b/kernels/src/kernels/utils.py index cd2f10002..e0a6d9845 100644 --- a/kernels/src/kernels/utils.py +++ b/kernels/src/kernels/utils.py @@ -1,15 +1,10 @@ import functools -import hashlib import importlib -import importlib.metadata -import inspect -import json import os import platform import sys import warnings from dataclasses import dataclass -from importlib.metadata import Distribution from pathlib import Path from types import ModuleType @@ -22,7 +17,10 @@ from kernels.backends import _backend, _select_backend from kernels.compat import has_torch, has_tvm_ffi from kernels.deps import validate_dependencies -from kernels.lockfile import KernelLock, VariantLock +from kernels.locking import ( + get_caller_locked_kernel_revision, + get_locked_kernel_revision, +) from kernels.status import resolve_status from kernels.variants import ( Decision, @@ -268,7 +266,6 @@ def install_kernel( revision: str, local_files_only: bool = False, backend: str | None = None, - variant_locks: dict[str, VariantLock] | None = None, user_agent: str | dict | None = None, validate_dependencies: bool = False, ) -> Path: @@ -308,7 +305,6 @@ def install_kernel( repo_id, revision=revision, backend=backend, - variant_locks=variant_locks, ) # For locally downloaded kernels, we run the validation after resolving the path if validate_dependencies: @@ -357,7 +353,6 @@ def install_kernel( return _find_kernel_in_repo_path( repo_path, variant=variant, - variant_locks=variant_locks, ) except FileNotFoundError: raise FileNotFoundError(f"Cannot install kernel from repo {repo_id} (revision: {revision})") @@ -369,7 +364,6 @@ def _resolve_local_variant_path( *, revision: str, backend: str | None = None, - variant_locks: dict[str, VariantLock] | None = None, ) -> Path: """Resolve a kernel variant path from the local Hugging Face cache only. @@ -382,6 +376,7 @@ def _resolve_local_variant_path( api.snapshot_download( repo_id, repo_type="kernel", + ignore_patterns=_BYTECODE_IGNORE_PATTERNS, cache_dir=CACHE_DIR, revision=revision, local_files_only=True, @@ -402,41 +397,22 @@ def _resolve_local_variant_path( f"Cannot find a build variant for this system in {repo_id} (revision: {revision}):\n\n{variants_trace_str(status)}" ) - allow_patterns = [f"build/{variant.variant_str}/*"] - ignore_patterns = _BYTECODE_IGNORE_PATTERNS - repo_path = Path( - str( - api.snapshot_download( - repo_id, - repo_type="kernel", - allow_patterns=allow_patterns, - ignore_patterns=ignore_patterns, - cache_dir=CACHE_DIR, - revision=revision, - local_files_only=True, - ) - ) + return _find_kernel_in_repo_path( + local_repo_path, + variant=variant, ) - return _find_kernel_in_repo_path(repo_path, variant=variant, variant_locks=variant_locks) def _find_kernel_in_repo_path( repo_path: Path, *, variant: Variant, - variant_locks: dict[str, VariantLock] | None = None, ) -> Path: variant_str = variant.variant_str variant_path = repo_path / "build" / variant_str if not variant_path.exists(): raise FileNotFoundError(f"Variant path does not exist: `{variant_path}`") - if variant_locks is not None: - variant_lock = variant_locks.get(variant_str) - if variant_lock is None: - raise ValueError(f"No lock found for build variant: {variant}") - validate_kernel(repo_path=repo_path, variant=variant_str, hash=variant_lock.hash) - return variant_path @@ -444,8 +420,6 @@ def install_kernel_all_variants( repo_id: str, *, revision: str, - local_files_only: bool = False, - variant_locks: dict[str, VariantLock] | None = None, ) -> Path: api = _get_hf_api() @@ -458,21 +432,10 @@ def install_kernel_all_variants( ignore_patterns=_BYTECODE_IGNORE_PATTERNS, cache_dir=CACHE_DIR, revision=revision, - local_files_only=local_files_only, ) ) ) - if variant_locks is not None: - for entry in (repo_path / "build").iterdir(): - variant = entry.parts[-1] - - variant_lock = variant_locks.get(variant) - if variant_lock is None: - raise ValueError(f"No lock found for build variant: {variant}") - - validate_kernel(repo_path=repo_path, variant=variant, hash=variant_lock.hash) - return repo_path / "build" @@ -676,7 +639,6 @@ def load_kernel( *, lockfile: Path | None, backend: str | None = None, - revision: str | None = None, ) -> ModuleType: """ Get a pre-downloaded, locked kernel. @@ -691,22 +653,15 @@ def load_kernel( backend (`str`, *optional*): The backend to load the kernel for. Can only be `cpu` or the backend that Torch is compiled for. The backend will be detected automatically if not provided. - revision (`str`, *optional*): - The specific revision (branch, tag, or commit) to download. Cannot be used together with `version`. Returns: `ModuleType`: The imported kernel module. """ - if lockfile is not None and revision is not None: - raise ValueError("`lockfile` and `revision` both cannot be specified at the same time.") - - if lockfile is None and revision is None: - locked_sha = _get_caller_locked_kernel(repo_id) - elif revision is not None: - locked_sha = revision - elif lockfile is not None: + if lockfile is None: + locked_sha = get_caller_locked_kernel_revision(repo_id) + else: with open(lockfile, "r") as f: - locked_sha = _get_locked_kernel(repo_id, f.read()) + locked_sha = get_locked_kernel_revision(repo_id, f.read()) if locked_sha is None: raise ValueError( @@ -744,7 +699,7 @@ def get_locked_kernel(repo_id: str, local_files_only: bool = False) -> ModuleTyp Returns: `ModuleType`: The imported kernel module. """ - locked_sha = _get_caller_locked_kernel(repo_id) + locked_sha = get_caller_locked_kernel_revision(repo_id) if locked_sha is None: raise ValueError(f"Kernel `{repo_id}` is not locked") @@ -759,103 +714,6 @@ def get_locked_kernel(repo_id: str, local_files_only: bool = False) -> ModuleTyp return _import_from_path(variant_path) -def _get_caller_locked_kernel(repo_id: str) -> str | None: - for dist in _get_caller_distributions(): - lock_json = dist.read_text("kernels.lock") - if lock_json is None: - continue - locked_sha = _get_locked_kernel(repo_id, lock_json) - if locked_sha is not None: - return locked_sha - return None - - -def _get_locked_kernel(repo_id: str, lock_json: str) -> str | None: - for kernel_lock_json in json.loads(lock_json): - kernel_lock = KernelLock.from_json(kernel_lock_json) - if kernel_lock.repo_id == repo_id: - return kernel_lock.sha - return None - - -def _get_caller_distributions() -> list[Distribution]: - module = _get_caller_module() - if module is None: - return [] - - # Look up all possible distributions that this module could be from. - package = module.__name__.split(".")[0] - dist_names = importlib.metadata.packages_distributions().get(package) - if dist_names is None: - return [] - - return [importlib.metadata.distribution(dist_name) for dist_name in dist_names] - - -def _get_caller_module() -> ModuleType | None: - stack = inspect.stack() - # Get first module in the stack that is not the current module. - first_module = inspect.getmodule(stack[0][0]) - for frame in stack[1:]: - module = inspect.getmodule(frame[0]) - if module is not None and module != first_module: - return module - return first_module - - -def validate_kernel(*, repo_path: Path, variant: str, hash: str): - """Validate the given build variant of a kernel against a hash.""" - variant_path = repo_path / "build" / variant - - # Get the file paths. The first element is a byte-encoded relative path - # used for sorting. The second element is the absolute path. - files: list[tuple[bytes, Path]] = [] - # Ideally we'd use Path.walk, but it's only available in Python 3.12. - for dirpath, _, filenames in os.walk(variant_path): - for filename in filenames: - file_abs = Path(dirpath) / filename - - # Python likes to create files when importing modules from the - # cache, only hash files that are symlinked blobs. - if file_abs.is_symlink(): - files.append( - ( - file_abs.relative_to(variant_path).as_posix().encode("utf-8"), - file_abs, - ) - ) - - m = hashlib.sha256() - - for filename_bytes, full_path in sorted(files): - m.update(filename_bytes) - - blob_filename = full_path.resolve().name - if len(blob_filename) == 40: - # SHA-1 hashed, so a Git blob. - m.update(git_hash_object(full_path.read_bytes())) - elif len(blob_filename) == 64: - # SHA-256 hashed, so a Git LFS blob. - m.update(hashlib.sha256(full_path.read_bytes()).digest()) - else: - raise ValueError(f"Unexpected blob filename length: {len(blob_filename)}") - - computedHash = f"sha256-{m.hexdigest()}" - if computedHash != hash: - raise ValueError( - f"Lock file specifies kernel with hash {hash}, but downloaded kernel has hash: {computedHash}" - ) - - -def git_hash_object(data: bytes, object_type: str = "blob"): - """Calculate git SHA1 of data.""" - header = f"{object_type} {len(data)}\0".encode() - m = hashlib.sha1() - m.update(header) - m.update(data) - return m.digest() - - def _platform() -> str: cpu = platform.machine() os = platform.system().lower()