python3Packages.tensorflow_2: 2.3.2 -> 2.4.0

Also the following related changes: * Removed Python 2 support because it's not supported by TF and related packages for a long time. * Upgraded tensorboard and estimator packages to the required versions. * Added extra plugins for tensorboard to support profiling. * In the previous derivation versions, TF_SYSTEM_LIBS didn't have any effect because it was reset at repo fetching stage, so TF always used its own dependencies. Made TF_SYSTEM_LIBS actually work and fixed the errors caused by enabling it. * Enabled tensorboard by default (but still keeping an option to disable it if needed).
2024-11-15 13:37:21 +03:00 · 2021-01-10 11:50:33 +01:00 · 2021-01-10 11:50:33 +01:00 · 69d1e62b36
commit 69d1e62b36
parent db0d9e8510
11 changed files with 237 additions and 117 deletions
--- a/pkgs/development/python-modules/flatbuffers/default.nix
+++ b/pkgs/development/python-modules/flatbuffers/default.nix
@ -8,6 +8,10 @@ buildPythonPackage rec {

  sourceRoot = "source/python";

+  # flatbuffers needs VERSION environment variable for setting the correct
+  # version, otherwise it uses the current date.
+  VERSION = "${version}";
+
  pythonImportsCheck = [ "flatbuffers" ];

  meta = flatbuffers.meta // {
--- a/pkgs/development/python-modules/gviz-api/default.nix
+++ b/pkgs/development/python-modules/gviz-api/default.nix
@ -0,0 +1,26 @@
+{ lib, fetchPypi, buildPythonPackage
+, six
+}:
+
+buildPythonPackage rec {
+  pname = "gviz_api";
+  version = "1.9.0";
+  format = "wheel";
+
+  src = fetchPypi {
+    inherit pname version;
+    format = "wheel";
+    sha256 = "1yag559lpmwfdxpxn679a6ajifcbpgljr5n6k5b7rrj38k2xq7jg";
+  };
+
+  propagatedBuildInputs = [
+    six
+  ];
+
+  meta = with lib; {
+    description = "Python API for Google Visualization";
+    homepage = https://developers.google.com/chart/interactive/docs/dev/gviz_api_lib;
+    license = licenses.asl20;
+    maintainers = with maintainers; [ ndl ];
+  };
+}
--- a/pkgs/development/python-modules/tensorboard-plugin-profile/default.nix
+++ b/pkgs/development/python-modules/tensorboard-plugin-profile/default.nix
@ -0,0 +1,31 @@
+{ lib, fetchPypi, buildPythonPackage
+, gviz-api
+, protobuf
+, werkzeug
+}:
+
+buildPythonPackage rec {
+  pname = "tensorboard_plugin_profile";
+  version = "2.4.0";
+  format = "wheel";
+
+  src = fetchPypi {
+    inherit pname version;
+    format = "wheel";
+    python = "py3";
+    sha256 = "0z6dcjvkk3pzmmmjxi2ybawnfshz5qa3ga92kqj69ld1g9k3i9bj";
+  };
+
+  propagatedBuildInputs = [
+    gviz-api
+    protobuf
+    werkzeug
+  ];
+
+  meta = with lib; {
+    description = "Profile Tensorboard Plugin.";
+    homepage = http://tensorflow.org;
+    license = licenses.asl20;
+    maintainers = with maintainers; [ ndl ];
+  };
+}
--- a/pkgs/development/python-modules/tensorboard-plugin-wit/default.nix
+++ b/pkgs/development/python-modules/tensorboard-plugin-wit/default.nix
@ -0,0 +1,22 @@
+{ lib, fetchPypi, buildPythonPackage
+}:
+
+buildPythonPackage rec {
+  pname = "tensorboard_plugin_wit";
+  version = "1.7.0";
+  format = "wheel";
+
+  src = fetchPypi {
+    inherit pname version;
+    format = "wheel";
+    python = "py3";
+    sha256 = "0nv855qm2fav70lndsrv810pqgg41sbmd70fk86wk18ih825yxzf";
+  };
+
+  meta = with lib; {
+    description = "What-If Tool TensorBoard plugin.";
+    homepage = http://tensorflow.org;
+    license = licenses.asl20;
+    maintainers = with maintainers; [ ndl ];
+  };
+}
--- a/pkgs/development/python-modules/tensorflow-estimator/default.nix
+++ b/pkgs/development/python-modules/tensorflow-estimator/default.nix
@ -6,13 +6,13 @@

 buildPythonPackage rec {
  pname = "tensorflow-estimator";
-  version = "2.3.0";
+  version = "2.4.0";
  format = "wheel";

  src = fetchPypi {
    pname = "tensorflow_estimator";
    inherit version format;
-    sha256 = "11n4sl9wfr00fv1i837b7a36ink86ggmlsgj7i06kcfc011h6pmp";
+    sha256 = "1w0pkcslm6934qqd6m5gxyjdlnb4pbl47k6s99wsh6dyvvr7nysv";
  };

  propagatedBuildInputs = [ mock numpy absl-py ];
--- a/pkgs/development/python-modules/tensorflow-tensorboard/default.nix
+++ b/pkgs/development/python-modules/tensorflow-tensorboard/default.nix
@ -5,9 +5,10 @@
 , protobuf
 , grpcio
 , markdown
-, futures
 , absl-py
 , google-auth-oauthlib
+, tensorboard-plugin-wit
+, tensorboard-plugin-profile
 }:

 # tensorflow/tensorboard is built from a downloaded wheel, because
@ -16,19 +17,16 @@

 buildPythonPackage rec {
  pname = "tensorflow-tensorboard";
-  version = "2.1.0";
+  version = "2.4.0";
  format = "wheel";
+  disabled = !isPy3k;

-  src = fetchPypi ({
+  src = fetchPypi {
    pname = "tensorboard";
    inherit version format;
-  } // (if isPy3k then {
    python = "py3";
-    sha256 = "1wpjdzhjpcdkyaahzd4bl71k4l30z5c55280ndiwj32hw70lxrp6";
-  } else {
-    python = "py2";
-    sha256 = "1f805839xa36wxb7xac9fyxzaww92vw4d50vs6g61wnlr4byp00w";
-  }));
+    sha256 = "0f17h6i398n8maam0r3rssqvdqnqbwjyf96nnhf482anm1iwdq6d";
+  };

  propagatedBuildInputs = [
    numpy
@ -38,10 +36,12 @@ buildPythonPackage rec {
    grpcio
    absl-py
    google-auth-oauthlib
+    tensorboard-plugin-profile
+    tensorboard-plugin-wit
    # not declared in install_requires, but used at runtime
    # https://github.com/NixOS/nixpkgs/issues/73840
    wheel
-  ] ++ lib.optional (!isPy3k) futures;
+  ];

  # in the absence of a real test suite, run cli and imports
  checkPhase = ''
