diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 8bb14b1..01fcbe5 100644 --- a/.github/workflows/test.yml +++ b/.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 }} diff --git a/src/kernels/deps.py b/src/kernels/deps.py new file mode 100644 index 0000000..32b6a57 --- /dev/null +++ b/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}" + ) diff --git a/src/kernels/utils.py b/src/kernels/utils.py index e35d8e0..bfa3f55 100644 --- a/src/kernels/utils.py +++ b/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" diff --git a/tests/test_deps.py b/tests/test_deps.py new file mode 100644 index 0000000..63c9c91 --- /dev/null +++ b/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")