Skip to content

update TensorFlow easyblock to put Bazel build files in build directory + avoid unnecessary runtime patching #2166

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

Merged
boegel merged 5 commits into from
Sep 24, 2020
Merged

update TensorFlow easyblock to put Bazel build files in build directory + avoid unnecessary runtime patching #2166

Changes from 4 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
bbfabf2
Unify usages of tmpdir into a function
Flamefire Sep 9, 2020
a506dd5
Make the Bazel temporary folders be a subdir of builddir
Flamefire Sep 9, 2020
756e4bf
Remove setting of XDG_CACHE_HOME (done in PythonPackage)
Flamefire Sep 9, 2020
5348b22
Remove some patches which are no longer required
Flamefire Sep 9, 2020
2ac6f7e
Update formatting
Flamefire Sep 16, 2020
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
73 changes: 43 additions & 30 deletions easybuild/easyblocks/t/tensorflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -402,6 +402,17 @@ def get_system_libs(self):

return system_libs, cpaths, libpaths

def setup_build_dirs(self):
"""Setup temporary build directories"""
# Tensorflow/Bazel needs a couple of directories where it stores build cache and artefacts
tmpdir = tempfile.mkdtemp(suffix='-bazel-tf', dir=self.builddir)
self.output_root_dir = os.path.join(tmpdir, 'output_root')
self.output_base_dir = os.path.join(tmpdir, 'output_base')
self.output_user_root_dir = os.path.join(tmpdir, 'output_user_root')
self.wrapper_dir = os.path.join(tmpdir, 'wrapper_bin')
# This (likely) needs to be a subdir of output_base
self.install_base_dir = os.path.join(self.output_base_dir, 'inst_base')

def configure_step(self):
"""Custom configuration procedure for TensorFlow."""

Expand All @@ -410,8 +421,6 @@ def configure_step(self):
raise EasyBuildError("Failed to determine installation prefix for binutils")
self.binutils_bin_path = os.path.join(binutils_root, 'bin')

tmpdir = tempfile.mkdtemp(suffix='-bazel-configure')

# filter out paths from CPATH and LIBRARY_PATH. This is needed since bazel will pull some dependencies that
# might conflict with dependencies on the system and/or installed with EB. For example: protobuf
path_filter = self.cfg['path_filter']
Expand All @@ -423,14 +432,14 @@ def configure_step(self):
filtered_path = os.pathsep.join([p for fil in path_filter for p in path if fil not in p])
env.setvar(var, filtered_path)

wrapper_dir = os.path.join(tmpdir, 'bin')
use_wrapper = False
self.setup_build_dirs()

use_wrapper = False
if self.toolchain.comp_family() == toolchain.INTELCOMP:
# put wrappers for Intel C/C++ compilers in place (required to make sure license server is found)
# cfr. https://github.com/bazelbuild/bazel/issues/663
for compiler in ('icc', 'icpc'):
self.write_wrapper(wrapper_dir, compiler, 'NOT-USED-WITH-ICC')
self.write_wrapper(self.wrapper_dir, compiler, 'NOT-USED-WITH-ICC')
use_wrapper = True

use_mpi = self.toolchain.options.get('usempi', False)
Expand All @@ -441,7 +450,7 @@ def configure_step(self):
# put wrappers for Intel MPI compiler wrappers in place
# (required to make sure license server and I_MPI_ROOT are found)
for compiler in (os.getenv('MPICC'), os.getenv('MPICXX')):
self.write_wrapper(wrapper_dir, compiler, os.getenv('I_MPI_ROOT'))
self.write_wrapper(self.wrapper_dir, compiler, os.getenv('I_MPI_ROOT'))
use_wrapper = True
# set correct value for MPI_HOME
mpi_home = os.path.join(impi_root, 'intel64')
Expand All @@ -452,7 +461,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([wrapper_dir, os.getenv('PATH')]))
env.setvar('PATH', os.pathsep.join([self.wrapper_dir, os.getenv('PATH')]))

