diff --git a/pkgs/development/libraries/science/math/tensorflow/bin.nix b/pkgs/development/libraries/science/math/tensorflow/bin.nix index 3db7a6f1d1e2..f9f300883a16 100644 --- 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 index e46d02989531..db2308a0be8a 100644 --- 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 index 1895f2bfd0a8..e180c9d70e94 100644 --- 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 @@ -296,6 +297,12 @@ 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; { @@ -347,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)