Add local mode to gcp-metadata modifier

This allows for it to contiune to work in environments where we might not have pdsh, such as GKE
GoogleCloudPlatform · Oct 9, 2024 · fde413f · fde413f
1 parent a11a7fd
commit fde413f
Showing 1 changed file with 47 additions and 15 deletions.
diff --git a/var/ramble/repos/builtin/modifiers/gcp-metadata/modifier.py b/var/ramble/repos/builtin/modifiers/gcp-metadata/modifier.py
@@ -27,12 +27,42 @@ class GcpMetadata(BasicModifier):
     maintainers("rfbgo")
 
     mode("standard", description="Standard execution mode")
+    mode(
+        "local", description="Local execution (disables parallel prefix/pdssh)"
+    )
     default_mode("standard")
 
     software_spec("pdsh", pkg_spec="pdsh", package_manager="spack*")
 
     required_variable("hostlist")
 
+    modifier_variable(
+        "metadata_parallel_prefix",
+        default="pdsh -R ssh -N -w {hostlist} '",
+        modes=["standard"],
+        description="Express how parlalelism should be done between nodes",
+    )
+    modifier_variable(
+        "metadata_parallel_prefix",
+        default="",
+        modes=["local"],
+        description="Express how parlalelism should be done between nodes",
+    )
+
+    # Need to close any open `'` we leave in the prefix
+    modifier_variable(
+        "metadata_parallel_suffix",
+        default="'",
+        modes=["standard"],
+        description="Optional suffix for {metadata_parallel_prefix}",
+    )
+    modifier_variable(
+        "metadata_parallel_suffix",
+        default="",
+        modes=["local"],
+        description="Optional suffix for {metadata_parallel_prefix}",
+    )
+
     executable_modifier("gcp_metadata_exec")
 
     def gcp_metadata_exec(self, executable_name, executable, app_inst=None):
@@ -46,15 +76,16 @@ def gcp_metadata_exec(self, executable_name, executable, app_inst=None):
         post_cmds = []
         pre_cmds = []
 
-        pre_cmds.append(
-            CommandExecutable(
-                "save-old-loglevel",
-                template=[
-                    'old_pdsh_args="$PDSH_SSH_ARGS_APPEND"',
-                    'export PDSH_SSH_ARGS_APPEND="-q"',
-                ],
+        if self._usage_mode != "local":
+            pre_cmds.append(
+                CommandExecutable(
+                    "save-old-loglevel",
+                    template=[
+                        'old_pdsh_args="$PDSH_SSH_ARGS_APPEND"',
+                        'export PDSH_SSH_ARGS_APPEND="-q"',
+                    ],
+                )
             )
-        )
 
         payloads = [
             # type, end point, per_node
@@ -74,8 +105,8 @@ def gcp_metadata_exec(self, executable_name, executable, app_inst=None):
             prefix = ""
             suffix = ""
             if per_node:
-                prefix = "pdsh -R ssh -N -w {hostlist} '"
-                suffix = "'"
+                prefix = self.expander.expand_var("{metadata_parallel_prefix}")
+                suffix = self.expander.expand_var("{metadata_parallel_suffix}")
             log_name = end_point.split("/")[-1]
             pre_cmds.append(
                 CommandExecutable(
@@ -92,12 +123,13 @@ def gcp_metadata_exec(self, executable_name, executable, app_inst=None):
                 )
             )
 
-        pre_cmds.append(
-            CommandExecutable(
-                "restore-old-loglevel",
-                template=['export PDSH_SSH_ARGS_APPEND="$old_pdsh_args"'],
+        if self._usage_mode != "local":
+            pre_cmds.append(
+                CommandExecutable(
+                    "restore-old-loglevel",
+                    template=['export PDSH_SSH_ARGS_APPEND="$old_pdsh_args"'],
+                )
             )
-        )
 
         return pre_cmds, post_cmds