diff --git a/.github/workflows/test_kernels.yaml b/.github/workflows/test_kernels.yaml index de18b32ce..ca5e8f2a4 100644 --- a/.github/workflows/test_kernels.yaml +++ b/.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 diff --git a/kernels/src/kernels/__init__.py b/kernels/src/kernels/__init__.py index 3dae0dc9a..cc6c31123 100644 --- a/kernels/src/kernels/__init__.py +++ b/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", diff --git a/kernels/src/kernels/deps.py b/kernels/src/kernels/deps.py index 4723f418d..3e0b3c6e0 100644 --- a/kernels/src/kernels/deps.py +++ b/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 diff --git a/kernels/src/kernels/importer.py b/kernels/src/kernels/importer.py index 3df5807e8..a62e74bd5 100644 --- a/kernels/src/kernels/importer.py +++ b/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 diff --git a/kernels/src/kernels/install.py b/kernels/src/kernels/install.py index 08d598b91..16f358e7f 100644 --- a/kernels/src/kernels/install.py +++ b/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, diff --git a/kernels/src/kernels/load.py b/kernels/src/kernels/load.py index 2a6a40766..28a13df92 100644 --- a/kernels/src/kernels/load.py +++ b/kernels/src/kernels/load.py @@ -6,7 +6,6 @@ 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 @@ -14,6 +13,7 @@ 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 `." ) 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={}) diff --git a/kernels/src/kernels/python_deps.py b/kernels/src/kernels/python_deps.py new file mode 100644 index 000000000..4723f418d --- /dev/null +++ b/kernels/src/kernels/python_deps.py @@ -0,0 +1,106 @@ +import importlib.util +import json +from dataclasses import dataclass, field +from pathlib import Path + +from huggingface_hub.dataclasses import strict +from kernels_data import Metadata + +from kernels.backends import Backend, _backend + + +@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"), + ) + + +@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", [])], + ) + + +@strict +@dataclass +class DependencyData: + general: dict = field(default_factory=dict) + backends: dict = field(default_factory=dict) + + @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) + + +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?") + + +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 validate_dependencies(kernel_module_name: str, dependencies: list[str], backend: Backend): + """ + Validate a list of dependencies to ensure they are installed. + + Args: + dependencies (`list[str]`): A list of dependency strings to validate. + backend (`str`): The backend to validate dependencies for. + """ + 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}" + ) diff --git a/kernels/src/kernels/resolver.py b/kernels/src/kernels/resolver.py index 67cd32742..0c6227d60 100644 --- a/kernels/src/kernels/resolver.py +++ b/kernels/src/kernels/resolver.py @@ -1,6 +1,6 @@ from dataclasses import dataclass from pathlib import Path -from typing import Protocol, Self, runtime_checkable +from typing import Protocol, runtime_checkable from huggingface_hub.errors import LocalEntryNotFoundError from huggingface_hub.hf_api import HfApi @@ -36,7 +36,7 @@ class LocalKernel: metadata: Metadata origin: "RemoteKernel | None" = None - def install(self, *, api: HfApi, all_variants: bool = False) -> Self: + def install(self, *, api: HfApi) -> "LocalKernel": # Local kernels are already installed, so we just return self. return self @@ -50,38 +50,22 @@ class RemoteKernel: metadata: Metadata variant: Variant - def install(self, *, api: HfApi, all_variants: bool = False) -> LocalKernel: - + def install(self, *, api: HfApi) -> 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, - ) + 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, diff --git a/kernels/tests/test_deps.py b/kernels/tests/test_deps.py index 3eaa7bce2..5b7a7de72 100644 --- a/kernels/tests/test_deps.py +++ b/kernels/tests/test_deps.py @@ -1,69 +1,161 @@ -from importlib.util import find_spec +import json +from dataclasses import dataclass from pathlib import Path +from types import ModuleType import pytest -from huggingface_hub import HfApi, snapshot_download +from kernels_data import KernelDependency, KernelVersion, Metadata -from kernels import get_kernel, get_local_kernel, install_kernel +from kernels.deps import get_kernel_dep, resolve_kernel_tree, use_kernel_deps +from kernels.resolver import LocalKernel -@pytest.mark.cuda_only -@pytest.mark.parametrize("dependency", ["einops", "nvidia-cutlass-dsl"]) -def test_python_deps(dependency): - must_raise = find_spec(dependency.replace("-", "_")) is None - if must_raise: - with pytest.raises( - ImportError, - match=r"Kernel module `python_dep` requires Python dependency `(einops|nvidia-cutlass-dsl)`", - ): - get_kernel("kernels-test/python-dep", revision="main") - else: - get_kernel("kernels-test/python-dep", revision="main") +def _module(name: str) -> ModuleType: + return ModuleType(name) -def test_illegal_dep(): - with pytest.raises(ValueError, match=r"Kernel module `python_invalid_dep` uses.*kepler-22b"): - get_kernel("kernels-test/python-invalid-dep", revision="main") +def _dep(repo_id: str = "test/kernel", version: int = 1) -> KernelDependency: + return KernelDependency(repo_id=repo_id, version=KernelVersion.Version(version)) -def test_deps_validated_before_download(monkeypatch): - """With `validate_dependencies=True`, deps are checked *before* the build - variant is downloaded. +def _metadata(*deps: KernelDependency) -> Metadata: + return Metadata.from_bytes( + json.dumps( + { + "id": "test_1_cpu", + "name": "test", + "version": 1, + "license": "Apache-2.0", + "python-depends": ["torch"], + "kernel-depends": [{"repo-id": dep.repo_id, "version": dep.version.version} for dep in deps], + "backend": {"type": "cpu"}, + } + ).encode() + ) + + +def _local_kernel(repo_id: str, *deps: KernelDependency) -> LocalKernel: + return LocalKernel(variant_path=Path("/repo") / repo_id, metadata=_metadata(*deps)) + + +@dataclass +class _MapResolver: + """Test resolver that resolves kernels from a repo ID mapping.""" + + kernels: dict[str, LocalKernel] + + def resolve(self, *, api, backend, kernel) -> LocalKernel | None: + return self.kernels.get(kernel.repo_id) + - `snapshot_download` (the full-variant download) is patched to fail, so the - dependency error can only surface if validation runs first — guarding the - early-bail ordering against regressions that move validation back after the - download. - """ +def test_resolve_kernel_tree_no_dependencies(): + kernel = _local_kernel("test/kernel") + resolver = _MapResolver({"test/kernel": kernel}) - class _Downloaded(RuntimeError): - pass + tree = resolve_kernel_tree(api=None, backend=None, kernel=_dep(), resolver=resolver) - def _fail(*_args, **_kwargs): - raise _Downloaded("build variant was downloaded before dependency validation") + assert tree.location is kernel + assert tree.deps == {} - monkeypatch.setattr(HfApi, "snapshot_download", _fail) - with pytest.raises(ValueError, match=r"Kernel module `python_invalid_dep` uses.*kepler-22b"): - install_kernel("kernels-test/python-invalid-dep", revision="main", validate_dependencies=True) +def test_resolve_kernel_tree_builds_tree(): + a, b, c = _dep("test/a"), _dep("test/b"), _dep("test/c") + kernel_a = _local_kernel("test/a", b) + kernel_b = _local_kernel("test/b", c) + kernel_c = _local_kernel("test/c") + resolver = _MapResolver({"test/a": kernel_a, "test/b": kernel_b, "test/c": kernel_c}) + tree = resolve_kernel_tree(api=None, backend=None, kernel=a, resolver=resolver) -def test_install_kernel_skips_validation_by_default(): - """Validation is opt-in: with the default `validate_dependencies=False`, - `install_kernel` downloads an invalid-dep kernel without raising.""" - variant_path = install_kernel("kernels-test/python-invalid-dep", revision="main") - assert (variant_path / "metadata.json").exists() + assert tree.location is kernel_a + assert set(tree.deps) == {"test/b"} + node_b = tree.deps["test/b"] + assert node_b.location is kernel_b + assert set(node_b.deps) == {"test/c"} + node_c = node_b.deps["test/c"] + assert node_c.location is kernel_c + assert node_c.deps == {} -def test_local_kernel_validates_deps(tmp_path): - """`get_local_kernel` validates dependencies even though it never goes through - `install_kernel`'s pre-download check, so removing the gate from - `_import_from_path` must not drop validation for local kernels.""" - repo_path = snapshot_download( - "kernels-test/python-invalid-dep", - repo_type="kernel", - revision="main", - cache_dir=tmp_path, +@pytest.mark.parametrize( + "kernel, message", + [ + (_dep("test/kernel", version=1), "version: 1"), + (KernelDependency(repo_id="test/kernel", version=KernelVersion.Revision("main")), "revision: main"), + ], +) +@pytest.mark.parametrize("resolver", [None, _MapResolver({})]) +def test_resolve_kernel_tree_unresolvable(kernel, resolver, message): + with pytest.raises(ValueError, match="Could not resolve kernel") as exc_info: + resolve_kernel_tree(api=None, backend=None, kernel=kernel, resolver=resolver) + + assert "test/kernel" in str(exc_info.value) + assert message in str(exc_info.value) + + +def test_resolve_kernel_tree_detects_cycle(): + a, b = _dep("test/a"), _dep("test/b") + resolver = _MapResolver( + { + "test/a": _local_kernel("test/a", b), + "test/b": _local_kernel("test/b", a), + } + ) + + with pytest.raises(ValueError, match="Cyclic kernel dependency detected: test/a"): + resolve_kernel_tree(api=None, backend=None, kernel=a, resolver=resolver) + + +def test_resolve_kernel_tree_diamond_is_not_a_cycle(): + a, b, c, d = (_dep(f"test/{name}") for name in "abcd") + kernel_d = _local_kernel("test/d") + resolver = _MapResolver( + { + "test/a": _local_kernel("test/a", b, c), + "test/b": _local_kernel("test/b", d), + "test/c": _local_kernel("test/c", d), + "test/d": kernel_d, + } ) - with pytest.raises(ValueError, match=r"Kernel module `python_invalid_dep` uses.*kepler-22b"): - get_local_kernel(Path(repo_path)) + + tree = resolve_kernel_tree(api=None, backend=None, kernel=a, resolver=resolver) + + assert set(tree.deps) == {"test/b", "test/c"} + assert tree.deps["test/b"].deps["test/d"].location is kernel_d + assert tree.deps["test/c"].deps["test/d"].location is kernel_d + + +def test_get_kernel_dep_outside_loading_context(): + with pytest.raises(RuntimeError, match="only works during kernel loading"): + get_kernel_dep("test/kernel") + + +def test_get_kernel_dep_returns_module(): + module = _module("kernel") + + with use_kernel_deps({"test/kernel": module}): + assert get_kernel_dep("test/kernel") is module + + +def test_get_kernel_dep_unknown_repo_id(): + with use_kernel_deps({"test/other": _module("other")}): + with pytest.raises(RuntimeError, match="not found in kernel dependencies") as exc_info: + get_kernel_dep("test/kernel") + + assert "test/kernel" in str(exc_info.value) + + +def test_use_kernel_deps_restores_previous_state(): + outer_module = _module("outer") + inner_module = _module("inner") + + with use_kernel_deps({"test/kernel": outer_module}): + assert get_kernel_dep("test/kernel") is outer_module + + with use_kernel_deps({"test/kernel": inner_module}): + assert get_kernel_dep("test/kernel") is inner_module + + assert get_kernel_dep("test/kernel") is outer_module + + with pytest.raises(RuntimeError, match="only works during kernel loading"): + get_kernel_dep("test/kernel") diff --git a/kernels/tests/test_dirty_provenance.py b/kernels/tests/test_dirty_provenance.py index b9119d63d..a5bf17aa8 100644 --- a/kernels/tests/test_dirty_provenance.py +++ b/kernels/tests/test_dirty_provenance.py @@ -65,7 +65,7 @@ def test_import_from_path_warns_on_dirty(tmp_path): _loaded_kernels.pop(variant_dir, None) try: with pytest.warns(UserWarning, match="dirty git tree"): - module = _import_from_path(variant_dir) + module = _import_from_path(variant_dir, deps={}) assert module.value == 42 finally: _loaded_kernels.pop(variant_dir, None) diff --git a/kernels/tests/test_python_deps.py b/kernels/tests/test_python_deps.py new file mode 100644 index 000000000..854815819 --- /dev/null +++ b/kernels/tests/test_python_deps.py @@ -0,0 +1,72 @@ +from importlib.util import find_spec +from pathlib import Path + +import pytest +from huggingface_hub import HfApi, snapshot_download + +from kernels import get_kernel, get_local_kernel, install_kernel + + +@pytest.mark.cuda_only +@pytest.mark.parametrize("dependency", ["einops", "nvidia-cutlass-dsl"]) +def test_python_deps(dependency): + must_raise = find_spec(dependency.replace("-", "_")) is None + if must_raise: + with pytest.raises( + ImportError, + match=r"Kernel module `python_dep` requires Python dependency `(einops|nvidia-cutlass-dsl)`", + ): + get_kernel("kernels-test/python-dep", revision="main") + else: + get_kernel("kernels-test/python-dep", revision="main") + + +def test_illegal_dep(): + with pytest.raises(ValueError, match=r"Kernel module `python_invalid_dep` uses.*kepler-22b"): + get_kernel("kernels-test/python-invalid-dep", revision="main") + + +def test_deps_validated_before_download(monkeypatch): + """With `validate_dependencies=True`, deps are checked *before* the build + variant is downloaded. + + `snapshot_download` (the full-variant download) is patched to fail, so the + dependency error can only surface if validation runs first — guarding the + early-bail ordering against regressions that move validation back after the + download. + """ + + class _Downloaded(RuntimeError): + pass + + def _fail(*_args, **_kwargs): + raise _Downloaded("build variant was downloaded before dependency validation") + + monkeypatch.setattr(HfApi, "snapshot_download", _fail) + + with pytest.raises(ValueError, match=r"Kernel module `python_invalid_dep` uses.*kepler-22b"): + get_kernel( + "kernels-test/python-invalid-dep", + revision="main", + ) + + +def test_install_kernel_skips_validation_by_default(): + """Validation is opt-in: with the default `validate_dependencies=False`, + `install_kernel` downloads an invalid-dep kernel without raising.""" + variant_path = install_kernel("kernels-test/python-invalid-dep", revision="main") + assert (variant_path / "metadata.json").exists() + + +def test_local_kernel_validates_deps(tmp_path): + """`get_local_kernel` validates dependencies even though it never goes through + `install_kernel`'s pre-download check, so removing the gate from + `_import_from_path` must not drop validation for local kernels.""" + repo_path = snapshot_download( + "kernels-test/python-invalid-dep", + repo_type="kernel", + revision="main", + cache_dir=tmp_path, + ) + with pytest.raises(ValueError, match=r"Kernel module `python_invalid_dep` uses.*kepler-22b"): + get_local_kernel(Path(repo_path)) diff --git a/kernels/tests/test_resolver.py b/kernels/tests/test_resolver.py index 3ccb3d338..f599e3131 100644 --- a/kernels/tests/test_resolver.py +++ b/kernels/tests/test_resolver.py @@ -1,6 +1,6 @@ import json import re -from dataclasses import dataclass +from dataclasses import FrozenInstanceError, dataclass from pathlib import Path import pytest @@ -31,6 +31,7 @@ SequentialResolver, _locked_revision, ) +from kernels.variants import parse_variant @pytest.fixture(scope="module") @@ -75,6 +76,11 @@ def _local_kernel(variant_path: str) -> LocalKernel: return LocalKernel(variant_path=Path(variant_path), metadata=_metadata()) +def _variant(): + # A noarch variant so the test does not depend on a specific backend. + return parse_variant("torch-cpu") + + 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 @@ -94,6 +100,54 @@ def _write_variant(repo_path: Path, variant: str = "torch-cpu") -> Path: return variant_dir +def test_local_kernel_eq_and_hash(): + m = _metadata() + a = LocalKernel(variant_path=Path("/tmp/a"), metadata=m) + b = LocalKernel(variant_path=Path("/tmp/a"), metadata=m) + c = LocalKernel(variant_path=Path("/tmp/b"), metadata=m) + + assert a == b + assert hash(a) == hash(b) + assert a != c + + assert {a, b, c} == {a, c} + assert {a: 1}[b] == 1 + + +def test_remote_kernel_eq_and_hash(): + v = _variant() + m = _metadata() + a = RemoteKernel(repo_id="foo/bar", revision="deadbeef", variant=v, metadata=m) + b = RemoteKernel(repo_id="foo/bar", revision="deadbeef", variant=v, metadata=m) + c = RemoteKernel(repo_id="foo/bar", revision="cafef00d", variant=v, metadata=m) + d = RemoteKernel(repo_id="other/repo", revision="deadbeef", variant=v, metadata=m) + + assert a == b + assert hash(a) == hash(b) + assert a != c + assert a != d + + assert {a, b, c, d} == {a, c, d} + assert {a: 1}[b] == 1 + + +def test_local_kernel_is_immutable(): + kernel = LocalKernel(variant_path=Path("/tmp/a"), metadata=_metadata()) + with pytest.raises(FrozenInstanceError): + kernel.variant_path = Path("/tmp/b") + + +def test_remote_kernel_is_immutable(): + v = _variant() + kernel = RemoteKernel(repo_id="foo/bar", revision="deadbeef", variant=v, metadata=_metadata()) + with pytest.raises(FrozenInstanceError): + kernel.repo_id = "other/repo" + with pytest.raises(FrozenInstanceError): + kernel.revision = "cafef00d" + with pytest.raises(FrozenInstanceError): + kernel.variant = parse_variant("torch-cuda") + + def test_noop_resolver_never_resolves(api): assert NoopResolver().resolve(api=api, backend=None, kernel=_dep()) is None