|
| 1 | +# Copyright 2025 Google LLC |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +"""A DAG to test JobSet Time-To-Recover (TTR) metric by triggering a node reboot.""" |
| 16 | + |
| 17 | +import datetime |
| 18 | +from datetime import timedelta |
| 19 | + |
| 20 | +from airflow import models |
| 21 | +from airflow.models.baseoperator import chain |
| 22 | +from airflow.utils.trigger_rule import TriggerRule |
| 23 | + |
| 24 | +from dags import composer_env |
| 25 | +from dags.tpu_observability.utils import jobset_util as jobset |
| 26 | +from dags.tpu_observability.utils import node_pool_util as node_pool |
| 27 | +from dags.tpu_observability.utils.jobset_util import Workload |
| 28 | +from dags.tpu_observability.configs.common import ( |
| 29 | + MachineConfigMap, |
| 30 | + GCS_CONFIG_PATH, |
| 31 | + GCS_JOBSET_CONFIG_PATH, |
| 32 | +) |
| 33 | +from dags.common.scheduling_helper.scheduling_helper import ( |
| 34 | + SchedulingHelper, |
| 35 | + get_dag_timeout, |
| 36 | +) |
| 37 | +from dags.common.task_group_with_timeout import TaskGroupWithTimeout |
| 38 | + |
| 39 | + |
| 40 | +DAG_ID = "jobset_ttr_node_reboot" |
| 41 | +DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID) |
| 42 | +SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) |
| 43 | + |
| 44 | + |
| 45 | +# Keyword arguments are generated dynamically at runtime (pylint does not |
| 46 | +# know this signature). |
| 47 | +with models.DAG( # pylint: disable=unexpected-keyword-arg |
| 48 | + dag_id=DAG_ID, |
| 49 | + start_date=datetime.datetime(2026, 1, 21), |
| 50 | + schedule=SCHEDULE if composer_env.is_prod_env() else None, |
| 51 | + dagrun_timeout=DAGRUN_TIMEOUT, |
| 52 | + catchup=False, |
| 53 | + tags=[ |
| 54 | + "cloud-ml-auto-solutions", |
| 55 | + "jobset", |
| 56 | + "time-to-recover", |
| 57 | + "tpu-observability", |
| 58 | + "node_reboot", |
| 59 | + "TPU", |
| 60 | + "v6e-16", |
| 61 | + ], |
| 62 | + description=( |
| 63 | + "Tests JobSet TTR metric by rebooting a random TPU node to trigger " |
| 64 | + "recovery, then polls the metric for updates." |
| 65 | + ), |
| 66 | + doc_md=""" |
| 67 | + # JobSet Time-To-Recover (TTR) Test Using Random Node Reboot |
| 68 | +
|
| 69 | + ### Description |
| 70 | + This DAG verifies that a TPU JobSet can recover from a hardware-level failure. |
| 71 | + It launches a JobSet, executes a `reboot` command on a random node via a |
| 72 | + privileged container (using nsenter), and uses a sensor to confirm that |
| 73 | + the TTR (Time-To-Recover) metric is recorded. |
| 74 | +
|
| 75 | + ### Prerequisites |
| 76 | + This test requires an existing cluster to run. |
| 77 | + GKE Cluster with TPU v6e support. |
| 78 | + The JobSet must be configured with |
| 79 | + privileged: True to allow node-level interaction. |
| 80 | +
|
| 81 | + ### Procedures |
| 82 | + First the node-pool is created, a jobset yaml is then launched on the cluster. |
| 83 | + After initialization, a random node reboot is triggered by executing |
| 84 | + `nsenter` through a privileged pod. This simulates a hardware failure |
| 85 | + by taking one of the TPU nodes offline. A sensor finally polls Cloud |
| 86 | + Monitoring to confirm the jobset TTR metric is updated. |
| 87 | + """, |
| 88 | +) as dag: |
| 89 | + for machine in MachineConfigMap: |
| 90 | + config = machine.value |
| 91 | + |
| 92 | + # Keyword arguments are generated dynamically at runtime (pylint does not |
| 93 | + # know this signature). |
| 94 | + with TaskGroupWithTimeout( # pylint: disable=unexpected-keyword-arg |
| 95 | + group_id=f"v{config.tpu_version.value}", |
| 96 | + timeout=timedelta(minutes=90), |
| 97 | + ): |
| 98 | + selector = jobset.generate_node_pool_selector("jobset-ttr-node-reboot") |
| 99 | + |
| 100 | + jobset_config = jobset.build_jobset_from_gcs_yaml( |
| 101 | + gcs_path=GCS_JOBSET_CONFIG_PATH, |
| 102 | + dag_name=DAG_ID, |
| 103 | + node_pool_selector=selector, |
| 104 | + privileged=True, |
| 105 | + ) |
| 106 | + |
| 107 | + cluster_info = node_pool.build_node_pool_info_from_gcs_yaml.override( |
| 108 | + task_id="build_node_pool_info_from_gcs_yaml" |
| 109 | + )( |
| 110 | + gcs_path=GCS_CONFIG_PATH, |
| 111 | + dag_name=DAG_ID, |
| 112 | + is_prod=composer_env.is_prod_env(), |
| 113 | + machine_type=config.machine_version.value, |
| 114 | + tpu_topology=config.tpu_topology, |
| 115 | + node_pool_selector=selector, |
| 116 | + ) |
| 117 | + |
| 118 | + create_node_pool = node_pool.create.override(task_id="create_node_pool")( |
| 119 | + node_pool=cluster_info, |
| 120 | + ) |
| 121 | + |
| 122 | + startup = jobset.create_jobset_startup_tasks( |
| 123 | + node_pool=cluster_info, |
| 124 | + jobset_config=jobset_config, |
| 125 | + workload_type=Workload.JAX_TPU_BENCHMARK, |
| 126 | + ) |
| 127 | + |
| 128 | + target_pod = jobset.draw_random_pod( |
| 129 | + node_pool=cluster_info, |
| 130 | + jobset_config=jobset_config, |
| 131 | + ) |
| 132 | + |
| 133 | + reboot_node = jobset.operate_pod.override(task_id="reboot_node")( |
| 134 | + node_pool=cluster_info, |
| 135 | + operation=jobset.PodOperationSpec.Reboot(), |
| 136 | + pod_name=target_pod, |
| 137 | + namespace="default", |
| 138 | + ) |
| 139 | + |
| 140 | + wait_for_metric_upload = jobset.wait_for_jobset_ttr_to_be_found.override( |
| 141 | + task_id="wait_for_jobset_ttr_to_be_found" |
| 142 | + )( |
| 143 | + node_pool=cluster_info, |
| 144 | + jobset_config=jobset_config, |
| 145 | + ) |
| 146 | + |
| 147 | + cleanup_workload = jobset.end_workload.override( |
| 148 | + task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE |
| 149 | + )(node_pool=cluster_info, jobset_config=jobset_config).as_teardown( |
| 150 | + setups=startup.jobset_start_time |
| 151 | + ) |
| 152 | + |
| 153 | + cleanup_node_pool = node_pool.delete.override( |
| 154 | + task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE |
| 155 | + )(node_pool=cluster_info).as_teardown( |
| 156 | + setups=create_node_pool, |
| 157 | + ) |
| 158 | + |
| 159 | + chain( |
| 160 | + selector, |
| 161 | + jobset_config, |
| 162 | + cluster_info, |
| 163 | + create_node_pool, |
| 164 | + *startup.tasks, |
| 165 | + target_pod, |
| 166 | + reboot_node, |
| 167 | + wait_for_metric_upload, |
| 168 | + cleanup_workload, |
| 169 | + cleanup_node_pool, |
| 170 | + ) |
0 commit comments