tjh.dev
1{ stdenv, bazel_3, buildBazelPackage, isPy3k, lib, fetchFromGitHub, symlinkJoin
2, addOpenGLRunpath
3# Python deps
4, buildPythonPackage, pythonOlder, pythonAtLeast, python
5# Python libraries
6, numpy, tensorflow-tensorboard_2, absl-py
7, future, setuptools, wheel, keras-preprocessing, google-pasta
8, opt-einsum, astunparse, h5py
9, termcolor, grpcio, six, wrapt, protobuf, tensorflow-estimator_2
10, dill, flatbuffers-python, tblib, typing-extensions
11# Common deps
12, git, pybind11, which, binutils, glibcLocales, cython, perl
13# Common libraries
14, jemalloc, mpi, gast, grpc, sqlite, boringssl, jsoncpp
15, curl, snappy, flatbuffers-core, lmdb-core, icu, double-conversion, libpng, libjpeg_turbo, 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, cudatoolkit ? null, cudnn ? null, nccl ? null
21, mklSupport ? false, mkl ? null
22, tensorboardSupport ? true
23# XLA without CUDA is broken
24, xlaSupport ? cudaSupport
25# Default from ./configure script
26, cudaCapabilities ? [ "sm_35" "sm_50" "sm_60" "sm_70" "sm_75" "compute_80" ]
27, sse42Support ? stdenv.hostPlatform.sse4_2Support
28, avx2Support ? stdenv.hostPlatform.avx2Support
29, fmaSupport ? stdenv.hostPlatform.fmaSupport
30# Darwin deps
31, Foundation, Security
32}:
33
34assert cudaSupport -> 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") || tensorboardSupport;
44
45 cudatoolkit_joined = symlinkJoin {
46 name = "${cudatoolkit.name}-merged";
47 paths = [
48 cudatoolkit.lib
49 cudatoolkit.out
50 ] ++ lib.optionals (lib.versionOlder cudatoolkit.version "11") [
51 # for some reason some of the required libs are in the targets/x86_64-linux
52 # directory; not sure why but this works around it
53 "${cudatoolkit}/targets/${stdenv.system}"
54 ];
55 };
56
57 cudatoolkit_cc_joined = symlinkJoin {
58 name = "${cudatoolkit.cc.name}-merged";
59 paths = [
60 cudatoolkit.cc
61 binutils.bintools # for ar, dwp, nm, objcopy, objdump, strip
62 ];
63 };
64
65 # Needed for _some_ system libraries, grep INCLUDEDIR.
66 includes_joined = symlinkJoin {
67 name = "tensorflow-deps-merged";
68 paths = [
69 jsoncpp
70 ];
71 };
72
73 tfFeature = x: if x then "1" else "0";
74
75 version = "2.4.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 # This list can likely be shortened, but each trial takes multiple hours so won't bother for now.
82 absl-py
83 astunparse
84 dill
85 flatbuffers-python
86 gast
87 google-pasta
88 grpcio
89 h5py
90 keras-preprocessing
91 numpy
92 opt-einsum
93 protobuf
94 setuptools
95 six
96 tblib
97 tensorflow-estimator_2
98 tensorflow-tensorboard_2
99 termcolor
100 typing-extensions
101 wheel
102 wrapt
103 ]);
104
105 bazel-build = buildBazelPackage {
106 name = "${pname}-${version}";
107 bazel = bazel_3;
108
109 src = fetchFromGitHub {
110 owner = "tensorflow";
111 repo = "tensorflow";
112 rev = "v${version}";
113 sha256 = "0yl06aypfxrcs35828xf04mkidz1x0j89v0q5h4d2xps1cb5rv3f";
114 };
115
116 patches = [
117 # Relax too strict Python packages versions dependencies.
118 ./relax-dependencies.patch
119 # Add missing `io_bazel_rules_docker` dependency.
120 ./workspace.patch
121 ];
122
123 # On update, it can be useful to steal the changes from gentoo
124 # https://gitweb.gentoo.org/repo/gentoo.git/tree/sci-libs/tensorflow
125
126 nativeBuildInputs = [
127 which pythonEnv cython perl
128 ] ++ lib.optional cudaSupport addOpenGLRunpath;
129
130 buildInputs = [
131 jemalloc
132 mpi
133 glibcLocales
134 git
135
136 # libs taken from system through the TF_SYS_LIBS mechanism
137 grpc
138 sqlite
139 boringssl
140 jsoncpp
141 curl
142 pybind11
143 snappy
144 flatbuffers-core
145 icu
146 double-conversion
147 libpng
148 libjpeg_turbo
149 giflib
150 lmdb-core
151 ] ++ lib.optionals cudaSupport [
152 cudatoolkit
153 cudnn
154 ] ++ lib.optionals mklSupport [
155 mkl
156 ] ++ lib.optionals stdenv.isDarwin [
157 Foundation
158 Security
159 ];
160
161 # arbitrarily set to the current latest bazel version, overly careful
162 TF_IGNORE_MAX_BAZEL_VERSION = true;
163
164 # Take as many libraries from the system as possible. Keep in sync with
165 # list of valid syslibs in
166 # https://github.com/tensorflow/tensorflow/blob/master/third_party/systemlibs/syslibs_configure.bzl
167 TF_SYSTEM_LIBS = lib.concatStringsSep "," [
168 "absl_py"
169 "astor_archive"
170 "astunparse_archive"
171 "boringssl"
172 # Not packaged in nixpkgs
173 # "com_github_googleapis_googleapis"
174 # "com_github_googlecloudplatform_google_cloud_cpp"
175 "com_github_grpc_grpc"
176 # Multiple issues with custom protobuf.
177 # First `com_github_googleapis` fails to configure. Can be worked around by disabling `com_github_googleapis`
178 # and related functionality, but then the next error is about "dangling symbolic link", and in general
179 # looks like that's only the beginning: see
180 # https://stackoverflow.com/questions/55578884/how-to-build-tensorflow-1-13-1-with-custom-protobuf
181 # "com_google_protobuf"
182 # Fails with the error: external/org_tensorflow/tensorflow/core/profiler/utils/tf_op_utils.cc:46:49: error: no matching function for call to 're2::RE2::FullMatch(absl::lts_2020_02_25::string_view&, re2::RE2&)'
183 # "com_googlesource_code_re2"
184 "curl"
185 "cython"
186 "dill_archive"
187 "double_conversion"
188 "enum34_archive"
189 "flatbuffers"
190 "functools32_archive"
191 "gast_archive"
192 "gif"
193 "hwloc"
194 "icu"
195 "jsoncpp_git"
196 "libjpeg_turbo"
197 "lmdb"
198 "nasm"
199 # "nsync" # not packaged in nixpkgs
200 "opt_einsum_archive"
201 "org_sqlite"
202 "pasta"
203 "pcre"
204 "png"
205 "pybind11"
206 "six_archive"
207 "snappy"
208 "tblib_archive"
209 "termcolor_archive"
210 "typing_extensions_archive"
211 "wrapt"
212 "zlib"
213 ];
214
215 INCLUDEDIR = "${includes_joined}/include";
216
217 PYTHON_BIN_PATH = pythonEnv.interpreter;
218
219 TF_NEED_GCP = true;
220 TF_NEED_HDFS = true;
221 TF_ENABLE_XLA = tfFeature xlaSupport;
222
223 CC_OPT_FLAGS = " ";
224
225 # https://github.com/tensorflow/tensorflow/issues/14454
226 TF_NEED_MPI = tfFeature cudaSupport;
227
228 TF_NEED_CUDA = tfFeature cudaSupport;
229 TF_CUDA_PATHS = lib.optionalString cudaSupport "${cudatoolkit_joined},${cudnn},${nccl}";
230 GCC_HOST_COMPILER_PREFIX = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin";
231 GCC_HOST_COMPILER_PATH = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin/gcc";
232 TF_CUDA_COMPUTE_CAPABILITIES = lib.concatStringsSep "," cudaCapabilities;
233
234 postPatch = ''
235 # bazel 3.3 should work just as well as bazel 3.1
236 rm -f .bazelversion
237 '' + lib.optionalString (!withTensorboard) ''
238 # Tensorboard pulls in a bunch of dependencies, some of which may
239 # include security vulnerabilities. So we make it optional.
240 # https://github.com/tensorflow/tensorflow/issues/20280#issuecomment-400230560
241 sed -i '/tensorboard ~=/d' tensorflow/tools/pip_package/setup.py
242 '';
243
244 # https://github.com/tensorflow/tensorflow/pull/39470
245 NIX_CFLAGS_COMPILE = [ "-Wno-stringop-truncation" ];
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 hardeningDisable = [ "format" ];
276
277 bazelBuildFlags = [
278 "--config=opt" # optimize using the flags set in the configure phase
279 ]
280 ++ lib.optionals (mklSupport) [ "--config=mkl" ];
281
282 bazelTarget = "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow";
283
284 removeRulesCC = false;
285 # Without this Bazel complaints about sandbox violations.
286 dontAddBazelOpts = true;
287
288 fetchAttrs = {
289 # cudaSupport causes fetch of ncclArchive, resulting in different hashes
290 sha256 = if cudaSupport then
291 "1i7z2a7bc2q1vn1h9nx1xc6g1r1cby2xvbcs20fj9h6c2fgaw9j4"
292 else
293 "0s8q5rxq8abr50c5jpwv96ncfc0k8jw7w70ri8viqy031g9v9v45";
294 };
295
296 buildAttrs = {
297 outputs = [ "out" "python" ];
298
299 preBuild = ''
300 patchShebangs .
301 '';
302
303 installPhase = ''
304 mkdir -p "$out"
305 tar -xf bazel-bin/tensorflow/tools/lib_package/libtensorflow.tar.gz -C "$out"
306 # Write pkgconfig file.
307 mkdir "$out/lib/pkgconfig"
308 cat > "$out/lib/pkgconfig/tensorflow.pc" << EOF
309 Name: TensorFlow
310 Version: ${version}
311 Description: Library for computation using data flow graphs for scalable machine learning
312 Requires:
313 Libs: -L$out/lib -ltensorflow
314 Cflags: -I$out/include/tensorflow
315 EOF
316
317 # build the source code, then copy it to $python (build_pip_package
318 # actually builds a symlink farm so we must dereference them).
319 bazel-bin/tensorflow/tools/pip_package/build_pip_package --src "$PWD/dist"
320 cp -Lr "$PWD/dist" "$python"
321 '';
322
323 postFixup = lib.optionalString cudaSupport ''
324 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
325 addOpenGLRunpath "$lib"
326 done
327 '';
328
329 requiredSystemFeatures = [
330 "big-parallel"
331 ];
332 };
333
334 meta = with lib; {
335 description = "Computation using data flow graphs for scalable machine learning";
336 homepage = "http://tensorflow.org";
337 license = licenses.asl20;
338 maintainers = with maintainers; [ jyp abbradar ];
339 platforms = with platforms; linux ++ darwin;
340 broken = !(xlaSupport -> cudaSupport);
341 };
342 };
343
344in buildPythonPackage {
345 inherit version pname;
346 disabled = !isPy3k;
347
348 src = bazel-build.python;
349
350 # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
351 # and the propagated input tensorflow-tensorboard, which causes environment collisions.
352 # Another possibility would be to have tensorboard only in the buildInputs
353 # https://github.com/tensorflow/tensorflow/blob/v1.7.1/tensorflow/tools/pip_package/setup.py#L79
354 postInstall = ''
355 rm $out/bin/tensorboard
356 '';
357
358 setupPyGlobalFlags = [ "--project_name ${pname}" ];
359
360 # tensorflow/tools/pip_package/setup.py
361 propagatedBuildInputs = [
362 absl-py
363 astunparse
364 dill
365 flatbuffers-python
366 gast
367 google-pasta
368 grpcio
369 h5py
370 keras-preprocessing
371 numpy
372 opt-einsum
373 protobuf
374 six
375 tblib
376 tensorflow-estimator_2
377 termcolor
378 typing-extensions
379 wrapt
380 ] ++ lib.optionals withTensorboard [
381 tensorflow-tensorboard_2
382 ];
383
384 nativeBuildInputs = lib.optional cudaSupport addOpenGLRunpath;
385
386 postFixup = lib.optionalString cudaSupport ''
387 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
388 addOpenGLRunpath "$lib"
389
390 patchelf --set-rpath "${cudatoolkit}/lib:${cudatoolkit.lib}/lib:${cudnn}/lib:${nccl}/lib:$(patchelf --print-rpath "$lib")" "$lib"
391 done
392 '';
393
394 # Actual tests are slow and impure.
395 # TODO try to run them anyway
396 # TODO better test (files in tensorflow/tools/ci_build/builds/*test)
397 checkPhase = ''
398 ${python.interpreter} <<EOF
399 # A simple "Hello world"
400 import tensorflow as tf
401 hello = tf.constant("Hello, world!")
402 tf.print(hello)
403
404 # Fit a simple model to random data
405 import numpy as np
406 np.random.seed(0)
407 tf.random.set_seed(0)
408 model = tf.keras.models.Sequential([
409 tf.keras.layers.Dense(1, activation="linear")
410 ])
411 model.compile(optimizer="sgd", loss="mse")
412
413 x = np.random.uniform(size=(1,1))
414 y = np.random.uniform(size=(1,))
415 model.fit(x, y, epochs=1)
416 EOF
417 '';
418 # Regression test for #77626 removed because not more `tensorflow.contrib`.
419
420 passthru = {
421 deps = bazel-build.deps;
422 libtensorflow = bazel-build.out;
423 };
424
425 inherit (bazel-build) meta;
426}