diff --git a/lib/ramble/ramble/expander.py b/lib/ramble/ramble/expander.py index c5c96accbe..dd81319039 100644 --- a/lib/ramble/ramble/expander.py +++ b/lib/ramble/ramble/expander.py @@ -18,7 +18,7 @@ import sys from contextlib import contextmanager from enum import Enum -from typing import Dict, FrozenSet, List, Optional, Union +from typing import Any, Dict, FrozenSet, List, Optional, Set, Union import ramble.config import ramble.error @@ -223,11 +223,134 @@ def _maybe(expander, var_name, default=""): } +def pow2_range(start, stop=None, inclusive=True): + """Generate a sequence of numbers by doubling from start up to stop. + + Args: + start (int): Starting value (if stop is provided), or stop value (if stop is None). + stop (int): Optional ending value. If None, start is treated as stop, + and sequence starts at 1. + inclusive (bool): Whether stop is inclusive. Defaults to True. + + Returns: + list[int]: Sequence of doubling values / powers of 2. + """ + if stop is None: + start, stop = 1, start + + start = int(start) + stop = int(stop) + if isinstance(inclusive, str): + inclusive = inclusive.lower() in ("true", "1", "yes") + + if start <= 0 or stop < start: + return [] + + values = [] + current = start + if inclusive: + while current <= stop: + values.append(current) + current *= 2 + else: + while current < stop: + values.append(current) + current *= 2 + + return values + + supported_list_function_pointers = { "range": range, + "pow2_range": pow2_range, + "power2_range": pow2_range, } +def is_dynamic_list_expression(val) -> bool: + """Check if a variable value is an unexpanded call to a supported list function. + + Returns True if val is a string matching a call to any function registered in + `supported_list_function_pointers`. + """ + if not isinstance(val, str): + return False + if not supported_list_function_pointers: + return False + func_names = "|".join(re.escape(fn) for fn in supported_list_function_pointers) + return bool(re.search(rf"\b(?:{func_names})\s*\(", val)) + + +def extract_var_refs(expr) -> Set[str]: + """Extract variable references enclosed in braces from an expression string.""" + if not isinstance(expr, str): + return set() + refs = set(re.findall(r"\{\s*([a-zA-Z_][a-zA-Z0-9_]*)", expr)) + in_brace = 0 + cur_ident: List[str] = [] + for ch in expr: + if ch == "{": + in_brace += 1 + elif ch == "}": + if cur_ident: + refs.add("".join(cur_ident)) + cur_ident = [] + in_brace = max(0, in_brace - 1) + elif in_brace > 0: + if ch.isalnum() or ch == "_": + cur_ident.append(ch) + else: + if cur_ident: + refs.add("".join(cur_ident)) + cur_ident = [] + if cur_ident and in_brace > 0: + refs.add("".join(cur_ident)) + return {r for r in refs if r and not r[0].isdigit()} + + +def find_dependent_vector_vars( + dynamic_list_vars: Dict[str, Any], + variables: Dict[str, Any], + expander: Optional["Expander"] = None, +) -> Set[str]: + """Find all vector variables in `variables` that `dynamic_list_vars` transitively depend on.""" + dependent_vector_vars = set() + visited = set() + queue = list(dynamic_list_vars.values()) + + if expander is not None: + for val in dynamic_list_vars.values(): + saved_used = expander._used_variables.copy() + expander._used_variables = set() + try: + expander.expand_var(val) + except Exception: + pass + used = expander._used_variables + expander._used_variables = saved_used + for var_ref in used: + if var_ref in variables and var_ref not in dynamic_list_vars: + if isinstance(variables[var_ref], list): + dependent_vector_vars.add(var_ref) + elif var_ref not in visited: + visited.add(var_ref) + queue.append(str(variables[var_ref])) + + while queue: + expr = queue.pop(0) + if not isinstance(expr, str): + continue + for var_ref in extract_var_refs(expr): + if var_ref in variables and var_ref not in dynamic_list_vars: + if isinstance(variables[var_ref], list): + dependent_vector_vars.add(var_ref) + elif var_ref not in visited: + visited.add(var_ref) + queue.append(str(variables[var_ref])) + + return dependent_vector_vars + + supported_modules = { "math": math, } diff --git a/lib/ramble/ramble/experiment_set.py b/lib/ramble/ramble/experiment_set.py index 15efce6784..635e40696f 100644 --- a/lib/ramble/ramble/experiment_set.py +++ b/lib/ramble/ramble/experiment_set.py @@ -285,6 +285,9 @@ def _prepare_experiment( app_inst = self._setup_experiment_minimal(workload_template_name, variables, context) + if getattr(app_inst, "has_dynamic_range_variables", False): + return app_inst + final_wl_name = app_inst.expander.expand_var_name( self.keywords.workload_name, allow_passthrough=False ) @@ -345,6 +348,9 @@ def _get_used_variables( ): app_inst = self._setup_experiment_minimal(workload_template_name, variables, context) + if getattr(app_inst, "has_dynamic_range_variables", False): + return set() + # The `_get_used_variables` is only called for the base experiment, # so no need to consider repeat suffix. exp_name = app_inst.expander.expand_var(exp_template_name, allow_passthrough=False) @@ -369,6 +375,7 @@ def _process_render_object( """Helper to render a base and its repeated experiments, for parallel execution.""" experiment_vars, repeats = render_item processed_experiments = [] + dynamic_range_experiments = [] wl_stats = {} # Expand and prepare base and repeated experiments # TODO: Exploit the relationship between base and repeated experiments, @@ -390,6 +397,10 @@ def _process_render_object( cur_repeats, ) + if getattr(app_inst, "has_dynamic_range_variables", False): + dynamic_range_experiments.append(app_inst) + break + final_exp_name = app_inst.expander.expand_var_name(self.keywords.experiment_name) final_exp_namespace = app_inst.expander.expand_var_name( self.keywords.experiment_namespace @@ -427,7 +438,7 @@ def _process_render_object( if active: app_inst.read_status() processed_experiments.append((app_inst, final_exp_namespace, n == 0)) - return processed_experiments, wl_stats + return processed_experiments, dynamic_range_experiments, wl_stats def render_experiment_set( self, @@ -654,8 +665,10 @@ def _ingest_experiments( results = list(executor.map(worker_func, render_list)) overall_wl_stats = {} - for processed_experiments, wl_stats in results: + all_dynamic_range_experiments = [] + for processed_experiments, dynamic_range_experiments, wl_stats in results: all_processed_experiments.extend(processed_experiments) + all_dynamic_range_experiments.extend(dynamic_range_experiments) for wl_name, stats in wl_stats.items(): if wl_name not in overall_wl_stats: overall_wl_stats[wl_name] = { @@ -666,6 +679,27 @@ def _ingest_experiments( overall_wl_stats[wl_name]["passed_global"] += stats["passed_global"] overall_wl_stats[wl_name]["dropped_wl"] += stats["dropped_wl"] + if all_dynamic_range_experiments: + saved_contexts = { + self._contexts.application: self._context[self._contexts.application], + self._contexts.workload: self._context[self._contexts.workload], + self._contexts.experiment: self._context[self._contexts.experiment], + } + try: + for dyn_inst in all_dynamic_range_experiments: + range_rendered = dyn_inst.render_range_experiments( + saved_contexts[self._contexts.experiment], + warn_validation=warn_validation, + die_on_validate_error=die_on_validate_error, + chained=chained, + ) + rendered_instances.extend(range_rendered) + for inst in range_rendered: + workload_names.add(inst.expander.workload_name) + finally: + for ctx_key, ctx_val in saved_contexts.items(): + self._set_context(ctx_key, ctx_val) + # 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: logger.debug(f" Final name: {final_exp_namespace}") diff --git a/lib/ramble/ramble/renderer.py b/lib/ramble/ramble/renderer.py index 9f49c06383..6c180dffa8 100644 --- a/lib/ramble/ramble/renderer.py +++ b/lib/ramble/ramble/renderer.py @@ -468,6 +468,101 @@ def render_objects(self, render_group, exclude_where=None, ignore_used=True, fat # Also expand all variables that generate lists object_variables = self._expand_variables(variables, expander) + if render_group.object == "experiment": + dynamic_list_vars = { + k: v + for k, v in object_variables.items() + if ramble.expander.is_dynamic_list_expression(v) + } + if dynamic_list_vars: + dependent_vector_vars = ramble.expander.find_dependent_vector_vars( + dynamic_list_vars, object_variables, expander + ) + + if not dependent_vector_vars: + yield object_variables, ramble.repeats.Repeats() + return + + # Expand only dependent_vector_vars (and any variables zipped with them) + # to produce separate seed instances for each vector value. + vars_to_scalarize = set(dependent_vector_vars) + if render_group.zips: + changed = True + while changed: + changed = False + for z_name, z_vars in render_group.zips.items(): + if z_name in vars_to_scalarize or any( + v in vars_to_scalarize for v in z_vars + ): + for v in z_vars: + if ( + v not in dynamic_list_vars + and isinstance(object_variables.get(v), list) + and v not in vars_to_scalarize + ): + vars_to_scalarize.add(v) + changed = True + + if len(render_group.matrices) > 1: + if any( + any(v in vars_to_scalarize for v in mat) for mat in render_group.matrices + ): + for mat in render_group.matrices: + for v in mat: + if v not in dynamic_list_vars and isinstance( + object_variables.get(v), list + ): + vars_to_scalarize.add(v) + + explicit_matrix_and_zip_vars = set() + for mat in render_group.matrices: + explicit_matrix_and_zip_vars.update(mat) + for z_name, z_vars in render_group.zips.items(): + explicit_matrix_and_zip_vars.add(z_name) + explicit_matrix_and_zip_vars.update(z_vars) + + if any(v not in explicit_matrix_and_zip_vars for v in vars_to_scalarize): + for k, v in object_variables.items(): + if ( + k not in dynamic_list_vars + and isinstance(v, list) + and k not in explicit_matrix_and_zip_vars + ): + vars_to_scalarize.add(k) + + sub_vars = {} + for k, v in object_variables.items(): + if k in vars_to_scalarize or not isinstance(v, list): + if k not in dynamic_list_vars: + sub_vars[k] = v + + sub_group = RenderGroup(render_group.object, render_group.action) + sub_group.variables = sub_vars + + sub_zips = {} + for z_name, z_vars in render_group.zips.items(): + filtered_z = [v for v in z_vars if v in vars_to_scalarize] + if len(filtered_z) > 1: + sub_zips[z_name] = filtered_z + sub_group.zips = sub_zips + + sub_matrices = [] + for mat in render_group.matrices: + filtered_mat = [v for v in mat if v in vars_to_scalarize or v in sub_zips] + if len(filtered_mat) >= 1: + sub_matrices.append(filtered_mat) + sub_group.matrices = sub_matrices + sub_group.n_repeats = 1 + sub_group.used_variables = render_group.used_variables.union(vars_to_scalarize) + + for rendered_sub_vars, _ in self.render_objects( + sub_group, ignore_used=False, fatal=fatal + ): + seed_vars = object_variables.copy() + seed_vars.update(rendered_sub_vars) + yield seed_vars, ramble.repeats.Repeats() + return + # Expand zip and matrix members to allow indirections like # ``` # variables: diff --git a/lib/ramble/ramble/test/expander.py b/lib/ramble/ramble/test/expander.py index 74c3dfe127..2c4feba1d0 100644 --- a/lib/ramble/ramble/test/expander.py +++ b/lib/ramble/ramble/test/expander.py @@ -417,3 +417,32 @@ class StatusEnum(str, Enum): enum_expander = ramble.expander.Expander({"status": StatusEnum.SETUP}, None) assert enum_expander.expand_var("{status}") == "SETUP" + + +def test_pow2_range_function(): + from ramble.expander import pow2_range + + assert pow2_range(8) == [1, 2, 4, 8] + assert pow2_range(1, 16) == [1, 2, 4, 8, 16] + assert pow2_range(2, 16) == [2, 4, 8, 16] + assert pow2_range(1, 15) == [1, 2, 4, 8] + assert pow2_range(1, 16, inclusive=False) == [1, 2, 4, 8] + assert pow2_range(1, 1) == [1] + assert pow2_range(1, 1, inclusive=False) == [] + assert pow2_range(16, 4) == [] + assert pow2_range(0) == [] + assert pow2_range(-1, 8) == [] + assert pow2_range(3, 24) == [3, 6, 12, 24] + + +def test_pow2_range_in_expander(): + expander = ramble.expander.Expander({"max_nodes": "16"}, None) + + assert expander.expand_lists("pow2_range(1, 16)") == [1, 2, 4, 8, 16] + assert expander.expand_lists("power2_range(1, 16)") == [1, 2, 4, 8, 16] + + res = expander.expand_var("pow2_range(1, {max_nodes})", typed=True) + assert res == [1, 2, 4, 8, 16] + + res_alias = expander.expand_var("power2_range(1, {max_nodes})", typed=True) + assert res_alias == [1, 2, 4, 8, 16] diff --git a/lib/ramble/ramble/test/experiment_set.py b/lib/ramble/ramble/test/experiment_set.py index befee56e32..44e4af9e14 100644 --- a/lib/ramble/ramble/test/experiment_set.py +++ b/lib/ramble/ramble/test/experiment_set.py @@ -2149,3 +2149,838 @@ def test_modifiers_no_version_set_correctly(workspace_name, mock_modifiers): assert mod_def["mode"] in expected_modifier_modes expected_modifier_modes.remove(mod_def["mode"]) assert not expected_modifier_modes + + +def test_dynamic_range_in_experiment_variables(workspace_name): + workspace("create", workspace_name) + + assert workspace_name in workspace("list") + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + application_context = ramble.context.Context() + application_context.context_name = "basic" + application_context.variables = { + "app_var1": "1", + "app_var2": "2", + "processes_per_node": "2", + "mpi_command": "", + "batch_submit": "", + } + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + workload_context.variables = { + "wl_var1": "1", + "wl_var2": "2", + } + + experiment_context = ramble.context.Context() + experiment_context.context_name = "series_{n_nodes}" + experiment_context.variables = { + "max_nodes": "3", + "n_nodes": "range(1, {max_nodes})", + "n_ranks": "{n_nodes} * 2", + } + + exp_set.set_application_context(application_context) + exp_set.set_workload_context(workload_context) + exp_set.set_experiment_context(experiment_context) + exp_set.build_experiment_chains() + + assert "basic.test_wl.series_1" in exp_set.experiments + assert "basic.test_wl.series_2" in exp_set.experiments + assert len(exp_set.experiments) == 2 + assert exp_set.get_var_from_experiment("basic.test_wl.series_1", "{n_ranks}") == "2" + assert exp_set.get_var_from_experiment("basic.test_wl.series_2", "{n_ranks}") == "4" + + +def test_dynamic_range_in_workload_variables(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + application_context = ramble.context.Context() + application_context.context_name = "basic" + application_context.variables = { + "app_var1": "1", + "app_var2": "2", + "processes_per_node": "2", + "mpi_command": "", + "batch_submit": "", + } + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + workload_context.variables = { + "wl_var1": "1", + "wl_var2": "2", + "max_nodes": "4", + "n_nodes": "range(1, {max_nodes})", + } + + experiment_context = ramble.context.Context() + experiment_context.context_name = "series_{n_nodes}" + experiment_context.variables = { + "n_ranks": "{n_nodes} * {processes_per_node}", + } + + exp_set.set_application_context(application_context) + exp_set.set_workload_context(workload_context) + exp_set.set_experiment_context(experiment_context) + exp_set.build_experiment_chains() + + assert "basic.test_wl.series_1" in exp_set.experiments + assert "basic.test_wl.series_2" in exp_set.experiments + assert "basic.test_wl.series_3" in exp_set.experiments + assert len(exp_set.experiments) == 3 + assert exp_set.get_var_from_experiment("basic.test_wl.series_1", "{n_ranks}") == "2" + assert exp_set.get_var_from_experiment("basic.test_wl.series_3", "{n_ranks}") == "6" + + +def test_dynamic_range_with_exclude_where(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + application_context = ramble.context.Context() + application_context.context_name = "basic" + application_context.variables = { + "app_var1": "1", + "app_var2": "2", + "processes_per_node": "2", + "mpi_command": "", + "batch_submit": "", + } + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + workload_context.variables = { + "wl_var1": "1", + "wl_var2": "2", + } + + experiment_context = ramble.context.Context() + experiment_context.context_name = "series_{n_nodes}" + experiment_context.variables = { + "max_nodes": "5", + "n_nodes": "range(1, {max_nodes})", + "n_ranks": "{n_nodes} * 2", + } + experiment_context.exclude = {"where": ["'{n_nodes}' == '2'"]} + + exp_set.set_application_context(application_context) + exp_set.set_workload_context(workload_context) + exp_set.set_experiment_context(experiment_context) + exp_set.build_experiment_chains() + + assert "basic.test_wl.series_1" in exp_set.experiments + assert "basic.test_wl.series_2" not in exp_set.experiments + assert "basic.test_wl.series_3" in exp_set.experiments + assert "basic.test_wl.series_4" in exp_set.experiments + assert len(exp_set.experiments) == 3 + + +def test_dynamic_range_with_chained_experiments(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + application_context = ramble.context.Context() + application_context.context_name = "basic" + application_context.variables = { + "app_var1": "1", + "app_var2": "2", + "processes_per_node": "2", + "mpi_command": "", + "batch_submit": "", + } + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + workload_context.variables = { + "wl_var1": "1", + "wl_var2": "2", + } + + experiment1_context = ramble.context.Context() + experiment1_context.context_name = "test1" + experiment1_context.variables = {"n_ranks": "2"} + + experiment2_context = ramble.context.Context() + experiment2_context.context_name = "series_{n_nodes}" + experiment2_context.variables = { + "max_nodes": "3", + "n_nodes": "range(1, {max_nodes})", + "n_ranks": "{n_nodes} * 2", + } + experiment2_context.chained_experiments = [ + { + "name": "basic.test_wl.test1", + "order": "after_root", + "command": "{execute_experiment}", + "variables": {}, + }, + ] + + exp_set.set_application_context(application_context) + exp_set.set_workload_context(workload_context) + exp_set.set_experiment_context(experiment1_context) + exp_set.set_experiment_context(experiment2_context) + exp_set.build_experiment_chains() + + assert "basic.test_wl.series_1" in exp_set.experiments + assert "basic.test_wl.series_2" in exp_set.experiments + assert "basic.test_wl.test1" in exp_set.experiments + assert "basic.test_wl.series_1.chain.0.basic.test_wl.test1" in exp_set.chained_experiments + assert "basic.test_wl.series_2.chain.0.basic.test_wl.test1" in exp_set.chained_experiments + + +def test_dynamic_range_workspace_yaml(make_workspace_from_config): + test_config = """ +ramble: + variables: + mpi_command: '' + batch_submit: '' + processes_per_node: '1' + applications: + basic: + workloads: + test_wl: + experiments: + exp_{n_nodes}: + variables: + max_nodes: '3' + n_nodes: 'range(1, {max_nodes})' + n_ranks: '{n_nodes} * 2' +""" + ws, ws_name = make_workspace_from_config(test_config, activate=True) + exp_set = ws.build_experiment_set() + assert "basic.test_wl.exp_1" in exp_set.experiments + assert "basic.test_wl.exp_2" in exp_set.experiments + assert len(exp_set.experiments) == 2 + + +def test_dynamic_range_in_application_variables(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + application_context = ramble.context.Context() + application_context.context_name = "basic" + application_context.variables = { + "app_var1": "1", + "app_var2": "2", + "processes_per_node": "2", + "mpi_command": "", + "batch_submit": "", + "min_nodes": "1", + "max_nodes": "6", + "stride": "2", + "n_nodes": "range({min_nodes}, {max_nodes}, {stride})", + } + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + workload_context.variables = { + "wl_var1": "1", + "wl_var2": "2", + } + + experiment_context = ramble.context.Context() + experiment_context.context_name = "series_{n_nodes}" + experiment_context.variables = { + "n_ranks": "{n_nodes} * {processes_per_node}", + } + + exp_set.set_application_context(application_context) + exp_set.set_workload_context(workload_context) + exp_set.set_experiment_context(experiment_context) + exp_set.build_experiment_chains() + + assert "basic.test_wl.series_1" in exp_set.experiments + assert "basic.test_wl.series_3" in exp_set.experiments + assert "basic.test_wl.series_5" in exp_set.experiments + assert len(exp_set.experiments) == 3 + assert exp_set.get_var_from_experiment("basic.test_wl.series_1", "{n_ranks}") == "2" + assert exp_set.get_var_from_experiment("basic.test_wl.series_5", "{n_ranks}") == "10" + + +def test_dynamic_range_dry_run(make_workspace_from_config): + test_config = """ +ramble: + variables: + mpi_command: '' + batch_submit: '{execute_experiment}' + processes_per_node: '1' + applications: + basic: + workloads: + test_wl: + experiments: + exp_{n_nodes}: + variables: + max_nodes: '3' + n_nodes: 'range(1, {max_nodes})' + n_ranks: '{n_nodes} * 2' +""" + ws, ws_name = make_workspace_from_config(test_config, activate=True) + workspace("setup", "--dry-run", global_args=["-w", ws_name]) + + exp1_script = os.path.join( + ws.experiment_dir, "basic", "test_wl", "exp_1", "execute_experiment" + ) + exp2_script = os.path.join( + ws.experiment_dir, "basic", "test_wl", "exp_2", "execute_experiment" + ) + assert os.path.exists(exp1_script) + assert os.path.exists(exp2_script) + + +def test_chained_experiment_with_dynamic_range(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + app_context = ramble.context.Context() + app_context.context_name = "basic" + exp_set.set_application_context(app_context) + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + exp_set.set_workload_context(workload_context) + + experiment_context = ramble.context.Context() + experiment_context.context_name = "test_{n_nodes}" + experiment_context.variables = { + "max_nodes": "3", + "n_nodes": "range(1, {max_nodes})", + "n_ranks": "1", + "processes_per_node": "1", + } + rendered_instances = exp_set.render_experiment_set( + "basic", "test_wl", experiment_context, chained=True + ) + + assert len(rendered_instances) == 2 + assert "test_1" in exp_set.chained_experiments + assert "test_2" in exp_set.chained_experiments + assert "basic.test_wl.test_1" not in exp_set.experiments + assert "basic.test_wl.test_2" not in exp_set.experiments + + +def test_dynamic_range_with_repeats(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + app_context = ramble.context.Context() + app_context.context_name = "basic" + app_context.variables = { + "processes_per_node": "1", + "mpi_command": "", + "batch_submit": "", + } + exp_set.set_application_context(app_context) + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + exp_set.set_workload_context(workload_context) + + experiment_context = ramble.context.Context() + experiment_context.context_name = "test_{n_nodes}" + experiment_context.n_repeats = 2 + experiment_context.variables = { + "max_nodes": "3", + "n_nodes": "range(1, {max_nodes})", + "n_ranks": "1", + } + rendered = exp_set.set_experiment_context(experiment_context) + assert len(rendered) == 6 + assert "basic.test_wl.test_1" in exp_set.experiments + assert "basic.test_wl.test_1.1" in exp_set.experiments + assert "basic.test_wl.test_1.2" in exp_set.experiments + assert "basic.test_wl.test_2" in exp_set.experiments + assert "basic.test_wl.test_2.1" in exp_set.experiments + assert "basic.test_wl.test_2.2" in exp_set.experiments + + +def test_dynamic_range_with_whitespace(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + app_context = ramble.context.Context() + app_context.context_name = "basic" + app_context.variables = { + "processes_per_node": "1", + "mpi_command": "", + "batch_submit": "", + } + exp_set.set_application_context(app_context) + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + exp_set.set_workload_context(workload_context) + + experiment_context = ramble.context.Context() + experiment_context.context_name = "test_{n_nodes}" + experiment_context.variables = { + "max_nodes": "3", + "n_nodes": "range ( 1 , {max_nodes} )", + "n_ranks": "1", + } + rendered = exp_set.set_experiment_context(experiment_context) + assert len(rendered) == 2 + assert "basic.test_wl.test_1" in exp_set.experiments + assert "basic.test_wl.test_2" in exp_set.experiments + + +def test_dynamic_range_with_matrix(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + app_context = ramble.context.Context() + app_context.context_name = "basic" + app_context.variables = { + "processes_per_node": "1", + "mpi_command": "", + "batch_submit": "", + } + exp_set.set_application_context(app_context) + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + exp_set.set_workload_context(workload_context) + + experiment_context = ramble.context.Context() + experiment_context.context_name = "sweep_{nodes}_{cores}" + experiment_context.variables = { + "max_nodes": "3", + "nodes": "range(1, {max_nodes})", + "cores": [1, 2], + "n_ranks": "1", + } + experiment_context.matrices = [["nodes", "cores"]] + rendered = exp_set.set_experiment_context(experiment_context) + assert len(rendered) == 4 + assert "basic.test_wl.sweep_1_1" in exp_set.experiments + assert "basic.test_wl.sweep_1_2" in exp_set.experiments + assert "basic.test_wl.sweep_2_1" in exp_set.experiments + assert "basic.test_wl.sweep_2_2" in exp_set.experiments + + +def test_dynamic_range_with_two_ranges_in_matrix(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + app_context = ramble.context.Context() + app_context.context_name = "basic" + app_context.variables = { + "processes_per_node": "1", + "mpi_command": "", + "batch_submit": "", + } + exp_set.set_application_context(app_context) + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + exp_set.set_workload_context(workload_context) + + experiment_context = ramble.context.Context() + experiment_context.context_name = "sweep_{nodes}_{cores}" + experiment_context.variables = { + "max_nodes": "3", + "max_cores": "4", + "nodes": "range(1, {max_nodes})", + "cores": "range(1, {max_cores})", + "n_ranks": "1", + } + experiment_context.matrices = [["nodes", "cores"]] + rendered = exp_set.set_experiment_context(experiment_context) + assert len(rendered) == 6 + assert "basic.test_wl.sweep_1_1" in exp_set.experiments + assert "basic.test_wl.sweep_1_2" in exp_set.experiments + assert "basic.test_wl.sweep_1_3" in exp_set.experiments + assert "basic.test_wl.sweep_2_1" in exp_set.experiments + assert "basic.test_wl.sweep_2_2" in exp_set.experiments + assert "basic.test_wl.sweep_2_3" in exp_set.experiments + + +def test_dynamic_range_with_zips(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + app_context = ramble.context.Context() + app_context.context_name = "basic" + app_context.variables = { + "processes_per_node": "1", + "mpi_command": "", + "batch_submit": "", + } + exp_set.set_application_context(app_context) + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + exp_set.set_workload_context(workload_context) + + experiment_context = ramble.context.Context() + experiment_context.context_name = "sweep_{nodes}_{cores}" + experiment_context.variables = { + "max_nodes": "3", + "nodes": "range(1, {max_nodes})", + "cores": "range(10, 10 + {max_nodes} - 1)", + "n_ranks": "1", + } + experiment_context.zips = {"my_zip": ["nodes", "cores"]} + rendered = exp_set.set_experiment_context(experiment_context) + assert len(rendered) == 2 + assert "basic.test_wl.sweep_1_10" in exp_set.experiments + assert "basic.test_wl.sweep_2_11" in exp_set.experiments + + +def test_dynamic_range_with_static_vector_implicit_zip(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + app_context = ramble.context.Context() + app_context.context_name = "basic" + app_context.variables = { + "processes_per_node": "1", + "mpi_command": "", + "batch_submit": "", + } + exp_set.set_application_context(app_context) + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + exp_set.set_workload_context(workload_context) + + experiment_context = ramble.context.Context() + experiment_context.context_name = "sweep_{nodes}_{cores}" + experiment_context.variables = { + "max_nodes": "3", + "nodes": "range(1, {max_nodes})", + "cores": [10, 20], + "n_ranks": "1", + } + rendered = exp_set.set_experiment_context(experiment_context) + assert len(rendered) == 2 + assert "basic.test_wl.sweep_1_10" in exp_set.experiments + assert "basic.test_wl.sweep_2_20" in exp_set.experiments + + +def test_dynamic_custom_list_function(workspace_name, monkeypatch): + def custom_seq(start, stop): + return list(range(start, stop)) + + monkeypatch.setitem(ramble.expander.supported_list_function_pointers, "custom_seq", custom_seq) + + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + app_context = ramble.context.Context() + app_context.context_name = "basic" + app_context.variables = { + "processes_per_node": "1", + "mpi_command": "", + "batch_submit": "", + } + exp_set.set_application_context(app_context) + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + exp_set.set_workload_context(workload_context) + + experiment_context = ramble.context.Context() + experiment_context.context_name = "sweep_{nodes}" + experiment_context.variables = { + "max_nodes": "4", + "nodes": "custom_seq(1, {max_nodes})", + "n_ranks": "1", + } + rendered = exp_set.set_experiment_context(experiment_context) + assert len(rendered) == 3 + assert "basic.test_wl.sweep_1" in exp_set.experiments + assert "basic.test_wl.sweep_2" in exp_set.experiments + assert "basic.test_wl.sweep_3" in exp_set.experiments + + +def test_dynamic_pow2_range_in_experiment_set(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + app_context = ramble.context.Context() + app_context.context_name = "basic" + app_context.variables = { + "processes_per_node": "1", + "mpi_command": "", + "batch_submit": "", + } + exp_set.set_application_context(app_context) + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + exp_set.set_workload_context(workload_context) + + experiment_context = ramble.context.Context() + experiment_context.context_name = "sweep_{nodes}" + experiment_context.variables = { + "max_nodes": "16", + "nodes": "pow2_range(1, {max_nodes})", + "n_ranks": "1", + } + rendered = exp_set.set_experiment_context(experiment_context) + assert len(rendered) == 5 + assert "basic.test_wl.sweep_1" in exp_set.experiments + assert "basic.test_wl.sweep_2" in exp_set.experiments + assert "basic.test_wl.sweep_4" in exp_set.experiments + assert "basic.test_wl.sweep_8" in exp_set.experiments + assert "basic.test_wl.sweep_16" in exp_set.experiments + + +def test_dynamic_range_dependent_on_vector(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + app_context = ramble.context.Context() + app_context.context_name = "basic" + app_context.variables = { + "processes_per_node": "1", + "mpi_command": "", + "batch_submit": "", + } + exp_set.set_application_context(app_context) + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + exp_set.set_workload_context(workload_context) + + experiment_context = ramble.context.Context() + experiment_context.context_name = "sweep_{max_nodes}_{nodes}" + experiment_context.variables = { + "max_nodes": [2, 4], + "nodes": "range(1, {max_nodes} + 1)", + "n_ranks": "1", + } + rendered = exp_set.set_experiment_context(experiment_context) + assert len(rendered) == 6 + assert "basic.test_wl.sweep_2_1" in exp_set.experiments + assert "basic.test_wl.sweep_2_2" in exp_set.experiments + assert "basic.test_wl.sweep_4_1" in exp_set.experiments + assert "basic.test_wl.sweep_4_2" in exp_set.experiments + assert "basic.test_wl.sweep_4_3" in exp_set.experiments + assert "basic.test_wl.sweep_4_4" in exp_set.experiments + + +def test_dynamic_range_dependent_on_vector_with_matrix(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + app_context = ramble.context.Context() + app_context.context_name = "basic" + app_context.variables = { + "processes_per_node": "1", + "mpi_command": "", + "batch_submit": "", + } + exp_set.set_application_context(app_context) + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + exp_set.set_workload_context(workload_context) + + experiment_context = ramble.context.Context() + experiment_context.context_name = "sweep_{max_nodes}_{nodes}" + experiment_context.variables = { + "max_nodes": [2, 4], + "nodes": "range(1, {max_nodes} + 1)", + "n_ranks": "1", + } + experiment_context.matrices = [["max_nodes", "nodes"]] + rendered = exp_set.set_experiment_context(experiment_context) + assert len(rendered) == 6 + assert "basic.test_wl.sweep_2_1" in exp_set.experiments + assert "basic.test_wl.sweep_2_2" in exp_set.experiments + assert "basic.test_wl.sweep_4_1" in exp_set.experiments + assert "basic.test_wl.sweep_4_2" in exp_set.experiments + assert "basic.test_wl.sweep_4_3" in exp_set.experiments + assert "basic.test_wl.sweep_4_4" in exp_set.experiments + + +def test_dynamic_pow2_range_dependent_on_vector(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + app_context = ramble.context.Context() + app_context.context_name = "basic" + app_context.variables = { + "processes_per_node": "1", + "mpi_command": "", + "batch_submit": "", + } + exp_set.set_application_context(app_context) + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + exp_set.set_workload_context(workload_context) + + experiment_context = ramble.context.Context() + experiment_context.context_name = "sweep_{max_nodes}_{nodes}" + experiment_context.variables = { + "max_nodes": [4, 8], + "nodes": "pow2_range(1, {max_nodes})", + "n_ranks": "1", + } + rendered = exp_set.set_experiment_context(experiment_context) + assert len(rendered) == 7 + assert "basic.test_wl.sweep_4_1" in exp_set.experiments + assert "basic.test_wl.sweep_4_2" in exp_set.experiments + assert "basic.test_wl.sweep_4_4" in exp_set.experiments + assert "basic.test_wl.sweep_8_1" in exp_set.experiments + assert "basic.test_wl.sweep_8_2" in exp_set.experiments + assert "basic.test_wl.sweep_8_4" in exp_set.experiments + assert "basic.test_wl.sweep_8_8" in exp_set.experiments + + +def test_chained_multi_level_dependent_ranges(workspace_name): + workspace("create", workspace_name) + + with ramble.workspace.read(workspace_name) as ws: + exp_set = ramble.experiment_set.ExperimentSet(ws) + + app_context = ramble.context.Context() + app_context.context_name = "basic" + app_context.variables = { + "processes_per_node": "1", + "mpi_command": "", + "batch_submit": "", + } + exp_set.set_application_context(app_context) + + workload_context = ramble.context.Context() + workload_context.context_name = "test_wl" + exp_set.set_workload_context(workload_context) + + experiment_context = ramble.context.Context() + experiment_context.context_name = "sweep_{foo}_{bar}_{baz}" + experiment_context.variables = { + "foo": [1, 2, 4], + "bar": "range(1, {foo} + 1)", + "baz": "range(1, {bar} + 1)", + "n_ranks": "1", + } + rendered = exp_set.set_experiment_context(experiment_context) + assert len(rendered) == 14 + assert "basic.test_wl.sweep_1_1_1" in exp_set.experiments + assert "basic.test_wl.sweep_2_1_1" in exp_set.experiments + assert "basic.test_wl.sweep_2_2_1" in exp_set.experiments + assert "basic.test_wl.sweep_2_2_2" in exp_set.experiments + assert "basic.test_wl.sweep_4_1_1" in exp_set.experiments + assert "basic.test_wl.sweep_4_2_1" in exp_set.experiments + assert "basic.test_wl.sweep_4_2_2" in exp_set.experiments + assert "basic.test_wl.sweep_4_3_1" in exp_set.experiments + assert "basic.test_wl.sweep_4_3_2" in exp_set.experiments + assert "basic.test_wl.sweep_4_3_3" in exp_set.experiments + assert "basic.test_wl.sweep_4_4_1" in exp_set.experiments + assert "basic.test_wl.sweep_4_4_2" in exp_set.experiments + assert "basic.test_wl.sweep_4_4_3" in exp_set.experiments + assert "basic.test_wl.sweep_4_4_4" in exp_set.experiments + + +def test_dynamic_range_indirect_vector_dependency_with_mpi_matrix(make_workspace_from_config): + test_config = """ +ramble: + variables: + processes_per_node: 4 + mpi_command: 'mpirun -n {n_ranks}' + batch_submit: '{execute_experiment}' + vm_family: [h4d] + mpi_provider: [tcp, rxm] + Alltoall_info: + max_algo: 4 + Alltoallv_info: + max_algo: 2 + applications: + basic: + workloads: + test_wl2: + experiments: + '{vm_family}-{mpi_provider}-{n_nodes}node-{collective_type}-{algo_setting}': + variables: + n_nodes: [16, 32] + max_algo: '{collective_type}_info["max_algo"]' + algo_setting: 'range(1, {max_algo})' + n_ranks: '{processes_per_node}*{n_nodes}' + collective_type: [Alltoall, Alltoallv] + matrix: + - n_nodes + - algo_setting + - collective_type + - vm_family + - mpi_provider + exclude: + where: + - '{algo_setting} > {max_algo}' +""" + ws, _ = make_workspace_from_config(test_config, activate=True) + exp_set = ws.build_experiment_set() + # Alltoall has range(1, 4) -> [1, 2, 3] (3 algos) * 2 n_nodes * 1 vm_family * 2 mpi = 12 + # Alltoallv has range(1, 2) -> [1] (1 algo) * 2 n_nodes * 1 vm_family * 2 mpi = 4 + # Total = 16 experiments + assert len(exp_set.experiments) == 16 + assert "basic.test_wl2.h4d-tcp-16node-Alltoall-1" in exp_set.experiments + assert "basic.test_wl2.h4d-rxm-32node-Alltoall-3" in exp_set.experiments + assert "basic.test_wl2.h4d-tcp-16node-Alltoallv-1" in exp_set.experiments + assert "basic.test_wl2.h4d-tcp-16node-Alltoallv-2" not in exp_set.experiments + + +def test_dynamic_range_on_n_nodes_mpi(make_workspace_from_config): + test_config = """ +ramble: + variables: + processes_per_node: 4 + mpi_command: 'mpirun -n {n_ranks}' + batch_submit: '{execute_experiment}' + applications: + basic: + workloads: + test_wl2: + experiments: + node_sweep_{n_nodes}: + variables: + max_nodes: 4 + n_nodes: 'range(1, {max_nodes})' +""" + ws, _ = make_workspace_from_config(test_config, activate=True) + exp_set = ws.build_experiment_set() + assert len(exp_set.experiments) == 3 + assert exp_set.experiments["basic.test_wl2.node_sweep_1"].variables["n_ranks"] == "4" + assert exp_set.experiments["basic.test_wl2.node_sweep_2"].variables["n_ranks"] == "8" + assert exp_set.experiments["basic.test_wl2.node_sweep_3"].variables["n_ranks"] == "12" diff --git a/var/ramble/repos/builtin/base_classes/application-base/base_class.py b/var/ramble/repos/builtin/base_classes/application-base/base_class.py index 7ca66d7c69..b20eb8ff52 100644 --- a/var/ramble/repos/builtin/base_classes/application-base/base_class.py +++ b/var/ramble/repos/builtin/base_classes/application-base/base_class.py @@ -259,6 +259,7 @@ def __init__(self, file_path): # A dict storing fom values, currently it only stores inmem FOMs self._fom_map = {} self._template_paths_defined = False + self._dynamic_range_variables = None # Ensure we always have the application name, and this is never empty self._file_path = file_path @@ -730,6 +731,7 @@ def set_variables_and_variants( self.expander = ramble.expander.Expander( self.variables, self.experiment_set ) + self._dynamic_range_variables = None # Set application version or use preferred version if none specified _, _, maybe_version = self.expander.application_spec.partition("@") @@ -892,8 +894,9 @@ def non_reserved_variables( for key in remove_keys: cleaned_variables.pop(key, None) - for template_name, _ in workspace.all_templates(): - cleaned_variables.pop(template_name, None) + if workspace: + for template_name, _ in workspace.all_templates(): + cleaned_variables.pop(template_name, None) for _, tpl_configs in self._object_templates(): for tpl_config in tpl_configs: @@ -901,6 +904,143 @@ def non_reserved_variables( return cleaned_variables + @property + def has_dynamic_range_variables(self) -> bool: + """Check if any variables define dynamic ranges that evaluate to lists or depend on vectors.""" + if self.dynamic_range_variables(): + return True + dyn_vars = { + var: val + for var, val in self.variables.items() + if ramble.expander.is_dynamic_list_expression(val) + } + if not dyn_vars: + return False + return bool( + ramble.expander.find_dependent_vector_vars( + dyn_vars, self.variables, self.expander + ) + ) + + def dynamic_range_variables(self) -> Dict[str, list]: + """Identify any variables defined as dynamic ranges that can now be evaluated into lists. + + Returns: + dict: Mapping of variable name to evaluated list + """ + if self._dynamic_range_variables is None: + ranges = {} + for var, val in self.variables.items(): + if ramble.expander.is_dynamic_list_expression(val): + try: + expanded = self.expander.expand_var(val, typed=True) + if isinstance(expanded, list): + ranges[var] = expanded + except Exception: + pass + self._dynamic_range_variables = ranges + return self._dynamic_range_variables + + def render_range_experiments( + self, + experiment_context, + warn_validation=True, + die_on_validate_error=True, + chained=False, + ) -> list: + """Render range experiments using the finalized variables. + + Args: + experiment_context (ramble.context.Context): Context object for the experiment + warn_validation (bool): Whether validation warnings should print + die_on_validate_error (bool): Whether validation errors should be fatal + chained (bool): Whether the experiments are chained experiments or not + + Returns: + list: List of application instances from the rendered set of experiments + """ + ranges = self.dynamic_range_variables() + sub_context = copy.deepcopy(experiment_context) + if not ranges: + dyn_vars = { + var: val + for var, val in self.variables.items() + if ramble.expander.is_dynamic_list_expression(val) + } + if dyn_vars: + added = False + for var_name, var_val in self.variables.items(): + if ( + var_name not in sub_context.variables + and var_val is not None + ): + sub_context.variables[var_name] = var_val + added = True + if added: + return self.experiment_set.set_experiment_context( + sub_context, + warn_validation=warn_validation, + die_on_validate_error=die_on_validate_error, + chained=chained, + ) + return [] + + matrix_and_zip_vars = set() + if sub_context.matrices: + for mat in sub_context.matrices: + matrix_and_zip_vars.update(mat) + if sub_context.zips: + for z_vars in sub_context.zips.values(): + matrix_and_zip_vars.update(z_vars) + + # Propagate any variables that were resolved to scalars in this seed instance + for var_name, var_val in self.variables.items(): + if ( + var_name in sub_context.variables + or var_name in matrix_and_zip_vars + ) and not isinstance(var_val, list): + sub_context.variables[var_name] = var_val + + for range_var, range_list in ranges.items(): + sub_context.variables[range_var] = range_list + + effective_vars = self.variables.copy() + effective_vars.update(sub_context.variables) + + # Remove any variables from zips that became scalars in this seed + if sub_context.zips: + new_zips = {} + for z_name, z_vars in sub_context.zips.items(): + new_z = [ + v + for v in z_vars + if isinstance(effective_vars.get(v), list) + ] + if len(new_z) > 1: + new_zips[z_name] = new_z + sub_context.zips = new_zips + + # Remove any variables from matrices that became scalars in this seed + if sub_context.matrices: + new_matrices = [] + for mat in sub_context.matrices: + new_mat = [ + v + for v in mat + if isinstance(effective_vars.get(v), list) + or (sub_context.zips and v in sub_context.zips) + ] + if len(new_mat) >= 1: + new_matrices.append(new_mat) + sub_context.matrices = new_matrices + + return self.experiment_set.set_experiment_context( + sub_context, + warn_validation=warn_validation, + die_on_validate_error=die_on_validate_error, + chained=chained, + ) + def register_missing_command_variable(self, var): """Register a missing command variable, so we can report it later in the correct log file. @@ -2228,11 +2368,14 @@ def _define_commands(self, exec_graph=None, success_list=None): n_nodes = self.expander.expand_var_name( self.keywords.n_nodes ) - n_nodes = ( - 1 - if n_nodes in ("{n_nodes}", None, "") - else int(n_nodes) - ) + try: + n_nodes = ( + 1 + if n_nodes in ("{n_nodes}", None, "") + else int(n_nodes) + ) + except (ValueError, TypeError): + n_nodes = 1 if not raw_mpi_cmd and n_nodes > 1: logger.warn( f"Command {cmd_conf.name} requires a non-empty `mpi_command` " @@ -4677,9 +4820,14 @@ def define_mpi_vars(): value = None # If two variables are defined, use the formula to compute the missing ones. if len(mpi_vars_defined) >= 2: - value = self.expander.expand_var( + val = self.expander.expand_var( formula, allow_passthrough=False ) + try: + int(val) + value = val + except (ValueError, TypeError): + value = None # If there is not enough information to use the formulas, or they are not required. # Set missing vars to 0 elif not mpi_required: