Skip to content

Use build dir for big files/folders while building TensorFlow #3371

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
branfosj merged 3 commits into from
Jul 3, 2024
Merged

Use build dir for big files/folders while building TensorFlow #3371

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
1a297ea
Use build dir for building TensorFlow pip package
Flamefire Jun 20, 2024
2399cbd
Already change the temp dir during build
Flamefire Jun 27, 2024
e53a3e2
Also set $TMPDIR for the tests
Flamefire Jun 27, 2024
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
27 changes: 21 additions & 6 deletions easybuild/easyblocks/t/tensorflow.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
import re
import stat
import tempfile
from contextlib import contextmanager
from itertools import chain

import easybuild.tools.environment as env
Expand Down Expand Up @@ -451,6 +452,18 @@ def setup_build_dirs(self):
self.wrapper_dir = os.path.join(parent_dir, 'wrapper_bin')
mkdir(self.wrapper_dir)

@contextmanager
def set_tmp_dir(self):
# TF uses the temporary folder, which becomes quite large (~2 GB) so use the build folder explicitely.
old_tmpdir = os.environ['TMPDIR']
tmpdir = os.path.join(self.builddir, 'tmpdir')
mkdir(tmpdir)
os.environ['TMPDIR'] = tmpdir
try:
yield tmpdir
finally:
os.environ['TMPDIR'] = old_tmpdir

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

Expand Down Expand Up @@ -940,11 +953,12 @@ def build_step(self):
+ ['//tensorflow/tools/pip_package:build_pip_package']
)

run_cmd(' '.join(cmd), log_all=True, simple=True, log_ok=True)
with self.set_tmp_dir():
run_cmd(' '.join(cmd), log_all=True, simple=True, log_ok=True)

# run generated 'build_pip_package' script to build the .whl
cmd = "bazel-bin/tensorflow/tools/pip_package/build_pip_package %s" % self.builddir
run_cmd(cmd, log_all=True, simple=True, log_ok=True)
# run generated 'build_pip_package' script to build the .whl
cmd = "bazel-bin/tensorflow/tools/pip_package/build_pip_package %s" % self.builddir
run_cmd(cmd, log_all=True, simple=True, log_ok=True)

def test_step(self):
"""Run TensorFlow unit tests"""
Expand Down Expand Up @@ -1067,7 +1081,8 @@ def test_step(self):
+ test_targets
)

stdouterr, ec = run_cmd(cmd, log_ok=False, simple=False)
with self.set_tmp_dir():
stdouterr, ec = run_cmd(cmd, log_ok=False, simple=False)
if ec:
fail_msg = 'Tests on %s (cmd: %s) failed with exit code %s and output:\n%s' % (
device, cmd, ec, stdouterr)
Expand Down Expand Up @@ -1104,7 +1119,7 @@ def test_step(self):
def install_step(self):
"""Custom install procedure for TensorFlow."""
# find .whl file that was built, and install it using 'pip install'
if ("-rc" in self.version):
if "-rc" in self.version:
whl_version = self.version.replace("-rc", "rc")
else:
whl_version = self.version
Expand Down