Merge pull request #1329 from douglasjacobsen/exec-mod-restrictions

Add usage filters to executable_modifiers

Author: linsword13

lib/ramble/ramble/language/modifier_language.py

@@ -131,7 +131,7 @@ def _execute_variable_modification(mod):
 
 
 @modifier_directive("executable_modifiers")
-def executable_modifier(name, when=None, **kwargs):
+def executable_modifier(name, usage_filter=None, when=None, **kwargs):
     """Register an executable modifier
 
     Executable modifiers can modify various aspects of non-builtin application
@@ -166,6 +166,9 @@ def write_exec_name(self, executable_name, executable, app_inst=None):
     Args:
         name (str): Name of executable modifier to use. Should be the name of a
                     class method.
+        usage_filter (str): Filters the application of this executable modifier.
+                            Modifiers can register filters to select how to apply this.
+                            Valid default options include: None, "once", "first_mpi", "all_mpi"
         when (list | None): List of when conditions this executable modifier should apply in
 
     Each executable modifier needs to return:
@@ -185,6 +188,7 @@ def _executable_modifier(mod):
             mod.executable_modifiers[when_set] = {}
 
         mod.executable_modifiers[when_set][name] = {
+            "usage_filter": usage_filter,
             "when": when_list,
         }
 

lib/ramble/ramble/test/modifier_functionality/executable_modifier_usage_filters.py

@@ -0,0 +1,164 @@
+# Copyright 2022-2025 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.
+
+import os
+import re
+
+import pytest
+
+import ramble.workspace
+from ramble.error import RambleCommandError
+from ramble.main import RambleCommand
+
+workspace = RambleCommand("workspace")
+
+
+@pytest.mark.parametrize(
+    "filter_type,expected_count",
+    [
+        ("none", -1),
+        ("once", 1),
+        ("all_mpi", 6),
+        ("first_mpi", 1),
+    ],
+)
+def test_executable_modifier_usage_filters(
+    mutable_mock_workspace_path,
+    mutable_applications,
+    mock_modifiers,
+    workspace_name,
+    filter_type,
+    expected_count,
+):
+
+    global_args = ["-w", workspace_name]
+
+    pre_str = f"exec_mod_{filter_type}_pre_applied"
+    post_str = f"exec_mod_{filter_type}_post_applied"
+
+    with ramble.workspace.create(workspace_name) as ws:
+        ws.write()
+
+        workspace(
+            "manage",
+            "experiments",
+            "openfoam",
+            "--wf",
+            "motorbike_20m",
+            "-v",
+            "n_ranks=1",
+            "-v",
+            "n_nodes=1",
+            "-v",
+            "openfoam_path=/not/needed",
+            global_args=global_args,
+        )
+
+        workspace(
+            "manage",
+            "modifiers",
+            "--add",
+            "--scope",
+            "workspace",
+            "--name",
+            "exec-mod-usage-filters",
+            global_args=global_args,
+        )
+
+        with open(os.path.join(ws.config_dir, "variants.yaml"), "w+") as f:
+            f.write(
+                f"""variants:
+  usage_filter_type: {filter_type}"""
+            )
+
+        ws._re_read()
+
+        workspace("setup", "--dry-run", global_args=global_args)
+
+        exec_path = os.path.join(
+            ws.experiment_dir, "openfoam", "motorbike_20m", "generated", "execute_experiment"
+        )
+
+        assert os.path.isfile(exec_path)
+
+        pre_regex = re.compile(pre_str)
+        post_regex = re.compile(post_str)
+
+        with open(exec_path) as f:
+            pre_count = 0
+            post_count = 0
+
+            for line in f.readlines():
+                pre_m = pre_regex.search(line)
+                if pre_m:
+                    pre_count += 1
+
+                post_m = post_regex.search(line)
+                if post_m:
+                    post_count += 1
+
+        if expected_count >= 0:
+            assert pre_count == expected_count
+            assert post_count == expected_count
+        else:
+            assert pre_count > 0
+            assert post_count > 0
+
+
+def test_executable_modifier_usage_filters_broken_errors(
+    mutable_mock_workspace_path,
+    mutable_applications,
+    mock_modifiers,
+    workspace_name,
+):
+
+    expected_err = (
+        "When extracting a usage_filter for an executable_modifier "
+        "on modifier exec-mod-usage-filters the filter __broken__ does not exist"
+    )
+    global_args = ["-w", workspace_name]
+
+    with ramble.workspace.create(workspace_name) as ws:
+        ws.write()
+
+        workspace(
+            "manage",
+            "experiments",
+            "openfoam",
+            "--wf",
+            "motorbike_20m",
+            "-v",
+            "n_ranks=1",
+            "-v",
+            "n_nodes=1",
+            "-v",
+            "openfoam_path=/not/needed",
+            global_args=global_args,
+        )
+
+        workspace(
+            "manage",
+            "modifiers",
+            "--add",
+            "--scope",
+            "workspace",
+            "--name",
+            "exec-mod-usage-filters",
+            global_args=global_args,
+        )
+
+        with open(os.path.join(ws.config_dir, "variants.yaml"), "w+") as f:
+            f.write(
+                """variants:
+  usage_filter_type: broken"""
+            )
+
+        ws._re_read()
+
+        with pytest.raises(RambleCommandError, match=expected_err):
+            workspace("setup", "--dry-run", global_args=global_args)

var/ramble/repos/builtin.mock/modifiers/exec-mod-usage-filters/modifier.py

@@ -0,0 +1,146 @@
+# Copyright 2022-2025 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 *
+
+
+class ExecModUsageFilters(BasicModifier):
+    """Define a modifier to test usage filters on executable modifiers"""
+
+    name = "exec-mod-usage-filters"
+
+    tags("test")
+
+    mode("test", description="This is a test mode")
+    default_mode("test")
+
+    modifier_conflict(MODIFIER_CONFLICT.name_mode_executables)
+
+    variant(
+        "usage_filter_type",
+        default="none",
+        values=["none", "once", "all_mpi", "first_mpi", "broken"],
+        description="Control which usage filter to use on exec mods",
+    )
+
+    executable_modifier(
+        "exec_mod_none", usage_filter=None, when=["usage_filter_type=none"]
+    )
+
+    def exec_mod_none(self, executable_name, executable, app_inst=None):
+        from ramble.util.executable import CommandExecutable
+
+        pre_cmds = [
+            CommandExecutable(
+                "exec-mod-none-pre", template=["exec_mod_none_pre_applied"]
+            )
+        ]
+
+        post_cmds = [
+            CommandExecutable(
+                "exec-mod-once-post", template=["exec_mod_none_post_applied"]
+            )
+        ]
+
+        return pre_cmds, post_cmds
+
+    executable_modifier(
+        "exec_mod_once", usage_filter="once", when=["usage_filter_type=once"]
+    )
+
+    def exec_mod_once(self, executable_name, executable, app_inst=None):
+        from ramble.util.executable import CommandExecutable
+
+        pre_cmds = [
+            CommandExecutable(
+                "exec-mod-once-pre", template=["exec_mod_once_pre_applied"]
+            )
+        ]
+
+        post_cmds = [
+            CommandExecutable(
+                "exec-mod-once-post", template=["exec_mod_once_post_applied"]
+            )
+        ]
+
+        return pre_cmds, post_cmds
+
+    executable_modifier(
+        "exec_mod_first_mpi",
+        usage_filter="first_mpi",
+        when=["usage_filter_type=first_mpi"],
+    )
+
+    def exec_mod_first_mpi(self, executable_name, executable, app_inst=None):
+        from ramble.util.executable import CommandExecutable
+
+        pre_cmds = [
+            CommandExecutable(
+                "exec-mod-first-mpi-pre",
+                template=["exec_mod_first_mpi_pre_applied"],
+            )
+        ]
+
+        post_cmds = [
+            CommandExecutable(
+                "exec-mod-first-mpi-post",
+                template=["exec_mod_first_mpi_post_applied"],
+            )
+        ]
+
+        return pre_cmds, post_cmds
+
+    executable_modifier(
+        "exec_mod_all_mpi",
+        usage_filter="all_mpi",
+        when=["usage_filter_type=all_mpi"],
+    )
+
+    def exec_mod_all_mpi(self, executable_name, executable, app_inst=None):
+        from ramble.util.executable import CommandExecutable
+
+        pre_cmds = [
+            CommandExecutable(
+                "exec-mod-all-mpi-pre",
+                template=["exec_mod_all_mpi_pre_applied"],
+            )
+        ]
+
+        post_cmds = [
+            CommandExecutable(
+                "exec-mod-all-mpi-post",
+                template=["exec_mod_all_mpi_post_applied"],
+            )
+        ]
+
+        return pre_cmds, post_cmds
+
+    executable_modifier(
+        "exec_mod_broken",
+        usage_filter="__broken__",
+        when=["usage_filter_type=broken"],
+    )
+
+    def exec_mod_broken(self, executable_name, executable, app_inst=None):
+        from ramble.util.executable import CommandExecutable
+
+        pre_cmds = [
+            CommandExecutable(
+                "exec-mod-all-mpi-pre",
+                template=["exec_mod_all_mpi_pre_applied"],
+            )
+        ]
+
+        post_cmds = [
+            CommandExecutable(
+                "exec-mod-all-mpi-post",
+                template=["exec_mod_all_mpi_post_applied"],
+            )
+        ]
+
+        return pre_cmds, post_cmds