fix: Update cluster names to differentiate between production and development environments (GoogleCloudPlatform#1062)

This change switches the following DAGs to a dedicated, production-isolated cluster to prevent development activity from interfering with test results:
- `gke_node_pool_status`
- `multi-host-availability-rollback`
- `tpu_info_format_validation_dag`
- `gke_node_pool_label_update`

Author: yuna-tzeng

dags/tpu_observability/multi_host_nodepool_rollback_dag.py

@@ -24,6 +24,7 @@
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.trigger_rule import TriggerRule
 
+from dags import composer_env
 from dags.map_reproducibility.utils import constants
 from dags.common.vm_resource import Region, Zone
 from dags.tpu_observability.utils import node_pool_util as node_pool
@@ -72,11 +73,11 @@
 ) as dag:
   for machine in MachineConfigMap:
     config = machine.value
+    cluster_name = "tpu-observability-automation"
+    cluster_name += "-prod" if composer_env.is_prod_env() else "-dev"
     node_pool_info = node_pool.Info(
         project_id="cienet-cmcs",
-        cluster_name=Variable.get(
-            "CLUSTER_NAME", default_var="tpu-observability-automation"
-        ),
+        cluster_name=cluster_name,
         node_pool_name=Variable.get(
             "NODE_POOL_NAME", default_var="multi-host-nodepool-rollback-auto"
         ),

dags/tpu_observability/node_pool_status.py

@@ -21,6 +21,7 @@
 from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.task_group import TaskGroup
 
+from dags import composer_env
 from dags.map_reproducibility.utils import constants
 from dags.common.vm_resource import Region, Zone
 from dags.tpu_observability.utils import node_pool_util as node_pool
@@ -60,11 +61,11 @@
 ) as dag:
   for machine in MachineConfigMap:
     config = machine.value
+    cluster_name = "tpu-observability-automation"
+    cluster_name += "-prod" if composer_env.is_prod_env() else "-dev"
     node_pool_info = node_pool.Info(
         project_id=models.Variable.get("PROJECT_ID", default_var="cienet-cmcs"),
-        cluster_name=models.Variable.get(
-            "CLUSTER_NAME", default_var="tpu-observability-automation"
-        ),
+        cluster_name=cluster_name,
         node_pool_name=models.Variable.get(
             "NODE_POOL_NAME", default_var="node-pool-status-v6e-autotest"
         ),

dags/tpu_observability/tpu_info_format_validation_dags.py

@@ -29,6 +29,8 @@
 from airflow.exceptions import AirflowFailException
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.trigger_rule import TriggerRule
+
+from dags import composer_env
 from dags.common.vm_resource import Region, Zone
 from dags.map_reproducibility.utils import constants
 from dags.tpu_observability.configs.common import MachineConfigMap, TpuConfig
@@ -324,13 +326,13 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]):
 ) as dag:
   for machine in MachineConfigMap:
     config = machine.value
+    cluster_name = "tpu-observability-automation"
+    cluster_name += "-prod" if composer_env.is_prod_env() else "-dev"
     cluster_info = node_pool.Info(
         project_id=models.Variable.get(
             "TFV_PROJECT_ID", default_var="cienet-cmcs"
         ),
-        cluster_name=models.Variable.get(
-            "TFV_CLUSTER_NAME", default_var="tpu-observability-automation"
-        ),
+        cluster_name=cluster_name,
         node_pool_name=models.Variable.get(
             "TFV_NODE_POOL_NAME", default_var="tpu-info-fromat-test-v6e"
         ),

dags/tpu_observability/update_node_pool_label.py

@@ -21,6 +21,8 @@
 from airflow import models
 from airflow.utils.task_group import TaskGroup
 from airflow.utils.trigger_rule import TriggerRule
+
+from dags import composer_env
 from dags.common.vm_resource import Region, Zone
 from dags.map_reproducibility.utils import constants
 from dags.tpu_observability.configs.common import MachineConfigMap
@@ -57,11 +59,11 @@
 ) as dag:
   for machine in MachineConfigMap:
     config = machine.value
+    cluster_name = "tpu-observability-automation"
+    cluster_name += "-prod" if composer_env.is_prod_env() else "-dev"
     node_pool_info = node_pool.Info(
         project_id="cienet-cmcs",
-        cluster_name=models.Variable.get(
-            "CLUSTER_NAME", default_var="tpu-observability-automation"
-        ),
+        cluster_name=cluster_name,
         node_pool_name=models.Variable.get(
             "NODE_POOL_NAME", default_var="update-node-pool-label-v6e-autotest"
         ),

dags/tpu_observability/utils/node_pool_util.py

@@ -82,8 +82,8 @@ def create(
 
   composer.log_metadata_for_xlml_dashboard({
       "cluster_project": node_pool.project_id,
-      "region": node_pool.region,
-      "zone": node_pool.zone,
+      "region": node_pool.location,
+      "zone": node_pool.node_locations,
       "cluster_name": node_pool.cluster_name,
       "node_pool_name": node_pool.node_pool_name,
       "accelerator_type": node_pool.machine_type,