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, packaging
13, protobuf
14, absl-py
15, grpcio
16, mock
17, scipy
18, wheel
19, jax
20, opt-einsum
21, tensorflow-estimator-bin
22, tensorboard
23, config
24, cudaSupport ? config.cudaSupport
25, cudaPackages ? {}
26, zlib
27, python
28, keras-applications
29, keras-preprocessing
30, addOpenGLRunpath
31, astunparse
32, flatbuffers
33, h5py
34, typing-extensions
35}:
36
37# We keep this binary build for two reasons:
38# - the source build doesn't work on Darwin.
39# - the source build is currently brittle and not easy to maintain
40
41# unsupported combination
42assert ! (stdenv.isDarwin && cudaSupport);
43
44let
45 packages = import ./binary-hashes.nix;
46 inherit (cudaPackages) cudatoolkit cudnn;
47in buildPythonPackage {
48 pname = "tensorflow" + lib.optionalString cudaSupport "-gpu";
49 inherit (packages) version;
50 format = "wheel";
51
52 src = let
53 pyVerNoDot = lib.strings.stringAsChars (x: lib.optionalString (x != ".") x) python.pythonVersion;
54 platform = if stdenv.isDarwin then "mac" else "linux";
55 unit = if cudaSupport then "gpu" else "cpu";
56 key = "${platform}_py_${pyVerNoDot}_${unit}";
57 in fetchurl (packages.${key} or {});
58
59 propagatedBuildInputs = [
60 astunparse
61 flatbuffers
62 typing-extensions
63 packaging
64 protobuf
65 numpy
66 scipy
67 jax
68 termcolor
69 grpcio
70 six
71 astor
72 absl-py
73 gast
74 opt-einsum
75 google-pasta
76 wrapt
77 tensorflow-estimator-bin
78 tensorboard
79 keras-applications
80 keras-preprocessing
81 h5py
82 ] ++ lib.optional (!isPy3k) mock;
83
84 nativeBuildInputs = [ wheel ] ++ lib.optionals cudaSupport [ addOpenGLRunpath ];
85
86 preConfigure = ''
87 unset SOURCE_DATE_EPOCH
88
89 # Make sure that dist and the wheel file are writable.
90 chmod u+rwx -R ./dist
91
92 pushd dist
93
94 orig_name="$(echo ./*.whl)"
95 wheel unpack --dest unpacked ./*.whl
96 rm ./*.whl
97 (
98 cd unpacked/tensorflow*
99 # Adjust dependency requirements:
100 # - Relax flatbuffers, gast, protobuf, tensorboard, and tensorflow-estimator version requirements that don't match what we have packaged
101 # - The purpose of python3Packages.libclang is not clear at the moment and we don't have it packaged yet
102 # - keras and tensorlow-io-gcs-filesystem will be considered as optional for now.
103 # - numpy was pinned to fix some internal tests: https://github.com/tensorflow/tensorflow/issues/60216
104 sed -i *.dist-info/METADATA \
105 -e "/Requires-Dist: flatbuffers/d" \
106 -e "/Requires-Dist: gast/d" \
107 -e "/Requires-Dist: keras/d" \
108 -e "/Requires-Dist: libclang/d" \
109 -e "/Requires-Dist: protobuf/d" \
110 -e "/Requires-Dist: tensorboard/d" \
111 -e "/Requires-Dist: tensorflow-estimator/d" \
112 -e "/Requires-Dist: tensorflow-io-gcs-filesystem/d" \
113 -e "s/Requires-Dist: numpy (.*)/Requires-Dist: numpy/"
114 )
115 wheel pack ./unpacked/tensorflow*
116 mv *.whl $orig_name # avoid changes to the _os_arch.whl suffix
117
118 popd
119 '';
120
121 # Note that we need to run *after* the fixup phase because the
122 # libraries are loaded at runtime. If we run in preFixup then
123 # patchelf --shrink-rpath will remove the cuda libraries.
124 postFixup =
125 let
126 # rpaths we only need to add if CUDA is enabled.
127 cudapaths = lib.optionals cudaSupport [
128 cudatoolkit.out
129 cudatoolkit.lib
130 cudnn
131 ];
132
133 libpaths = [
134 stdenv.cc.cc.lib
135 zlib
136 ];
137
138 rpath = lib.makeLibraryPath (libpaths ++ cudapaths);
139 in
140 lib.optionalString stdenv.isLinux ''
141 # This is an array containing all the directories in the tensorflow2
142 # package that contain .so files.
143 #
144 # TODO: Create this list programmatically, and remove paths that aren't
145 # actually needed.
146 rrPathArr=(
147 "$out/${python.sitePackages}/tensorflow/"
148 "$out/${python.sitePackages}/tensorflow/core/kernels"
149 "$out/${python.sitePackages}/tensorflow/compiler/tf2tensorrt/"
150 "$out/${python.sitePackages}/tensorflow/compiler/tf2xla/ops/"
151 "$out/${python.sitePackages}/tensorflow/lite/experimental/microfrontend/python/ops/"
152 "$out/${python.sitePackages}/tensorflow/lite/python/analyzer_wrapper/"
153 "$out/${python.sitePackages}/tensorflow/lite/python/interpreter_wrapper/"
154 "$out/${python.sitePackages}/tensorflow/lite/python/metrics/"
155 "$out/${python.sitePackages}/tensorflow/lite/python/optimize/"
156 "$out/${python.sitePackages}/tensorflow/python/"
157 "$out/${python.sitePackages}/tensorflow/python/autograph/impl/testing"
158 "$out/${python.sitePackages}/tensorflow/python/client"
159 "$out/${python.sitePackages}/tensorflow/python/data/experimental/service"
160 "$out/${python.sitePackages}/tensorflow/python/framework"
161 "$out/${python.sitePackages}/tensorflow/python/grappler"
162 "$out/${python.sitePackages}/tensorflow/python/lib/core"
163 "$out/${python.sitePackages}/tensorflow/python/lib/io"
164 "$out/${python.sitePackages}/tensorflow/python/platform"
165 "$out/${python.sitePackages}/tensorflow/python/profiler/internal"
166 "$out/${python.sitePackages}/tensorflow/python/saved_model"
167 "$out/${python.sitePackages}/tensorflow/python/util"
168 "$out/${python.sitePackages}/tensorflow/tsl/python/lib/core"
169 "${rpath}"
170 )
171
172 # The the bash array into a colon-separated list of RPATHs.
173 rrPath=$(IFS=$':'; echo "''${rrPathArr[*]}")
174 echo "about to run patchelf with the following rpath: $rrPath"
175
176 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
177 echo "about to patchelf $lib..."
178 chmod a+rx "$lib"
179 patchelf --set-rpath "$rrPath" "$lib"
180 ${lib.optionalString cudaSupport ''
181 addOpenGLRunpath "$lib"
182 ''}
183 done
184 '';
185
186 # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
187 # and the propagated input tensorboard, which causes environment collisions.
188 # Another possibility would be to have tensorboard only in the buildInputs
189 # See https://github.com/NixOS/nixpkgs/pull/44381 for more information.
190 postInstall = ''
191 rm $out/bin/tensorboard
192 '';
193
194 pythonImportsCheck = [
195 "tensorflow"
196 "tensorflow.python"
197 "tensorflow.python.framework"
198 ];
199
200 passthru = {
201 inherit cudaPackages;
202 };
203
204 meta = with lib; {
205 description = "Computation using data flow graphs for scalable machine learning";
206 homepage = "http://tensorflow.org";
207 sourceProvenance = with sourceTypes; [ binaryNativeCode ];
208 license = licenses.asl20;
209 maintainers = with maintainers; [ jyp abbradar ];
210 platforms = [ "x86_64-linux" "x86_64-darwin" ];
211 };
212}