From 2a6f7e8f1b0cd35002511c0747eeccccad173787 Mon Sep 17 00:00:00 2001
From: Douglas Jacobsen <dwjacobsen@google.com>
Date: Wed, 15 Jan 2025 12:43:16 -0700
Subject: [PATCH] Allow templates registered into other workspace paths

This commit allows workspace path variables to be replaced in
destination paths for templates that are registered in objects.

This allows a template to be rendered into the shared directory of a
workspace, as an example.
---
 lib/ramble/ramble/application.py                   | 14 +++++++++++---
 lib/ramble/ramble/test/end_to_end/test_template.py |  6 ++++++
 .../applications/template/application.py           |  7 +++++++
 .../builtin.mock/applications/template/script.sh   | 12 ++++++++++++
 4 files changed, 36 insertions(+), 3 deletions(-)
 create mode 100644 var/ramble/repos/builtin.mock/applications/template/script.sh

diff --git a/lib/ramble/ramble/application.py b/lib/ramble/ramble/application.py
index 24d60a67d..5af8008af 100644
--- a/lib/ramble/ramble/application.py
+++ b/lib/ramble/ramble/application.py
@@ -48,6 +48,7 @@
 import ramble.util.stats
 import ramble.util.graph
 import ramble.util.class_attributes
+import ramble.util.path
 import ramble.util.lock as lk
 from ramble.util.foms import FomType
 from ramble.util.logger import logger
@@ -1389,7 +1390,7 @@ def _make_experiments(self, workspace, app_inst=None):
                     )
                 os.chmod(expand_path, _DEFAULT_CONTENT_PERM)
 
-            self._render_object_templates(exec_vars)
+            self._render_object_templates(exec_vars, replacement_vars=workspace.workspace_paths())
 
             experiment_script = workspace.experiments_script
             experiment_script.write(self.expander.expand_var("{batch_submit}\n"))
@@ -2311,8 +2312,11 @@ def _get_template_config(obj, tpl_config, obj_type):
             for tpl_conf in obj.templates.values():
                 yield _get_template_config(obj, tpl_conf, obj_type=obj_type)
 
-    def _render_object_templates(self, extra_vars):
+    def _render_object_templates(self, extra_vars, replacement_vars=None):
         run_dir = self.expander.experiment_run_dir
+        replacements = {}
+        if replacement_vars is not None:
+            replacements = replacement_vars
         for obj, tpl_config in self._object_templates():
             src_path = tpl_config["src_path"]
             with open(src_path) as f_in:
@@ -2325,7 +2329,11 @@ def _render_object_templates(self, extra_vars):
                 extra_vars_func = getattr(obj, extra_vars_func_name)
                 extra_vars.update(extra_vars_func())
             rendered = self.expander.expand_var(content, extra_vars=extra_vars)
-            out_path = os.path.join(run_dir, tpl_config["dest_name"])
+            out_path = ramble.util.path.substitute_path_variables(
+                tpl_config["dest_name"], local_replacements=replacements
+            )
+            if not os.path.isabs(out_path):
+                out_path = os.path.join(run_dir, out_path)
             perm = tpl_config.get("content_perm", _DEFAULT_CONTENT_PERM)
             with open(out_path, "w+") as f_out:
                 f_out.write(rendered)
diff --git a/lib/ramble/ramble/test/end_to_end/test_template.py b/lib/ramble/ramble/test/end_to_end/test_template.py
index a2f5ad694..9d9d5f31f 100644
--- a/lib/ramble/ramble/test/end_to_end/test_template.py
+++ b/lib/ramble/ramble/test/end_to_end/test_template.py
@@ -58,6 +58,9 @@ def test_template():
         content = f.read()
         assert script_path in content
 
+    script_path = os.path.join(ws.shared_dir, "script.sh")
+    assert os.path.isfile(script_path)
+
 
 def test_template_inherited():
     test_config = """
@@ -93,3 +96,6 @@ def test_template_inherited():
     with open(execute_path) as f:
         content = f.read()
         assert script_path in content
+
+    script_path = os.path.join(ws.shared_dir, "script.sh")
+    assert os.path.isfile(script_path)
diff --git a/var/ramble/repos/builtin.mock/applications/template/application.py b/var/ramble/repos/builtin.mock/applications/template/application.py
index 9a48807d8..4ced53659 100644
--- a/var/ramble/repos/builtin.mock/applications/template/application.py
+++ b/var/ramble/repos/builtin.mock/applications/template/application.py
@@ -41,3 +41,10 @@ def _bar_vars(self):
         expander = self.expander
         val = expander.expand_var('"hello {hello_name}"')
         return {"dynamic_hello_world": val}
+
+    register_template(
+        name="test",
+        src_name="script.sh",
+        dest_name="$workspace_shared/script.sh",
+        output_perm="755",
+    )
diff --git a/var/ramble/repos/builtin.mock/applications/template/script.sh b/var/ramble/repos/builtin.mock/applications/template/script.sh
new file mode 100644
index 000000000..3975dfb21
--- /dev/null
+++ b/var/ramble/repos/builtin.mock/applications/template/script.sh
@@ -0,0 +1,12 @@
+#!/bin/bash
+
+# This is a test script that should end up in a rendered location
+
+if [ ! -f file_list ]; then
+  exit 1
+fi
+
+for FILE in `cat file_list`
+do
+  wc -l $FILE
+done