pyrox.dev
1{ stdenv, bazel_5, buildBazelPackage, isPy3k, lib, fetchFromGitHub, symlinkJoin
2, addOpenGLRunpath, fetchpatch, patchelfUnstable
3# Python deps
4, buildPythonPackage, pythonOlder, python
5# Python libraries
6, numpy, tensorboard, absl-py
7, packaging, setuptools, wheel, keras, keras-preprocessing, google-pasta
8, opt-einsum, astunparse, h5py
9, termcolor, grpcio, six, wrapt, protobuf-python, tensorflow-estimator
10, dill, flatbuffers-python, portpicker, 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, nsync
15, curl, snappy, flatbuffers-core, lmdb-core, icu, double-conversion, libpng, libjpeg_turbo, giflib, protobuf-core
16# Upstream 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, cudaPackages ? {}
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, cctools, llvmPackages_11
32}:
33
34let
35 inherit (cudaPackages) cudatoolkit cudnn nccl;
36in
37
38assert cudaSupport -> cudatoolkit != null
39 && cudnn != null;
40
41# unsupported combination
42assert ! (stdenv.isDarwin && cudaSupport);
43
44assert mklSupport -> mkl != null;
45
46let
47 withTensorboard = (pythonOlder "3.6") || tensorboardSupport;
48
49 cudatoolkit_joined = symlinkJoin {
50 name = "${cudatoolkit.name}-merged";
51 paths = [
52 cudatoolkit.lib
53 cudatoolkit.out
54 ] ++ lib.optionals (lib.versionOlder cudatoolkit.version "11") [
55 # for some reason some of the required libs are in the targets/x86_64-linux
56 # directory; not sure why but this works around it
57 "${cudatoolkit}/targets/${stdenv.system}"
58 ];
59 };
60
61 cudatoolkit_cc_joined = symlinkJoin {
62 name = "${cudatoolkit.cc.name}-merged";
63 paths = [
64 cudatoolkit.cc
65 binutils.bintools # for ar, dwp, nm, objcopy, objdump, strip
66 ];
67 };
68
69 # Needed for _some_ system libraries, grep INCLUDEDIR.
70 includes_joined = symlinkJoin {
71 name = "tensorflow-deps-merged";
72 paths = [
73 jsoncpp
74 ];
75 };
76
77 tfFeature = x: if x then "1" else "0";
78
79 version = "2.10.0";
80 variant = if cudaSupport then "-gpu" else "";
81 pname = "tensorflow${variant}";
82
83 pythonEnv = python.withPackages (_:
84 [ # python deps needed during wheel build time (not runtime, see the buildPythonPackage part for that)
85 # This list can likely be shortened, but each trial takes multiple hours so won't bother for now.
86 absl-py
87 astunparse
88 dill
89 flatbuffers-python
90 gast
91 google-pasta
92 grpcio
93 h5py
94 keras-preprocessing
95 numpy
96 opt-einsum
97 packaging
98 protobuf-python
99 setuptools
100 six
101 tblib
102 tensorboard
103 tensorflow-estimator
104 termcolor
105 typing-extensions
106 wheel
107 wrapt
108 ]);
109
110 rules_cc_darwin_patched = stdenv.mkDerivation {
111 name = "rules_cc-${pname}-${version}";
112
113 src = _bazel-build.deps;
114
115 prePatch = "pushd rules_cc";
116 patches = [
117 # https://github.com/bazelbuild/rules_cc/issues/122
118 (fetchpatch {
119 name = "tensorflow-rules_cc-libtool-path.patch";
120 url = "https://github.com/bazelbuild/rules_cc/commit/8c427ab30bf213630dc3bce9d2e9a0e29d1787db.diff";
121 sha256 = "sha256-C4v6HY5+jm0ACUZ58gBPVejCYCZfuzYKlHZ0m2qDHCk=";
122 })
123
124 # https://github.com/bazelbuild/rules_cc/pull/124
125 (fetchpatch {
126 name = "tensorflow-rules_cc-install_name_tool-path.patch";
127 url = "https://github.com/bazelbuild/rules_cc/commit/156497dc89100db8a3f57b23c63724759d431d05.diff";
128 sha256 = "sha256-NES1KeQmMiUJQVoV6dS4YGRxxkZEjOpFSCyOq9HZYO0=";
129 })
130 ];
131 postPatch = "popd";
132
133 dontConfigure = true;
134 dontBuild = true;
135
136 installPhase = ''
137 runHook preInstall
138
139 mv rules_cc/ "$out"
140
141 runHook postInstall
142 '';
143 };
144 llvm-raw_darwin_patched = stdenv.mkDerivation {
145 name = "llvm-raw-${pname}-${version}";
146
147 src = _bazel-build.deps;
148
149 prePatch = "pushd llvm-raw";
150 patches = [
151 # Fix a vendored config.h that requires the 10.13 SDK
152 ./llvm_bazel_fix_macos_10_12_sdk.patch
153 ];
154 postPatch = ''
155 touch {BUILD,WORKSPACE}
156 popd
157 '';
158
159 dontConfigure = true;
160 dontBuild = true;
161
162 installPhase = ''
163 runHook preInstall
164
165 mv llvm-raw/ "$out"
166
167 runHook postInstall
168 '';
169 };
170 bazel-build = if stdenv.isDarwin then _bazel-build.overrideAttrs (prev: {
171 bazelBuildFlags = prev.bazelBuildFlags ++ [
172 "--override_repository=rules_cc=${rules_cc_darwin_patched}"
173 "--override_repository=llvm-raw=${llvm-raw_darwin_patched}"
174 ];
175 preBuild = ''
176 export AR="${cctools}/bin/libtool"
177 '';
178 }) else _bazel-build;
179
180 _bazel-build = (buildBazelPackage.override (lib.optionalAttrs stdenv.isDarwin {
181 # clang 7 fails to emit a symbol for
182 # __ZN4llvm11SmallPtrSetIPKNS_10AllocaInstELj8EED1Ev in any of the
183 # translation units, so the build fails at link time
184 stdenv = llvmPackages_11.stdenv;
185 })) {
186 name = "${pname}-${version}";
187 bazel = bazel_5;
188
189 src = fetchFromGitHub {
190 owner = "tensorflow";
191 repo = "tensorflow";
192 rev = "v${version}";
193 hash = "sha256-Y6cujiMoQXKQlsLBr7d0T278ltdd00IfsTRycJbRVN4=";
194 };
195
196 # On update, it can be useful to steal the changes from gentoo
197 # https://gitweb.gentoo.org/repo/gentoo.git/tree/sci-libs/tensorflow
198
199 nativeBuildInputs = [
200 which pythonEnv cython perl protobuf-core
201 ] ++ lib.optional cudaSupport addOpenGLRunpath;
202
203 buildInputs = [
204 jemalloc
205 mpi
206 glibcLocales
207 git
208
209 # libs taken from system through the TF_SYS_LIBS mechanism
210 boringssl
211 curl
212 double-conversion
213 flatbuffers-core
214 giflib
215 grpc
216 icu
217 jsoncpp
218 libjpeg_turbo
219 libpng
220 lmdb-core
221 pybind11
222 snappy
223 sqlite
224 ] ++ lib.optionals cudaSupport [
225 cudatoolkit
226 cudnn
227 ] ++ lib.optionals mklSupport [
228 mkl
229 ] ++ lib.optionals stdenv.isDarwin [
230 Foundation
231 Security
232 ] ++ lib.optionals (!stdenv.isDarwin) [
233 nsync
234 ];
235
236 # arbitrarily set to the current latest bazel version, overly careful
237 TF_IGNORE_MAX_BAZEL_VERSION = true;
238
239 # Take as many libraries from the system as possible. Keep in sync with
240 # list of valid syslibs in
241 # https://github.com/tensorflow/tensorflow/blob/master/third_party/systemlibs/syslibs_configure.bzl
242 TF_SYSTEM_LIBS = lib.concatStringsSep "," ([
243 "absl_py"
244 "astor_archive"
245 "astunparse_archive"
246 "boringssl"
247 # Not packaged in nixpkgs
248 # "com_github_googleapis_googleapis"
249 # "com_github_googlecloudplatform_google_cloud_cpp"
250 "com_github_grpc_grpc"
251 "com_google_protobuf"
252 # 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&)'
253 # "com_googlesource_code_re2"
254 "curl"
255 "cython"
256 "dill_archive"
257 "double_conversion"
258 "flatbuffers"
259 "functools32_archive"
260 "gast_archive"
261 "gif"
262 "hwloc"
263 "icu"
264 "jsoncpp_git"
265 "libjpeg_turbo"
266 "lmdb"
267 "nasm"
268 "opt_einsum_archive"
269 "org_sqlite"
270 "pasta"
271 "png"
272 "pybind11"
273 "six_archive"
274 "snappy"
275 "tblib_archive"
276 "termcolor_archive"
277 "typing_extensions_archive"
278 "wrapt"
279 "zlib"
280 ] ++ lib.optionals (!stdenv.isDarwin) [
281 "nsync" # fails to build on darwin
282 ]);
283
284 INCLUDEDIR = "${includes_joined}/include";
285
286 # This is needed for the Nix-provided protobuf dependency to work,
287 # as otherwise the rule `link_proto_files` tries to create the links
288 # to `/usr/include/...` which results in build failures.
289 PROTOBUF_INCLUDE_PATH = "${protobuf-core}/include";
290
291 PYTHON_BIN_PATH = pythonEnv.interpreter;
292
293 TF_NEED_GCP = true;
294 TF_NEED_HDFS = true;
295 TF_ENABLE_XLA = tfFeature xlaSupport;
296
297 CC_OPT_FLAGS = " ";
298
299 # https://github.com/tensorflow/tensorflow/issues/14454
300 TF_NEED_MPI = tfFeature cudaSupport;
301
302 TF_NEED_CUDA = tfFeature cudaSupport;
303 TF_CUDA_PATHS = lib.optionalString cudaSupport "${cudatoolkit_joined},${cudnn},${nccl}";
304 GCC_HOST_COMPILER_PREFIX = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin";
305 GCC_HOST_COMPILER_PATH = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin/gcc";
306 TF_CUDA_COMPUTE_CAPABILITIES = lib.concatStringsSep "," cudaCapabilities;
307
308 postPatch = ''
309 # bazel 3.3 should work just as well as bazel 3.1
310 rm -f .bazelversion
311 '' + lib.optionalString (!withTensorboard) ''
312 # Tensorboard pulls in a bunch of dependencies, some of which may
313 # include security vulnerabilities. So we make it optional.
314 # https://github.com/tensorflow/tensorflow/issues/20280#issuecomment-400230560
315 sed -i '/tensorboard ~=/d' tensorflow/tools/pip_package/setup.py
316 '';
317
318 # https://github.com/tensorflow/tensorflow/pull/39470
319 NIX_CFLAGS_COMPILE = [ "-Wno-stringop-truncation" ];
320
321 preConfigure = let
322 opt_flags = []
323 ++ lib.optionals sse42Support ["-msse4.2"]
324 ++ lib.optionals avx2Support ["-mavx2"]
325 ++ lib.optionals fmaSupport ["-mfma"];
326 in ''
327 patchShebangs configure
328
329 # dummy ldconfig
330 mkdir dummy-ldconfig
331 echo "#!${stdenv.shell}" > dummy-ldconfig/ldconfig
332 chmod +x dummy-ldconfig/ldconfig
333 export PATH="$PWD/dummy-ldconfig:$PATH"
334
335 export PYTHON_LIB_PATH="$NIX_BUILD_TOP/site-packages"
336 export CC_OPT_FLAGS="${lib.concatStringsSep " " opt_flags}"
337 mkdir -p "$PYTHON_LIB_PATH"
338
339 # To avoid mixing Python 2 and Python 3
340 unset PYTHONPATH
341 '';
342
343 configurePhase = ''
344 runHook preConfigure
345 ./configure
346 runHook postConfigure
347 '';
348
349 hardeningDisable = [ "format" ];
350
351 bazelBuildFlags = [
352 "--config=opt" # optimize using the flags set in the configure phase
353 ]
354 ++ lib.optionals stdenv.cc.isClang [
355 "--cxxopt=-x" "--cxxopt=c++"
356 "--host_cxxopt=-x" "--host_cxxopt=c++"
357
358 # workaround for https://github.com/bazelbuild/bazel/issues/15359
359 "--spawn_strategy=sandboxed"
360 ]
361 ++ lib.optionals (mklSupport) [ "--config=mkl" ];
362
363 bazelTarget = "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow";
364
365 removeRulesCC = false;
366 # Without this Bazel complaints about sandbox violations.
367 dontAddBazelOpts = true;
368
369 fetchAttrs = {
370 # cudaSupport causes fetch of ncclArchive, resulting in different hashes
371 sha256 = if cudaSupport then
372 "sha256-KtVReqHL3zxE8TPrqIerSOt59Mgke/ftoFZKMzgX/u8="
373 else
374 if stdenv.isDarwin then
375 # FIXME: this checksum is currently wrong, since the tensorflow dependency fetch is broken on darwin
376 "sha256-j2k9Q+k41nq5nP1VjjkkNjXRov1uAda4RCMDMAthjr0="
377 else
378 "sha256-zH3xNFEU2JR0Ww8bpD4mCiorGtao0WVPP4vklVMgS4A=";
379 };
380
381 buildAttrs = {
382 outputs = [ "out" "python" ];
383
384 preBuild = ''
385 patchShebangs .
386 '';
387
388 installPhase = ''
389 mkdir -p "$out"
390 tar -xf bazel-bin/tensorflow/tools/lib_package/libtensorflow.tar.gz -C "$out"
391 # Write pkgconfig file.
392 mkdir "$out/lib/pkgconfig"
393 cat > "$out/lib/pkgconfig/tensorflow.pc" << EOF
394 Name: TensorFlow
395 Version: ${version}
396 Description: Library for computation using data flow graphs for scalable machine learning
397 Requires:
398 Libs: -L$out/lib -ltensorflow
399 Cflags: -I$out/include/tensorflow
400 EOF
401
402 # build the source code, then copy it to $python (build_pip_package
403 # actually builds a symlink farm so we must dereference them).
404 bazel-bin/tensorflow/tools/pip_package/build_pip_package --src "$PWD/dist"
405 cp -Lr "$PWD/dist" "$python"
406 '';
407
408 postFixup = lib.optionalString cudaSupport ''
409 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
410 addOpenGLRunpath "$lib"
411 done
412 '';
413
414 requiredSystemFeatures = [
415 "big-parallel"
416 ];
417 };
418
419 meta = with lib; {
420 description = "Computation using data flow graphs for scalable machine learning";
421 homepage = "http://tensorflow.org";
422 license = licenses.asl20;
423 maintainers = with maintainers; [ jyp abbradar ];
424 platforms = with platforms; linux ++ darwin;
425 broken = !(xlaSupport -> cudaSupport);
426 } // lib.optionalAttrs stdenv.isDarwin {
427 timeout = 86400; # 24 hours
428 maxSilent = 14400; # 4h, double the default of 7200s
429 };
430 };
431
432in buildPythonPackage {
433 inherit version pname;
434 disabled = !isPy3k;
435
436 src = bazel-build.python;
437
438 # Adjust dependency requirements:
439 # - Relax flatbuffers and gast version requirements
440 # - The purpose of python3Packages.libclang is not clear at the moment and we don't have it packaged yet
441 # - keras and tensorlow-io-gcs-filesystem will be considered as optional for now.
442 postPatch = ''
443 sed -i setup.py \
444 -e "s/'flatbuffers[^']*',/'flatbuffers',/" \
445 -e "s/'gast[^']*',/'gast',/" \
446 -e "/'libclang[^']*',/d" \
447 -e "/'keras[^']*')\?,/d" \
448 -e "/'tensorflow-io-gcs-filesystem[^']*',/d" \
449 -e "s/'protobuf[^']*',/'protobuf',/" \
450 '';
451
452 # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
453 # and the propagated input tensorboard, which causes environment collisions.
454 # Another possibility would be to have tensorboard only in the buildInputs
455 # https://github.com/tensorflow/tensorflow/blob/v1.7.1/tensorflow/tools/pip_package/setup.py#L79
456 postInstall = ''
457 rm $out/bin/tensorboard
458 '';
459
460 setupPyGlobalFlags = [ "--project_name ${pname}" ];
461
462 # tensorflow/tools/pip_package/setup.py
463 propagatedBuildInputs = [
464 absl-py
465 astunparse
466 flatbuffers-python
467 gast
468 google-pasta
469 grpcio
470 h5py
471 keras-preprocessing
472 numpy
473 opt-einsum
474 packaging
475 protobuf-python
476 six
477 tensorflow-estimator
478 termcolor
479 typing-extensions
480 wrapt
481 ] ++ lib.optionals withTensorboard [
482 tensorboard
483 ];
484
485 # remove patchelfUnstable once patchelf 0.14 with https://github.com/NixOS/patchelf/pull/256 becomes the default
486 nativeBuildInputs = lib.optionals cudaSupport [ addOpenGLRunpath patchelfUnstable ];
487
488 postFixup = lib.optionalString cudaSupport ''
489 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
490 addOpenGLRunpath "$lib"
491
492 patchelf --set-rpath "${cudatoolkit}/lib:${cudatoolkit.lib}/lib:${cudnn}/lib:${nccl}/lib:$(patchelf --print-rpath "$lib")" "$lib"
493 done
494 '';
495
496 # Actual tests are slow and impure.
497 # TODO try to run them anyway
498 # TODO better test (files in tensorflow/tools/ci_build/builds/*test)
499 # TEST_PACKAGES in tensorflow/tools/pip_package/setup.py
500 checkInputs = [
501 dill
502 keras
503 portpicker
504 tblib
505 ];
506 checkPhase = ''
507 ${python.interpreter} <<EOF
508 # A simple "Hello world"
509 import tensorflow as tf
510 hello = tf.constant("Hello, world!")
511 tf.print(hello)
512
513 # Fit a simple model to random data
514 import numpy as np
515 np.random.seed(0)
516 tf.random.set_seed(0)
517 model = tf.keras.models.Sequential([
518 tf.keras.layers.Dense(1, activation="linear")
519 ])
520 model.compile(optimizer="sgd", loss="mse")
521
522 x = np.random.uniform(size=(1,1))
523 y = np.random.uniform(size=(1,))
524 model.fit(x, y, epochs=1)
525 EOF
526 '';
527 # Regression test for #77626 removed because not more `tensorflow.contrib`.
528
529 passthru = {
530 inherit cudaPackages;
531 deps = bazel-build.deps;
532 libtensorflow = bazel-build.out;
533 };
534
535 inherit (bazel-build) meta;
536}