Merge pull request #31507 from abbradar/tensorflow-new

Fixes for CUDA and Tensorflow update

5 files changed, 168 insertions, 126 deletions

diff --git a/pkgs/build-support/build-bazel-package/default.nix b/pkgs/build-support/build-bazel-package/default.nix
new file mode 100644
index 000000000000..349a7e5aa333
--- /dev/null
+++ b/pkgs/build-support/build-bazel-package/default.nix
@@ -0,0 +1,78 @@
+{ stdenv, bazel }:
+
+args@{ name, bazelFlags ? [], bazelTarget, buildAttrs, fetchAttrs, ... }:
+
+let
+  fArgs = removeAttrs args [ "buildAttrs" "fetchAttrs" ];
+  fBuildAttrs = fArgs // buildAttrs;
+  fFetchAttrs = fArgs // removeAttrs fetchAttrs [ "sha256" ];
+
+in stdenv.mkDerivation (fBuildAttrs // {
+  inherit name bazelFlags bazelTarget;
+
+  deps = stdenv.mkDerivation (fFetchAttrs // {
+    name = "${name}-deps";
+    inherit bazelFlags bazelTarget;
+
+    nativeBuildInputs = fFetchAttrs.nativeBuildInputs or [] ++ [ bazel ];
+
+    preHook = fFetchAttrs.preHook or "" + ''
+      export bazelOut="$NIX_BUILD_TOP/output"
+      export HOME="$NIX_BUILD_TOP"
+    '';
+
+    buildPhase = fFetchAttrs.buildPhase or ''
+      runHook preBuild
+
+      bazel --output_base="$bazelOut" fetch $bazelFlags $bazelTarget
+
+      runHook postBuild
+    '';
+
+    installPhase = fFetchAttrs.installPhase or ''
+      runHook preInstall
+
+      # Patching markers to make them deterministic
+      for i in $bazelOut/external/\@*.marker; do
+        sed -i 's, -\?[0-9][0-9]*$, 1,' "$i"
+      done
+      # Patching symlinks to remove build directory reference
+      find $bazelOut/external -type l | while read symlink; do
+        ln -sf $(readlink "$symlink" | sed "s,$NIX_BUILD_TOP,NIX_BUILD_TOP,") "$symlink"
+      done
+
+      cp -r $bazelOut/external $out
+
+      runHook postInstall
+    '';
+
+    dontFixup = true;
+    outputHashMode = "recursive";
+    outputHashAlgo = "sha256";
+    outputHash = fetchAttrs.sha256;
+  });
+
+  nativeBuildInputs = fBuildAttrs.nativeBuildInputs or [] ++ [ (bazel.override { enableNixHacks = true; }) ];
+
+  preHook = fBuildAttrs.preHook or "" + ''
+    export bazelOut="$NIX_BUILD_TOP/output"
+    export HOME="$NIX_BUILD_TOP"
+  '';
+
+  preConfigure = ''
+    mkdir -p $bazelOut/external
+    cp -r $deps/* $bazelOut/external
+    chmod -R +w $bazelOut
+    find $bazelOut -type l | while read symlink; do
+      ln -sf $(readlink "$symlink" | sed "s,NIX_BUILD_TOP,$NIX_BUILD_TOP,") "$symlink"
+    done
+  '' + fBuildAttrs.preConfigure or "";
+
+  buildPhase = fBuildAttrs.buildPhase or ''
+    runHook preBuild
+
+    bazel --output_base="$bazelOut" build -j $NIX_BUILD_CORES $bazelFlags $bazelTarget
+
+    runHook postBuild
+  '';
+})
diff --git a/pkgs/development/python-modules/tensorflow-tensorboard/default.nix b/pkgs/development/python-modules/tensorflow-tensorboard/default.nix
index 9f48b8fb971a..f0b4e6f341d4 100644
--- a/pkgs/development/python-modules/tensorflow-tensorboard/default.nix
+++ b/pkgs/development/python-modules/tensorflow-tensorboard/default.nix
@@ -1,22 +1,19 @@
-{ stdenv
-, fetchPypi
-, buildPythonPackage
-, isPy3k
+{ stdenv, lib, fetchPypi, buildPythonPackage, isPy3k
 , bleach_1_5_0
 , numpy
 , werkzeug
 , protobuf
 , markdown
+, futures
 }:
 
-# tensorflow is built from a downloaded wheel, because the upstream
-# project's build system is an arcane beast based on
-# bazel. Untangling it and building the wheel from source is an open
-# problem.
+# tensorflow is built from a downloaded wheel, because
+# https://github.com/tensorflow/tensorboard/issues/719
+# blocks buildBazelPackage.
 
 buildPythonPackage rec {
   pname = "tensorflow-tensorboard";
-  version = "0.1.5";
+  version = "1.5.1";
   name = "${pname}-${version}";
   format = "wheel";
 
@@ -26,16 +23,16 @@ buildPythonPackage rec {
     format = "wheel";
   } // (if isPy3k then {
     python = "py3";
-    sha256 = "0sfia05y1mzgy371faj96vgzhag1rgpa3gnbz9w1fay13jryw26x";
+    sha256 = "1cydgvrr0s05xqz1v9z2wdiv60gzbs8wv9wvbflw5700a2llb63l";
   } else {
     python = "py2";
-    sha256 = "0qx4f55zp54x079kxir4zz5b1ckiglsdcb9afz5wcdj6af4a6czg";
+    sha256 = "0dhljddlirq6nr84zg4yrk5k69gj3x2abb6wg3crgrparb6qbya7";
   }));
 
-  propagatedBuildInputs = [ bleach_1_5_0 numpy werkzeug protobuf markdown ];
+  propagatedBuildInputs = [ bleach_1_5_0 numpy werkzeug protobuf markdown ] ++ lib.optional (!isPy3k) futures;
 
   meta = with stdenv.lib; {
-    description = "TensorFlow helps the tensors flow";
+    description = "TensorFlow's Visualization Toolkit";
     homepage = http://tensorflow.org;
     license = licenses.asl20;
     maintainers = with maintainers; [ abbradar ];
diff --git a/pkgs/development/python-modules/tensorflow/default.nix b/pkgs/development/python-modules/tensorflow/default.nix
index 63630c1e1463..e51ceb313ef2 100644
--- a/pkgs/development/python-modules/tensorflow/default.nix
+++ b/pkgs/development/python-modules/tensorflow/default.nix
@@ -1,8 +1,8 @@
-{ stdenv, lib, fetchFromGitHub, fetchpatch, symlinkJoin, buildPythonPackage, isPy3k, pythonOlder
-, bazel, which, swig, binutils, glibcLocales
+{ stdenv, buildBazelPackage, lib, fetchFromGitHub, fetchpatch, symlinkJoin
+, buildPythonPackage, isPy3k, pythonOlder, pythonAtLeast
+, which, swig, binutils, glibcLocales
 , python, jemalloc, openmpi
-, numpy, six, protobuf, tensorflow-tensorboard, backports_weakref
-, wheel, mock, scipy
+, numpy, six, protobuf, tensorflow-tensorboard, backports_weakref, mock, enum34, absl-py
 , xlaSupport ? true
 , cudaSupport ? false, nvidia_x11 ? null, cudatoolkit ? null, cudnn ? null
 # Default from ./configure script
@@ -12,7 +12,8 @@
 , fmaSupport ? false
 }:
 
-assert cudaSupport -> cudatoolkit != null
+assert cudaSupport -> nvidia_x11 != null
+                   && cudatoolkit != null
                    && cudnn != null;
 
 # unsupported combination
@@ -27,40 +28,49 @@ let
     paths = [ cudatoolkit.out cudatoolkit.lib ];
   };
 
-  cudaLibPath = lib.makeLibraryPath [ cudatoolkit.out cudatoolkit.lib nvidia_x11 cudnn ];
-
   tfFeature = x: if x then "1" else "0";
 
-  common = rec {
-    version = "1.3.1";
+  version = "1.5.0";
+
+  pkg = buildBazelPackage rec {
+    name = "tensorflow-build-${version}";
 
     src = fetchFromGitHub {
       owner = "tensorflow";
       repo = "tensorflow";
       rev = "v${version}";
-      sha256 = "0gvi32dvv4ynr05p0gg5i0a6c55pig48k5qm7zslcqnp4sifwx0i";
+      sha256 = "1c4djsaip901nasm7a6dsimr02bsv70a7b1g0kysb4n39qpdh22q";
     };
 
-    nativeBuildInputs = [ swig which wheel scipy ];
+    patches = [
+      # Fix build with Bazel >= 0.10
+      (fetchpatch {
+        url = "https://github.com/tensorflow/tensorflow/commit/6fcfab770c2672e2250e0f5686b9545d99eb7b2b.patch";
+        sha256 = "0p61za1mx3a7gj1s5lsps16fcw18iwnvq2b46v1kyqfgq77a12vb";
+      })
+      (fetchpatch {
+        url = "https://github.com/tensorflow/tensorflow/commit/3f57956725b553d196974c9ad31badeb3eabf8bb.patch";
+        sha256 = "11dja5gqy0qw27sc9b6yw9r0lfk8dznb32vrqqfcnypk2qmv26va";
+      })
+    ];
 
-    buildInputs = [ python jemalloc openmpi glibcLocales ]
-      ++ lib.optionals cudaSupport [ cudatoolkit cudnn ];
+    nativeBuildInputs = [ swig which ];
 
-    propagatedBuildInputs = [ numpy six protobuf ]
-                            ++ lib.optional (!isPy3k) mock
-                            ++ lib.optional (pythonOlder "3.4") backports_weakref
-                            ++ lib.optional withTensorboard tensorflow-tensorboard;
+    buildInputs = [ python jemalloc openmpi glibcLocales numpy ]
+      ++ lib.optionals cudaSupport [ cudatoolkit cudnn nvidia_x11 ];
 
     preConfigure = ''
       patchShebangs configure
-      export HOME="$NIX_BUILD_TOP"
 
       export PYTHON_BIN_PATH="${python.interpreter}"
+      export PYTHON_LIB_PATH="$NIX_BUILD_TOP/site-packages"
       export TF_NEED_GCP=1
       export TF_NEED_HDFS=1
-      export TF_NEED_CUDA=${tfFeature cudaSupport}
-      export TF_NEED_MPI=1
       export TF_ENABLE_XLA=${tfFeature xlaSupport}
+      export CC_OPT_FLAGS=" "
+      # https://github.com/tensorflow/tensorflow/issues/14454
+      export TF_NEED_MPI=${tfFeature cudaSupport}
+      export TF_NEED_CUDA=${tfFeature cudaSupport}
       ${lib.optionalString cudaSupport ''
         export CUDA_TOOLKIT_PATH=${cudatoolkit_joined}
         export TF_CUDA_VERSION=${cudatoolkit.majorVersion}
@@ -70,13 +80,11 @@ let
         export TF_CUDA_COMPUTE_CAPABILITIES=${lib.concatStringsSep "," cudaCapabilities}
       ''}
 
-      # There is _no_ non-interactive mode of configure.
-      sed -i \
-        -e 's,read -p,echo,g' \
-        -e 's,lib64,lib,g' \
-        configure
+      mkdir -p "$PYTHON_LIB_PATH"
     '';
 
+    NIX_LDFLAGS = lib.optionals cudaSupport [ "-lcublas" "-lcudnn" "-lcuda" "-lcudart" ];
+
     hardeningDisable = [ "all" ];
 
     bazelFlags = [ "--config=opt" ]
@@ -87,101 +95,59 @@ let
 
     bazelTarget = "//tensorflow/tools/pip_package:build_pip_package";
 
-    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 ];
-      platforms = with platforms; if cudaSupport then linux else linux ++ darwin;
-    };
-  };
+    fetchAttrs = {
+      preInstall = ''
+        rm -rf $bazelOut/external/{bazel_tools,\@bazel_tools.marker,local_*,\@local_*}
+      '';
 
-in buildPythonPackage (common // {
-  pname = "tensorflow";
-  version = common.version;
-  name = "tensorflow-${common.version}";
-
-  deps = stdenv.mkDerivation (common // {
-    name = "tensorflow-external-${common.version}";
-
-    nativeBuildInputs = common.nativeBuildInputs ++ [ bazel ];
-
-    preConfigure = common.preConfigure + ''
-      export PYTHON_LIB_PATH="$(pwd)/site-packages"
-    '';
-
-    buildPhase = ''
-      mkdir site-packages
-      bazel --output_base="$(pwd)/output" fetch $bazelFlags $bazelTarget
-    '';
+      sha256 = "1nc98aqrp14q7llypcwaa0kdn9xi7r0p1mnd3vmmn1m299py33ca";
+    };
 
-    installPhase = ''
-      rm -rf output/external/{bazel_tools,\@bazel_tools.marker,local_*,\@local_*}
-      # Patching markers to make them deterministic
-      for i in output/external/\@*.marker; do
-        sed -i 's, -\?[0-9][0-9]*$, 1,' "$i"
-      done
-      # Patching symlinks to remove build directory reference
-      find output/external -type l | while read symlink; do
-        ln -sf $(readlink "$symlink" | sed "s,$NIX_BUILD_TOP,NIX_BUILD_TOP,") "$symlink"
-      done
-
-      cp -r output/external $out
-    '';
+    buildAttrs = {
+      preBuild = ''
+        patchShebangs .
+        find -type f -name CROSSTOOL\* -exec sed -i \
+          -e 's,/usr/bin/ar,${binutils.bintools}/bin/ar,g' \
+          {} \;
+      '';
+
+      installPhase = ''
+        sed -i 's,.*bdist_wheel.*,cp -rL . "$out"; exit 0,' bazel-bin/tensorflow/tools/pip_package/build_pip_package 
+        bazel-bin/tensorflow/tools/pip_package/build_pip_package $PWD/dist
+      '';
+    };
 
     dontFixup = true;
+  };
 
-    outputHashMode = "recursive";
-    outputHashAlgo = "sha256";
-    outputHash = "0xs2n061gnpizfcnhs5jjpfk2av634j1l2l17zhy10bbmrwn3vrp";
-  });
-
-  nativeBuildInputs = common.nativeBuildInputs ++ [ (bazel.override { enableNixHacks = true; }) ];
-
-  configurePhase = ''
-    runHook preConfigure
-    export PYTHON_LIB_PATH="$out/${python.sitePackages}"
-    ./configure
-    runHook postConfigure
-  '';
-
-  buildPhase = ''
-    mkdir -p output/external
-    cp -r $deps/* output/external
-    chmod -R +w output
-    find output -type l | while read symlink; do
-      ln -sf $(readlink "$symlink" | sed "s,NIX_BUILD_TOP,$NIX_BUILD_TOP,") "$symlink"
-    done
-
-    patchShebangs .
-    find -type f -name CROSSTOOL\* -exec sed -i \
-      -e 's,/usr/bin/ar,${binutils}/bin/ar,g' \
-      {} \;
-
-    mkdir -p $out/${python.sitePackages}
-    bazel --output_base="$(pwd)/output" build $bazelFlags $bazelTarget
+in buildPythonPackage rec {
+  pname = "tensorflow";
+  inherit version;
+  name = "${pname}-${version}";
 
-    bazel-bin/tensorflow/tools/pip_package/build_pip_package $PWD/dist
-  '';
+  src = pkg;
 
-  # tensorflow depends on tensorflow_tensorboard, which cannot be
-  # built at the moment (some of its dependencies do not build
-  # [htlm5lib9999999 (seven nines) -> tensorboard], and it depends on an old version of
-  # bleach) Hence we disable dependency checking for now.
   installFlags = lib.optional (!withTensorboard) "--no-dependencies";
 
-  # Tests are slow and impure.
-  doCheck = false;
-
-  # For some reason, CUDA is not retained in RPATH.
-  postFixup = lib.optionalString cudaSupport ''
-    libPath="$out/${python.sitePackages}/tensorflow/python/_pywrap_tensorflow_internal.so"
-    patchelf --set-rpath "$(patchelf --print-rpath "$libPath"):${cudaLibPath}" "$libPath"
+  postPatch = lib.optionalString (pythonAtLeast "3.4") ''
+    sed -i '/enum34/d' setup.py
   '';
 
-  doInstallCheck = true;
-  installCheckPhase = ''
-    cd $NIX_BUILD_TOP
+  propagatedBuildInputs = [ numpy six protobuf absl-py ]
+                 ++ lib.optional (!isPy3k) mock
+                 ++ lib.optionals (pythonOlder "3.4") [ backports_weakref enum34 ]
+                 ++ lib.optional withTensorboard tensorflow-tensorboard;
+
+  # Actual tests are slow and impure.
+  checkPhase = ''
     ${python.interpreter} -c "import tensorflow"
   '';
-})
+
+  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 ];
+    platforms = with platforms; if cudaSupport then linux else linux ++ darwin;
+  };
+}
diff --git a/pkgs/top-level/all-packages.nix b/pkgs/top-level/all-packages.nix
index 88c87a134c5d..1e710216e5dd 100644
--- a/pkgs/top-level/all-packages.nix
+++ b/pkgs/top-level/all-packages.nix
@@ -7339,6 +7339,8 @@ with pkgs;
   bazel_0_4 = callPackage ../development/tools/build-managers/bazel/0.4.nix { };
   bazel = callPackage ../development/tools/build-managers/bazel { };
 
+  buildBazelPackage = callPackage ../build-support/build-bazel-package { };
+
   bear = callPackage ../development/tools/build-managers/bear { };
 
   bin_replace_string = callPackage ../development/tools/misc/bin_replace_string { };
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index c3e2e99c9432..2f39e8b0d36e 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -21251,11 +21251,10 @@ EOF
   tensorflow-tensorboard = callPackage ../development/python-modules/tensorflow-tensorboard { };
 
   tensorflow = callPackage ../development/python-modules/tensorflow rec {
-    bazel = pkgs.bazel_0_4;
     cudaSupport = pkgs.config.cudaSupport or false;
     inherit (pkgs.linuxPackages) nvidia_x11;
-    cudatoolkit = pkgs.cudatoolkit8;
-    cudnn = pkgs.cudnn6_cudatoolkit8;
+    cudatoolkit = pkgs.cudatoolkit9;
+    cudnn = pkgs.cudnn_cudatoolkit9;
   };
 
   tensorflowWithoutCuda = self.tensorflow.override {