Skip to content

Commit 5fb66f1

Browse files
leagueleague
committed
nix: Big improvement generating bifrost pkgs for various configs
1 parent 925634d commit 5fb66f1

File tree

1 file changed

+138
-135
lines changed

1 file changed

+138
-135
lines changed

‎flake.nix

+138-135
Original file line numberDiff line numberDiff line change
@@ -17,32 +17,6 @@
1717
let
1818
inherit (nixpkgs) lib;
1919

20-
eachSystem = do:
21-
lib.genAttrs [ "x86_64-linux" "x86_64-darwin" ] (system:
22-
do (import nixpkgs {
23-
inherit system;
24-
config.allowUnfree = true;
25-
overlays = [ self.overlay ];
26-
}));
27-
28-
isPython = name: builtins.match "python3[0-9]*" name != null;
29-
shortenPy = lib.replaceStrings [ "thon" ] [ "" ];
30-
pythonAttrs = lib.filterAttrs (name: _: isPython name);
31-
pythonNames = pkgs: lib.filter isPython (lib.attrNames pkgs);
32-
concatForEach = args: f: lib.concatMap f args;
33-
concatEachBool = concatForEach [ true false ];
34-
concatEachPython = pkgs: concatForEach (pythonNames pkgs);
35-
36-
# Make a package name, embedding debug flag, CUDA version, or other
37-
# variations.
38-
mkpname = basename:
39-
{ enableDebug ? false, enableCuda ? false, cudatoolkit ? null }:
40-
basename + lib.optionalString enableCuda "-cuda"
41-
+ lib.optionalString (enableCuda && cudatoolkit != null)
42-
(lib.versions.major cudatoolkit.version
43-
+ lib.versions.minor cudatoolkit.version)
44-
+ lib.optionalString enableDebug "-debug";
45-
4620
# Parse the version info in the AC_INIT declaration.
4721
acVersion = lib.head
4822
(builtins.match "AC_INIT\\(\\[bifrost], *\\[([.0-9]+)].*" (lib.head
@@ -54,36 +28,60 @@
5428
version = "${acVersion}.dev"
5529
+ lib.optionalString (self ? shortRev) "+g${self.shortRev}";
5630

57-
libbifrost = { stdenv, ctags, ncurses, cudatoolkit, util-linuxMinimal
58-
, enableDebug ? false, enableCuda ? false }:
31+
bifrost = { stdenv, ctags, ncurses, file, enableDebug ? false
32+
, enablePython ? true, python3, enableCuda ? false, cudatoolkit
33+
, util-linuxMinimal }:
5934
let
60-
basename = "libbifrost";
61-
pname =
62-
mkpname basename { inherit enableDebug enableCuda cudatoolkit; };
35+
pname = lib.optionalString (!enablePython) "lib" + "bifrost"
36+
+ lib.optionalString enablePython
37+
"-py${lib.versions.majorMinor python3.version}"
38+
+ lib.optionalString enableCuda
39+
"-cuda${lib.versions.majorMinor cudatoolkit.version}"
40+
+ lib.optionalString enableDebug "-debug";
6341
in stdenv.mkDerivation {
6442
name = "${pname}-${version}";
6543
inherit version;
6644
src = ./.;
67-
# ncurses provides tput, used in C++ build messages.
68-
# util-linux provides rev, used in configure if GPU.
69-
buildInputs = [ ctags ncurses ]
70-
++ lib.optionals enableCuda [ cudatoolkit util-linuxMinimal ];
71-
configureFlags = lib.optionals enableCuda [
72-
"--with-cuda-home=${cudatoolkit}"
73-
"--with-gpu-archs=60"
74-
"LDFLAGS=-L${cudatoolkit}/lib/stubs"
75-
] ++ lib.optionals enableDebug [ "--enable-debug" ];
45+
buildInputs = [ ctags ncurses ] ++ lib.optionals enablePython [
46+
python3
47+
python3.pkgs.ctypesgen
48+
python3.pkgs.setuptools
49+
python3.pkgs.pip
50+
python3.pkgs.wheel
51+
] ++ lib.optionals enableCuda [ cudatoolkit util-linuxMinimal ];
52+
propagatedBuildInputs = lib.optionals enablePython [
53+
python3.pkgs.contextlib2
54+
python3.pkgs.graphviz
55+
python3.pkgs.matplotlib
56+
python3.pkgs.numpy
57+
python3.pkgs.pint
58+
python3.pkgs.scipy
59+
python3.pkgs.simplejson
60+
];
61+
patchPhase = ''
62+
sed -i -e 's:\\r::g' -e 's:echo -n:echo:' \
63+
-e 's:\(CLEAR_LINE = \).*:\1:' src/autodep.mk src/Makefile.in
64+
sed -i 's:/usr/bin/file:${file}/bin/file:' configure
65+
sed -i 's/ctypesgen==1.0.2/ctypesgen/' python/setup.py
66+
'';
67+
configureFlags = lib.optionals enablePython
68+
[ "--with-pyinstall-flags=--prefix=${placeholder "out"}" ]
69+
++ lib.optionals enableDebug [ "--enable-debug" ]
70+
++ lib.optionals enableCuda [
71+
"--with-cuda-home=${cudatoolkit}"
72+
"--with-gpu-archs=60"
73+
"LDFLAGS=-L${cudatoolkit}/lib/stubs"
74+
];
75+
preBuild = lib.optionalString enablePython ''
76+
make -C python bifrost/libbifrost_generated.py
77+
sed -i \
78+
-e "s:^add_library_search_dirs(\[:&'$out/lib':" \
79+
python/bifrost/libbifrost_generated.py
80+
'';
7681
makeFlags =
7782
lib.optionals enableCuda [ "CUDA_LIBDIR64=$(CUDA_HOME)/lib" ];
78-
installPhase = let
79-
ext = if stdenv.isDarwin then "dylib" else "so";
80-
base = "${basename}.${ext}";
81-
dll = "${base}.${lib.versions.majorMinor version}";
82-
in ''
83-
install -D -t $out/lib lib/${dll}
84-
ln -s ${dll} $out/lib/${base}.${lib.versions.major version}
85-
ln -s ${dll} $out/lib/${base}
86-
install -m 644 -D -t $out/include/bifrost src/bifrost/*.{h,hpp}
83+
preInstall = ''
84+
mkdir -p "$out/lib"
8785
'';
8886
};
8987

@@ -123,46 +121,16 @@
123121
checkPhase = "python -m unittest -v tests/testsuite.py";
124122
};
125123

126-
bifrost = { buildPythonPackage, contextlib2, ctags, ctypesgen, graphviz
127-
, libbifrost, matplotlib, numpy, pint, scipy, simplejson }:
128-
buildPythonPackage {
129-
pname = mkpname "bifrost" {
130-
enableDebug = lib.elem "--enable-debug" libbifrost.configureFlags;
131-
enableCuda = lib.any (lib.hasPrefix "--with-cuda-home")
132-
libbifrost.configureFlags;
133-
};
134-
inherit version;
135-
src = ./.;
136-
# Run configure here to get the rules in ‘python/Makefile.in’ for
137-
# generating and tweaking the generated ‘libbifrost’ interface. Then
138-
# also patch ‘add_library_search_dirs’.
139-
patchPhase = ''
140-
./configure ${lib.concatStringsSep " " libbifrost.configureFlags}
141-
cd python
142-
make bifrost/libbifrost_generated.py
143-
sed -i \
144-
-e 's:^add_library_search_dirs(\[:&"${libbifrost}/lib":' \
145-
bifrost/libbifrost_generated.py
146-
sed -i 's/ctypesgen==1.0.2/ctypesgen/' setup.py
147-
echo '__version__ = "${version}"' > bifrost/version/__init__.py
148-
'';
149-
nativeBuildInputs = [ ctags ];
150-
buildInputs = [ libbifrost ctypesgen ];
151-
propagatedBuildInputs =
152-
[ contextlib2 graphviz matplotlib numpy pint scipy simplejson ];
153-
# Tests currently fail when run by nix build, because they assume
154-
# they can access network to download data, and home directory for
155-
# telemetry.
156-
doCheck = false;
157-
};
158-
159124
pyOverlay = self: _: {
160125
ctypesgen = self.callPackage ctypesgen { };
161-
bifrost = self.callPackage bifrost { };
126+
bifrost = self.toPythonModule (self.callPackage bifrost {
127+
enablePython = true;
128+
python3 = self.python;
129+
});
162130
};
163131

164132
bifrost-doc =
165-
{ stdenv, python3, ctags, doxygen, destDir ? "/share/doc/bifrost" }:
133+
{ stdenv, python3, ctags, doxygen, docDir ? "/share/doc/bifrost" }:
166134
let pname = "bifrost-doc";
167135
in stdenv.mkDerivation {
168136
name = "${pname}-${version}";
@@ -171,7 +139,10 @@
171139
buildInputs = [
172140
ctags
173141
doxygen
174-
(python3.withPackages (p: [ p.bifrost p.sphinx p.breathe ]))
142+
python3
143+
python3.pkgs.bifrost
144+
python3.pkgs.sphinx
145+
python3.pkgs.breathe
175146
];
176147
buildPhase = ''
177148
make doc
@@ -188,59 +159,88 @@
188159
cd ../../..
189160
'';
190161
installPhase = ''
191-
mkdir -p "$out${destDir}"
192-
cp -r docs/build/html "$out${destDir}"
162+
mkdir -p "$out${docDir}"
163+
cp -r docs/build/html "$out${docDir}"
193164
'';
194165
};
195166

167+
# Enable pre-configured packages for these systems.
168+
eachSystem = do:
169+
lib.genAttrs [ "x86_64-linux" "x86_64-darwin" ] (system:
170+
do (import nixpkgs {
171+
inherit system;
172+
config.allowUnfree = true;
173+
overlays = [ self.overlay ];
174+
}));
175+
176+
# Which python3 packages should be modified by the overlay?
177+
isPython = name: builtins.match "python3[0-9]*" name != null;
178+
pythonAttrs = lib.filterAttrs (name: _: isPython name);
179+
196180
in {
197181
overlay = final: prev:
198182
{
199-
libbifrost = final.callPackage libbifrost { };
183+
bifrost = final.callPackage bifrost { };
200184
bifrost-doc = final.callPackage bifrost-doc { };
201185
}
202186
# Apply the python overlay to every python package set we find.
203187
// lib.mapAttrs (_: py: py.override { packageOverrides = pyOverlay; })
204188
(pythonAttrs prev);
205189

206190
packages = eachSystem (pkgs:
207-
lib.listToAttrs ([{
208-
name = "bifrost-doc";
209-
value = pkgs.bifrost-doc;
210-
}] ++
211-
# Generate a runnable ctypesgen corresponding to each python.
212-
concatEachPython pkgs (python: [{
213-
name = "ctypesgen-${shortenPy python}";
214-
value = pkgs.${python}.pkgs.ctypesgen;
215-
}]) ++
216-
# For each bifrost configuration, generate library package.
217-
concatEachBool (enableDebug:
218-
concatForEach
219-
([ false ] ++ lib.optionals pkgs.stdenv.isLinux [ true ])
220-
(enableCuda:
221-
let
222-
config = { inherit enableDebug enableCuda; };
223-
libbifrost = pkgs.libbifrost.override config;
224-
in [{
225-
name = mkpname "libbifrost" config;
226-
value = libbifrost;
227-
}] ++
228-
# And then corresponding python packages.
229-
concatEachPython pkgs (python:
230-
let
231-
bipy = pkgs.${python}.pkgs.bifrost.override {
232-
inherit libbifrost;
233-
};
234-
in [
235-
{
236-
name = mkpname "bifrost-${shortenPy python}" config;
237-
value = bipy;
238-
}
239-
{
240-
name = mkpname "${python}-bifrost" config;
241-
value = pkgs.${python}.withPackages (_: [ bipy ]);
242-
}
243-
])))));
191+
let
192+
shortenPy = lib.replaceStrings [ "thon" ] [ "" ];
193+
194+
# Which cuda versions should be target by the packages? Let's just do
195+
# the default 10 and 11. It's easy to generate other point releases
196+
# from the overlay. (Versions prior to 10 are not supported anymore by
197+
# nixpkgs.)
198+
isCuda = name: builtins.match "cudatoolkit(_1[01])" name != null;
199+
shortenCuda = lib.replaceStrings [ "toolkit" "_" ] [ "" "" ];
200+
cudaAttrs = lib.filterAttrs
201+
(name: pkg: isCuda name && lib.elem pkgs.system pkg.meta.platforms)
202+
pkgs;
203+
204+
eachBool = f: lib.concatMap f [ true false ];
205+
eachCuda = f: lib.concatMap f ([ null ] ++ lib.attrNames cudaAttrs);
206+
eachConfig = f:
207+
eachBool (enableDebug:
208+
eachCuda (cuda:
209+
f (lib.optionalString (cuda != null) "-${shortenCuda cuda}"
210+
+ lib.optionalString enableDebug "-debug") {
211+
inherit enableDebug;
212+
enableCuda = cuda != null;
213+
cudatoolkit = pkgs.${cuda};
214+
}));
215+
216+
# Runnable ctypesgen per python. Though it's just the executable we
217+
# need, it's possible something about ctypes library could change
218+
# between releases.
219+
cgens = lib.mapAttrs' (name: py: {
220+
name = "ctypesgen-${shortenPy name}";
221+
value = py.pkgs.ctypesgen;
222+
}) (pythonAttrs pkgs);
223+
224+
# The whole set of bifrost packages, with or without python (and each
225+
# python version), and for each configuration.
226+
bfs = lib.listToAttrs (eachConfig (suffix: config:
227+
[{
228+
name = "libbifrost${suffix}";
229+
value =
230+
pkgs.bifrost.override (config // { enablePython = false; });
231+
}] ++ lib.mapAttrsToList (name: py: {
232+
name = "bifrost-${shortenPy name}${suffix}";
233+
value = py.pkgs.bifrost.override config;
234+
}) (pythonAttrs pkgs)));
235+
236+
# Now generate pythons with bifrost packaged.
237+
pys = lib.listToAttrs (eachConfig (suffix: config:
238+
lib.mapAttrsToList (name: py: {
239+
name = "${name}-bifrost${suffix}";
240+
value = py.withPackages (p: [ (p.bifrost.override config) ]);
241+
}) (pythonAttrs pkgs)));
242+
243+
in { inherit (pkgs) bifrost-doc; } // cgens // bfs // pys);
244244

245245
devShell = eachSystem (pkgs:
246246
let
@@ -254,16 +254,19 @@
254254
in pkgs.mkShell {
255255
inherit (pre-commit) shellHook;
256256

257-
buildInputs = pkgs.libbifrost.buildInputs ++ [
258-
pkgs.black
259-
pkgs.doxygen
260-
pkgs.nixfmt
261-
pkgs.nix-linter
262-
pkgs.yamllint
263-
(pkgs.python3.withPackages (p:
264-
p.bifrost.propagatedBuildInputs
265-
++ [ p.sphinx p.breathe p.ctypesgen ]))
266-
];
257+
# Tempting to include bifrost-doc.buildInputs here, but that requires
258+
# bifrost to already be built.
259+
buildInputs = pkgs.bifrost.buildInputs
260+
++ pkgs.bifrost.propagatedBuildInputs ++ [
261+
pkgs.black
262+
pkgs.ctags
263+
pkgs.doxygen
264+
pkgs.nixfmt
265+
pkgs.nix-linter
266+
pkgs.python3.pkgs.breathe
267+
pkgs.python3.pkgs.sphinx
268+
pkgs.yamllint
269+
];
267270
});
268271
};
269272
}

0 commit comments

Comments
 (0)