Add doGetKernelCheck option to genFlakeOutputs (#188)

This is a 'backdoor' for kernels that require a GPU during import (which
is not available in the build sandbox).

Author: danieldk

docs/nix.md

@@ -128,6 +128,17 @@ The available packages can be found on [search.nixos.org](https://search.nixos.o
 Keep in mind that these additional dependencies will only be available to
 the Nix shells, not the final kernel uploaded to the Hub.
 
+## Skipping the `get_kernel` check
+
+`kernel-builder` verifies that a kernel can be
+imported with the [`kernels`](https://github.com/huggingface/kernels)
+package. This check can be disabled by passing `doGetKernelCheck = false`
+to `genFlakeOutputs`. **Warning:** it is strongly recommended to keep
+this check enabled, as it is one of the checks that validates that a kernel
+is compliant. This option is primarily intended for kernels with
+`triton.autotune` decorators, which can fail because there is no GPU available
+in the build sandbox.
+
 ## Building a kernel without `flake.nix`
 
 If a kernels source directory does not have a `flake.nix` file, you can build the

flake.nix

@@ -72,6 +72,13 @@
             path,
             rev,
 
+            # This option is not documented on purpose. You should not use it,
+            # if a kernel cannot be imported, it is non-compliant. This is for
+            # one exceptional case: packaging a third-party kernel (where you
+            # want to stay close to upstream) where importing the kernel will
+            # fail in a GPU-less sandbox. Even in that case, it's better to lazily
+            # load the part with this functionality.
+            doGetKernelCheck ? true,
             pythonCheckInputs ? pkgs: [ ],
             pythonNativeCheckInputs ? pkgs: [ ],
             torchVersions ? torchVersions',
@@ -94,22 +101,32 @@
                 default = devShells.${shellTorch};
                 test = testShells.${shellTorch};
                 devShells = build.torchDevShells {
-                  inherit path pythonCheckInputs pythonNativeCheckInputs;
+                  inherit
+                    path
+                    doGetKernelCheck
+                    pythonCheckInputs
+                    pythonNativeCheckInputs
+                    ;
                   rev = revUnderscored;
                 };
                 testShells = build.torchExtensionShells {
-                  inherit path pythonCheckInputs pythonNativeCheckInputs;
+                  inherit
+                    path
+                    doGetKernelCheck
+                    pythonCheckInputs
+                    pythonNativeCheckInputs
+                    ;
                   rev = revUnderscored;
                 };
               };
               packages = rec {
                 default = bundle;
                 bundle = build.buildTorchExtensionBundle {
-                  inherit path;
+                  inherit path doGetKernelCheck;
                   rev = revUnderscored;
                 };
                 redistributable = build.buildDistTorchExtensions {
-                  inherit path;
+                  inherit path doGetKernelCheck;
                   buildSets = buildSetPerSystem'.${system};
                   rev = revUnderscored;
                 };

lib/build.nix

@@ -95,6 +95,7 @@ rec {
     {
       path,
       rev,
+      doGetKernelCheck,
       stripRPath ? false,
       oldLinuxCompat ? false,
     }:
@@ -128,12 +129,18 @@ rec {
     if buildConfig.general.universal then
       # No torch extension sources? Treat it as a noarch package.
       pkgs.callPackage ./torch-extension-noarch ({
-        inherit src rev torch;
+        inherit
+          src
+          rev
+          torch
+          doGetKernelCheck
+          ;
         extensionName = buildConfig.general.name;
       })
     else
       pkgs.callPackage ./torch-extension ({
         inherit
+          doGetKernelCheck
           extraDeps
           nvccThreads
           src
@@ -166,14 +173,15 @@ rec {
       buildSets,
       path,
       rev,
+      doGetKernelCheck,
     }:
     let
       extensionForTorch =
         { path, rev }:
         buildSet: {
           name = torchBuildVersion buildSet;
           value = buildTorchExtension buildSet {
-            inherit path rev;
+            inherit path rev doGetKernelCheck;
             stripRPath = true;
             oldLinuxCompat = true;
           };
@@ -183,13 +191,17 @@ rec {
     builtins.listToAttrs (lib.map (extensionForTorch { inherit path rev; }) filteredBuildSets);
 
   buildTorchExtensionBundle =
-    { path, rev }:
+    {
+      path,
+      rev,
+      doGetKernelCheck,
+    }:
     let
       # We just need to get any nixpkgs for use by the path join.
       pkgs = (builtins.head buildSets).pkgs;
       upstreamBuildSets = builtins.filter (buildSet: buildSet.upstreamVariant) buildSets;
       extensions = buildDistTorchExtensions {
-        inherit path rev;
+        inherit path rev doGetKernelCheck;
         buildSets = upstreamBuildSets;
       };
       buildConfig = readBuildConfig path;
@@ -211,6 +223,7 @@ rec {
     {
       path,
       rev,
+      doGetKernelCheck,
       pythonCheckInputs,
       pythonNativeCheckInputs,
     }:
@@ -237,7 +250,9 @@ rec {
               ]
               ++ (pythonCheckInputs python3.pkgs);
             shellHook = ''
-              export PYTHONPATH=''${PYTHONPATH}:${buildTorchExtension buildSet { inherit path rev; }}
+              export PYTHONPATH=''${PYTHONPATH}:${
+                buildTorchExtension buildSet { inherit path rev doGetKernelCheck; }
+              }
             '';
           };
         };
@@ -249,6 +264,7 @@ rec {
     {
       path,
       rev,
+      doGetKernelCheck,
       pythonCheckInputs,
       pythonNativeCheckInputs,
     }:
@@ -272,7 +288,7 @@ rec {
               ]
               ++ (pythonNativeCheckInputs python3.pkgs);
             buildInputs = with pkgs; [ python3.pkgs.pytest ] ++ (pythonCheckInputs python3.pkgs);
-            inputsFrom = [ (buildTorchExtension buildSet { inherit path rev; }) ];
+            inputsFrom = [ (buildTorchExtension buildSet { inherit path rev doGetKernelCheck; }) ];
             env = lib.optionalAttrs rocmSupport {
               PYTORCH_ROCM_ARCH = lib.concatStringsSep ";" buildSet.torch.rocmArchs;
               HIP_PATH = pkgs.rocmPackages.clr;

lib/torch-extension-noarch/default.nix

@@ -3,6 +3,10 @@
   extensionName,
   rev,
 
+  # Whether to run get-kernel-check.
+  doGetKernelCheck ? true,
+
+  lib,
   build2cmake,
   get-kernel-check,
   torch,
@@ -19,10 +23,13 @@ stdenv.mkDerivation (prevAttrs: {
   # also get torch as a build input.
   buildInputs = [ torch ];
 
-  nativeBuildInputs = [
-    build2cmake
-    get-kernel-check
-  ];
+  nativeBuildInputs =
+    [
+      build2cmake
+    ]
+    ++ lib.optionals doGetKernelCheck [
+      get-kernel-check
+    ];
 
   dontBuild = true;
 

lib/torch-extension/default.nix

@@ -3,6 +3,9 @@
   nvccThreads,
   rev,
 
+  # Whether to run get-kernel-check.
+  doGetKernelCheck ? true,
+
   # Wheter to strip rpath for non-nix use.
   stripRPath ? false,
 
@@ -81,12 +84,14 @@ stdenv.mkDerivation (prevAttrs: {
 
   nativeBuildInputs =
     [
-      get-kernel-check
       kernel-abi-check
       cmake
       ninja
       build2cmake
     ]
+    ++ lib.optionals doGetKernelCheck [
+      get-kernel-check
+    ]
     ++ lib.optionals cudaSupport [
       cmakeNvccThreadsHook
       cudaPackages.cuda_nvcc