diff --git a/user_tools/src/spark_rapids_pytools/cloud_api/dataproc.py b/user_tools/src/spark_rapids_pytools/cloud_api/dataproc.py index bd17016bd..d3a24e648 100644 --- a/user_tools/src/spark_rapids_pytools/cloud_api/dataproc.py +++ b/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