Add kernel dependency tree resolver

.github/workflows/test_kernels.yaml

@@ -84,7 +84,7 @@ jobs:
           HF_HUB_DOWNLOAD_TIMEOUT: 60
         run: |
           uv pip install einops nvidia-cutlass-dsl
-          uv run pytest tests/test_deps.py
+          uv run pytest tests/test_python_deps.py
 
       - name: Import check without torch
         working-directory: ./kernels

kernels/src/kernels/__init__.py

@@ -6,6 +6,7 @@
 
 from kernels._windows import _add_additional_dll_paths
 from kernels.benchmark import Benchmark
+from kernels.deps import get_kernel_dep
 from kernels.hf_hub import RepoInfo
 from kernels.importer import LoadedKernel, get_loaded_kernels
 from kernels.install import install_kernel
@@ -62,6 +63,7 @@
     "VariantAccepted",
     "VariantRejected",
     "get_kernel",
+    "get_kernel_dep",
     "get_kernel_variants",
     "get_loaded_kernels",
     "get_local_kernel",

kernels/src/kernels/deps.py

@@ -1,106 +1,172 @@
-import importlib.util
-import json
-from dataclasses import dataclass, field
-from pathlib import Path
+from contextvars import ContextVar
+from dataclasses import dataclass
+from types import ModuleType
+from typing import Generic, TypeVar
 
-from huggingface_hub.dataclasses import strict
-from kernels_data import Metadata
+from huggingface_hub.hf_api import HfApi
+from kernels_data import KernelDependency, KernelVersion
 
-from kernels.backends import Backend, _backend
+from kernels.backends import _backend
+from kernels.hf_hub import RepoInfo
+from kernels.importer import _import_from_path
+from kernels.python_deps import validate_dependencies
+from kernels.resolver import (
+    LocalKernel,
+    RemoteKernel,
+    Resolver,
+)
+
+# Default state is `None`, to signal that we are not in a kernel
+# loading context.
+_KERNEL_DEPS: ContextVar[dict[str, ModuleType] | None] = ContextVar("_KERNEL_DEPS", default=None)
+
+
+# Tree node type variable.
+#
+# `LocalKernel | RemoteKernel`: the initial tree.
+# `LocalKernel`: the tree after installing kernels.
+T = TypeVar("T", LocalKernel | RemoteKernel, LocalKernel)
 
 
-@strict
 @dataclass
