Skip to content
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.

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

add run_shell_cmd hook #4323

Merged
merged 11 commits into from
Aug 15, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
12 changes: 9 additions & 3 deletions easybuild/tools/hooks.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,8 @@
CRASH = 'crash'
FAIL = 'fail'

RUN_SHELL_CMD = 'run_shell_cmd'

PRE_PREF = 'pre_'
POST_PREF = 'post_'
HOOK_SUFF = '_hook'
Expand Down Expand Up @@ -108,6 +110,8 @@
CANCEL,
CRASH,
FAIL,
PRE_PREF + RUN_SHELL_CMD,
POST_PREF + RUN_SHELL_CMD,
]
KNOWN_HOOKS = [h + HOOK_SUFF for h in HOOK_NAMES]

Expand Down Expand Up @@ -205,7 +209,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.

Expand All @@ -221,6 +225,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
Expand All @@ -232,6 +238,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
2 changes: 1 addition & 1 deletion easybuild/tools/options.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
44 changes: 41 additions & 3 deletions easybuild/tools/run.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand All @@ -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()

Expand Down Expand Up @@ -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,
Expand All @@ -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=''):
Expand Down Expand Up @@ -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.

Expand All @@ -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
Expand Down Expand Up @@ -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)))

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)))

Expand Down
24 changes: 13 additions & 11 deletions easybuild/tools/systemtools.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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())

Expand Down Expand Up @@ -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', {})
Expand All @@ -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]
Expand Down Expand Up @@ -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
Expand Down
Loading
Loading