add --dns-falg for imo-a3ultra cluster (GoogleCloudPlatform#811)

Author: RUEI4341

dags/map_reproducibility/utils/common_utils.py

@@ -216,8 +216,7 @@ def configure_project_and_cluster(cluster: str, cluster_region: str):
   set_project_command = (
       f"gcloud config set project {PROJECT}",
       "sudo chown -R airflow:airflow /home/airflow/composer_kube_config",
-      "gcloud container clusters get-credentials "
-      f"{cluster} --region {cluster_region}",
+      get_credentials_command(cluster, cluster_region),
   )
   return set_project_command
 
@@ -799,7 +798,8 @@ def get_nemo_metrics_cmds(
 
 def cleanup_all_runs_cmds(cluster, cluster_region, prefix="cml-"):
   cleanup_cmds = (
-      f"echo 'Getting credentials for cluster {cluster}...' && gcloud container clusters get-credentials {cluster} --region {cluster_region} --project {PROJECT}",
+      f"echo 'Uninstalling jobs with prefix {prefix}...'",
+      get_credentials_command(cluster, cluster_region, PROJECT),
       f"echo 'Uninstalling jobs with prefix {prefix}...' && "
       f"JOBS=$(kubectl get job -n default | grep '^{prefix}' | awk '{{print $1}}' | sed 's/-workload-0$//' | sort -u) && "
       f'if [ -n "$JOBS" ]; then '
@@ -1076,6 +1076,28 @@ def get_cluster(hardware: str = "a3ultra"):
     return "gke-a4-sbrg1", "us-east4"
 
 
+def is_dns_endpoint(cluster: str):
+  """Indicate if the cluster needs dns-endpoint flag."""
+  if cluster == "imo-a3ultra":
+    return True
+  return False
+
+
+def get_credentials_command(
+    cluster: str, cluster_region: str, project: None | str = None
+):
+  """
+  Returns the gcloud get-credentials command string.
+  Includes the --dns-endpoint flag if required.
+  """
+  base_command = f"gcloud container clusters get-credentials {cluster} --region {cluster_region}"
+  if project is not None:
+    base_command += f" --project {project}"
+  if is_dns_endpoint(cluster):
+    return f"{base_command} --dns-endpoint"
+  return base_command
+
+
 def get_scheduled_time(hardware: str, model: str, framework: str):
   """
   Returns a cron expression for the DAG schedule based on

dags/map_reproducibility/utils/sample_workload_utils.py

@@ -31,6 +31,7 @@
     get_internal_pre_workload_job_name,
     get_bq_writer_path,
     get_cluster,
+    get_credentials_command,
     calculate_maxtext_metrics,
     copy_bucket_cmds_maxtext,
     get_job_gcs_bucket_folder,
@@ -124,8 +125,7 @@ def execute_workload_commands(commands: list, cwd: str) -> Tuple[bool, list]:
 def sample_job_configure_project_and_cluster(cluster: str, cluster_region: str):
   set_project_command = (
       f"gcloud config set project {PROJECT}",
-      "gcloud container clusters get-credentials "
-      f"{cluster} --region {cluster_region}",
+      get_credentials_command(cluster, cluster_region),
   )
   return set_project_command