Add 'when' to input_file directive

Author: dapomeroy

lib/ramble/ramble/application.py

@@ -1349,19 +1349,41 @@ def _inputs_and_fetchers(self, workload=None):
 
         workload_names = [workload] if workload else self.workloads.keys()
 
+        # Batch 'when' evaluation to avoid repeat expander calls
+        when_satisfied = set()
+        for when_set in self.inputs.keys():
+            if self.expander.satisfies(when_set, variant_set=self.object_variants):
+                when_satisfied.add(when_set)
+
         inputs = {}
         for workload_name in workload_names:
             workload = self.workloads[workload_name]
 
             for input_file in workload.inputs:
-                if input_file not in self.inputs:
+                inputs_found = 0
+                active_inputs = 0
+                input_conf = {}
+                for when_set, app_inputs in self.inputs.items():
+                    if input_file in app_inputs:
+                        inputs_found += 1
+                        if when_set in when_satisfied:
+                            active_inputs += 1
+                            input_conf = app_inputs[input_file].copy()
+
+                if not inputs_found:
                     logger.die(
                         f"Workload {workload_name} references a non-existent input file "
                         f"{input_file}.\n"
                         f"Make sure this input file is defined before using it in a workload."
                     )
-
-                input_conf = self.inputs[input_file].copy()
+                if active_inputs == 0:
+                    logger.debug(f"Skipping input {input_file}. `When` conditions not satisfied.")
+                    continue
+                elif active_inputs > 1:
+                    logger.die(
+                        f"Input files {input_file} are defined with overlapping 'when' "
+                        f"conditions. Make sure that conditions are mutually exclusive."
+                    )
 
                 # Expand input value as it may be a var
                 expanded_url = self.expander.expand_var(input_conf["url"])

lib/ramble/ramble/language/application_language.py

@@ -166,6 +166,7 @@ def input_file(
     sha256=None,
     extension=None,
     expand=True,
+    when=None,
     **kwargs,
 ):
     """Adds an input file definition to this application
@@ -184,16 +185,23 @@ def input_file(
         extension (str): Optiona, the extension to use for the input, if it isn't part of the
                               file name.
         expand (bool): Optional. Whether the input should be expanded or not. Defaults to True
+        when (list | None): List of when conditions to apply to directive
     """
 
     def _execute_input_file(app):
-        app.inputs[name] = {
+        when_list = ramble.language.language_helpers.build_when_list(when, app, name, "input_file")
+        when_set = frozenset(when_list)
+        if when_set not in app.inputs:
+            app.inputs[when_set] = {}
+
+        app.inputs[when_set][name] = {
             "url": url,
             "description": description,
             "target_dir": target_dir,
             "sha256": sha256,
             "extension": extension,
             "expand": expand,
+            "when": when_list,
         }
 
     return _execute_input_file

lib/ramble/ramble/language/language_base.py

@@ -239,6 +239,7 @@ def _wrapper(*args, **_kwargs):
                         "package_manager_variable",
                         "workflow_manager_variable",
                         "executable",
+                        "input_file",
                     ]:
                         msg = (
                             'directive "{0}" cannot be used within a "when"'

lib/ramble/ramble/test/application.py

@@ -113,9 +113,9 @@ def test_basic_app(mutable_mock_apps_repo):
     assert fom_conf["group_name"] == "test"
     assert fom_conf["units"] == "s"
 
-    assert "input" in basic_inst.inputs
-    assert basic_inst.inputs["input"]["url"] == "file:///tmp/test_file.log"
-    assert basic_inst.inputs["input"]["description"] == "Not a file"
+    assert "input" in basic_inst.inputs[_FS]
+    assert basic_inst.inputs[_FS]["input"]["url"] == "file:///tmp/test_file.log"
+    assert basic_inst.inputs[_FS]["input"]["description"] == "Not a file"
 
 
 @pytest.mark.parametrize("app_name", ["basic", "zlib"])
@@ -399,7 +399,7 @@ def test_set_default_experiment_variables(mutable_mock_apps_repo):
 
     executable_application_instance.internals = {}
 
-    executable_application_instance.inputs = {"input": {"target_dir": "."}}
+    executable_application_instance.inputs[_FS] = {"input": {"target_dir": "."}}
     executable_application_instance.variables = {}
 
     executable_application_instance._set_default_experiment_variables()
@@ -424,7 +424,7 @@ def test_define_commands(mutable_mock_apps_repo):
 
     executable_application_instance.internals = {}
 
-    executable_application_instance.inputs = {"input": {"target_dir": "."}}
+    executable_application_instance.inputs[_FS] = {"input": {"target_dir": "."}}
     executable_application_instance.variables = {}
 
     exec_graph = executable_application_instance._get_executable_graph("test_wl2")
@@ -495,7 +495,7 @@ def test_derive_variables_for_template_path(mutable_mock_apps_repo):
 
     executable_application_instance.internals = {}
 
-    executable_application_instance.inputs = {"input": {"target_dir": "."}}
+    executable_application_instance.inputs[_FS] = {"input": {"target_dir": "."}}
     executable_application_instance.variables = {}
 
     exec_graph = executable_application_instance._get_executable_graph("test_wl2")

lib/ramble/ramble/test/application_inheritance.py

@@ -58,12 +58,12 @@ def test_basic_inheritance(mutable_mock_apps_repo):
     assert fom_conf["group_name"] == "test"
     assert fom_conf["units"] == "s"
 
-    assert "input" in app_inst.inputs
-    assert app_inst.inputs["input"]["url"] == "file:///tmp/test_file.log"
-    assert app_inst.inputs["input"]["description"] == "Not a file"
-    assert "inherited_input" in app_inst.inputs
-    assert app_inst.inputs["inherited_input"]["url"] == "file:///tmp/inherited_file.log"
-    assert app_inst.inputs["inherited_input"]["description"] == "Again, not a file"
+    assert "input" in app_inst.inputs[_FS]
+    assert app_inst.inputs[_FS]["input"]["url"] == "file:///tmp/test_file.log"
+    assert app_inst.inputs[_FS]["input"]["description"] == "Not a file"
+    assert "inherited_input" in app_inst.inputs[_FS]
+    assert app_inst.inputs[_FS]["inherited_input"]["url"] == "file:///tmp/inherited_file.log"
+    assert app_inst.inputs[_FS]["inherited_input"]["description"] == "Again, not a file"
 
     possible_vars = app_inst.workloads["test_wl"].find_variable("my_base_var")
     assert len(possible_vars) == 1

lib/ramble/ramble/test/application_language.py

@@ -289,14 +289,14 @@ def test_input_file_directive(app_class):
 
     assert hasattr(app_inst, "inputs")
     for input_name, conf in test_defs.items():
-        assert input_name in app_inst.inputs
+        assert input_name in app_inst.inputs[_FS]
 
         for conf_name, conf_val in conf.items():
-            assert conf_name in app_inst.inputs[input_name]
-            assert app_inst.inputs[input_name][conf_name] == conf_val
+            assert conf_name in app_inst.inputs[_FS][input_name]
+            assert app_inst.inputs[_FS][input_name][conf_name] == conf_val
 
-        assert "extension" in app_inst.inputs[input_name]
-        assert "expand" in app_inst.inputs[input_name]
+        assert "extension" in app_inst.inputs[_FS][input_name]
+        assert "expand" in app_inst.inputs[_FS][input_name]
 
 
 @pytest.mark.parametrize("app_class", app_types)

lib/ramble/ramble/test/mirror.py

@@ -34,6 +34,8 @@
     ),
 ]
 
