From 53a52d74c59a1a6e8ab0884e932a4f5ccb1007d1 Mon Sep 17 00:00:00 2001 From: Douglas Jacobsen Date: Thu, 25 Jun 2026 08:14:34 -0700 Subject: [PATCH] Add (exclude_)where args to workload directive This commit adds support for filtering experiments at the workload directive level. This can be used to silently reduce the possible experiment definitions, without creating errors. Tests and documentation are added as well. --- .../docs/dev_guides/application_dev_guide.rst | 8 ++ lib/ramble/ramble/experiment_set.py | 44 ++++++- .../ramble/language/application_language.py | 4 +- .../ramble/test/application_language.py | 13 ++ .../test/end_to_end/workload_where_clauses.py | 120 ++++++++++++++++++ lib/ramble/ramble/workload.py | 12 +- .../workload-where-mock/application.py | 29 +++++ 7 files changed, 225 insertions(+), 5 deletions(-) create mode 100644 lib/ramble/ramble/test/end_to_end/workload_where_clauses.py create mode 100644 var/ramble/repos/builtin.mock/applications/workload-where-mock/application.py diff --git a/lib/ramble/docs/dev_guides/application_dev_guide.rst b/lib/ramble/docs/dev_guides/application_dev_guide.rst index 5feb0522d..1186e8aac 100644 --- a/lib/ramble/docs/dev_guides/application_dev_guide.rst +++ b/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 `` 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 ^^^^^^^^^^^^^^^ diff --git a/lib/ramble/ramble/experiment_set.py b/lib/ramble/ramble/experiment_set.py index 31b7b668f..8451b52f0 100644 --- a/lib/ramble/ramble/experiment_set.py +++ b/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 diff --git a/lib/ramble/ramble/language/application_language.py b/lib/ramble/ramble/language/application_language.py index 1d674ef3f..0a59d3e61 100644 --- a/lib/ramble/ramble/language/application_language.py +++ b/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 diff --git a/lib/ramble/ramble/test/application_language.py b/lib/ramble/ramble/test/application_language.py index 64d783bff..02f107e10 100644 --- a/lib/ramble/ramble/test/application_language.py +++ b/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"] diff --git a/lib/ramble/ramble/test/end_to_end/workload_where_clauses.py b/lib/ramble/ramble/test/end_to_end/workload_where_clauses.py new file mode 100644 index 000000000..9c5757b85 --- /dev/null +++ b/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 or the MIT license +# , 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 diff --git a/lib/ramble/ramble/workload.py b/lib/ramble/ramble/workload.py index c1011fb95..c1f48dbca 100644 --- a/lib/ramble/ramble/workload.py +++ b/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"), ] diff --git a/var/ramble/repos/builtin.mock/applications/workload-where-mock/application.py b/var/ramble/repos/builtin.mock/applications/workload-where-mock/application.py new file mode 100644 index 000000000..f2123439a --- /dev/null +++ b/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 or the MIT license +# , 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'"], + )