From a93579970c3f4eae962f3e31e6fcaeb9d5f2d5f8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Thu, 9 Oct 2025 10:13:34 +0000
Subject: [PATCH] Add `kernels` package/command to generated flake output

This makes it possible to do run `kernels` inside a kernel to e.g.
upload a kernel:

```
nix run .#kernels -- upload ...
```
---
 .github/workflows/test_extra_commands.yaml    | 27 +++++
 lib/gen-flake-outputs.nix                     | 98 ++++++++++++-------
 overlay.nix                                   |  9 ++
 .../kernel-abi-check/default.nix              | 48 +++++++++
 tests/run-kernels/pyproject.toml              |  2 +
 5 files changed, 147 insertions(+), 37 deletions(-)
 create mode 100644 .github/workflows/test_extra_commands.yaml
 create mode 100644 pkgs/python-modules/kernel-abi-check/default.nix
 create mode 100644 tests/run-kernels/pyproject.toml

diff --git a/.github/workflows/test_extra_commands.yaml b/.github/workflows/test_extra_commands.yaml
new file mode 100644
index 00000000..7f26ccc4
--- /dev/null
+++ b/.github/workflows/test_extra_commands.yaml
@@ -0,0 +1,27 @@
+name: "Test kernel flake commands"
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+    types: [opened, synchronize, reopened] # trigger on PRs
+  workflow_dispatch:
+
+jobs:
+  build:
+    name: Build kernel
+    runs-on:
+      group: aws-g6-12xlarge-plus
+    steps:
+      - uses: actions/checkout@v4
+      - uses: cachix/install-nix-action@v27
+        with:
+          nix_path: nixpkgs=channel:nixos-unstable
+      - uses: cachix/cachix-action@v14
+        with:
+          name: huggingface
+          #authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
+        env:
+          USER: github_runner
+      - name: Test nix run .#kernels
+        run: ( cd examples/relu ; nix run .#kernels -- lock ../../tests/run-kernels )
diff --git a/lib/gen-flake-outputs.nix b/lib/gen-flake-outputs.nix
index 121dabed..58b8c1be 100644
--- a/lib/gen-flake-outputs.nix
+++ b/lib/gen-flake-outputs.nix
@@ -46,21 +46,29 @@ let
 
   # Enrich the build configs with generic attributes for framework
   # order/version. Also make bundleBuild attr explicit.
-  buildConfigs = map (
+  buildSets = map (
     set:
     let
       inherit (set) buildConfig;
     in
-    buildConfig
+    set
     // {
-      bundleBuild = buildConfig.bundleBuild or false;
-      frameworkOrder = if buildConfig ? cudaVersion then 0 else 1;
-      frameworkVersion =
-        buildConfig.cudaVersion or buildConfig.rocmVersion or buildConfig.xpuVersion or "0.0";
+      buildConfig =
+
+        buildConfig // {
+          bundleBuild = buildConfig.bundleBuild or false;
+          frameworkOrder = if buildConfig ? cudaVersion then 0 else 1;
+          frameworkVersion =
+            buildConfig.cudaVersion or buildConfig.rocmVersion or buildConfig.xpuVersion or "0.0";
+        };
     }
   ) (build.applicableBuildSets path);
   configCompare =
-    a: b:
+    setA: setB:
+    let
+      a = setA.buildConfig;
+      b = setB.buildConfig;
+    in
     if a.bundleBuild != b.bundleBuild then
       a.bundleBuild
     else if a.frameworkOrder != b.frameworkOrder then
@@ -69,14 +77,14 @@ let
       builtins.compareVersions a.torchVersion b.torchVersion > 0
     else
       builtins.compareVersions a.frameworkVersion b.frameworkVersion < 0;
-  buildConfigsSorted = lib.sort configCompare buildConfigs;
-  shellTorch =
-    if buildConfigsSorted == [ ] then
+  buildSetsSorted = lib.sort configCompare buildSets;
+  bestBuildSet =
+    if buildSetsSorted == [ ] then
       throw "No build variant is compatible with this system"
     else
-      buildName (builtins.head buildConfigsSorted);
+      builtins.head buildSetsSorted;
+  shellTorch = buildName bestBuildSet.buildConfig;
 in
-
 {
   devShells = rec {
     default = devShells.${shellTorch};
@@ -100,37 +108,53 @@ in
       rev = revUnderscored;
     };
   };
