Add reservation specification to create function and standardize DAGs (GoogleCloudPlatform#915)

Modify `node_pool_status` and `multi_host_nodepool_rollback_dag` DAG to use same cluster and adding use reservation in `create` task.

Author: yuna-tzeng

dags/common/vm_resource.py

@@ -90,6 +90,7 @@ class Region(enum.Enum):
   # used for GKE
   US_CENTRAL1 = "us-central1"
   ASIA_NORTHEAST1 = "asia-northeast1"
+  US_EAST5 = "us-east5"
 
 
 class Zone(enum.Enum):

dags/tpu_observability/multi_host_nodepool_rollback_dag.py

@@ -49,23 +49,23 @@
   node_pool_info = node_pool.Info(
       project_id=Project.TPU_PROD_ENV_ONE_VM.value,
       cluster_name=Variable.get(
-          "CLUSTER_NAME", default_var="qmcgarry-auto-test"
+          "CLUSTER_NAME", default_var="tpu-observability-automation"
       ),
       node_pool_name=Variable.get(
-          "NODE_POOL_NAME", default_var="nodepool-auto"
-      ),
-      location=Variable.get(
-          "LOCATION", default_var=Region.ASIA_NORTHEAST1.value
+          "NODE_POOL_NAME", default_var="multi_host_nodepool_rollback_auto"
       ),
+      location=Variable.get("LOCATION", default_var=Region.US_EAST5.value),
       node_locations=Variable.get(
-          "NODE_LOCATIONS", default_var=Zone.ASIA_NORTHEAST1_B.value
+          "NODE_LOCATIONS", default_var=Zone.US_EAST5_B.value
       ),
       num_nodes=Variable.get("NUM_NODES", default_var=4),
       machine_type=Variable.get("MACHINE_TYPE", default_var="ct6e-standard-4t"),
       tpu_topology=Variable.get("TPU_TOPOLOGY", default_var="4x4"),
   )
 
-  create_node_pool = node_pool.create(node_pool=node_pool_info)
+  create_node_pool = node_pool.create(
+      node_pool=node_pool_info, reservation="cloudtpu-20250131131310-2118578099"
+  )
 
   wait_node_pool_available = node_pool.wait_for_availability(
       node_pool=node_pool_info, availability=True

dags/tpu_observability/node_pool_status.py

@@ -45,16 +45,16 @@
           "PROJECT_ID", default_var=Project.TPU_PROD_ENV_ONE_VM.value
       ),
       cluster_name=models.Variable.get(
-          "CLUSTER_NAME", default_var="yuna-xpk-v6e-2"
+          "CLUSTER_NAME", default_var="tpu-observability-automation"
       ),
       node_pool_name=models.Variable.get(
-          "NODE_POOL_NAME", default_var="yuna-v6e-autotest"
+          "NODE_POOL_NAME", default_var="node-pool-status-v6e-autotest"
       ),
       location=models.Variable.get(
-          "LOCATION", default_var=Region.ASIA_NORTHEAST1.value
+          "LOCATION", default_var=Region.US_EAST5.value
       ),
       node_locations=models.Variable.get(
-          "NODE_LOCATIONS", default_var=Zone.ASIA_NORTHEAST1_B.value
+          "NODE_LOCATIONS", default_var=Zone.US_EAST5_B.value
       ),
       num_nodes=models.Variable.get("NUM_NODES", default_var=4),
       machine_type=models.Variable.get(
@@ -74,7 +74,7 @@
 
   task_id = "create_node_pool"
   create_node_pool = node_pool.create.override(task_id=task_id)(
-      node_pool=node_pool_info
+      node_pool=node_pool_info, reservation="cloudtpu-20250131131310-2118578099"
   )
 
   task_id = "wait_for_provisioning"

dags/tpu_observability/utils/node_pool_util.py

@@ -51,21 +51,29 @@ class Info:
 
 
 @task
-def create(node_pool: Info, ignore_failure: bool = False) -> None:
+def create(
+    node_pool: Info,
+    reservation: str = None,
+    ignore_failure: bool = False,
+) -> None:
   """Creates a GKE node pool by the given node pool information."""
 
   command = (
       f"gcloud container node-pools create {node_pool.node_pool_name} "
       f"--project={node_pool.project_id} "
       f"--cluster={node_pool.cluster_name} "
       f"--location={node_pool.location} "
-      f"--node-locations {node_pool.node_locations} "
+      f"--node-locations={node_pool.node_locations} "
       f"--num-nodes={node_pool.num_nodes} "
       f"--machine-type={node_pool.machine_type} "
-      f"--tpu-topology={node_pool.tpu_topology}"
+      f"--tpu-topology={node_pool.tpu_topology} "
   )
+
+  if reservation:
+    command += f"--reservation-affinity=specific --reservation={reservation} "
+
   if ignore_failure:
-    command += " 2>&1 || true"
+    command += "2>&1 || true "
 
   process = subprocess.run(
       command, shell=True, check=True, capture_output=True, text=True