forked from mirrors/nixpkgs
361 lines
11 KiB
Nix
361 lines
11 KiB
Nix
{ stdenv, buildBazelPackage, lib, fetchFromGitHub, fetchpatch, symlinkJoin
|
|
, git
|
|
, buildPythonPackage, isPy3k, pythonOlder, pythonAtLeast
|
|
, which, binutils, glibcLocales
|
|
, python, jemalloc, openmpi
|
|
, numpy, tensorflow-tensorboard, backports_weakref, mock, enum34, absl-py
|
|
, future
|
|
, keras-preprocessing
|
|
, keras-applications
|
|
, astor
|
|
, gast
|
|
, google-pasta
|
|
, termcolor
|
|
, 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
|
|
# XLA without CUDA is broken
|
|
, xlaSupport ? cudaSupport
|
|
# Default from ./configure script
|
|
, cudaCapabilities ? [ "3.5" "5.2" ]
|
|
, sse42Support ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") ["westmere" "sandybridge" "ivybridge" "haswell" "broadwell" "skylake" "skylake-avx512"]
|
|
, avx2Support ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") [ "haswell" "broadwell" "skylake" "skylake-avx512"]
|
|
, fmaSupport ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") [ "haswell" "broadwell" "skylake" "skylake-avx512"]
|
|
}:
|
|
|
|
assert cudaSupport -> nvidia_x11 != null
|
|
&& cudatoolkit != null
|
|
&& cudnn != null;
|
|
|
|
# unsupported combination
|
|
assert ! (stdenv.isDarwin && cudaSupport);
|
|
|
|
let
|
|
withTensorboard = pythonOlder "3.6";
|
|
|
|
cudatoolkit_joined = symlinkJoin {
|
|
name = "${cudatoolkit.name}-unsplit";
|
|
paths = [ cudatoolkit.out cudatoolkit.lib ];
|
|
};
|
|
|
|
bazel_cuda_toolkit = symlinkJoin {
|
|
name = "bazel-gcc-toolkit";
|
|
paths = [
|
|
cudatoolkit.cc
|
|
cudatoolkit.out
|
|
cudatoolkit.lib
|
|
binutils.bintools # for ar, dwp, nm, objcopy, objdump, strip
|
|
];
|
|
};
|
|
|
|
tfFeature = x: if x then "1" else "0";
|
|
|
|
version = "1.14";
|
|
variant = if cudaSupport then "_gpu" else "";
|
|
pname = "tensorflow${variant}";
|
|
|
|
bazel-wheel = buildBazelPackage rec {
|
|
# indicate which configuration of the wheel is being built
|
|
name = let
|
|
python_tag = with lib.versions; "cp${major python.version}${minor python.version}"; # cp37
|
|
arch_platform = lib.splitString "-" stdenv.hostPlatform.system;
|
|
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 {
|
|
owner = "tensorflow";
|
|
repo = "tensorflow";
|
|
rev = "r${version}";
|
|
sha256 = "071rndqh3xy4y5idlykqv656lggp25vy4qx1divz4id3im1zlr13";
|
|
};
|
|
|
|
patches = [
|
|
# Work around https://github.com/tensorflow/tensorflow/issues/24752
|
|
./no-saved-proto.patch
|
|
|
|
# https://github.com/tensorflow/tensorflow/pull/29673
|
|
(fetchpatch {
|
|
name = "fix-compile-with-cuda-and-mpi.patch";
|
|
url = "https://github.com/tensorflow/tensorflow/pull/29673/commits/498e35a3bfe38dd75cf1416a1a23c07c3b59e6af.patch";
|
|
sha256 = "1m2qmwv1ysqa61z6255xggwbq6mnxbig749bdvrhnch4zydxb4di";
|
|
})
|
|
|
|
# https://github.com/tensorflow/tensorflow/issues/29220
|
|
(fetchpatch {
|
|
name = "bazel-0.27.patch";
|
|
url = "https://github.com/tensorflow/tensorflow/commit/cfccbdb8c4a92dd26382419dceb4d934c2380391.patch";
|
|
sha256 = "1l56wjia2c4685flsfkkgy471wx3c66wyv8khspv06zchj0k0liw";
|
|
})
|
|
];
|
|
|
|
# On update, it can be useful to steal the changes from gentoo
|
|
# https://gitweb.gentoo.org/repo/gentoo.git/tree/sci-libs/tensorflow
|
|
|
|
nativeBuildInputs = [ swig which ];
|
|
|
|
buildInputs = [
|
|
python
|
|
jemalloc
|
|
openmpi
|
|
glibcLocales
|
|
git
|
|
|
|
# python deps needed during wheel build time
|
|
numpy
|
|
keras-preprocessing
|
|
|
|
# libs taken from system through the TF_SYS_LIBS mechanism
|
|
absl-py
|
|
|
|
# for building the wheel
|
|
setuptools
|
|
wheel
|
|
] ++ lib.optionals (!isPy3k) [
|
|
future
|
|
mock
|
|
] ++ lib.optionals cudaSupport [
|
|
cudatoolkit
|
|
cudnn
|
|
nvidia_x11
|
|
];
|
|
|
|
# Take as many libraries from the system as possible. Keep in sync with
|
|
# list of valid syslibs in
|
|
# https://github.com/perfinion/tensorflow/blob/master/third_party/systemlibs/syslibs_configure.bzl
|
|
SYSLIBS= [
|
|
"absl_py"
|
|
"astor_archive"
|
|
"boringssl"
|
|
"com_github_googleapis_googleapis"
|
|
"com_github_googlecloudplatform_google_cloud_cpp"
|
|
"com_google_protobuf"
|
|
"com_google_protobuf_cc"
|
|
"com_googlesource_code_re2"
|
|
"curl"
|
|
"cython"
|
|
"double_conversion"
|
|
"enum34_archive"
|
|
"flatbuffers"
|
|
"gast_archive"
|
|
"gif_archive"
|
|
"grpc"
|
|
"hwloc"
|
|
"icu"
|
|
"jpeg"
|
|
"jsoncpp_git"
|
|
"keras_applications_archive"
|
|
"lmdb"
|
|
"nasm"
|
|
# "nsync" # not packaged in nixpkgs
|
|
"sqlite"
|
|
"pasta"
|
|
"pcre"
|
|
"png_archive"
|
|
"protobuf_archive"
|
|
"six_archive"
|
|
"snappy"
|
|
"swig"
|
|
"termcolor_archive"
|
|
"wrapt"
|
|
"zlib_archive"
|
|
];
|
|
|
|
preConfigure = ''
|
|
patchShebangs configure
|
|
|
|
# dummy ldconfig
|
|
mkdir dummy-ldconfig
|
|
echo "#!${stdenv.shell}" > dummy-ldconfig/ldconfig
|
|
chmod +x dummy-ldconfig/ldconfig
|
|
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 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"
|
|
'';
|
|
|
|
# ar, cpp, gcov, strip, gcc, compat-ld, ld, objdump, dwp, nm, objcopy
|
|
configurePhase = ''
|
|
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
|
|
|
|
runHook postConfigure
|
|
'';
|
|
|
|
# FIXME
|
|
NIX_LDFLAGS = lib.optionals cudaSupport [ "-lcublas" "-lcudnn" "-lcuda" "-lcudart" ];
|
|
|
|
hardeningDisable = [ "all" ];
|
|
|
|
bazelFlags = [
|
|
# temporary fixes to make the build work with bazel 0.27
|
|
"--incompatible_no_support_tools_in_action_inputs=false"
|
|
] ++ lib.optional sse42Support "--copt=-msse4.2"
|
|
++ lib.optional avx2Support "--copt=-mavx2"
|
|
++ lib.optional fmaSupport "--copt=-mfma";
|
|
|
|
bazelTarget = "//tensorflow/tools/pip_package:build_pip_package";
|
|
|
|
fetchAttrs = {
|
|
preInstall = ''
|
|
rm -rf $bazelOut/external/{bazel_tools,\@bazel_tools.marker,local_*,\@local_*}
|
|
'';
|
|
|
|
# cudaSupport causes fetch of ncclArchive, resulting in different hashes
|
|
sha256 = if cudaSupport then
|
|
"0q1pmw7fzn6l554ap576r48m0zgwb7n1ljhyy1p36708z94scdh4"
|
|
else
|
|
"14n26h2r4w7wd5sddy4w0s51s2qcwf276n3hvv2505iysa8wqlc3";
|
|
};
|
|
|
|
buildAttrs = {
|
|
preBuild = ''
|
|
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 = ''
|
|
# work around timestamp issues
|
|
# https://github.com/NixOS/nixpkgs/issues/270#issuecomment-467583872
|
|
export SOURCE_DATE_EPOCH=315532800
|
|
|
|
# bulid the wheel, then move it to $out (building directly to $out
|
|
# would actually put it into a *directory* called $out, but we want the
|
|
# file itself in $out)
|
|
bazel-bin/tensorflow/tools/pip_package/build_pip_package --dst $PWD/dist
|
|
mv dist/*.whl "$out"
|
|
'';
|
|
};
|
|
};
|
|
|
|
in buildPythonPackage rec {
|
|
inherit version pname;
|
|
|
|
src = bazel-wheel;
|
|
|
|
format = "wheel";
|
|
|
|
# Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
|
|
# and the propagated input tensorflow-tensorboard, which causes environment collisions.
|
|
# Another possibility would be to have tensorboard only in the buildInputs
|
|
# https://github.com/tensorflow/tensorflow/blob/v1.7.1/tensorflow/tools/pip_package/setup.py#L79
|
|
postInstall = ''
|
|
rm $out/bin/tensorboard
|
|
'';
|
|
|
|
# tensorflow/tools/pip_package/setup.py
|
|
propagatedBuildInputs = [
|
|
absl-py
|
|
astor
|
|
gast
|
|
google-pasta
|
|
keras-applications
|
|
keras-preprocessing
|
|
numpy
|
|
six
|
|
protobuf
|
|
tensorflow-estimator
|
|
termcolor
|
|
wrapt
|
|
grpcio
|
|
] ++ lib.optionals (!isPy3k) [
|
|
mock
|
|
future # FIXME
|
|
] ++ lib.optionals (pythonOlder "3.4") [
|
|
backports_weakref enum34
|
|
] ++ lib.optionals withTensorboard [
|
|
tensorflow-tensorboard
|
|
];
|
|
|
|
# Actual tests are slow and impure.
|
|
# TODO try to run them anyway
|
|
# TODO better test (files in tensorflow/tools/ci_build/builds/*test)
|
|
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 = platforms.linux;
|
|
broken = !(xlaSupport -> cudaSupport);
|
|
};
|
|
}
|