Feat: add jobset_healthiness_validation DAG

.github/workflows/dag-check.yml

@@ -3,7 +3,6 @@ name: DAG Check
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
 
   push:

.github/workflows/pyink-check.yml

@@ -2,11 +2,12 @@ name: Formatter
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
   push:
     branches: [master]
 
+  workflow_dispatch: {}
+
 jobs:
   format_check:
     runs-on: ubuntu-latest

.github/workflows/pylint-check.yml

@@ -2,12 +2,13 @@ name: Linter
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
 
   push:
     branches: [master]
 
+  workflow_dispatch: {}
+
 jobs:
   linting_check:
     runs-on: ubuntu-latest

.github/workflows/require-checklist.yml

@@ -2,10 +2,13 @@ name: Require Checklist
 on:
   pull_request:
     types: [opened, edited, synchronize]
+
+  workflow_dispatch: {}
+
 jobs:
   check_pr_body:
     runs-on: ubuntu-latest
     steps:
       - uses: mheap/require-checklist-action@v2
         with:
-          requireChecklist: false # If this is true and there are no checklists detected, the action will fail
+          requireChecklist: false # If this is true and there are no checklists detected, the action will fail

.github/workflows/unit-test.yml

@@ -3,7 +3,6 @@ name: Unit Test
 
 on:
   pull_request:
-    branches: [master]
     types: [opened, synchronize, edited]
 
   push:

dags/common/quarantined_tests.py

