tjh.dev
1{ stdenv
2, lib
3, fetchurl
4, buildPythonPackage
5, isPy3k, pythonOlder, isPy38
6, astor
7, gast
8, google-pasta
9, wrapt
10, numpy
11, six
12, termcolor
13, protobuf
14, absl-py
15, grpcio
16, mock
17, backports_weakref
18, tensorflow-estimator_1
19, tensorflow-tensorboard
20, cudaSupport ? false
21, cudatoolkit ? null
22, cudnn ? null
23, nvidia_x11 ? null
24, zlib
25, python
26, symlinkJoin
27, keras-applications
28, keras-preprocessing
29, addOpenGLRunpath
30}:
31
32# We keep this binary build for two reasons:
33# - the source build doesn't work on Darwin.
34# - the source build is currently brittle and not easy to maintain
35
36assert cudaSupport -> cudatoolkit != null
37 && cudnn != null
38 && nvidia_x11 != null;
39
40# unsupported combination
41assert ! (stdenv.isDarwin && cudaSupport);
42
43let
44 packages = import ./binary-hashes.nix;
45
46 variant = if cudaSupport then "-gpu" else "";
47 pname = "tensorflow${variant}";
48
49in buildPythonPackage {
50 inherit pname;
51 inherit (packages) version;
52 format = "wheel";
53
54 disabled = isPy38;
55
56 src = let
57 pyVerNoDot = lib.strings.stringAsChars (x: if x == "." then "" else x) python.pythonVersion;
58 pyver = if stdenv.isDarwin then builtins.substring 0 1 pyVerNoDot else pyVerNoDot;
59 platform = if stdenv.isDarwin then "mac" else "linux";
60 unit = if cudaSupport then "gpu" else "cpu";
61 key = "${platform}_py_${pyver}_${unit}";
62 in fetchurl packages.${key};
63
64 propagatedBuildInputs = [
65 protobuf
66 numpy
67 termcolor
68 grpcio
69 six
70 astor
71 absl-py
72 gast
73 google-pasta
74 wrapt
75 tensorflow-estimator_1
76 tensorflow-tensorboard
77 keras-applications
78 keras-preprocessing
79 ] ++ lib.optional (!isPy3k) mock
80 ++ lib.optionals (pythonOlder "3.4") [ backports_weakref ];
81
82 nativeBuildInputs = lib.optional cudaSupport addOpenGLRunpath;
83
84 # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
85 # and the propageted input tensorflow-tensorboard which causes environment collisions.
86 # another possibility would be to have tensorboard only in the buildInputs
87 # https://github.com/tensorflow/tensorflow/blob/v1.7.1/tensorflow/tools/pip_package/setup.py#L79
88 postInstall = ''
89 rm $out/bin/tensorboard
90 '';
91
92 # Note that we need to run *after* the fixup phase because the
93 # libraries are loaded at runtime. If we run in preFixup then
94 # patchelf --shrink-rpath will remove the cuda libraries.
95 postFixup = let
96 rpath = stdenv.lib.makeLibraryPath
97 ([ stdenv.cc.cc.lib zlib ] ++ lib.optionals cudaSupport [ cudatoolkit.out cudatoolkit.lib cudnn nvidia_x11 ]);
98 in
99 lib.optionalString stdenv.isLinux ''
100 rrPath="$out/${python.sitePackages}/tensorflow/:$out/${python.sitePackages}/tensorflow/contrib/tensor_forest/:${rpath}"
101 internalLibPath="$out/${python.sitePackages}/tensorflow/python/_pywrap_tensorflow_internal.so"
102 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
103 patchelf --set-rpath "$rrPath" "$lib"
104 ${lib.optionalString cudaSupport ''
105 addOpenGLRunpath "$lib"
106 ''}
107 done
108 '';
109
110
111 meta = with stdenv.lib; {
112 description = "Computation using data flow graphs for scalable machine learning";
113 homepage = http://tensorflow.org;
114 license = licenses.asl20;
115 maintainers = with maintainers; [ jyp abbradar ];
116 platforms = [ "x86_64-linux" "x86_64-darwin" ];
117 # Python 2.7 build uses different string encoding.
118 # See https://github.com/NixOS/nixpkgs/pull/37044#issuecomment-373452253
119 broken = stdenv.isDarwin && !isPy3k;
120 };
121}