self.prepare_python()
self.handle_jemalloc()
Expand Down Expand Up @@ -599,23 +608,22 @@ def configure_step(self):
for (key, val) in sorted(config_env_vars.items()):
env.setvar(key, val)

# patch configure.py (called by configure script) to avoid that Bazel abuses $HOME/.cache/bazel
regex_subs = [(r"(run_shell\(\['bazel')",
r"\1, '--output_base=%s', '--install_base=%s'" % (tmpdir, os.path.join(tmpdir, 'inst_base')))]
apply_regex_substitutions('configure.py', regex_subs)
# Does no longer apply (and might not be required at all) since 1.12.0
if LooseVersion(self.version) < LooseVersion('1.12.0'):
# patch configure.py (called by configure script) to avoid that Bazel abuses $HOME/.cache/bazel
regex_subs = [(r"(run_shell\(\['bazel')",
r"\1, '--output_base=%s', '--install_base=%s'" % (self.output_base_dir,
self.install_base_dir))]
apply_regex_substitutions('configure.py', regex_subs)

# Tell Bazel to not use $HOME/.cache/bazel at all
# See https://docs.bazel.build/versions/master/output_directories.html
env.setvar('TEST_TMPDIR', os.path.join(tmpdir, 'output_root'))
env.setvar('TEST_TMPDIR', self.output_root_dir)
cmd = self.cfg['preconfigopts'] + './configure ' + self.cfg['configopts']
run_cmd(cmd, log_all=True, simple=True)

def build_step(self):
"""Custom build procedure for TensorFlow."""

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

def patch_crosstool_files(self):
"""Patches the CROSSTOOL files to include EasyBuild provided compiler paths"""
inc_paths, lib_paths = [], []

gcc_root = get_software_root('GCCcore') or get_software_root('GCC')
Expand Down Expand Up @@ -681,13 +689,24 @@ def build_step(self):
self.log.info("Patching %s", full_path)
apply_regex_substitutions(full_path, regex_subs)

tmpdir = tempfile.mkdtemp(suffix='-bazel-build')
user_root_tmpdir = tempfile.mkdtemp(suffix='-user_root')
def build_step(self):
"""Custom build procedure for TensorFlow."""

# 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'):
self.patch_crosstool_files()

# compose "bazel build" command with all its options...
cmd = [self.cfg['prebuildopts'], 'bazel', '--output_base=%s' % tmpdir,
'--install_base=%s' % os.path.join(tmpdir, 'inst_base'),
'--output_user_root=%s' % user_root_tmpdir, 'build']
cmd = [self.cfg['prebuildopts'],
'bazel',
'--output_base=%s' % self.output_base_dir,
'--install_base=%s' % self.install_base_dir,
'--output_user_root=%s' % self.output_user_root_dir,
'build'
]

# build with optimization enabled
# cfr. https://docs.bazel.build/versions/master/user-manual.html#flag--compilation_mode
Expand Down Expand Up @@ -749,7 +768,7 @@ def build_step(self):

# TF 2 (final) sets this in configure
if LooseVersion(self.version) < LooseVersion('2.0'):
if cuda_root:
if get_software_root('CUDA'):
cmd.append('--config=cuda')

# if mkl-dnn is listed as a dependency it is used. Otherwise downloaded if with_mkl_dnn is true
Expand Down Expand Up @@ -778,12 +797,6 @@ def test_step(self):

def install_step(self):
"""Custom install procedure for TensorFlow."""

# avoid that pip (ab)uses $HOME/.cache/pip
# cfr. https://pip.pypa.io/en/stable/reference/pip_install/#caching
env.setvar('XDG_CACHE_HOME', tempfile.gettempdir())
self.log.info("Using %s as pip cache directory", os.environ['XDG_CACHE_HOME'])

# find .whl file that was built, and install it using 'pip install'
if ("-rc" in self.version):
whl_version = self.version.replace("-rc", "rc")
Expand Down