tensorflow: cleanup and merge with libtensorflow

This merges work done by yorickvP and timokau in #63208 and #63616 respectively.
Now the derivation builds both libtensorflow and the Python package and puts them into
different outputs.

Quite a bit of improvements were done on the top, including:

* Use official tag revision as source, not a branch;
* Use all system libraries possible (before only one was actually used);
* Move various environment variables to the derivation itself from hooks;
* Use source Python build instead of wheel build to ensure fixup hooks do their important jobs on libraries;
* And more that I forgot!
This commit is contained in:
Nikolay Amiantov 2019-07-13 22:42:50 +03:00
parent 0a1bf4734f
commit 019c13616b
4 changed files with 169 additions and 142 deletions

View file

@ -1,41 +1,15 @@
{ stdenv, buildBazelPackage, lib, fetchFromGitHub, fetchpatch, symlinkJoin { stdenv, pkgs, buildBazelPackage, lib, fetchFromGitHub, fetchpatch, symlinkJoin
, git # Python deps
, buildPythonPackage, isPy3k, pythonOlder, pythonAtLeast , buildPythonPackage, isPy3k, pythonOlder, pythonAtLeast, python
, which, binutils, glibcLocales # Python libraries
, python, jemalloc, openmpi
, numpy, tensorflow-tensorboard, backports_weakref, mock, enum34, absl-py , numpy, tensorflow-tensorboard, backports_weakref, mock, enum34, absl-py
, future , future, setuptools, wheel, keras-preprocessing, keras-applications, google-pasta
, keras-preprocessing , termcolor, grpcio, six, wrapt, protobuf, tensorflow-estimator
, keras-applications # Common deps
, astor , git, swig, which, binutils, glibcLocales, cython
, gast # Common libraries
, google-pasta , jemalloc, openmpi, astor, gast, grpc, sqlite, openssl, jsoncpp, re2
, termcolor , curl, snappy, flatbuffers, icu, double-conversion, libpng, libjpeg, giflib
, cython
, flatbuffers
, giflib
, libjpeg
, grpc
, grpcio
, hwloc
, icu
, jsoncpp
, lmdb
, nasm
, sqlite
, pcre
, libpng
, six
, snappy
, swig
, wrapt
, zlib
, protobuf
, protobuf_cc
, curl
, tensorflow-estimator
, setuptools
, wheel
, cudaSupport ? false, nvidia_x11 ? null, cudatoolkit ? null, cudnn ? null, nccl ? null , cudaSupport ? false, nvidia_x11 ? null, cudatoolkit ? null, cudnn ? null, nccl ? null
# XLA without CUDA is broken # XLA without CUDA is broken
, xlaSupport ? cudaSupport , xlaSupport ? cudaSupport
@ -57,48 +31,51 @@ let
withTensorboard = pythonOlder "3.6"; withTensorboard = pythonOlder "3.6";
cudatoolkit_joined = symlinkJoin { cudatoolkit_joined = symlinkJoin {
name = "${cudatoolkit.name}-unsplit"; name = "${cudatoolkit.name}-merged";
paths = [ cudatoolkit.out cudatoolkit.lib ]; paths = [ cudatoolkit.out cudatoolkit.lib ];
}; };
bazel_cuda_toolkit = symlinkJoin { cudatoolkit_cc_joined = symlinkJoin {
name = "bazel-gcc-toolkit"; name = "${cudatoolkit.cc.name}-merged";
paths = [ paths = [
cudatoolkit.cc cudatoolkit.cc
cudatoolkit.out
cudatoolkit.lib
binutils.bintools # for ar, dwp, nm, objcopy, objdump, strip binutils.bintools # for ar, dwp, nm, objcopy, objdump, strip
]; ];
}; };
# Needed for _some_ system libraries, grep INCLUDEDIR.
includes_joined = symlinkJoin {
name = "tensorflow-deps-merged";
paths = [
pkgs.protobuf
jsoncpp
];
};
tfFeature = x: if x then "1" else "0"; tfFeature = x: if x then "1" else "0";
version = "1.14"; version = "1.14.0";
variant = if cudaSupport then "_gpu" else ""; variant = if cudaSupport then "-gpu" else "";
pname = "tensorflow${variant}"; pname = "tensorflow${variant}";
bazel-wheel = buildBazelPackage rec { # TODO: remove after there's support for setupPyDistFlags
# indicate which configuration of the wheel is being built setuppy = ../../../development/interpreters/python/run_setup.py;
name = let
python_tag = with lib.versions; "cp${major python.version}${minor python.version}"; # cp37 bazel-build = buildBazelPackage rec {
arch_platform = lib.splitString "-" stdenv.hostPlatform.system; name = "${pname}-${version}";
platform = lib.last arch_platform; # "linux"
arch = lib.head arch_platform; # "x86_64"
in
# https://www.python.org/dev/peps/pep-0427/#file-name-convention
# tensorflow_gpu-1.14.0-cp27-none-linux_x86_64.whl
"${pname}-${version}-${python_tag}-none-${platform}_${arch}.whl";
src = fetchFromGitHub { src = fetchFromGitHub {
owner = "tensorflow"; owner = "tensorflow";
repo = "tensorflow"; repo = "tensorflow";
rev = "r${version}"; rev = "v${version}";
sha256 = "071rndqh3xy4y5idlykqv656lggp25vy4qx1divz4id3im1zlr13"; sha256 = "06jvwlsm14b8rqwd8q8796r0vmn0wk64s4ps2zg0sapkmp9vvcmi";
}; };
patches = [ patches = [
# Work around https://github.com/tensorflow/tensorflow/issues/24752 # Work around https://github.com/tensorflow/tensorflow/issues/24752
./no-saved-proto.patch ./no-saved-proto.patch
# Fixes for NixOS jsoncpp
./system-jsoncpp.patch
# https://github.com/tensorflow/tensorflow/pull/29673 # https://github.com/tensorflow/tensorflow/pull/29673
(fetchpatch { (fetchpatch {
@ -118,7 +95,9 @@ let
# On update, it can be useful to steal the changes from gentoo # On update, it can be useful to steal the changes from gentoo
# https://gitweb.gentoo.org/repo/gentoo.git/tree/sci-libs/tensorflow # https://gitweb.gentoo.org/repo/gentoo.git/tree/sci-libs/tensorflow
nativeBuildInputs = [ swig which ]; nativeBuildInputs = [
swig which cython
];
buildInputs = [ buildInputs = [
python python
@ -130,9 +109,30 @@ let
# python deps needed during wheel build time # python deps needed during wheel build time
numpy numpy
keras-preprocessing keras-preprocessing
protobuf
wrapt
gast
astor
absl-py
termcolor
keras-applications
# libs taken from system through the TF_SYS_LIBS mechanism # libs taken from system through the TF_SYS_LIBS mechanism
absl-py grpc
sqlite
openssl
jsoncpp
pkgs.protobuf
curl
snappy
flatbuffers
icu
double-conversion
libpng
libjpeg
giflib
re2
pkgs.lmdb
# for building the wheel # for building the wheel
setuptools setuptools
@ -146,22 +146,26 @@ let
nvidia_x11 nvidia_x11
]; ];
# arbitrarily set to the current latest bazel version, overly careful
TF_IGNORE_MAX_BAZEL_VERSION = true;
# Take as many libraries from the system as possible. Keep in sync with # Take as many libraries from the system as possible. Keep in sync with
# list of valid syslibs in # list of valid syslibs in
# https://github.com/perfinion/tensorflow/blob/master/third_party/systemlibs/syslibs_configure.bzl # https://github.com/tensorflow/tensorflow/blob/master/third_party/systemlibs/syslibs_configure.bzl
SYSLIBS= [ TF_SYSTEM_LIBS = lib.concatStringsSep "," [
"absl_py" "absl_py"
"astor_archive" "astor_archive"
"boringssl" "boringssl"
"com_github_googleapis_googleapis" # Not packaged in nixpkgs
"com_github_googlecloudplatform_google_cloud_cpp" # "com_github_googleapis_googleapis"
# "com_github_googlecloudplatform_google_cloud_cpp"
"com_google_protobuf" "com_google_protobuf"
"com_google_protobuf_cc" "com_google_protobuf_cc"
"com_googlesource_code_re2" "com_googlesource_code_re2"
"curl" "curl"
"cython" "cython"
"double_conversion" "double_conversion"
"enum34_archive"
"flatbuffers" "flatbuffers"
"gast_archive" "gast_archive"
"gif_archive" "gif_archive"
@ -174,7 +178,7 @@ let
"lmdb" "lmdb"
"nasm" "nasm"
# "nsync" # not packaged in nixpkgs # "nsync" # not packaged in nixpkgs
"sqlite" "org_sqlite"
"pasta" "pasta"
"pcre" "pcre"
"png_archive" "png_archive"
@ -187,6 +191,34 @@ let
"zlib_archive" "zlib_archive"
]; ];
INCLUDEDIR = "${includes_joined}/include";
PYTHON_BIN_PATH = python.interpreter;
TF_NEED_GCP = true;
TF_NEED_HDFS = true;
TF_ENABLE_XLA = tfFeature xlaSupport;
CC_OPT_FLAGS = " ";
# https://github.com/tensorflow/tensorflow/issues/14454
TF_NEED_MPI = tfFeature cudaSupport;
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";
TF_CUDA_COMPUTE_CAPABILITIES = lib.concatStringsSep "," cudaCapabilities;
postPatch = ''
# https://github.com/tensorflow/tensorflow/issues/20919
sed -i '/androidndk/d' tensorflow/lite/kernels/internal/BUILD
# Tensorboard pulls in a bunch of dependencies, some of which may
# include security vulnerabilities. So we make it optional.
# https://github.com/tensorflow/tensorflow/issues/20280#issuecomment-400230560
sed -i '/tensorboard >=/d' tensorflow/tools/pip_package/setup.py
'';
preConfigure = '' preConfigure = ''
patchShebangs configure patchShebangs configure
@ -196,52 +228,21 @@ let
chmod +x dummy-ldconfig/ldconfig chmod +x dummy-ldconfig/ldconfig
export PATH="$PWD/dummy-ldconfig:$PATH" export PATH="$PWD/dummy-ldconfig:$PATH"
# arbitrarily set to the current latest bazel version, overly careful
export TF_IGNORE_MAX_BAZEL_VERSION=1
# don't rebuild the world
export TF_SYSTEM_LIBS=${lib.concatStringsSep " " SYSLIBS}
export PYTHON_BIN_PATH="${python.interpreter}"
export PYTHON_LIB_PATH="$NIX_BUILD_TOP/site-packages" export PYTHON_LIB_PATH="$NIX_BUILD_TOP/site-packages"
export TF_NEED_GCP=1
export TF_NEED_HDFS=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 TF_CUDA_PATHS="${cudatoolkit_joined},${cudnn},${nccl}"
export TF_CUDA_VERSION=${cudatoolkit.majorVersion}
export TF_CUDNN_VERSION=${cudnn.majorVersion}
export GCC_HOST_COMPILER_PATH=${bazel_cuda_toolkit}/bin/gcc
export GCC_HOST_COMPILER_PREFIX="${bazel_cuda_toolkit}/bin"
export TF_CUDA_COMPUTE_CAPABILITIES=${lib.concatStringsSep "," cudaCapabilities}
''}
# https://github.com/tensorflow/tensorflow/issues/20919
sed -i '/androidndk/d' tensorflow/lite/kernels/internal/BUILD
mkdir -p "$PYTHON_LIB_PATH" mkdir -p "$PYTHON_LIB_PATH"
''; '';
# ar, cpp, gcov, strip, gcc, compat-ld, ld, objdump, dwp, nm, objcopy
configurePhase = '' configurePhase = ''
runHook preConfigure runHook preConfigure
# no flags (options provided by previously set environment variables)
export AR="${binutils.bintools}/bin/ar"
export BAZEL_TOOLS_OVERRIDE="ar=${binutils.bintools}/bin/ar"
# crosstool used with cuda
./configure ./configure
runHook postConfigure runHook postConfigure
''; '';
# FIXME # FIXME: Tensorflow uses dlopen() for CUDA libraries.
NIX_LDFLAGS = lib.optionals cudaSupport [ "-lcublas" "-lcudnn" "-lcuda" "-lcudart" ]; # No idea why gpr isn't linked properly; perhaps Tensorflow expects a static library?
NIX_LDFLAGS = [ "-lgpr" ] ++ lib.optionals cudaSupport [ "-lcudart" "-lcublas" "-lcufft" "-lcurand" "-lcusolver" "-lcusparse" "-lcudnn" ];
hardeningDisable = [ "all" ]; hardeningDisable = [ "format" ];
bazelFlags = [ bazelFlags = [
# temporary fixes to make the build work with bazel 0.27 # temporary fixes to make the build work with bazel 0.27
@ -250,7 +251,7 @@ let
++ lib.optional avx2Support "--copt=-mavx2" ++ lib.optional avx2Support "--copt=-mavx2"
++ lib.optional fmaSupport "--copt=-mfma"; ++ lib.optional fmaSupport "--copt=-mfma";
bazelTarget = "//tensorflow/tools/pip_package:build_pip_package"; bazelTarget = "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow";
fetchAttrs = { fetchAttrs = {
preInstall = '' preInstall = ''
@ -259,46 +260,36 @@ let
# cudaSupport causes fetch of ncclArchive, resulting in different hashes # cudaSupport causes fetch of ncclArchive, resulting in different hashes
sha256 = if cudaSupport then sha256 = if cudaSupport then
"0q1pmw7fzn6l554ap576r48m0zgwb7n1ljhyy1p36708z94scdh4" "0gmwwcld8hm1nfxmx2w2gqxjxq5b5ln7vbxn0x63x9yy7avs3dkp"
else else
"14n26h2r4w7wd5sddy4w0s51s2qcwf276n3hvv2505iysa8wqlc3"; "1d67h0n461cyfimzlsj6bv9p95iv9jhglr0izdc6y0n6w09z8pw0";
}; };
buildAttrs = { buildAttrs = {
outputs = [ "out" "python" ];
preBuild = '' preBuild = ''
patchShebangs . patchShebangs .
# beautiful bash to iterate over files containing a string
# https://github.com/bazelbuild/bazel/issues/5915#issuecomment-505100422
# .. to make sure the output directory is covered
grep -lrZ '/usr/bin/ar\b' .. | while IFS="" read -r -d "" file; do
# patch /usr/bin/ar to the proper location
echo "File is $file"
sed -i \
-e 's,/usr/bin/ar\b,${binutils.bintools}/bin/ar,g' \
"$file"
done
# Tensorboard pulls in a bunch of dependencies, some of which may
# include security vulnerabilities. So we make it optional.
# https://github.com/tensorflow/tensorflow/issues/20280#issuecomment-400230560
sed -i '/tensorboard >=/d' tensorflow/tools/pip_package/setup.py
''; '';
# Could alternatively use --src instead of --dst to output the sources
# instead of a whee. Generating a wheel makes it easier to unify handling
# of source and binary build though.
installPhase = '' installPhase = ''
# work around timestamp issues mkdir -p "$out"
# https://github.com/NixOS/nixpkgs/issues/270#issuecomment-467583872 tar -xf bazel-bin/tensorflow/tools/lib_package/libtensorflow.tar.gz -C "$out"
export SOURCE_DATE_EPOCH=315532800 # Write pkgconfig file.
mkdir "$out/lib/pkgconfig"
cat > "$out/lib/pkgconfig/tensorflow.pc" << EOF
Name: TensorFlow
Version: ${version}
Description: Library for computation using data flow graphs for scalable machine learning
Requires:
Libs: -L$out/lib -ltensorflow
Cflags: -I$out/include/tensorflow
EOF
# bulid the wheel, then move it to $out (building directly to $out # build the source code, then copy it to $python (build_pip_package
# would actually put it into a *directory* called $out, but we want the # actually builds a symlink farm so we must dereference them).
# file itself in $out) bazel-bin/tensorflow/tools/pip_package/build_pip_package --src "$PWD/dist"
bazel-bin/tensorflow/tools/pip_package/build_pip_package --dst $PWD/dist cp -Lr "$PWD/dist" "$python"
mv dist/*.whl "$out"
''; '';
}; };
}; };
@ -306,9 +297,7 @@ let
in buildPythonPackage rec { in buildPythonPackage rec {
inherit version pname; inherit version pname;
src = bazel-wheel; src = bazel-build.python;
format = "wheel";
# Upstream has a pip hack that results in bin/tensorboard being in both tensorflow # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
# and the propagated input tensorflow-tensorboard, which causes environment collisions. # and the propagated input tensorflow-tensorboard, which causes environment collisions.
@ -318,6 +307,14 @@ in buildPythonPackage rec {
rm $out/bin/tensorboard rm $out/bin/tensorboard
''; '';
# TODO: remove after there's support for setupPyDistFlags
buildPhase = ''
runHook preBuild
cp ${setuppy} nix_run_setup
${python.interpreter} nix_run_setup --project_name ${pname} bdist_wheel
runHook postBuild
'';
# tensorflow/tools/pip_package/setup.py # tensorflow/tools/pip_package/setup.py
propagatedBuildInputs = [ propagatedBuildInputs = [
absl-py absl-py
@ -349,6 +346,8 @@ in buildPythonPackage rec {
${python.interpreter} -c "import tensorflow" ${python.interpreter} -c "import tensorflow"
''; '';
passthru.libtensorflow = bazel-build.out;
meta = with stdenv.lib; { meta = with stdenv.lib; {
description = "Computation using data flow graphs for scalable machine learning"; description = "Computation using data flow graphs for scalable machine learning";
homepage = http://tensorflow.org; homepage = http://tensorflow.org;

View file

@ -0,0 +1,21 @@
diff --git a/third_party/systemlibs/jsoncpp.BUILD b/third_party/systemlibs/jsoncpp.BUILD
index 526fd0c418..646f3fdcea 100644
--- a/third_party/systemlibs/jsoncpp.BUILD
+++ b/third_party/systemlibs/jsoncpp.BUILD
@@ -7,6 +7,7 @@ filegroup(
HEADERS = [
"include/json/autolink.h",
+ "include/json/allocator.h",
"include/json/config.h",
"include/json/features.h",
"include/json/forwards.h",
@@ -23,7 +24,7 @@ genrule(
cmd = """
for i in $(OUTS); do
i=$${i##*/}
- ln -sf $(INCLUDEDIR)/jsoncpp/json/$$i $(@D)/include/json/$$i
+ ln -sf $(INCLUDEDIR)/json/$$i $(@D)/include/json/$$i
done
""",
)

View file

@ -1093,12 +1093,15 @@ in
tensor = libsForQt5.callPackage ../applications/networking/instant-messengers/tensor { }; tensor = libsForQt5.callPackage ../applications/networking/instant-messengers/tensor { };
libtensorflow = callPackage ../development/libraries/science/math/tensorflow/bin.nix { libtensorflow =
cudaSupport = pkgs.config.cudaSupport or false; if python.pkgs.tensorflow ? libtensorflow
inherit (linuxPackages) nvidia_x11; then python.pkgs.tensorflow.libtensorflow
cudatoolkit = cudatoolkit_10_0; else callPackage ../development/libraries/science/math/tensorflow/bin.nix {
cudnn = cudnn_cudatoolkit_10_0; cudaSupport = pkgs.config.cudaSupport or false;
}; inherit (linuxPackages) nvidia_x11;
cudatoolkit = cudatoolkit_10_0;
cudnn = cudnn_cudatoolkit_10_0;
};
behdad-fonts = callPackage ../data/fonts/behdad-fonts { }; behdad-fonts = callPackage ../data/fonts/behdad-fonts { };

View file

@ -5702,14 +5702,18 @@ in {
tensorflow = tensorflow =
if stdenv.isDarwin if stdenv.isDarwin
then callPackage ../development/python-modules/tensorflow/bin.nix { } then callPackage ../development/python-modules/tensorflow/bin.nix {
cudaSupport = pkgs.config.cudaSupport or false;
inherit (pkgs.linuxPackages) nvidia_x11;
cudatoolkit = pkgs.cudatoolkit_10;
cudnn = pkgs.cudnn_cudatoolkit_10;
}
else callPackage ../development/python-modules/tensorflow rec { else callPackage ../development/python-modules/tensorflow rec {
cudaSupport = pkgs.config.cudaSupport or false; cudaSupport = pkgs.config.cudaSupport or false;
inherit (pkgs.linuxPackages) nvidia_x11; inherit (pkgs.linuxPackages) nvidia_x11;
cudatoolkit = pkgs.cudatoolkit_10; cudatoolkit = pkgs.cudatoolkit_10;
cudnn = pkgs.cudnn_cudatoolkit_10; cudnn = pkgs.cudnn_cudatoolkit_10;
nccl = pkgs.nccl_cudatoolkit_10; nccl = pkgs.nccl_cudatoolkit_10;
protobuf_cc = pkgs.protobuf; # not the python version
}; };
tensorflowWithoutCuda = self.tensorflow.override { tensorflowWithoutCuda = self.tensorflow.override {