From 298917a0b7f4cd9663c7eea53e3ba0aa7ffcfba6 Mon Sep 17 00:00:00 2001 From: Lin Guo Date: Thu, 18 Jul 2024 00:31:55 -0700 Subject: [PATCH 1/3] Add a `run_in_background` option for executables --- lib/ramble/ramble/application.py | 7 ++++++- lib/ramble/ramble/language/application_language.py | 3 ++- lib/ramble/ramble/test/application_language.py | 12 ++++++++---- .../ramble/test/end_to_end/custom_executables.py | 3 ++- lib/ramble/ramble/util/executable.py | 5 +++++ .../builtin/applications/hostname/application.py | 8 ++++++++ 6 files changed, 31 insertions(+), 7 deletions(-) diff --git a/lib/ramble/ramble/application.py b/lib/ramble/ramble/application.py index 176c016aa..d03d0627e 100644 --- a/lib/ramble/ramble/application.py +++ b/lib/ramble/ramble/application.py @@ -1060,8 +1060,13 @@ def _define_commands( if ramble.config.get("config:shell") in ["bash", "sh"]: redirect += " 2>&1" + if cmd_conf.run_in_background: + bg_cmd = " &" + else: + bg_cmd = "" + for part in cmd_conf.template: - command_part = f"{mpi_cmd}{part}{redirect}" + command_part = f"{mpi_cmd}{part}{redirect}{bg_cmd}" self._command_list.append( self.expander.expand_var(command_part, exec_vars) ) diff --git a/lib/ramble/ramble/language/application_language.py b/lib/ramble/ramble/language/application_language.py index ea7032e51..bb099f3d1 100644 --- a/lib/ramble/ramble/language/application_language.py +++ b/lib/ramble/ramble/language/application_language.py @@ -134,7 +134,8 @@ def executable(name, template, **kwargs): defaults to {log_file} output_capture (Optional): Declare which ouptu (stdout, stderr, both) to capture. Defaults to stdout - + run_in_background (Optional): Declare if the command should run in the background. + Defaults to False """ def _execute_executable(app): diff --git a/lib/ramble/ramble/test/application_language.py b/lib/ramble/ramble/test/application_language.py index cbbb4ed81..902f840fa 100644 --- a/lib/ramble/ramble/test/application_language.py +++ b/lib/ramble/ramble/test/application_language.py @@ -74,7 +74,7 @@ def add_workload(app_inst, wl_num=1, func_type=func_types.directive): def add_executable(app_inst, exe_num=1, func_type=func_types.directive): - nompi_exec_name = "SerialExe%s" % exe_num + nompi_bg_exec_name = "SerialExe%s" % exe_num mpi_exec_name = "MpiExe%s" % exe_num nompi_list_exec_name = "MultiLineSerialExe%s" % exe_num mpi_list_exec_name = "MultiLineMpiExe%s" % exe_num @@ -84,11 +84,12 @@ def add_executable(app_inst, exe_num=1, func_type=func_types.directive): if func_type == func_types.directive: executable( - nompi_exec_name, + nompi_bg_exec_name, template, # noqa: F405 use_mpi=False, redirect=redirect_test, output_capture=output_capture, + run_in_background=True, )(app_inst) executable(mpi_exec_name, template, use_mpi=True)(app_inst) # noqa: F405 @@ -108,11 +109,12 @@ def add_executable(app_inst, exe_num=1, func_type=func_types.directive): )(app_inst) elif func_type == func_types.method: app_inst.executable( - nompi_exec_name, + nompi_bg_exec_name, template, # noqa: F405 use_mpi=False, redirect=redirect_test, output_capture=output_capture, + run_in_background=True, ) app_inst.executable(mpi_exec_name, template, use_mpi=True) # noqa: F405 @@ -134,16 +136,18 @@ def add_executable(app_inst, exe_num=1, func_type=func_types.directive): assert False exec_def = { - nompi_exec_name: { + nompi_bg_exec_name: { "template": [template], "mpi": False, "redirect": redirect_test, "output_capture": output_capture, + "run_in_background": True, }, mpi_exec_name: { "template": [template], "mpi": True, "redirect": "{log_file}", # Default value + "run_in_background": False, # Default }, nompi_list_exec_name: { "template": [template, template, template], diff --git a/lib/ramble/ramble/test/end_to_end/custom_executables.py b/lib/ramble/ramble/test/end_to_end/custom_executables.py index c41214fc6..061b93c14 100644 --- a/lib/ramble/ramble/test/end_to_end/custom_executables.py +++ b/lib/ramble/ramble/test/end_to_end/custom_executables.py @@ -45,6 +45,7 @@ def test_custom_executables(mutable_config, mutable_mock_workspace_path, mock_ap use_mpi: false redirect: '{log_file}' output_capture: '>>' + run_in_background: true before_all: template: - 'echo "before all"' @@ -102,7 +103,7 @@ def test_custom_executables(mutable_config, mutable_mock_workspace_path, mock_ap import re - custom_regex = re.compile("lscpu >>") + custom_regex = re.compile(r"^lscpu >> .* &$") export_regex = re.compile(r"export MY_VAR=TEST") cmd_regex = re.compile("foo >>") diff --git a/lib/ramble/ramble/util/executable.py b/lib/ramble/ramble/util/executable.py index c5ee71b8d..d7e8e1ef5 100644 --- a/lib/ramble/ramble/util/executable.py +++ b/lib/ramble/ramble/util/executable.py @@ -72,6 +72,7 @@ def __init__( variables={}, redirect="{log_file}", output_capture=OUTPUT_CAPTURE.DEFAULT, + run_in_background=False, **kwargs, ): """Create a CommandExecutable instance @@ -85,6 +86,7 @@ def __init__( - variables (dict): dictionary of variable definitions to use for this executable only - redirect: File to redirect output of template into - output_capture: Operator to use when capturing output + - run_in_background: If true, run the command in background """ if isinstance(template, str): @@ -100,6 +102,7 @@ def __init__( self.mpi = use_mpi or mpi self.redirect = redirect self.output_capture = output_capture + self.run_in_background = run_in_background self.variables = variables.copy() def copy(self): @@ -111,6 +114,7 @@ def copy(self): redirect=self.redirect, variables=self.variables, output_capture=self.output_capture, + run_in_background=self.run_in_background, ) return new_inst @@ -123,6 +127,7 @@ def __str__(self): + f" variables: {self.variables}\n" + f" redirect: {self.redirect}\n" + f" output_capture: {self.output_capture}\n" + + f" run_in_background: {self.run_in_background}\n" ) return self_str diff --git a/var/ramble/repos/builtin/applications/hostname/application.py b/var/ramble/repos/builtin/applications/hostname/application.py index 39c61e42c..873568d30 100644 --- a/var/ramble/repos/builtin/applications/hostname/application.py +++ b/var/ramble/repos/builtin/applications/hostname/application.py @@ -26,6 +26,13 @@ class Hostname(ExecutableApplication): executable( "local", "hostname", use_mpi=False, output_capture=OUTPUT_CAPTURE.ALL ) + executable( + "local_bg", + "hostname", + use_mpi=False, + output_capture=OUTPUT_CAPTURE.ALL, + run_in_background=True, + ) executable( "serial", "/usr/bin/time hostname", @@ -40,6 +47,7 @@ class Hostname(ExecutableApplication): ) workload("local", executable="local") + workload("local_bg", executable="local_bg") workload("serial", executable="serial") workload("parallel", executable="parallel") From 290c112dc247a648d05f49536de425c3ba964f12 Mon Sep 17 00:00:00 2001 From: Lin Guo Date: Thu, 18 Jul 2024 23:31:58 -0700 Subject: [PATCH 2/3] Add a modifier to wait for bg jobs completion --- lib/ramble/ramble/test/util/shell_vars.py | 23 ++++++++ lib/ramble/ramble/util/shell_vars.py | 19 ++++++ .../applications/hostname/application.py | 11 +++- .../modifiers/wait-for-bg-jobs/modifier.py | 58 +++++++++++++++++++ 4 files changed, 109 insertions(+), 2 deletions(-) create mode 100644 lib/ramble/ramble/test/util/shell_vars.py create mode 100644 lib/ramble/ramble/util/shell_vars.py create mode 100644 var/ramble/repos/builtin/modifiers/wait-for-bg-jobs/modifier.py diff --git a/lib/ramble/ramble/test/util/shell_vars.py b/lib/ramble/ramble/test/util/shell_vars.py new file mode 100644 index 000000000..665395309 --- /dev/null +++ b/lib/ramble/ramble/test/util/shell_vars.py @@ -0,0 +1,23 @@ +# Copyright 2022-2024 The Ramble Authors +# +# Licensed under the Apache License, Version 2.0 or the MIT license +# , at your +# option. This file may not be copied, modified, or distributed +# except according to those terms. + +from ramble.util.shell_vars import last_pid_var + +import pytest + + +@pytest.mark.parametrize( + "var_func,shell,expect", + [ + (last_pid_var, "fish", "$last_pid"), + (last_pid_var, "bash", "$!"), + (last_pid_var, "unknown_shell", "$!"), + ], +) +def test_shell_vars(var_func, shell, expect): + assert var_func(shell) == expect diff --git a/lib/ramble/ramble/util/shell_vars.py b/lib/ramble/ramble/util/shell_vars.py new file mode 100644 index 000000000..b1b774700 --- /dev/null +++ b/lib/ramble/ramble/util/shell_vars.py @@ -0,0 +1,19 @@ +# Copyright 2022-2024 The Ramble Authors +# +# Licensed under the Apache License, Version 2.0 or the MIT license +# , at your +# option. This file may not be copied, modified, or distributed +# except according to those terms. + +"""Utils for getting common shell special variables""" + + +def last_pid_var(shell: str) -> str: + if shell == "fish": + # Based on https://fishshell.com/docs/current/fish_for_bash_users.html#special-variables + last_pid_var = "$last_pid" + else: + # This works for sh/bash and (t)csh + last_pid_var = "$!" + return last_pid_var diff --git a/var/ramble/repos/builtin/applications/hostname/application.py b/var/ramble/repos/builtin/applications/hostname/application.py index 873568d30..f6e364eb5 100644 --- a/var/ramble/repos/builtin/applications/hostname/application.py +++ b/var/ramble/repos/builtin/applications/hostname/application.py @@ -28,7 +28,14 @@ class Hostname(ExecutableApplication): ) executable( "local_bg", - "hostname", + "(sleep 5; hostname)", + use_mpi=False, + output_capture=OUTPUT_CAPTURE.ALL, + run_in_background=True, + ) + executable( + "local_bg2", + "(sleep 10; hostname)", use_mpi=False, output_capture=OUTPUT_CAPTURE.ALL, run_in_background=True, @@ -47,7 +54,7 @@ class Hostname(ExecutableApplication): ) workload("local", executable="local") - workload("local_bg", executable="local_bg") + workload("local_bg", executables=["local_bg", "local_bg2"]) workload("serial", executable="serial") workload("parallel", executable="parallel") diff --git a/var/ramble/repos/builtin/modifiers/wait-for-bg-jobs/modifier.py b/var/ramble/repos/builtin/modifiers/wait-for-bg-jobs/modifier.py new file mode 100644 index 000000000..99e686a0e --- /dev/null +++ b/var/ramble/repos/builtin/modifiers/wait-for-bg-jobs/modifier.py @@ -0,0 +1,58 @@ +# Copyright 2022-2024 The Ramble Authors +# +# Licensed under the Apache License, Version 2.0 or the MIT license +# , at your +# option. This file may not be copied, modified, or distributed +# except according to those terms. + +from ramble.modkit import * +import ramble.config +from ramble.util.executable import CommandExecutable +from ramble.util.shell_vars import last_pid_var + + +class WaitForBgJobs(BasicModifier): + """Define a modifier for waiting for background jobs. + + This only waits for background jobs started from the experiment. + """ + + name = "wait-for-bg-jobs" + + tags("info") + + mode("standard", description="Wait for background jobs using PIDs") + + default_mode("standard") + + register_builtin("setup_pid_array") + + def setup_pid_array(self): + return ["wait_for_pid_list=()"] + + executable_modifier("store_bg_pid") + + def store_bg_pid(self, executable_name, executable, app_inst=None): + post_exec = [] + + if executable.run_in_background: + shell = ramble.config.get("config:shell") + last_pid_str = last_pid_var(shell) + post_exec.append( + CommandExecutable( + f"store-bg-pid-{executable_name}", + template=[f"wait_for_pid_list+=({last_pid_str})"], + ) + ) + + return [], post_exec + + register_builtin("wait_for_completion", injection_method="append") + + def wait_for_completion(self): + return [ + 'for wait_pid in "${wait_for_pid_list[@]}"; do', + " wait ${wait_pid}", + "done", + ] From 1f7652d2808bf5ee43247ccb770ac6e2e0edbde2 Mon Sep 17 00:00:00 2001 From: Lin Guo Date: Fri, 19 Jul 2024 14:40:48 -0700 Subject: [PATCH 3/3] Add bash-only check --- .../repos/builtin/modifiers/wait-for-bg-jobs/modifier.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/var/ramble/repos/builtin/modifiers/wait-for-bg-jobs/modifier.py b/var/ramble/repos/builtin/modifiers/wait-for-bg-jobs/modifier.py index 99e686a0e..e9ba2d55b 100644 --- a/var/ramble/repos/builtin/modifiers/wait-for-bg-jobs/modifier.py +++ b/var/ramble/repos/builtin/modifiers/wait-for-bg-jobs/modifier.py @@ -10,12 +10,14 @@ import ramble.config from ramble.util.executable import CommandExecutable from ramble.util.shell_vars import last_pid_var +import ramble.util.logger as logger class WaitForBgJobs(BasicModifier): """Define a modifier for waiting for background jobs. This only waits for background jobs started from the experiment. + For now, this modifier works with bash only. """ name = "wait-for-bg-jobs" @@ -38,6 +40,10 @@ def store_bg_pid(self, executable_name, executable, app_inst=None): if executable.run_in_background: shell = ramble.config.get("config:shell") + if shell != "bash": + logger.die( + f"waif-for-bg-jobs does not support {shell}, it's currently bash-only" + ) last_pid_str = last_pid_var(shell) post_exec.append( CommandExecutable(