use run_shell_cmd to install extensions in parallel

Author: boegel

easybuild/framework/easyblock.py

@@ -87,7 +87,7 @@
 from easybuild.tools.hooks import MODULE_STEP, MODULE_WRITE, PACKAGE_STEP, PATCH_STEP, PERMISSIONS_STEP, POSTITER_STEP
 from easybuild.tools.hooks import POSTPROC_STEP, PREPARE_STEP, READY_STEP, SANITYCHECK_STEP, SOURCE_STEP
 from easybuild.tools.hooks import SINGLE_EXTENSION, TEST_STEP, TESTCASES_STEP, load_hooks, run_hook
-from easybuild.tools.run import RunShellCmdError, run_shell_cmd
+from easybuild.tools.run import RunShellCmdError, raise_run_shell_cmd_error, run_shell_cmd
 from easybuild.tools.jenkins import write_to_xml
 from easybuild.tools.module_generator import ModuleGeneratorLua, ModuleGeneratorTcl, module_generator, dependencies_for
 from easybuild.tools.module_naming_scheme.utilities import det_full_ec_version
@@ -1961,6 +1961,8 @@ def install_extensions_parallel(self, install=True):
         """
         self.log.info("Installing extensions in parallel...")
 
+        thread_pool = ThreadPoolExecutor(max_workers=self.cfg['parallel'])
+
         running_exts = []
         installed_ext_names = []
 
@@ -1997,16 +1999,23 @@ def update_exts_progress_bar_helper(running_exts, progress_size):
 
             # check for extension installations that have completed
             if running_exts:
-                self.log.info("Checking for completed extension installations (%d running)...", len(running_exts))
+                self.log.info(f"Checking for completed extension installations ({len(running_exts)} running)...")
                 for ext in running_exts[:]:
-                    if self.dry_run or ext.async_cmd_check():
-                        self.log.info("Installation of %s completed!", ext.name)
-                        ext.postrun()
-                        running_exts.remove(ext)
-                        installed_ext_names.append(ext.name)
-                        update_exts_progress_bar_helper(running_exts, 1)
+                    if self.dry_run or ext.async_cmd_task.done():
+                        res = ext.async_cmd_task.result()
+                        if res.exit_code == 0:
+                            self.log.info(f"Installation of extension {ext.name} completed!")
+                            # run post-install method for extension from same working dir as installation of extension
+                            cwd = change_dir(res.work_dir)
+                            ext.postrun()
+                            change_dir(cwd)
+                            running_exts.remove(ext)
+                            installed_ext_names.append(ext.name)
+                            update_exts_progress_bar_helper(running_exts, 1)
+                        else:
+                            raise_run_shell_cmd_error(res)
                     else:
-                        self.log.debug("Installation of %s is still running...", ext.name)
+                        self.log.debug(f"Installation of extension {ext.name} is still running...")
 
             # try to start as many extension installations as we can, taking into account number of available cores,
             # but only consider first 100 extensions still in the queue
@@ -2073,9 +2082,9 @@ def update_exts_progress_bar_helper(running_exts, progress_size):
                                           rpath_filter_dirs=self.rpath_filter_dirs)
                     if install:
                         ext.prerun()
-                        ext.run_async()
+                        ext.async_cmd_task = ext.run_async(thread_pool)
                         running_exts.append(ext)
-                        self.log.info("Started installation of extension %s in the background...", ext.name)
+                        self.log.info(f"Started installation of extension {ext.name} in the background...")
                         update_exts_progress_bar_helper(running_exts, 0)
 
             # print progress info after every iteration (unless that info is already shown via progress bar)
@@ -2088,6 +2097,8 @@ def update_exts_progress_bar_helper(running_exts, progress_size):
                     running_ext_names = ', '.join(x.name for x in running_exts[:3]) + ", ..."
                 print_msg(msg % (installed_cnt, exts_cnt, queued_cnt, running_cnt, running_ext_names), log=self.log)
 
+        thread_pool.shutdown()
+
     #
     # MISCELLANEOUS UTILITY FUNCTIONS
     #

easybuild/framework/extension.py

@@ -42,7 +42,7 @@
 from easybuild.framework.easyconfig.templates import TEMPLATE_NAMES_EASYBLOCK_RUN_STEP, template_constant_dict
 from easybuild.tools.build_log import EasyBuildError, raise_nosupport
 from easybuild.tools.filetools import change_dir
-from easybuild.tools.run import check_async_cmd, run_cmd, run_shell_cmd
+from easybuild.tools.run import run_shell_cmd
 
 
 def resolve_exts_filter_template(exts_filter, ext):
@@ -150,12 +150,7 @@ def __init__(self, mself, ext, extra_params=None):
         self.sanity_check_module_loaded = False
         self.fake_mod_data = None
 
-        self.async_cmd_info = None
-        self.async_cmd_output = None
-        self.async_cmd_check_cnt = None
-        # initial read size should be relatively small,
-        # to avoid hanging for a long time until desired output is available in async_cmd_check
-        self.async_cmd_read_size = 1024
+        self.async_cmd_task = None
 
     @property
     def name(self):
@@ -195,44 +190,6 @@ def postrun(self):
         """
         self.master.run_post_install_commands(commands=self.cfg.get('postinstallcmds', []))
 
