diff --git a/kernels/src/kernels/_versions.py b/kernels/src/kernels/_versions.py index 3bcdee538..9b12c5a57 100644 --- a/kernels/src/kernels/_versions.py +++ b/kernels/src/kernels/_versions.py @@ -5,15 +5,16 @@ from huggingface_hub import constants from huggingface_hub.file_download import repo_folder_name from huggingface_hub.hf_api import GitRefInfo +from kernels_data import KernelDependency, KernelVersion logger = logging.getLogger(__name__) -def _get_available_versions(repo_id: str) -> dict[int, GitRefInfo]: +def _get_available_versions(repo_id: str, *, local_files_only: bool) -> dict[int, GitRefInfo]: """Get kernel versions that are available in the repository.""" from kernels.hf_hub import _get_hf_api - if constants.HF_HUB_OFFLINE: + if local_files_only: return _get_available_versions_from_cache(repo_id) refs = _get_hf_api().list_repo_refs(repo_id=repo_id, repo_type="kernel") @@ -60,15 +61,15 @@ def _get_available_versions_from_cache(repo_id: str) -> dict[int, GitRefInfo]: return versions -def resolve_version_spec_as_ref(repo_id: str, version_spec: int) -> GitRefInfo: +def resolve_version_spec_as_ref(repo_id: str, version_spec: int, local_files_only: bool) -> GitRefInfo: """ Get the ref for a kernel with the given version. """ - versions = _get_available_versions(repo_id) + versions = _get_available_versions(repo_id, local_files_only=local_files_only) ref = versions.get(version_spec, None) if ref is None: - if constants.HF_HUB_OFFLINE and not versions: + if local_files_only and not versions: raise ValueError( f"Version {version_spec} of '{repo_id}' is not available in the local cache " "and Hugging Face Hub is in offline mode. Download the kernel " @@ -90,18 +91,47 @@ def resolve_version_spec_as_ref(repo_id: str, version_spec: int) -> GitRefInfo: return ref +# TODO: maybe we should make this a KernelVersion factory method? +def revision_or_version(*, revision: str | None, version: int | None) -> KernelVersion: + if revision is not None and version is not None: + raise ValueError("Only one of `revision` or `version` must be specified.") + + if revision is not None: + return KernelVersion.Revision(revision) + elif version is not None: + return KernelVersion.Version(version) + + raise ValueError( + "A kernel version or revision must be specified. " + "Use `version=` for a stable kernel API version or `revision=` " + "for an explicit Hub revision. See: https://huggingface.co/docs/kernels/migration" + ) + + +def resolve_kernel_version(kernel: KernelDependency, local_files_only: bool) -> str: + if isinstance(kernel.version, KernelVersion.Version): + return resolve_version_spec_as_ref( + kernel.repo_id, kernel.version.version, local_files_only=local_files_only + ).target_commit + elif isinstance(kernel.version, KernelVersion.Revision): + return kernel.version.revision + else: + raise ValueError(f"Invalid version type: {kernel.version}") + + def select_revision_or_version( repo_id: str, *, revision: str | None, version: int | None, + local_files_only: bool, ) -> str: if revision is not None and version is not None: raise ValueError("Only one of `revision` or `version` must be specified.") elif revision is not None: return revision elif version is not None: - return resolve_version_spec_as_ref(repo_id, version).target_commit + return resolve_version_spec_as_ref(repo_id, version, local_files_only=local_files_only).target_commit else: raise ValueError( "A kernel version or revision must be specified. " diff --git a/kernels/src/kernels/cli/info.py b/kernels/src/kernels/cli/info.py index 36aae73d3..6619ab50c 100644 --- a/kernels/src/kernels/cli/info.py +++ b/kernels/src/kernels/cli/info.py @@ -3,6 +3,7 @@ from pathlib import Path from typing import Any +from huggingface_hub import constants from kernels_data import Metadata from kernels._versions import _get_available_versions, resolve_version_spec_as_ref @@ -26,7 +27,10 @@ def print_kernel_info( path = Path(kernel) if path.is_dir(): if revision is not None or version is not None: - print("--revision and --version cannot be used with a local path", file=sys.stderr) + print( + "--revision and --version cannot be used with a local path", + file=sys.stderr, + ) sys.exit(1) info = _local_kernel_info(path) else: @@ -49,15 +53,18 @@ def _hub_kernel_info( sys.exit(1) if version is not None: - revision = resolve_version_spec_as_ref(repo_id, version).name + revision = resolve_version_spec_as_ref(repo_id, version, local_files_only=constants.HF_HUB_OFFLINE).name elif revision is None: - versions = _get_available_versions(repo_id) + versions = _get_available_versions(repo_id, local_files_only=constants.HF_HUB_OFFLINE) revision = versions[max(versions.keys())].name if versions else "main" api = _get_hf_api() variants = get_variants(api, repo_id=repo_id, revision=revision) if not variants: - print(f"No build variants found in {repo_id} (revision: {revision})", file=sys.stderr) + print( + f"No build variants found in {repo_id} (revision: {revision})", + file=sys.stderr, + ) sys.exit(1) # The metadata fields that describe the kernel (rather than a specific diff --git a/kernels/src/kernels/cli/verify_signature.py b/kernels/src/kernels/cli/verify_signature.py index 5afe2550d..c96e5af1b 100644 --- a/kernels/src/kernels/cli/verify_signature.py +++ b/kernels/src/kernels/cli/verify_signature.py @@ -1,6 +1,8 @@ import argparse import sys +from huggingface_hub import constants + if sys.version_info >= (3, 11): from typing import assert_never else: @@ -13,7 +15,9 @@ def verify_signature(args: argparse.Namespace) -> None: - revision = select_revision_or_version(args.repo_id, revision=None, version=args.version) + revision = select_revision_or_version( + args.repo_id, revision=None, version=args.version, local_files_only=constants.HF_HUB_OFFLINE + ) if args.all_variants: repo_path = install_kernel_all_variants(args.repo_id, revision=revision) diff --git a/kernels/src/kernels/cli/versions.py b/kernels/src/kernels/cli/versions.py index 7b091848a..38846f216 100644 --- a/kernels/src/kernels/cli/versions.py +++ b/kernels/src/kernels/cli/versions.py @@ -1,3 +1,5 @@ +from huggingface_hub import constants + from kernels._versions import _get_available_versions from kernels.hf_hub import _get_hf_api from kernels.variants import ( @@ -9,7 +11,7 @@ def print_kernel_versions(repo_id: str): api = _get_hf_api() - versions = _get_available_versions(repo_id) + versions = _get_available_versions(repo_id, local_files_only=constants.HF_HUB_OFFLINE) if not versions: print(f"Repository does not support kernel versions: {repo_id}") return diff --git a/kernels/src/kernels/hf_hub.py b/kernels/src/kernels/hf_hub.py index 309013c40..91dcc3c3b 100644 --- a/kernels/src/kernels/hf_hub.py +++ b/kernels/src/kernels/hf_hub.py @@ -92,7 +92,7 @@ class RepoInfo: revision: str -def _check_trust_remote_code(repo_id: str, trust_remote_code: bool | list[str]) -> None: +def _check_trust_remote_code(repo_id: str, local_files_only: bool, trust_remote_code: bool | list[str]) -> None: """Check whether a kernel repository is trusted. When ``trust_remote_code`` is ``False`` (the default), only repositories @@ -119,7 +119,7 @@ def _check_trust_remote_code(repo_id: str, trust_remote_code: bool | list[str]) stacklevel=3, ) - if constants.HF_HUB_OFFLINE: + if local_files_only: # Publisher trust cannot be verified offline. The user opted into # offline mode and the kernel must already be in the local cache, # so trust was established when it was originally downloaded. diff --git a/kernels/src/kernels/layer/func.py b/kernels/src/kernels/layer/func.py index 70049bda5..06b1a7e93 100644 --- a/kernels/src/kernels/layer/func.py +++ b/kernels/src/kernels/layer/func.py @@ -4,6 +4,8 @@ from types import ModuleType from typing import TYPE_CHECKING, Protocol, Type +from huggingface_hub import constants + from .._versions import select_revision_or_version from ..load import ( get_kernel, @@ -99,6 +101,7 @@ def _resolve_revision(self) -> str: repo_id=self._repo_id, revision=self._revision, version=self._version, + local_files_only=constants.HF_HUB_OFFLINE, ) def load(self) -> Type["nn.Module"]: diff --git a/kernels/src/kernels/layer/layer.py b/kernels/src/kernels/layer/layer.py index 3c822dd1d..2b2b71eb1 100644 --- a/kernels/src/kernels/layer/layer.py +++ b/kernels/src/kernels/layer/layer.py @@ -9,6 +9,8 @@ from types import MethodType, ModuleType from typing import TYPE_CHECKING, Callable, Protocol, Type +from huggingface_hub import constants + from .._versions import select_revision_or_version from ..load import ( get_kernel, @@ -93,6 +95,7 @@ def _resolve_revision(self) -> str: repo_id=self._repo_id, revision=self._revision, version=self._version, + local_files_only=constants.HF_HUB_OFFLINE, ) def load(self) -> Type["nn.Module"]: diff --git a/kernels/src/kernels/load.py b/kernels/src/kernels/load.py index 9bc74f34a..2a6a40766 100644 --- a/kernels/src/kernels/load.py +++ b/kernels/src/kernels/load.py @@ -96,9 +96,18 @@ def get_kernel( if override is not None: return get_local_kernel(override) - _check_trust_remote_code(repo_id, trust_remote_code) + _check_trust_remote_code( + repo_id=repo_id, + local_files_only=constants.HF_HUB_OFFLINE, + trust_remote_code=trust_remote_code, + ) - revision = select_revision_or_version(repo_id, revision=revision, version=version) + 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, @@ -174,7 +183,12 @@ def has_kernel( Returns: `bool`: `True` if a kernel is available for the current environment. """ - revision = select_revision_or_version(repo_id, revision=revision, version=version) + 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) diff --git a/kernels/src/kernels/locking.py b/kernels/src/kernels/locking.py index cb185eb2d..844611f3e 100644 --- a/kernels/src/kernels/locking.py +++ b/kernels/src/kernels/locking.py @@ -7,6 +7,7 @@ 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 @@ -48,7 +49,7 @@ def get_kernel_locks(repo_id: str, version_spec: int) -> KernelLock: # final destination repo. repo_id, _ = resolve_status(api, repo_id, "main") - tag_for_newest = resolve_version_spec_as_ref(repo_id, version_spec) + 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 diff --git a/kernels/src/kernels/resolver.py b/kernels/src/kernels/resolver.py new file mode 100644 index 000000000..67cd32742 --- /dev/null +++ b/kernels/src/kernels/resolver.py @@ -0,0 +1,381 @@ +from dataclasses import dataclass +from pathlib import Path +from typing import Protocol, Self, runtime_checkable + +from huggingface_hub.errors import LocalEntryNotFoundError +from huggingface_hub.hf_api import HfApi +from kernels_data import KernelDependency, KernelLocks, KernelPaths, Metadata + +from kernels._versions import resolve_kernel_version +from kernels.hf_hub import CACHE_DIR, _check_trust_remote_code +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__/**"] + + +@dataclass(frozen=True) +class LocalKernel: + """ + A kernel that can be loaded from a local path. + + `origin` is `None` if the kernel was a local kernel before installation, + otherwise it will point to the remote kernel that was intalled. + """ + + variant_path: Path + metadata: Metadata + origin: "RemoteKernel | None" = None + + def install(self, *, api: HfApi, all_variants: bool = False) -> Self: + # Local kernels are already installed, so we just return self. + return self + + +@dataclass(frozen=True) +class RemoteKernel: + """A kernel that can be loaded from a remote path.""" + + repo_id: str + revision: str + metadata: Metadata + variant: Variant + + def install(self, *, api: HfApi, all_variants: bool = False) -> LocalKernel: + + allow_patterns = [f"build/{self.variant.variant_str}/*"] + + if all_variants: + repo_path = Path( + str( + api.snapshot_download( + self.repo_id, + repo_type="kernel", + allow_patterns="build/*", + ignore_patterns=_BYTECODE_IGNORE_PATTERNS, + cache_dir=CACHE_DIR, + revision=self.revision, + local_files_only=False, + ) + ) + ) + else: + repo_path = Path( + str( + api.snapshot_download( + self.repo_id, + repo_type="kernel", + allow_patterns=allow_patterns, + ignore_patterns=_BYTECODE_IGNORE_PATTERNS, + cache_dir=CACHE_DIR, + revision=self.revision, + local_files_only=False, + ) + ) + ) + + return LocalKernel( + variant_path=repo_path / "build" / self.variant.variant_str, + metadata=self.metadata, + origin=self, + ) + + +def resolve_local_kernel(repo_path: Path, *, backend: str | None) -> LocalKernel: + variant_path = None + + for base_path in [repo_path, repo_path / "build"]: + variants = get_variants_local(base_path) + variant, trace = resolve_variant(variants, backend) + if variant is not None: + variant_path = base_path / variant.variant_str + break + + # The local kernel path might be pointing directly to the variant. + if variant_path is None: + metadata_path = repo_path / "metadata.json" + if metadata_path.exists(): + variant_path = repo_path + + if variant_path is None: + raise FileNotFoundError( + f"Cannot find a build variant for this system in {repo_path}:\n\n{variants_trace_str(trace)}" + ) + + metadata = Metadata.read_from_file(variant_path / "metadata.json") + location = LocalKernel(variant_path=variant_path, metadata=metadata) + + return location + + +def resolve_hub_kernel( + repo_id: str, + *, + api: HfApi, + backend: str | None, + revision: str, +) -> RemoteKernel: + 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)}" + ) + + metadata_path = 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, + ) + ) + + metadata = Metadata.read_from_file(metadata_path) + location = RemoteKernel(repo_id=repo_id, revision=revision, variant=variant, metadata=metadata) + + return location + + +def resolve_hub_cache_kernel( + api: HfApi, + repo_id: str, + *, + revision: str, + backend: str | None, +) -> LocalKernel: + """Resolve a kernel variant path from the local Hugging Face cache only. + + Used by `load_kernel` (which always operates on a pre-downloaded, locked + kernel) and by the offline branch of `install_kernel`. + """ + try: + repo_path = Path( + str( + api.snapshot_download( + repo_id, + repo_type="kernel", + ignore_patterns=_BYTECODE_IGNORE_PATTERNS, + cache_dir=CACHE_DIR, + revision=revision, + local_files_only=True, + ) + ) + ) + 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(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)}" + ) + + variant_path = repo_path / "build" / variant.variant_str + if not variant_path.exists(): + raise FileNotFoundError(f"Variant path does not exist: `{variant_path}`") + + metadata = Metadata.read_from_file(variant_path / "metadata.json") + location = LocalKernel(variant_path=variant_path, metadata=metadata) + + return location + + +@runtime_checkable +class Resolver(Protocol): + def resolve( + self, *, api: HfApi, backend: str | None, kernel: KernelDependency + ) -> LocalKernel | RemoteKernel | None: + """Resolve a kernel dependency to a location. + + Returns `None` when the kernel cannot be resolved.""" + ... + + +@dataclass +class SequentialResolver: + """Sequential kernel solver. + + This solver tries the embedded solvers sequentially until one succeeds.""" + + resolvers: list[Resolver] + + def resolve( + self, *, api: HfApi, backend: str | None, kernel: KernelDependency + ) -> LocalKernel | RemoteKernel | None: + for resolver in self.resolvers: + location = resolver.resolve(api=api, backend=backend, kernel=kernel) + if location is not None: + return location + + return None + + +@dataclass +class HubResolver: + """Hugging Face Hub kernel solver. + + This solver solves a kernel dependency by finding the kernel on the Hub.""" + + trust_remote_code: bool | list[str] + + def resolve(self, *, api: HfApi, backend: str | None, kernel: KernelDependency) -> RemoteKernel: + # 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=self.trust_remote_code, + ) + + # Resolve the revision that we need. + revision = resolve_kernel_version(kernel, local_files_only=False) + + # Get the kernel metadata for the revision. + return resolve_hub_kernel(kernel.repo_id, api=api, revision=revision, backend=backend) + + +@dataclass +class HubCacheResolver: + """Hugging Face Hub cache kernel solver. + + This solver solves a kernel dependency from the local Hugging Face Hub + cache. The kernel must already be present in the cache.""" + + trust_remote_code: bool | list[str] + + def resolve(self, *, api: HfApi, backend: str | None, kernel: KernelDependency) -> LocalKernel: + _check_trust_remote_code( + repo_id=kernel.repo_id, + local_files_only=True, + trust_remote_code=self.trust_remote_code, + ) + + # Resolve the revision that we need. + revision = resolve_kernel_version(kernel, local_files_only=True) + + # Get the kernel metadata for the revision. + return resolve_hub_cache_kernel( + api, + kernel.repo_id, + revision=revision, + backend=backend, + ) + + +def _locked_revision(kernel_locks: KernelLocks, kernel: KernelDependency) -> str: + kernel_lock = kernel_locks.get(kernel, None) + if kernel_lock is None: + raise ValueError( + f"Kernel `{kernel.repo_id}` is not locked. Please lock it with `kernels lock ` and then reinstall the project." + ) + return kernel_lock.commit + + +@dataclass +class LockedHubResolver: + """Lock file Hugging Face Hub kernel solver. + + This solver uses a lock file to solve kernel dependencies from the Hub.""" + + kernel_locks: KernelLocks + trust_remote_code: bool | list[str] + + def resolve(self, *, api: HfApi, backend: str | None, kernel: KernelDependency) -> RemoteKernel: + # 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=self.trust_remote_code, + ) + + revision = _locked_revision(self.kernel_locks, kernel) + + # Get the kernel metadata for the revision. + return resolve_hub_kernel(kernel.repo_id, api=api, revision=revision, backend=backend) + + +@dataclass +class LockedHubCacheResolver: + """Lock file Hugging Face Hub cache kernel solver. + + This solver uses a lock file to solve kernel dependencies from the local + Hugging Face Hub cache. The kernel must already be present in the cache.""" + + kernel_locks: KernelLocks + trust_remote_code: bool | list[str] + + def resolve(self, *, api: HfApi, backend: str | None, kernel: KernelDependency) -> LocalKernel: + _check_trust_remote_code( + repo_id=kernel.repo_id, + local_files_only=True, + trust_remote_code=self.trust_remote_code, + ) + + revision = _locked_revision(self.kernel_locks, kernel) + + # Get the kernel metadata for the revision. + return resolve_hub_cache_kernel( + api, + kernel.repo_id, + revision=revision, + backend=backend, + ) + + +@dataclass +class KernelPathsResolver: + """Path-based kernel solver. + + This solver uses a kernel dependency to path mapping to solve kernel + dependencies.""" + + kernel_paths: KernelPaths + + def resolve(self, *, api: HfApi, backend: str | None, kernel: KernelDependency) -> LocalKernel | None: + repo_path = self.kernel_paths.get(kernel, None) + if repo_path is None: + return None + + return resolve_local_kernel(repo_path, backend=backend) + + +@dataclass +class RepoPathsResolver: + """Path-based kernel solver. + + This solver uses a repo ID to path mapping to solve kernel dependencies.""" + + local_kernels: dict[str, Path] + + def resolve(self, *, api: HfApi, backend: str | None, kernel: KernelDependency) -> LocalKernel | None: + repo_path = self.local_kernels.get(kernel.repo_id, None) + if repo_path is None: + return None + + return resolve_local_kernel(repo_path, backend=backend) + + +@dataclass +class NoopResolver: + def resolve(self, *, api: HfApi, backend: str | None, kernel: KernelDependency) -> LocalKernel | None: + return None diff --git a/kernels/src/kernels/variants.py b/kernels/src/kernels/variants.py index 50ba0b1e6..443cd9211 100644 --- a/kernels/src/kernels/variants.py +++ b/kernels/src/kernels/variants.py @@ -8,7 +8,7 @@ from pathlib import Path from typing import ClassVar -from huggingface_hub import HfApi +from huggingface_hub import HfApi, constants from huggingface_hub.dataclasses import strict from huggingface_hub.hf_api import RepoFolder from packaging.version import Version, parse @@ -557,6 +557,7 @@ def sort_key(vs: Decision) -> tuple[int, ...]: def get_kernel_variants( repo_id: str, + *, revision: str | None = None, version: int | None = None, backend: str | None = None, @@ -598,7 +599,12 @@ def get_kernel_variants( """ from kernels.hf_hub import _get_hf_api - revision = select_revision_or_version(repo_id, revision=revision, version=version) + 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) diff --git a/kernels/tests/test_basic.py b/kernels/tests/test_basic.py index 93d66ea28..fadbedc39 100644 --- a/kernels/tests/test_basic.py +++ b/kernels/tests/test_basic.py @@ -306,7 +306,11 @@ def test_version_resolution_offline_missing(monkeypatch): 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) + resolve_version_spec_as_ref( + "kernels-test/this-repo-should-not-exist", + 1, + local_files_only=constants.HF_HUB_OFFLINE, + ) def silu_and_mul_torch(x: torch.Tensor): diff --git a/kernels/tests/test_resolver.py b/kernels/tests/test_resolver.py new file mode 100644 index 000000000..3ccb3d338 --- /dev/null +++ b/kernels/tests/test_resolver.py @@ -0,0 +1,323 @@ +import json +import re +from dataclasses import dataclass +from pathlib import Path + +import pytest +from kernels_data import ( + KernelDependency, + KernelLock, + KernelLocks, + KernelName, + KernelPaths, + KernelVersion, + Metadata, +) + +from kernels._versions import resolve_version_spec_as_ref +from kernels.hf_hub import _get_hf_api +from kernels.install import install_kernel +from kernels.resolver import ( + HubCacheResolver, + HubResolver, + KernelPathsResolver, + LocalKernel, + LockedHubCacheResolver, + LockedHubResolver, + NoopResolver, + RemoteKernel, + RepoPathsResolver, + Resolver, + SequentialResolver, + _locked_revision, +) + + +@pytest.fixture(scope="module") +def api(): + return _get_hf_api() + + +@pytest.fixture(scope="module") +def installed_relu_cpu(): + """Install the relu kernel CPU variant, so that it is in the local cache.""" + return install_kernel("kernels-community/relu", revision="v1", backend="cpu") + + +@pytest.fixture(scope="module") +def relu_locks(): + """Locks for `kernels-community/relu` version 1, mirroring a lock file.""" + commit = resolve_version_spec_as_ref("kernels-community/relu", 1, local_files_only=False).target_commit + dep = KernelDependency(repo_id="kernels-community/relu", version=KernelVersion.Version(1)) + return dep, KernelLocks({dep: KernelLock(commit=commit)}), commit + + +def _dep(repo_id: str = "test/kernel", version: int = 1) -> KernelDependency: + return KernelDependency(repo_id=repo_id, version=KernelVersion.Version(version)) + + +def _metadata() -> Metadata: + return Metadata.from_bytes( + json.dumps( + { + "id": "test_1_cpu", + "name": "test", + "version": 1, + "license": "Apache-2.0", + "python-depends": ["torch"], + "backend": {"type": "cpu"}, + } + ).encode() + ) + + +def _local_kernel(variant_path: str) -> LocalKernel: + return LocalKernel(variant_path=Path(variant_path), metadata=_metadata()) + + +def _write_variant(repo_path: Path, variant: str = "torch-cpu") -> Path: + """Write a fake kernel build variant (CPU noarch) under `repo_path/build`.""" + variant_dir = repo_path / "build" / variant + variant_dir.mkdir(parents=True) + (variant_dir / "metadata.json").write_text( + json.dumps( + { + "id": "test_1_cpu", + "name": "test", + "version": 1, + "license": "Apache-2.0", + "python-depends": ["torch"], + "backend": {"type": "cpu"}, + } + ) + ) + return variant_dir + + +def test_noop_resolver_never_resolves(api): + assert NoopResolver().resolve(api=api, backend=None, kernel=_dep()) is None + + +@dataclass +class _FixedResolver: + """Test resolver that always returns a fixed result.""" + + result: LocalKernel | None + + def resolve(self, *, api, backend, kernel) -> LocalKernel | None: + return self.result + + +def test_sequential_resolver_returns_first_match(api): + first = _local_kernel("/tmp/first") + second = _local_kernel("/tmp/second") + resolver = SequentialResolver(resolvers=[_FixedResolver(first), _FixedResolver(second)]) + + assert resolver.resolve(api=api, backend=None, kernel=_dep()) is first + + +def test_sequential_resolver_skips_non_matching(api): + match = _local_kernel("/tmp/match") + resolver = SequentialResolver(resolvers=[NoopResolver(), _FixedResolver(None), _FixedResolver(match)]) + + assert resolver.resolve(api=api, backend=None, kernel=_dep()) is match + + +def test_sequential_resolver_returns_none_when_nothing_matches(api): + resolver = SequentialResolver(resolvers=[NoopResolver(), _FixedResolver(None)]) + + assert resolver.resolve(api=api, backend=None, kernel=_dep()) is None + + +@pytest.mark.parametrize("path_kind", ["repo", "build", "variant"]) +def test_repo_paths_resolver_path_layouts(api, tmp_path, path_kind): + variant_dir = _write_variant(tmp_path) + repo_path = { + "repo": tmp_path, + "build": tmp_path / "build", + "variant": variant_dir, + }[path_kind] + + resolver = RepoPathsResolver(local_kernels={"test/kernel": repo_path}) + location = resolver.resolve(api=api, backend="cpu", kernel=_dep("test/kernel")) + + assert isinstance(location, LocalKernel) + assert location.variant_path == variant_dir + assert location.metadata.name == KernelName("test") + + +def test_repo_paths_resolver_unknown_repo(api, tmp_path): + _write_variant(tmp_path) + resolver = RepoPathsResolver(local_kernels={"test/kernel": tmp_path}) + + assert resolver.resolve(api=api, backend="cpu", kernel=_dep("other/repo")) is None + + +def test_repo_paths_resolver_no_matching_variant(api, tmp_path): + resolver = RepoPathsResolver(local_kernels={"test/kernel": tmp_path}) + + with pytest.raises(FileNotFoundError, match="Cannot find a build variant"): + resolver.resolve(api=api, backend="cpu", kernel=_dep("test/kernel")) + + +def test_kernel_paths_resolver_resolves_known_dep(api, tmp_path): + variant_dir = _write_variant(tmp_path) + dep = _dep("test/kernel", version=2) + resolver = KernelPathsResolver(kernel_paths=KernelPaths({dep: tmp_path})) + + location = resolver.resolve(api=api, backend="cpu", kernel=dep) + + assert isinstance(location, LocalKernel) + assert location.variant_path == variant_dir + + +def test_kernel_paths_resolver_unknown_dep(api, tmp_path): + _write_variant(tmp_path) + resolver = KernelPathsResolver(kernel_paths=KernelPaths({_dep("test/kernel", version=1): tmp_path})) + assert resolver.resolve(api=api, backend="cpu", kernel=_dep("test/kernel", version=2)) is None + + +def test_hub_resolver_resolves_remote_kernel(api): + location = HubResolver(trust_remote_code=False).resolve( + api=api, backend="cpu", kernel=_dep("kernels-community/relu", version=1) + ) + + assert isinstance(location, RemoteKernel) + assert location.repo_id == "kernels-community/relu" + assert re.fullmatch(r"[0-9a-f]{40}", location.revision) + assert location.metadata.name == KernelName("relu") + + +def test_hub_resolver_revision_passthrough(api): + location = HubResolver(trust_remote_code=False).resolve( + api=api, + backend="cpu", + kernel=KernelDependency(repo_id="kernels-community/relu", version=KernelVersion.Revision("v1")), + ) + + assert location.revision == "v1" + + +def test_hub_resolver_blocks_untrusted_org(api): + with pytest.raises(ValueError, match="not from a trusted publisher"): + HubResolver(trust_remote_code=False).resolve( + api=api, backend="cpu", kernel=_dep("kernels-test-untrusted/not-a-trused-org-kernel", version=1) + ) + + +@pytest.mark.cuda_only +def test_hub_resolver_trust_remote_code_bypasses_check(api): + location = HubResolver(trust_remote_code=True).resolve( + api=api, backend="cuda", kernel=_dep("kernels-test-untrusted/ci-test-kernel", version=1) + ) + + assert isinstance(location, RemoteKernel) + + +def test_hub_resolver_no_matching_variant(api): + with pytest.raises(FileNotFoundError, match="Cannot find a build variant"): + HubResolver(trust_remote_code=False).resolve( + api=api, + backend="cpu", + kernel=KernelDependency(repo_id="kernels-test/only-torch-2.4", version=KernelVersion.Revision("main")), + ) + + +def test_hub_cache_resolver_resolves_cached_kernel(api, installed_relu_cpu): + location = HubCacheResolver(trust_remote_code=False).resolve( + api=api, backend="cpu", kernel=_dep("kernels-community/relu", version=1) + ) + + assert isinstance(location, LocalKernel) + assert location.variant_path == installed_relu_cpu + + +def test_hub_cache_resolver_revision_passthrough(api, installed_relu_cpu): + location = HubCacheResolver(trust_remote_code=False).resolve( + api=api, + backend="cpu", + kernel=KernelDependency(repo_id="kernels-community/relu", version=KernelVersion.Revision("v1")), + ) + + assert location.variant_path == installed_relu_cpu + + +def test_hub_cache_resolver_uncached_repo(api): + with pytest.raises(FileNotFoundError, match="local snapshot"): + HubCacheResolver(trust_remote_code=False).resolve( + api=api, + backend="cpu", + kernel=KernelDependency( + repo_id="kernels-test/this-repo-should-not-exist", + version=KernelVersion.Revision("0" * 40), + ), + ) + + +def test_locked_hub_resolver_resolves_locked_revision(api, relu_locks): + dep, locks, commit = relu_locks + + location = LockedHubResolver(kernel_locks=locks, trust_remote_code=False).resolve( + api=api, backend="cpu", kernel=dep + ) + + assert isinstance(location, RemoteKernel) + assert location.revision == commit + + +def test_locked_hub_resolver_requires_lock(api, relu_locks): + _, locks, _ = relu_locks + + with pytest.raises(ValueError, match="is not locked"): + LockedHubResolver(kernel_locks=locks, trust_remote_code=False).resolve( + api=api, backend="cpu", kernel=_dep("kernels-community/relu", version=2) + ) + + +def test_locked_hub_cache_resolver_resolves_locked_kernel(api, relu_locks, installed_relu_cpu): + dep, locks, _ = relu_locks + + location = LockedHubCacheResolver(kernel_locks=locks, trust_remote_code=False).resolve( + api=api, backend="cpu", kernel=dep + ) + + assert isinstance(location, LocalKernel) + assert location.variant_path == installed_relu_cpu + + +def test_locked_hub_cache_resolver_requires_lock(api, relu_locks): + _, locks, _ = relu_locks + + with pytest.raises(ValueError, match="is not locked"): + LockedHubCacheResolver(kernel_locks=locks, trust_remote_code=False).resolve( + api=api, backend="cpu", kernel=_dep("kernels-community/relu", version=2) + ) + + +def test_locked_revision_returns_commit(): + dep = _dep("test/kernel", version=1) + locks = KernelLocks({dep: KernelLock(commit="a" * 40)}) + + assert _locked_revision(locks, dep) == "a" * 40 + + +def test_locked_revision_requires_lock(): + with pytest.raises(ValueError, match="is not locked"): + _locked_revision(KernelLocks({}), _dep("test/kernel", version=1)) + + +@pytest.mark.parametrize( + "resolver", + [ + NoopResolver(), + SequentialResolver(resolvers=[]), + RepoPathsResolver(local_kernels={}), + KernelPathsResolver(kernel_paths=KernelPaths({})), + HubResolver(trust_remote_code=False), + HubCacheResolver(trust_remote_code=False), + LockedHubResolver(kernel_locks=KernelLocks({}), trust_remote_code=False), + LockedHubCacheResolver(kernel_locks=KernelLocks({}), trust_remote_code=False), + ], +) +def test_resolver_protocol_conformance(resolver): + assert isinstance(resolver, Resolver) diff --git a/kernels/tests/test_verify.py b/kernels/tests/test_verify.py index c1f1b97b6..caf819179 100644 --- a/kernels/tests/test_verify.py +++ b/kernels/tests/test_verify.py @@ -11,13 +11,13 @@ def test_correctly_signed_kernel_passes_with_default_policy(): - revision = select_revision_or_version("kernels-community/relu", revision=None, version=1) + revision = select_revision_or_version("kernels-community/relu", revision=None, version=1, local_files_only=False) variant_path = install_kernel("kernels-community/relu", revision=revision) assert verify_variant(variant_path) == VerificationResult.Success() def test_correctly_signed_kernel_passes(): - revision = select_revision_or_version("kernels-test/signatures", revision=None, version=1) + revision = select_revision_or_version("kernels-test/signatures", revision=None, version=1, local_files_only=False) variant_path = install_kernel("kernels-test/signatures", revision=revision) assert ( verify_variant(