Merge pull request #4323 from boegel/run_cmd_hook

add `run_shell_cmd` hook

Author: jfgrimm

easybuild/tools/hooks.py

@@ -69,6 +69,8 @@
 CANCEL = 'cancel'
 FAIL = 'fail'
 
+RUN_SHELL_CMD = 'run_shell_cmd'
+
 PRE_PREF = 'pre_'
 POST_PREF = 'post_'
 HOOK_SUFF = '_hook'
@@ -106,6 +108,8 @@
     END,
     CANCEL,
     FAIL,
+    PRE_PREF + RUN_SHELL_CMD,
+    POST_PREF + RUN_SHELL_CMD,
 ]
 KNOWN_HOOKS = [h + HOOK_SUFF for h in HOOK_NAMES]
 
@@ -203,7 +207,7 @@ def find_hook(label, hooks, pre_step_hook=False, post_step_hook=False):
     return res
 
 
-def run_hook(label, hooks, pre_step_hook=False, post_step_hook=False, args=None, msg=None):
+def run_hook(label, hooks, pre_step_hook=False, post_step_hook=False, args=None, kwargs=None, msg=None):
     """
     Run hook with specified label and return result of calling the hook or None.
 
@@ -219,6 +223,8 @@ def run_hook(label, hooks, pre_step_hook=False, post_step_hook=False, args=None,
     if hook:
         if args is None:
             args = []
+        if kwargs is None:
+            kwargs = {}
 
         if pre_step_hook:
             label = 'pre-' + label
@@ -230,6 +236,6 @@ def run_hook(label, hooks, pre_step_hook=False, post_step_hook=False, args=None,
         if build_option('debug'):
             print_msg(msg)
 
-        _log.info("Running '%s' hook function (arguments: %s)...", hook.__name__, args)
-        res = hook(*args)
+        _log.info("Running '%s' hook function (args: %s, keyword args: %s)...", hook.__name__, args, kwargs)
+        res = hook(*args, **kwargs)
     return res

easybuild/tools/options.py

@@ -1894,7 +1894,7 @@ def set_tmpdir(tmpdir=None, raise_error=False):
             os.close(fd)
             os.chmod(tmptest_file, 0o700)
             if not run_cmd(tmptest_file, simple=True, log_ok=False, regexp=False, force_in_dry_run=True, trace=False,
-                           stream_output=False):
+                           stream_output=False, with_hooks=False):
                 msg = "The temporary directory (%s) does not allow to execute files. " % tempfile.gettempdir()
                 msg += "This can cause problems in the build process, consider using --tmpdir."
                 if raise_error:

easybuild/tools/run.py

@@ -50,6 +50,7 @@
 from easybuild.base import fancylogger
 from easybuild.tools.build_log import EasyBuildError, dry_run_msg, print_msg, time_str_since
 from easybuild.tools.config import ERROR, IGNORE, WARN, build_option
+from easybuild.tools.hooks import RUN_SHELL_CMD, load_hooks, run_hook
 from easybuild.tools.py2vs3 import string_type
 from easybuild.tools.utilities import trace_msg
 
@@ -131,7 +132,8 @@ def get_output_from_process(proc, read_size=None, asynchronous=False):
 
 @run_cmd_cache
 def run_cmd(cmd, log_ok=True, log_all=False, simple=False, inp=None, regexp=True, log_output=False, path=None,
-            force_in_dry_run=False, verbose=True, shell=None, trace=True, stream_output=None, asynchronous=False):
+            force_in_dry_run=False, verbose=True, shell=None, trace=True, stream_output=None, asynchronous=False,
+            with_hooks=True):
     """
     Run specified command (in a subshell)
     :param cmd: command to run
