diff --git a/kernels/src/kernels/cli/__init__.py b/kernels/src/kernels/cli/__init__.py index 55811d8b5..131c9b4e9 100644 --- a/kernels/src/kernels/cli/__init__.py +++ b/kernels/src/kernels/cli/__init__.py @@ -1,18 +1,12 @@ import argparse -import dataclasses -import json import sys from pathlib import Path +from kernels.cli.download import download_kernels from kernels.cli.info import print_kernel_info +from kernels.cli.lock import lock_kernels from kernels.cli.verify_signature import verify_signature from kernels.cli.versions import print_kernel_versions -from kernels.compat import tomllib -from kernels.install import ( - install_kernel, - install_kernel_all_variants, -) -from kernels.locking import KernelLock, get_kernel_locks def main(): @@ -131,43 +125,6 @@ def main(): args.func(args) -def download_kernels(args): - lock_path = args.project_dir / "kernels.lock" - - if not lock_path.exists(): - print(f"No kernels.lock file found in: {args.project_dir}", file=sys.stderr) - sys.exit(1) - - with open(args.project_dir / "kernels.lock", "r") as f: - lock_json = json.load(f) - - all_successful = True - - for kernel_lock_json in lock_json: - kernel_lock = KernelLock.from_json(kernel_lock_json) - print( - f"Downloading `{kernel_lock.repo_id}` with SHA: {kernel_lock.sha}", - file=sys.stderr, - ) - if args.all_variants: - install_kernel_all_variants( - kernel_lock.repo_id, - revision=kernel_lock.sha, - ) - else: - try: - install_kernel( - kernel_lock.repo_id, - revision=kernel_lock.sha, - ) - except FileNotFoundError as e: - print(e, file=sys.stderr) - all_successful = False - - if not all_successful: - sys.exit(1) - - def kernel_info(args): print_kernel_info( args.kernel, @@ -181,27 +138,6 @@ def kernel_versions(args): print_kernel_versions(args.repo_id) -def lock_kernels(args): - with open(args.project_dir / "pyproject.toml", "rb") as f: - data = tomllib.load(f) - - kernel_versions = data.get("tool", {}).get("kernels", {}).get("dependencies", None) - - all_locks = [] - for kernel, version in kernel_versions.items(): - all_locks.append(get_kernel_locks(kernel, version)) - - with open(args.project_dir / "kernels.lock", "w") as f: - json.dump(all_locks, f, cls=_JSONEncoder, indent=2) - - -class _JSONEncoder(json.JSONEncoder): - def default(self, o): - if dataclasses.is_dataclass(o): - return dataclasses.asdict(o) - return super().default(o) - - def _check_moved(_args): print( "`kernels check` has moved to `kernel-builder check-abi`", diff --git a/kernels/src/kernels/cli/download.py b/kernels/src/kernels/cli/download.py new file mode 100644 index 000000000..7860a6eee --- /dev/null +++ b/kernels/src/kernels/cli/download.py @@ -0,0 +1,46 @@ +import sys + +from kernels_data import KernelLocks + +from kernels.hf_hub import CACHE_DIR, _get_hf_api +from kernels.resolver import _BYTECODE_IGNORE_PATTERNS, resolve_hub_kernel + + +def download_kernels(args): + lock_path = args.project_dir / "kernels.lock" + + if not lock_path.exists(): + print(f"No kernels.lock file found in: {args.project_dir}", file=sys.stderr) + sys.exit(1) + + with open(args.project_dir / "kernels.lock", "r") as f: + kernel_locks = KernelLocks.from_json(f.read()) + + all_successful = True + + api = _get_hf_api() + for dep, lock in kernel_locks.items(): + print( + f"Downloading kernel: {dep.repo_id} (revision: {lock.commit})", + file=sys.stderr, + ) + + try: + if args.all_variants: + api.snapshot_download( + dep.repo_id, + repo_type="kernel", + allow_patterns="build/*", + ignore_patterns=_BYTECODE_IGNORE_PATTERNS, + cache_dir=CACHE_DIR, + revision=lock.commit, + ) + else: + location = resolve_hub_kernel(dep.repo_id, api=api, backend=None, revision=lock.commit) + location.install(api=api) + except FileNotFoundError as e: + print(e, file=sys.stderr) + all_successful = False + + if not all_successful: + sys.exit(1) diff --git a/kernels/src/kernels/cli/lock.py b/kernels/src/kernels/cli/lock.py new file mode 100644 index 000000000..f7b19ff3e --- /dev/null +++ b/kernels/src/kernels/cli/lock.py @@ -0,0 +1,24 @@ +from kernels_data import KernelDependency, KernelVersion + +from kernels.compat import tomllib +from kernels.hf_hub import _get_hf_api +from kernels.locking import extract_dependency_locks + + +def lock_kernels(args): + project_dir = args.project_dir + + with open(project_dir / "pyproject.toml", "rb") as f: + data = tomllib.load(f) + + kernel_versions = data.get("tool", {}).get("kernels", {}).get("dependencies", None) + + depends = [ + KernelDependency(repo_id=kernel, version=KernelVersion.Version(version)) + for kernel, version in kernel_versions.items() + ] + + locks = extract_dependency_locks(depends, api=_get_hf_api()) + + with open(project_dir / "kernels.lock", "w") as f: + f.write(locks.to_json()) diff --git a/kernels/src/kernels/cli/verify_signature.py b/kernels/src/kernels/cli/verify_signature.py index c96e5af1b..4f2bec8b6 100644 --- a/kernels/src/kernels/cli/verify_signature.py +++ b/kernels/src/kernels/cli/verify_signature.py @@ -16,7 +16,10 @@ def verify_signature(args: argparse.Namespace) -> None: revision = select_revision_or_version( - args.repo_id, revision=None, version=args.version, local_files_only=constants.HF_HUB_OFFLINE + args.repo_id, + revision=None, + version=args.version, + local_files_only=constants.HF_HUB_OFFLINE, ) if args.all_variants: diff --git a/kernels/src/kernels/importer.py b/kernels/src/kernels/importer.py index a62e74bd5..4b39498cb 100644 --- a/kernels/src/kernels/importer.py +++ b/kernels/src/kernels/importer.py @@ -116,7 +116,12 @@ def _import_from_path( if module is None: raise ImportError(f"Cannot load module {module_name} from spec") sys.modules[metadata.id] = module - spec.loader.exec_module(module) # type: ignore + + # Avoid an import cycle. + from kernels.deps import use_kernel_deps + + with use_kernel_deps(deps): + spec.loader.exec_module(module) # type: ignore _loaded_kernels[variant_path] = LoadedKernel( metadata=metadata, diff --git a/kernels/src/kernels/install.py b/kernels/src/kernels/install.py index 16f358e7f..41bafe1ad 100644 --- a/kernels/src/kernels/install.py +++ b/kernels/src/kernels/install.py @@ -1,34 +1,23 @@ from pathlib import Path -from huggingface_hub import HfApi -from huggingface_hub.errors import LocalEntryNotFoundError +from kernels_data import KernelDependency +from kernels._versions import revision_or_version +from kernels.deps import resolve_kernel_tree from kernels.hf_hub import CACHE_DIR, _get_hf_api -from kernels.python_deps import validate_variant_dependencies -from kernels.status import resolve_status -from kernels.variants import ( - Variant, - get_variants, - get_variants_local, - resolve_variant, - variants_trace_str, -) - -# Exclude pattern for bytecode. These are not included in kernel builds, -# but builds not done using kernel-builder might accidentally include -# bytecode. So these patterns are used to ensure that they are never -# downloaded. -_BYTECODE_IGNORE_PATTERNS = ["*.pyc", "**/__pycache__/**"] +from kernels.locking import extract_dependency_locks +from kernels.resolver import _BYTECODE_IGNORE_PATTERNS, HubCacheResolver, HubResolver def install_kernel( repo_id: str, *, - revision: str, - local_files_only: bool = False, + revision: str | None = None, + version: int | None = None, backend: str | None = None, + local_files_only: bool = False, user_agent: str | dict | None = None, - validate_dependencies: bool = False, + trust_remote_code: bool | list[str] = False, ) -> Path: """ Download a kernel for the current environment to the cache. @@ -39,160 +28,75 @@ def install_kernel( repo_id (`str`): The Hub repository containing the kernel. revision (`str`): - The specific revision (branch, tag, or commit) to download. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only use local files and not download from the Hub. + The specific revision (branch, tag, or commit) to download. Cannot be used together with `version`. + version (`int`, *optional*): + The kernel version to download. Cannot be used together with `revision`. + Either `version` or `revision` must be specified. 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. + local_files_only (`bool`, *optional*, defaults to `False`): + Whether to only use local files and not download from the Hub. user_agent (`Union[str, dict]`, *optional*): The `user_agent` info to pass to `snapshot_download()` for internal telemetry. - validate_dependencies (`bool`, defaults to False): - When set to True, performs dependency validation. Useful for kernels that have - extra Python dependencies. + trust_remote_code (`bool | list[str]`, *optional*, defaults to `False`): + Whether to allow loading kernels from untrusted organisations. When ``False``, + only kernels from trusted organisations are allowed. When ``True``, all + repositories are allowed. A list of strings will be used to verify signing + identities in a future release; for now it emits a warning and falls + back to the default trust check. Returns: `Path`: The path to the variant directory. """ api = _get_hf_api(user_agent=user_agent) - if local_files_only: - # Same local-cache resolution path used by `load_kernel`, which is - # always offline. Sharing the helper avoids the network dependency - # that `get_variants` would otherwise introduce. - variant_path = _resolve_local_variant_path( - api, - repo_id, - revision=revision, - backend=backend, - ) - # For locally downloaded kernels, we run the validation after resolving the path - if validate_dependencies: - validate_variant_dependencies(variant_path) - return variant_path - - repo_id, revision = resolve_status(api, repo_id, revision) - variants = get_variants(api, repo_id=repo_id, revision=revision) - variant, trace = resolve_variant(variants, backend) - - if variant is None: - raise FileNotFoundError( - f"Cannot find a build variant for this system in {repo_id} (revision: {revision}):\n\n{variants_trace_str(trace)}" - ) - # Validate Python dependencies before downloading the variant. - if validate_dependencies: - metadata_path = api.hf_hub_download( - repo_id, - repo_type="kernel", - filename=f"build/{variant.variant_str}/metadata.json", - cache_dir=CACHE_DIR, - revision=revision, - local_files_only=False, - ) - validate_variant_dependencies(Path(metadata_path).parent) - - 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=False, - ) - ) + kernel_version = revision_or_version(revision=revision, version=version) + + resolver = ( + HubCacheResolver(trust_remote_code=trust_remote_code) + if local_files_only + else HubResolver(trust_remote_code=trust_remote_code) ) - try: - return _find_kernel_in_repo_path( - repo_path, - variant=variant, - ) - except FileNotFoundError: - raise FileNotFoundError(f"Cannot install kernel from repo {repo_id} (revision: {revision})") + tree = resolve_kernel_tree( + api=api, + backend=backend, + kernel=KernelDependency(repo_id=repo_id, version=kernel_version), + resolver=resolver, + ) + + return tree.install(api=api).location.variant_path -def _resolve_local_variant_path( - api: HfApi, +def install_kernel_all_variants( repo_id: str, *, - revision: str, - backend: str | None = None, + revision: str | None = None, + version: int | None = None, ) -> Path: - """Resolve a kernel variant path from the local Hugging Face cache only. + kernel_dep = KernelDependency(repo_id=repo_id, version=revision_or_version(revision=revision, version=version)) - Used by `load_kernel` (which always operates on a pre-downloaded, locked - kernel) and by the offline branch of `install_kernel`. - """ - try: - local_repo_path = Path( + # Use locking code path to recursively get dependencies. + api = _get_hf_api() + locks = extract_dependency_locks( + [kernel_dep], + api=api, + ) + + paths = {} + for dep, lock in locks.items(): + paths[dep] = Path( str( api.snapshot_download( - repo_id, + dep.repo_id, repo_type="kernel", + allow_patterns="build/*", ignore_patterns=_BYTECODE_IGNORE_PATTERNS, cache_dir=CACHE_DIR, - revision=revision, - local_files_only=True, + revision=lock.commit, ) ) ) - except LocalEntryNotFoundError as e: - raise FileNotFoundError( - f"Cannot find a local snapshot for {repo_id} (revision: {revision}). " - "When Hugging Face Hub is in offline mode the kernel must already " - "be present in the local cache." - ) from e - - variants = get_variants_local(local_repo_path / "build") - variant, status = resolve_variant(variants, backend) - if variant is None: - raise FileNotFoundError( - f"Cannot find a build variant for this system in {repo_id} (revision: {revision}):\n\n{variants_trace_str(status)}" - ) - - return _find_kernel_in_repo_path( - local_repo_path, - variant=variant, - ) - - -def _find_kernel_in_repo_path( - repo_path: Path, - *, - variant: Variant, -) -> 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}`") - - return variant_path - - -def install_kernel_all_variants( - repo_id: str, - *, - revision: str, -) -> Path: - api = _get_hf_api() - - repo_path = Path( - str( - api.snapshot_download( - repo_id, - repo_type="kernel", - allow_patterns="build/*", - ignore_patterns=_BYTECODE_IGNORE_PATTERNS, - cache_dir=CACHE_DIR, - revision=revision, - ) - ) - ) - return repo_path / "build" + return paths[kernel_dep] / "build" diff --git a/kernels/src/kernels/layer/func.py b/kernels/src/kernels/layer/func.py index 06b1a7e93..597aaca3d 100644 --- a/kernels/src/kernels/layer/func.py +++ b/kernels/src/kernels/layer/func.py @@ -5,10 +5,15 @@ from typing import TYPE_CHECKING, Protocol, Type from huggingface_hub import constants +from kernels_data import KernelDependency, KernelLocks + +from kernels.resolver import LockedHubCacheResolver, LockedHubResolver from .._versions import select_revision_or_version +from ..hf_hub import _get_hf_api from ..load import ( get_kernel, + get_kernel_with_resolver, get_local_kernel, ) from ..locking import ( @@ -299,25 +304,33 @@ def __init__( self._lockfile = lockfile self.func_name = func_name self._trust_remote_code = trust_remote_code - self._revision = self._resolve_revision() + kernel_locks, kernel_dep = self._get_lock() + self.kernel_locks = kernel_locks + self.kernel_dep = kernel_dep - def _resolve_revision(self) -> str: + def _get_lock(self) -> tuple[KernelLocks, KernelDependency]: if self._lockfile is None: - locked_sha = get_caller_locked_kernel_revision(self._repo_id) + return get_caller_locked_kernel_revision(self._repo_id) else: - with open(self._lockfile, "r") as f: - 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") - - return locked_sha + return get_locked_kernel_revision(self._repo_id, self._lockfile) def load(self) -> Type["nn.Module"]: - kernel = get_kernel( - repo_id=self._repo_id, - revision=self._revision, - trust_remote_code=self._trust_remote_code, + resolver = ( + LockedHubCacheResolver( + kernel_locks=self.kernel_locks, + trust_remote_code=self._trust_remote_code, + ) + if constants.HF_HUB_OFFLINE + else LockedHubResolver( + kernel_locks=self.kernel_locks, + trust_remote_code=self._trust_remote_code, + ) + ) + kernel = get_kernel_with_resolver( + api=_get_hf_api(), + backend=None, + kernel=self.kernel_dep, + resolver=resolver, ) return _get_kernel_func(self, kernel) @@ -326,15 +339,25 @@ def __eq__(self, other): isinstance(other, LockedFuncRepository) and self.func_name == other.func_name and self._repo_id == other._repo_id - and self._revision == other._revision + and self.kernel_dep == other.kernel_dep + and self.kernel_locks == other.kernel_locks and self._trust_remote_code == other._trust_remote_code ) def __hash__(self): - return hash((self.func_name, self._repo_id, self._revision, self._trust_remote_code)) + return hash( + ( + self.func_name, + self._repo_id, + self.kernel_dep, + self.kernel_locks, + self._trust_remote_code, + ) + ) def __str__(self) -> str: - return f"`{self._repo_id}` (revision: {self._revision}), function `{self.func_name}`" + commit = self.kernel_locks[self.kernel_dep].commit + return f"`{self._repo_id}` (revision: {commit}), function `{self.func_name}`" def _get_kernel_func(repo: FuncRepositoryProtocol, kernel: ModuleType) -> Type["nn.Module"]: diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py index 2b2b71eb1..2e7aae205 100644 --- a/kernels/src/kernels/layer/layer.py +++ b/kernels/src/kernels/layer/layer.py @@ -10,10 +10,15 @@ from typing import TYPE_CHECKING, Callable, Protocol, Type from huggingface_hub import constants +from kernels_data import KernelDependency, KernelLocks + +from kernels.resolver import LockedHubCacheResolver, LockedHubResolver from .._versions import select_revision_or_version +from ..hf_hub import _get_hf_api from ..load import ( get_kernel, + get_kernel_with_resolver, get_local_kernel, ) from ..locking import ( @@ -208,25 +213,33 @@ def __init__( self._lockfile = lockfile self.layer_name = layer_name self._trust_remote_code = trust_remote_code - self._revision = self._resolve_revision() + kernel_locks, kernel_dep = self._get_lock() + self.kernel_locks = kernel_locks + self.kernel_dep = kernel_dep - def _resolve_revision(self) -> str: + def _get_lock(self) -> tuple[KernelLocks, KernelDependency]: if self._lockfile is None: - locked_sha = get_caller_locked_kernel_revision(self._repo_id) + return get_caller_locked_kernel_revision(self._repo_id) else: - with open(self._lockfile, "r") as f: - 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") - - return locked_sha + return get_locked_kernel_revision(self._repo_id, self._lockfile) def load(self) -> Type["nn.Module"]: - kernel = get_kernel( - repo_id=self._repo_id, - revision=self._revision, - trust_remote_code=self._trust_remote_code, + resolver = ( + LockedHubCacheResolver( + kernel_locks=self.kernel_locks, + trust_remote_code=self._trust_remote_code, + ) + if constants.HF_HUB_OFFLINE + else LockedHubResolver( + kernel_locks=self.kernel_locks, + trust_remote_code=self._trust_remote_code, + ) + ) + kernel = get_kernel_with_resolver( + api=_get_hf_api(), + backend=None, + kernel=self.kernel_dep, + resolver=resolver, ) return _get_kernel_layer(self, kernel) @@ -235,15 +248,25 @@ def __eq__(self, other): isinstance(other, LockedLayerRepository) and self.layer_name == other.layer_name and self._repo_id == other._repo_id - and self._revision == other._revision + and self.kernel_dep == other.kernel_dep + and self.kernel_locks == other.kernel_locks and self._trust_remote_code == other._trust_remote_code ) def __hash__(self): - return hash((self.layer_name, self._repo_id, self._revision, self._trust_remote_code)) + return hash( + ( + self.layer_name, + self._repo_id, + self.kernel_dep, + self.kernel_locks, + self._trust_remote_code, + ) + ) def __str__(self) -> str: - return f"`{self._repo_id}` (revision: {self._revision}), layer `{self.layer_name}`" + commit = self.kernel_locks[self.kernel_dep].commit + return f"`{self._repo_id}` (revision: {commit}), layer `{self.layer_name}`)" _CACHED_LAYER: dict[RepositoryProtocol, Type["nn.Module"]] = {} diff --git a/kernels/src/kernels/load.py b/kernels/src/kernels/load.py index 28a13df92..bddb0c01d 100644 --- a/kernels/src/kernels/load.py +++ b/kernels/src/kernels/load.py @@ -3,30 +3,35 @@ from pathlib import Path from types import ModuleType -from huggingface_hub import constants +from huggingface_hub import HfApi, constants +from kernels_data import KernelDependency, KernelVersion -from kernels._versions import select_revision_or_version -from kernels.hf_hub import RepoInfo, _check_trust_remote_code, _get_hf_api -from kernels.importer import _import_from_path -from kernels.install import install_kernel +from kernels._versions import revision_or_version +from kernels.deps import resolve_kernel_tree +from kernels.hf_hub import _get_hf_api from kernels.locking import ( get_caller_locked_kernel_revision, get_locked_kernel_revision, ) -from kernels.python_deps import validate_variant_dependencies -from kernels.variants import ( - get_variants, - get_variants_local, - resolve_variant, +from kernels.resolver import ( + HubCacheResolver, + HubResolver, + LockedHubCacheResolver, + LockedHubResolver, + NoopResolver, + RepoPathsResolver, + Resolver, + SequentialResolver, ) -def _get_local_kernel_overrides() -> dict[str, Path]: +def _get_local_kernel_overrides() -> Resolver: """Returns list local overrides for kernels.""" local_kerels = os.environ.get("LOCAL_KERNELS", None) if local_kerels is None: - return dict() - return _parse_local_kernel_overrides(local_kerels) + return NoopResolver() + overrides = _parse_local_kernel_overrides(local_kerels) + return RepoPathsResolver(local_kernels=overrides) @functools.lru_cache(maxsize=1) @@ -44,8 +49,44 @@ def _parse_local_kernel_overrides(local_kernels: str) -> dict[str, Path]: return overrides +def get_kernel_with_resolver( + *, + api: HfApi, + backend: str | None, + kernel: KernelDependency, + resolver: Resolver | None, +) -> ModuleType: + """ + Load a kernel and its (transitive) dependencies using the given resolver. + + Args: + api (`HfApi`): + The Hugging Face Hub API client. + 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. + kernel (`KernelDependency`): + The kernel to load. + resolver (`Resolver`, *optional*): + The resolver used to resolve the kernel and its (transitive) dependencies. + + Returns: + `ModuleType`: The imported kernel module. + """ + tree = resolve_kernel_tree( + api=api, + backend=backend, + kernel=kernel, + resolver=resolver, + ) + tree.validate_dependencies() + tree_only_local = tree.install(api=api) + return tree_only_local.load() + + def get_kernel( repo_id: str, + *, revision: str | None = None, version: int | None = None, backend: str | None = None, @@ -92,81 +133,92 @@ def get_kernel( result = activation.relu(out, x) ``` """ - override = _get_local_kernel_overrides().get(repo_id, None) - if override is not None: - return get_local_kernel(override) - - _check_trust_remote_code( - repo_id=repo_id, - local_files_only=constants.HF_HUB_OFFLINE, - trust_remote_code=trust_remote_code, - ) + api = _get_hf_api(user_agent=user_agent) - revision = select_revision_or_version( - repo_id, - revision=revision, - version=version, - local_files_only=constants.HF_HUB_OFFLINE, - ) - repo_info = RepoInfo( - repo_id=repo_id, - revision=revision, - ) - variant_path = install_kernel( - repo_id, + kernel_version = revision_or_version(revision=revision, version=version) + + resolvers = [ + _get_local_kernel_overrides(), + ( + HubCacheResolver(trust_remote_code=trust_remote_code) + if constants.HF_HUB_OFFLINE + else HubResolver(trust_remote_code=trust_remote_code) + ), + ] + + return get_kernel_with_resolver( + api=api, backend=backend, - revision=revision, - user_agent=user_agent, - validate_dependencies=True, - local_files_only=constants.HF_HUB_OFFLINE, + kernel=KernelDependency(repo_id=repo_id, version=kernel_version), + resolver=SequentialResolver(resolvers=resolvers), ) - return _import_from_path(variant_path, deps={}, repo_info=repo_info) def get_local_kernel( repo_path: Path, + *, backend: str | None = None, + trust_remote_code: bool | list[str] = False, + user_agent: str | dict | None = None, ) -> ModuleType: """ Import a kernel from a local kernel repository path. + If the kernel has any (transitive) dependencies, they will be downloaded. + Args: repo_path (`Path`): The local path to the kernel repository. 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. + user_agent (`Union[str, dict]`, *optional*): + The `user_agent` info to pass to `snapshot_download()` for internal telemetry. + trust_remote_code (`bool | list[str]`, *optional*, defaults to `False`): + Whether to allow loading kernels from untrusted organisations. When ``False``, + only kernels from trusted organisations are allowed. When ``True``, all + repositories are allowed. A list of strings will be used to verify signing + identities in a future release; for now it emits a warning and falls + back to the default trust check. + + Returns: `ModuleType`: The imported kernel module. """ - for base_path in [repo_path, repo_path / "build"]: - variants = get_variants_local(base_path) - variant, _ = resolve_variant(variants, backend) - - if variant is not None: - variant_path = base_path / variant.variant_str - validate_variant_dependencies(variant_path) - return _import_from_path(variant_path, deps={}) - - # If we didn't find the package in the repo we may have a explicit - # package path. - variant_path = repo_path - if variant_path.exists(): - validate_variant_dependencies(variant_path) - return _import_from_path(variant_path, deps={}) - - raise FileNotFoundError(f"Could not find kernel in {repo_path}") + api = _get_hf_api(user_agent=user_agent) + + resolvers = [ + RepoPathsResolver(local_kernels={str(repo_path): repo_path}), + _get_local_kernel_overrides(), + ( + HubCacheResolver(trust_remote_code=trust_remote_code) + if constants.HF_HUB_OFFLINE + else HubResolver(trust_remote_code=trust_remote_code) + ), + ] + + return get_kernel_with_resolver( + api=api, + backend=backend, + # We don't have a name for the kernel, so let's just use the path. + kernel=KernelDependency(repo_id=str(repo_path), version=KernelVersion.Version(0)), + resolver=SequentialResolver(resolvers), + ) def has_kernel( repo_id: str, + *, revision: str | None = None, version: int | None = None, backend: str | None = None, + trust_remote_code: bool | list[str] = False, ) -> bool: """ - Check whether a kernel build exists for the current environment (Torch version and compute framework). + Check whether a kernel build exists for the current environment (framework version and backend). + + If the kernel has any (transitive) dependencies, they will be checked as well. Args: repo_id (`str`): @@ -179,30 +231,40 @@ def has_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. + trust_remote_code (`bool | list[str]`, *optional*, defaults to `False`): + Whether to allow loading kernels from untrusted organisations. When ``False``, + only kernels from trusted organisations are allowed. When ``True``, all + repositories are allowed. A list of strings will be used to verify signing + identities in a future release; for now it emits a warning and falls + back to the default trust check. Returns: `bool`: `True` if a kernel is available for the current environment. """ - revision = select_revision_or_version( - repo_id, - revision=revision, - version=version, - local_files_only=constants.HF_HUB_OFFLINE, - ) - api = _get_hf_api() - variants = get_variants(api, repo_id=repo_id, revision=revision) - variant, _ = resolve_variant(variants, backend) - if variant is None: + kernel_version = revision_or_version(revision=revision, version=version) + + resolvers = [ + _get_local_kernel_overrides(), + ( + HubCacheResolver(trust_remote_code=trust_remote_code) + if constants.HF_HUB_OFFLINE + else HubResolver(trust_remote_code=trust_remote_code) + ), + ] + + try: + resolve_kernel_tree( + api=api, + backend=backend, + kernel=KernelDependency(repo_id=repo_id, version=kernel_version), + resolver=SequentialResolver(resolvers=resolvers), + ) + except FileNotFoundError: return False - return api.file_exists( - repo_id, - repo_type="kernel", - revision=revision, - filename=f"build/{variant.variant_str}/metadata.json", - ) + return True def load_kernel( @@ -214,6 +276,9 @@ def load_kernel( """ Get a pre-downloaded, locked kernel. + This function will never download anything and will fail when a kernel is + not available locally. + If `lockfile` is not specified, the lockfile will be loaded from the caller's package metadata. Args: @@ -229,50 +294,57 @@ def load_kernel( `ModuleType`: The imported kernel module. """ if lockfile is None: - locked_sha = get_caller_locked_kernel_revision(repo_id) + kernel_locks, kernel_dep = get_caller_locked_kernel_revision(repo_id) else: - with open(lockfile, "r") as f: - locked_sha = get_locked_kernel_revision(repo_id, f.read()) + kernel_locks, kernel_dep = get_locked_kernel_revision(repo_id, lockfile) - if locked_sha is None: - raise ValueError( - f"Kernel `{repo_id}` is not locked. Please lock it with `kernels lock ` and then reinstall the project." - ) + resolver = LockedHubCacheResolver(kernel_locks=kernel_locks, trust_remote_code=False) - try: - variant_path = install_kernel(repo_id, revision=locked_sha, backend=backend, local_files_only=True) - except FileNotFoundError as e: - raise FileNotFoundError( - f"Locked kernel `{repo_id}` was not downloaded or does not have an " - "applicable variant. Make sure it's downloaded locally via " - "`kernels download `." - ) from e - return _import_from_path(variant_path, deps={}) + return get_kernel_with_resolver( + # We need to provide the API, but it is never used, so no need for a + # user agent. + api=_get_hf_api(), + backend=backend, + kernel=kernel_dep, + resolver=resolver, + ) -def get_locked_kernel(repo_id: str) -> ModuleType: +def get_locked_kernel( + repo_id: str, + *, + lockfile: Path | None, + trust_remote_code: bool | list[str] = False, + user_agent: str | dict | None = None, +) -> ModuleType: """ Get a kernel using a lock file. Args: repo_id (`str`): The Hub repository containing the kernel. - local_files_only (`bool`, *optional*, defaults to `False`): - Whether to only use local files and not download from the Hub. + lockfile (`Path`, *optional*): + Path to the lockfile. If not provided, the lockfile will be loaded from the caller's package metadata. + user_agent (`Union[str, dict]`, *optional*): + The `user_agent` info to pass to `snapshot_download()` for internal telemetry. Returns: `ModuleType`: The imported kernel module. """ - locked_sha = get_caller_locked_kernel_revision(repo_id) - - if locked_sha is None: - raise ValueError(f"Kernel `{repo_id}` is not locked") + if lockfile is None: + kernel_locks, kernel_dep = get_caller_locked_kernel_revision(repo_id) + else: + kernel_locks, kernel_dep = get_locked_kernel_revision(repo_id, lockfile) - variant_path = install_kernel( - repo_id, - revision=locked_sha, - local_files_only=constants.HF_HUB_OFFLINE, - validate_dependencies=True, + resolver = ( + LockedHubCacheResolver(kernel_locks=kernel_locks, trust_remote_code=trust_remote_code) + if constants.HF_HUB_OFFLINE + else LockedHubResolver(kernel_locks=kernel_locks, trust_remote_code=trust_remote_code) ) - return _import_from_path(variant_path, deps={}) + return get_kernel_with_resolver( + api=_get_hf_api(user_agent=user_agent), + backend=None, + kernel=kernel_dep, + resolver=resolver, + ) diff --git a/kernels/src/kernels/locking.py b/kernels/src/kernels/locking.py index 844611f3e..0a393c5b9 100644 --- a/kernels/src/kernels/locking.py +++ b/kernels/src/kernels/locking.py @@ -1,113 +1,101 @@ -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 import constants -from huggingface_hub.dataclasses import strict -from huggingface_hub.hf_api import RepoFile +from huggingface_hub.hf_api import HfApi +from kernels_data import ( + KernelDependency, + KernelLock, + KernelLocks, + Metadata, +) -from kernels._versions import resolve_version_spec_as_ref +from kernels._versions import resolve_kernel_version from kernels.compat import tomllib -from kernels.status import resolve_status - - -@strict -@dataclass -class VariantLock: - hash: str - hash_type: str = "git_lfs_concat" - - -@strict -@dataclass -class KernelLock: - repo_id: str - sha: str - variants: dict[str, VariantLock] - - @classmethod - def from_json(cls, o: dict): - variants = {variant: VariantLock(**lock) for variant, lock in o["variants"].items()} - return cls(repo_id=o["repo_id"], sha=o["sha"], variants=variants) - - -def get_kernel_locks(repo_id: str, version_spec: int) -> KernelLock: - """ - Get the locks for a kernel with the given version. - """ - from kernels.hf_hub import _get_hf_api - - api = _get_hf_api() - - # NOTE: the destination of a redirect is respected but we still use - # resolve_version_spec_as_ref to resolve the version specifier of the - # final destination repo. - repo_id, _ = resolve_status(api, repo_id, "main") - - tag_for_newest = resolve_version_spec_as_ref(repo_id, version_spec, local_files_only=constants.HF_HUB_OFFLINE) - - revision = tag_for_newest.target_commit - - r = api.repo_info( - repo_id=repo_id, - repo_type="kernel", - revision=revision, +from kernels.hf_hub import CACHE_DIR, _check_trust_remote_code +from kernels.variants import get_variants + + +def lock_kernel_tree( + kernel: KernelDependency, + *, + api: HfApi, + seen: set[KernelDependency] | None = None, + kernel_locks: dict[KernelDependency, KernelLock] | None = None, +) -> dict[KernelDependency, KernelLock]: + """Lock all kernels in the dependency tree""" + + # NOTE: this function does not use resolve_kernel_tree, since that function + # constructs a tree using the best-matching variant for a system. In + # this function we need to explore all variants, since we want to lock + # every possible variant (and its dependencies). + + if kernel_locks is None: + kernel_locks = dict() + + # If the kernel is already in the locks, we have already resolved + # the kernel and its dependencies. + if kernel in kernel_locks: + return kernel_locks + + if seen is None: + seen = set() + + # Check for cycles. + if kernel in seen: + raise ValueError(f"Cyclic kernel dependency detected: {kernel.repo_id}") + seen.add(kernel) + + # Check if the repo is trusted before downloading anything. + _check_trust_remote_code( + repo_id=kernel.repo_id, + local_files_only=False, + trust_remote_code=False, ) - if r.sha is None: - raise ValueError(f"Cannot get commit SHA for repo {repo_id} for tag {tag_for_newest.name}") - - siblings = [ - f - for f in api.list_repo_tree( - repo_id=repo_id, - repo_type="kernel", - revision=revision, - recursive=True, - ) - if isinstance(f, RepoFile) - ] - variant_files: dict[str, list[tuple[bytes, str]]] = {} - for sibling in siblings: - if sibling.rfilename.startswith("build/torch"): - if sibling.blob_id is None: - raise ValueError(f"Cannot get blob ID for {sibling.rfilename}") + revision = resolve_kernel_version(kernel, local_files_only=False) + + for variant in get_variants(api, repo_id=kernel.repo_id, revision=revision): + metadata_path = Path( + api.hf_hub_download( + kernel.repo_id, + repo_type="kernel", + filename=f"build/{variant.variant_str}/metadata.json", + cache_dir=CACHE_DIR, + revision=revision, + local_files_only=False, + ) + ) + metadata = Metadata.read_from_file(metadata_path) - # Exclude Python bytecode. If bytecode exists, it is generated - # by the interpreter, since we exclude bytecode from builds - # and downloads. - if sibling.rfilename.endswith(".pyc") or "__pycache__" in sibling.rfilename.split("/"): - continue + kernel_deps = {} + for dep in metadata.kernel_depends: + kernel_deps[dep] = lock_kernel_tree(dep, api=api, seen=seen, kernel_locks=kernel_locks) - path = Path(sibling.rfilename) - variant = path.parts[1] - filename = Path(*path.parts[2:]) + seen.remove(kernel) - hash = sibling.lfs.sha256 if sibling.lfs is not None else sibling.blob_id + kernel_locks[kernel] = KernelLock( + commit=revision, + ) - files = variant_files.setdefault(variant, []) + return kernel_locks - # Encode as posix for consistent slash handling, then encode - # as utf-8 for byte-wise sorting later. - files.append((filename.as_posix().encode("utf-8"), hash)) - variant_locks = {} - for variant, files in variant_files.items(): - m = hashlib.sha256() - for filename_bytes, hash in sorted(files): - # Filename as bytes. - m.update(filename_bytes) - # Git blob or LFS file hash as bytes. - m.update(bytes.fromhex(hash)) +def extract_dependency_locks( + kernels: list[KernelDependency], + *, + api: HfApi, +) -> KernelLocks: + kernel_locks = None + for kernel in kernels: + kernel_locks = lock_kernel_tree(kernel, api=api, kernel_locks=kernel_locks) - variant_locks[variant] = VariantLock(hash=f"sha256-{m.hexdigest()}") + if kernel_locks is None: + raise ValueError("No kernels found to lock.") - return KernelLock(repo_id=repo_id, sha=r.sha, variants=variant_locks) + return KernelLocks(locks=kernel_locks) def write_egg_lockfile(cmd, basename, filename): @@ -137,23 +125,40 @@ def write_egg_lockfile(cmd, basename, filename): 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_locked_kernel_revision(repo_id: str, lock_json: str) -> tuple[KernelLocks, KernelDependency]: + kernel_locks = KernelLocks.from_json(lock_json) + + # Lock files are keyed `KernelDependency`, but for project-locked + # kenels we only have one version at the top level, so we have to + # do a linear search. + kernel_dep = next((dep for dep in kernel_locks.keys() if dep.repo_id == repo_id), None) + + if kernel_dep is None: + raise ValueError( + f"Kernel `{repo_id}` is not locked. Please lock it with `kernels lock ` and then reinstall the project." + ) + + return kernel_locks, kernel_dep + + +def get_locked_kernel_revision(repo_id: str, lockfile: Path) -> tuple[KernelLocks, KernelDependency]: + with open(lockfile, "r") as f: + return _get_locked_kernel_revision(repo_id, f.read()) -def get_caller_locked_kernel_revision(repo_id: str) -> str | None: +def get_caller_locked_kernel_revision( + repo_id: str, +) -> tuple[KernelLocks, KernelDependency]: 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 + + return _get_locked_kernel_revision(repo_id, lock_json) + + raise ValueError( + "Could not find a `kernels.lock` file in the caller's package metadata. Please lock kernels with `kernels lock ` and then reinstall the project." + ) def _get_caller_distributions() -> list[Distribution]: diff --git a/kernels/tests/kernel_locking/kernels.lock b/kernels/tests/kernel_locking/kernels.lock index dfad45099..b6a69d1cb 100644 --- a/kernels/tests/kernel_locking/kernels.lock +++ b/kernels/tests/kernel_locking/kernels.lock @@ -1,238 +1,38 @@ [ { - "repo_id": "kernels-community/relu", - "sha": "d649efb56fb249ac8f7a57fa1866728ad0c60e52", - "variants": { - "torch210-cpu-aarch64-darwin": { - "hash": "sha256-d39147b57a44ceda745fb101e028e87030c33b8375a049b75ab9ec5dd606fa64", - "hash_type": "git_lfs_concat" - }, - "torch210-cxx11-cpu-aarch64-linux": { - "hash": "sha256-bceaead15cecf8727794a6d991d33eb29786993d72d08b1e91fd504028f9a539", - "hash_type": "git_lfs_concat" - }, - "torch210-cxx11-cpu-x86_64-linux": { - "hash": "sha256-0b762a6b76aaf4f198e966839f8c6cb0969710aef96d8f86796042c78eb994eb", - "hash_type": "git_lfs_concat" - }, - "torch210-cxx11-cu126-aarch64-linux": { - "hash": "sha256-1859c8492abb2156413cf674eed7133c3e210ed4c0e82ad8bdfe41175a5fddaa", - "hash_type": "git_lfs_concat" - }, - "torch210-cxx11-cu126-x86_64-linux": { - "hash": "sha256-80a5d86f6adc0a9ebeaacfe20bcc639891c8e187445e7c2798099af059564fc4", - "hash_type": "git_lfs_concat" - }, - "torch210-cxx11-cu128-aarch64-linux": { - "hash": "sha256-1dfffa162dae607736b0035df9f6f7dd5178229d6bf41b88dac6d6c2ec6edd56", - "hash_type": "git_lfs_concat" - }, - "torch210-cxx11-cu128-x86_64-linux": { - "hash": "sha256-6608f0269bdf34745e8614383e92e03242fe5b7b1ecc2c1c0072aaa1c793faa9", - "hash_type": "git_lfs_concat" - }, - "torch210-cxx11-cu130-aarch64-linux": { - "hash": "sha256-d068a7ef00ea383f3462e6bf964065629edf43ab7ee9f5682c65df1aa4209569", - "hash_type": "git_lfs_concat" - }, - "torch210-cxx11-cu130-x86_64-linux": { - "hash": "sha256-213c98984f1200cfca87046bb8bce99ef1dd30567a0ee50851bd877a159cfa12", - "hash_type": "git_lfs_concat" - }, - "torch210-cxx11-rocm70-x86_64-linux": { - "hash": "sha256-1f1278182cb6849851f4088ed3d8f67300797dd29da83ced98fd6dab32c65eb0", - "hash_type": "git_lfs_concat" - }, - "torch210-cxx11-rocm71-x86_64-linux": { - "hash": "sha256-30c9f507a7a677e841a24659c48eb4a11d07c027d31122837ccf115d05f32f4a", - "hash_type": "git_lfs_concat" - }, - "torch210-cxx11-xpu20253-x86_64-linux": { - "hash": "sha256-df61ad8b2ebd21da965f809ed22549f9a61016085d49006808ffbb5d0b6c5284", - "hash_type": "git_lfs_concat" - }, - "torch210-metal-aarch64-darwin": { - "hash": "sha256-8cef707a99f64179ae03f56d7b46b39ee62a34cd45b5f41b91c26bb68d1963f9", - "hash_type": "git_lfs_concat" - }, - "torch211-cpu-aarch64-darwin": { - "hash": "sha256-188febb002663d80080c1e930ac74bdfa370807c77df075e47d82f2bf5c07d14", - "hash_type": "git_lfs_concat" - }, - "torch211-cxx11-cpu-aarch64-linux": { - "hash": "sha256-34b7e51dd3e7d1fd70a227ed107d76dbb9e8e34e12b321cbe8d565a78f6af2fc", - "hash_type": "git_lfs_concat" - }, - "torch211-cxx11-cpu-x86_64-linux": { - "hash": "sha256-eaf133dc0d831ceb7dd338d45e2ddcc8e2d568e3ac66c0e5f6615f8d4a9cf86f", - "hash_type": "git_lfs_concat" - }, - "torch211-cxx11-cu126-aarch64-linux": { - "hash": "sha256-6f061a86ddf6335de4d29b621232463edf5d5f60a8695d99489f72df699d0130", - "hash_type": "git_lfs_concat" - }, - "torch211-cxx11-cu126-x86_64-linux": { - "hash": "sha256-ce24704d15452c41647c747dd79e7ace5dfa61732c14eb5bafe892957c136208", - "hash_type": "git_lfs_concat" - }, - "torch211-cxx11-cu128-aarch64-linux": { - "hash": "sha256-2a89c6544a38d535d8afd8aceb4865f609112d6df49f4ec813a0db257a63a6fc", - "hash_type": "git_lfs_concat" - }, - "torch211-cxx11-cu128-x86_64-linux": { - "hash": "sha256-ebee39f6f391accb4bdf69a557cdb9e2b06c8e8baa3d0e98b940304f050127a4", - "hash_type": "git_lfs_concat" - }, - "torch211-cxx11-cu130-aarch64-linux": { - "hash": "sha256-f6a4195e68dc94251a4c2f0cd0029779bb45f06a806d36aa4efd1bb65bb40507", - "hash_type": "git_lfs_concat" - }, - "torch211-cxx11-cu130-x86_64-linux": { - "hash": "sha256-df54dd5c9c17b29bf27fa3c1aac6f533ff0aa28ca12f2b73fe67b2f1078349e2", - "hash_type": "git_lfs_concat" - }, - "torch211-cxx11-rocm71-x86_64-linux": { - "hash": "sha256-d69995638031b7c0caef408af86fad0707b93252ac26d466f72dee17445108ed", - "hash_type": "git_lfs_concat" - }, - "torch211-cxx11-rocm72-x86_64-linux": { - "hash": "sha256-667ed8cf85057c3ee0130a9b10802ef757ae89b67ece75381c4a3f59e00a2903", - "hash_type": "git_lfs_concat" - }, - "torch211-cxx11-xpu20253-x86_64-linux": { - "hash": "sha256-12bf712ad3b26abe24a793e588758579255b2dd6d643de302ee99d28f65ecd30", - "hash_type": "git_lfs_concat" - }, - "torch211-metal-aarch64-darwin": { - "hash": "sha256-9ae17b8ffef95a4583c5fa88631428e7afa9af51e5ffde1f96851ff66442cd77", - "hash_type": "git_lfs_concat" - }, - "torch212-cpu-aarch64-darwin": { - "hash": "sha256-306977881be7fa4ad72b67e6602242281784d8594935287ecefc3051f8d8fcbb", - "hash_type": "git_lfs_concat" - }, - "torch212-cxx11-cpu-aarch64-linux": { - "hash": "sha256-431a6a1d8dcc05bb22c11b30825b37e9078c348196e0f87087510a78071d3d6d", - "hash_type": "git_lfs_concat" - }, - "torch212-cxx11-cpu-x86_64-linux": { - "hash": "sha256-d7db2d05b924238145901c5e7aa816697b24366a598e5251756c4323925241ca", - "hash_type": "git_lfs_concat" - }, - "torch212-cxx11-cu126-aarch64-linux": { - "hash": "sha256-ab69dba825a3639f1909dbfa02e8bf4d830ce6a3b3234cc69c3ae986dc0e3e2f", - "hash_type": "git_lfs_concat" - }, - "torch212-cxx11-cu126-x86_64-linux": { - "hash": "sha256-edc9e4c111eee46709bfbf15a769cca4b868a96b0db7e9fa77e02115f344f71b", - "hash_type": "git_lfs_concat" - }, - "torch212-cxx11-cu130-aarch64-linux": { - "hash": "sha256-1469f1a6b2131781c7c851d4fabea9d6951921df95cf07ad39ba93674fd81334", - "hash_type": "git_lfs_concat" - }, - "torch212-cxx11-cu130-x86_64-linux": { - "hash": "sha256-c5b1c7e9ab0b5f55503004c5782806e7dadc5ac81eafad52149c24b30c92c012", - "hash_type": "git_lfs_concat" - }, - "torch212-cxx11-cu132-aarch64-linux": { - "hash": "sha256-fc1326612613bb7a4e7e903902aa21e333716269f63bf17e869d95b79c583ef7", - "hash_type": "git_lfs_concat" - }, - "torch212-cxx11-cu132-x86_64-linux": { - "hash": "sha256-912633eacf4605505a0c5573f62f42be74e2e716cf274462d55a51b8250f6cf8", - "hash_type": "git_lfs_concat" - }, - "torch212-cxx11-rocm71-x86_64-linux": { - "hash": "sha256-01a05ff87b2b5267e2cdb2e64e4cab8e2c7710fbb4cde06e70de0d1f9dd27011", - "hash_type": "git_lfs_concat" - }, - "torch212-cxx11-rocm72-x86_64-linux": { - "hash": "sha256-661f414130df56fc7a5a010c136ab06af013eae6d211161e8946c2e0ead886d6", - "hash_type": "git_lfs_concat" - }, - "torch212-cxx11-xpu20253-x86_64-linux": { - "hash": "sha256-9e10388a71a93e40e6e54afa728f78d5c75947fa5cbfa37f491f6a7d90deb2b1", - "hash_type": "git_lfs_concat" - }, - "torch212-metal-aarch64-darwin": { - "hash": "sha256-d913c015cd6d23ca8e787230094c38d839282afc5c2416323445120436cf21f0", - "hash_type": "git_lfs_concat" - }, - "torch213-cpu-aarch64-darwin": { - "hash": "sha256-07236f9d7590f8b2c55fc69f2e395d1840fa3da858969ed0c1d66ee1e317251d", - "hash_type": "git_lfs_concat" - }, - "torch213-cxx11-cpu-aarch64-linux": { - "hash": "sha256-6d723fa34eb7546614947392834061082433e951d5ffbd5c77d6c17828e65009", - "hash_type": "git_lfs_concat" - }, - "torch213-cxx11-cpu-x86_64-linux": { - "hash": "sha256-fdbd9eabb0bfa17482c3040af263e60be2446c1e727b13db45723a5eccc74348", - "hash_type": "git_lfs_concat" - }, - "torch213-cxx11-cu126-aarch64-linux": { - "hash": "sha256-4d83197824a66de23669ed8236e54d51a755042d347a598816d183721112461b", - "hash_type": "git_lfs_concat" - }, - "torch213-cxx11-cu126-x86_64-linux": { - "hash": "sha256-3ec65d53dec68f093e6ae87413f963989803b15761fc87af8378ab4906f77e91", - "hash_type": "git_lfs_concat" - }, - "torch213-cxx11-cu130-aarch64-linux": { - "hash": "sha256-620199006ffd33992f52ec80220ecab244eb385a78cc8266d3c25fdb2b78fac1", - "hash_type": "git_lfs_concat" - }, - "torch213-cxx11-cu130-x86_64-linux": { - "hash": "sha256-bee06cbb5f5d06fe0d11e2690b7c677713d59b9c7c2fbd012b5c4b51bbb8bc20", - "hash_type": "git_lfs_concat" - }, - "torch213-cxx11-cu132-aarch64-linux": { - "hash": "sha256-42de5a7390d98aeabc72d9504978b8647db6c7bd87a62978fb79525d0ee222df", - "hash_type": "git_lfs_concat" - }, - "torch213-cxx11-cu132-x86_64-linux": { - "hash": "sha256-33676e27fe45b1a9d27ca4f8ed49d21d4cd1307836ae827faf6a6880a868c943", - "hash_type": "git_lfs_concat" - }, - "torch213-cxx11-rocm71-x86_64-linux": { - "hash": "sha256-85537035d062b111f7150cff2411a525fab4476fe1e0b731c2d605a374be070f", - "hash_type": "git_lfs_concat" - }, - "torch213-cxx11-rocm72-x86_64-linux": { - "hash": "sha256-d58690fbf71726ef75a61acbff0074acfd660cc44cbf968479b780c5284e60be", - "hash_type": "git_lfs_concat" - }, - "torch213-cxx11-xpu20260-x86_64-linux": { - "hash": "sha256-f32c78628caebc9d7f3ef6754d4381c7fc63c633a2814b79c074065dd892e407", - "hash_type": "git_lfs_concat" - }, - "torch213-metal-aarch64-darwin": { - "hash": "sha256-9c67ea09f557711e33ec5b82a8cb39df6be3892ca2da4f92aceee457dacb9d74", - "hash_type": "git_lfs_concat" - } + "dependency": { + "repo-id": "kernels-community/einops", + "version": 1 + }, + "lock": { + "commit": "0bd00775195dda6f84c16c31f6d94564232d08b3" } }, { - "repo_id": "kernels-test/versions", - "sha": "f609e51b856b3d874b0ae8445913e200f02c1735", - "variants": { - "torch-cpu": { - "hash": "sha256-d9fced49c3beeb47ce5b996c28f760bec579cf444f9fe1040725686b13b1336e", - "hash_type": "git_lfs_concat" - }, - "torch-cuda": { - "hash": "sha256-b6b786e496802bc3d1ae59cb48acb419d99ca1ffbeee0fecbf16238e36c00842", - "hash_type": "git_lfs_concat" - }, - "torch-rocm": { - "hash": "sha256-0ca53208b80507e0b592cbfb7c0c9b5f22a4def6413b809c551dabbea1e2ac86", - "hash_type": "git_lfs_concat" - }, - "torch-xpu": { - "hash": "sha256-09b9d84cb273d71a081daef373c42bec0691620591822e44a6811833f896b810", - "hash_type": "git_lfs_concat" - } + "dependency": { + "repo-id": "kernels-community/relu", + "version": 1 + }, + "lock": { + "commit": "d649efb56fb249ac8f7a57fa1866728ad0c60e52" + } + }, + { + "dependency": { + "repo-id": "kernels-test/kernel-deps", + "version": 1 + }, + "lock": { + "commit": "b1c7a501401316a27653d46107b98259ee79d888" + } + }, + { + "dependency": { + "repo-id": "kernels-test/versions", + "version": 2 + }, + "lock": { + "commit": "f609e51b856b3d874b0ae8445913e200f02c1735" } } ] \ No newline at end of file diff --git a/kernels/tests/kernel_locking/pyproject.toml b/kernels/tests/kernel_locking/pyproject.toml index e50f22e0f..c044ba13b 100644 --- a/kernels/tests/kernel_locking/pyproject.toml +++ b/kernels/tests/kernel_locking/pyproject.toml @@ -1,3 +1,4 @@ [tool.kernels.dependencies] "kernels-community/relu" = 1 "kernels-test/versions" = 2 +"kernels-test/kernel-deps" = 1 diff --git a/kernels/tests/layer_locking/kernels.lock b/kernels/tests/layer_locking/kernels.lock index a102598cf..f87340313 100644 --- a/kernels/tests/layer_locking/kernels.lock +++ b/kernels/tests/layer_locking/kernels.lock @@ -1,24 +1,11 @@ [ { - "repo_id": "kernels-test/versions", - "sha": "f609e51b856b3d874b0ae8445913e200f02c1735", - "variants": { - "torch-cpu": { - "hash": "sha256-d9fced49c3beeb47ce5b996c28f760bec579cf444f9fe1040725686b13b1336e", - "hash_type": "git_lfs_concat" - }, - "torch-cuda": { - "hash": "sha256-b6b786e496802bc3d1ae59cb48acb419d99ca1ffbeee0fecbf16238e36c00842", - "hash_type": "git_lfs_concat" - }, - "torch-rocm": { - "hash": "sha256-0ca53208b80507e0b592cbfb7c0c9b5f22a4def6413b809c551dabbea1e2ac86", - "hash_type": "git_lfs_concat" - }, - "torch-xpu": { - "hash": "sha256-09b9d84cb273d71a081daef373c42bec0691620591822e44a6811833f896b810", - "hash_type": "git_lfs_concat" - } + "dependency": { + "repo-id": "kernels-test/versions", + "version": 2 + }, + "lock": { + "commit": "f609e51b856b3d874b0ae8445913e200f02c1735" } } ] \ No newline at end of file diff --git a/kernels/tests/test_basic.py b/kernels/tests/test_basic.py index fadbedc39..8f72480d9 100644 --- a/kernels/tests/test_basic.py +++ b/kernels/tests/test_basic.py @@ -210,7 +210,10 @@ def test_local_overrides(monkeypatch, local_kernel_path): "LOCAL_KERNELS", "kernels-test/non-existing2=/non/existing:kernels-test/activation=/non/existing", ) - with pytest.raises(FileNotFoundError, match=r"Could not find kernel in /non/existing"): + with pytest.raises( + FileNotFoundError, + match=r"Cannot find a build variant for this system in /non/existing", + ): get_kernel("kernels-test/activation", revision="main") with monkeypatch.context() as m: @@ -274,7 +277,7 @@ def _fail(*_args, **_kwargs): monkeypatch.setattr("huggingface_hub.hf_api.get_session", _fail) # Online path must touch the Hub via get_session and therefore fail. - with pytest.raises(_NoNetwork): + with pytest.raises(ValueError, match=r"could not verify publisher trust status"): install_kernel("kernels-community/relu", revision="v1") # Offline mode resolves entirely from the local cache, so get_session is @@ -301,16 +304,10 @@ def test_install_kernel_offline_uncached_revision(): ) -def test_version_resolution_offline_missing(monkeypatch): +def test_version_resolution_offline_missing(): """resolve_version_spec_as_ref should raise a clear error when offline and no cache.""" - monkeypatch.setattr(constants, "HF_HUB_OFFLINE", True) - with pytest.raises(ValueError, match=r"offline mode"): - resolve_version_spec_as_ref( - "kernels-test/this-repo-should-not-exist", - 1, - local_files_only=constants.HF_HUB_OFFLINE, - ) + resolve_version_spec_as_ref("kernels-test/this-repo-should-not-exist", 1, local_files_only=True) def silu_and_mul_torch(x: torch.Tensor): diff --git a/kernels/tests/test_deps.py b/kernels/tests/test_deps.py index 5b7a7de72..837e5b2a0 100644 --- a/kernels/tests/test_deps.py +++ b/kernels/tests/test_deps.py @@ -81,7 +81,10 @@ def test_resolve_kernel_tree_builds_tree(): "kernel, message", [ (_dep("test/kernel", version=1), "version: 1"), - (KernelDependency(repo_id="test/kernel", version=KernelVersion.Revision("main")), "revision: main"), + ( + KernelDependency(repo_id="test/kernel", version=KernelVersion.Revision("main")), + "revision: main", + ), ], ) @pytest.mark.parametrize("resolver", [None, _MapResolver({})]) diff --git a/kernels/tests/test_kernel_locking.py b/kernels/tests/test_kernel_locking.py index 9dc32ece5..038c60a7b 100644 --- a/kernels/tests/test_kernel_locking.py +++ b/kernels/tests/test_kernel_locking.py @@ -15,6 +15,7 @@ use_kernel_mapping, ) from kernels.cli import download_kernels +from kernels.load import get_locked_kernel # Mock download arguments class. @@ -37,6 +38,14 @@ def test_load_locked(): load_kernel("kernels-community/relu", lockfile=project_dir / "kernels.lock") +@pytest.mark.cuda_only +def test_get_locked_kernel(): + project_dir = Path(__file__).parent / "kernel_locking" + # Also validates that hashing works correctly. + get_locked_kernel("kernels-community/relu", lockfile=project_dir / "kernels.lock") + get_locked_kernel("kernels-test/kernel-deps", lockfile=project_dir / "kernels.lock") + + def test_layer_locked(device): project_dir = Path(__file__).parent / "layer_locking" diff --git a/kernels/tests/test_verify.py b/kernels/tests/test_verify.py index caf819179..b6493fc4f 100644 --- a/kernels/tests/test_verify.py +++ b/kernels/tests/test_verify.py @@ -1,8 +1,12 @@ +from pathlib import Path + from kernels_data import DigestViolation from sigstore.verify import policy from kernels import install_kernel from kernels._versions import select_revision_or_version +from kernels.hf_hub import CACHE_DIR, _get_hf_api +from kernels.resolver import _BYTECODE_IGNORE_PATTERNS from kernels.verify import VerificationResult, verify_variant TEST_POLICIES: list[policy.VerificationPolicy] = [ @@ -43,10 +47,34 @@ def test_invalid_digest_fails(): def test_invalid_metadata_fails(): - variant_path = install_kernel("kernels-test/signatures", revision="invalid-metadata") + # We cannot use regular code paths, because they require valid metadata. + revision = select_revision_or_version( + "kernels-test/signatures", + revision="invalid-metadata", + version=None, + local_files_only=False, + ) + + api = _get_hf_api() + variant_paths = ( + Path( + str( + api.snapshot_download( + "kernels-test/signatures", + repo_type="kernel", + allow_patterns="build/*", + ignore_patterns=_BYTECODE_IGNORE_PATTERNS, + cache_dir=CACHE_DIR, + revision=revision, + ) + ) + ) + / "build" + ) match verify_variant( - variant_path, + # No CUDA dependency, we are only checking metadata. + variant_paths / "torch-cuda", policies=TEST_POLICIES, ): case VerificationResult.MetadataInvalid(reason=reason): @@ -67,10 +95,35 @@ def test_missing_digest_fails(): def test_missing_metadata_fails(): - variant_path = install_kernel("kernels-test/signatures", revision="missing-metadata") + # We cannot use regular code paths, because they require valid metadata. + revision = select_revision_or_version( + "kernels-test/signatures", + revision="missing-metadata", + version=None, + local_files_only=False, + ) + + api = _get_hf_api() + variant_paths = ( + Path( + str( + api.snapshot_download( + "kernels-test/signatures", + repo_type="kernel", + allow_patterns="build/*", + ignore_patterns=_BYTECODE_IGNORE_PATTERNS, + cache_dir=CACHE_DIR, + revision=revision, + ) + ) + ) + / "build" + ) + assert ( verify_variant( - variant_path, + # No CUDA dependency, we are only checking metadata. + variant_paths / "torch-cuda", policies=TEST_POLICIES, ) == VerificationResult.MetadataMissing()