Add (exclude_)where args to workload directive [AI]

lib/ramble/docs/dev_guides/application_dev_guide.rst

@@ -210,6 +210,14 @@ workload in an ``application.py`` allows it to be used within a
 :ref:`workspace-config` and will be shown when executing ``ramble info <app>``
 on the named application.
 
+Workloads can also define internal filters using the ``where`` and ``exclude_where``
+arguments. These arguments take lists of logical expressions (similar to the
+``--where`` flag in the CLI). If defined, these conditions are evaluated for each
+experiment generated by the workload. Any experiment that evaluates to ``False`` for
+any expression in ``where``, or ``True`` for any expression in ``exclude_where``,
+will be silently dropped during workspace setup. If a workload drops all of its
+experiments because of these internal clauses, Ramble will emit a warning.
+
 ^^^^^^^^^^^^^^^
 Workload Groups
 ^^^^^^^^^^^^^^^

lib/ramble/ramble/experiment_set.py

@@ -366,6 +366,7 @@ def _process_render_object(
         """Helper to render a base and its repeated experiments, for parallel execution."""
         experiment_vars, repeats = render_item
         processed_experiments = []
+        wl_stats = {}
         # Expand and prepare base and repeated experiments
         # TODO: Exploit the relationship between base and repeated experiments,
         # to save up redundant works.
@@ -391,6 +392,10 @@ def _process_render_object(
                 self.keywords.experiment_namespace
             )
 
+            wl_name = app_inst.expander.workload_name
+            if wl_name not in wl_stats:
+                wl_stats[wl_name] = {"passed_global": 0, "dropped_wl": 0}
+
             # Skip explicitly excluded experiments
             if final_exp_name not in excluded_experiments:
                 active = True
@@ -400,10 +405,26 @@ def _process_render_object(
                             active = False
                             break
 
+                if active:
+                    wl_stats[wl_name]["passed_global"] += 1
+                    wl_inst = app_inst.get_workload()
+                    if wl_inst:
+                        for expression in wl_inst.where:
+                            if not app_inst.expander.evaluate_predicate(expression):
+                                active = False
+                                wl_stats[wl_name]["dropped_wl"] += 1
+                                break
+                        if active:
+                            for expression in wl_inst.exclude_where:
+                                if app_inst.expander.evaluate_predicate(expression):
+                                    active = False
+                                    wl_stats[wl_name]["dropped_wl"] += 1
+                                    break
+
                 if active:
                     app_inst.read_status()
                     processed_experiments.append((app_inst, final_exp_namespace, n == 0))
-        return processed_experiments
+        return processed_experiments, wl_stats
 
     def render_experiment_set(
         self,
@@ -629,8 +650,18 @@ def _ingest_experiments(
             with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
                 results = list(executor.map(worker_func, render_list))
 
-        for processed_experiments in results:
+        overall_wl_stats = {}
+        for processed_experiments, wl_stats in results:
             all_processed_experiments.extend(processed_experiments)
+            for wl_name, stats in wl_stats.items():
+                if wl_name not in overall_wl_stats:
+                    overall_wl_stats[wl_name] = {
+                        "passed_global": 0,
+                        "dropped_wl": 0,
+                        "processed": 0,
+                    }
+                overall_wl_stats[wl_name]["passed_global"] += stats["passed_global"]
+                overall_wl_stats[wl_name]["dropped_wl"] += stats["dropped_wl"]
 
         # The results are now processed serially to update the experiment set state
         for app_inst, final_exp_namespace, is_base_experiment in all_processed_experiments:
@@ -761,6 +792,8 @@ def _ingest_experiments(
                         ) from None
 
             workload_names.add(app_inst.expander.workload_name)
+            if app_inst.expander.workload_name in overall_wl_stats:
+                overall_wl_stats[app_inst.expander.workload_name]["processed"] += 1
 
             app_inst.define_variable(
                 self.keywords.experiment_index,
@@ -775,6 +808,13 @@ def _ingest_experiments(
             else:
                 self.add_chained_experiment(app_inst.expander.experiment_name, app_inst)
 
+        for wl_name, stats in overall_wl_stats.items():
+            if stats["passed_global"] > 0 and stats["processed"] == 0 and stats["dropped_wl"] > 0:
+                logger.warn(
+                    f"Workload {wl_name} generated zero valid experiments because they were "
+                    "all filtered out by the workload's internal clauses."
+                )
+
         self.define_scoped_tables(workload_names, experiment_template_name)
         return rendered_instances
 

lib/ramble/ramble/language/application_language.py

@@ -59,6 +59,8 @@ def workload(
     inputs=None,
     tags=None,
     when=None,
+    where=None,
+    exclude_where=None,
     **kwargs,
 ):
     """Adds a workload to this application
@@ -90,7 +92,7 @@ def _execute_workload(app):
             app.workloads[when_set] = {}
 
         app.workloads[when_set][name] = ramble.workload.Workload(
-            name, all_execs, all_inputs, tags, when_list
+            name, all_execs, all_inputs, tags, where, exclude_where, when_list
         )
 
     return _execute_workload

lib/ramble/ramble/test/application_language.py

@@ -657,3 +657,16 @@ def test_add_variable_validator_numeric_conversion(var_value, expected_result):
 
     expander = ramble.expander.Expander({"test_var": var_value}, None)
     assert expander.evaluate_predicate(predicate) == expected_result
+
+
+@pytest.mark.parametrize("app_class", app_types)
+def test_workload_directive_where(app_class):
+    app_inst = app_class("/not/a/path")
+    wl_name = "test_where_wl"
+    app_inst.workload(
+        wl_name, executable="test_exec", where=["{n_nodes} <= 4"], exclude_where=["{n_gpus} == 0"]
+    )
+    assert hasattr(app_inst, "workloads")
+    assert wl_name in app_inst.workloads[_FS]
+    assert app_inst.workloads[_FS][wl_name].where == ["{n_nodes} <= 4"]
+    assert app_inst.workloads[_FS][wl_name].exclude_where == ["{n_gpus} == 0"]

lib/ramble/ramble/test/end_to_end/workload_where_clauses.py

@@ -0,0 +1,120 @@
+# Copyright 2022-2026 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
+
+from ramble.main import RambleCommand
+
+# everything here should be mocked if possible
+pytestmark = pytest.mark.usefixtures("mutable_config", "mutable_mock_workspace_path")
+
+workspace = RambleCommand("workspace")
+
+
+@pytest.mark.parametrize(
+    "workload,experiments,expected_generated,expected_dropped",
+    [
+        ("always", [("test1", "foo")], ["test1"], []),
+        (
+            "only_when_var_is_foo",
+            [("test_foo", "foo"), ("test_bar", "bar")],
+            ["test_foo"],
+            ["test_bar"],
+        ),
+        (
+            "exclude_when_var_is_bar",
+            [("test_foo", "foo"), ("test_bar", "bar")],
+            ["test_foo"],
+            ["test_bar"],
+        ),
+    ],
+)
+def test_workload_where_clauses(
+    workload,
+    experiments,
+    expected_generated,
+    expected_dropped,
+    mutable_config,
+    mutable_mock_workspace_path,
+    mutable_mock_apps_repo,
+):
+    """Test workload where and exclude_where clauses."""
+
+    import ramble.workspace
+
+    workspace_name = f"test_workload_where_clauses_{workload}"
+    ws = ramble.workspace.create(workspace_name)
+    ws.write()
+
+    for exp, var in experiments:
+        workspace(
+            "manage",
+            "experiments",
+            "workload-where-mock",
+            "--workload-filter",
+            workload,
+            "--experiment-name",
+            exp,
+            "-v",
+            f"test_var={var}",
+            global_args=["-D", ws.root],
+        )
+
+    out = workspace("setup", "--dry-run", global_args=["-D", ws.root])
+
+    for exp in expected_generated:
+        assert f"workload-where-mock.{workload}.{exp}" in out
+
+    for exp in expected_dropped:
+        assert f"workload-where-mock.{workload}.{exp}" not in out
+
+
+@pytest.mark.parametrize(
+    "workload,experiments,expected_warning",
+    [
+        (
+            "only_when_var_is_foo",
+            [("test_bar", "bar"), ("test_baz", "baz")],
+            "Workload only_when_var_is_foo generated zero valid experiments because they were "
+            "all filtered out by the workload's internal clauses.",
+        ),
+    ],
+)
+def test_workload_where_clauses_warnings(
+    workload,
+    experiments,
+    expected_warning,
+    mutable_config,
+    mutable_mock_workspace_path,
+    mutable_mock_apps_repo,
+):
+    """Test workload where and exclude_where clauses."""
+
+    import ramble.workspace
+
+    workspace_name = f"test_workload_where_clauses_warnings_{workload}"
+    ws = ramble.workspace.create(workspace_name)
+    ws.write()
+
+    for exp, var in experiments:
+        workspace(
+            "manage",
+            "experiments",
+            "workload-where-mock",
+            "--workload-filter",
+            workload,
+            "--experiment-name",
+            exp,
+            "-v",
+            f"test_var={var}",
+            global_args=["-D", ws.root],
+        )
+
+    out = workspace("setup", "--dry-run", global_args=["-D", ws.root])
+
+    assert expected_warning in out

lib/ramble/ramble/workload.py

@@ -27,6 +27,8 @@ def __init__(
         executables: List[str],
         inputs: Optional[List[str]] = None,
         tags: Optional[List[str]] = None,
+        where: Optional[List[str]] = None,
+        exclude_where: Optional[List[str]] = None,
         when: Optional[List[str]] = None,
     ):
         """Constructor for a workload
@@ -41,6 +43,10 @@ def __init__(
             inputs = []
         if tags is None:
             tags = []
+        if where is None:
+            where = []
+        if exclude_where is None:
+            exclude_where = []
         if when is None:
             when = []
 
@@ -49,8 +55,8 @@ def __init__(
         self.environment_variables: Dict[FrozenSet[str], List[EnvironmentVariable]] = {}
         self.when = when
 
-        attr_names = ["executables", "inputs", "tags"]
-        attr_vals = [executables, inputs, tags]
+        attr_names = ["executables", "inputs", "tags", "where", "exclude_where"]
+        attr_vals = [executables, inputs, tags, where, exclude_where]
 
         for attr, vals in zip(attr_names, attr_vals):
             if isinstance(vals, list):
@@ -79,6 +85,8 @@ def as_str(self, n_indent: int = 0, verbose: bool = False):
             ("Executables", "executables"),
             ("Inputs", "inputs"),
             ("Tags", "tags"),
+            ("Where", "where"),
+            ("Exclude Where", "exclude_where"),
             ("When", "when"),
         ]
 

var/ramble/repos/builtin.mock/applications/workload-where-mock/application.py

@@ -0,0 +1,29 @@
+# Copyright 2022-2026 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 WorkloadWhereMock(ExecutableApplication):
+    name = "workload-where-mock"
+
+    executable("test_exec", "echo '{test_var}'", use_mpi=False)
+
+    workload_variable("test_var", default="foo", workloads=["*"])
+
+    workload("always", executable="test_exec")
+    workload(
+        "only_when_var_is_foo",
+        executable="test_exec",
+        where=["'{test_var}' == 'foo'"],
+    )
+    workload(
+        "exclude_when_var_is_bar",
+        executable="test_exec",
+        exclude_where=["'{test_var}' == 'bar'"],
+    )