Skip to content

Update TensorFlow easyblock for TensorFlow 2.19+ #4205

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Flamefire wants to merge 8 commits into
base: develop
Choose a base branch
from
+73 −53
Open

Update TensorFlow easyblock for TensorFlow 2.19+ #4205

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
e56904c
Update TensorFlow easyblock for TensorFlow 2.19+
Flamefire Aug 13, 2026
fe940b4
Add flag to use our compilers in TF 2.21+
Flamefire Aug 13, 2026
e1eedf7
Avoid repeated construction of LooseVersion
Flamefire Aug 13, 2026
7a28fbe
Set HOME for tests
Flamefire Aug 13, 2026
85f5fd8
Also set test timeout
Flamefire Aug 13, 2026
b49d70e
Don't require TensorBoard for 2.21+
Flamefire Aug 13, 2026
cc5f421
Add alternative EB dependency names
Flamefire Aug 26, 2026
6d79598
Use
Flamefire Aug 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
126 changes: 73 additions & 53 deletions easybuild/easyblocks/t/tensorflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@
import tempfile
from contextlib import contextmanager
from itertools import chain
from typing import Dict, List, Set, Tuple, Union

import easybuild.tools.environment as env
import easybuild.tools.toolchain as toolchain
Expand All @@ -52,6 +53,7 @@
from easybuild.tools.run import run_shell_cmd
from easybuild.tools.systemtools import AARCH64, X86_64, get_cpu_architecture, get_os_name, get_os_version
from easybuild.tools.toolchain.toolchain import RPATH_WRAPPERS_SUBDIR
from easybuild.tools.utilities import nub


CPU_DEVICE = 'cpu'
Expand Down Expand Up @@ -82,7 +84,7 @@
KNOWN_BINUTILS = ('ar', 'as', 'dwp', 'ld', 'ld.bfd', 'ld.gold', 'nm', 'objcopy', 'objdump', 'strip')


def split_tf_libs_txt(valid_libs_txt):
def split_tf_libs_txt(valid_libs_txt: str) -> List[str]:
"""Split the VALID_LIBS entry from the TF file into single names"""
entries = valid_libs_txt.split(',')
# Remove double quotes and whitespace
Expand All @@ -93,7 +95,7 @@ def split_tf_libs_txt(valid_libs_txt):
return result


def get_system_libs_from_tf(source_dir):
def get_system_libs_from_tf(source_dir: str) -> List[str]:
"""Return the valid values for TF_SYSTEM_LIBS from the TensorFlow source directory"""
syslibs_path = os.path.join(source_dir, 'third_party', 'systemlibs', 'syslibs_configure.bzl')
result = []
Expand All @@ -106,7 +108,9 @@ def get_system_libs_from_tf(source_dir):
return result


def get_system_libs_for_version(tf_version, as_valid_libs=False):
def get_system_libs_for_version(tf_version, as_valid_libs=False) -> Union[str,
Tuple[Dict[Tuple[str, ...], str],
Dict[str, str]]]:
"""
Determine valid values for $TF_SYSTEM_LIBS for the given TF version

Expand All @@ -120,13 +124,13 @@ def get_system_libs_for_version(tf_version, as_valid_libs=False):
"""
tf_version = LooseVersion(tf_version)

def is_version_ok(version_range):
def is_version_ok(version_range: str) -> bool:
"""Return True if the TF version to be installed matches the version_range"""
min_version, max_version = version_range.split(':')
result = True
if min_version and tf_version < LooseVersion(min_version):
if min_version and tf_version < min_version:
result = False
if max_version and tf_version >= LooseVersion(max_version):
if max_version and tf_version >= max_version:
result = False
return result

Expand All @@ -135,12 +139,13 @@ def is_version_ok(version_range):
# if it does something "strange" (e.g. link hardcoded headers)

