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 ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") ["westmere" "sandybridge" "ivybridge" "haswell" "broadwell" "skylake" "skylake-avx512"]
27, avx2Support ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") [ "haswell" "broadwell" "skylake" "skylake-avx512"]
28, fmaSupport ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") [ "haswell" "broadwell" "skylake" "skylake-avx512"]
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 # cuda 10.2 does not have "-bin2c-path" option anymore
136 # https://github.com/tensorflow/tensorflow/issues/34429
137 ../cuda-10.2-no-bin2c-path.patch
138 ];
139
140 # On update, it can be useful to steal the changes from gentoo
141 # https://gitweb.gentoo.org/repo/gentoo.git/tree/sci-libs/tensorflow
142
143 nativeBuildInputs = [
144 swig which pythonEnv
145 ] ++ lib.optional cudaSupport addOpenGLRunpath;
146
147 buildInputs = [
148 jemalloc
149 openmpi
150 glibcLocales
151 git
152
153 # libs taken from system through the TF_SYS_LIBS mechanism
154 # grpc
155 sqlite
156 openssl
157 jsoncpp
158 pkgs.protobuf
159 curl
160 snappy
161 flatbuffers
162 icu
163 double-conversion
164 libpng
165 libjpeg
166 giflib
167 re2
168 pkgs.lmdb
169 ] ++ lib.optionals cudaSupport [
170 cudatoolkit
171 cudnn
172 nvidia_x11
173 ] ++ lib.optionals mklSupport [
174 mkl
175 ] ++ lib.optionals stdenv.isDarwin [
176 Foundation
177 Security
178 ];
179
180 # arbitrarily set to the current latest bazel version, overly careful
181 TF_IGNORE_MAX_BAZEL_VERSION = true;
182
183 # Take as many libraries from the system as possible. Keep in sync with
184 # list of valid syslibs in
185 # https://github.com/tensorflow/tensorflow/blob/master/third_party/systemlibs/syslibs_configure.bzl
186 TF_SYSTEM_LIBS = lib.concatStringsSep "," [
187 "absl_py"
188 "astor_archive"
189 "boringssl"
190 # Not packaged in nixpkgs
191 # "com_github_googleapis_googleapis"
192 # "com_github_googlecloudplatform_google_cloud_cpp"
193 "com_google_protobuf"
194 "com_googlesource_code_re2"
195 "curl"
196 "cython"
197 "double_conversion"
198 "flatbuffers"
199 "gast_archive"
200 "gif_archive"
201 # Lots of errors, requires an older version
202 # "grpc"
203 "hwloc"
204 "icu"
205 "jpeg"
206 "jsoncpp_git"
207 "keras_applications_archive"
208 "lmdb"
209 "nasm"
210 # "nsync" # not packaged in nixpkgs
211 "opt_einsum_archive"
212 "org_sqlite"
213 "pasta"
214 "pcre"
215 "png_archive"
216 "six_archive"
217 "snappy"
218 "swig"
219 "termcolor_archive"
220 "wrapt"
221 "zlib_archive"
222 ];
223
224 INCLUDEDIR = "${includes_joined}/include";
225
226 PYTHON_BIN_PATH = pythonEnv.interpreter;
227
228 TF_NEED_GCP = true;
229 TF_NEED_HDFS = true;
230 TF_ENABLE_XLA = tfFeature xlaSupport;
231
232 CC_OPT_FLAGS = " ";
233
234 # https://github.com/tensorflow/tensorflow/issues/14454
235 TF_NEED_MPI = tfFeature cudaSupport;
236
237 TF_NEED_CUDA = tfFeature cudaSupport;
238 TF_CUDA_PATHS = lib.optionalString cudaSupport "${cudatoolkit_joined},${cudnn},${nccl}";
239 GCC_HOST_COMPILER_PREFIX = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin";
240 GCC_HOST_COMPILER_PATH = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin/gcc";
241 TF_CUDA_COMPUTE_CAPABILITIES = lib.concatStringsSep "," cudaCapabilities;
242
243 postPatch = ''
244 # https://github.com/tensorflow/tensorflow/issues/20919
245 sed -i '/androidndk/d' tensorflow/lite/kernels/internal/BUILD
246
247 # Tensorboard pulls in a bunch of dependencies, some of which may
248 # include security vulnerabilities. So we make it optional.
249 # https://github.com/tensorflow/tensorflow/issues/20280#issuecomment-400230560
250 sed -i '/tensorboard >=/d' tensorflow/tools/pip_package/setup.py
251 '';
252
253 preConfigure = let
254 opt_flags = []
255 ++ lib.optionals sse42Support ["-msse4.2"]
256 ++ lib.optionals avx2Support ["-mavx2"]
257 ++ lib.optionals fmaSupport ["-mfma"];
258 in ''
259 patchShebangs configure
260
261 # dummy ldconfig
262 mkdir dummy-ldconfig
263 echo "#!${stdenv.shell}" > dummy-ldconfig/ldconfig
264 chmod +x dummy-ldconfig/ldconfig
265 export PATH="$PWD/dummy-ldconfig:$PATH"
266
267 export PYTHON_LIB_PATH="$NIX_BUILD_TOP/site-packages"
268 export CC_OPT_FLAGS="${lib.concatStringsSep " " opt_flags}"
269 mkdir -p "$PYTHON_LIB_PATH"
270
271 # To avoid mixing Python 2 and Python 3
272 unset PYTHONPATH
273 '';
274
275 configurePhase = ''
276 runHook preConfigure
277 ./configure
278 runHook postConfigure
279 '';
280
281 # FIXME: Tensorflow uses dlopen() for CUDA libraries.
282 NIX_LDFLAGS = lib.optionalString cudaSupport "-lcudart -lcublas -lcufft -lcurand -lcusolver -lcusparse -lcudnn";
283
284 hardeningDisable = [ "format" ];
285
286 bazelFlags = [
287 # temporary fixes to make the build work with bazel 0.27
288 "--incompatible_no_support_tools_in_action_inputs=false"
289 ];
290 bazelBuildFlags = [
291 "--config=opt" # optimize using the flags set in the configure phase
292 ]
293 ++ lib.optionals (mklSupport) [ "--config=mkl" ];
294
295 bazelTarget = "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow";
296
297 fetchAttrs = {
298 # So that checksums don't depend on these.
299 TF_SYSTEM_LIBS = null;
300
301 # cudaSupport causes fetch of ncclArchive, resulting in different hashes
302 sha256 = if cudaSupport then
303 "1qygfcvvn9vysap9nk6xccxi9mgmzyxiywz6k456f811l1v70p2c"
304 else
305 "0kfjanw0mfbh30vi1ms2xlg8yp429cbyfriik6yxd5cla2pncg2j";
306 };
307
308 buildAttrs = {
309 outputs = [ "out" "python" ];
310
311 preBuild = ''
312 patchShebangs .
313 '';
314
315 installPhase = ''
316 mkdir -p "$out"
317 tar -xf bazel-bin/tensorflow/tools/lib_package/libtensorflow.tar.gz -C "$out"
318 # Write pkgconfig file.
319 mkdir "$out/lib/pkgconfig"
320 cat > "$out/lib/pkgconfig/tensorflow.pc" << EOF
321 Name: TensorFlow
322 Version: ${version}
323 Description: Library for computation using data flow graphs for scalable machine learning
324 Requires:
325 Libs: -L$out/lib -ltensorflow
326 Cflags: -I$out/include/tensorflow
327 EOF
328
329 # build the source code, then copy it to $python (build_pip_package
330 # actually builds a symlink farm so we must dereference them).
331 bazel-bin/tensorflow/tools/pip_package/build_pip_package --src "$PWD/dist"
332 cp -Lr "$PWD/dist" "$python"
333 '';
334
335 postFixup = lib.optionalString cudaSupport ''
336 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
337 addOpenGLRunpath "$lib"
338 done
339 '';
340 };
341
342 meta = with stdenv.lib; {
343 description = "Computation using data flow graphs for scalable machine learning";
344 homepage = http://tensorflow.org;
345 license = licenses.asl20;
346 maintainers = with maintainers; [ jyp abbradar ];
347 platforms = with platforms; linux ++ darwin;
348 # The py2 build fails due to some issue importing protobuf. Possibly related to the fix in
349 # https://github.com/akesandgren/easybuild-easyblocks/commit/1f2e517ddfd1b00a342c6abb55aef3fd93671a2b
350 broken = !(xlaSupport -> cudaSupport) || !isPy3k;
351 };
352 };
353
354in buildPythonPackage {
355 inherit version pname;
356 disabled = isPy27 || (pythonAtLeast "3.8");
357
358 src = bazel-build.python;
359
360 # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
361 # and the propagated input tensorflow-tensorboard, which causes environment collisions.
362 # Another possibility would be to have tensorboard only in the buildInputs
363 # https://github.com/tensorflow/tensorflow/blob/v1.7.1/tensorflow/tools/pip_package/setup.py#L79
364 postInstall = ''
365 rm $out/bin/tensorboard
366 '';
367
368 setupPyGlobalFlags = [ "--project_name ${pname}" ];
369
370 # tensorflow/tools/pip_package/setup.py
371 propagatedBuildInputs = [
372 absl-py
373 astor
374 gast
375 google-pasta
376 keras-applications
377 keras-preprocessing
378 numpy
379 six
380 protobuf
381 tensorflow-estimator_1
382 termcolor
383 wrapt
384 grpcio
385 opt-einsum
386 ] ++ lib.optionals (!isPy3k) [
387 mock
388 future
389 functools32
390 ] ++ lib.optionals (pythonOlder "3.4") [
391 backports_weakref enum34
392 ] ++ lib.optionals withTensorboard [
393 tensorflow-tensorboard_1
394 ];
395
396 nativeBuildInputs = lib.optional cudaSupport addOpenGLRunpath;
397
398 postFixup = lib.optionalString cudaSupport ''
399 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
400 addOpenGLRunpath "$lib"
401 done
402 '';
403
404 # Actual tests are slow and impure.
405 # TODO try to run them anyway
406 # TODO better test (files in tensorflow/tools/ci_build/builds/*test)
407 checkPhase = ''
408 ${python.interpreter} <<EOF
409 # A simple "Hello world"
410 import tensorflow as tf
411 hello = tf.constant("Hello, world!")
412 sess = tf.Session()
413 sess.run(hello)
414
415 # Fit a simple model to random data
416 import numpy as np
417 np.random.seed(0)
418 tf.random.set_random_seed(0)
419 model = tf.keras.models.Sequential([
420 tf.keras.layers.Dense(1, activation="linear")
421 ])
422 model.compile(optimizer="sgd", loss="mse")
423
424 x = np.random.uniform(size=(1,1))
425 y = np.random.uniform(size=(1,))
426 model.fit(x, y, epochs=1)
427
428 # regression test for #77626
429 from tensorflow.contrib import tensor_forest
430 EOF
431 '';
432
433 passthru.libtensorflow = bazel-build.out;
434
435 inherit (bazel-build) meta;
436}