@@ -54,22 +54,36 @@ def match_quarantine_patterns(
   return False
 
 
+def is_ci_environment() -> bool:
+  """Checks if running in a CI environment.
+
+  Identifies environments where framework-specific runtime dependencies do not
+  exist. For example, GitHub Actions lacks the actual Airflow environment and
+  its components (such as database context or cloud credentials), requiring code
+  to bypass these dependencies and use mocks instead.
+
+  Returns:
+    bool: True if in CI (GitHub Actions), False otherwise.
+  """
+  return os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
+
+
 def safe_get_from_variable(key: str, default_var: str):
   """
-  Check whether the current runtime is GitHub Actions. Skip retrieving variables in GitHub Actions to avoid excessive log output.
+  Check whether the current runtime is GitHub Actions. Skip retrieving variables
+  in GitHub Actions to avoid excessive log output.
   """
   value = default_var
-  is_ci_env = os.getenv("GITHUB_ACTIONS", "false").lower() == "true"
-  if is_ci_env:
+  if is_ci_environment():
     logging.info("In GitHub Actions, skip getting variables")
   else:
     value = Variable.get(key, default_var=default_var)
   return value
 
 
 """
-The quarantine list is defined by a set of UNIX Shell Glob Patterns. 
-These patterns are used to match and quarantine tests. 
+The quarantine list is defined by a set of UNIX Shell Glob Patterns.
+These patterns are used to match and quarantine tests.
 
 The patterns are stored in the Airflow Variable named 'quarantine_patterns'.
 
@@ -90,9 +104,10 @@ def is_quarantined(test_name) -> bool:
     """
     Checks if a test is quarantined using both legacy and current methods.
 
-    The legacy method checks if `test_name` is present in `QuarantineTests.tests`.
-    The current method checks against a runtime quarantine list fetched from Airflow Variables
-    (key: 'quarantine_list') using `is_in_runtime_quarantine_list()`.
+    The legacy method checks if `test_name` is present in
+    `QuarantineTests.tests`. The current method checks against a runtime
+    quarantine list fetched from Airflow Variables (key: 'quarantine_list')
+    using `is_in_runtime_quarantine_list()`.
 
     The test is considered quarantined if it's found by either method.
     """

dags/common/scheduling_helper/scheduling_helper.py

@@ -72,6 +72,7 @@ class DayOfWeek(enum.Enum):
         "jobset_ttr_kill_process": dt.timedelta(minutes=90),
         "jobset_uptime_validation": dt.timedelta(minutes=90),
         "jobset_ttr_drain_restart": DefaultTimeout,
+        "jobset_healthiness_validation": dt.timedelta(minutes=90),
         "tpu_info_metrics_verification": DefaultTimeout,
         "jobset_ttr_node_reboot": dt.timedelta(minutes=90),
     },

dags/tpu_observability/jobset_healthiness_validation.py

@@ -0,0 +1,264 @@
+# 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 test "Jobset Suspended Healthiness" metric."""
+
+import datetime
+from datetime import timedelta
+
+from airflow import models
+from airflow.utils.trigger_rule import TriggerRule
+from airflow.utils.task_group import TaskGroup
+from airflow.models.baseoperator import chain
+
+from dags import composer_env
+from dags.tpu_observability.utils import jobset_util as jobset
+from dags.tpu_observability.utils import node_pool_util as node_pool
+from dags.common.task_group_with_timeout import TaskGroupWithTimeout
+from dags.tpu_observability.utils.jobset_util import Workload, JobSetHealthiness
+from dags.tpu_observability.configs.common import (
+    MachineConfigMap,
+    GCS_CONFIG_PATH,
+    GCS_JOBSET_CONFIG_PATH,
+)
+from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
+
+
+DAG_ID = "jobset_healthiness_validation"
+DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)
+SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID)
+
+FAIL_WORKLOAD = "python3 -c 'import logging; import sys; logging.error(\"Simulating Failure\"); sys.exit(1)'"
+SUCCESS_WORKLOAD = "python3 -c 'import logging; import sys; logging.info(\"Simulating Success\"); sys.exit(0)'"
+
+# Keyword arguments are generated dynamically at runtime (pylint does not
+# know this signature).
+with models.DAG(  # pylint: disable=unexpected-keyword-arg
+    dag_id=DAG_ID,
+    start_date=datetime.datetime(2025, 8, 10),
+    schedule=SCHEDULE if composer_env.is_prod_env() else None,
+    dagrun_timeout=DAGRUN_TIMEOUT,
+    catchup=False,
+    tags=[
+        "cloud-ml-auto-solutions",
+        "jobset",
+        "healthiness",
+        "tpu-obervability",
+        "TPU",
+        "v6e-16",
+    ],
+    description=(
+        "This DAG tests the 'Suspended' status of jobset healthiness by "
+        "comparing the number of 'suspended' replicas before and after "
+        "a jobset is running."
+    ),
+    doc_md="""
+      # JobSet Healthiness Test For the "Suspended" Status
+      ### Description
+      This DAG automates node-pool creation and validates JobSet healthiness
+      by examining replica-based metrics: Specified, Active, Ready,
+      Suspended, Succeeded, and Failed. It ensures the JobSet controller
+      accurately reports these states during startup, maintenance,
+      and failure scenarios.
+      ### Prerequisites
+      This test requires an existing cluster to run.
+      ### Procedures
+      First a node-pool is created. This test uses a State-Trigger-Observe pattern:
+      it triggers lifecycle transitions (e.g., suspension, failure or succeeded)
+      and verifies that GKE telemetry reflects these shifts. Using sensors,
+      the DAG polls for eventual consistency to account for ingestion latency,
+      dynamically matching runtime JobSet configurations against normalized monitoring
+      data types to ensure accurate state validation.
+      """,
+) as dag:
+  for machine in MachineConfigMap:
+    config = machine.value
+
+    # Keyword arguments are generated dynamically at runtime (pylint does not
+    # know this signature).
+    with TaskGroupWithTimeout(  # pylint: disable=unexpected-keyword-arg
+        group_id=f"v{config.tpu_version.value}",
+        timeout=timedelta(minutes=90),
+    ):
+      cluster_info = node_pool.build_node_pool_info_from_gcs_yaml(
+          gcs_path=GCS_CONFIG_PATH,
+          dag_name=DAG_ID,
+          is_prod=composer_env.is_prod_env(),
+          machine_type=config.machine_version.value,
+          tpu_topology=config.tpu_topology,
+      )
+
+      jobset_config = jobset.build_jobset_from_gcs_yaml(
+          gcs_path=GCS_JOBSET_CONFIG_PATH,
+          dag_name=DAG_ID,
+      )
+
+      selector = jobset.generate_node_pool_selector(DAG_ID)
+      jobset_name = jobset.generate_jobset_name(jobset_config.dag_id_prefix)
+
+      create_node_pool = node_pool.create.override(task_id="create_node_pool")(
+          node_pool=cluster_info,
+          node_pool_selector=selector,
+      )
+
+      startup = jobset.create_jobset_startup_tasks(
+          node_pool=cluster_info,
+          jobset_config=jobset_config,
+          jobset_name=jobset_name,
+          node_pool_selector=selector,
+          workload_type=Workload.JAX_TPU_BENCHMARK,
+      )
+
+      running_metrics = [
+          (JobSetHealthiness.SPECIFIED, jobset_config.replicas),
+          (JobSetHealthiness.ACTIVE, jobset_config.replicas),
+          (JobSetHealthiness.READY, jobset_config.replicas),
+          (JobSetHealthiness.FAILED, 0),
+          (JobSetHealthiness.SUCCEEDED, 0),
+          (JobSetHealthiness.SUSPENDED, 0),
+      ]
+      validate_running_tasks = []
+      for status, expected in running_metrics:
+        t = jobset.wait_for_jobset_metrics.override(
+            task_id=f"validate_running_wait_{status.value}"
+        )(
+            metric_name=status,
+            expected_value=expected,
+            node_pool=cluster_info,
+            jobset_name=jobset_name,
+        )
+        validate_running_tasks.append(t)
+
+      suspend_action = jobset.suspended_jobset.override(
+          task_id="suspend_jobset"
+      )(
+          node_pool=cluster_info,
+          jobset_config=jobset_config,
+          jobset_name=jobset_name,
+      )
+
+      suspended_metrics = [
+          (JobSetHealthiness.ACTIVE, 0),
+          (JobSetHealthiness.SUSPENDED, jobset_config.replicas),
+      ]
+      validate_suspended_tasks = []
+      for status, expected in suspended_metrics:
+        t = jobset.wait_for_jobset_metrics.override(
+            task_id=f"validate_suspended_wait_after_suspend_{status.value}"
+        )(
+            metric_name=status,
+            expected_value=expected,
+            node_pool=cluster_info,
+            jobset_name=jobset_name,
+        )
+        validate_suspended_tasks.append(t)
+
+      resume_action = jobset.resume_jobset.override(task_id="resume_jobset")(
+          node_pool=cluster_info,
+          jobset_config=jobset_config,
+          jobset_name=jobset_name,
+      )
+
+      cleanup_for_success = jobset.end_workload.override(
+          task_id="cleanup_before_success_injection"
+      )(
+          node_pool=cluster_info,
+          jobset_config=jobset_config,
+          jobset_name=jobset_name,
+      )
+
+      start_success_job = jobset.run_workload.override(
+          task_id="start_success_job"
+      )(
+          node_pool=cluster_info,
+          jobset_config=jobset_config,
+          jobset_name=jobset_name,
+          workload_type=SUCCESS_WORKLOAD,
+      )
+
+      validate_succeeded_metric = jobset.wait_for_jobset_metrics.override(
+          task_id="wait_for_succeeded_count"
+      )(
+          metric_name=JobSetHealthiness.SUCCEEDED,
+          expected_value=jobset_config.replicas,
+          node_pool=cluster_info,
+          jobset_name=jobset_name,
+      )
+
+      cleanup_for_failure = jobset.end_workload.override(
+          task_id="cleanup_before_failure_injection"
+      )(
+          node_pool=cluster_info,
+          jobset_config=jobset_config,
+          jobset_name=jobset_name,
+      )
+
+      start_fail_job = jobset.run_workload.override(task_id="start_fail_job")(
+          node_pool=cluster_info,
+          jobset_config=jobset_config,
+          jobset_name=jobset_name,
+          workload_type=FAIL_WORKLOAD,
+      )
+
+      validate_failed_metric = jobset.wait_for_jobset_metrics.override(
+          task_id="wait_for_failed_count"
+      )(
+          metric_name=JobSetHealthiness.FAILED,
+          expected_value=jobset_config.replicas,
+          node_pool=cluster_info,
+          jobset_name=jobset_name,
+      )
+
+      cleanup_workload = jobset.end_workload.override(
+          task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE
+      )(
+          node_pool=cluster_info,
+          jobset_config=jobset_config,
+          jobset_name=jobset_name,
+      ).as_teardown(
+          setups=startup.jobset_start_time
+      )
+
+      cleanup_node_pool = node_pool.delete.override(
+          task_id="cleanup_node_pool", trigger_rule=TriggerRule.ALL_DONE
+      )(node_pool=cluster_info).as_teardown(
+          setups=create_node_pool,
+      )
+
+      chain(
+          selector,
+          jobset_name,
+          create_node_pool,
+          *startup.tasks,
+          *validate_running_tasks,
+          suspend_action,
+          *validate_suspended_tasks,
+          resume_action,
+          cleanup_for_success,
+      )
+
+      chain(
+          cleanup_for_success,
+          start_success_job,
+          validate_succeeded_metric,
+          cleanup_for_failure,
+      )
+
+      chain(
+          cleanup_for_failure,
+          start_fail_job,
+          validate_failed_metric,
+          cleanup_workload,
+          cleanup_node_pool,
+      )