# 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
import ramble.expander
import ramble.variants
[docs]
def exp_dict():
return {
"application_name": "foo",
"application_spec": "foo@1.2",
"workload_name": "bar",
"experiment_name": "baz",
"application_input_dir": "/workspace/inputs/foo",
"workload_input_dir": "/workspace/inputs/foo/bar",
"application_run_dir": "/workspace/experiments/foo",
"workload_run_dir": "/workspace/experiments/foo/bar",
"experiment_run_dir": "/workspace/experiments/foo/bar/baz",
"env_name": "spack_foo.bar",
"n_ranks": "4",
"processes_per_node": "2",
"n_nodes": "2",
"cores_per_node": "16",
"n_threads": "4",
"var1": "{var2}",
"var2": "{var3}",
"var3": "3",
"decimal.06.var": "foo",
"size": '"0000.96"', # Escaped as a string
"test_mask": '"0x0"',
"max_len": 9,
"test_dict": {"test_key1": "test_val1", "test_key2": "test_val2"},
"experiment_index": 5,
"var4": "a-b",
"test::var": "value",
}
[docs]
def build_variant_set():
variant_set = ramble.variants.VariantSet()
variants = {
"test_variant": "defined",
"package_manager": "spack",
"workflow_manager": "slurm",
"compiler": "gcc@14.2.0",
"bool_true": True,
"bool_false": False,
}
multi_value_variants = [
("package_manager_family", "spack"),
("workflow_manager_family", "slurm"),
]
for name, value in variants.items():
variant_set.default_variant(name, value)
for name, value in multi_value_variants:
variant_set.multi_value_variant(name, value)
return variant_set
[docs]
@pytest.mark.parametrize(
"input,output,no_expand_vars,passes",
[
("{var1}", "3", set(), 1),
("{var2}", "3", set(), 1),
("{var3}", "3", set(), 1),
("{application_name}", "foo", set(), 1),
("{n_nodes}", "2", set(), 1),
("{processes_per_node}", "2", set(), 1),
("{n_nodes}*{processes_per_node}", "4", set(), 1),
("2**4", "16", set(), 1),
("{((((16-10+2)/4)**2)*4)}", "16.0", set(), 1),
("gromacs +blas", "gromacs +blas", set(), 1),
("range(0, 5)", "[0, 1, 2, 3, 4]", set(), 1),
("{decimal.06.var}", "foo", set(), 1),
("{}", "{}", set(), 1),
("{{n_ranks}+2}", "6", set(), 1),
("{{n_ranks}*{var{processes_per_node}}:05d}", "00012", set(), 1),
("{{n_ranks}-1}", "3", set(), 1),
("{{{n_ranks}/2}:0.0f}", "2", set(), 1),
("{size}", "0000.96", {"size"}, 1),
("CPU(s)", "CPU(s)", set(), 1),
("str(1.5)", "1.5", set(), 1),
("int(1.5)", "1", set(), 1),
("float(1.5)", "1.5", set(), 1),
("ceil(0.6)", "1", set(), 1),
("floor(0.6)", "0", set(), 1),
("max(1, 5)", "5", set(), 1),
("min(1, 5)", "1", set(), 1),
('simplify_str("a.b_c")', "a-b-c", set(), 1),
(r"\{experiment_name\}", "{experiment_name}", set(), 1),
(r"\{experiment_name\}", "baz", set(), 2),
(r"{\{experiment_name\}}", "{{experiment_name}}", set(), 1),
(r"\\{experiment_name\\}", r"\{experiment_name\}", set(), 1),
(r"\\{experiment_name\\}", "{experiment_name}", set(), 2),
(r"\\{experiment_name\\}", "baz", set(), 3),
('"2.1.1" in ["2.1.1", "3.1.1", "4.2.1"]', "True", set(), 1),
('"2.1.2" in ["2.1.1", "3.1.1", "4.2.1"]', "False", set(), 1),
('"3.1.2" not in ["2.1.1", "3.1.1", "4.2.1"]', "True", set(), 1),
('"2.1.1" not in ["2.1.1", "3.1.1", "4.2.1"]', "False", set(), 1),
("{test_mask}", "0x0", {"test_mask"}, 1),
('re_search(r"bz$", {experiment_name})', "False", set(), 1),
('re_search("o+\\\\.b", {env_name})', "True", set(), 1),
('"foo" in "{env_name}"', "True", set(), 1),
('"c" in "{experiment_name}"', "False", set(), 1),
('"foo123" not in "{env_name}"', "True", set(), 1),
('"foo" not in "{env_name}"', "False", set(), 1),
('"{env_name}"[:{max_len}:1]', "spack_foo", set(), 1),
("not_a_slice[0]", "not_a_slice[0]", set(), 1),
("not_a_valid_slice[0:a]", "not_a_valid_slice[0:a]", set(), 1),
("{test_dict}", "{'test_key1': 'test_val1', 'test_key2': 'test_val2'}", set(), 1),
("{test_dict['test_key1']}", "test_val1", set(), 1),
("{test_dict['test_key2']}", "test_val2", set(), 1),
("{test_dict['missing_key']}", "{test_dict['missing_key']}", set(), 1),
("{test_dict[None]}", "{test_dict[None]}", set(), 1),
("maybe(env_name, foo)", "spack_foo.bar", set(), 1),
("maybe(not_a_var, foo)", "foo", set(), 1),
("maybe(not_a_var)", "", set(), 1),
("2.1.a", "2.1.a", set(), 1),
("{'key1':'val1'}", "{'key1':'val1'}", set(), 1),
("{'key1': 'val1'}", "{'key1': 'val1'}", set(), 1),
("{'key1': 'val1', 'key2': 'val2'}", "{'key1': 'val1', 'key2': 'val2'}", set(), 1),
("${job_id}", "${job_id}", set(), 1),
("${job_id:-}", "${job_id:-}", set(), 1),
("join_str(range(0, 5, 2), ': ')", "0: 2: 4", set(), 1),
("join_str(range(0, {cores_per_node}, {n_threads}))", "0,4,8,12", set(), 1),
('int("0b101110", 2)', "46", set(), 1),
("0b10 & 0b01", "0", set(), 1),
("0b10 | 0b01", "3", set(), 1),
("0b10 ^ 0b01", "3", set(), 1),
("~0", "-1", set(), 1),
("~2", "-3", set(), 1),
("0b10 << 1", "4", set(), 1),
("0b10 >> 1", "1", set(), 1),
# Can be a handy way to select experiments to run
("(1 << {experiment_index} & 0b1011010) == 0", "True", set(), 1),
("$HOSTNAME", "$HOSTNAME", set(), 1),
("${HOSTNAME}", "${HOSTNAME}", set(), 1),
("log2(8)", "3.0", set(), 1),
("log10(100)", "2.0", set(), 1),
("sqrt(16)", "4.0", set(), 1),
# Can also reference functions available in the math module directly
("math_sqrt(64)", "8.0", set(), 1),
("math_log(9, 3)", "2.0", set(), 1),
("math_not_exist(1)", "math_not_exist(1)", set(), 1),
("str_upper('foo')", "FOO", set(), 1),
("str_lower('FOO')", "foo", set(), 1),
("str_upper(foobar)", "FOOBAR", set(), 1),
("str_capitalize('foo')", "Foo", set(), 1),
("str_lstrip('AAAbbb', 'A')", "bbb", set(), 1),
("str_join('.', str_split('a b c 1'))", "a.b.c.1", set(), 1),
("str_no_such_method('a')", "str_no_such_method('a')", set(), 1),
("replace('abc', 'a', 'd')", "dbc", set(), 1),
("replace('{application_name}', 'f', 'F')", "Foo", set(), 1),
("str_upper({var4})", "A-B", set(), 1),
("{test::var}", "value", set(), 1),
("{n_ranks:05d}", "00004", set(), 1),
("1_01", "1_01", set(), 1),
("0x10", "0x10", set(), 1),
("{application_name}-1_01", "foo-1_01", set(), 1),
],
)
def test_expansions(input, output, no_expand_vars, passes):
expansion_vars = exp_dict()
expander = ramble.expander.Expander(expansion_vars, None, no_expand_vars=no_expand_vars)
step_input = input
for _ in range(passes):
step_input = expander.expand_var(step_input)
final_output = step_input
assert final_output == output
[docs]
@pytest.mark.parametrize(
"input,output,no_expand_vars,passes",
[
("{var1}", 3, set(), 1),
("{var2}", 3, set(), 1),
("{var3}", 3, set(), 1),
("{application_name}", "foo", set(), 1),
("{n_nodes}", 2, set(), 1),
("{processes_per_node}", 2, set(), 1),
("{n_nodes}*{processes_per_node}", 4, set(), 1),
("2**4", 16, set(), 1),
("{((((16-10+2)/4)**2)*4)}", 16.0, set(), 1),
('"gromacs +blas"', "gromacs +blas", set(), 1),
("range(0, 5)", [0, 1, 2, 3, 4], set(), 1),
("{decimal.06.var}", "foo", set(), 1),
("{}", {}, set(), 1),
("{{n_ranks}+2}", 6, set(), 1),
("{{n_ranks}*{var{processes_per_node}}:05d}", "00012", set(), 1),
("{{n_ranks}-1}", 3, set(), 1),
("{{{n_ranks}/2}:0.0f}", 2, set(), 1),
("{size}", 0.96, {"size"}, 1),
("CPU(s)", "CPU(s)", set(), 1),
("str(1.5)", 1.5, set(), 1),
("int(1.5)", 1, set(), 1),
("float(1.5)", 1.5, set(), 1),
("ceil(0.6)", 1, set(), 1),
("floor(0.6)", 0, set(), 1),
("max(1, 5)", 5, set(), 1),
("min(1, 5)", 1, set(), 1),
('simplify_str("a.b_c")', "a-b-c", set(), 1),
(r"\{experiment_name\}", "{experiment_name}", set(), 1),
(r"\{experiment_name\}", "baz", set(), 2),
(r"{\{experiment_name\}}", "{{experiment_name}}", set(), 1),
(r"\\{experiment_name\\}", r"\{experiment_name\}", set(), 1),
(r"\\{experiment_name\\}", "{experiment_name}", set(), 2),
(r"\\{experiment_name\\}", "baz", set(), 3),
('"2.1.1" in ["2.1.1", "3.1.1", "4.2.1"]', True, set(), 1),
('"2.1.2" in ["2.1.1", "3.1.1", "4.2.1"]', False, set(), 1),
("{test_mask}", 0, {"test_mask"}, 1),
("{var3} // {processes_per_node}", 1, set(), 1),
("-5 % 3", 1, set(), 1),
("2 and 1", 1, set(), 1),
("2 or 1", 2, set(), 1),
("randrange(2, 3, 1)", 2, set(), 1),
("randint(3, 3)", 3, set(), 1),
("~0", -1, set(), 1),
("~2", -3, set(), 1),
],
)
def test_typed_expansions(input, output, no_expand_vars, passes):
expansion_vars = exp_dict()
expander = ramble.expander.Expander(expansion_vars, None, no_expand_vars=no_expand_vars)
step_input = input
for _ in range(passes):
step_input = expander.expand_var(step_input, typed=True)
final_output = step_input
assert final_output == output
[docs]
@pytest.mark.parametrize(
"input,output",
[
("application_name", "foo"),
("workload_name", "bar"),
("experiment_name", "baz"),
("var1", "3"),
("var2", "3"),
("var3", "3"),
],
)
def test_expand_var_name(input, output):
expansion_vars = exp_dict()
expander = ramble.expander.Expander(expansion_vars, None)
assert expander.expand_var_name(input) == output
[docs]
def test_expansion_namespaces():
expansion_vars = exp_dict()
expander = ramble.expander.Expander(expansion_vars, None)
assert expander.application_name == "foo"
assert expander.application_namespace == "foo@1.2"
assert expander.workload_namespace == "foo@1.2.bar"
assert expander.experiment_namespace == "foo@1.2.bar.baz"
[docs]
@pytest.mark.parametrize(
"input_list,output",
[
(["package_manager=spack"], True),
(["test_variant=defined", "package_manager=spack"], True),
(["test_variant=undefined", "package_manager=spack"], False),
(["test_variant=undefined", "package_manager=spack", "workflow_manager=slurm"], False),
(["test_variant=defined", "package_manager=spack", "workflow_manager=slurm"], True),
(["compiler=gcc@14.2.0"], True),
(["+bool_true"], True),
(["bool_true=True"], True),
(["bool_true=true"], True),
(["bool_true=False"], False),
(["bool_true=false"], False),
(["~bool_true"], False),
(["~bool_false"], True),
(["bool_false=False"], True),
(["bool_false=false"], True),
(["bool_false=True"], False),
(["bool_false=true"], False),
(["+bool_false"], False),
],
)
def test_satisfies_works(input_list, output):
variants = build_variant_set()
expansion_vars = exp_dict()
expander = ramble.expander.Expander(expansion_vars, None)
satisfied = expander.satisfies(input_list, variant_set=variants)
assert satisfied == output
[docs]
def test_variable_mutation():
"""Test that variable changes are reflected in expansions"""
expansion_vars = {"my_dict": {"key": 1}}
expander = ramble.expander.Expander(expansion_vars, None)
in_str = "{my_dict['key']}"
res1 = expander.expand_var(in_str)
assert res1 == "1"
# sneaky mutation
expander._variables["my_dict"]["key"] = 2
res2 = expander.expand_var(in_str)
assert res2 == "2"