+_FS = frozenset()
+
 
 class MockFetcher:
     """Mock fetcher object which implements the necessary functionality for
@@ -85,15 +87,15 @@ def create_archive(archive_dir, app_class):
                 with open(archive_name, "rb") as f:
                     bytes = f.read()
                     conf["fetcher"].digest = hashlib.sha256(bytes).hexdigest()
-                    app_class.inputs[conf["input_name"]]["sha256"] = conf["fetcher"].digest
+                    app_class.inputs[_FS][conf["input_name"]]["sha256"] = conf["fetcher"].digest
         else:
             with open(input_name, "w+") as f:
                 f.write("Input file\n")
 
             with open(input_name, "rb") as f:
                 bytes = f.read()
                 conf["fetcher"].digest = hashlib.sha256(bytes).hexdigest()
-                app_class.inputs[conf["input_name"]]["sha256"] = conf["fetcher"].digest
+                app_class.inputs[_FS][conf["input_name"]]["sha256"] = conf["fetcher"].digest
 
 
 def check_mirror(mirror_path, app_name, app_class):

lib/ramble/ramble/test/when.py

@@ -1008,3 +1008,49 @@ def test_executable_errors_when_overlapping_conditions(request):
 
             captured = workspace("setup", global_args=global_args)
             assert "test_exec_def_when is defined for overlapping `when` conditions" in captured
+
+
+@pytest.mark.parametrize(
+    "input_when,expected_input_file",
+    [
+        (False, "input1_false"),
+        (True, "input1_true"),
+    ],
+)
+def test_input_when(request, input_when, expected_input_file):
+    ws_name = request.node.name.replace("[", "_").replace("]", "_")
+
+    global_args = ["-w", ws_name]
+
+    with ramble.workspace.create(ws_name) as ws:
+        workspace(
+            "manage",
+            "experiments",
+            "when-directives",
+            "--wf",
+            "test_inputs_wl",
+            "-v",
+            "n_ranks=1",
+            "-v",
+            "n_nodes=1",
+            "-v",
+            "processes_per_node=1",
+            global_args=global_args,
+        )
+
+        config("add", f"variants:input_when:{input_when}", global_args=global_args)
+
+        ws._re_read()
+        workspace("setup", "--dry-run", global_args=global_args)
+
+        log_file = os.path.join(
+            ws.path,
+            "logs",
+            "setup.latest.out",
+        )
+
+        with open(log_file) as f:
+            data = f.read()
+
+            assert expected_input_file in data
+            assert ("input2" in data) == input_when

var/ramble/repos/builtin.mock/applications/when-directives/application.py

@@ -6,6 +6,8 @@
 # option. This file may not be copied, modified, or distributed
 # except according to those terms.
 
+import os
+
 from ramble.appkit import *
 
 
@@ -213,3 +215,41 @@ def test_builtin_when(self):
         "exec_when_wl",
         executables=["test_exec", "test_exec_skipped", "test_exec_def_when"],
     )
+
+    # for input_file()
+    cwd = os.getcwd()
+    variant(
+        "input_when",
+        default=False,
+        values=[True, False],
+        description="Add input file using when",
+    )
+
+    input_file(
+        "test-input1",
+        url=f"file://{cwd}/input1_false",
+        description="Test input when false",
+        when=["~input_when"],
+    )
+
+    input_file(
+        "test-input1",
+        url=f"file://{cwd}/input1_true",
+        expand=False,
+        description="Test input when true",
+        when=["+input_when"],
+    )
+
+    input_file(
+        "test-input2",
+        url=f"file://{cwd}/input2",
+        expand=False,
+        description="Test input skipped when false",
+        when=["+input_when"],
+    )
+
+    workload(
+        "test_inputs_wl",
+        executables=["test_exec"],
+        inputs=["test-input1", "test-input2"],
+    )