node_pool_ttr_update_label.py

# Copyright 2025 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""A DAG to validate GKE node pool Times To Recover(TTR) metrics by triggering a label update."""

import datetime

from airflow import models
from airflow.models.baseoperator import chain
from airflow.utils.task_group import TaskGroup
from airflow.utils.trigger_rule import TriggerRule

from dags import composer_env
from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH
from dags.tpu_observability.utils import node_pool_util as node_pool


with models.DAG(
    dag_id="node_pool_ttr_update_label",
    start_date=datetime.datetime(2025, 9, 30),
    schedule="30 21 * * *" if composer_env.is_prod_env() else None,
    catchup=False,
    tags=[
        "gke",
        "tpu-observability",
        "node-pool-ttr-update-label",
        "TPU",
        "v6e-16",
    ],
    description=(
        "This DAG verifies the GKE node pool's Times To Recover(TTR) metrics "
        "by triggering a label update and confirming the recovery time "
        "is recorded"
    ),
    doc_md="""
      # GKE Node Pool Times To Recover(TTR) Metric Validation DAG

      ### Description
      This DAG automates the validation of GKE node pool Times To Recover(TTR) metrics.
      It creates a node pool and updates its labels then verifies that the TTR metric
      is correctly generated and reported to Google Cloud Monitoring.

      ### Prerequisites
      This test requires an existing GKE cluster.

      ### Procedures
      1. Create a temporary node pool.
      2. Wait for the node pool to be RUNNING.
      3. Update the node pool label.
      4. Wait for the Times To Recover(TTR) metrics to appear in Google Cloud Monitoring.
      5. Clean up the node pool after the tests.
    """,
) as dag:
  for machine in MachineConfigMap:
    config = machine.value
    LABELS_TO_UPDATE = {"test_key": "test_val"}

    with TaskGroup(group_id=f"v{config.tpu_version.value}"):
      node_pool_info = node_pool.build_node_pool_info_from_gcs_yaml.override(
          task_id="build_node_pool_info_from_gcs_yaml"
      )(
          gcs_path=GCS_CONFIG_PATH,
          dag_name="node_pool_ttr_update_label",
          is_prod=composer_env.is_prod_env(),
          machine_type=config.machine_version.value,
          tpu_topology=config.tpu_topology,
      )

      task_id = "create_node_pool"
      create_node_pool = node_pool.create.override(task_id=task_id)(
          node_pool=node_pool_info,
      )

      task_id = "wait_for_provisioning"
      wait_for_provisioning = node_pool.wait_for_status.override(
          task_id=task_id
      )(node_pool=node_pool_info, status=node_pool.Status.PROVISIONING)

      task_id = "wait_for_running"
      wait_for_running = node_pool.wait_for_status.override(task_id=task_id)(
          node_pool=node_pool_info, status=node_pool.Status.RUNNING
      )

      task_id = "update_node_pool_label"
      update_node_pool_label = node_pool.update.override(task_id=task_id)(
          node_pool=node_pool_info,
          spec=node_pool.NodePoolUpdateSpec.Label(delta=LABELS_TO_UPDATE),
      )

      task_id = "wait_for_recovered"
      wait_for_recovered = node_pool.wait_for_status.override(task_id=task_id)(
          node_pool=node_pool_info, status=node_pool.Status.RUNNING
      )

      task_id = "wait_for_ttr"
      wait_for_ttr = node_pool.wait_for_ttr.override(task_id=task_id)(
          node_pool=node_pool_info, operation_start_time=update_node_pool_label
      )

      task_id = "cleanup_node_pool"
      cleanup_node_pool = node_pool.delete.override(
          task_id=task_id, trigger_rule=TriggerRule.ALL_DONE
      )(node_pool=node_pool_info).as_teardown(
          setups=create_node_pool,
      )

      chain(
          node_pool_info,
          create_node_pool,
          wait_for_provisioning,
          wait_for_running,
          update_node_pool_label,
          wait_for_recovered,
          wait_for_ttr,
          cleanup_node_pool,
      )