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'"], + )