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