diff options
Diffstat (limited to 'pkgs/development/python-modules/tensorflow/bin.nix')
| -rw-r--r-- | pkgs/development/python-modules/tensorflow/bin.nix | 106 |
1 files changed, 82 insertions, 24 deletions
diff --git a/pkgs/development/python-modules/tensorflow/bin.nix b/pkgs/development/python-modules/tensorflow/bin.nix index 38fec9adc7d..e81dc2dbd00 100644 --- a/pkgs/development/python-modules/tensorflow/bin.nix +++ b/pkgs/development/python-modules/tensorflow/bin.nix @@ -14,6 +14,9 @@ , absl-py , grpcio , mock +, scipy +, wheel +, opt-einsum , backports_weakref , tensorflow-estimator , tensorflow-tensorboard @@ -55,21 +58,22 @@ in buildPythonPackage { src = let pyVerNoDot = lib.strings.stringAsChars (x: if x == "." then "" else x) python.pythonVersion; - pyver = if stdenv.isDarwin then builtins.substring 0 1 pyVerNoDot else pyVerNoDot; platform = if stdenv.isDarwin then "mac" else "linux"; unit = if cudaSupport then "gpu" else "cpu"; - key = "${platform}_py_${pyver}_${unit}"; + key = "${platform}_py_${pyVerNoDot}_${unit}"; in fetchurl packages.${key}; propagatedBuildInputs = [ protobuf numpy + scipy termcolor grpcio six astor absl-py gast + opt-einsum google-pasta wrapt tensorflow-estimator @@ -79,40 +83,94 @@ in buildPythonPackage { ] ++ lib.optional (!isPy3k) mock ++ lib.optionals (pythonOlder "3.4") [ backports_weakref ]; - nativeBuildInputs = lib.optional cudaSupport addOpenGLRunpath; + nativeBuildInputs = [ wheel ] ++ lib.optional cudaSupport addOpenGLRunpath; - # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow - # and the propageted input tensorflow-tensorboard which causes environment collisions. - # another possibility would be to have tensorboard only in the buildInputs - # https://github.com/tensorflow/tensorflow/blob/v1.7.1/tensorflow/tools/pip_package/setup.py#L79 - postInstall = '' - rm $out/bin/tensorboard + preConfigure = '' + unset SOURCE_DATE_EPOCH + + # Make sure that dist and the wheel file are writable. + chmod u+rwx -R ./dist + + pushd dist + + # Unpack the wheel file. + wheel unpack --dest unpacked ./*.whl + + # Tensorflow has a hard dependency on gast==0.2.2, but we relax it to + # gast==0.3.2. + substituteInPlace ./unpacked/tensorflow*/tensorflow_core/tools/pip_package/setup.py --replace "gast == 0.2.2" "gast == 0.3.2" + substituteInPlace ./unpacked/tensorflow*/tensorflow_*.dist-info/METADATA --replace "gast (==0.2.2)" "gast (==0.3.2)" + + # Pack the wheel file back up. + wheel pack ./unpacked/tensorflow* + + popd ''; # Note that we need to run *after* the fixup phase because the # libraries are loaded at runtime. If we run in preFixup then # patchelf --shrink-rpath will remove the cuda libraries. - postFixup = let - rpath = stdenv.lib.makeLibraryPath - ([ stdenv.cc.cc.lib zlib ] ++ lib.optionals cudaSupport [ cudatoolkit.out cudatoolkit.lib cudnn nvidia_x11 ]); - in - lib.optionalString stdenv.isLinux '' - rrPath="$out/${python.sitePackages}/tensorflow/:$out/${python.sitePackages}/tensorflow/contrib/tensor_forest/:${rpath}" - internalLibPath="$out/${python.sitePackages}/tensorflow/python/_pywrap_tensorflow_internal.so" - find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do - patchelf --set-rpath "$rrPath" "$lib" - ${lib.optionalString cudaSupport '' - addOpenGLRunpath "$lib" - ''} - done - ''; + postFixup = + let + # rpaths we only need to add if CUDA is enabled. + cudapaths = lib.optionals cudaSupport [ + cudatoolkit.out + cudatoolkit.lib + cudnn + nvidia_x11 + ]; + + libpaths = [ + stdenv.cc.cc.lib + zlib + ]; + + rpath = stdenv.lib.makeLibraryPath (libpaths ++ cudapaths); + in + lib.optionalString stdenv.isLinux '' + # This is an array containing all the directories in the tensorflow2 + # package that contain .so files. + # + # TODO: Create this list programmatically, and remove paths that aren't + # actually needed. + rrPathArr=( + "$out/${python.sitePackages}/tensorflow_core/" + "$out/${python.sitePackages}/tensorflow_core/compiler/tf2tensorrt/" + "$out/${python.sitePackages}/tensorflow_core/compiler/tf2xla/ops/" + "$out/${python.sitePackages}/tensorflow_core/lite/experimental/microfrontend/python/ops/" + "$out/${python.sitePackages}/tensorflow_core/lite/python/interpreter_wrapper/" + "$out/${python.sitePackages}/tensorflow_core/lite/python/optimize/" + "$out/${python.sitePackages}/tensorflow_core/python/" + "$out/${python.sitePackages}/tensorflow_core/python/framework/" + "${rpath}" + ) + + # The the bash array into a colon-separated list of RPATHs. + rrPath=$(IFS=$':'; echo "''${rrPathArr[*]}") + echo "about to run patchelf with the following rpath: $rrPath" + + find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do + echo "about to patchelf $lib..." + chmod a+rx "$lib" + patchelf --set-rpath "$rrPath" "$lib" + ${lib.optionalString cudaSupport '' + addOpenGLRunpath "$lib" + ''} + done + ''; + pythonImportsCheck = [ + "tensorflow" + "tensorflow.keras" + "tensorflow.python" + "tensorflow.python.framework" + ]; meta = with stdenv.lib; { description = "Computation using data flow graphs for scalable machine learning"; homepage = http://tensorflow.org; license = licenses.asl20; - maintainers = with maintainers; [ jyp abbradar ]; + maintainers = with maintainers; [ jyp abbradar cdepillabout ]; platforms = [ "x86_64-linux" "x86_64-darwin" ]; # Python 2.7 build uses different string encoding. # See https://github.com/NixOS/nixpkgs/pull/37044#issuecomment-373452253 |