2021-01-15 22:18:51 +00:00
{ stdenv , bazel_3 , buildBazelPackage , isPy3k , lib , fetchFromGitHub , symlinkJoin
2021-08-05 21:33:18 +00:00
, addOpenGLRunpath , fetchpatch
2020-04-24 23:36:52 +00:00
# Python deps
2021-01-15 22:18:51 +00:00
, buildPythonPackage , pythonOlder , pythonAtLeast , python
2020-04-24 23:36:52 +00:00
# Python libraries
2021-01-15 22:18:51 +00:00
, numpy , tensorflow-tensorboard_2 , absl-py
, future , setuptools , wheel , keras-preprocessing , google-pasta
2020-10-07 09:15:18 +00:00
, opt-einsum , astunparse , h5py
2020-04-24 23:36:52 +00:00
, termcolor , grpcio , six , wrapt , protobuf , tensorflow-estimator_2
2021-01-15 22:18:51 +00:00
, dill , flatbuffers-python , tblib , typing-extensions
2020-04-24 23:36:52 +00:00
# Common deps
2021-01-15 22:18:51 +00:00
, git , pybind11 , which , binutils , glibcLocales , cython , perl
2020-04-24 23:36:52 +00:00
# Common libraries
2021-02-05 17:12:51 +00:00
, jemalloc , mpi , gast , grpc , sqlite , boringssl , jsoncpp
2021-01-15 22:18:51 +00:00
, curl , snappy , flatbuffers-core , lmdb-core , icu , double-conversion , libpng , libjpeg_turbo , giflib
2020-04-24 23:36:52 +00:00
# Upsteam by default includes cuda support since tensorflow 1.15. We could do
# that in nix as well. It would make some things easier and less confusing, but
# it would also make the default tensorflow package unfree. See
# https://groups.google.com/a/tensorflow.org/forum/#!topic/developers/iRCt5m4qUz0
2020-10-07 09:15:18 +00:00
, cudaSupport ? false , cudatoolkit ? null , cudnn ? null , nccl ? null
2020-04-24 23:36:52 +00:00
, mklSupport ? false , mkl ? null
2021-01-15 22:18:51 +00:00
, tensorboardSupport ? true
2020-04-24 23:36:52 +00:00
# XLA without CUDA is broken
, xlaSupport ? cudaSupport
# Default from ./configure script
2020-11-03 02:18:15 +00:00
, cudaCapabilities ? [ " s m _ 3 5 " " s m _ 5 0 " " s m _ 6 0 " " s m _ 7 0 " " s m _ 7 5 " " c o m p u t e _ 8 0 " ]
2020-09-25 04:45:31 +00:00
, sse42Support ? stdenv . hostPlatform . sse4_2Support
, avx2Support ? stdenv . hostPlatform . avx2Support
, fmaSupport ? stdenv . hostPlatform . fmaSupport
2020-04-24 23:36:52 +00:00
# Darwin deps
, Foundation , Security
} :
2020-10-07 09:15:18 +00:00
assert cudaSupport -> cudatoolkit != null
2020-04-24 23:36:52 +00:00
&& cudnn != null ;
# unsupported combination
assert ! ( stdenv . isDarwin && cudaSupport ) ;
assert mklSupport -> mkl != null ;
let
2021-01-15 22:18:51 +00:00
withTensorboard = ( pythonOlder " 3 . 6 " ) || tensorboardSupport ;
2020-04-24 23:36:52 +00:00
cudatoolkit_joined = symlinkJoin {
name = " ${ cudatoolkit . name } - m e r g e d " ;
paths = [
cudatoolkit . lib
cudatoolkit . out
2020-10-07 09:15:18 +00:00
] ++ lib . optionals ( lib . versionOlder cudatoolkit . version " 1 1 " ) [
2020-04-24 23:36:52 +00:00
# for some reason some of the required libs are in the targets/x86_64-linux
# directory; not sure why but this works around it
" ${ cudatoolkit } / t a r g e t s / ${ stdenv . system } "
] ;
} ;
cudatoolkit_cc_joined = symlinkJoin {
name = " ${ cudatoolkit . cc . name } - m e r g e d " ;
paths = [
cudatoolkit . cc
binutils . bintools # for ar, dwp, nm, objcopy, objdump, strip
] ;
} ;
# Needed for _some_ system libraries, grep INCLUDEDIR.
includes_joined = symlinkJoin {
name = " t e n s o r f l o w - d e p s - m e r g e d " ;
paths = [
jsoncpp
] ;
} ;
tfFeature = x : if x then " 1 " else " 0 " ;
2021-08-05 21:33:18 +00:00
version = " 2 . 4 . 2 " ;
2020-04-24 23:36:52 +00:00
variant = if cudaSupport then " - g p u " else " " ;
pname = " t e n s o r f l o w ${ variant } " ;
pythonEnv = python . withPackages ( _ :
[ # python deps needed during wheel build time (not runtime, see the buildPythonPackage part for that)
2021-01-15 22:18:51 +00:00
# This list can likely be shortened, but each trial takes multiple hours so won't bother for now.
absl-py
astunparse
dill
flatbuffers-python
gast
google-pasta
grpcio
h5py
2020-04-24 23:36:52 +00:00
keras-preprocessing
2021-01-15 22:18:51 +00:00
numpy
opt-einsum
2020-04-24 23:36:52 +00:00
protobuf
setuptools
2021-01-15 22:18:51 +00:00
six
tblib
tensorflow-estimator_2
tensorflow-tensorboard_2
termcolor
typing-extensions
2020-04-24 23:36:52 +00:00
wheel
2021-01-15 22:18:51 +00:00
wrapt
2020-04-24 23:36:52 +00:00
] ) ;
bazel-build = buildBazelPackage {
name = " ${ pname } - ${ version } " ;
2020-10-07 09:15:18 +00:00
bazel = bazel_3 ;
2020-04-24 23:36:52 +00:00
src = fetchFromGitHub {
owner = " t e n s o r f l o w " ;
repo = " t e n s o r f l o w " ;
rev = " v ${ version } " ;
2021-08-05 21:33:18 +00:00
sha256 = " 0 7 a 2 y 0 5 h i x c h 1 b j a g 5 p z w 3 p 1 m 7 b d j 3 b q 4 g d v m s f k 2 x r a z 4 9 b 1 p i 8 " ;
2020-04-24 23:36:52 +00:00
} ;
patches = [
2021-08-05 21:33:18 +00:00
# included from 2.6.0 onwards
( fetchpatch {
name = " f i x - n u m p y - 1 . 2 0 - n o t i m p l e m e n t e d e r r o r . p a t c h " ;
url = " h t t p s : / / g i t h u b . c o m / t e n s o r f l o w / t e n s o r f l o w / c o m m i t / b 2 5 8 9 4 1 5 2 5 f 4 9 6 7 6 3 d 4 2 7 7 0 4 5 b 6 5 1 3 c 8 1 5 7 2 0 e 3 a . p a t c h " ;
sha256 = " 1 9 f 9 b z r c f s y n k 1 1 s 2 h q v s c i n 5 c 6 5 z f 7 r 6 g 3 n b 1 0 j n i m w 7 9 v a f i r y " ;
} )
2021-01-15 22:18:51 +00:00
# Relax too strict Python packages versions dependencies.
2020-12-25 13:55:36 +00:00
./relax-dependencies.patch
2021-01-15 22:18:51 +00:00
# Add missing `io_bazel_rules_docker` dependency.
./workspace.patch
2020-04-24 23:36:52 +00:00
] ;
# On update, it can be useful to steal the changes from gentoo
# https://gitweb.gentoo.org/repo/gentoo.git/tree/sci-libs/tensorflow
nativeBuildInputs = [
2021-01-15 22:18:51 +00:00
which pythonEnv cython perl
2020-04-24 23:36:52 +00:00
] ++ lib . optional cudaSupport addOpenGLRunpath ;
buildInputs = [
jemalloc
2021-02-05 17:12:51 +00:00
mpi
2020-04-24 23:36:52 +00:00
glibcLocales
git
# libs taken from system through the TF_SYS_LIBS mechanism
2020-10-07 09:15:18 +00:00
grpc
2020-04-24 23:36:52 +00:00
sqlite
2021-01-15 22:18:51 +00:00
boringssl
2020-04-24 23:36:52 +00:00
jsoncpp
curl
2021-01-15 22:18:51 +00:00
pybind11
2020-04-24 23:36:52 +00:00
snappy
2021-01-15 22:18:51 +00:00
flatbuffers-core
2020-04-24 23:36:52 +00:00
icu
double-conversion
libpng
2021-01-15 22:18:51 +00:00
libjpeg_turbo
2020-04-24 23:36:52 +00:00
giflib
2021-01-15 22:18:51 +00:00
lmdb-core
2020-04-24 23:36:52 +00:00
] ++ lib . optionals cudaSupport [
cudatoolkit
cudnn
] ++ lib . optionals mklSupport [
mkl
] ++ lib . optionals stdenv . isDarwin [
Foundation
Security
] ;
# 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
# list of valid syslibs in
# https://github.com/tensorflow/tensorflow/blob/master/third_party/systemlibs/syslibs_configure.bzl
TF_SYSTEM_LIBS = lib . concatStringsSep " , " [
" a b s l _ p y "
" a s t o r _ a r c h i v e "
2020-10-07 09:15:18 +00:00
" a s t u n p a r s e _ a r c h i v e "
2020-04-24 23:36:52 +00:00
" b o r i n g s s l "
# Not packaged in nixpkgs
# "com_github_googleapis_googleapis"
# "com_github_googlecloudplatform_google_cloud_cpp"
2020-10-07 09:15:18 +00:00
" c o m _ g i t h u b _ g r p c _ g r p c "
2021-01-15 22:18:51 +00:00
# Multiple issues with custom protobuf.
# First `com_github_googleapis` fails to configure. Can be worked around by disabling `com_github_googleapis`
# and related functionality, but then the next error is about "dangling symbolic link", and in general
# looks like that's only the beginning: see
# https://stackoverflow.com/questions/55578884/how-to-build-tensorflow-1-13-1-with-custom-protobuf
# "com_google_protobuf"
# Fails with the error: external/org_tensorflow/tensorflow/core/profiler/utils/tf_op_utils.cc:46:49: error: no matching function for call to 're2::RE2::FullMatch(absl::lts_2020_02_25::string_view&, re2::RE2&)'
# "com_googlesource_code_re2"
2020-04-24 23:36:52 +00:00
" c u r l "
" c y t h o n "
2021-01-15 22:18:51 +00:00
" d i l l _ a r c h i v e "
2020-04-24 23:36:52 +00:00
" d o u b l e _ c o n v e r s i o n "
2020-10-07 09:15:18 +00:00
" e n u m 3 4 _ a r c h i v e "
2020-04-24 23:36:52 +00:00
" f l a t b u f f e r s "
2020-10-07 09:15:18 +00:00
" f u n c t o o l s 3 2 _ a r c h i v e "
2020-04-24 23:36:52 +00:00
" g a s t _ a r c h i v e "
2020-10-07 09:15:18 +00:00
" g i f "
2020-04-24 23:36:52 +00:00
" h w l o c "
" i c u "
" j s o n c p p _ g i t "
2020-10-07 09:15:18 +00:00
" l i b j p e g _ t u r b o "
2020-04-24 23:36:52 +00:00
" l m d b "
" n a s m "
# "nsync" # not packaged in nixpkgs
" o p t _ e i n s u m _ a r c h i v e "
" o r g _ s q l i t e "
" p a s t a "
" p c r e "
2020-10-07 09:15:18 +00:00
" p n g "
" p y b i n d 1 1 "
2020-04-24 23:36:52 +00:00
" s i x _ a r c h i v e "
" s n a p p y "
2021-01-15 22:18:51 +00:00
" t b l i b _ a r c h i v e "
2020-04-24 23:36:52 +00:00
" t e r m c o l o r _ a r c h i v e "
2021-01-15 22:18:51 +00:00
" t y p i n g _ e x t e n s i o n s _ a r c h i v e "
2020-04-24 23:36:52 +00:00
" w r a p t "
2020-10-07 09:15:18 +00:00
" z l i b "
2020-04-24 23:36:52 +00:00
] ;
INCLUDEDIR = " ${ includes_joined } / i n c l u d e " ;
PYTHON_BIN_PATH = pythonEnv . 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 } / b i n " ;
GCC_HOST_COMPILER_PATH = lib . optionalString cudaSupport " ${ cudatoolkit_cc_joined } / b i n / g c c " ;
TF_CUDA_COMPUTE_CAPABILITIES = lib . concatStringsSep " , " cudaCapabilities ;
postPatch = ''
2021-01-15 22:18:51 +00:00
# bazel 3.3 should work just as well as bazel 3.1
rm - f . bazelversion
'' + l i b . o p t i o n a l S t r i n g ( ! w i t h T e n s o r b o a r d ) ''
2020-04-24 23:36:52 +00:00
# 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
2021-01-15 22:18:51 +00:00
sed - i ' /tensorboard ~ = /d ' tensorflow/tools/pip_package/setup.py
2020-04-24 23:36:52 +00:00
'' ;
2021-01-09 10:05:03 +00:00
# https://github.com/tensorflow/tensorflow/pull/39470
NIX_CFLAGS_COMPILE = [ " - W n o - s t r i n g o p - t r u n c a t i o n " ] ;
2020-04-24 23:36:52 +00:00
preConfigure = let
opt_flags = [ ]
++ lib . optionals sse42Support [ " - m s s e 4 . 2 " ]
++ lib . optionals avx2Support [ " - m a v x 2 " ]
++ lib . optionals fmaSupport [ " - m f m a " ] ;
in ''
patchShebangs configure
# dummy ldconfig
mkdir dummy-ldconfig
echo " # ! ${ stdenv . shell } " > dummy-ldconfig/ldconfig
chmod + x dummy-ldconfig/ldconfig
export PATH = " $ P W D / d u m m y - l d c o n f i g : $ P A T H "
export PYTHON_LIB_PATH = " $ N I X _ B U I L D _ T O P / s i t e - p a c k a g e s "
export CC_OPT_FLAGS = " ${ lib . concatStringsSep " " opt_flags } "
mkdir - p " $ P Y T H O N _ L I B _ P A T H "
# To avoid mixing Python 2 and Python 3
unset PYTHONPATH
'' ;
configurePhase = ''
runHook preConfigure
./configure
runHook postConfigure
'' ;
hardeningDisable = [ " f o r m a t " ] ;
bazelBuildFlags = [
" - - c o n f i g = o p t " # optimize using the flags set in the configure phase
]
++ lib . optionals ( mklSupport ) [ " - - c o n f i g = m k l " ] ;
bazelTarget = " / / t e n s o r f l o w / t o o l s / p i p _ p a c k a g e : b u i l d _ p i p _ p a c k a g e / / t e n s o r f l o w / t o o l s / l i b _ p a c k a g e : l i b t e n s o r f l o w " ;
2020-10-07 09:15:18 +00:00
removeRulesCC = false ;
2021-01-15 22:18:51 +00:00
# Without this Bazel complaints about sandbox violations.
dontAddBazelOpts = true ;
2020-10-07 09:15:18 +00:00
2020-04-24 23:36:52 +00:00
fetchAttrs = {
# cudaSupport causes fetch of ncclArchive, resulting in different hashes
sha256 = if cudaSupport then
2021-06-28 23:13:55 +00:00
" 1 0 m 6 q j 3 k c h g x f g b 6 q h 5 9 v c 5 1 k n m 9 r 9 p k n g 8 b f 9 0 h 0 0 d n g g v v 8 2 3 4 "
2020-04-24 23:36:52 +00:00
else
2021-06-28 23:13:55 +00:00
" 0 4 a 9 8 y r p 0 9 n d 0 p 1 7 k 0 j b z k g j p p x s 0 y m a 7 m 5 z k f r w g v r 4 g 0 w 7 1 v 6 8 " ;
2020-04-24 23:36:52 +00:00
} ;
buildAttrs = {
outputs = [ " o u t " " p y t h o n " ] ;
preBuild = ''
patchShebangs .
'' ;
installPhase = ''
mkdir - p " $ o u t "
tar - xf bazel-bin/tensorflow/tools/lib_package/libtensorflow.tar.gz - C " $ o u t "
# Write pkgconfig file.
mkdir " $ o u t / l i b / p k g c o n f i g "
cat > " $ o u t / l i b / p k g c o n f i g / t e n s o r f l o w . p c " < < 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
# build the source code, then copy it to $python (build_pip_package
# actually builds a symlink farm so we must dereference them).
bazel-bin/tensorflow/tools/pip_package/build_pip_package - - src " $ P W D / d i s t "
cp - Lr " $ P W D / d i s t " " $ p y t h o n "
'' ;
postFixup = lib . optionalString cudaSupport ''
find $ out - type f \ ( - name ' * . so' - or - name ' * . so . * ' \ ) | while read lib ; do
addOpenGLRunpath " $ l i b "
done
'' ;
2021-04-26 19:14:03 +00:00
requiredSystemFeatures = [
" b i g - p a r a l l e l "
] ;
2020-04-24 23:36:52 +00:00
} ;
2021-01-15 22:18:51 +00:00
meta = with lib ; {
2020-04-24 23:36:52 +00:00
description = " C o m p u t a t i o n u s i n g d a t a f l o w g r a p h s f o r s c a l a b l e m a c h i n e l e a r n i n g " ;
homepage = " h t t p : / / t e n s o r f l o w . o r g " ;
license = licenses . asl20 ;
maintainers = with maintainers ; [ jyp abbradar ] ;
platforms = with platforms ; linux ++ darwin ;
2021-01-15 22:18:51 +00:00
broken = ! ( xlaSupport -> cudaSupport ) ;
2020-04-24 23:36:52 +00:00
} ;
} ;
in buildPythonPackage {
inherit version pname ;
2021-01-15 22:18:51 +00:00
disabled = ! isPy3k ;
2020-04-24 23:36:52 +00:00
src = bazel-build . python ;
# 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
'' ;
setupPyGlobalFlags = [ " - - p r o j e c t _ n a m e ${ pname } " ] ;
# tensorflow/tools/pip_package/setup.py
propagatedBuildInputs = [
absl-py
2021-01-15 22:18:51 +00:00
astunparse
dill
flatbuffers-python
2020-04-24 23:36:52 +00:00
gast
google-pasta
2021-01-15 22:18:51 +00:00
grpcio
h5py
2020-04-24 23:36:52 +00:00
keras-preprocessing
numpy
2021-01-15 22:18:51 +00:00
opt-einsum
2020-04-24 23:36:52 +00:00
protobuf
2021-01-15 22:18:51 +00:00
six
tblib
2020-04-24 23:36:52 +00:00
tensorflow-estimator_2
termcolor
2021-01-15 22:18:51 +00:00
typing-extensions
2020-04-24 23:36:52 +00:00
wrapt
] ++ lib . optionals withTensorboard [
tensorflow-tensorboard_2
] ;
nativeBuildInputs = lib . optional cudaSupport addOpenGLRunpath ;
postFixup = lib . optionalString cudaSupport ''
find $ out - type f \ ( - name ' * . so' - or - name ' * . so . * ' \ ) | while read lib ; do
addOpenGLRunpath " $ l i b "
2020-10-07 09:15:18 +00:00
patchelf - - set-rpath " ${ cudatoolkit } / l i b : ${ cudatoolkit . lib } / l i b : ${ cudnn } / l i b : ${ nccl } / l i b : $ ( p a t c h e l f - - p r i n t - r p a t h " $ lib " ) " " $ l i b "
2020-04-24 23:36:52 +00:00
done
'' ;
# 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 } < < EOF
# A simple "Hello world"
import tensorflow as tf
hello = tf . constant ( " H e l l o , w o r l d ! " )
tf . print ( hello )
# Fit a simple model to random data
import numpy as np
np . random . seed ( 0 )
tf . random . set_seed ( 0 )
model = tf . keras . models . Sequential ( [
tf . keras . layers . Dense ( 1 , activation = " l i n e a r " )
] )
model . compile ( optimizer = " s g d " , loss = " m s e " )
x = np . random . uniform ( size = ( 1 , 1 ) )
y = np . random . uniform ( size = ( 1 , ) )
model . fit ( x , y , epochs = 1 )
EOF
'' ;
# Regression test for #77626 removed because not more `tensorflow.contrib`.
2020-05-29 06:06:01 +00:00
passthru = {
deps = bazel-build . deps ;
libtensorflow = bazel-build . out ;
} ;
2020-04-24 23:36:52 +00:00
inherit ( bazel-build ) meta ;
}