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 6 commits into
base: develop
Choose a base branch
from
+35 −25
Open

Update TensorFlow easyblock for TensorFlow 2.19+ #4205

Changes from all commits
Commits
Show all changes
6 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
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
60 changes: 35 additions & 25 deletions easybuild/easyblocks/t/tensorflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,9 +124,9 @@ def is_version_ok(version_range):
"""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 @@ -140,7 +140,7 @@ def is_version_ok(version_range):
# <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,7 +153,7 @@ 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',
Expand Down Expand Up @@ -516,7 +516,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 +546,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 +851,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 +866,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 +898,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 +918,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 +953,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 +969,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 +1015,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 +1042,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 +1250,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 +1266,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