pyrox.dev
1{ stdenv
2, lib
3, fetchurl
4, buildPythonPackage
5, isPy3k, pythonOlder, pythonAtLeast, astor
6, gast
7, google-pasta
8, wrapt
9, numpy
10, six
11, termcolor
12, protobuf
13, absl-py
14, grpcio
15, mock
16, scipy
17, wheel
18, opt-einsum
19, backports_weakref
20, tensorflow-estimator_2
21, tensorflow-tensorboard_2
22, cudaSupport ? false
23, cudatoolkit
24, cudnn
25, zlib
26, python
27, keras-applications
28, keras-preprocessing
29, addOpenGLRunpath
30, astunparse
31, flatbuffers
32, h5py
33, typing-extensions
34}:
35
36# We keep this binary build for two reasons:
37# - the source build doesn't work on Darwin.
38# - the source build is currently brittle and not easy to maintain
39
40# unsupported combination
41assert ! (stdenv.isDarwin && cudaSupport);
42
43let
44 packages = import ./binary-hashes.nix;
45 metadataPatch = ./relax-dependencies-metadata.patch;
46 patch = ./relax-dependencies.patch;
47in buildPythonPackage {
48 pname = "tensorflow" + lib.optionalString cudaSupport "-gpu";
49 inherit (packages) version;
50 format = "wheel";
51
52 disabled = pythonAtLeast "3.9";
53
54 src = let
55 pyVerNoDot = lib.strings.stringAsChars (x: if x == "." then "" else x) python.pythonVersion;
56 platform = if stdenv.isDarwin then "mac" else "linux";
57 unit = if cudaSupport then "gpu" else "cpu";
58 key = "${platform}_py_${pyVerNoDot}_${unit}";
59 in fetchurl packages.${key};
60
61 propagatedBuildInputs = [
62 astunparse
63 flatbuffers
64 typing-extensions
65 protobuf
66 numpy
67 scipy
68 termcolor
69 grpcio
70 six
71 astor
72 absl-py
73 gast
74 opt-einsum
75 google-pasta
76 wrapt
77 tensorflow-estimator_2
78 tensorflow-tensorboard_2
79 keras-applications
80 keras-preprocessing
81 h5py
82 ] ++ lib.optional (!isPy3k) mock
83 ++ lib.optionals (pythonOlder "3.4") [ backports_weakref ];
84
85 nativeBuildInputs = [ wheel ] ++ lib.optional cudaSupport addOpenGLRunpath;
86
87 preConfigure = ''
88 unset SOURCE_DATE_EPOCH
89
90 # Make sure that dist and the wheel file are writable.
91 chmod u+rwx -R ./dist
92
93 pushd dist
94
95 wheel unpack --dest unpacked ./*.whl
96 (
97 cd unpacked/tensorflow*
98 # relax too strict versions in setup.py
99 patch -p 1 < ${patch}
100 cd *.dist-info
101 # relax too strict versions in *.dist-info/METADATA
102 patch -p 3 < ${metadataPatch}
103 )
104 wheel pack ./unpacked/tensorflow*
105
106 popd
107 '';
108
109 # Note that we need to run *after* the fixup phase because the
110 # libraries are loaded at runtime. If we run in preFixup then
111 # patchelf --shrink-rpath will remove the cuda libraries.
112 postFixup =
113 let
114 # rpaths we only need to add if CUDA is enabled.
115 cudapaths = lib.optionals cudaSupport [
116 cudatoolkit.out
117 cudatoolkit.lib
118 cudnn
119 ];
120
121 libpaths = [
122 stdenv.cc.cc.lib
123 zlib
124 ];
125
126 rpath = lib.makeLibraryPath (libpaths ++ cudapaths);
127 in
128 lib.optionalString stdenv.isLinux ''
129 # This is an array containing all the directories in the tensorflow2
130 # package that contain .so files.
131 #
132 # TODO: Create this list programmatically, and remove paths that aren't
133 # actually needed.
134 rrPathArr=(
135 "$out/${python.sitePackages}/tensorflow/"
136 "$out/${python.sitePackages}/tensorflow/core/kernels"
137 "$out/${python.sitePackages}/tensorflow/compiler/tf2tensorrt/"
138 "$out/${python.sitePackages}/tensorflow/compiler/tf2xla/ops/"
139 "$out/${python.sitePackages}/tensorflow/lite/experimental/microfrontend/python/ops/"
140 "$out/${python.sitePackages}/tensorflow/lite/python/interpreter_wrapper/"
141 "$out/${python.sitePackages}/tensorflow/lite/python/optimize/"
142 "$out/${python.sitePackages}/tensorflow/python/"
143 "$out/${python.sitePackages}/tensorflow/python/framework/"
144 "$out/${python.sitePackages}/tensorflow/python/autograph/impl/testing"
145 "$out/${python.sitePackages}/tensorflow/python/data/experimental/service"
146 "$out/${python.sitePackages}/tensorflow/python/framework"
147 "$out/${python.sitePackages}/tensorflow/python/profiler/internal"
148 "${rpath}"
149 )
150
151 # The the bash array into a colon-separated list of RPATHs.
152 rrPath=$(IFS=$':'; echo "''${rrPathArr[*]}")
153 echo "about to run patchelf with the following rpath: $rrPath"
154
155 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
156 echo "about to patchelf $lib..."
157 chmod a+rx "$lib"
158 patchelf --set-rpath "$rrPath" "$lib"
159 ${lib.optionalString cudaSupport ''
160 addOpenGLRunpath "$lib"
161 ''}
162 done
163 '';
164
165 # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
166 # and the propagated input tensorflow-tensorboard, which causes environment collisions.
167 # Another possibility would be to have tensorboard only in the buildInputs
168 # See https://github.com/NixOS/nixpkgs/pull/44381 for more information.
169 postInstall = ''
170 rm $out/bin/tensorboard
171 '';
172
173 pythonImportsCheck = [
174 "tensorflow"
175 "tensorflow.keras"
176 "tensorflow.python"
177 "tensorflow.python.framework"
178 ];
179
180 meta = with lib; {
181 description = "Computation using data flow graphs for scalable machine learning";
182 homepage = "http://tensorflow.org";
183 license = licenses.asl20;
184 maintainers = with maintainers; [ jyp abbradar cdepillabout ];
185 platforms = [ "x86_64-linux" "x86_64-darwin" ];
186 };
187}