fix: Patch jobset_uptime_validation scheduling/selector and standardize DAG_ID (#1188)

This change provides critical patches for the `jobset_uptime_validation` DAG to ensure proper scheduling and resource selection. Additionally, it refactors the tpu_observability module to use a unified DAG_ID constant, eliminating hard-coded strings.

Author: chengpinglin

dags/common/scheduling_helper/scheduling_helper.py

@@ -45,12 +45,13 @@ class DayOfWeek(enum.Enum):
         "jobset_rollback_ttr": dt.timedelta(minutes=90),
         "jobset_ttr_node_pool_resize": dt.timedelta(minutes=90),
         "jobset_ttr_pod_delete": dt.timedelta(minutes=90),
-        "multi-host-availability-rollback": dt.timedelta(minutes=30),
+        "multi_host_nodepool_rollback": dt.timedelta(minutes=30),
         "node_pool_ttr_disk_size": dt.timedelta(minutes=90),
         "node_pool_ttr_update_label": dt.timedelta(minutes=90),
         "tpu_info_format_validation_dag": dt.timedelta(minutes=30),
         "tpu_sdk_monitoring_validation": dt.timedelta(minutes=30),
         "jobset_ttr_kill_process": dt.timedelta(minutes=90),
+        "jobset_uptime_validation": dt.timedelta(minutes=90),
     },
 }
 

dags/tpu_observability/jobset_ttr_kill_process.py

@@ -131,15 +131,15 @@ def kill_tpu_pod_workload(info: node_pool.Info, pod_name: str) -> None:
 
       jobset_config = jobset.build_jobset_from_gcs_yaml(
           gcs_path=GCS_JOBSET_CONFIG_PATH,
-          dag_name="jobset_ttr_kill_process",
+          dag_name=DAG_ID,
           node_pool_selector=selector,
       )
 
       cluster_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="jobset_ttr_kill_process",
+          dag_name=DAG_ID,
           is_prod=composer_env.is_prod_env(),
           machine_type=config.machine_version.value,
           tpu_topology=config.tpu_topology,

dags/tpu_observability/jobset_ttr_node_pool_resize.py

@@ -95,15 +95,15 @@
 
       jobset_config = jobset.build_jobset_from_gcs_yaml(
           gcs_path=GCS_JOBSET_CONFIG_PATH,
-          dag_name="jobset_ttr_node_pool_resize",
+          dag_name=DAG_ID,
           node_pool_selector=selector,
       )
 
       cluster_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="jobset_ttr_node_pool_resize",
+          dag_name=DAG_ID,
           is_prod=composer_env.is_prod_env(),
           machine_type=config.machine_version.value,
           tpu_topology=config.tpu_topology,

dags/tpu_observability/jobset_ttr_pod_delete.py

@@ -90,15 +90,15 @@
 
       jobset_config = jobset.build_jobset_from_gcs_yaml(
           gcs_path=GCS_JOBSET_CONFIG_PATH,
-          dag_name="jobset_ttr_pod_delete",
+          dag_name=DAG_ID,
           node_pool_selector=selector,
       )
 
       cluster_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="jobset_ttr_pod_delete",
+          dag_name=DAG_ID,
           is_prod=composer_env.is_prod_env(),
           machine_type=config.machine_version.value,
           tpu_topology=config.tpu_topology,

dags/tpu_observability/jobset_ttr_rollback.py

@@ -92,15 +92,15 @@
 
       jobset_config = jobset.build_jobset_from_gcs_yaml(
           gcs_path=GCS_JOBSET_CONFIG_PATH,
-          dag_name="jobset_rollback_ttr",
+          dag_name=DAG_ID,
           node_pool_selector=selector,
       )
 
       cluster_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="jobset_rollback_ttr",
+          dag_name=DAG_ID,
           is_prod=composer_env.is_prod_env(),
           machine_type=config.machine_version.value,
           tpu_topology=config.tpu_topology,

dags/tpu_observability/jobset_uptime_validation.py

