Merge pull request #3371 from Flamefire/tensorflow-tmp-dir

Use build dir for big files/folders while building TensorFlow

Author: branfosj

easybuild/easyblocks/t/tensorflow.py

@@ -35,6 +35,7 @@
 import re
 import stat
 import tempfile
+from contextlib import contextmanager
 from itertools import chain
 
 import easybuild.tools.environment as env
@@ -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."""
 
@@ -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"""
@@ -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)
@@ -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