tjh.dev
1{ stdenv, pkgs, buildBazelPackage, lib, fetchFromGitHub, fetchpatch, symlinkJoin
2, addOpenGLRunpath
3# Python deps
4, buildPythonPackage, isPy3k, pythonOlder, pythonAtLeast, python
5# Python libraries
6, numpy, tensorflow-tensorboard, backports_weakref, mock, enum34, absl-py
7, future, setuptools, wheel, keras-preprocessing, keras-applications, google-pasta
8, termcolor, grpcio, six, wrapt, protobuf, tensorflow-estimator
9# Common deps
10, git, swig, which, binutils, glibcLocales, cython
11# Common libraries
12, jemalloc, openmpi, astor, gast, grpc, sqlite, openssl, jsoncpp, re2
13, curl, snappy, flatbuffers, icu, double-conversion, libpng, libjpeg, giflib
14, cudaSupport ? false, nvidia_x11 ? null, cudatoolkit ? null, cudnn ? null, nccl ? null
15# XLA without CUDA is broken
16, xlaSupport ? cudaSupport
17# Default from ./configure script
18, cudaCapabilities ? [ "3.5" "5.2" ]
19, sse42Support ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") ["westmere" "sandybridge" "ivybridge" "haswell" "broadwell" "skylake" "skylake-avx512"]
20, avx2Support ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") [ "haswell" "broadwell" "skylake" "skylake-avx512"]
21, fmaSupport ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") [ "haswell" "broadwell" "skylake" "skylake-avx512"]
22}:
23
24assert cudaSupport -> nvidia_x11 != null
25 && cudatoolkit != null
26 && cudnn != null;
27
28# unsupported combination
29assert ! (stdenv.isDarwin && cudaSupport);
30
31let
32 withTensorboard = pythonOlder "3.6";
33
34 cudatoolkit_joined = symlinkJoin {
35 name = "${cudatoolkit.name}-merged";
36 paths = [ cudatoolkit.out cudatoolkit.lib ];
37 };
38
39 cudatoolkit_cc_joined = symlinkJoin {
40 name = "${cudatoolkit.cc.name}-merged";
41 paths = [
42 cudatoolkit.cc
43 binutils.bintools # for ar, dwp, nm, objcopy, objdump, strip
44 ];
45 };
46
47 # Needed for _some_ system libraries, grep INCLUDEDIR.
48 includes_joined = symlinkJoin {
49 name = "tensorflow-deps-merged";
50 paths = [
51 pkgs.protobuf
52 jsoncpp
53 ];
54 };
55
56 tfFeature = x: if x then "1" else "0";
57
58 version = "1.14.0";
59 variant = if cudaSupport then "-gpu" else "";
60 pname = "tensorflow${variant}";
61
62 pythonEnv = python.withPackages (_:
63 [ # python deps needed during wheel build time
64 numpy
65 keras-preprocessing
66 protobuf
67 wrapt
68 gast
69 astor
70 absl-py
71 termcolor
72 keras-applications
73 setuptools
74 wheel
75 ] ++ lib.optionals (!isPy3k)
76 [ future
77 mock
78 ]);
79
80 bazel-build = buildBazelPackage {
81 name = "${pname}-${version}";
82
83 src = fetchFromGitHub {
84 owner = "tensorflow";
85 repo = "tensorflow";
86 rev = "v${version}";
87 sha256 = "06jvwlsm14b8rqwd8q8796r0vmn0wk64s4ps2zg0sapkmp9vvcmi";
88 };
89
90 patches = [
91 # Work around https://github.com/tensorflow/tensorflow/issues/24752
92 ./no-saved-proto.patch
93 # Fixes for NixOS jsoncpp
94 ./system-jsoncpp.patch
95
96 # https://github.com/tensorflow/tensorflow/pull/29673
97 (fetchpatch {
98 name = "fix-compile-with-cuda-and-mpi.patch";
99 url = "https://github.com/tensorflow/tensorflow/pull/29673/commits/498e35a3bfe38dd75cf1416a1a23c07c3b59e6af.patch";
100 sha256 = "1m2qmwv1ysqa61z6255xggwbq6mnxbig749bdvrhnch4zydxb4di";
101 })
102
103 # https://github.com/tensorflow/tensorflow/issues/29220
104 (fetchpatch {
105 name = "bazel-0.27.patch";
106 url = "https://github.com/tensorflow/tensorflow/commit/cfccbdb8c4a92dd26382419dceb4d934c2380391.patch";
107 sha256 = "1l56wjia2c4685flsfkkgy471wx3c66wyv8khspv06zchj0k0liw";
108 })
109 ];
110
111 # On update, it can be useful to steal the changes from gentoo
112 # https://gitweb.gentoo.org/repo/gentoo.git/tree/sci-libs/tensorflow
113
114 nativeBuildInputs = [
115 swig which pythonEnv
116 ] ++ lib.optional cudaSupport addOpenGLRunpath;
117
118 buildInputs = [
119 jemalloc
120 openmpi
121 glibcLocales
122 git
123
124 # libs taken from system through the TF_SYS_LIBS mechanism
125 # grpc
126 sqlite
127 openssl
128 jsoncpp
129 pkgs.protobuf
130 curl
131 snappy
132 flatbuffers
133 icu
134 double-conversion
135 libpng
136 libjpeg
137 giflib
138 re2
139 pkgs.lmdb
140 ] ++ lib.optionals cudaSupport [
141 cudatoolkit
142 cudnn
143 nvidia_x11
144 ];
145
146 # arbitrarily set to the current latest bazel version, overly careful
147 TF_IGNORE_MAX_BAZEL_VERSION = true;
148
149 # Take as many libraries from the system as possible. Keep in sync with
150 # list of valid syslibs in
151 # https://github.com/tensorflow/tensorflow/blob/master/third_party/systemlibs/syslibs_configure.bzl
152 TF_SYSTEM_LIBS = lib.concatStringsSep "," [
153 "absl_py"
154 "astor_archive"
155 "boringssl"
156 # Not packaged in nixpkgs
157 # "com_github_googleapis_googleapis"
158 # "com_github_googlecloudplatform_google_cloud_cpp"
159 "com_google_protobuf"
160 "com_google_protobuf_cc"
161 "com_googlesource_code_re2"
162 "curl"
163 "cython"
164 "double_conversion"
165 "flatbuffers"
166 "gast_archive"
167 "gif_archive"
168 # Lots of errors, requires an older version
169 # "grpc"
170 "hwloc"
171 "icu"
172 "jpeg"
173 "jsoncpp_git"
174 "keras_applications_archive"
175 "lmdb"
176 "nasm"
177 # "nsync" # not packaged in nixpkgs
178 "org_sqlite"
179 "pasta"
180 "pcre"
181 "png_archive"
182 "protobuf_archive"
183 "six_archive"
184 "snappy"
185 "swig"
186 "termcolor_archive"
187 "wrapt"
188 "zlib_archive"
189 ];
190
191 INCLUDEDIR = "${includes_joined}/include";
192
193 PYTHON_BIN_PATH = pythonEnv.interpreter;
194
195 TF_NEED_GCP = true;
196 TF_NEED_HDFS = true;
197 TF_ENABLE_XLA = tfFeature xlaSupport;
198
199 CC_OPT_FLAGS = " ";
200
201 # https://github.com/tensorflow/tensorflow/issues/14454
202 TF_NEED_MPI = tfFeature cudaSupport;
203
204 TF_NEED_CUDA = tfFeature cudaSupport;
205 TF_CUDA_PATHS = lib.optionalString cudaSupport "${cudatoolkit_joined},${cudnn},${nccl}";
206 GCC_HOST_COMPILER_PREFIX = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin";
207 GCC_HOST_COMPILER_PATH = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin/gcc";
208 TF_CUDA_COMPUTE_CAPABILITIES = lib.concatStringsSep "," cudaCapabilities;
209
210 postPatch = ''
211 # https://github.com/tensorflow/tensorflow/issues/20919
212 sed -i '/androidndk/d' tensorflow/lite/kernels/internal/BUILD
213
214 # Tensorboard pulls in a bunch of dependencies, some of which may
215 # include security vulnerabilities. So we make it optional.
216 # https://github.com/tensorflow/tensorflow/issues/20280#issuecomment-400230560
217 sed -i '/tensorboard >=/d' tensorflow/tools/pip_package/setup.py
218 '';
219
220 preConfigure = let
221 opt_flags = []
222 ++ lib.optionals sse42Support ["-msse4.2"]
223 ++ lib.optionals avx2Support ["-mavx2"]
224 ++ lib.optionals fmaSupport ["-mfma"];
225 in ''
226 patchShebangs configure
227
228 # dummy ldconfig
229 mkdir dummy-ldconfig
230 echo "#!${stdenv.shell}" > dummy-ldconfig/ldconfig
231 chmod +x dummy-ldconfig/ldconfig
232 export PATH="$PWD/dummy-ldconfig:$PATH"
233
234 export PYTHON_LIB_PATH="$NIX_BUILD_TOP/site-packages"
235 export CC_OPT_FLAGS="${lib.concatStringsSep " " opt_flags}"
236 mkdir -p "$PYTHON_LIB_PATH"
237
238 # To avoid mixing Python 2 and Python 3
239 unset PYTHONPATH
240 '';
241
242 configurePhase = ''
243 runHook preConfigure
244 ./configure
245 runHook postConfigure
246 '';
247
248 # FIXME: Tensorflow uses dlopen() for CUDA libraries.
249 NIX_LDFLAGS = lib.optionals cudaSupport [ "-lcudart" "-lcublas" "-lcufft" "-lcurand" "-lcusolver" "-lcusparse" "-lcudnn" ];
250
251 hardeningDisable = [ "format" ];
252
253 bazelFlags = [
254 # temporary fixes to make the build work with bazel 0.27
255 "--incompatible_no_support_tools_in_action_inputs=false"
256 ];
257 bazelBuildFlags = [
258 "--config=opt" # optimize using the flags set in the configure phase
259 ];
260
261 bazelTarget = "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow";
262
263 fetchAttrs = {
264 # So that checksums don't depend on these.
265 TF_SYSTEM_LIBS = null;
266
267 # cudaSupport causes fetch of ncclArchive, resulting in different hashes
268 sha256 = if cudaSupport then
269 "196pm3ynfafqlcxah07hkvphf536hpix1ydgsynr1yg08aynlvvx"
270 else
271 "138r85n27ijzwxfwb5pcfyb79v14368jpckw0vmciz6pwf11bd9g";
272 };
273
274 buildAttrs = {
275 outputs = [ "out" "python" ];
276
277 preBuild = ''
278 patchShebangs .
279 '';
280
281 installPhase = ''
282 mkdir -p "$out"
283 tar -xf bazel-bin/tensorflow/tools/lib_package/libtensorflow.tar.gz -C "$out"
284 # Write pkgconfig file.
285 mkdir "$out/lib/pkgconfig"
286 cat > "$out/lib/pkgconfig/tensorflow.pc" << EOF
287 Name: TensorFlow
288 Version: ${version}
289 Description: Library for computation using data flow graphs for scalable machine learning
290 Requires:
291 Libs: -L$out/lib -ltensorflow
292 Cflags: -I$out/include/tensorflow
293 EOF
294
295 # build the source code, then copy it to $python (build_pip_package
296 # actually builds a symlink farm so we must dereference them).
297 bazel-bin/tensorflow/tools/pip_package/build_pip_package --src "$PWD/dist"
298 cp -Lr "$PWD/dist" "$python"
299 '';
300
301 postFixup = lib.optionalString cudaSupport ''
302 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
303 addOpenGLRunpath "$lib"
304 done
305 '';
306 };
307
308 meta = with stdenv.lib; {
309 description = "Computation using data flow graphs for scalable machine learning";
310 homepage = http://tensorflow.org;
311 license = licenses.asl20;
312 maintainers = with maintainers; [ jyp abbradar ];
313 platforms = platforms.linux;
314 broken = !(xlaSupport -> cudaSupport);
315 };
316 };
317
318in buildPythonPackage {
319 inherit version pname;
320
321 src = bazel-build.python;
322
323 # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
324 # and the propagated input tensorflow-tensorboard, which causes environment collisions.
325 # Another possibility would be to have tensorboard only in the buildInputs
326 # https://github.com/tensorflow/tensorflow/blob/v1.7.1/tensorflow/tools/pip_package/setup.py#L79
327 postInstall = ''
328 rm $out/bin/tensorboard
329 '';
330
331 setupPyGlobalFlags = [ "--project_name ${pname}" ];
332
333 # tensorflow/tools/pip_package/setup.py
334 propagatedBuildInputs = [
335 absl-py
336 astor
337 gast
338 google-pasta
339 keras-applications
340 keras-preprocessing
341 numpy
342 six
343 protobuf
344 tensorflow-estimator
345 termcolor
346 wrapt
347 grpcio
348 ] ++ lib.optionals (!isPy3k) [
349 mock
350 future # FIXME
351 ] ++ lib.optionals (pythonOlder "3.4") [
352 backports_weakref enum34
353 ] ++ lib.optionals withTensorboard [
354 tensorflow-tensorboard
355 ];
356
357 nativeBuildInputs = lib.optional cudaSupport addOpenGLRunpath;
358
359 postFixup = lib.optionalString cudaSupport ''
360 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
361 addOpenGLRunpath "$lib"
362 done
363 '';
364
365 # Actual tests are slow and impure.
366 # TODO try to run them anyway
367 # TODO better test (files in tensorflow/tools/ci_build/builds/*test)
368 checkPhase = ''
369 ${python.interpreter} -c "import tensorflow"
370 '';
371
372 passthru.libtensorflow = bazel-build.out;
373
374 inherit (bazel-build) meta;
375}