-    def async_cmd_start(self, cmd, inp=None):
-        """
-        Start installation asynchronously using specified command.
-        """
-        self.async_cmd_output = ''
-        self.async_cmd_check_cnt = 0
-        self.async_cmd_info = run_cmd(cmd, log_all=True, simple=False, inp=inp, regexp=False, asynchronous=True)
-
-    def async_cmd_check(self):
-        """
-        Check progress of installation command that was started asynchronously.
-
-        :return: True if command completed, False otherwise
-        """
-        if self.async_cmd_info is None:
-            raise EasyBuildError("No installation command running asynchronously for %s", self.name)
-        elif self.async_cmd_info is False:
-            self.log.info("No asynchronous command was started for extension %s", self.name)
-            return True
-        else:
-            self.log.debug("Checking on installation of extension %s...", self.name)
-            # use small read size, to avoid waiting for a long time until sufficient output is produced
-            res = check_async_cmd(*self.async_cmd_info, output_read_size=self.async_cmd_read_size)
-            self.async_cmd_output += res['output']
-            if res['done']:
-                self.log.info("Installation of extension %s completed!", self.name)
-                self.async_cmd_info = None
-            else:
-                self.async_cmd_check_cnt += 1
-                self.log.debug("Installation of extension %s still running (checked %d times)",
-                               self.name, self.async_cmd_check_cnt)
-                # increase read size after sufficient checks,
-                # to avoid that installation hangs due to output buffer filling up...
-                if self.async_cmd_check_cnt % 10 == 0 and self.async_cmd_read_size < (1024 ** 2):
-                    self.async_cmd_read_size *= 2
-
-            return res['done']
-
     @property
     def required_deps(self):
         """Return list of required dependencies for this extension."""
@@ -273,7 +230,6 @@ def sanity_check_step(self):
             self.log.info("modulename set to False for '%s' extension, so skipping sanity check", self.name)
         elif exts_filter:
             cmd, stdin = resolve_exts_filter_template(exts_filter, self)
-            # set log_ok to False so we can catch the error instead of run_cmd
             cmd_res = run_shell_cmd(cmd, fail_on_error=False, stdin=stdin)
 
             if cmd_res.exit_code:

easybuild/tools/run.py

@@ -207,6 +207,7 @@ def run_shell_cmd(cmd, fail_on_error=True, split_stderr=False, stdin=None, env=N
     :param output_file: collect command output in temporary output file
     :param stream_output: stream command output to stdout (auto-enabled with --logtostdout if None)
     :param asynchronous: indicate that command is being run asynchronously
+    :param task_id: task ID for specified shell command (included in return value)
     :param with_hooks: trigger pre/post run_shell_cmd hooks (if defined)
     :param qa_patterns: list of 2-tuples with patterns for questions + corresponding answers
     :param qa_wait_patterns: list of 2-tuples with patterns for non-questions

test/framework/sandbox/easybuild/easyblocks/generic/toy_extension.py

@@ -27,6 +27,7 @@
 
 @author: Kenneth Hoste (Ghent University)
 """
+import os
 
 from easybuild.framework.easyconfig import CUSTOM
 from easybuild.framework.extensioneasyblock import ExtensionEasyBlock
@@ -81,22 +82,24 @@ def prerun(self):
             super(Toy_Extension, self).run(unpack_src=True)
             EB_toy.configure_step(self.master, name=self.name, cfg=self.cfg)
 
-    def run_async(self):
+    def run_async(self, thread_pool):
         """
         Install toy extension asynchronously.
         """
+        task_id = f'ext_{self.name}_{self.version}'
         if self.src:
             cmd = compose_toy_build_cmd(self.cfg, self.name, self.cfg['prebuildopts'], self.cfg['buildopts'])
-            self.async_cmd_start(cmd)
         else:
-            self.async_cmd_info = False
+            cmd = f"echo 'no sources for {self.name}'"
+
+        return thread_pool.submit(run_shell_cmd, cmd, asynchronous=True, env=os.environ.copy(),
+                                  fail_on_error=False, task_id=task_id)
 
     def postrun(self):
         """
         Wrap up installation of toy extension.
         """
         super(Toy_Extension, self).postrun()
-
         EB_toy.install_step(self.master, name=self.name)
 
     def sanity_check_step(self, *args, **kwargs):

test/framework/sandbox/easybuild/easyblocks/t/toy.py

@@ -163,12 +163,14 @@ def run(self):
         """
         self.build_step()
 
-    def run_async(self):
+    def run_async(self, thread_pool):
         """
         Asynchronous installation of toy as extension.
         """
         cmd = compose_toy_build_cmd(self.cfg, self.name, self.cfg['prebuildopts'], self.cfg['buildopts'])
-        self.async_cmd_start(cmd)
+        task_id = f'ext_{self.name}_{self.version}'
+        return thread_pool.submit(run_shell_cmd, cmd, asynchronous=True, env=os.environ.copy(),
+                                  fail_on_error=False, task_id=task_id)
 
     def postrun(self):
         """