-class PythonPackage:
-    pkg: str
-    import_name: str | None
-
-    @staticmethod
-    def from_dict(data: dict) -> "PythonPackage":
-        return PythonPackage(
-            pkg=data["pkg"],
-            import_name=data.get("import"),
+class DepTreeNode(Generic[T]):
+    """Node in a kernel dependency tree.
+
+    The content of the node is the location of the kernel (`location`). The
+    outgoing edges (`deps`) are dependencies of the kernel."""
+
+    location: T
+    deps: dict[str, "DepTreeNode[T]"]
+
+    def install(self, *, api: HfApi) -> "DepTreeNode[LocalKernel]":
+        """Install the kernel and its dependencies.
+
+        This will download the kernel to the Hub cache if it is a remote
+        kernel. Otherwise, installation is a no-op."""
+        local_deps = {}
+        for repo_id, node in self.deps.items():
+            local_deps[repo_id] = node.install(api=api)
+
+        return DepTreeNode(
+            location=self.location.install(api=api),
+            deps=local_deps,
         )
 
+    def load(self) -> ModuleType:
+        if isinstance(self.location, RemoteKernel):
+            raise RuntimeError("Can only load installed kernels, run `install()` on the kernel dependency tree first.")
 
-@strict
-@dataclass
-class DependencyInfo:
-    nix: list[str]
-    python: list[PythonPackage]
-
-    @staticmethod
-    def from_dict(data: dict) -> "DependencyInfo":
-        return DependencyInfo(
-            nix=data.get("nix", []),
-            python=[PythonPackage.from_dict(p) for p in data.get("python", [])],
+        deps = {repo_id: dep.load() for repo_id, dep in self.deps.items()}
+
+        repo_info = (
+            RepoInfo(
+                repo_id=self.location.origin.repo_id,
+                revision=self.location.origin.revision,
+            )
+            if self.location.origin
+            else None
         )
 
+        return _import_from_path(self.location.variant_path, repo_info=repo_info, deps=deps)
 
-@strict
-@dataclass
-class DependencyData:
-    general: dict = field(default_factory=dict)
-    backends: dict = field(default_factory=dict)
+    def validate_dependencies(self) -> None:
+        """Validate that the dependencies of this kernel are satisfied."""
 
-    @staticmethod
-    def from_dict(data: dict) -> "DependencyData":
-        general = {name: DependencyInfo.from_dict(info) for name, info in data.get("general", {}).items()}
-        backends = {
-            backend_name: {name: DependencyInfo.from_dict(info) for name, info in deps.items()}
-            for backend_name, deps in data.get("backends", {}).items()
-        }
-        return DependencyData(general=general, backends=backends)
+        validate_dependencies(
+            self.location.metadata.name.python_name,
+            self.location.metadata.python_depends,
+            _backend(),
+        )
 
+        for node in self.deps.values():
+            node.validate_dependencies()
 
-try:
-    with open(Path(__file__).parent / "python_depends.json", "r") as f:
-        _DEPENDENCY_DATA = DependencyData.from_dict(json.load(f))
-except FileNotFoundError:
-    raise FileNotFoundError("Cannot load dependency data, is `kernels` correctly installed?")
 
+LoadCache = dict[KernelDependency, DepTreeNode]
 
-def validate_variant_dependencies(variant_path: Path) -> None:
-    metadata = Metadata.read_from_file(variant_path / "metadata.json")
-    validate_dependencies(metadata.name.python_name, metadata.python_depends, _backend())
 
+def resolve_kernel_tree(
+    api: HfApi,
+    kernel: KernelDependency,
+    backend: str | None,
+    resolver: Resolver | None,
+    seen: set[KernelDependency] | None = None,
+) -> DepTreeNode[LocalKernel | RemoteKernel]:
+    """
+    Recursively solve kernel dependencies.
 
-def validate_dependencies(kernel_module_name: str, dependencies: list[str], backend: Backend):
+    Constructs a tree where nodes encode kernel information (e.g. location and
+    metadata) and edges kernel-kernel dependendies.
     """
-    Validate a list of dependencies to ensure they are installed.
+    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)
+
+    location: LocalKernel | RemoteKernel | None
+
+    location = resolver.resolve(api=api, backend=backend, kernel=kernel) if resolver else None
+
+    if location is None:
+        match kernel.version:
+            case KernelVersion.Version(version=version):
+                version_str = f"version: {version}"
+            case KernelVersion.Revision(revision=revision):
+                version_str = f"revision: {revision}"
+        raise ValueError(f"Could not resolve kernel: {kernel.repo_id} ({version_str})")
+
+    # Recurse into dependencies.
+    kernel_deps = {}
+    for dep in location.metadata.kernel_depends:
+        kernel_deps[dep.repo_id] = resolve_kernel_tree(
+            api=api,
+            backend=backend,
+            seen=seen,
+            kernel=dep,
+            resolver=resolver,
+        )
+
+    seen.remove(kernel)
+
+    return DepTreeNode(
+        location=location,
+        deps=kernel_deps,
+    )
+
+
+def use_kernel_deps(deps: dict[str, ModuleType]):
+    class ContextManager:
+        def __enter__(self):
+            self.token = _KERNEL_DEPS.set(deps)
+
+        def __exit__(self, exc_type, exc_value, traceback):
+            _KERNEL_DEPS.reset(self.token)
+
+    return ContextManager()
+
+
+def get_kernel_dep(repo_id: str) -> ModuleType:
+    """Get a kernel dependency.
+
+    This function can be used by a kernel to get one of its dependencies.
+    The exact version/revision of the dependency must be encoded in the
+    kernel metadata. This function can only be called during kernel loading.
 
     Args:
-        dependencies (`list[str]`): A list of dependency strings to validate.
-        backend (`str`): The backend to validate dependencies for.
+        repo_id (`str`):
+            The Hub kernel repository containing the dependency.
     """
-    general_deps = _DEPENDENCY_DATA.general
-    backend_deps = _DEPENDENCY_DATA.backends.get(backend.name, {})
-
-    # Validate each dependency
-    for dependency in dependencies:
-        # Look up dependency in general dependencies first, then backend-specific
-        if dependency in general_deps:
-            dep_info = general_deps[dependency]
-        elif dependency in backend_deps:
-            dep_info = backend_deps[dependency]
-        else:
-            # Dependency not found in general or backend-specific dependencies
-            raise ValueError(f"Kernel module `{kernel_module_name}` uses unsupported kernel dependency: {dependency}")
-
-        # Check if each python package is installed
-        for python_package in dep_info.python:
-            pkg_name = python_package.pkg
-            # Assertion because this should not happen and is a bug.
-            assert pkg_name is not None, f"Invalid dependency data for `{dependency}`: missing `pkg` field."
-
-            module_name = python_package.import_name
-            if module_name is None:
-                # These are typically packages that do not provide any Python
-                # code, but get installed to Python's library dirctory. E.g.
-                # OneAPI.
-                continue
-
-            if importlib.util.find_spec(module_name) is None:
-                raise ImportError(
-                    f"Kernel module `{kernel_module_name}` requires Python dependency `{pkg_name}`. Please install with: pip install {pkg_name}"
-                )
+    deps = _KERNEL_DEPS.get()
+    if deps is None:
+        raise RuntimeError("`get_kernel_dep` only works during kernel loading.")
+
+    module = deps.get(repo_id)
+    if module is None:
+        raise RuntimeError(
+            f"Dependency '{repo_id}' not found in kernel dependencies. Ensure the dependency is added to `kernel-deps` in build.toml."
+        )
+
+    return module

kernels/src/kernels/importer.py

@@ -91,7 +91,11 @@ def _warn_if_dirty(metadata: Metadata, variant_str: str) -> None:
     )
 
 
