Merge pull request #994 from dapomeroy/register_phase_when

Add 'when' to register_phase directive

Author: douglasjacobsen

lib/ramble/ramble/application.py

@@ -461,10 +461,11 @@ def get_pipeline_phases(self, pipeline, phase_filters=None):
         final_added_index = None
         if pipeline in self._pipeline_graphs:
             for idx, phase in enumerate(self._pipeline_graphs[pipeline].walk()):
-                for phase_filter in phase_filters:
-                    if fnmatch.fnmatch(phase.key, phase_filter):
-                        phases.add(phase)
-                        final_added_index = idx
+                if self.expander.satisfies(phase.when, variant_set=self.object_variants):
+                    for phase_filter in phase_filters:
+                        if fnmatch.fnmatch(phase.key, phase_filter):
+                            phases.add(phase)
+                            final_added_index = idx
 
         include_phase_deps = ramble.config.get("config:include_phase_dependencies")
         if include_phase_deps:

lib/ramble/ramble/language/language_base.py

@@ -213,6 +213,7 @@ def _wrapper(*args, **_kwargs):
                         "required_package",
                         "define_compiler",
                         "package_manager_config",
+                        "register_phase",
                     ]:
                         msg = (
                             'directive "{0}" cannot be used within a "when"'

lib/ramble/ramble/language/shared_language.py

@@ -138,6 +138,7 @@ def define_compiler(
         compiler (str): Package name to use for compilation
         package_manager (str): Glob supported pattern to match package managers
                                this compiler applies to
+        when (list | None): List of when conditions to apply to directive
     """
 
     def _execute_define_compiler(obj):
@@ -184,6 +185,7 @@ def software_spec(
         compiler (str): Package name to use as compiler for compiling this package
         package_manager (str): Glob supported pattern to match package managers
                                this package applies to
+        when (list | None): List of when conditions to apply to directive
     """
 
     def _execute_software_spec(obj):
@@ -221,6 +223,7 @@ def package_manager_config(name, config, package_manager=None, when=None, **kwar
         name (str): Name of this configuration
         config (str): Configuration option to set
         package_manager (str): Name of the package manager this config should be used with
+        when (list | None): List of when conditions to apply to directive
     """
 
     def _execute_package_manager_config(obj):
@@ -251,6 +254,7 @@ def required_package(name, package_manager=None, when=None, **kwargs):
     Args:
         name (str): Name of required package
         package_manager (str): Glob package manager name to apply this required package to
+        when (list | None): List of when conditions to apply to directive
     """
 
     def _execute_required_package(obj):
@@ -419,7 +423,7 @@ def _store_builtin(obj):
 
 
 @shared_directive("phase_definitions")
-def register_phase(name, pipeline=None, run_before=None, run_after=None, **kwargs):
+def register_phase(name, pipeline=None, run_before=None, run_after=None, when=None, **kwargs):
     """Register a phase
 
     Phases are portions of a pipeline that will execute when
@@ -436,6 +440,7 @@ def register_phase(name, pipeline=None, run_before=None, run_after=None, **kwarg
       pipeline (str): The name of the pipeline this phase should be registered into.
       run_before (list(str) | None): A list of phase names this phase should run before
       run_after (list(str) | None): A list of phase names this phase should run after
+      when (list | None): List of when conditions to apply to directive
     """
     if run_before is None:
         run_before = []
@@ -445,6 +450,10 @@ def register_phase(name, pipeline=None, run_before=None, run_after=None, **kwarg
     def _execute_register_phase(obj):
         import ramble.util.graph
 
+        when_list = ramble.language.language_helpers.build_when_list(
+            when, obj, name, "register_phase"
+        )
+
         if pipeline not in obj._pipelines:
             raise ramble.language.language_base.DirectiveError(
                 "Directive register_phase was "
@@ -490,6 +499,8 @@ def _execute_register_phase(obj):
         for after in run_after:
             phase_node.order_after(after)
 
+        phase_node.when = when_list
+
         obj.phase_definitions[pipeline][name] = phase_node
 
     return _execute_register_phase

lib/ramble/ramble/test/when.py

@@ -0,0 +1,54 @@
+# 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 pytest
+
+import ramble.workspace
+from ramble.main import RambleCommand
+
+pytestmark = pytest.mark.usefixtures("mutable_mock_workspace_path", "mutable_mock_apps_repo")
+
+config = RambleCommand("config")
+workspace = RambleCommand("workspace")
+
+
+def test_register_phase_when(request):
+    ws_name = request.node.name
+
+    global_args = ["-w", ws_name]
+
+    with ramble.workspace.create(ws_name) as ws:
+        workspace(
+            "manage",
+            "experiments",
+            "when-directives",
+            "--wf",
+            "test_wl",
+            "-v",
+            "n_ranks=1",
+            "-v",
+            "n_nodes=1",
+            "-v",
+            "processes_per_node=1",
+            global_args=global_args,
+        )
+
+        config("add", "variants:register_phase_when:true", global_args=global_args)
+
+        ws._re_read()
+        output = workspace("setup", "--dry-run", global_args=global_args)
+
+        assert "Test Phase" in output
+
+        config("remove", "variants:register_phase_when:true", global_args=global_args)
+        config("add", "variants:register_phase_when:false", global_args=global_args)
+
+        ws._re_read()
+        output = workspace("setup", "--dry-run", global_args=global_args)
+
+        assert "Test Phase" not in output

lib/ramble/ramble/util/graph.py

@@ -28,6 +28,7 @@ def __init__(self, key, attribute=None, obj_inst=None):
         self._order_before = []
         self._order_after = []
         self.obj_inst = obj_inst
+        self.when = []
 
     def set_attribute(self, attr):
         """Sets the attribute of a graph node

var/ramble/repos/builtin.mock/applications/when-directives/application.py

@@ -0,0 +1,41 @@
+# 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.appkit import *
+
+
+class WhenDirectives(ExecutableApplication):
+    name = "when-directives"
+
+    executable("test_exec", "echo '{test_variable}'", use_mpi=False)
+
+    workload("test_wl", executable="test_exec")
+
+    with default_args(workload="test_wl"):
+        workload_variable(
+            "test_variable",
+            default="Test",
+            description="Variable to print for testing",
+        )
+
+    variant(
+        "register_phase_when",
+        default=False,
+        values=[True, False],
+        description="Register additional phase using when",
+    )
+
+    with when("+register_phase_when"):
+        register_phase(
+            "test_phase",
+            pipeline="setup",
+            run_before=["get_inputs"],
+        )
+
+    def _test_phase(self, workspace, app_inst):
+        print("Test Phase")