use concurrent.futures.wait in EasyBlock.skip_extensions_parallel

Author: boegel

easybuild/framework/easyblock.py

@@ -41,7 +41,6 @@
 * Davide Vanzo (Vanderbilt University)
 * Caspar van Leeuwen (SURF)
 """
-
 import copy
 import glob
 import inspect
@@ -52,7 +51,7 @@
 import tempfile
 import time
 import traceback
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
 from datetime import datetime
 
 import easybuild.tools.environment as env
@@ -1821,42 +1820,41 @@ def skip_extensions_parallel(self, exts_filter):
 
         thread_pool = ThreadPoolExecutor(max_workers=self.cfg['parallel'])
 
-        async_shell_cmd_tasks = {}
-        running_checks_ids = []
+        async_shell_cmd_tasks = []
         installed_exts_ids = []
         exts_queue = list(enumerate(self.ext_instances[:]))
         checked_exts_cnt = 0
         exts_cnt = len(self.ext_instances)
+        done_tasks = []
 
         # asynchronously run checks to see whether extensions are already installed
-        while exts_queue or running_checks_ids:
+        while exts_queue or async_shell_cmd_tasks:
 
             # first handle completed checks
-            for idx in running_checks_ids[:]:
+            for task in done_tasks:
+                async_shell_cmd_tasks.remove(task)
+                res = task.result()
+                idx = res.task_id
                 ext_name = self.ext_instances[idx].name
-                # check whether command completed
-                task = async_shell_cmd_tasks[idx]
-                if task.done():
-                    res = task.result()
-                    self.log.info(f"exts_filter result for {ext_name}: exit code {res.exit_code}; output: {res.output}")
-                    running_checks_ids.remove(idx)
-                    if res.exit_code == 0:
-                        print_msg(f"skipping extension {ext_name}", log=self.log)
-                        installed_exts_ids.append(idx)
-
-                    checked_exts_cnt += 1
-                    exts_pbar_label = "skipping installed extensions "
-                    exts_pbar_label += "(%d/%d checked)" % (checked_exts_cnt, exts_cnt)
-                    self.update_exts_progress_bar(exts_pbar_label)
+                self.log.info(f"exts_filter result for {ext_name}: exit code {res.exit_code}; output: {res.output}")
+                if res.exit_code == 0:
+                    print_msg(f"skipping extension {ext_name}", log=self.log)
+                    installed_exts_ids.append(idx)
+
+                checked_exts_cnt += 1
+                exts_pbar_label = "skipping installed extensions "
+                exts_pbar_label += "(%d/%d checked)" % (checked_exts_cnt, exts_cnt)
+                self.update_exts_progress_bar(exts_pbar_label)
 
             # start additional checks asynchronously
             while exts_queue:
                 idx, ext = exts_queue.pop(0)
                 cmd, stdin = resolve_exts_filter_template(exts_filter, ext)
                 task = thread_pool.submit(run_shell_cmd, cmd, stdin=stdin, hidden=True,
-                                          fail_on_error=False, asynchronous=True)
-                async_shell_cmd_tasks[idx] = task
-                running_checks_ids.append(idx)
+                                          fail_on_error=False, asynchronous=True, task_id=idx)
+                async_shell_cmd_tasks.append(task)
+
+            (done_tasks, _) = wait(async_shell_cmd_tasks, timeout=1, return_when=FIRST_COMPLETED)
 
         # compose new list of extensions, skip over the ones that are already installed;
         # note: original order in extensions list should be preserved!

easybuild/tools/run.py

@@ -80,7 +80,7 @@
 
 
 RunShellCmdResult = namedtuple('RunShellCmdResult', ('cmd', 'exit_code', 'output', 'stderr', 'work_dir',
-                                                     'out_file', 'err_file', 'thread_id'))
+                                                     'out_file', 'err_file', 'thread_id', 'task_id'))
 
 
 class RunShellCmdError(BaseException):
@@ -184,7 +184,7 @@ def cache_aware_func(cmd, *args, **kwargs):
 @run_shell_cmd_cache
 def run_shell_cmd(cmd, fail_on_error=True, split_stderr=False, stdin=None, env=None,
                   hidden=False, in_dry_run=False, verbose_dry_run=False, work_dir=None, use_bash=True,
-                  output_file=True, stream_output=None, asynchronous=False, with_hooks=True,
+                  output_file=True, stream_output=None, asynchronous=False, task_id=None, with_hooks=True,
                   qa_patterns=None, qa_wait_patterns=None):
     """
     Run specified (interactive) shell command, and capture output + exit code.
