From b0c431fee475ecd27d545d5e1c5e2dea372cf1fc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= Date: Thu, 25 Sep 2025 19:05:29 +0200 Subject: [PATCH] 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 --- .github/workflows/test.yml | 5 ++ docs/source/cli.md | 19 ++++- flake.nix | 2 + nix/kernel-abi-check.nix | 27 +++++++ pyproject.toml | 1 + src/kernels/check.py | 141 +++++++++++++++++++++++++++++++++++++ src/kernels/cli.py | 46 ++++++++++++ 7 files changed, 240 insertions(+), 1 deletion(-) create mode 100644 nix/kernel-abi-check.nix create mode 100644 src/kernels/check.py diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 4d39d77..da47cbe 100644 --- a/.github/workflows/test.yml +++ b/.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 diff --git a/docs/source/cli.md b/docs/source/cli.md index 4641439..a6b746b 100644 --- a/docs/source/cli.md +++ b/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. - diff --git a/flake.nix b/flake.nix index 6c1aff0..07c0e83 100644 --- a/flake.nix +++ b/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 diff --git a/nix/kernel-abi-check.nix b/nix/kernel-abi-check.nix new file mode 100644 index 0000000..b8553f8 --- /dev/null +++ b/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 ]; +} diff --git a/pyproject.toml b/pyproject.toml index f1c4651..60eec8b 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/kernels/check.py b/src/kernels/check.py new file mode 100644 index 0000000..1e091f5 --- /dev/null +++ b/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) diff --git a/src/kernels/cli.py b/src/kernels/cli.py index 90f5db5..b1f1be4 100644 --- a/src/kernels/cli.py +++ b/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 @@ class _JSONEncoder(json.JSONEncoder): 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, + )