Merge pull request #590 from linsword13/expand-zips

Add expansion for zip and matrix members
GoogleCloudPlatform · Aug 2, 2024 · 3d631a3 · 3d631a3
2 parents 5d52320 + 702a724
commit 3d631a3
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 2 deletions.
diff --git a/lib/ramble/ramble/renderer.py b/lib/ramble/ramble/renderer.py
@@ -163,6 +163,24 @@ def render_objects(self, render_group, exclude_where=None, ignore_used=True, fat
             value = expander.expand_lists(unexpanded)
             object_variables[name] = value
 
+        # Expand zip and matrix members to allow indirections like
+        # ```
+        # variables:
+        #   my_vec: [1, 2, 3]
+        #   my_vec_ref: 'my_vec'
+        #   matrix:
+        #   - '{my_vec_ref}'
+        # ```
+        if matrices:
+            for matrix in matrices:
+                for i, unexpanded_var in enumerate(matrix):
+                    var = expander.expand_var(unexpanded_var)
+                    matrix[i] = var
+        if zips:
+            for zip_group, group_def in zips.items():
+                for i, unexpanded_var in enumerate(group_def):
+                    group_def[i] = expander.expand_var(unexpanded_var)
+
         new_objects = []
         defined_zips = {}
         consumed_zips = set()

diff --git a/lib/ramble/ramble/test/end_to_end/explicit_zips.py b/lib/ramble/ramble/test/end_to_end/explicit_zips.py
@@ -33,12 +33,13 @@ def test_wrfv4_explicit_zips(mutable_config, mutable_mock_workspace_path):
     mpi_command: 'mpirun -n {n_ranks} -ppn {processes_per_node}'
     batch_submit: 'batch_submit {execute_experiment}'
     partition: ['part1', 'part2']
+    partition_ref: partition
     processes_per_node: ['16', '36']
     n_ranks: '{processes_per_node}*{n_nodes}'
     n_threads: '1'
   zips:
     partitions:
-    - partition
+    - '{partition_ref}'
     - processes_per_node
   applications:
     wrfv4:
@@ -73,13 +74,17 @@ def test_wrfv4_explicit_zips(mutable_config, mutable_mock_workspace_path):
                 - TEST_VAR
               variables:
                 n_nodes: ['1', '2', '4', '8', '16']
+                n_nodes_ref: n_nodes
+                matrix_var: [1]
+                matrix_var_ref: matrix_var
               zips:
                 nodes:
-                - n_nodes
+                - '{n_nodes_ref}'
               matrix:
               - nodes
               - environments
               - partitions
+              - '{matrix_var_ref}'
   software:
     packages:
       gcc: