Add the kernels check subcommand (#158)

* Add the `kernels check` subcommand

This subcommand checks a given kernel. Currently it applies the same ABI
checks as `kernel-abi-check` in `kernel-builder`.

* Print an error when `build` contains files

* Forgot to update has_issues in two places

Author: danieldk

.github/workflows/test.yml

@@ -73,6 +73,11 @@ jobs:
         run: |
           uv run kernels generate-readme kernels-community/triton-layer-norm
 
+      - name: Check kernel check
+        run: |
+          uv pip install kernel-abi-check
+          kernels check kernels-community/activation
+
       - name: Import check without torch
         run: |
           uv pip uninstall torch

docs/source/cli.md

@@ -2,6 +2,24 @@
 
 ## Main Functions
 
+### kernels check
+
+You can use `kernels check` to test compliance of a kernel on the Hub.
+This currently checks that the kernel:
+
+- Supports the currently-required Python ABI version.
+- Works on supported operating system versions.
+
+For example:
+
+```bash
+$ kernels check kernels-community/flash-attn3
+Checking variant: torch28-cxx11-cu128-aarch64-linux
+  🐍 Python ABI 3.9 compatible
+  🐧 manylinux_2_28 compatible
+[...]
+```
+
 ### kernels to-wheel
 
 We strongly recommend downloading kernels from the Hub using the `kernels`
@@ -38,4 +56,3 @@ your kernel builds to the Hub.
 - If a repo with the `repo_id` already exists and if it contains a `build` with the build variant
   being uploaded, it will attempt to delete the files existing under it.
 - Make sure to be authenticated (run `hf auth login` if not) to be able to perform uploads to the Hub.
-

flake.nix

@@ -24,6 +24,7 @@
       in
       {
         formatter = pkgs.nixfmt-tree;
+        packages.kernel-abi-check = pkgs.python3.pkgs.callPackage ./nix/kernel-abi-check.nix {};
         devShells = with pkgs; rec {
           default = mkShell {
             nativeBuildInputs = [
@@ -40,6 +41,7 @@
               ++ (with python3.pkgs; [
                 docutils
                 huggingface-hub
+                (callPackage ./nix/kernel-abi-check.nix {})
                 mktestdocs
                 pytest
                 pytest-benchmark

nix/kernel-abi-check.nix

@@ -0,0 +1,27 @@
+{
+  buildPythonPackage,
+  fetchPypi,
+  rustPlatform,
+}:
+
+buildPythonPackage rec {
+  pname = "kernel-abi-check";
+  version = "0.6.2";
+
+  src = fetchPypi {
+    inherit version;
+    pname = "kernel_abi_check";
+    hash = "sha256-goWC7SK79FVNEvkp3bISBwbOqdSrmobANtrWIve9/Ys=";
+  };
+
+  cargoDeps = rustPlatform.fetchCargoVendor {
+    inherit pname version src sourceRoot;
+    hash = "sha256-+1jdbKsDKmG+bf0NEVYMv8t7Meuge1z2cgYfbdB9q8A=";
+  };
+
+  sourceRoot = "kernel_abi_check-${version}/bindings/python";
+
+  pyproject = true;
+
+  nativeBuildInputs = with rustPlatform; [ cargoSetupHook maturinBuildHook ];
+}

pyproject.toml

@@ -34,6 +34,7 @@ dev = [
 ]
 
 [project.optional-dependencies]
+abi-check = ["kernel-abi-check>=0.6.2,<0.7.0"]
 torch = ["torch"]
 docs = [
   "hf-doc-builder",

src/kernels/check.py

@@ -0,0 +1,141 @@
+from pathlib import Path
+import sys
+
+from huggingface_hub import snapshot_download
+from kernels.utils import CACHE_DIR
+from kernel_abi_check import (
+    BinaryFormat,
+    IncompatibleMacOSVersion,
+    ObjectFile,
+    IncompatibleAbi3Symbol,
+    NonAbi3Symbol,
+    IncompatibleManylinuxSymbol,
+    MissingMacOSVersion,
+)
+
+
+def check_kernel(
+    *, macos: str, manylinux: str, python_abi: str, repo_id: str, revision: str
+):
+    variants_path = (
+        Path(
+            snapshot_download(
+                repo_id,
+                allow_patterns=["build/*"],
+                cache_dir=CACHE_DIR,
+                revision=revision,
+            )
+        )
+        / "build"
+    )
+
+    has_issues = False
+    for variant_path in variants_path.iterdir():
+        if not variant_path.is_dir():
+            print(
+                f"⛔ `build/` must only contain directories, found: {variant_path.name}",
+                file=sys.stderr,
+            )
+            has_issues = True
+            continue
+
+        print(f"Checking variant: {variant_path.name}", file=sys.stderr)
+
+        indent = 2
+
+        for dylib_path in variant_path.rglob("*.so"):
+            print_with_indent(
+                indent,
+                f"Dynamic library {dylib_path.relative_to(variant_path)}:",
+            )
+
+            o = ObjectFile(dylib_path)
+            has_issues |= check_abi3(o, python_abi, indent + 2)
+
+            # TODO: also check operating system
+            if o.format() == BinaryFormat.ELF:
+                has_issues |= check_manylinux(o, manylinux, indent + 2)
+            elif o.format() == BinaryFormat.MACH_O:
+                has_issues |= check_macos(o, macos, indent + 2)
+
+    if has_issues:
+        sys.exit(1)
+
+
+def check_abi3(object_file: ObjectFile, python_abi: str, indent: int) -> bool:
+    has_issues = False
+    violations = object_file.check_python_abi(python_abi)
+    if violations != []:
+        has_issues = True
+        print_with_indent(
+            indent,
+            f"⛔ Found symbols that are incompatible with Python ABI {python_abi}:",
+        )
+        for violation in violations:
+            if isinstance(violation, IncompatibleAbi3Symbol):
+                print_with_indent(
+                    indent + 3,
+                    f"{violation.name}: {violation.version_added}",
+                )
+            elif isinstance(violation, NonAbi3Symbol):
+                print_with_indent(
+                    indent + 3,
+                    f"{violation.name}",
+                )
+    else:
+        print_with_indent(indent, f"🐍 Python ABI {python_abi} compatible")
+
+    return has_issues
+
+
+def check_macos(object_file: ObjectFile, macos: str, indent: int) -> bool:
+    has_issues = False
+    violations = object_file.check_macos(macos)
+    if violations != []:
+        has_issues = True
+        print_with_indent(
+            indent,
+            f"⛔ Found incompatibility with macOS {macos}:",
+        )
+
+        for violation in violations:
+            if isinstance(violation, MissingMacOSVersion):
+                print_with_indent(
+                    indent + 3,
+                    "shared library does not contain macOS version",
+                )
+            elif isinstance(violation, IncompatibleMacOSVersion):
+                print_with_indent(
+                    indent + 3,
+                    f"shared library requires macOS {violation.version}",
+                )
+    else:
+        print_with_indent(indent, f"🍏 compatible with macOS {macos}")
+
+    return has_issues
+
+
+def check_manylinux(object_file: ObjectFile, manylinux: str, indent: int) -> bool:
+    has_issues = False
+    violations = object_file.check_manylinux(manylinux)
+    if violations != []:
+        has_issues = True
+        print_with_indent(
+            indent,
+            f"⛔ Found symbols that are incompatible with {manylinux}:",
+        )
+
+        for violation in violations:
+            if isinstance(violation, IncompatibleManylinuxSymbol):
+                print_with_indent(
+                    indent + 3,
+                    f"{violation.name}_{violation.dep}: {violation.version}",
+                )
+    else:
+        print_with_indent(indent, f"🐧 {manylinux} compatible")
+
+    return has_issues
+
+
+def print_with_indent(indent: int, message: str):
+    print(f"{' ' * indent}{message}", file=sys.stderr)

src/kernels/cli.py

@@ -20,6 +20,31 @@ def main():
     )
     subparsers = parser.add_subparsers(required=True)
 
+    check_parser = subparsers.add_parser("check", help="Check a kernel for compliance")
+    check_parser.add_argument("repo_id", type=str, help="The kernel repo ID")
+    check_parser.add_argument(
+        "--revision",
+        type=str,
+        default="main",
+        help="The kernel revision (branch, tag, or commit SHA, defaults to 'main')",
+    )
+    check_parser.add_argument("--macos", type=str, help="macOS version", default="15.0")
+    check_parser.add_argument(
+        "--manylinux", type=str, help="Manylinux version", default="manylinux_2_28"
+    )
+    check_parser.add_argument(
+        "--python-abi", type=str, help="Python ABI version", default="3.9"
+    )
+    check_parser.set_defaults(
+        func=lambda args: check_kernel(
+            macos=args.macos,
+            manylinux=args.manylinux,
+            python_abi=args.python_abi,
+            repo_id=args.repo_id,
+            revision=args.revision,
+        )
+    )
+
     download_parser = subparsers.add_parser("download", help="Download locked kernels")
     download_parser.add_argument(
         "project_dir",
@@ -205,3 +230,24 @@ def default(self, o):
         if dataclasses.is_dataclass(o):
             return dataclasses.asdict(o)
         return super().default(o)
+
+
+def check_kernel(
+    *, macos: str, manylinux: str, python_abi: str, repo_id: str, revision: str
+):
+    try:
+        import kernels.check
+    except ImportError:
+        print(
+            "`kernels check` requires the `kernel-abi-check` package: pip install kernel-abi-check",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+    kernels.check.check_kernel(
+        macos=macos,
+        manylinux=manylinux,
+        python_abi=python_abi,
+        repo_id=repo_id,
+        revision=revision,
+    )