tjh.dev
1{ stdenv, pkgs, bazel_0_26, buildBazelPackage, lib, fetchFromGitHub, fetchpatch, symlinkJoin
2, addOpenGLRunpath
3# Python deps
4, buildPythonPackage, isPy3k, isPy27, pythonOlder, pythonAtLeast, python
5# Python libraries
6, numpy, tensorflow-tensorboard_1, backports_weakref, mock, enum34, absl-py
7, future, setuptools, wheel, keras-preprocessing, keras-applications, google-pasta
8, functools32
9, opt-einsum
10, termcolor, grpcio, six, wrapt, protobuf, tensorflow-estimator_1
11# Common deps
12, git, swig, which, binutils, glibcLocales, cython
13# Common libraries
14, jemalloc, openmpi, astor, gast, grpc, sqlite, openssl, jsoncpp, re2
15, curl, snappy, flatbuffers, icu, double-conversion, libpng, libjpeg, giflib
16# Upsteam by default includes cuda support since tensorflow 1.15. We could do
17# that in nix as well. It would make some things easier and less confusing, but
18# it would also make the default tensorflow package unfree. See
19# https://groups.google.com/a/tensorflow.org/forum/#!topic/developers/iRCt5m4qUz0
20, cudaSupport ? false, nvidia_x11 ? null, cudatoolkit ? null, cudnn ? null, nccl ? null
21, mklSupport ? false, mkl ? null
22# XLA without CUDA is broken
23, xlaSupport ? cudaSupport
24# Default from ./configure script
25, cudaCapabilities ? [ "3.5" "5.2" ]
26, sse42Support ? stdenv.hostPlatform.sse4_2Support
27, avx2Support ? stdenv.hostPlatform.avx2Support
28, fmaSupport ? stdenv.hostPlatform.fmaSupport
29# Darwin deps
30, Foundation, Security
31}:
32
33assert cudaSupport -> nvidia_x11 != null
34 && cudatoolkit != null
35 && cudnn != null;
36
37# unsupported combination
38assert ! (stdenv.isDarwin && cudaSupport);
39
40assert mklSupport -> mkl != null;
41
42let
43 withTensorboard = pythonOlder "3.6";
44
45 cudatoolkit_joined = symlinkJoin {
46 name = "${cudatoolkit.name}-merged";
47 paths = [
48 cudatoolkit.lib
49 cudatoolkit.out
50 # for some reason some of the required libs are in the targets/x86_64-linux
51 # directory; not sure why but this works around it
52 "${cudatoolkit}/targets/${stdenv.system}"
53 ];
54 };
55
56 cudatoolkit_cc_joined = symlinkJoin {
57 name = "${cudatoolkit.cc.name}-merged";
58 paths = [
59 cudatoolkit.cc
60 binutils.bintools # for ar, dwp, nm, objcopy, objdump, strip
61 ];
62 };
63
64 # Needed for _some_ system libraries, grep INCLUDEDIR.
65 includes_joined = symlinkJoin {
66 name = "tensorflow-deps-merged";
67 paths = [
68 pkgs.protobuf
69 jsoncpp
70 ];
71 };
72
73 tfFeature = x: if x then "1" else "0";
74
75 version = "1.15.2";
76 variant = if cudaSupport then "-gpu" else "";
77 pname = "tensorflow${variant}";
78
79 pythonEnv = python.withPackages (_:
80 [ # python deps needed during wheel build time (not runtime, see the buildPythonPackage part for that)
81 numpy
82 keras-preprocessing
83 protobuf
84 wrapt
85 gast
86 astor
87 absl-py
88 termcolor
89 keras-applications
90 setuptools
91 wheel
92 ] ++ lib.optionals (!isPy3k)
93 [ future
94 functools32
95 mock
96 ]);
97
98 bazel-build = buildBazelPackage {
99 name = "${pname}-${version}";
100 bazel = bazel_0_26;
101
102 src = fetchFromGitHub {
103 owner = "tensorflow";
104 repo = "tensorflow";
105 rev = "v${version}";
106 sha256 = "1q0848drjvnaaa38dgns8knmpmkj5plzsc98j20m5ybv68s55w78";
107 };
108
109 patches = [
110 # Work around https://github.com/tensorflow/tensorflow/issues/24752
111 ../no-saved-proto.patch
112 # Fixes for NixOS jsoncpp
113 ../system-jsoncpp.patch
114
115 # https://github.com/tensorflow/tensorflow/pull/29673
116 (fetchpatch {
117 name = "fix-compile-with-cuda-and-mpi.patch";
118 url = "https://github.com/tensorflow/tensorflow/pull/29673/commits/498e35a3bfe38dd75cf1416a1a23c07c3b59e6af.patch";
119 sha256 = "1m2qmwv1ysqa61z6255xggwbq6mnxbig749bdvrhnch4zydxb4di";
120 })
121 (fetchpatch {
122 name = "backport-pr-18950.patch";
123 url = "https://github.com/tensorflow/tensorflow/commit/73640aaec2ab0234d9fff138e3c9833695570c0a.patch";
124 sha256 = "1n9ypbrx36fc1kc9cz5b3p9qhg15xxhq4nz6ap3hwqba535nakfz";
125 })
126
127 (fetchpatch {
128 # be compatible with gast >0.2 instead of only gast 0.2.2
129 name = "gast-update.patch";
130 url = "https://github.com/tensorflow/tensorflow/commit/85751ad6c7f5fd12c6c79545d96896cba92fa8b4.patch";
131 sha256 = "077cpj0kzyqxzdya1dwh8df17zfzhqn7c685hx6iskvw2979zg2n";
132 })
133 ./lift-gast-restriction.patch
134
135 (fetchpatch {
136 # fix compilation with numpy >= 1.19
137 name = "add-const-overload.patch";
138 url = "https://github.com/tensorflow/tensorflow/commit/75ea0b31477d6ba9e990e296bbbd8ca4e7eebadf.patch";
139 sha256 = "1xp1icacig0xm0nmb05sbrf4nw4xbln9fhc308birrv8286zx7wv";
140 })
141
142 # cuda 10.2 does not have "-bin2c-path" option anymore
143 # https://github.com/tensorflow/tensorflow/issues/34429
144 ../cuda-10.2-no-bin2c-path.patch
145 ];
146
147 # On update, it can be useful to steal the changes from gentoo
148 # https://gitweb.gentoo.org/repo/gentoo.git/tree/sci-libs/tensorflow
149
150 nativeBuildInputs = [
151 swig which pythonEnv
152 ] ++ lib.optional cudaSupport addOpenGLRunpath;
153
154 buildInputs = [
155 jemalloc
156 openmpi
157 glibcLocales
158 git
159
160 # libs taken from system through the TF_SYS_LIBS mechanism
161 # grpc
162 sqlite
163 openssl
164 jsoncpp
165 pkgs.protobuf
166 curl
167 snappy
168 flatbuffers
169 icu
170 double-conversion
171 libpng
172 libjpeg
173 giflib
174 re2
175 pkgs.lmdb
176 ] ++ lib.optionals cudaSupport [
177 cudatoolkit
178 cudnn
179 nvidia_x11
180 ] ++ lib.optionals mklSupport [
181 mkl
182 ] ++ lib.optionals stdenv.isDarwin [
183 Foundation
184 Security
185 ];
186
187 # arbitrarily set to the current latest bazel version, overly careful
188 TF_IGNORE_MAX_BAZEL_VERSION = true;
189
190 # Take as many libraries from the system as possible. Keep in sync with
191 # list of valid syslibs in
192 # https://github.com/tensorflow/tensorflow/blob/master/third_party/systemlibs/syslibs_configure.bzl
193 TF_SYSTEM_LIBS = lib.concatStringsSep "," [
194 "absl_py"
195 "astor_archive"
196 "boringssl"
197 # Not packaged in nixpkgs
198 # "com_github_googleapis_googleapis"
199 # "com_github_googlecloudplatform_google_cloud_cpp"
200 "com_google_protobuf"
201 "com_googlesource_code_re2"
202 "curl"
203 "cython"
204 "double_conversion"
205 "flatbuffers"
206 "gast_archive"
207 "gif_archive"
208 # Lots of errors, requires an older version
209 # "grpc"
210 "hwloc"
211 "icu"
212 "jpeg"
213 "jsoncpp_git"
214 "keras_applications_archive"
215 "lmdb"
216 "nasm"
217 # "nsync" # not packaged in nixpkgs
218 "opt_einsum_archive"
219 "org_sqlite"
220 "pasta"
221 "pcre"
222 "png_archive"
223 "six_archive"
224 "snappy"
225 "swig"
226 "termcolor_archive"
227 "wrapt"
228 "zlib_archive"
229 ];
230
231 INCLUDEDIR = "${includes_joined}/include";
232
233 PYTHON_BIN_PATH = pythonEnv.interpreter;
234
235 TF_NEED_GCP = true;
236 TF_NEED_HDFS = true;
237 TF_ENABLE_XLA = tfFeature xlaSupport;
238
239 CC_OPT_FLAGS = " ";
240
241 # https://github.com/tensorflow/tensorflow/issues/14454
242 TF_NEED_MPI = tfFeature cudaSupport;
243
244 TF_NEED_CUDA = tfFeature cudaSupport;
245 TF_CUDA_PATHS = lib.optionalString cudaSupport "${cudatoolkit_joined},${cudnn},${nccl}";
246 GCC_HOST_COMPILER_PREFIX = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin";
247 GCC_HOST_COMPILER_PATH = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin/gcc";
248 TF_CUDA_COMPUTE_CAPABILITIES = lib.concatStringsSep "," cudaCapabilities;
249
250 postPatch = ''
251 # https://github.com/tensorflow/tensorflow/issues/20919
252 sed -i '/androidndk/d' tensorflow/lite/kernels/internal/BUILD
253
254 # Tensorboard pulls in a bunch of dependencies, some of which may
255 # include security vulnerabilities. So we make it optional.
256 # https://github.com/tensorflow/tensorflow/issues/20280#issuecomment-400230560
257 sed -i '/tensorboard >=/d' tensorflow/tools/pip_package/setup.py
258 '';
259
260 preConfigure = let
261 opt_flags = []
262 ++ lib.optionals sse42Support ["-msse4.2"]
263 ++ lib.optionals avx2Support ["-mavx2"]
264 ++ lib.optionals fmaSupport ["-mfma"];
265 in ''
266 patchShebangs configure
267
268 # dummy ldconfig
269 mkdir dummy-ldconfig
270 echo "#!${stdenv.shell}" > dummy-ldconfig/ldconfig
271 chmod +x dummy-ldconfig/ldconfig
272 export PATH="$PWD/dummy-ldconfig:$PATH"
273
274 export PYTHON_LIB_PATH="$NIX_BUILD_TOP/site-packages"
275 export CC_OPT_FLAGS="${lib.concatStringsSep " " opt_flags}"
276 mkdir -p "$PYTHON_LIB_PATH"
277
278 # To avoid mixing Python 2 and Python 3
279 unset PYTHONPATH
280 '';
281
282 configurePhase = ''
283 runHook preConfigure
284 ./configure
285 runHook postConfigure
286 '';
287
288 # FIXME: Tensorflow uses dlopen() for CUDA libraries.
289 NIX_LDFLAGS = lib.optionalString cudaSupport "-lcudart -lcublas -lcufft -lcurand -lcusolver -lcusparse -lcudnn";
290
291 hardeningDisable = [ "format" ];
292
293 bazelFlags = [
294 # temporary fixes to make the build work with bazel 0.27
295 "--incompatible_no_support_tools_in_action_inputs=false"
296 ];
297 bazelBuildFlags = [
298 "--config=opt" # optimize using the flags set in the configure phase
299 ]
300 ++ lib.optionals (mklSupport) [ "--config=mkl" ];
301
302 bazelTarget = "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow";
303
304 fetchAttrs = {
305 # So that checksums don't depend on these.
306 TF_SYSTEM_LIBS = null;
307
308 # cudaSupport causes fetch of ncclArchive, resulting in different hashes
309 sha256 = if cudaSupport then
310 "09j57w6kc0vkfcdwr0qggy3qgrgq82kfa2jrwvvcnij4bl3wj40j"
311 else
312 "14g8z49qz7d8n1c2mcsfhr7yqpcy7mhmpm6hgmqvpgb8vm7yvwrc";
313 };
314
315 buildAttrs = {
316 outputs = [ "out" "python" ];
317
318 preBuild = ''
319 patchShebangs .
320 '';
321
322 installPhase = ''
323 mkdir -p "$out"
324 tar -xf bazel-bin/tensorflow/tools/lib_package/libtensorflow.tar.gz -C "$out"
325 # Write pkgconfig file.
326 mkdir "$out/lib/pkgconfig"
327 cat > "$out/lib/pkgconfig/tensorflow.pc" << EOF
328 Name: TensorFlow
329 Version: ${version}
330 Description: Library for computation using data flow graphs for scalable machine learning
331 Requires:
332 Libs: -L$out/lib -ltensorflow
333 Cflags: -I$out/include/tensorflow
334 EOF
335
336 # build the source code, then copy it to $python (build_pip_package
337 # actually builds a symlink farm so we must dereference them).
338 bazel-bin/tensorflow/tools/pip_package/build_pip_package --src "$PWD/dist"
339 cp -Lr "$PWD/dist" "$python"
340 '';
341
342 postFixup = lib.optionalString cudaSupport ''
343 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
344 addOpenGLRunpath "$lib"
345 done
346 '';
347 };
348
349 meta = with stdenv.lib; {
350 description = "Computation using data flow graphs for scalable machine learning";
351 homepage = "http://tensorflow.org";
352 license = licenses.asl20;
353 maintainers = with maintainers; [ jyp abbradar ];
354 platforms = with platforms; linux ++ darwin;
355 # The py2 build fails due to some issue importing protobuf. Possibly related to the fix in
356 # https://github.com/akesandgren/easybuild-easyblocks/commit/1f2e517ddfd1b00a342c6abb55aef3fd93671a2b
357 broken = !(xlaSupport -> cudaSupport) || !isPy3k;
358 };
359 };
360
361in buildPythonPackage {
362 inherit version pname;
363 disabled = isPy27 || (pythonAtLeast "3.8");
364
365 src = bazel-build.python;
366
367 # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
368 # and the propagated input tensorflow-tensorboard, which causes environment collisions.
369 # Another possibility would be to have tensorboard only in the buildInputs
370 # https://github.com/tensorflow/tensorflow/blob/v1.7.1/tensorflow/tools/pip_package/setup.py#L79
371 postInstall = ''
372 rm $out/bin/tensorboard
373 '';
374
375 setupPyGlobalFlags = [ "--project_name ${pname}" ];
376
377 # tensorflow/tools/pip_package/setup.py
378 propagatedBuildInputs = [
379 absl-py
380 astor
381 gast
382 google-pasta
383 keras-applications
384 keras-preprocessing
385 numpy
386 six
387 protobuf
388 tensorflow-estimator_1
389 termcolor
390 wrapt
391 grpcio
392 opt-einsum
393 ] ++ lib.optionals (!isPy3k) [
394 mock
395 future
396 functools32
397 ] ++ lib.optionals (pythonOlder "3.4") [
398 backports_weakref enum34
399 ] ++ lib.optionals withTensorboard [
400 tensorflow-tensorboard_1
401 ];
402
403 nativeBuildInputs = lib.optional cudaSupport addOpenGLRunpath;
404
405 postFixup = lib.optionalString cudaSupport ''
406 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
407 addOpenGLRunpath "$lib"
408 done
409 '';
410
411 # Actual tests are slow and impure.
412 # TODO try to run them anyway
413 # TODO better test (files in tensorflow/tools/ci_build/builds/*test)
414 checkPhase = ''
415 ${python.interpreter} <<EOF
416 # A simple "Hello world"
417 import tensorflow as tf
418 hello = tf.constant("Hello, world!")
419 sess = tf.Session()
420 sess.run(hello)
421
422 # Fit a simple model to random data
423 import numpy as np
424 np.random.seed(0)
425 tf.random.set_random_seed(0)
426 model = tf.keras.models.Sequential([
427 tf.keras.layers.Dense(1, activation="linear")
428 ])
429 model.compile(optimizer="sgd", loss="mse")
430
431 x = np.random.uniform(size=(1,1))
432 y = np.random.uniform(size=(1,))
433 model.fit(x, y, epochs=1)
434
435 # regression test for #77626
436 from tensorflow.contrib import tensor_forest
437 EOF
438 '';
439
440 passthru = {
441 deps = bazel-build.deps;
442 libtensorflow = bazel-build.out;
443 };
444
445 inherit (bazel-build) meta;
446}