-  packages = rec {
-    default = bundle;
+  packages =
+    let
+      bundle = build.buildTorchExtensionBundle {
+        inherit path doGetKernelCheck;
+        rev = revUnderscored;
+      };
+    in
+    {
+      inherit bundle;
 
-    build-and-copy = writeScriptBin "build-and-copy" ''
-      #!/usr/bin/env bash
-      set -euo pipefail
+      default = bundle;
 
-      if [ ! -d build ]; then
-        mkdir build
-      fi
+      build-and-copy = writeScriptBin "build-and-copy" ''
+        #!/usr/bin/env bash
+        set -euo pipefail
 
-      for build_variant in ${bundle}/*; do
-        build_variant=$(basename $build_variant)
-        if [ -e build/$build_variant ]; then
-          rm -rf build/$build_variant
+        if [ ! -d build ]; then
+          mkdir build
         fi
 
-        cp -r ${bundle}/$build_variant build/
-      done
+        for build_variant in ${bundle}/*; do
+          build_variant=$(basename $build_variant)
+          if [ -e build/$build_variant ]; then
+            rm -rf build/$build_variant
+          fi
 
-      chmod -R +w build
-    '';
+          cp -r ${bundle}/$build_variant build/
+        done
 
-    bundle = build.buildTorchExtensionBundle {
-      inherit path doGetKernelCheck;
-      rev = revUnderscored;
-    };
-    redistributable = build.buildDistTorchExtensions {
-      inherit path doGetKernelCheck;
-      bundleOnly = false;
-      rev = revUnderscored;
+        chmod -R +w build
+      '';
+
+      kernels =
+        bestBuildSet.pkgs.python3.withPackages (
+          ps: with ps; [
+            kernel-abi-check
+            kernels
+          ]
+        )
+        // {
+          meta.mainProgram = "kernels";
+        };
+
+      redistributable = build.buildDistTorchExtensions {
+        inherit path doGetKernelCheck;
+        bundleOnly = false;
+        rev = revUnderscored;
+      };
     };
-  };
 }
diff --git a/overlay.nix b/overlay.nix
index 34e69680..4b40f978 100644
--- a/overlay.nix
+++ b/overlay.nix
@@ -12,4 +12,13 @@ final: prev: {
   rewrite-nix-paths-macho = prev.callPackage ./pkgs/rewrite-nix-paths-macho { };
 
   stdenvGlibc_2_27 = prev.callPackage ./pkgs/stdenv-glibc-2_27 { };
+
+  # Python packages
+  pythonPackagesExtensions = prev.pythonPackagesExtensions ++ [
+    (
+      python-self: python-super: with python-self; {
+        kernel-abi-check = callPackage ./pkgs/python-modules/kernel-abi-check { };
+      }
+    )
+  ];
 }
diff --git a/pkgs/python-modules/kernel-abi-check/default.nix b/pkgs/python-modules/kernel-abi-check/default.nix
new file mode 100644
index 00000000..8f30f7fa
--- /dev/null
+++ b/pkgs/python-modules/kernel-abi-check/default.nix
@@ -0,0 +1,48 @@
+{
+  lib,
+  buildPythonPackage,
+  rustPlatform,
+}:
+
+let
+  version =
+    (builtins.fromTOML (builtins.readFile ../../../kernel-abi-check/kernel-abi-check/Cargo.toml))
+    .package.version;
+in
+buildPythonPackage {
+  pname = "kernel-abi-check";
+  inherit version;
+  format = "pyproject";
+
+  src =
+    let
+      sourceFiles =
+        file:
+        file.name == "Cargo.toml"
+        || file.name == "Cargo.lock"
+        || file.name == "manylinux-policy.json"
+        || file.hasExt "pyi"
+        || file.name == "pyproject.toml"
+        || file.hasExt "rs"
+        || file.name == "stable_abi.toml";
+    in
+    lib.fileset.toSource {
+      root = ../../../kernel-abi-check;
+      fileset = lib.fileset.fileFilter sourceFiles ../../../kernel-abi-check;
+    };
+
+  cargoDeps = rustPlatform.importCargoLock {
+    lockFile = ../../../kernel-abi-check/bindings/python/Cargo.lock;
+  };
+
+  sourceRoot = "source/bindings/python";
+
+  build-system = [
+    rustPlatform.cargoSetupHook
+    rustPlatform.maturinBuildHook
+  ];
+
+  meta = with lib; {
+    description = "Check ABI compliance of Hugging Face Hub kernels";
+  };
+}
diff --git a/tests/run-kernels/pyproject.toml b/tests/run-kernels/pyproject.toml
new file mode 100644
index 00000000..92a0bc11
--- /dev/null
+++ b/tests/run-kernels/pyproject.toml
@@ -0,0 +1,2 @@
+[tool.kernels.dependencies]
+"kernels-test/versions" = ">=0.1.0,<0.2.0"