@@ -32,6 +32,12 @@
 from dags.tpu_observability.utils import node_pool_util as node_pool
 from dags.tpu_observability.utils.jobset_util import Workload
 from dags.tpu_observability.utils.time_util import TimeUtil
+from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
+
+
+DAG_ID = "jobset_uptime_validation"
+DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)
+SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID)
 
 
 @task
@@ -44,10 +50,11 @@ def get_current_time() -> TimeUtil:
 # Keyword arguments are generated dynamically at runtime (pylint does not
 # know this signature).
 with models.DAG(  # pylint: disable=unexpected-keyword-arg
-    dag_id="jobset_uptime_validation",
+    dag_id=DAG_ID,
     start_date=datetime.datetime(2025, 8, 15),
     default_args={"retries": 0},
-    schedule="30 16 * * *" if composer_env.is_prod_env() else None,
+    schedule=SCHEDULE if composer_env.is_prod_env() else None,
+    dagrun_timeout=DAGRUN_TIMEOUT,
     catchup=False,
     tags=[
         "cloud-ml-auto-solutions",
@@ -93,19 +100,23 @@ def get_current_time() -> TimeUtil:
     with TaskGroup(  # pylint: disable=unexpected-keyword-arg
         group_id=f"v{config.tpu_version.value}"
     ):
+      selector = jobset.generate_node_pool_selector("jobset-rollback-ttr")
+
       jobset_config = jobset.build_jobset_from_gcs_yaml(
           gcs_path=GCS_JOBSET_CONFIG_PATH,
-          dag_name="jobset_uptime_validation",
+          dag_name=DAG_ID,
+          node_pool_selector=selector,
       )
 
       cluster_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="jobset_uptime_validation",
+          dag_name=DAG_ID,
           is_prod=composer_env.is_prod_env(),
           machine_type=config.machine_version.value,
           tpu_topology=config.tpu_topology,
+          node_pool_selector=selector,
       )
 
       create_node_pool = node_pool.create.override(task_id="create_node_pool")(

dags/tpu_observability/multi_host_nodepool_rollback_dag.py

@@ -31,7 +31,7 @@
 from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
 
 
-DAG_ID = "multi-host-availability-rollback"
+DAG_ID = "multi_host_nodepool_rollback"
 DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)
 SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID)
 
@@ -91,7 +91,7 @@
           task_id="build_node_pool_info_from_gcs_yaml"
       )(
           gcs_path=GCS_CONFIG_PATH,
-          dag_name="multi_host_nodepool_rollback",
+          dag_name=DAG_ID,
           is_prod=composer_env.is_prod_env(),
           machine_type=config.machine_version.value,
           tpu_topology=config.tpu_topology,

dags/tpu_observability/node_pool_status.py

@@ -91,7 +91,7 @@ def generate_problematic_node_location(
           task_id="build_node_pool_info_from_gcs_yaml"
       )(
           gcs_path=GCS_CONFIG_PATH,
-          dag_name="gke_node_pool_status",
+          dag_name=DAG_ID,
           is_prod=composer_env.is_prod_env(),
           machine_type=config.machine_version.value,
           tpu_topology=config.tpu_topology,

dags/tpu_observability/node_pool_ttr_disk_size.py

@@ -78,7 +78,7 @@
     with TaskGroup(group_id=f"v{config.tpu_version.value}"):
       node_pool_info = node_pool.build_node_pool_info_from_gcs_yaml(
           gcs_path=GCS_CONFIG_PATH,
-          dag_name="node_pool_ttr_disk_size",
+          dag_name=DAG_ID,
           is_prod=composer_env.is_prod_env(),
           machine_type=config.machine_version.value,
           tpu_topology=config.tpu_topology,

dags/tpu_observability/node_pool_ttr_update_label.py

@@ -77,7 +77,7 @@
           task_id="build_node_pool_info_from_gcs_yaml"
       )(
           gcs_path=GCS_CONFIG_PATH,
-          dag_name="node_pool_ttr_update_label",
+          dag_name=DAG_ID,
           is_prod=composer_env.is_prod_env(),
           machine_type=config.machine_version.value,
           tpu_topology=config.tpu_topology,