# Software which is added as a dependency in the EC
available_system_libs = {
available_system_libs: Dict[Tuple[Union[str, Tuple[str, ...], str]], str] = {
# Format: (<EB name>, <version range>): <TF name>
# <EB name> can be a tuple of alternative names
# <version range> is '<min version>:<exclusive max version>'
('Abseil', '2.9.0:'): 'com_google_absl',
('cURL', '2.0.0:'): 'curl',
('double-conversion', '2.0.0:'): 'double_conversion',
('double-conversion', '2.0.0:2.19.0'): 'double_conversion',
('flatbuffers', '2.0.0:'): 'flatbuffers',
('giflib', '2.0.0:2.1.0'): 'gif_archive',
('giflib', '2.1.0:'): 'gif',
Expand All @@ -153,19 +158,19 @@ def is_version_ok(version_range):
('libpng', '2.1.0:'): 'png',
('LMDB', '2.0.0:2.13.0'): 'lmdb',
('NASM', '2.0.0:'): 'nasm',
('nsync', '2.0.0:'): 'nsync',
('nsync', '2.0.0:2.19.0'): 'nsync',
('PCRE', '2.0.0:2.6.0'): 'pcre',
('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:2.4.0'): 'swig',
('zlib', '2.0.0:2.2.0'): 'zlib_archive',
('zlib', '2.2.0:'): 'zlib',
(('zlib', 'zlib-ng'), '2.0.0:2.2.0'): 'zlib_archive',
(('zlib', 'zlib-ng'), '2.2.0:'): 'zlib',
}
# Software recognized by TF but which is always disabled (usually because no EC is known)
# Format: <TF name>: <version range>
unused_system_libs = {
unused_system_libs: Dict[str, str] = {
'boringssl': '2.0.0:', # Implied by cURL and existence of OpenSSL anywhere in the dependency chain
'com_github_googleapis_googleapis': '2.0.0:2.5.0',
'com_github_googlecloudplatform_google_cloud_cpp': '2.0.0:', # Not used due to $TF_NEED_GCP=0
Expand All @@ -176,7 +181,7 @@ def is_version_ok(version_range):
# Python packages installed as extensions or in the Python module
# Will be checked for availabilitly
# Format: (<package name>, <version range>): <TF name>
python_system_libs = {
python_system_libs: Dict[Tuple[str, str], str] = {
('absl', '2.0.0:'): 'absl_py',
('astor', '2.0.0:'): 'astor_archive',
('astunparse', '2.2.0:'): 'astunparse_archive',
Expand All @@ -197,16 +202,21 @@ def is_version_ok(version_range):
('wrapt', '2.0.0:'): 'wrapt',
}

dependency_mapping = {dep_name: tf_name
for (dep_name, version_range), tf_name in available_system_libs.items()
if is_version_ok(version_range)}
python_mapping = {pkg_name: tf_name
for (pkg_name, version_range), tf_name in python_system_libs.items()
if is_version_ok(version_range)}
def maybe_to_tuple(s: Union[str, Tuple[str]]) -> Tuple[str, ...]:
return (s,) if isinstance(s, str) else s

dependency_mapping: Dict[Tuple[str, ...], str] = {
maybe_to_tuple(dep_name): tf_name
for (dep_name, version_range), tf_name in available_system_libs.items()
if is_version_ok(version_range)}
python_mapping: Dict[str, str] = {
pkg_name: tf_name
for (pkg_name, version_range), tf_name in python_system_libs.items()
if is_version_ok(version_range)}

if as_valid_libs:
tf_names = [tf_name for tf_name, version_range in unused_system_libs.items()
if is_version_ok(version_range)]
tf_names: List[str] = [tf_name for tf_name, version_range in unused_system_libs.items()
if is_version_ok(version_range)]
tf_names.extend(dependency_mapping.values())
tf_names.extend(python_mapping.values())
result = '\n'.join([' "%s",' % name for name in sorted(tf_names)])
Expand Down Expand Up @@ -365,9 +375,10 @@ def get_system_libs(self):
ignored_system_deps = []

# Check direct dependencies
dep_names = self.cfg.dependency_names()
for dep_name, tf_name in sorted(dependency_mapping.items(), key=lambda i: i[0].lower()):
if dep_name in dep_names:
dep_names: Set[str] = self.cfg.dependency_names()
for eb_dep_names, tf_name in sorted(dependency_mapping.items(), key=lambda i: [name.lower() for name in i[0]]):
dep_name = next((eb_dep_name for eb_dep_name in eb_dep_names if eb_dep_name in dep_names), None)
if dep_name:
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)
Expand Down Expand Up @@ -397,11 +408,11 @@ def get_system_libs(self):
env.setvar('INCLUDEDIR', incpath)
else:
env.setvar('PROTOBUF_INCLUDE_PATH', incpath)
libpath = get_software_libdir(dep_name)
libpath = get_software_libdir(dep_name, full_path=True)
if libpath:
libpaths.append(os.path.join(sw_root, libpath))
libpaths.append(libpath)
else:
ignored_system_deps.append('%s (Dependency %s)' % (tf_name, dep_name))
ignored_system_deps.append('%s (Dependency %s)' % (tf_name, ' or '.join(eb_dep_names)))

for pkg_name, tf_name in sorted(python_mapping.items(), key=lambda i: i[0].lower()):
if self.python_pkg_exists(pkg_name):
Expand All @@ -414,14 +425,13 @@ def get_system_libs(self):
# If we use OpenSSL (potentially as a wrapper) somewhere in the chain we must tell TF to use it too
openssl_root = get_software_root('OpenSSL')
if openssl_root:
if 'boringssl' not in system_libs:
system_libs.append('boringssl')
system_libs.append('boringssl')
incpath = os.path.join(openssl_root, 'include')
if os.path.exists(incpath):
cpaths.append(incpath)
libpath = get_software_libdir('OpenSSL')
libpath = get_software_libdir('OpenSSL', full_path=True)
if libpath:
libpaths.append(os.path.join(openssl_root, libpath))
libpaths.append(libpath)

if ignored_system_deps:
print_warning('%d TensorFlow dependencies have not been resolved by EasyBuild. '
Expand All @@ -436,7 +446,7 @@ def get_system_libs(self):
else:
self.log.info("All known TensorFlow $TF_SYSTEM_LIBS dependencies resolved via EasyBuild!")

return system_libs, cpaths, libpaths
return nub(system_libs), cpaths, libpaths

def setup_build_dirs(self):
"""Setup temporary build directories"""
Expand Down Expand Up @@ -516,7 +526,7 @@ def configure_step(self):
for var in ['CPATH', 'LIBRARY_PATH']:
path = os.getenv(var).split(os.pathsep)
self.log.info("$%s old value was %s" % (var, path))
filtered_path = os.pathsep.join([p for fil in path_filter for p in path if fil not in p])
filtered_path = env.join_path_var(p for fil in path_filter for p in path if fil not in p)
env.setvar(var, filtered_path)

use_wrapper = False
Expand Down Expand Up @@ -546,7 +556,7 @@ def configure_step(self):
self.log.debug("Derived value for MPI_HOME: %s", mpi_home)

if use_wrapper:
env.setvar('PATH', os.pathsep.join([self.wrapper_dir, os.getenv('PATH')]))
env.setvar('PATH', env.join_path_var([self.wrapper_dir, os.getenv('PATH')]))

self.prepare_python()

Expand Down Expand Up @@ -851,7 +861,7 @@ def patch_crosstool_files(self):
regex_subs.append((os.path.join('/usr', 'bin', tool), path))

# -fPIE/-pie and -fPIC are not compatible, so patch out hardcoded occurences of -fPIE/-pie if -fPIC is used
if self.toolchain.options.get('pic', None):
if self.toolchain.options.get('pic'):
regex_subs.extend([('-fPIE', '-fPIC'), ('"-pie"', '"-fPIC"')])

# patch all CROSSTOOL* scripts to fix hardcoding of locations of binutils/GCC binaries
Expand All @@ -866,12 +876,13 @@ def build_step(self):
"""Custom build procedure for TensorFlow."""

bazel_version = get_bazel_version()
tf_version = LooseVersion(self.version)

# pre-create target installation directory
mkdir(os.path.join(self.installdir, self.pylibdir), parents=True)

# This seems to be no longer required since at least 2.0, likely also for older versions
if LooseVersion(self.version) < LooseVersion('2.0'):
if tf_version < '2.0':
self.patch_crosstool_files()

# Options passed to the bazel command
Expand All @@ -897,7 +908,7 @@ def build_step(self):
'--host_jvm_args=-Xmx%sm' % jvm_max_memory
])

if self.toolchain.options.get('debug', None):
if self.toolchain.options.get('debug'):
self.target_opts.append('--strip=never')
self.target_opts.append('--compilation_mode=dbg')
self.target_opts.append('--copt="-Og"')
Expand All @@ -917,22 +928,24 @@ def build_step(self):

self.target_opts.append(f'--jobs={self.cfg.parallel}')

if self.toolchain.options.get('pic', None):
if tf_version >= '2.21':
self.target_opts.append('--config=clang_local')

if self.toolchain.options.get('pic'):
self.target_opts.append('--copt="-fPIC"')

# include install location of Python packages in $PYTHONPATH,
# and specify that value of $PYTHONPATH should be passed down into Bazel build environment;
# this is required to make sure that Python packages included as extensions are found at build time;
# see also https://github.com/tensorflow/tensorflow/issues/22395
pythonpath = os.getenv('PYTHONPATH', '')
action_pythonpath = [os.path.join(self.installdir, self.pylibdir), pythonpath]
if LooseVersion(self.version) >= LooseVersion('2.14') and 'EBPYTHONPREFIXES' in os.environ:
action_pythonpath = [os.path.join(self.installdir, self.pylibdir), os.getenv('PYTHONPATH')]
if tf_version >= '2.14' and 'EBPYTHONPREFIXES' in os.environ:
# Since TF 2.14 the build uses hermetic python, which ignores sitecustomize.py from EB python;
# explicity include our site-packages here to respect EBPYTHONPREFIXERS, if that's prefered.
pyshortver = '.'.join(get_software_version('Python').split('.')[:2])
eb_pythonpath = os.path.join(os.getenv('EBROOTPYTHON'), 'lib', 'python' + pyshortver, 'site-packages')
action_pythonpath.append(eb_pythonpath)
env.setvar('PYTHONPATH', os.pathsep.join(action_pythonpath))
env.setvar('PYTHONPATH', env.join_path_var(action_pythonpath))

# Make TF find our modules. LD_LIBRARY_PATH gets automatically added by configure.py
cpaths, libpaths = self.system_libs_info[1:]
Expand All @@ -950,10 +963,10 @@ def build_step(self):
action_env.update(PY_ENV_VARS)

# TF 2 (final) sets this in configure
if (LooseVersion(self.version) < LooseVersion('2.0')) and self._with_cuda:
if (tf_version < '2.0') and self._with_cuda:
self.target_opts.append('--config=cuda')
# TF 2.18 with CUDA needs to set cuda_wheel to config
if (LooseVersion(self.version) >= LooseVersion('2.18')) and self._with_cuda:
if (tf_version >= '2.18') and self._with_cuda:
self.target_opts.append('--config=cuda_wheel')

# note: using --config=mkl results in a significantly different build, with a different
Expand All @@ -966,7 +979,7 @@ def build_step(self):
# auto-enable use of MKL-DNN/oneDNN and --config=mkl when possible if with_mkl_dnn is left unspecified;
# only do this for TensorFlow versions older than 2.4.0, since more recent versions
# oneDNN is used automatically for x86_64 systems (and mkl-dnn is no longer a dependency);
if self.cfg['with_mkl_dnn'] is None and LooseVersion(self.version) < LooseVersion('2.4.0'):
if self.cfg['with_mkl_dnn'] is None and tf_version < '2.4.0':
cpu_arch = get_cpu_architecture()
if cpu_arch == X86_64:
# Supported on x86 since forever
Expand Down Expand Up @@ -1012,16 +1025,16 @@ def build_step(self):
+ self.target_opts
+ [self.cfg['buildopts']]
)
if LooseVersion(self.version) < '2.16':
if tf_version < '2.16':
cmd += ['//tensorflow/tools/pip_package:build_pip_package']
elif LooseVersion(self.version) < '2.17': # for v2.16.x
elif tf_version < '2.17': # for v2.16.x
cmd += ['//tensorflow/tools/pip_package:v2/wheel']
else:
cmd += ['//tensorflow/tools/pip_package:wheel']

with self.set_tmp_dir():
run_shell_cmd(' '.join(cmd))
if LooseVersion(self.version) < LooseVersion('2.16'):
if tf_version < '2.16':
# run generated 'build_pip_package' script to build the .whl
cmd = "bazel-bin/tensorflow/tools/pip_package/build_pip_package %s" % self.builddir
run_shell_cmd(cmd)
Expand All @@ -1039,8 +1052,12 @@ def test_step(self):
test_targets = test_targets.split(' ')

test_opts = self.target_opts
test_opts.append('--test_output=errors') # (Additionally) show logs from failed tests
test_opts.append('--build_tests_only') # Don't build tests which won't be executed
test_opts.extend([
'--test_output=errors', # (Additionally) show logs from failed tests
'--build_tests_only', # Don't build tests which won't be executed
f"--test_env=HOME='{self.home_dir}'",
"--test_timeout=3600",
])

# determine number of cores/GPUs to use for tests
max_num_test_jobs = self.cfg['test_max_parallel'] or self.cfg.parallel
Expand Down Expand Up @@ -1243,9 +1260,12 @@ def sanity_check_step(self):
self.prepare_python()

custom_paths = {
'files': ['bin/tensorboard'],
'files': [],
'dirs': [self.pylibdir],
}
# Tensorboard was removed from TensorFlow in 2.21
if LooseVersion(self.version) < '2.21':
custom_paths['files'].append('bin/tensorboard')

custom_commands = [
"%s -c 'import tensorflow'" % self.python_cmd,
Expand All @@ -1256,8 +1276,8 @@ def sanity_check_step(self):

# test installation using MNIST tutorial examples
if self.cfg['runtest']:
pythonpath = os.getenv('PYTHONPATH', '')
env.setvar('PYTHONPATH', os.pathsep.join([os.path.join(self.installdir, self.pylibdir), pythonpath]))
env.setvar('PYTHONPATH',
env.join_path_var([os.path.join(self.installdir, self.pylibdir), os.getenv('PYTHONPATH')]))

mnist_pys = []

Expand Down