Enable Dynamic 'Zone' Configuration for Dataproc User Tools (#629)

user_tools/src/spark_rapids_pytools/cloud_api/dataproc.py

@@ -422,9 +422,20 @@ def _init_nodes(self):
             SparkNodeType.MASTER: master_node
         }
 
+    def _set_zone_from_props(self, prop_container: JSONPropertiesContainer):
+        """
+        Extracts the 'zoneUri' from the properties container and updates the environment variable dictionary.
+        """
+        if prop_container:
+            zone_uri = prop_container.get_value_silent('config', 'gceClusterConfig', 'zoneUri')
+            if zone_uri:
+                self.cli.env_vars['zone'] = FSUtil.get_resource_name(zone_uri)
+
     def _init_connection(self, cluster_id: str = None,
                          props: str = None) -> dict:
         cluster_args = super()._init_connection(cluster_id=cluster_id, props=props)
+        # extract and update zone to the environment variable
+        self._set_zone_from_props(cluster_args['props'])
         # propagate zone to the cluster
         cluster_args.setdefault('zone', self.cli.get_env_var('zone'))
         return cluster_args