NixOS · Oct 2, 2019 · Oct 2, 2019 · Oct 2, 2019
Showing with 42 additions and 11 deletions.

+6 −1 pkgs/development/libraries/science/math/tensorflow/bin.nix

+9 −1 pkgs/development/python-modules/tensorflow/bin.nix

+27 −9 pkgs/development/python-modules/tensorflow/default.nix
diff --git a/pkgs/development/libraries/science/math/tensorflow/bin.nix b/pkgs/development/libraries/science/math/tensorflow/bin.nix
@@ -1,6 +1,6 @@
 { stdenv
 , fetchurl
-, patchelf
+, addOpenGLRunpath
 , cudaSupport ? false, symlinkJoin, cudatoolkit, cudnn, nvidia_x11
 }:
 
@@ -35,6 +35,9 @@ let
     else ''
       patchelf --set-rpath "${rpath}:$out/lib" $out/lib/libtensorflow.so
       patchelf --set-rpath "${rpath}" $out/lib/libtensorflow_framework.so
+      ${optionalString cudaSupport ''
+        addOpenGLRunpath $out/lib/libtensorflow.so $out/lib/libtensorflow_framework.so
+      ''}
     '';
 
 in stdenv.mkDerivation rec {
@@ -43,6 +46,8 @@ in stdenv.mkDerivation rec {
 
   src = fetchurl url;
 
+  nativeBuildInputs = optional cudaSupport addOpenGLRunpath;
+
   # Patch library to use our libc, libstdc++ and others
   buildCommand = ''
     mkdir -pv $out

diff --git a/pkgs/development/python-modules/tensorflow/bin.nix b/pkgs/development/python-modules/tensorflow/bin.nix
@@ -26,6 +26,7 @@
 , symlinkJoin
 , keras-applications
 , keras-preprocessing
+, addOpenGLRunpath
 }:
 
 # We keep this binary build for two reasons:
@@ -76,6 +77,8 @@ in buildPythonPackage {
   ] ++ lib.optional (!isPy3k) mock
     ++ lib.optionals (pythonOlder "3.4") [ backports_weakref ];
 
+  nativeBuildInputs = 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
@@ -94,7 +97,12 @@ in buildPythonPackage {
   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 \( -name '*.so' -or -name '*.so.*' \) -exec patchelf --set-rpath "$rrPath" {} \;
+    find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
+      patchelf --set-rpath "$rrPath" "$lib"
+      ${lib.optionalString cudaSupport ''
+        addOpenGLRunpath "$lib"
+      ''}
+    done
   '';
 
 

diff --git a/pkgs/development/python-modules/tensorflow/default.nix b/pkgs/development/python-modules/tensorflow/default.nix
@@ -1,4 +1,5 @@
 { stdenv, pkgs, buildBazelPackage, lib, fetchFromGitHub, fetchpatch, symlinkJoin
+, addOpenGLRunpath
 # Python deps
 , buildPythonPackage, isPy3k, pythonOlder, pythonAtLeast, python
 # Python libraries
@@ -112,7 +113,7 @@ let
 
     nativeBuildInputs = [
       swig which pythonEnv
-    ];
+    ] ++ lib.optional cudaSupport addOpenGLRunpath;
 
     buildInputs = [
       jemalloc
@@ -203,6 +204,7 @@ let
     TF_NEED_CUDA = tfFeature cudaSupport;
     TF_CUDA_PATHS = lib.optionalString cudaSupport "${cudatoolkit_joined},${cudnn},${nccl}";
     GCC_HOST_COMPILER_PREFIX = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin";
+    GCC_HOST_COMPILER_PATH = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin/gcc";
     TF_CUDA_COMPUTE_CAPABILITIES = lib.concatStringsSep "," cudaCapabilities;
 
     postPatch = ''
@@ -295,6 +297,21 @@ let
         bazel-bin/tensorflow/tools/pip_package/build_pip_package --src "$PWD/dist"
         cp -Lr "$PWD/dist" "$python"
       '';
+
+      postFixup = lib.optionalString cudaSupport ''
+        find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
+          addOpenGLRunpath "$lib"
+        done
+      '';
+    };
+
+    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 = platforms.linux;
+      broken = !(xlaSupport -> cudaSupport);
     };
   };
 
@@ -337,6 +354,14 @@ in buildPythonPackage {
     tensorflow-tensorboard
   ];
 
+  nativeBuildInputs = lib.optional cudaSupport addOpenGLRunpath;
+
+  postFixup = lib.optionalString cudaSupport ''
+    find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
+      addOpenGLRunpath "$lib"
+    done
+  '';
+
   # Actual tests are slow and impure.
   # TODO try to run them anyway
   # TODO better test (files in tensorflow/tools/ci_build/builds/*test)
@@ -346,12 +371,5 @@ in buildPythonPackage {
 
   passthru.libtensorflow = bazel-build.out;
 
-  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 = platforms.linux;
-    broken = !(xlaSupport -> cudaSupport);
-  };
+  inherit (bazel-build) meta;
 }