Merge pull request #2225 from Flamefire/tf2.4

Update TensorFlow easyblock for TensorFlow 2.4

Author: boegel

easybuild/easyblocks/t/tensorflow.py

@@ -148,11 +148,11 @@ def is_version_ok(version_range):
         ('NASM', '2.0.0:'): 'nasm',
         ('nsync', '2.0.0:'): 'nsync',
         ('PCRE', '2.0.0:'): 'pcre',
-        ('protobuf-python', '2.0.0:'): 'com_google_protobuf',
+        ('protobuf', '2.0.0:'): 'com_google_protobuf',
         ('pybind11', '2.2.0:'): 'pybind11',
         ('snappy', '2.0.0:'): 'snappy',
         ('SQLite', '2.0.0:'): 'org_sqlite',
-        ('SWIG', '2.0.0:'): 'swig',
+        ('SWIG', '2.0.0:2.4.0'): 'swig',
         ('zlib', '2.0.0:2.2.0'): 'zlib_archive',
         ('zlib', '2.2.0:'): 'zlib',
     }
@@ -174,14 +174,19 @@ def is_version_ok(version_range):
         ('astor', '2.0.0:'): 'astor_archive',
         ('astunparse', '2.2.0:'): 'astunparse_archive',
         ('cython', '2.0.0:'): 'cython',  # Part of Python EC
+        ('dill', '2.4.0:'): 'dill_archive',
         ('enum', '2.0.0:'): 'enum34_archive',  # Part of Python3
+        ('flatbuffers', '2.4.0:'): 'flatbuffers',
         ('functools', '2.0.0:'): 'functools32_archive',  # Part of Python3
         ('gast', '2.0.0:'): 'gast_archive',
+        ('google.protobuf', '2.0.0:'): 'com_google_protobuf',
         ('keras_applications', '2.0.0:2.2.0'): 'keras_applications_archive',
         ('opt_einsum', '2.0.0:'): 'opt_einsum_archive',
         ('pasta', '2.0.0:'): 'pasta',
         ('six', '2.0.0:'): 'six_archive',  # Part of Python EC
+        ('tblib', '2.4.0:'): 'tblib_archive',
         ('termcolor', '2.0.0:'): 'termcolor_archive',
+        ('typing_extensions', '2.4.0:'): 'typing_extensions_archive',
         ('wrapt', '2.0.0:'): 'wrapt',
     }
     dependency_mapping = dict((dep_name, tf_name)
@@ -348,6 +353,9 @@ def get_system_libs(self):
         Returns a tuple of lists: $TF_SYSTEM_LIBS names, include paths, library paths
         """
         dependency_mapping, python_mapping = get_system_libs_for_version(self.version)
+        # Some TF dependencies require both a (usually C++) dependency and a Python package
+        deps_with_python_pkg = set(tf_name for tf_name in dependency_mapping.values()
+                                   if tf_name in python_mapping.values())
 
         system_libs = []
         cpaths = []
@@ -358,15 +366,17 @@ def get_system_libs(self):
         dep_names = set(dep['name'] for dep in self.cfg.dependencies())
         for dep_name, tf_name in sorted(dependency_mapping.items(), key=lambda i: i[0].lower()):
             if dep_name in dep_names:
+                if tf_name in deps_with_python_pkg:
+                    pkg_name = next(cur_pkg_name for cur_pkg_name, cur_tf_name in python_mapping.items()
+                                    if cur_tf_name == tf_name)
+                    # Simply ignore. Error reporting is done in the other loop
+                    if not self.python_pkg_exists(pkg_name):
+                        continue
                 system_libs.append(tf_name)
                 # When using cURL (which uses the system OpenSSL), we also need to use "boringssl"
                 # which essentially resolves to using OpenSSL as the API and library names are compatible
                 if dep_name == 'cURL':
                     system_libs.append('boringssl')
-                # For protobuf we need protobuf and protobuf-python where the latter depends on the former
-                # For includes etc. we need to get the values from protobuf
-                if dep_name == 'protobuf-python':
-                    dep_name = 'protobuf'
                 sw_root = get_software_root(dep_name)
                 # Dependency might be filtered via --filter-deps. In that case assume globally installed version
                 if not sw_root:
@@ -379,9 +389,12 @@ def get_system_libs(self):
                         # https://github.com/tensorflow/tensorflow/issues/42303
                         cpaths.append(sw_root)
                     if dep_name == 'protobuf':
-                        # Need to set INCLUDEDIR as TF wants to symlink headers from there:
-                        # https://github.com/tensorflow/tensorflow/issues/37835
-                        env.setvar('INCLUDEDIR', incpath)
+                        if LooseVersion(self.version) < LooseVersion('2.4'):
+                            # Need to set INCLUDEDIR as TF wants to symlink files from there:
+                            # https://github.com/tensorflow/tensorflow/issues/37835
+                            env.setvar('INCLUDEDIR', incpath)
+                        else:
+                            env.setvar('PROTOBUF_INCLUDE_PATH', incpath)
                 libpath = get_software_libdir(dep_name)
                 if libpath:
                     libpaths.append(os.path.join(sw_root, libpath))
@@ -390,7 +403,9 @@ def get_system_libs(self):
 
         for pkg_name, tf_name in sorted(python_mapping.items(), key=lambda i: i[0].lower()):
             if self.python_pkg_exists(pkg_name):
-                system_libs.append(tf_name)
+                # If it is in deps_with_python_pkg we already added it
+                if tf_name not in deps_with_python_pkg:
+                    system_libs.append(tf_name)
             else:
                 ignored_system_deps.append('%s (Python package %s)' % (tf_name, pkg_name))