Validate kernel Python dependencies (#182)

* Validate kernel Python dependencies

* Fix nit: deps will never be None

Author: danieldk

.github/workflows/test.yml

@@ -54,6 +54,11 @@ jobs:
         run: |
           uv run pytest tests
 
+      - name: Run dependency test with dependency installed
+        run: |
+          uv pip install nvidia-cutlass-dsl
+          uv run pytest tests/test_deps.py
+
       - name: Run staging tests
         env:
           HF_TOKEN: ${{ secrets.HF_STAGING_TOKEN }}

src/kernels/deps.py

@@ -0,0 +1,27 @@
+import importlib.util
+from typing import List, Set
+
+allowed_dependencies: Set[str] = {
+    "einops",
+    "nvidia-cutlass-dsl",
+}
+
+
+def validate_dependencies(dependencies: List[str]):
+    """
+    Validate a list of dependencies to ensure they are installed.
+
+    Args:
+        dependencies (`List[str]`): A list of dependency strings.
+    """
+    for dependency in dependencies:
+        if dependency not in allowed_dependencies:
+            allowed = ", ".join(sorted(allowed_dependencies))
+            raise ValueError(
+                f"Invalid dependency: {dependency}, allowed dependencies: {allowed}"
+            )
+
+        if importlib.util.find_spec(dependency.replace("-", "_")) is None:
+            raise ImportError(
+                f"Kernel requires dependency `{dependency}`. Please install with: pip install {dependency}"
+            )

src/kernels/utils.py

@@ -19,6 +19,7 @@
 from kernels._system import glibc_version
 from kernels._versions import select_revision_or_version
 from kernels.lockfile import KernelLock, VariantLock
+from kernels.deps import validate_dependencies
 
 ENV_VARS_TRUE_VALUES = {"1", "ON", "YES", "TRUE"}
 
@@ -89,6 +90,13 @@ def universal_build_variant() -> str:
 
 
 def _import_from_path(module_name: str, variant_path: Path) -> ModuleType:
+    metadata_path = variant_path / "metadata.json"
+    if metadata_path.exists():
+        with open(metadata_path, "r") as f:
+            metadata = json.load(f)
+            deps = metadata.get("python-depends", [])
+            validate_dependencies(deps)
+
     file_path = variant_path / "__init__.py"
     if not file_path.exists():
         file_path = variant_path / module_name / "__init__.py"

tests/test_deps.py

@@ -0,0 +1,21 @@
+from importlib.util import find_spec
+
+import pytest
+
+from kernels import get_kernel
+
+
+def test_python_deps():
+    must_raise = find_spec("nvidia_cutlass_dsl") is None
+    if must_raise:
+        with pytest.raises(
+            ImportError, match=r"Kernel requires dependency `nvidia-cutlass-dsl`"
+        ):
+            get_kernel("kernels-test/python-dep")
+    else:
+        get_kernel("kernels-test/python-dep")
+
+
+def test_illegal_dep():
+    with pytest.raises(ValueError, match=r"Invalid dependency: kepler-22b"):
+        get_kernel("kernels-test/python-invalid-dep")