Add kernel resolver infrastructure

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=<major>` for a stable kernel API version or `revision=<branch/tag/commit>` "
+        "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. "

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

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)

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

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.

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"]:

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"]:

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)

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