--- a/pkgs/development/python-modules/tensorflow/default.nix
+++ b/pkgs/development/python-modules/tensorflow/default.nix
@ -1,24 +1,25 @@
-{ stdenv, pkgs, bazel_3, buildBazelPackage, lib, fetchFromGitHub, fetchpatch, symlinkJoin
+{ stdenv, bazel_3, buildBazelPackage, isPy3k, lib, fetchFromGitHub, symlinkJoin
 , addOpenGLRunpath
 # Python deps
-, buildPythonPackage, isPy3k, isPy27, pythonOlder, pythonAtLeast, python
+, buildPythonPackage, pythonOlder, pythonAtLeast, python
 # Python libraries
-, numpy, tensorflow-tensorboard_2, backports_weakref, mock, enum34, absl-py
-, future, setuptools, wheel, keras-preprocessing, keras-applications, google-pasta
-, functools32
+, numpy, tensorflow-tensorboard_2, absl-py
+, future, setuptools, wheel, keras-preprocessing, google-pasta
 , opt-einsum, astunparse, h5py
 , termcolor, grpcio, six, wrapt, protobuf, tensorflow-estimator_2
+, dill, flatbuffers-python, tblib, typing-extensions
 # Common deps
-, git, swig, which, binutils, glibcLocales, cython
+, git, pybind11, which, binutils, glibcLocales, cython, perl
 # Common libraries
-, jemalloc, openmpi, astor, gast, grpc, sqlite, openssl, jsoncpp, re2
-, curl, snappy, flatbuffers, icu, double-conversion, libpng, libjpeg, giflib
+, jemalloc, openmpi, gast, grpc, sqlite, boringssl, jsoncpp
+, curl, snappy, flatbuffers-core, lmdb-core, icu, double-conversion, libpng, libjpeg_turbo, giflib
 # 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
 , cudaSupport ? false, cudatoolkit ? null, cudnn ? null, nccl ? null
 , mklSupport ? false, mkl ? null
+, tensorboardSupport ? true
 # XLA without CUDA is broken
 , xlaSupport ? cudaSupport
 # Default from ./configure script
@ -39,7 +40,7 @@ assert ! (stdenv.isDarwin && cudaSupport);
 assert mklSupport -> mkl != null;

 let
-  withTensorboard = pythonOlder "3.6";
+  withTensorboard = (pythonOlder "3.6") || tensorboardSupport;

  cudatoolkit_joined = symlinkJoin {
    name = "${cudatoolkit.name}-merged";
@ -65,34 +66,40 @@ let
  includes_joined = symlinkJoin {
    name = "tensorflow-deps-merged";
    paths = [
-      pkgs.protobuf
      jsoncpp
    ];
  };

  tfFeature = x: if x then "1" else "0";

-  version = "2.3.2";
+  version = "2.4.0";
  variant = if cudaSupport then "-gpu" else "";
  pname = "tensorflow${variant}";

  pythonEnv = python.withPackages (_:
    [ # python deps needed during wheel build time (not runtime, see the buildPythonPackage part for that)
-      numpy
-      keras-preprocessing
-      protobuf
-      wrapt
-      gast
-      astor
+      # This list can likely be shortened, but each trial takes multiple hours so won't bother for now.
      absl-py
-      termcolor
-      keras-applications
+      astunparse
+      dill
+      flatbuffers-python
+      gast
+      google-pasta
+      grpcio
+      h5py
+      keras-preprocessing
+      numpy
+      opt-einsum
+      protobuf
      setuptools
+      six
+      tblib
+      tensorflow-estimator_2
+      tensorflow-tensorboard_2
+      termcolor
+      typing-extensions
      wheel
-  ] ++ lib.optionals (!isPy3k)
-  [ future
-    functools32
-    mock
+      wrapt
  ]);

  bazel-build = buildBazelPackage {
@ -103,27 +110,21 @@ let
      owner = "tensorflow";
      repo = "tensorflow";
      rev = "v${version}";
-      sha256 = "sha256-ncwIkqLDqrB33pB9/FTlBklsIJUEvnDUmyAeUfufCFs=";
+      sha256 = "0yl06aypfxrcs35828xf04mkidz1x0j89v0q5h4d2xps1cb5rv3f";
    };

    patches = [
-      # Fixes for NixOS jsoncpp
-      ./system-jsoncpp.patch
-
+      # Relax too strict Python packages versions dependencies.
      ./relax-dependencies.patch
-
-      # see https://github.com/tensorflow/tensorflow/issues/40688
-      (fetchpatch {
-        url = "https://github.com/tensorflow/tensorflow/commit/75ea0b31477d6ba9e990e296bbbd8ca4e7eebadf.patch";
-        sha256 = "1xp1icacig0xm0nmb05sbrf4nw4xbln9fhc308birrv8286zx7wv";
-      })
+      # Add missing `io_bazel_rules_docker` dependency.
+      ./workspace.patch
    ];

    # 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 pythonEnv
+      which pythonEnv cython perl
    ] ++ lib.optional cudaSupport addOpenGLRunpath;

    buildInputs = [
@ -135,19 +136,18 @@ let
      # libs taken from system through the TF_SYS_LIBS mechanism
      grpc
      sqlite
-      openssl
+      boringssl
      jsoncpp
-      pkgs.protobuf
      curl
+      pybind11
      snappy
-      flatbuffers
+      flatbuffers-core
      icu
      double-conversion
      libpng
-      libjpeg
+      libjpeg_turbo
      giflib
-      re2
-      pkgs.lmdb
+      lmdb-core
    ] ++ lib.optionals cudaSupport [
      cudatoolkit
      cudnn
@ -173,10 +173,17 @@ let
      # "com_github_googleapis_googleapis"
      # "com_github_googlecloudplatform_google_cloud_cpp"
      "com_github_grpc_grpc"
-      "com_google_protobuf"
-      "com_googlesource_code_re2"
+      # 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"
      "curl"
      "cython"
+      "dill_archive"
      "double_conversion"
      "enum34_archive"
      "flatbuffers"
@ -198,8 +205,9 @@ let
      "pybind11"
      "six_archive"
      "snappy"
-      "swig"
+      "tblib_archive"
      "termcolor_archive"
+      "typing_extensions_archive"
      "wrapt"
      "zlib"
    ];
@ -224,16 +232,13 @@ let
    TF_CUDA_COMPUTE_CAPABILITIES = lib.concatStringsSep "," cudaCapabilities;

    postPatch = ''
+      # bazel 3.3 should work just as well as bazel 3.1
+      rm -f .bazelversion
+    '' + 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
-
-      # numpy 1.19 added in https://github.com/tensorflow/tensorflow/commit/75ea0b31477d6ba9e990e296bbbd8ca4e7eebadf.patch
-      sed -i 's/numpy >= 1.16.0, < 1.19.0/numpy >= 1.16.0/' tensorflow/tools/pip_package/setup.py
-
-      # bazel 3.3 should work just as well as bazel 3.1
-      rm -f .bazelversion
+      sed -i '/tensorboard ~=/d' tensorflow/tools/pip_package/setup.py
    '';

    # https://github.com/tensorflow/tensorflow/pull/39470
@ -277,16 +282,15 @@ let
    bazelTarget = "//tensorflow/tools/pip_package:build_pip_package //tensorflow/tools/lib_package:libtensorflow";

    removeRulesCC = false;
+    # Without this Bazel complaints about sandbox violations.
+    dontAddBazelOpts = true;

    fetchAttrs = {
-      # So that checksums don't depend on these.
-      TF_SYSTEM_LIBS = null;
-
      # cudaSupport causes fetch of ncclArchive, resulting in different hashes
      sha256 = if cudaSupport then
-        "sha256-lEdPA9vhYO6vd5FgPMbFp2PkRvDBurPidYsxtJLXcbQ="
+        "0vyy1hv0jy5pqwvnc8pxb9isgnbw07c4a4d4wn61db00np114crz"
      else
-        "sha256-ZEY/bWo5M3Juw1x3CwhXYXZHD4q5LzWDlhgXnh4P95U=";
+        "0vczv5f9s4dxgwdkmf1y9b9ybh5d3y1nllqhb5q8aj9kq73izyn9";
    };

    buildAttrs = {
@ -329,15 +333,13 @@ let
      license = licenses.asl20;
      maintainers = with maintainers; [ jyp abbradar ];
      platforms = with platforms; linux ++ darwin;
-      # The py2 build fails due to some issue importing protobuf. Possibly related to the fix in
-      # https://github.com/akesandgren/easybuild-easyblocks/commit/1f2e517ddfd1b00a342c6abb55aef3fd93671a2b
-      broken = !(xlaSupport -> cudaSupport) || !isPy3k;
+      broken = !(xlaSupport -> cudaSupport);
    };
  };

 in buildPythonPackage {
  inherit version pname;
-  disabled = isPy27;
+  disabled = !isPy3k;

  src = bazel-build.python;

@ -354,27 +356,23 @@ in buildPythonPackage {
  # tensorflow/tools/pip_package/setup.py
  propagatedBuildInputs = [
    absl-py
-    astor
+    astunparse
+    dill
+    flatbuffers-python
    gast
    google-pasta
-    keras-applications
+    grpcio
+    h5py
    keras-preprocessing
    numpy
-    six
+    opt-einsum
    protobuf
+    six
+    tblib
    tensorflow-estimator_2
    termcolor
+    typing-extensions
    wrapt
-    grpcio
-    opt-einsum
-    astunparse
-    h5py
-  ] ++ lib.optionals (!isPy3k) [
-    mock
-    future
-    functools32
-  ] ++ lib.optionals (pythonOlder "3.4") [
-    backports_weakref enum34
  ] ++ lib.optionals withTensorboard [
    tensorflow-tensorboard_2
  ];
--- a/pkgs/development/python-modules/tensorflow/relax-dependencies.patch
+++ b/pkgs/development/python-modules/tensorflow/relax-dependencies.patch
@ -1,16 +1,51 @@
 diff --git a/tensorflow/tools/pip_package/setup.py b/tensorflow/tools/pip_package/setup.py
-index 594e74f40c0..bfbf010144f 100644
+index 65133afdafe..8ef6364ff7e 100644
 --- a/tensorflow/tools/pip_package/setup.py
 +++ b/tensorflow/tools/pip_package/setup.py
-@@ -54,9 +54,9 @@ _VERSION = '2.3.1'
+@@ -75,23 +75,23 @@ if '--project_name' in sys.argv:
+ # comment the versioning scheme.
+ # NOTE: Please add test only packages to `TEST_PACKAGES` below.
 REQUIRED_PACKAGES = [
-     'absl-py >= 0.7.0',
-     'astunparse == 1.6.3',
+-    'absl-py ~= 0.10',
+-    'astunparse ~= 1.6.3',
+-    'flatbuffers ~= 1.12.0',
+-    'google_pasta ~= 0.2',
+-    'h5py ~= 2.10.0',
+-    'keras_preprocessing ~= 1.1.2',
+-    'numpy ~= 1.19.2',
+-    'opt_einsum ~= 3.3.0',
+    'absl-py >= 0.10',
+    'astunparse >= 1.6.3',
+    'flatbuffers >= 1.12.0',
+    'google_pasta >= 0.2',
+    'h5py >= 2.10.0',
+    'keras_preprocessing >= 1.1.2',
+    'numpy >= 1.19.1',
+    'opt_einsum >= 3.3.0',
+     'protobuf >= 3.9.2',
+-    'six ~= 1.15.0',
+-    'termcolor ~= 1.1.0',
+-    'typing_extensions ~= 3.7.4',
+-    'wheel ~= 0.35',
+-    'wrapt ~= 1.12.1',
+    'six >= 1.15.0',
+    'termcolor >= 1.1.0',
+    'typing_extensions >= 3.7.4',
+    'wheel >= 0.34.2',
+    'wrapt >= 1.12.1',
+     # These packages needs to be pinned exactly as newer versions are
+     # incompatible with the rest of the ecosystem
 -    'gast == 0.3.3',
 +    'gast >= 0.3.3',
-     'google_pasta >= 0.1.8',
-    'h5py >= 2.10.0, < 2.11.0',
-+    'h5py >= 2.10.0',
-     'keras_preprocessing >= 1.1.1, < 1.2',
-     # TODO(mihaimaruseac): numpy 1.19.0 has ABI breakage
-     # https://github.com/numpy/numpy/pull/15355
+     # TensorFlow ecosystem packages that TF exposes API for
+     # These need to be in sync with the existing TF version
+     # They are updated during the release process
+@@ -118,7 +118,7 @@ if 'tf_nightly' in project_name:
+ # BoringSSL support.
+ # See https://github.com/tensorflow/tensorflow/issues/17882.
+ if sys.byteorder == 'little':
+-  REQUIRED_PACKAGES.append('grpcio ~= 1.32.0')
+  REQUIRED_PACKAGES.append('grpcio >= 1.31.0')
+ 
+ 
+ # Packages which are only needed for testing code.
--- a/pkgs/development/python-modules/tensorflow/system-jsoncpp.patch
+++ b/pkgs/development/python-modules/tensorflow/system-jsoncpp.patch
@ -1,21 +0,0 @@
-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
-     """,
- )
--- a/pkgs/development/python-modules/tensorflow/workspace.patch
+++ b/pkgs/development/python-modules/tensorflow/workspace.patch
@ -0,0 +1,18 @@
+diff --git a/WORKSPACE b/WORKSPACE
+index 9db1d9b80eb..c46f13f4ca4 100644
+--- a/WORKSPACE
+++ b/WORKSPACE
+@@ -12,6 +12,13 @@ http_archive(
+     ],
+ )
+ 
+http_archive(
+    name = "io_bazel_rules_docker",
+    sha256 = "1698624e878b0607052ae6131aa216d45ebb63871ec497f26c67455b34119c80",
+    strip_prefix = "rules_docker-0.15.0",
+    urls = ["https://github.com/bazelbuild/rules_docker/releases/download/v0.15.0/rules_docker-v0.15.0.tar.gz"],
+)
+
+ # Load tf_repositories() before loading dependencies for other repository so
+ # that dependencies like com_google_protobuf won't be overridden.
+ load("//tensorflow:workspace.bzl", "tf_repositories")
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@ -2764,6 +2764,8 @@ in {

  guzzle_sphinx_theme = callPackage ../development/python-modules/guzzle_sphinx_theme { };

+  gviz-api = callPackage ../development/python-modules/gviz-api {};
+
  gwyddion = disabledIf isPy3k (toPythonModule (pkgs.gwyddion.override {
    pythonSupport = true;
    pythonPackages = self;
@ -7372,6 +7374,10 @@ in {

  tenacity = callPackage ../development/python-modules/tenacity { };

+  tensorboard-plugin-profile = callPackage ../development/python-modules/tensorboard-plugin-profile { };
+
+  tensorboard-plugin-wit = callPackage ../development/python-modules/tensorboard-plugin-wit {};
+
  tensorboardx = callPackage ../development/python-modules/tensorboardx { };

  tensorflow-bin_2 = callPackage ../development/python-modules/tensorflow/bin.nix {
@ -7388,9 +7394,10 @@ in {
    cudatoolkit = pkgs.cudatoolkit_11_0;
    cudnn = pkgs.cudnn_cudatoolkit_11_0;
    nccl = pkgs.nccl_cudatoolkit_11;
-    openssl = pkgs.openssl_1_1;
    inherit (pkgs.darwin.apple_sdk.frameworks) Foundation Security;
-    inherit (pkgs) flatbuffers;
+    flatbuffers-core = pkgs.flatbuffers;
+    flatbuffers-python = self.flatbuffers;
+    lmdb-core = pkgs.lmdb;
  };

  tensorflow-build = self.tensorflow-build_2;