pkgs/development/python-modules/tensorflow/bin.nix

   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
  42 assert ! (stdenv.isDarwin && cudaSupport);
  43
  44 let
  45   packages = import ./binary-hashes.nix;
  46   inherit (cudaPackages) cudatoolkit cudnn;
  47 in 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 }