Skip to content
This repository was archived by the owner on Apr 6, 2026. It is now read-only.

Commit 1774d44

Browse files
danieldkdanieldk
authored
Add kernels package/command to generated flake output (#251)
This makes it possible to do run `kernels` inside a kernel to e.g. upload a kernel: ``` nix run .#kernels -- upload ... ```
1 parent 378524d commit 1774d44

5 files changed

Lines changed: 147 additions & 37 deletions

File tree

‎.github/workflows/test_extra_commands.yaml‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
name: "Test kernel flake commands"
2+
on:
3+
push:
4+
branches: [main]
5+
pull_request:
6+
branches: [main]
7+
types: [opened, synchronize, reopened] # trigger on PRs
8+
workflow_dispatch:
9+
10+
jobs:
11+
build:
12+
name: Build kernel
13+
runs-on:
14+
group: aws-g6-12xlarge-plus
15+
steps:
16+
- uses: actions/checkout@v4
17+
- uses: cachix/install-nix-action@v27
18+
with:
19+
nix_path: nixpkgs=channel:nixos-unstable
20+
- uses: cachix/cachix-action@v14
21+
with:
22+
name: huggingface
23+
#authToken: "${{ secrets.CACHIX_AUTH_TOKEN }}"
24+
env:
25+
USER: github_runner
26+
- name: Test nix run .#kernels
27+
run: ( cd examples/relu ; nix run .#kernels -- lock ../../tests/run-kernels )

‎lib/gen-flake-outputs.nix‎

Lines changed: 61 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -46,21 +46,29 @@ let
4646

4747
# Enrich the build configs with generic attributes for framework
4848
# order/version. Also make bundleBuild attr explicit.
49-
buildConfigs = map (
49+
buildSets = map (
5050
set:
5151
let
5252
inherit (set) buildConfig;
5353
in
54-
buildConfig
54+
set
5555
// {
56-
bundleBuild = buildConfig.bundleBuild or false;
57-
frameworkOrder = if buildConfig ? cudaVersion then 0 else 1;
58-
frameworkVersion =
59-
buildConfig.cudaVersion or buildConfig.rocmVersion or buildConfig.xpuVersion or "0.0";
56+
buildConfig =
57+
58+
buildConfig // {
59+
bundleBuild = buildConfig.bundleBuild or false;
60+
frameworkOrder = if buildConfig ? cudaVersion then 0 else 1;
61+
frameworkVersion =
62+
buildConfig.cudaVersion or buildConfig.rocmVersion or buildConfig.xpuVersion or "0.0";
63+
};
6064
}
6165
) (build.applicableBuildSets path);
6266
configCompare =
63-
a: b:
67+
setA: setB:
68+
let
69+
a = setA.buildConfig;
70+
b = setB.buildConfig;
71+
in
6472
if a.bundleBuild != b.bundleBuild then
6573
a.bundleBuild
6674
else if a.frameworkOrder != b.frameworkOrder then
@@ -69,14 +77,14 @@ let
6977
builtins.compareVersions a.torchVersion b.torchVersion > 0
7078
else
7179
builtins.compareVersions a.frameworkVersion b.frameworkVersion < 0;
72-
buildConfigsSorted = lib.sort configCompare buildConfigs;
73-
shellTorch =
74-
if buildConfigsSorted == [ ] then
80+
buildSetsSorted = lib.sort configCompare buildSets;
81+
bestBuildSet =
82+
if buildSetsSorted == [ ] then
7583
throw "No build variant is compatible with this system"
7684
else
77-
buildName (builtins.head buildConfigsSorted);
85+
builtins.head buildSetsSorted;
86+
shellTorch = buildName bestBuildSet.buildConfig;
7887
in
79-
8088
{
8189
devShells = rec {
8290
default = devShells.${shellTorch};
@@ -100,37 +108,53 @@ in
100108
rev = revUnderscored;
101109
};
102110
};
103-
packages = rec {
104-
default = bundle;
111+
packages =
112+
let
113+
bundle = build.buildTorchExtensionBundle {
114+
inherit path doGetKernelCheck;
115+
rev = revUnderscored;
116+
};
117+
in
118+
{
119+
inherit bundle;
105120

106-
build-and-copy = writeScriptBin "build-and-copy" ''
107-
#!/usr/bin/env bash
108-
set -euo pipefail
121+
default = bundle;
109122

110-
if [ ! -d build ]; then
111-
mkdir build
112-
fi
123+
build-and-copy = writeScriptBin "build-and-copy" ''
124+
#!/usr/bin/env bash
125+
set -euo pipefail
113126
114-
for build_variant in ${bundle}/*; do
115-
build_variant=$(basename $build_variant)
116-
if [ -e build/$build_variant ]; then
117-
rm -rf build/$build_variant
127+
if [ ! -d build ]; then
128+
mkdir build
118129
fi
119130
120-
cp -r ${bundle}/$build_variant build/
121-
done
131+
for build_variant in ${bundle}/*; do
132+
build_variant=$(basename $build_variant)
133+
if [ -e build/$build_variant ]; then
134+
rm -rf build/$build_variant
135+
fi
122136
123-
chmod -R +w build
124-
'';
137+
cp -r ${bundle}/$build_variant build/
138+
done
125139
126-
bundle = build.buildTorchExtensionBundle {
127-
inherit path doGetKernelCheck;
128-
rev = revUnderscored;
129-
};
130-
redistributable = build.buildDistTorchExtensions {
131-
inherit path doGetKernelCheck;
132-
bundleOnly = false;
133-
rev = revUnderscored;
140+
chmod -R +w build
141+
'';
142+
143+
kernels =
144+
bestBuildSet.pkgs.python3.withPackages (
145+
ps: with ps; [
146+
kernel-abi-check
147+
kernels
148+
]
149+
)
150+
// {
151+
meta.mainProgram = "kernels";
152+
};
153+
154+
redistributable = build.buildDistTorchExtensions {
155+
inherit path doGetKernelCheck;
156+
bundleOnly = false;
157+
rev = revUnderscored;
158+
};
134159
};
135-
};
136160
}

‎overlay.nix‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,4 +12,13 @@ final: prev: {
1212
rewrite-nix-paths-macho = prev.callPackage ./pkgs/rewrite-nix-paths-macho { };
1313

1414
stdenvGlibc_2_27 = prev.callPackage ./pkgs/stdenv-glibc-2_27 { };
15+
16+
# Python packages
17+
pythonPackagesExtensions = prev.pythonPackagesExtensions ++ [
18+
(
19+
python-self: python-super: with python-self; {
20+
kernel-abi-check = callPackage ./pkgs/python-modules/kernel-abi-check { };
21+
}
22+
)
23+
];
1524
}

‎pkgs/python-modules/kernel-abi-check/default.nix‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
{
2+
lib,
3+
buildPythonPackage,
4+
rustPlatform,
5+
}:
6+
7+
let
8+
version =
9+
(builtins.fromTOML (builtins.readFile ../../../kernel-abi-check/kernel-abi-check/Cargo.toml))
10+
.package.version;
11+
in
12+
buildPythonPackage {
13+
pname = "kernel-abi-check";
14+
inherit version;
15+
format = "pyproject";
16+
17+
src =
18+
let
19+
sourceFiles =
20+
file:
21+
file.name == "Cargo.toml"
22+
|| file.name == "Cargo.lock"
23+
|| file.name == "manylinux-policy.json"
24+
|| file.hasExt "pyi"
25+
|| file.name == "pyproject.toml"
26+
|| file.hasExt "rs"
27+
|| file.name == "stable_abi.toml";
28+
in
29+
lib.fileset.toSource {
30+
root = ../../../kernel-abi-check;
31+
fileset = lib.fileset.fileFilter sourceFiles ../../../kernel-abi-check;
32+
};
33+
34+
cargoDeps = rustPlatform.importCargoLock {
35+
lockFile = ../../../kernel-abi-check/bindings/python/Cargo.lock;
36+
};
37+
38+
sourceRoot = "source/bindings/python";
39+
40+
build-system = [
41+
rustPlatform.cargoSetupHook
42+
rustPlatform.maturinBuildHook
43+
];
44+
45+
meta = with lib; {
46+
description = "Check ABI compliance of Hugging Face Hub kernels";
47+
};
48+
}

‎tests/run-kernels/pyproject.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
[tool.kernels.dependencies]
2+
"kernels-test/versions" = ">=0.1.0,<0.2.0"

0 commit comments

Comments
 (0)