@@ -267,7 +267,7 @@ def to_cmd_str(cmd):
             dry_run_msg(msg, silent=silent)
 
         return RunShellCmdResult(cmd=cmd_str, exit_code=0, output='', stderr=None, work_dir=work_dir,
-                                 out_file=cmd_out_fp, err_file=cmd_err_fp, thread_id=thread_id)
+                                 out_file=cmd_out_fp, err_file=cmd_err_fp, thread_id=thread_id, task_id=task_id)
 
     start_time = datetime.now()
     if not hidden:
@@ -344,7 +344,7 @@ def to_cmd_str(cmd):
             raise EasyBuildError(f"Failed to dump command output to temporary file: {err}")
 
     res = RunShellCmdResult(cmd=cmd_str, exit_code=proc.returncode, output=output, stderr=stderr, work_dir=work_dir,
-                            out_file=cmd_out_fp, err_file=cmd_err_fp, thread_id=thread_id)
+                            out_file=cmd_out_fp, err_file=cmd_err_fp, thread_id=thread_id, task_id=task_id)
 
     # always log command output
     cmd_name = cmd_str.split(' ')[0]

test/framework/run.py

@@ -911,7 +911,7 @@ def test_run_shell_cmd_cache(self):
         with self.mocked_stdout_stderr():
             cached_res = RunShellCmdResult(cmd=cmd, output="123456", exit_code=123, stderr=None,
                                            work_dir='/test_ulimit', out_file='/tmp/foo.out', err_file=None,
-                                           thread_id=None)
+                                           thread_id=None, task_id=None)
             run_shell_cmd.update_cache({(cmd, None): cached_res})
             res = run_shell_cmd(cmd)
         self.assertEqual(res.cmd, cmd)
@@ -931,7 +931,7 @@ def test_run_shell_cmd_cache(self):
         with self.mocked_stdout_stderr():
             cached_res = RunShellCmdResult(cmd=cmd, output="bar", exit_code=123, stderr=None,
                                            work_dir='/test_cat', out_file='/tmp/cat.out', err_file=None,
-                                           thread_id=None)
+                                           thread_id=None, task_id=None)
             run_shell_cmd.update_cache({(cmd, 'foo'): cached_res})
             res = run_shell_cmd(cmd, stdin='foo')
         self.assertEqual(res.cmd, cmd)

test/framework/systemtools.py

@@ -341,7 +341,7 @@ def mocked_run_shell_cmd(cmd, **kwargs):
     }
     if cmd in known_cmds:
         return RunShellCmdResult(cmd=cmd, exit_code=0, output=known_cmds[cmd], stderr=None, work_dir=os.getcwd(),
-                                 out_file=None, err_file=None, thread_id=None)
+                                 out_file=None, err_file=None, thread_id=None, task_id=None)
     else:
         return run_shell_cmd(cmd, **kwargs)
 
@@ -774,7 +774,7 @@ def test_gcc_version_darwin(self):
         out = "Apple LLVM version 7.0.0 (clang-700.1.76)"
         cwd = os.getcwd()
         mocked_run_res = RunShellCmdResult(cmd="gcc --version", exit_code=0, output=out, stderr=None, work_dir=cwd,
-                                           out_file=None, err_file=None, thread_id=None)
+                                           out_file=None, err_file=None, thread_id=None, task_id=None)
         st.run_shell_cmd = lambda *args, **kwargs: mocked_run_res
         self.assertEqual(get_gcc_version(), None)