Add a run_in_background option for executables

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)
                         )

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):

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],

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 >>")
 

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 <LICENSE-APACHE or
+# https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+# <LICENSE-MIT or https://opensource.org/licenses/MIT>, 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

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

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 <LICENSE-APACHE or
+# https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+# <LICENSE-MIT or https://opensource.org/licenses/MIT>, 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

var/ramble/repos/builtin/applications/hostname/application.py

@@ -26,6 +26,20 @@ class Hostname(ExecutableApplication):
     executable(
         "local", "hostname", use_mpi=False, output_capture=OUTPUT_CAPTURE.ALL
     )
+    executable(
+        "local_bg",
+        "(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,
+    )
     executable(
         "serial",
         "/usr/bin/time hostname",
@@ -40,6 +54,7 @@ class Hostname(ExecutableApplication):
     )
 
     workload("local", executable="local")
+    workload("local_bg", executables=["local_bg", "local_bg2"])
     workload("serial", executable="serial")
     workload("parallel", executable="parallel")
 

var/ramble/repos/builtin/modifiers/wait-for-bg-jobs/modifier.py

@@ -0,0 +1,64 @@
+# Copyright 2022-2024 The Ramble Authors
+#
+# Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+# https://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+# <LICENSE-MIT or https://opensource.org/licenses/MIT>, 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
+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"
+
+    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")
+            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(
+                    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",
+        ]