diff options
Diffstat (limited to 'pkgs/development/python-modules/tensorflow/default.nix')
| -rw-r--r-- | pkgs/development/python-modules/tensorflow/default.nix | 582 |
1 files changed, 337 insertions, 245 deletions
diff --git a/pkgs/development/python-modules/tensorflow/default.nix b/pkgs/development/python-modules/tensorflow/default.nix index 069b87b76f843..d311edc188ad6 100644 --- a/pkgs/development/python-modules/tensorflow/default.nix +++ b/pkgs/development/python-modules/tensorflow/default.nix @@ -1,35 +1,90 @@ -{ stdenv, bazel_5, buildBazelPackage, lib, fetchFromGitHub, symlinkJoin -, addOpenGLRunpath, fetchpatch, fetchzip, linkFarm -# Python deps -, buildPythonPackage, pythonOlder, python -# Python libraries -, numpy, tensorboard, abseil-cpp, absl-py -, packaging, setuptools, wheel, keras-preprocessing, google-pasta -, opt-einsum, astunparse, h5py -, termcolor, grpcio, six, wrapt, protobuf-python, tensorflow-estimator-bin -, dill, flatbuffers-python, portpicker, tblib, typing-extensions -# Common deps -, git, pybind11, which, binutils, glibcLocales, cython, perl -# Common libraries -, jemalloc, mpi, gast, grpc, sqlite, boringssl, jsoncpp, nsync -, curl, snappy, flatbuffers-core, icu, double-conversion, libpng, libjpeg_turbo, giflib, protobuf-core -# Upstream 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 -, config -, cudaSupport ? config.cudaSupport -, cudaPackagesGoogle -, cudaCapabilities ? cudaPackagesGoogle.cudaFlags.cudaCapabilities -, mklSupport ? false, mkl -, tensorboardSupport ? true -# XLA without CUDA is broken -, xlaSupport ? cudaSupport -, sse42Support ? stdenv.hostPlatform.sse4_2Support -, avx2Support ? stdenv.hostPlatform.avx2Support -, fmaSupport ? stdenv.hostPlatform.fmaSupport -# Darwin deps -, Foundation, Security, cctools, llvmPackages +{ + stdenv, + bazel_5, + buildBazelPackage, + lib, + fetchFromGitHub, + symlinkJoin, + addOpenGLRunpath, + fetchpatch, + fetchzip, + linkFarm, + # Python deps + buildPythonPackage, + pythonAtLeast, + pythonOlder, + python, + # Python libraries + numpy, + tensorboard, + abseil-cpp, + absl-py, + packaging, + setuptools, + wheel, + keras-preprocessing, + google-pasta, + opt-einsum, + astunparse, + h5py, + termcolor, + grpcio, + six, + wrapt, + protobuf-python, + tensorflow-estimator-bin, + dill, + flatbuffers-python, + portpicker, + tblib, + typing-extensions, + # Common deps + git, + pybind11, + which, + binutils, + glibcLocales, + cython, + perl, + # Common libraries + jemalloc, + mpi, + gast, + grpc, + sqlite, + boringssl, + jsoncpp, + nsync, + curl, + snappy, + flatbuffers-core, + icu, + double-conversion, + libpng, + libjpeg_turbo, + giflib, + protobuf-core, + # Upstream 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 + config, + cudaSupport ? config.cudaSupport, + cudaPackages, + cudaCapabilities ? cudaPackages.cudaFlags.cudaCapabilities, + mklSupport ? false, + mkl, + tensorboardSupport ? true, + # XLA without CUDA is broken + xlaSupport ? cudaSupport, + sse42Support ? stdenv.hostPlatform.sse4_2Support, + avx2Support ? stdenv.hostPlatform.avx2Support, + fmaSupport ? stdenv.hostPlatform.fmaSupport, + # Darwin deps + Foundation, + Security, + cctools, + llvmPackages, }: let @@ -50,21 +105,27 @@ let # __ZN4llvm11SmallPtrSetIPKNS_10AllocaInstELj8EED1Ev in any of the # translation units, so the build fails at link time stdenv = - if cudaSupport then cudaPackagesGoogle.backendStdenv - else if originalStdenv.isDarwin then llvmPackages.stdenv - else originalStdenv; - inherit (cudaPackagesGoogle) cudatoolkit nccl; + if cudaSupport then + cudaPackages.backendStdenv + else if originalStdenv.isDarwin then + llvmPackages.stdenv + else + originalStdenv; + inherit (cudaPackages) cudatoolkit nccl; # use compatible cuDNN (https://www.tensorflow.org/install/source#gpu) # cudaPackages.cudnn led to this: # https://github.com/tensorflow/tensorflow/issues/60398 cudnnAttribute = "cudnn_8_6"; - cudnn = cudaPackagesGoogle.${cudnnAttribute}; + cudnn = cudaPackages.${cudnnAttribute}; gentoo-patches = fetchzip { url = "https://dev.gentoo.org/~perfinion/patches/tensorflow-patches-2.12.0.tar.bz2"; hash = "sha256-SCRX/5/zML7LmKEPJkcM5Tebez9vv/gmE4xhT/jyqWs="; }; protobuf-extra = linkFarm "protobuf-extra" [ - { name = "include"; path = protobuf-core.src; } + { + name = "include"; + path = protobuf-core.src; + } ]; withTensorboard = (pythonOlder "3.6") || tensorboardSupport; @@ -72,14 +133,16 @@ let # FIXME: migrate to redist cudaPackages cudatoolkit_joined = symlinkJoin { name = "${cudatoolkit.name}-merged"; - paths = [ - cudatoolkit.lib - cudatoolkit.out - ] ++ lib.optionals (lib.versionOlder cudatoolkit.version "11") [ - # 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}/targets/${stdenv.system}" - ]; + paths = + [ + cudatoolkit.lib + cudatoolkit.out + ] + ++ lib.optionals (lib.versionOlder cudatoolkit.version "11") [ + # 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}/targets/${stdenv.system}" + ]; }; # Tensorflow expects bintools at hard-coded paths, e.g. /usr/bin/ar @@ -96,9 +159,7 @@ let # Needed for _some_ system libraries, grep INCLUDEDIR. includes_joined = symlinkJoin { name = "tensorflow-deps-merged"; - paths = [ - jsoncpp - ]; + paths = [ jsoncpp ]; }; tfFeature = x: if x then "1" else "0"; @@ -108,31 +169,31 @@ let variant = lib.optionalString cudaSupport "-gpu"; pname = "tensorflow${variant}"; - pythonEnv = python.withPackages (_: - [ # python deps needed during wheel build time (not runtime, see the buildPythonPackage part for that) - # 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 - keras-preprocessing - numpy - opt-einsum - packaging - protobuf-python - setuptools - six - tblib - tensorboard - tensorflow-estimator-bin - termcolor - typing-extensions - wheel - wrapt + pythonEnv = python.withPackages (_: [ + # python deps needed during wheel build time (not runtime, see the buildPythonPackage part for that) + # 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 + keras-preprocessing + numpy + opt-einsum + packaging + protobuf-python + setuptools + six + tblib + tensorboard + tensorflow-estimator-bin + termcolor + typing-extensions + wheel + wrapt ]); rules_cc_darwin_patched = stdenv.mkDerivation { @@ -195,15 +256,19 @@ let runHook postInstall ''; }; - bazel-build = if stdenv.isDarwin then _bazel-build.overrideAttrs (prev: { - bazelFlags = prev.bazelFlags ++ [ - "--override_repository=rules_cc=${rules_cc_darwin_patched}" - "--override_repository=llvm-raw=${llvm-raw_darwin_patched}" - ]; - preBuild = '' - export AR="${cctools}/bin/libtool" - ''; - }) else _bazel-build; + bazel-build = + if stdenv.isDarwin then + _bazel-build.overrideAttrs (prev: { + bazelFlags = prev.bazelFlags ++ [ + "--override_repository=rules_cc=${rules_cc_darwin_patched}" + "--override_repository=llvm-raw=${llvm-raw_darwin_patched}" + ]; + preBuild = '' + export AR="${cctools}/bin/libtool" + ''; + }) + else + _bazel-build; _bazel-build = buildBazelPackage.override { inherit stdenv; } { name = "${pname}-${version}"; @@ -220,42 +285,50 @@ let # https://gitweb.gentoo.org/repo/gentoo.git/tree/sci-libs/tensorflow nativeBuildInputs = [ - which pythonEnv cython perl protobuf-core protobuf-extra + which + pythonEnv + cython + perl + protobuf-core + protobuf-extra ] ++ lib.optional cudaSupport addOpenGLRunpath; - buildInputs = [ - jemalloc - mpi - glibcLocales - git - - # libs taken from system through the TF_SYS_LIBS mechanism - abseil-cpp - boringssl - curl - double-conversion - flatbuffers-core - giflib - grpc - # Necessary to fix the "`GLIBCXX_3.4.30' not found" error - (icu.override { inherit stdenv; }) - jsoncpp - libjpeg_turbo - libpng - (pybind11.overridePythonAttrs (_: { inherit stdenv; })) - snappy - sqlite - ] ++ lib.optionals cudaSupport [ - cudatoolkit - cudnn - ] ++ lib.optionals mklSupport [ - mkl - ] ++ lib.optionals stdenv.isDarwin [ - Foundation - Security - ] ++ lib.optionals (!stdenv.isDarwin) [ - nsync - ]; + buildInputs = + [ + jemalloc + mpi + glibcLocales + git + + # libs taken from system through the TF_SYS_LIBS mechanism + abseil-cpp + boringssl + curl + double-conversion + flatbuffers-core + giflib + grpc + # Necessary to fix the "`GLIBCXX_3.4.30' not found" error + (icu.override { inherit stdenv; }) + jsoncpp + libjpeg_turbo + libpng + (pybind11.overridePythonAttrs (_: { + inherit stdenv; + })) + snappy + sqlite + ] + ++ lib.optionals cudaSupport [ + cudatoolkit + cudnn + ] + ++ lib.optionals mklSupport [ mkl ] + ++ lib.optionals stdenv.isDarwin [ + Foundation + Security + ] + ++ lib.optionals (!stdenv.isDarwin) [ nsync ]; # arbitrarily set to the current latest bazel version, overly careful TF_IGNORE_MAX_BAZEL_VERSION = true; @@ -265,47 +338,50 @@ let # 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 "," ([ - "absl_py" - "astor_archive" - "astunparse_archive" - "boringssl" - "com_google_absl" - # Not packaged in nixpkgs - # "com_github_googleapis_googleapis" - # "com_github_googlecloudplatform_google_cloud_cpp" - "com_github_grpc_grpc" - "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" - "curl" - "cython" - "dill_archive" - "double_conversion" - "flatbuffers" - "functools32_archive" - "gast_archive" - "gif" - "hwloc" - "icu" - "jsoncpp_git" - "libjpeg_turbo" - "nasm" - "opt_einsum_archive" - "org_sqlite" - "pasta" - "png" - "pybind11" - "six_archive" - "snappy" - "tblib_archive" - "termcolor_archive" - "typing_extensions_archive" - "wrapt" - "zlib" - ] ++ lib.optionals (!stdenv.isDarwin) [ - "nsync" # fails to build on darwin - ]); + TF_SYSTEM_LIBS = lib.concatStringsSep "," ( + [ + "absl_py" + "astor_archive" + "astunparse_archive" + "boringssl" + "com_google_absl" + # Not packaged in nixpkgs + # "com_github_googleapis_googleapis" + # "com_github_googlecloudplatform_google_cloud_cpp" + "com_github_grpc_grpc" + "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" + "curl" + "cython" + "dill_archive" + "double_conversion" + "flatbuffers" + "functools32_archive" + "gast_archive" + "gif" + "hwloc" + "icu" + "jsoncpp_git" + "libjpeg_turbo" + "nasm" + "opt_einsum_archive" + "org_sqlite" + "pasta" + "png" + "pybind11" + "six_archive" + "snappy" + "tblib_archive" + "termcolor_archive" + "typing_extensions_archive" + "wrapt" + "zlib" + ] + ++ lib.optionals (!stdenv.isDarwin) [ + "nsync" # fails to build on darwin + ] + ); INCLUDEDIR = "${includes_joined}/include"; @@ -348,47 +424,51 @@ let ./protobuf_python.patch ./pybind11_protobuf_python_runtime_dep.patch ./pybind11_protobuf_newer_version.patch - ] ++ lib.optionals (stdenv.hostPlatform.system == "aarch64-darwin") [ - ./absl_to_std.patch - ]; + ] ++ lib.optionals (stdenv.hostPlatform.system == "aarch64-darwin") [ ./absl_to_std.patch ]; - postPatch = '' - # bazel 3.3 should work just as well as bazel 3.1 - rm -f .bazelversion - patchShebangs . - '' + lib.optionalString (stdenv.hostPlatform.system == "x86_64-darwin") '' - cat ${./com_google_absl_fix_macos.patch} >> third_party/absl/com_google_absl_fix_mac_and_nvcc_build.patch - '' + lib.optionalString (!withTensorboard) '' - # 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 - ''; + postPatch = + '' + # bazel 3.3 should work just as well as bazel 3.1 + rm -f .bazelversion + patchShebangs . + '' + + lib.optionalString (stdenv.hostPlatform.system == "x86_64-darwin") '' + cat ${./com_google_absl_fix_macos.patch} >> third_party/absl/com_google_absl_fix_mac_and_nvcc_build.patch + '' + + lib.optionalString (!withTensorboard) '' + # 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 + ''; # https://github.com/tensorflow/tensorflow/pull/39470 env.NIX_CFLAGS_COMPILE = toString [ "-Wno-stringop-truncation" ]; - preConfigure = let - opt_flags = [] - ++ lib.optionals sse42Support ["-msse4.2"] - ++ lib.optionals avx2Support ["-mavx2"] - ++ lib.optionals fmaSupport ["-mfma"]; - in '' - patchShebangs configure - - # dummy ldconfig - mkdir dummy-ldconfig - echo "#!${stdenv.shell}" > dummy-ldconfig/ldconfig - chmod +x dummy-ldconfig/ldconfig - export PATH="$PWD/dummy-ldconfig:$PATH" - - export PYTHON_LIB_PATH="$NIX_BUILD_TOP/site-packages" - export CC_OPT_FLAGS="${lib.concatStringsSep " " opt_flags}" - mkdir -p "$PYTHON_LIB_PATH" - - # To avoid mixing Python 2 and Python 3 - unset PYTHONPATH - ''; + preConfigure = + let + opt_flags = + [ ] + ++ lib.optionals sse42Support [ "-msse4.2" ] + ++ lib.optionals avx2Support [ "-mavx2" ] + ++ lib.optionals fmaSupport [ "-mfma" ]; + in + '' + patchShebangs configure + + # dummy ldconfig + mkdir dummy-ldconfig + echo "#!${stdenv.shell}" > dummy-ldconfig/ldconfig + chmod +x dummy-ldconfig/ldconfig + export PATH="$PWD/dummy-ldconfig:$PATH" + + export PYTHON_LIB_PATH="$NIX_BUILD_TOP/site-packages" + export CC_OPT_FLAGS="${lib.concatStringsSep " " opt_flags}" + mkdir -p "$PYTHON_LIB_PATH" + + # To avoid mixing Python 2 and Python 3 + unset PYTHONPATH + ''; configurePhase = '' runHook preConfigure @@ -398,39 +478,53 @@ let hardeningDisable = [ "format" ]; - bazelBuildFlags = [ - "--config=opt" # optimize using the flags set in the configure phase - ] - ++ lib.optionals stdenv.cc.isClang [ - "--cxxopt=-x" "--cxxopt=c++" - "--host_cxxopt=-x" "--host_cxxopt=c++" - - # workaround for https://github.com/bazelbuild/bazel/issues/15359 - "--spawn_strategy=sandboxed" - ] - ++ lib.optionals (mklSupport) [ "--config=mkl" ]; - - bazelTargets = [ "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow" ]; + bazelBuildFlags = + [ + "--config=opt" # optimize using the flags set in the configure phase + ] + ++ lib.optionals stdenv.cc.isClang [ + "--cxxopt=-x" + "--cxxopt=c++" + "--host_cxxopt=-x" + "--host_cxxopt=c++" + + # workaround for https://github.com/bazelbuild/bazel/issues/15359 + "--spawn_strategy=sandboxed" + ] + ++ lib.optionals (mklSupport) [ "--config=mkl" ]; + + bazelTargets = [ + "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow" + ]; removeRulesCC = false; # Without this Bazel complaints about sandbox violations. dontAddBazelOpts = true; fetchAttrs = { - sha256 = { - x86_64-linux = if cudaSupport - then "sha256-5VFMNHeLrUxW5RTr6EhT3pay9nWJ5JkZTGirDds5QkU=" - else "sha256-KzgWV69Btr84FdwQ5JI2nQEsqiPg1/+TWdbw5bmxXOE="; - aarch64-linux = if cudaSupport - then "sha256-ty5+51BwHWE1xR4/0WcWTp608NzSAS/iiyN+9zx7/wI=" - else "sha256-9btXrNHqd720oXTPDhSmFidv5iaZRLjCVX8opmrMjXk="; - x86_64-darwin = "sha256-gqb03kB0z2pZQ6m1fyRp1/Nbt8AVVHWpOJSeZNCLc4w="; - aarch64-darwin = "sha256-WdgAaFZU+ePwWkVBhLzjlNT7ELfGHOTaMdafcAMD5yo="; - }.${stdenv.hostPlatform.system} or (throw "unsupported system ${stdenv.hostPlatform.system}"); + sha256 = + { + x86_64-linux = + if cudaSupport then + "sha256-5VFMNHeLrUxW5RTr6EhT3pay9nWJ5JkZTGirDds5QkU=" + else + "sha256-KzgWV69Btr84FdwQ5JI2nQEsqiPg1/+TWdbw5bmxXOE="; + aarch64-linux = + if cudaSupport then + "sha256-ty5+51BwHWE1xR4/0WcWTp608NzSAS/iiyN+9zx7/wI=" + else + "sha256-9btXrNHqd720oXTPDhSmFidv5iaZRLjCVX8opmrMjXk="; + x86_64-darwin = "sha256-gqb03kB0z2pZQ6m1fyRp1/Nbt8AVVHWpOJSeZNCLc4w="; + aarch64-darwin = "sha256-WdgAaFZU+ePwWkVBhLzjlNT7ELfGHOTaMdafcAMD5yo="; + } + .${stdenv.hostPlatform.system} or (throw "unsupported system ${stdenv.hostPlatform.system}"); }; buildAttrs = { - outputs = [ "out" "python" ]; + outputs = [ + "out" + "python" + ]; # need to rebuild schemas since we use a different flatbuffers version preBuild = '' @@ -474,33 +568,34 @@ let done ''; - requiredSystemFeatures = [ - "big-parallel" - ]; + requiredSystemFeatures = [ "big-parallel" ]; }; - meta = with lib; { - badPlatforms = lib.optionals cudaSupport lib.platforms.darwin; - changelog = "https://github.com/tensorflow/tensorflow/releases/tag/v${version}"; - description = "Computation using data flow graphs for scalable machine learning"; - homepage = "http://tensorflow.org"; - license = licenses.asl20; - maintainers = with maintainers; [ abbradar ]; - platforms = with platforms; linux ++ darwin; - broken = - stdenv.isDarwin - || !(xlaSupport -> cudaSupport) - || !(cudaSupport -> builtins.hasAttr cudnnAttribute cudaPackagesGoogle) - || !(cudaSupport -> cudaPackagesGoogle ? cudatoolkit); - } // lib.optionalAttrs stdenv.isDarwin { - timeout = 86400; # 24 hours - maxSilent = 14400; # 4h, double the default of 7200s - }; + meta = + with lib; + { + badPlatforms = lib.optionals cudaSupport lib.platforms.darwin; + changelog = "https://github.com/tensorflow/tensorflow/releases/tag/v${version}"; + description = "Computation using data flow graphs for scalable machine learning"; + homepage = "http://tensorflow.org"; + license = licenses.asl20; + maintainers = with maintainers; [ abbradar ]; + platforms = with platforms; linux ++ darwin; + broken = + stdenv.isDarwin + || !(xlaSupport -> cudaSupport) + || !(cudaSupport -> builtins.hasAttr cudnnAttribute cudaPackages) + || !(cudaSupport -> cudaPackages ? cudatoolkit); + } + // lib.optionalAttrs stdenv.isDarwin { + timeout = 86400; # 24 hours + maxSilent = 14400; # 4h, double the default of 7200s + }; }; - -in buildPythonPackage { +in +buildPythonPackage { inherit version pname; - disabled = pythonOlder "3.8"; + disabled = pythonAtLeast "3.12"; src = bazel-build.python; @@ -550,9 +645,7 @@ in buildPythonPackage { termcolor typing-extensions wrapt - ] ++ lib.optionals withTensorboard [ - tensorboard - ]; + ] ++ lib.optionals withTensorboard [ tensorboard ]; nativeBuildInputs = lib.optionals cudaSupport [ addOpenGLRunpath ]; @@ -594,7 +687,6 @@ in buildPythonPackage { # Regression test for #77626 removed because not more `tensorflow.contrib`. passthru = { - cudaPackages = cudaPackagesGoogle; deps = bazel-build.deps; libtensorflow = bazel-build.out; }; |