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