@@ -148,6 +150,7 @@ def run_cmd(cmd, log_ok=True, log_all=False, simple=False, inp=None, regexp=True
     :param trace: print command being executed as part of trace output
     :param stream_output: enable streaming command output to stdout
     :param asynchronous: run command asynchronously (returns subprocess.Popen instance if set to True)
+    :param with_hooks: trigger pre/post run_shell_cmd hooks (if defined)
     """
     cwd = os.getcwd()
 
@@ -233,6 +236,13 @@ def run_cmd(cmd, log_ok=True, log_all=False, simple=False, inp=None, regexp=True
         else:
             raise EasyBuildError("Don't know how to prefix with /usr/bin/env for commands of type %s", type(cmd))
 
+    if with_hooks:
+        hooks = load_hooks(build_option('hooks'))
+        hook_res = run_hook(RUN_SHELL_CMD, hooks, pre_step_hook=True, args=[cmd], kwargs={'work_dir': os.getcwd()})
+        if isinstance(hook_res, string_type):
+            cmd, old_cmd = hook_res, cmd
+            _log.info("Command to run was changed by pre-%s hook: '%s' (was: '%s')", RUN_SHELL_CMD, cmd, old_cmd)
+
     _log.info('running cmd: %s ' % cmd)
     try:
         proc = subprocess.Popen(cmd, shell=shell, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
@@ -248,7 +258,7 @@ def run_cmd(cmd, log_ok=True, log_all=False, simple=False, inp=None, regexp=True
         return (proc, cmd, cwd, start_time, cmd_log)
     else:
         return complete_cmd(proc, cmd, cwd, start_time, cmd_log, log_ok=log_ok, log_all=log_all, simple=simple,
-                            regexp=regexp, stream_output=stream_output, trace=trace)
+                            regexp=regexp, stream_output=stream_output, trace=trace, with_hook=with_hooks)
 
 
 def check_async_cmd(proc, cmd, owd, start_time, cmd_log, fail_on_error=True, output_read_size=1024, output=''):
@@ -293,7 +303,7 @@ def check_async_cmd(proc, cmd, owd, start_time, cmd_log, fail_on_error=True, out
 
 
 def complete_cmd(proc, cmd, owd, start_time, cmd_log, log_ok=True, log_all=False, simple=False,
-                 regexp=True, stream_output=None, trace=True, output=''):
+                 regexp=True, stream_output=None, trace=True, output='', with_hook=True):
     """
     Complete running of command represented by passed subprocess.Popen instance.
 
@@ -308,6 +318,7 @@ def complete_cmd(proc, cmd, owd, start_time, cmd_log, log_ok=True, log_all=False
     :param regexp: regex used to check the output for errors;  if True it will use the default (see parse_log_for_error)
     :param stream_output: enable streaming command output to stdout
     :param trace: print command being executed as part of trace output
+    :param with_hook: trigger post run_shell_cmd hooks (if defined)
     """
     # use small read size when streaming output, to make it stream more fluently
     # read size should not be too small though, to avoid too much overhead
@@ -343,6 +354,15 @@ def complete_cmd(proc, cmd, owd, start_time, cmd_log, log_ok=True, log_all=False
         sys.stdout.write(output)
     stdouterr += output
 
+    if with_hook:
+        hooks = load_hooks(build_option('hooks'))
+        run_hook_kwargs = {
+            'exit_code': ec,
+            'output': stdouterr,
+            'work_dir': os.getcwd(),
+        }
+        run_hook(RUN_SHELL_CMD, hooks, post_step_hook=True, args=[cmd], kwargs=run_hook_kwargs)
+
     if trace:
         trace_msg("command completed: exit %s, ran in %s" % (ec, time_str_since(start_time)))
 
@@ -485,6 +505,17 @@ def check_answers_list(answers):
     # Part 2: Run the command and answer questions
     # - this needs asynchronous stdout
 
+    hooks = load_hooks(build_option('hooks'))
+    run_hook_kwargs = {
+        'interactive': True,
+        'work_dir': os.getcwd(),
+    }
+    hook_res = run_hook(RUN_SHELL_CMD, hooks, pre_step_hook=True, args=[cmd], kwargs=run_hook_kwargs)
+    if isinstance(hook_res, string_type):
+        cmd, old_cmd = hook_res, cmd
+        _log.info("Interactive command to run was changed by pre-%s hook: '%s' (was: '%s')",
+                  RUN_SHELL_CMD, cmd, old_cmd)
+
     # # Log command output
     if cmd_log:
         cmd_log.write("# output for interactive command: %s\n\n" % cmd)
@@ -599,6 +630,13 @@ def get_proc():
         except IOError as err:
             _log.debug("runqanda cmd %s: remaining data read failed: %s", cmd, err)
 
+    run_hook_kwargs.update({
+        'interactive': True,
+        'exit_code': ec,
+        'output': stdout_err,
+    })
+    run_hook(RUN_SHELL_CMD, hooks, post_step_hook=True, args=[cmd], kwargs=run_hook_kwargs)
+
     if trace:
         trace_msg("interactive command completed: exit %s, ran in %s" % (ec, time_str_since(start_time)))
 

easybuild/tools/systemtools.py

@@ -274,7 +274,7 @@ def get_avail_core_count():
         core_cnt = int(sum(sched_getaffinity()))
     else:
         # BSD-type systems
-        out, _ = run_cmd('sysctl -n hw.ncpu', force_in_dry_run=True, trace=False, stream_output=False)
+        out, _ = run_cmd('sysctl -n hw.ncpu', force_in_dry_run=True, trace=False, stream_output=False, with_hooks=False)
         try:
             if int(out) > 0:
                 core_cnt = int(out)
@@ -311,7 +311,7 @@ def get_total_memory():
     elif os_type == DARWIN:
         cmd = "sysctl -n hw.memsize"
         _log.debug("Trying to determine total memory size on Darwin via cmd '%s'", cmd)
-        out, ec = run_cmd(cmd, force_in_dry_run=True, trace=False, stream_output=False)
+        out, ec = run_cmd(cmd, force_in_dry_run=True, trace=False, stream_output=False, with_hooks=False)
         if ec == 0:
             memtotal = int(out.strip()) // (1024**2)
 
@@ -393,14 +393,15 @@ def get_cpu_vendor():
 
     elif os_type == DARWIN:
         cmd = "sysctl -n machdep.cpu.vendor"
-        out, ec = run_cmd(cmd, force_in_dry_run=True, trace=False, stream_output=False, log_ok=False)
+        out, ec = run_cmd(cmd, force_in_dry_run=True, trace=False, stream_output=False, log_ok=False, with_hooks=False)
         out = out.strip()
         if ec == 0 and out in VENDOR_IDS:
             vendor = VENDOR_IDS[out]
             _log.debug("Determined CPU vendor on DARWIN as being '%s' via cmd '%s" % (vendor, cmd))
         else:
             cmd = "sysctl -n machdep.cpu.brand_string"
-            out, ec = run_cmd(cmd, force_in_dry_run=True, trace=False, stream_output=False, log_ok=False)
+            out, ec = run_cmd(cmd, force_in_dry_run=True, trace=False, stream_output=False, log_ok=False,
+                              with_hooks=False)
             out = out.strip().split(' ')[0]
             if ec == 0 and out in CPU_VENDORS:
                 vendor = out
@@ -503,7 +504,7 @@ def get_cpu_model():
 
     elif os_type == DARWIN:
         cmd = "sysctl -n machdep.cpu.brand_string"
-        out, ec = run_cmd(cmd, force_in_dry_run=True, trace=False, stream_output=False)
+        out, ec = run_cmd(cmd, force_in_dry_run=True, trace=False, stream_output=False, with_hooks=False)
         if ec == 0:
             model = out.strip()
             _log.debug("Determined CPU model on Darwin using cmd '%s': %s" % (cmd, model))
@@ -548,7 +549,7 @@ def get_cpu_speed():
     elif os_type == DARWIN:
         cmd = "sysctl -n hw.cpufrequency_max"
         _log.debug("Trying to determine CPU frequency on Darwin via cmd '%s'" % cmd)
-        out, ec = run_cmd(cmd, force_in_dry_run=True, trace=False, stream_output=False)
+        out, ec = run_cmd(cmd, force_in_dry_run=True, trace=False, stream_output=False, with_hooks=False)
         out = out.strip()
         cpu_freq = None
         if ec == 0 and out:
@@ -596,7 +597,8 @@ def get_cpu_features():
         for feature_set in ['extfeatures', 'features', 'leaf7_features']:
             cmd = "sysctl -n machdep.cpu.%s" % feature_set
             _log.debug("Trying to determine CPU features on Darwin via cmd '%s'", cmd)
-            out, ec = run_cmd(cmd, force_in_dry_run=True, trace=False, stream_output=False, log_ok=False)
+            out, ec = run_cmd(cmd, force_in_dry_run=True, trace=False, stream_output=False, log_ok=False,
+                              with_hooks=False)
             if ec == 0:
                 cpu_feat.extend(out.strip().lower().split())
 
@@ -624,7 +626,7 @@ def get_gpu_info():
             cmd = "nvidia-smi --query-gpu=gpu_name,driver_version --format=csv,noheader"
             _log.debug("Trying to determine NVIDIA GPU info on Linux via cmd '%s'", cmd)
             out, ec = run_cmd(cmd, simple=False, log_ok=False, log_all=False,
-                              force_in_dry_run=True, trace=False, stream_output=False)
+                              force_in_dry_run=True, trace=False, stream_output=False, with_hooks=False)
             if ec == 0:
                 for line in out.strip().split('\n'):
                     nvidia_gpu_info = gpu_info.setdefault('NVIDIA', {})
@@ -643,14 +645,14 @@ def get_gpu_info():
             cmd = "rocm-smi --showdriverversion --csv"
             _log.debug("Trying to determine AMD GPU driver on Linux via cmd '%s'", cmd)
             out, ec = run_cmd(cmd, simple=False, log_ok=False, log_all=False,
-                              force_in_dry_run=True, trace=False, stream_output=False)
+                              force_in_dry_run=True, trace=False, stream_output=False, with_hooks=False)
             if ec == 0:
                 amd_driver = out.strip().split('\n')[1].split(',')[1]
 
             cmd = "rocm-smi --showproductname --csv"
             _log.debug("Trying to determine AMD GPU info on Linux via cmd '%s'", cmd)
             out, ec = run_cmd(cmd, simple=False, log_ok=False, log_all=False,
-                              force_in_dry_run=True, trace=False, stream_output=False)
+                              force_in_dry_run=True, trace=False, stream_output=False, with_hooks=False)
             if ec == 0:
                 for line in out.strip().split('\n')[1:]:
                     amd_card_series = line.split(',')[1]
@@ -898,7 +900,7 @@ def get_tool_version(tool, version_option='--version', ignore_ec=False):
     Output is returned as a single-line string (newlines are replaced by '; ').
     """
     out, ec = run_cmd(' '.join([tool, version_option]), simple=False, log_ok=False, force_in_dry_run=True,
-                      trace=False, stream_output=False)
+                      trace=False, stream_output=False, with_hooks=False)
     if not ignore_ec and ec:
         _log.warning("Failed to determine version of %s using '%s %s': %s" % (tool, tool, version_option, out))
         return UNKNOWN