feat: resolve repo type and fetch accordingly (#435)

* feat: resolve repo type and fetch accordingly

* fix: lint and format and prefer pinning hfh commit

* fix: pin hub library from git commit

* fix: update tests to specify repo type

* fix: bump overlay

* fix: prefer the hfh release

* fix: bump hf xet

* fix: improve _resolve_repo_type

* fix: remove debug line

* fix: prefer repo_info in _resolve_repo_type and handle 401

* fix: expect repo type model for the kernels-test org kernels

* fix: adjust error

* fix: adjust tests and reo type resolution

* fix: accept revision in _resolve_repo_type

* fix: get file metadata in repo info and remove extra try

* feat: run lock pytest as subprocess

* feat: prefer locking kernel repo hashes

* fix: avoid deprecation warning for now

* feat: make repo type optional and comment for siblings

* fix: remove repo_type from all public interfaces

* fix: remove backward compat of model repo types

* fix: improve error cases

* fix: bump lock files for kernel repos

* fix: avoid non version locking tests

* fix: avoid upload changes

* fix: apply lints after rebase

* fix: adjust repo not found error regression

* fix: adjust repo not found error regression in layer test too

* fix: revert back to hfhub error until 401 resolved

* fix: avoid the unneeded removal of repo type

* feat: enforce kwargs

* fix: adjust for kwarg

* fix: update test for kwarg changes

Author: drbh

kernels/pyproject.toml

@@ -8,9 +8,9 @@ authors = [
 ]
 license = { text = "Apache-2.0" }
 readme = "README.md"
-requires-python = ">= 3.9"
+requires-python = ">= 3.10"
 dependencies = [
-  "huggingface_hub>=1.3.0,<2.0",
+  "huggingface-hub>=1.10.0",
   "packaging>=20.0",
   "pyyaml>=6",
   "tomli>=2.0; python_version<'3.11'",

kernels/src/kernels/_versions.py

@@ -12,8 +12,10 @@ def _get_available_versions(repo_id: str) -> dict[int, GitRefInfo]:
     """Get kernel versions that are available in the repository."""
     from kernels.utils import _get_hf_api
 
+    refs = _get_hf_api().list_repo_refs(repo_id=repo_id, repo_type="kernel")
+
     versions = {}
-    for branch in _get_hf_api().list_repo_refs(repo_id).branches:
+    for branch in refs.branches:
         if not branch.name.startswith("v"):
             continue
         try:
@@ -33,7 +35,7 @@ def _get_available_versions_old(repo_id: str) -> dict[Version, GitRefInfo]:
     from kernels.utils import _get_hf_api
 
     versions = {}
-    for tag in _get_hf_api().list_repo_refs(repo_id).tags:
+    for tag in _get_hf_api().list_repo_refs(repo_id, repo_type="kernel").tags:
         if not tag.name.startswith("v"):
             continue
         try:
@@ -46,13 +48,14 @@ def _get_available_versions_old(repo_id: str) -> dict[Version, GitRefInfo]:
 
 def resolve_version_spec_as_ref(repo_id: str, version_spec: int | str) -> GitRefInfo:
     """
-    Get the locks for a kernel with the given version spec.
+    Get the ref for a kernel with the given version spec.
 
     The version specifier can be any valid Python version specifier:
     https://packaging.python.org/en/latest/specifications/version-specifiers/#version-specifiers
     """
     if isinstance(version_spec, int):
         versions = _get_available_versions(repo_id)
+
         ref = versions.get(version_spec, None)
         if ref is None:
             raise ValueError(

kernels/src/kernels/cli/__init__.py

@@ -164,14 +164,14 @@ def download_kernels(args):
         if args.all_variants:
             install_kernel_all_variants(
                 kernel_lock.repo_id,
-                kernel_lock.sha,
+                revision=kernel_lock.sha,
                 variant_locks=kernel_lock.variants,
             )
         else:
             try:
                 install_kernel(
                     kernel_lock.repo_id,
-                    kernel_lock.sha,
+                    revision=kernel_lock.sha,
                     variant_locks=kernel_lock.variants,
                 )
             except FileNotFoundError as e:

kernels/src/kernels/cli/check.py

@@ -15,13 +15,19 @@
 
 
 def check_kernel(
-    *, macos: str, manylinux: str, python_abi: str, repo_id: str, revision: str
+    *,
+    macos: str,
+    manylinux: str,
+    python_abi: str,
+    repo_id: str,
+    revision: str,
 ):
     variants_path = (
         Path(
             str(
                 _get_hf_api().snapshot_download(
                     repo_id,
+                    repo_type="kernel",
                     allow_patterns=["build/*"],
                     cache_dir=CACHE_DIR,
                     revision=revision,

kernels/src/kernels/cli/upload.py

@@ -70,7 +70,6 @@ def _upload_build_dir(
             repo_id=repo_id,
             operations=list(chunk),
             revision=revision,
-            repo_type="model",
             commit_message=commit_message,
         )
 

kernels/src/kernels/cli/versions.py

@@ -8,13 +8,12 @@
 
 def print_kernel_versions(repo_id: str):
     api = _get_hf_api()
-
-    versions = _get_available_versions(repo_id).items()
+    versions = _get_available_versions(repo_id)
     if not versions:
         print(f"Repository does not support kernel versions: {repo_id}")
         return
 
-    for version, ref in sorted(versions, key=lambda x: x[0]):
+    for version, ref in sorted(versions.items(), key=lambda x: x[0]):
         variants = get_variants(api, repo_id=repo_id, revision=ref.ref)
         resolved = resolve_variants(variants, None)
         best = resolved[0] if resolved else None

kernels/src/kernels/lockfile.py

@@ -3,6 +3,7 @@
 from pathlib import Path
 
 from huggingface_hub.dataclasses import strict
+from huggingface_hub.hf_api import RepoFile
 
 from kernels._versions import resolve_version_spec_as_ref
 from kernels.compat import tomllib
@@ -49,21 +50,31 @@ def get_kernel_locks(repo_id: str, version_spec: int | str) -> KernelLock:
 
     tag_for_newest = resolve_version_spec_as_ref(repo_id, version_spec)
 
+    revision = tag_for_newest.target_commit
+
     r = api.repo_info(
-        repo_id=repo_id, revision=tag_for_newest.target_commit, files_metadata=True
+        repo_id=repo_id,
+        repo_type="kernel",
+        revision=revision,
     )
     if r.sha is None:
         raise ValueError(
             f"Cannot get commit SHA for repo {repo_id} for tag {tag_for_newest.name}"
         )
 
-    if r.siblings is None:
-        raise ValueError(
-            f"Cannot get sibling information for {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 r.siblings:
+    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}")

kernels/src/kernels/status.py

@@ -4,6 +4,7 @@
 
 from huggingface_hub import HfApi
 from huggingface_hub.dataclasses import strict
+from huggingface_hub.errors import RepositoryNotFoundError
 from huggingface_hub.utils import EntryNotFoundError
 
 from kernels.compat import tomllib
@@ -54,11 +55,14 @@ def check_status(
     ) -> KernelStatusKind | None:
         try:
             path = api.hf_hub_download(
-                repo_id=repo_id, filename="kernel-status.toml", revision=revision
+                repo_id=repo_id,
+                repo_type="kernel",
+                filename="kernel-status.toml",
+                revision=revision,
             )
             with open(path, "r") as f:
                 return KernelStatus.from_toml(f.read())
-        except EntryNotFoundError:
+        except (EntryNotFoundError, RepositoryNotFoundError):
             return None
 
 

kernels/src/kernels/utils.py

@@ -131,6 +131,7 @@ def _import_from_path(
 
 def install_kernel(
     repo_id: str,
+    *,
     revision: str,
     local_files_only: bool = False,
     backend: str | None = None,
@@ -176,10 +177,12 @@ def install_kernel(
         )
 
     allow_patterns = [f"build/{variant.variant_str}/*"]
+
     repo_path = Path(
         str(
             api.snapshot_download(
                 repo_id,
+                repo_type="kernel",
                 allow_patterns=allow_patterns,
                 cache_dir=CACHE_DIR,
                 revision=revision,
@@ -234,15 +237,18 @@ def _find_kernel_in_repo_path(
 
 def install_kernel_all_variants(
     repo_id: str,
+    *,
     revision: str,
     local_files_only: bool = False,
     variant_locks: dict[str, VariantLock] | None = None,
 ) -> Path:
     api = _get_hf_api()
+
     repo_path = Path(
         str(
             api.snapshot_download(
                 repo_id,
+                repo_type="kernel",
                 allow_patterns="build/*",
                 cache_dir=CACHE_DIR,
                 revision=revision,
@@ -318,7 +324,10 @@ def get_kernel(
         backend=backend,
     )
     package_name, variant_path = install_kernel(
-        repo_id, revision=revision, backend=backend, user_agent=user_agent
+        repo_id,
+        revision=revision,
+        backend=backend,
+        user_agent=user_agent,
     )
     return _import_from_path(package_name, variant_path, _repo_infos=repo_infos)
 
@@ -396,6 +405,7 @@ def has_kernel(
     for init_file in ["__init__.py", f"{package_name}/__init__.py"]:
         if api.file_exists(
             repo_id,
+            repo_type="kernel",
             revision=revision,
             filename=f"build/{variant.variant_str}/{init_file}",
         ):
@@ -454,6 +464,7 @@ def load_kernel(
         str(
             api.snapshot_download(
                 repo_id,
+                repo_type="kernel",
                 allow_patterns=allow_patterns,
                 cache_dir=CACHE_DIR,
                 revision=locked_sha,
@@ -495,7 +506,7 @@ def get_locked_kernel(repo_id: str, local_files_only: bool = False) -> ModuleTyp
         raise ValueError(f"Kernel `{repo_id}` is not locked")
 
     package_name, variant_path = install_kernel(
-        repo_id, locked_sha, local_files_only=local_files_only
+        repo_id, revision=locked_sha, local_files_only=local_files_only
     )
 
     return _import_from_path(package_name, variant_path)

kernels/src/kernels/variants.py

@@ -295,7 +295,10 @@ def parse_variant(variant_str: str) -> Variant:
 def get_variants(api: HfApi, *, repo_id: str, revision: str) -> list[Variant]:
     """Get all the build variants available from a kernel repository."""
 
-    tree = api.list_repo_tree(repo_id, path_in_repo="build", revision=revision)
+    tree = api.list_repo_tree(
+        repo_id, path_in_repo="build", repo_type="kernel", revision=revision
+    )
+
     variant_strs = {
         item.path.split("/")[-1] for item in tree if isinstance(item, RepoFolder)
     }
@@ -333,6 +336,7 @@ def resolve_variant(
 ) -> Variant | None:
     """Return the best matching variant for the current system."""
     resolved = resolve_variants(variants, backend)
+
     return resolved[0] if resolved else None