-def _import_from_path(variant_path: Path, repo_info: RepoInfo | None = None) -> ModuleType:
+def _import_from_path(
+    variant_path: Path,
+    deps: dict[str, ModuleType],
+    repo_info: RepoInfo | None = None,
+) -> ModuleType:
     if (loaded_kernel := _loaded_kernels.get(variant_path)) is not None:
         return loaded_kernel.module
 

kernels/src/kernels/install.py

@@ -3,8 +3,8 @@
 from huggingface_hub import HfApi
 from huggingface_hub.errors import LocalEntryNotFoundError
 
-from kernels.deps import validate_variant_dependencies
 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,

kernels/src/kernels/load.py

@@ -6,14 +6,14 @@
 from huggingface_hub import constants
 
 from kernels._versions import select_revision_or_version
-from kernels.deps import validate_variant_dependencies
 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.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,
@@ -120,7 +120,7 @@ def get_kernel(
         validate_dependencies=True,
         local_files_only=constants.HF_HUB_OFFLINE,
     )
-    return _import_from_path(variant_path, repo_info=repo_info)
+    return _import_from_path(variant_path, deps={}, repo_info=repo_info)
 
 
 def get_local_kernel(
@@ -147,14 +147,14 @@ def get_local_kernel(
         if variant is not None:
             variant_path = base_path / variant.variant_str
             validate_variant_dependencies(variant_path)
-            return _import_from_path(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)
+        return _import_from_path(variant_path, deps={})
 
     raise FileNotFoundError(f"Could not find kernel in {repo_path}")
 
@@ -247,7 +247,7 @@ def load_kernel(
             "applicable variant. Make sure it's downloaded locally via "
             "`kernels download <project>`."
         ) from e
-    return _import_from_path(variant_path)
+    return _import_from_path(variant_path, deps={})
 
 
 def get_locked_kernel(repo_id: str) -> ModuleType:
@@ -275,4 +275,4 @@ def get_locked_kernel(repo_id: str) -> ModuleType:
         validate_dependencies=True,
     )
 
-    return _import_from_path(variant_path)
+    return _import_from_path(variant_path, deps={})