Merge pull request #1312 from douglasjacobsen/overwrite-stage-files

Update `stage_files` to overwrite named stages

Author: rfbgo

lib/ramble/ramble/language/application_language.py

@@ -431,9 +431,11 @@ def _define_cleanup(obj):
 
 @application_directive("executables")
 def stage_files(
-    src,
-    dst,
+    src=None,
+    dst=None,
+    stages=None,
     name=None,
+    method="user-defined",
     when=None,
     **kwargs,
 ):
@@ -447,21 +449,32 @@ def stage_files(
     option, which can be set to 'cp', 'rsync', 'symbolic_link', or 'hard_link'.
 
     Args:
-        src (str): The source path of the file or directory.
-        dst (str): The destination path.
+        src (str | None): The source path of the file or directory.
+        dst (str | None): The destination path. If src is passed in, and dst is
+                          not, dst defaults to the experiment_run_dir.
+        stages (list(tuple(str, str)) | None): A list of tuples describing pairs
+                                               of src, dest locations to stage.
         name (str | None): The name of the executable. Defaults to 'stage-files'.
+        method (str): The method to use for this stage. Can be one of:
+                      "user-defined", "cp", "rsync", "symbolic_link", "hard_link"
         when (list | None): List of when conditions to apply to this directive.
     """
 
+    valid_methods = ["user-defined", "cp", "rsync", "symbolic_link", "hard_link"]
+
+    method_map = {
+        "cp": "cp -Lr",
+        "rsync": "rsync -Lr",
+        "symbolic_link": "ln -sf",
+        "hard_link": "ln -f",
+    }
+
     def _execute_stage_files(app):
         import os
 
         import ramble.config
         from ramble.util.executable import CommandExecutable
 
-        cfg = ramble.config.config
-        stage_method = cfg.get("config", {}).get("stage_method", "cp")
-
         exec_name = name if name else "stage-files"
 
         when_list = ramble.language.language_helpers.build_when_list(
@@ -481,28 +494,40 @@ def _execute_stage_files(app):
                 "as an executable. Please provide a unique name attribute."
             )
 
-        # Prepare the core staging command
-        if stage_method == "rsync":
-            stage_cmd = f"rsync -Lr {src} {dst}"
-        elif stage_method == "hard_link":
-            stage_cmd = f"ln {src} {dst}"
-        elif stage_method == "symbolic_link":
-            stage_cmd = f"ln -s {src} {dst}"
-        else:  # stage_method == "cp"
-            stage_cmd = f"cp -Lr {src} {dst}"
-
-        template = [stage_cmd]
-
-        # Prepend mkdir if dst has a parent directory
-        parent_dir = os.path.dirname(dst)
-        if parent_dir and parent_dir != ".":
-            template.insert(0, f"mkdir -p {parent_dir}")
-
-        if exec_name in app.executables[when_set]:
-            app.executables[when_set][exec_name].add_template(template)
-        else:
-            app.executables[when_set][exec_name] = CommandExecutable(
-                name=exec_name, template=template, allow_extension=True, **kwargs
+        if method not in valid_methods:
+            raise DirectiveError(
+                f"stage_files directive on application {app.name} was given an "
+                f"invalid method argument of {method}.\n"
+                f"Valid methods include: {valid_methods}"
             )
 
+        stage_method = method
+        if stage_method == "user-defined":
+            cfg = ramble.config.config
+            stage_method = cfg.get("config", {}).get("stage_method", "cp")
+        stage_cmd = method_map[stage_method]
+
+        template = []
+
+        if src is not None:
+            if dst is not None:
+                parent_dir = os.path.dirname(dst)
+                if parent_dir and parent_dir != ".":
+                    template.insert(0, f"mkdir -p {parent_dir}")
+                template.append(f"{stage_cmd} {src} {dst}")
+            else:
+                template.append(f"{stage_cmd} {src} {{experiment_run_dir}}/.")
+
+        if isinstance(stages, list):
+            for pair_src, pair_dst in stages:
+                parent_dir = os.path.dirname(pair_dst)
+                if parent_dir and parent_dir != ".":
+                    template.append(f"mkdir -p {parent_dir}")
+
+                template.append(f"{stage_cmd} {pair_src} {pair_dst}")
+
+        app.executables[when_set][exec_name] = CommandExecutable(
+            name=exec_name, template=template, allow_extension=True, **kwargs
+        )
+
     return _execute_stage_files

lib/ramble/ramble/test/application_language.py

@@ -394,31 +394,106 @@ class BrokenWorkloadDefaults(ExecutableApplication):  # noqa: F405
         assert "workload_defaults cannot be used with workload, workloads" in err
 
 
+def test_stage_files_directive_no_dst():
+    import ramble.config
+
+    with ramble.config.override("config:stage_method", "cp"):
+
+        class TestApp(ExecutableApplication):  # noqa: F405
+            name = "test-app"
+
+        app_inst = TestApp("/not/a/path")
+        app_inst.stage_files(src="src")
+
+        assert "stage-files" in app_inst.executables[frozenset()]
+        exec = app_inst.executables[frozenset()]["stage-files"]
+        assert exec.template == ["cp -Lr src {experiment_run_dir}/."]
+
+
+def test_stage_files_directive_stages():
+    import ramble.config
+
+    with ramble.config.override("config:stage_method", "cp"):
+
+        class TestApp(ExecutableApplication):  # noqa: F405
+            name = "test-app"
+
+        app_inst = TestApp("/not/a/path")
+        app_inst.stage_files(stages=[("src1", "a/b"), ("src2", "c/d")])
+
+        assert "stage-files" in app_inst.executables[frozenset()]
+        exec = app_inst.executables[frozenset()]["stage-files"]
+        assert exec.template == [
+            "mkdir -p a",
+            "cp -Lr src1 a/b",
+            "mkdir -p c",
+            "cp -Lr src2 c/d",
+        ]
+
+
+def test_stage_files_directive_overwrite():
+    import ramble.config
+
+    with ramble.config.override("config:stage_method", "cp"):
+
+        class TestApp(ExecutableApplication):  # noqa: F405
+            name = "test-app"
+
+        app_inst = TestApp("/not/a/path")
+        app_inst.stage_files(src="src", dst="a/b")
+
+        assert "stage-files" in app_inst.executables[frozenset()]
+        exec = app_inst.executables[frozenset()]["stage-files"]
+        assert exec.template == ["mkdir -p a", "cp -Lr src a/b"]
+
+        app_inst.stage_files(src="src2", dst="c/d")
+        exec = app_inst.executables[frozenset()]["stage-files"]
+        assert exec.template == ["mkdir -p c", "cp -Lr src2 c/d"]
+
+
 @pytest.mark.parametrize(
     "stage_method,template_contents",
     [
-        ("cp", "cp -Lr src"),
-        ("rsync", "rsync -Lr src"),
-        ("symbolic_link", "ln -s src"),
-        ("hard_link", "ln src"),
+        ("cp", "cp -Lr src dst"),
+        ("rsync", "rsync -Lr src dst"),
+        ("symbolic_link", "ln -sf src dst"),
+        ("hard_link", "ln -f src dst"),
     ],
 )
-def test_stage_files_directive(stage_method, template_contents):
+def test_stage_files_directive_method(stage_method, template_contents):
+    class TestApp(ExecutableApplication):  # noqa: F405
+        name = "test-app"
+
+    app_inst = TestApp("/not/a/path")
+    app_inst.stage_files(src="src", dst="dst", method=stage_method)
+
+    assert "stage-files" in app_inst.executables[frozenset()]
+    exec = app_inst.executables[frozenset()]["stage-files"]
+
+    assert exec.template == [template_contents]
+
+
+def test_stage_files_directive_user_defined():
     import ramble.config
 
-    with ramble.config.override("config:stage_method", stage_method):
+    with ramble.config.override("config:stage_method", "rsync"):
 
         class TestApp(ExecutableApplication):  # noqa: F405
             name = "test-app"
 
         app_inst = TestApp("/not/a/path")
-        app_inst.stage_files(src="src", dst="dst")
+        app_inst.stage_files(src="src", dst="dst", method="user-defined")
 
         assert "stage-files" in app_inst.executables[frozenset()]
         exec = app_inst.executables[frozenset()]["stage-files"]
 
-        found = False
-        for line in exec.template:
-            if template_contents in line:
-                found = True
-        assert found
+        assert exec.template == ["rsync -Lr src dst"]
+
+
+def test_stage_files_directive_invalid_method():
+    class TestApp(ExecutableApplication):  # noqa: F405
+        name = "test-app"
+
+    app_inst = TestApp("/not/a/path")
+    with pytest.raises(DirectiveError):
+        app_inst.stage_files(src="src", dst="dst", method="invalid")

var/ramble/repos/builtin/applications/wrfv3/application.py

@@ -49,22 +49,23 @@ class Wrfv3(ExecutableApplication):
     )
 
     stage_files(
-        name="stage-run-files",
-        src="{wrf_path}/run/*",
-        dst="{experiment_run_dir}/.",
+        stages=[
+            ("{wrf_path}/run/*", "{experiment_run_dir}/."),
+            ("{input_path}/*", "{experiment_run_dir}/."),
+        ]
     )
 
     stage_files(
-        name="stage-input",
-        src="{input_path}/*",
+        name="stage-namelist",
+        method="cp",
+        src="{input_path}/namelist.*",
         dst="{experiment_run_dir}/.",
     )
 
     executable(
         "setup",
         template=[
             "rm -f rsl.* wrfout* namelist*",
-            "cp {input_path}/namelist.* {experiment_run_dir}/.",
         ],
         use_mpi=False,
         output_capture=OUTPUT_CAPTURE.ALL,
@@ -74,13 +75,13 @@ class Wrfv3(ExecutableApplication):
 
     workload(
         "CONUS_2p5km",
-        executables=["setup", "stage-input", "stage-run-files", "execute"],
+        executables=["stage-files", "stage-namelist", "setup", "execute"],
         input="CONUS_2p5km",
     )
 
     workload(
         "CONUS_12km",
-        executables=["setup", "stage-input", "stage-run-files", "execute"],
+        executables=["stage-files", "stage-namelist", "setup", "execute"],
         input="CONUS_12km",
     )
 

var/ramble/repos/builtin/applications/wrfv4/application.py

@@ -88,20 +88,23 @@ class Wrfv4(ExecutableApplication):
     )
 
     stage_files(
-        src="{wrf_path}/run/*",
-        dst="{experiment_run_dir}/.",
+        stages=[
+            ("{wrf_path}/run/*", "{experiment_run_dir}/."),
+            ("{input_path}/*", "{experiment_run_dir}/."),
+        ]
     )
 
     stage_files(
-        src="{input_path}/*",
+        name="stage-namelist",
+        src="{input_path}/namelist.*",
         dst="{experiment_run_dir}/.",
+        method="cp",
     )
 
     executable(
         "setup",
         template=[
             "rm -f rsl.* wrfout* namelist*",
-            "cp {input_path}/namelist.* {experiment_run_dir}/.",
         ],
         use_mpi=False,
         output_capture=OUTPUT_CAPTURE.ALL,
@@ -156,6 +159,7 @@ class Wrfv4(ExecutableApplication):
         "CONUS_2p5km",
         executables=[
             "stage-files",
+            "stage-namelist",
             "setup",
             "define_nproc_y",
             "define_nproc_x",
@@ -169,6 +173,7 @@ class Wrfv4(ExecutableApplication):
         "CONUS_12km",
         executables=[
             "stage-files",
+            "stage-namelist",
             "setup",
             "define_nproc_y",
             "define_nproc_x",
@@ -183,6 +188,7 @@ class Wrfv4(ExecutableApplication):
         "Maria_1km",
         executables=[
             "stage-files",
+            "stage-namelist",
             "setup",
             "define_nproc_y",
             "define_nproc_x",