From 0791ccb0427f21fb993e2fc5a807b447c7761bc1 Mon Sep 17 00:00:00 2001 From: Jim Tseng Date: Thu, 18 Jun 2026 08:37:15 +0800 Subject: [PATCH 1/7] fix: Skip GCS configuration download in GitHub Actions environment During GitHub Actions (CI) pipeline execution, native Airflow methods and cloud-dependent functions cause blocking errors due to the lack of an active database connection or cloud permissions. --- dags/common/quarantined_tests.py | 31 +++++++++++++++------ dags/tpu_observability/utils/jobset_util.py | 24 ++++++++-------- xlml/apis/gcs.py | 22 +++++++++++---- 3 files changed, 51 insertions(+), 26 deletions(-) diff --git a/dags/common/quarantined_tests.py b/dags/common/quarantined_tests.py index 6d44fcc33..d27a6cfd4 100644 --- a/dags/common/quarantined_tests.py +++ b/dags/common/quarantined_tests.py @@ -54,13 +54,27 @@ 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) @@ -68,8 +82,8 @@ def safe_get_from_variable(key: str, default_var: str): """ -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. """ diff --git a/dags/tpu_observability/utils/jobset_util.py b/dags/tpu_observability/utils/jobset_util.py index 112751ae8..8e8a82591 100644 --- a/dags/tpu_observability/utils/jobset_util.py +++ b/dags/tpu_observability/utils/jobset_util.py @@ -282,18 +282,18 @@ class JobSet: recreating pods. Defaults to False. """ - namespace: str - max_restarts: int - replicated_job_name: str - replicas: int - backoff_limit: int - completions: int - parallelism: int - tpu_accelerator_type: str - tpu_topology: str - container_name: str - image: str - tpu_cores_per_pod: int + namespace: str = "" + max_restarts: int = -1 + replicated_job_name: str = "" + replicas: int = -1 + backoff_limit: int = -1 + completions: int = -1 + parallelism: int = -1 + tpu_accelerator_type: str = "" + tpu_topology: str = "" + container_name: str = "" + image: str = "" + tpu_cores_per_pod: int = -1 privileged: bool = False dag_id_prefix: str = "" delay_recovery: bool = False diff --git a/xlml/apis/gcs.py b/xlml/apis/gcs.py index e2f483288..cab105ce6 100644 --- a/xlml/apis/gcs.py +++ b/xlml/apis/gcs.py @@ -14,16 +14,17 @@ """Functions for GCS Bucket""" +from absl import logging import os import re import tempfile from typing import List +import yaml -from absl import logging from airflow.decorators import task from airflow.hooks.subprocess import SubprocessHook from airflow.providers.google.cloud.operators.gcs import GCSHook -import yaml +from dags.common.quarantined_tests import is_ci_environment def obtain_file_list(gcs_path: str) -> List[str]: @@ -97,6 +98,14 @@ def load_yaml_from_gcs(gcs_path: str) -> dict: """Loads and parses the DAG configuration YAML file from GCS.""" logging.info(f"Attempting to load config from: {gcs_path}") + # Workflows triggered by fork PRs do not have access to upstream secrets. + # Bypass GCS loading strictly to allow the CI workflow to pass successfully. + if is_ci_environment(): + logging.info( + "In GitHub Actions, skip load_yaml_from_gcs and return an empty dict." + ) + return {} + if not gcs_path.startswith("gs://"): raise ValueError( f"Invalid GCS path: '{gcs_path}'. Path must start with 'gs://'." @@ -106,8 +115,8 @@ def load_yaml_from_gcs(gcs_path: str) -> dict: gcs_path.lower().endswith(".yaml") or gcs_path.lower().endswith(".yml") ): logging.warning( - f"GCS path '{gcs_path}' does not have a typical YAML extension (.yaml or .yml). " - "Proceeding, but be aware this might not be a YAML file." + f"GCS path '{gcs_path}' does not have a typical YAML extension (.yaml " + "or .yml). Proceeding, but be aware this might not be a YAML file." ) with tempfile.TemporaryDirectory() as tmpdir: @@ -123,8 +132,9 @@ def load_yaml_from_gcs(gcs_path: str) -> dict: if not os.path.exists(temp_file_path): logging.error( - f"gcloud storage cp command completed, but '{temp_file_path}' was not created. " - "This often means the copy failed. Check gcloud storage stdout/stderr in logs." + f"gcloud storage cp command completed, but '{temp_file_path}' " + "was not created. This often means the copy failed. " + "Check gcloud storage stdout/stderr in logs." ) raise FileNotFoundError( f"[Errno 2] Failed to download file from GCS path: {gcs_path}" From fdcbae901723e5672139aabb2b1e2e2f01bdebc3 Mon Sep 17 00:00:00 2001 From: Alfred Yu Date: Tue, 12 Aug 2025 10:03:26 +0800 Subject: [PATCH 2/7] [cienet-private] Enable workflow checks for the primary working branches --- .github/workflows/dag-check.yml | 1 - .github/workflows/pyink-check.yml | 3 +- .github/workflows/pylint-check.yml | 3 +- .github/workflows/require-checklist.yml | 5 ++- .github/workflows/unit-test.yml | 1 - scripts/code-style.sh | 41 +++++++++++++++++++------ 6 files changed, 40 insertions(+), 14 deletions(-) diff --git a/.github/workflows/dag-check.yml b/.github/workflows/dag-check.yml index b85625224..c1d3e7cf3 100644 --- a/.github/workflows/dag-check.yml +++ b/.github/workflows/dag-check.yml @@ -3,7 +3,6 @@ name: DAG Check on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: diff --git a/.github/workflows/pyink-check.yml b/.github/workflows/pyink-check.yml index ea7a8f250..c5cb1cd83 100644 --- a/.github/workflows/pyink-check.yml +++ b/.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 diff --git a/.github/workflows/pylint-check.yml b/.github/workflows/pylint-check.yml index 5e23d2812..02f4cb31e 100644 --- a/.github/workflows/pylint-check.yml +++ b/.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 diff --git a/.github/workflows/require-checklist.yml b/.github/workflows/require-checklist.yml index d15d19d99..4da288575 100644 --- a/.github/workflows/require-checklist.yml +++ b/.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 \ No newline at end of file + requireChecklist: false # If this is true and there are no checklists detected, the action will fail diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index 4f1723cb2..b771ca7e9 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -3,7 +3,6 @@ name: Unit Test on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: diff --git a/scripts/code-style.sh b/scripts/code-style.sh index 36cfa13e9..ea50709ac 100755 --- a/scripts/code-style.sh +++ b/scripts/code-style.sh @@ -19,14 +19,37 @@ set -e FOLDERS_TO_FORMAT=("dags" "xlml") -for folder in "${FOLDERS_TO_FORMAT[@]}" -do - pyink "$folder" --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80 --check --diff -done - -for folder in "${FOLDERS_TO_FORMAT[@]}" -do - pylint "./$folder" --fail-under=9.6 -done +HEAD_SHA="$(git rev-parse HEAD)" +BASE_BRANCH="dev" + +if ! git rev-parse --verify "$BASE_BRANCH" >/dev/null 2>&1; then + git fetch origin "$BASE_BRANCH":"$BASE_BRANCH" || { + echo "[code-style] base branch '$BASE_BRANCH' not found, skip diff-based check." + exit 0 + } +fi + +CHANGED_PY_FILES="$( + git diff --name-only --diff-filter=ACM "${BASE_BRANCH}" "${HEAD_SHA}" \ + | grep '\.py$' \ + | while read -r f; do + for folder in "${FOLDERS_TO_FORMAT[@]}"; do + if [[ "$f" == "$folder/"* ]]; then + echo "$f" + break + fi + done + done \ + | sort -u +)" + +if [[ -z "${CHANGED_PY_FILES}" ]]; then + echo "[pre-push hook] no changed files detected between ${HEAD_SHA} and ${BASE_BRANCH}" + exit 1 +fi + +pyink ${CHANGED_PY_FILES} --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80 --check --diff + +pylint ${CHANGED_PY_FILES} --fail-under=9.6 --disable=E1123 echo "Successfully clean up all codes." From 44a9518df349f0b6dff14956d8c421f489fe2ce2 Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Wed, 6 May 2026 08:44:26 +0000 Subject: [PATCH 3/7] feat(tpu): add JobSet healthiness validation DAG and enhance utils --- .../scheduling_helper/scheduling_helper.py | 1 + .../jobset_healthiness_validation.py | 220 ++++++++++++++++++ dags/tpu_observability/utils/jobset_util.py | 11 + 3 files changed, 232 insertions(+) create mode 100644 dags/tpu_observability/jobset_healthiness_validation.py diff --git a/dags/common/scheduling_helper/scheduling_helper.py b/dags/common/scheduling_helper/scheduling_helper.py index 8aef11b56..495d547db 100644 --- a/dags/common/scheduling_helper/scheduling_helper.py +++ b/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), }, diff --git a/dags/tpu_observability/jobset_healthiness_validation.py b/dags/tpu_observability/jobset_healthiness_validation.py new file mode 100644 index 000000000..31979cffb --- /dev/null +++ b/dags/tpu_observability/jobset_healthiness_validation.py @@ -0,0 +1,220 @@ +# 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 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.tpu_observability.utils.jobset_util import Workload, ReplicatedJobStatus +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)'" + +# 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 the process of creating node-pools, ensuring the + correct number of "Suspended" replicas appear, then launching a jobset on + multiple replicas to ensure the correct number begin running. + ### Prerequisites + This test requires an existing cluster to run. + ### Procedures + First a node-pool is created. The validation test is then run to + check if the number of "Suspended" replicas is 0. Once the jobset is + running the jobs should quickly enter the "Ready" state. Then using + command to suspend entire jobset. The number of found replicas is + tested against the number of replicas which should be "Suspended". + If they match the DAG is a success. + """, +) as dag: + for machine in MachineConfigMap: + config = machine.value + + # Keyword arguments are generated dynamically at runtime (pylint does not + # know this signature). + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id=f"v{config.tpu_version.value}" + ): + selector = jobset.generate_node_pool_selector( + "jobset-healthiness-validation" + ) + + jobset_config = jobset.build_jobset_from_gcs_yaml( + gcs_path=GCS_JOBSET_CONFIG_PATH, + 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=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")( + node_pool=cluster_info, + ) + + startup = jobset.create_jobset_startup_group( + node_pool=cluster_info, + jobset_config=jobset_config, + workload_type=Workload.JAX_TPU_BENCHMARK, + ) + + with TaskGroup(group_id="validate_running_metrics") as validate_running: + running_metrics = [ + (ReplicatedJobStatus.SPECIFIED, "USE_CONFIG_REPLICAS"), + (ReplicatedJobStatus.ACTIVE, "USE_CONFIG_REPLICAS"), + (ReplicatedJobStatus.READY, "USE_CONFIG_REPLICAS"), + (ReplicatedJobStatus.FAILED, 0), + (ReplicatedJobStatus.SUCCEEDED, 0), + (ReplicatedJobStatus.SUSPENDED, 0), + ] + for status, expected in running_metrics: + jobset.wait_for_jobset_metrics.override( + task_id=f"wait_{status.value}" + )( + metric_name=status, + expected_value=expected, + node_pool=cluster_info, + jobset_config=jobset_config, + ) + + suspend_action = jobset.suspended_jobset.override( + task_id="suspend_jobset" + )( + node_pool=cluster_info, + jobset_config=jobset_config, + ) + + with TaskGroup( + group_id="validate_suspended_metrics" + ) as validate_suspended: + suspended_metrics = [ + (ReplicatedJobStatus.ACTIVE, 0), + (ReplicatedJobStatus.SUSPENDED, "USE_CONFIG_REPLICAS"), + ] + for status, expected in suspended_metrics: + jobset.wait_for_jobset_metrics.override( + task_id=f"wait_after_suspend_{status.value}" + )( + metric_name=status, + expected_value=expected, + node_pool=cluster_info, + jobset_config=jobset_config, + ) + + resume_action = jobset.resume_jobset.override(task_id="resume_jobset")( + node_pool=cluster_info, + jobset_config=jobset_config, + ) + + with TaskGroup(group_id="inject_and_validate_failure") as failure_test: + cleanup_for_failure = jobset.end_workload.override( + task_id="cleanup_before_failure_injection" + )( + node_pool=cluster_info, + jobset_config=jobset_config, + ) + + start_fail_job = jobset.run_workload.override(task_id="start_fail_job")( + node_pool=cluster_info, + jobset_config=jobset_config, + workload_type=FAIL_WORKLOAD, + ) + + validate_failed_metric = jobset.wait_for_jobset_metrics.override( + task_id="wait_for_failed_count" + )( + metric_name=ReplicatedJobStatus.FAILED, + expected_value="USE_CONFIG_REPLICAS", + node_pool=cluster_info, + jobset_config=jobset_config, + ) + + chain(cleanup_for_failure, start_fail_job, validate_failed_metric) + + cleanup_workload = jobset.end_workload.override( + task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE + )( + node_pool=cluster_info, + jobset_config=jobset_config, + ).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_config, + cluster_info, + create_node_pool, + startup.task_group, + validate_running, + suspend_action, + validate_suspended, + resume_action, + failure_test, + cleanup_workload, + cleanup_node_pool, + ) diff --git a/dags/tpu_observability/utils/jobset_util.py b/dags/tpu_observability/utils/jobset_util.py index 8e8a82591..ec20f002b 100644 --- a/dags/tpu_observability/utils/jobset_util.py +++ b/dags/tpu_observability/utils/jobset_util.py @@ -367,6 +367,17 @@ class JobSetStartupOutput: jobset_start_time: XComArg +class ReplicatedJobStatus(enum.Enum): + """Defines status of a replicated job.""" + + READY = "ready" + SUSPENDED = "suspended" + ACTIVE = "active" + FAILED = "failed" + SUCCEEDED = "succeeded" + SPECIFIED = "specified" + + class Command: """ A collection of static methods to generate Kubernetes and gcloud commands. From b7f2b53e75928f9a9b6bd2fcfa627ed435476042 Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Thu, 7 May 2026 01:31:37 +0000 Subject: [PATCH 4/7] Fix: Add validation of success status --- .../jobset_healthiness_validation.py | 73 +++++++++++++------ dags/tpu_observability/utils/jobset_util.py | 7 +- 2 files changed, 57 insertions(+), 23 deletions(-) diff --git a/dags/tpu_observability/jobset_healthiness_validation.py b/dags/tpu_observability/jobset_healthiness_validation.py index 31979cffb..a77b9fc6e 100644 --- a/dags/tpu_observability/jobset_healthiness_validation.py +++ b/dags/tpu_observability/jobset_healthiness_validation.py @@ -24,7 +24,7 @@ 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.tpu_observability.utils.jobset_util import Workload, ReplicatedJobStatus +from dags.tpu_observability.utils.jobset_util import Workload, JobSetHealthiness from dags.tpu_observability.configs.common import ( MachineConfigMap, GCS_CONFIG_PATH, @@ -38,6 +38,7 @@ 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). @@ -63,18 +64,20 @@ doc_md=""" # JobSet Healthiness Test For the "Suspended" Status ### Description - This DAG automates the process of creating node-pools, ensuring the - correct number of "Suspended" replicas appear, then launching a jobset on - multiple replicas to ensure the correct number begin running. + 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. The validation test is then run to - check if the number of "Suspended" replicas is 0. Once the jobset is - running the jobs should quickly enter the "Ready" state. Then using - command to suspend entire jobset. The number of found replicas is - tested against the number of replicas which should be "Suspended". - If they match the DAG is a success. + 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: @@ -110,7 +113,7 @@ node_pool=cluster_info, ) - startup = jobset.create_jobset_startup_group( + startup = jobset.create_jobset_startup_tasks( node_pool=cluster_info, jobset_config=jobset_config, workload_type=Workload.JAX_TPU_BENCHMARK, @@ -118,12 +121,12 @@ with TaskGroup(group_id="validate_running_metrics") as validate_running: running_metrics = [ - (ReplicatedJobStatus.SPECIFIED, "USE_CONFIG_REPLICAS"), - (ReplicatedJobStatus.ACTIVE, "USE_CONFIG_REPLICAS"), - (ReplicatedJobStatus.READY, "USE_CONFIG_REPLICAS"), - (ReplicatedJobStatus.FAILED, 0), - (ReplicatedJobStatus.SUCCEEDED, 0), - (ReplicatedJobStatus.SUSPENDED, 0), + (JobSetHealthiness.SPECIFIED, "USE_CONFIG_REPLICAS"), + (JobSetHealthiness.ACTIVE, "USE_CONFIG_REPLICAS"), + (JobSetHealthiness.READY, "USE_CONFIG_REPLICAS"), + (JobSetHealthiness.FAILED, 0), + (JobSetHealthiness.SUCCEEDED, 0), + (JobSetHealthiness.SUSPENDED, 0), ] for status, expected in running_metrics: jobset.wait_for_jobset_metrics.override( @@ -146,8 +149,8 @@ group_id="validate_suspended_metrics" ) as validate_suspended: suspended_metrics = [ - (ReplicatedJobStatus.ACTIVE, 0), - (ReplicatedJobStatus.SUSPENDED, "USE_CONFIG_REPLICAS"), + (JobSetHealthiness.ACTIVE, 0), + (JobSetHealthiness.SUSPENDED, "USE_CONFIG_REPLICAS"), ] for status, expected in suspended_metrics: jobset.wait_for_jobset_metrics.override( @@ -164,6 +167,33 @@ jobset_config=jobset_config, ) + with TaskGroup(group_id="inject_and_validate_success") as success_test: + cleanup_for_success = jobset.end_workload.override( + task_id="cleanup_before_success_injection" + )( + node_pool=cluster_info, + jobset_config=jobset_config, + ) + + start_success_job = jobset.run_workload.override( + task_id="start_success_job" + )( + node_pool=cluster_info, + jobset_config=jobset_config, + 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="USE_CONFIG_REPLICAS", + node_pool=cluster_info, + jobset_config=jobset_config, + ) + + chain(cleanup_for_success, start_success_job, validate_succeeded_metric) + with TaskGroup(group_id="inject_and_validate_failure") as failure_test: cleanup_for_failure = jobset.end_workload.override( task_id="cleanup_before_failure_injection" @@ -181,7 +211,7 @@ validate_failed_metric = jobset.wait_for_jobset_metrics.override( task_id="wait_for_failed_count" )( - metric_name=ReplicatedJobStatus.FAILED, + metric_name=JobSetHealthiness.FAILED, expected_value="USE_CONFIG_REPLICAS", node_pool=cluster_info, jobset_config=jobset_config, @@ -209,11 +239,12 @@ jobset_config, cluster_info, create_node_pool, - startup.task_group, + *startup.tasks, validate_running, suspend_action, validate_suspended, resume_action, + success_test, failure_test, cleanup_workload, cleanup_node_pool, diff --git a/dags/tpu_observability/utils/jobset_util.py b/dags/tpu_observability/utils/jobset_util.py index ec20f002b..d4983e3a0 100644 --- a/dags/tpu_observability/utils/jobset_util.py +++ b/dags/tpu_observability/utils/jobset_util.py @@ -367,8 +367,11 @@ class JobSetStartupOutput: jobset_start_time: XComArg -class ReplicatedJobStatus(enum.Enum): - """Defines status of a replicated job.""" +class JobSetHealthiness(enum.Enum): + """Enumeration of JobSet healthiness metrics based on replica states. + The values represent the metric name suffix used in Prometheus/GKE + observability to track the count of replicas in each lifecycle state. + """ READY = "ready" SUSPENDED = "suspended" From 234847f9f1bf63c9e4e5a5811fafb2199478b13e Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Tue, 26 May 2026 03:39:45 +0000 Subject: [PATCH 5/7] fix --- dags/tpu_observability/utils/jobset_util.py | 60 +++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/dags/tpu_observability/utils/jobset_util.py b/dags/tpu_observability/utils/jobset_util.py index d4983e3a0..fa0c5aee6 100644 --- a/dags/tpu_observability/utils/jobset_util.py +++ b/dags/tpu_observability/utils/jobset_util.py @@ -913,6 +913,66 @@ def operate_pod( return TimeUtil.from_datetime(current_time_utc) +@task +def suspended_jobset(node_pool: node_pool_info, jobset_config: JobSet): + """ + Suspend a jobset from the GKE cluster. + + This task executes a bash script to: + 1. Authenticate `gcloud` with the specified GKE cluster. + 2. Suspend a JobSet. + + Args: + node_pool: Configuration object with cluster details. + jobset_name: The name of the JobSet to delete. + """ + + with tempfile.NamedTemporaryFile() as temp_config_file: + env = os.environ.copy() + env["KUBECONFIG"] = temp_config_file.name + + cmd = " && ".join([ + Command.get_credentials_command(node_pool), + Command.k8s_suspend_jobset_command( + temp_config_file.name, + jobset_config.jobset_name, + jobset_config.namespace, + ), + ]) + + subprocess.run_exec(cmd, env=env) + + +@task +def resume_jobset(node_pool: node_pool_info, jobset_config: JobSet): + """ + Resume a jobset from the GKE cluster. + + This task executes a bash script to: + 1. Authenticate `gcloud` with the specified GKE cluster. + 2. Resume a JobSet. + + Args: + node_pool: Configuration object with cluster details. + jobset_name: The name of the JobSet to delete. + """ + + with tempfile.NamedTemporaryFile() as temp_config_file: + env = os.environ.copy() + env["KUBECONFIG"] = temp_config_file.name + + cmd = " && ".join([ + Command.get_credentials_command(node_pool), + Command.k8s_resume_jobset_command( + temp_config_file.name, + jobset_config.jobset_name, + jobset_config.namespace, + ), + ]) + + subprocess.run_exec(cmd, env=env) + + @task.sensor(poke_interval=30, timeout=900, mode="poke") def wait_for_jobset_started( node_pool: node_pool_info, From 99acfc4ebf672f5875ce002576cf92f64325040a Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Tue, 26 May 2026 07:05:18 +0000 Subject: [PATCH 6/7] refactor: encapsulate JobSet metric mapping into JobSetHealthiness enum --- .../jobset_healthiness_validation.py | 54 ++++++----- dags/tpu_observability/utils/jobset_util.py | 91 ++++++++++++++++++- 2 files changed, 116 insertions(+), 29 deletions(-) diff --git a/dags/tpu_observability/jobset_healthiness_validation.py b/dags/tpu_observability/jobset_healthiness_validation.py index a77b9fc6e..8ef607fa9 100644 --- a/dags/tpu_observability/jobset_healthiness_validation.py +++ b/dags/tpu_observability/jobset_healthiness_validation.py @@ -88,42 +88,40 @@ with TaskGroup( # pylint: disable=unexpected-keyword-arg group_id=f"v{config.tpu_version.value}" ): - selector = jobset.generate_node_pool_selector( - "jobset-healthiness-validation" + 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, - 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=DAG_ID, - is_prod=composer_env.is_prod_env(), - machine_type=config.machine_version.value, - tpu_topology=config.tpu_topology, - node_pool_selector=selector, - ) + 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, ) with TaskGroup(group_id="validate_running_metrics") as validate_running: running_metrics = [ - (JobSetHealthiness.SPECIFIED, "USE_CONFIG_REPLICAS"), - (JobSetHealthiness.ACTIVE, "USE_CONFIG_REPLICAS"), - (JobSetHealthiness.READY, "USE_CONFIG_REPLICAS"), + (JobSetHealthiness.SPECIFIED, jobset_config.replicas), + (JobSetHealthiness.ACTIVE, jobset_config.replicas), + (JobSetHealthiness.READY, jobset_config.replicas), (JobSetHealthiness.FAILED, 0), (JobSetHealthiness.SUCCEEDED, 0), (JobSetHealthiness.SUSPENDED, 0), @@ -135,7 +133,7 @@ metric_name=status, expected_value=expected, node_pool=cluster_info, - jobset_config=jobset_config, + jobset_name=jobset_name, ) suspend_action = jobset.suspended_jobset.override( @@ -143,6 +141,7 @@ )( node_pool=cluster_info, jobset_config=jobset_config, + jobset_name=jobset_name, ) with TaskGroup( @@ -150,7 +149,7 @@ ) as validate_suspended: suspended_metrics = [ (JobSetHealthiness.ACTIVE, 0), - (JobSetHealthiness.SUSPENDED, "USE_CONFIG_REPLICAS"), + (JobSetHealthiness.SUSPENDED, jobset_config.replicas), ] for status, expected in suspended_metrics: jobset.wait_for_jobset_metrics.override( @@ -159,12 +158,13 @@ metric_name=status, expected_value=expected, node_pool=cluster_info, - jobset_config=jobset_config, + jobset_name=jobset_name, ) resume_action = jobset.resume_jobset.override(task_id="resume_jobset")( node_pool=cluster_info, jobset_config=jobset_config, + jobset_name=jobset_name, ) with TaskGroup(group_id="inject_and_validate_success") as success_test: @@ -173,6 +173,7 @@ )( node_pool=cluster_info, jobset_config=jobset_config, + jobset_name=jobset_name, ) start_success_job = jobset.run_workload.override( @@ -180,6 +181,7 @@ )( node_pool=cluster_info, jobset_config=jobset_config, + jobset_name=jobset_name, workload_type=SUCCESS_WORKLOAD, ) @@ -187,9 +189,9 @@ task_id="wait_for_succeeded_count" )( metric_name=JobSetHealthiness.SUCCEEDED, - expected_value="USE_CONFIG_REPLICAS", + expected_value=jobset_config.replicas, node_pool=cluster_info, - jobset_config=jobset_config, + jobset_name=jobset_name, ) chain(cleanup_for_success, start_success_job, validate_succeeded_metric) @@ -200,11 +202,13 @@ )( 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, ) @@ -212,9 +216,9 @@ task_id="wait_for_failed_count" )( metric_name=JobSetHealthiness.FAILED, - expected_value="USE_CONFIG_REPLICAS", + expected_value=jobset_config.replicas, node_pool=cluster_info, - jobset_config=jobset_config, + jobset_name=jobset_name, ) chain(cleanup_for_failure, start_fail_job, validate_failed_metric) @@ -224,6 +228,7 @@ )( node_pool=cluster_info, jobset_config=jobset_config, + jobset_name=jobset_name, ).as_teardown( setups=startup.jobset_start_time ) @@ -236,8 +241,7 @@ chain( selector, - jobset_config, - cluster_info, + jobset_name, create_node_pool, *startup.tasks, validate_running, diff --git a/dags/tpu_observability/utils/jobset_util.py b/dags/tpu_observability/utils/jobset_util.py index fa0c5aee6..4388ecb0c 100644 --- a/dags/tpu_observability/utils/jobset_util.py +++ b/dags/tpu_observability/utils/jobset_util.py @@ -380,6 +380,18 @@ class JobSetHealthiness(enum.Enum): SUCCEEDED = "succeeded" SPECIFIED = "specified" + @property + def healthiness_metric_type(self) -> str: + mapping = { + "ready": "prometheus.googleapis.com/kube_jobset_ready_replicas/gauge", + "active": "prometheus.googleapis.com/kube_jobset_active_replicas/gauge", + "suspended": "prometheus.googleapis.com/kube_jobset_suspended_replicas/gauge", + "succeeded": "prometheus.googleapis.com/kube_jobset_succeeded_replicas/gauge", + "failed": "prometheus.googleapis.com/kube_jobset_failed_replicas/gauge", + "specified": "prometheus.googleapis.com/kube_jobset_specified_replicas/gauge", + } + return mapping.get(self.value) + class Command: """ @@ -914,7 +926,9 @@ def operate_pod( @task -def suspended_jobset(node_pool: node_pool_info, jobset_config: JobSet): +def suspended_jobset( + node_pool: node_pool_info, jobset_config: JobSet, jobset_name: str +): """ Suspend a jobset from the GKE cluster. @@ -935,7 +949,7 @@ def suspended_jobset(node_pool: node_pool_info, jobset_config: JobSet): Command.get_credentials_command(node_pool), Command.k8s_suspend_jobset_command( temp_config_file.name, - jobset_config.jobset_name, + jobset_name, jobset_config.namespace, ), ]) @@ -944,7 +958,9 @@ def suspended_jobset(node_pool: node_pool_info, jobset_config: JobSet): @task -def resume_jobset(node_pool: node_pool_info, jobset_config: JobSet): +def resume_jobset( + node_pool: node_pool_info, jobset_config: JobSet, jobset_name: str +): """ Resume a jobset from the GKE cluster. @@ -965,7 +981,7 @@ def resume_jobset(node_pool: node_pool_info, jobset_config: JobSet): Command.get_credentials_command(node_pool), Command.k8s_resume_jobset_command( temp_config_file.name, - jobset_config.jobset_name, + jobset_name, jobset_config.namespace, ), ]) @@ -1055,6 +1071,73 @@ def verify_recovery_duration(start_time: TimeUtil, end_time: TimeUtil): "The 'jobset_time_to_recover' metric requires > 60s to be recorded. " "Failing fast to avoid waiting for a missing metric." ) +@task.sensor(poke_interval=60, timeout=3600, mode="poke") +def wait_for_jobset_metrics( + metric_name: JobSetHealthiness, + expected_value: int, + node_pool: node_pool_info, + jobset_name: str, + start_time: TimeUtil = None, +) -> bool: + """Polls Cloud Monitoring for a specific JobSet replicated job metric. + + A sensor task which polls various JobSet metrics (e.g., ready, active, + failed) every 60 seconds for 60 minutes. This sensor handles dynamic + scaling expectations and ensures compatibility with different data + types returned by the Prometheus exporter. + + Args: + metric_name (JobSetHealthiness): The status of the replicated job to + monitor. + expected_value (int): The target value to wait for. + node_pool (Info): An instance of the Info class containing GKE cluster + metadata. + jobset_name: The name of the JobSet. + start_time (TimeUtil, optional): The UTC timestamp to start polling + from. If not provided, defaults to 60 minutes before the current time. + + Returns: + bool: True if the current metric value matches the expected value, + False otherwise. + """ + + metric_type = metric_name.healthiness_metric_type + name_str = metric_name.value + + query_start = ( + start_time if start_time else TimeUtil.now() - timedelta(minutes=60) + ) + + time_series = list_time_series( + project_id=node_pool.project_id, + filter_str=( + f'metric.type="{metric_type}" ' + f'resource.type="prometheus_target" ' + f'resource.labels.cluster="{node_pool.cluster_name}" ' + f'metric.labels.jobset_name="{jobset_name}"' + ), + start_time=query_start, + end_time=TimeUtil.now(), + ) + + if not time_series or len(time_series) == 0 or not time_series[0].points: + return False + + point_value = time_series[0].points[0].value + if ( + hasattr(point_value, "double_value") + and point_value.double_value is not None + ): + latest_value = point_value.double_value + else: + latest_value = float(point_value.int64_value) + + logging.info( + f"Metric {name_str} for JobSet {jobset_name}: " + f"current={latest_value}, expected={expected_value}" + ) + + return float(latest_value) == float(expected_value) @task.sensor(poke_interval=60, timeout=3600, mode="poke", retries=0) From c5f8abf5353d5a0aaef2767ff0eb5fa1dab72e13 Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Mon, 15 Jun 2026 02:51:37 +0000 Subject: [PATCH 7/7] fix: Apply TaskGroupWithTimeout into DAG. --- .../jobset_healthiness_validation.py | 189 +++++++++--------- dags/tpu_observability/utils/jobset_util.py | 25 ++- 2 files changed, 116 insertions(+), 98 deletions(-) diff --git a/dags/tpu_observability/jobset_healthiness_validation.py b/dags/tpu_observability/jobset_healthiness_validation.py index 8ef607fa9..98924a545 100644 --- a/dags/tpu_observability/jobset_healthiness_validation.py +++ b/dags/tpu_observability/jobset_healthiness_validation.py @@ -15,6 +15,7 @@ """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 @@ -24,6 +25,7 @@ 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, @@ -85,8 +87,9 @@ # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). - with TaskGroup( # pylint: disable=unexpected-keyword-arg - group_id=f"v{config.tpu_version.value}" + 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, @@ -117,24 +120,25 @@ workload_type=Workload.JAX_TPU_BENCHMARK, ) - with TaskGroup(group_id="validate_running_metrics") as validate_running: - 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), - ] - for status, expected in running_metrics: - jobset.wait_for_jobset_metrics.override( - task_id=f"wait_{status.value}" - )( - metric_name=status, - expected_value=expected, - node_pool=cluster_info, - jobset_name=jobset_name, - ) + 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" @@ -144,22 +148,21 @@ jobset_name=jobset_name, ) - with TaskGroup( - group_id="validate_suspended_metrics" - ) as validate_suspended: - suspended_metrics = [ - (JobSetHealthiness.ACTIVE, 0), - (JobSetHealthiness.SUSPENDED, jobset_config.replicas), - ] - for status, expected in suspended_metrics: - jobset.wait_for_jobset_metrics.override( - task_id=f"wait_after_suspend_{status.value}" - )( - metric_name=status, - expected_value=expected, - node_pool=cluster_info, - 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, @@ -167,61 +170,55 @@ jobset_name=jobset_name, ) - with TaskGroup(group_id="inject_and_validate_success") as success_test: - 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_success = jobset.end_workload.override( + task_id="cleanup_before_success_injection" + )( + node_pool=cluster_info, + jobset_config=jobset_config, + jobset_name=jobset_name, + ) - chain(cleanup_for_success, start_success_job, validate_succeeded_metric) + 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, + ) - with TaskGroup(group_id="inject_and_validate_failure") as failure_test: - 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, - ) + 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, + ) - 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, - ) + 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, + ) - 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, - ) + 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, + ) - chain(cleanup_for_failure, start_fail_job, validate_failed_metric) + 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 @@ -244,12 +241,24 @@ jobset_name, create_node_pool, *startup.tasks, - validate_running, + *validate_running_tasks, suspend_action, - validate_suspended, + *validate_suspended_tasks, resume_action, - success_test, - failure_test, + 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, ) diff --git a/dags/tpu_observability/utils/jobset_util.py b/dags/tpu_observability/utils/jobset_util.py index 4388ecb0c..d6dd0fa8d 100644 --- a/dags/tpu_observability/utils/jobset_util.py +++ b/dags/tpu_observability/utils/jobset_util.py @@ -35,7 +35,9 @@ from airflow.models.baseoperator import chain from airflow.models.xcom_arg import XComArg from airflow.sensors.base import PokeReturnValue +from google.cloud import monitoring_v3 from google.cloud.monitoring_v3 import types +from google.protobuf.message import Message from websocket import WebSocketConnectionClosedException from dags.tpu_observability.utils import subprocess_util as subprocess @@ -1071,6 +1073,8 @@ def verify_recovery_duration(start_time: TimeUtil, end_time: TimeUtil): "The 'jobset_time_to_recover' metric requires > 60s to be recorded. " "Failing fast to avoid waiting for a missing metric." ) + + @task.sensor(poke_interval=60, timeout=3600, mode="poke") def wait_for_jobset_metrics( metric_name: JobSetHealthiness, @@ -1120,17 +1124,22 @@ def wait_for_jobset_metrics( end_time=TimeUtil.now(), ) - if not time_series or len(time_series) == 0 or not time_series[0].points: + if not time_series or not time_series[0].points: return False - point_value = time_series[0].points[0].value - if ( - hasattr(point_value, "double_value") - and point_value.double_value is not None - ): - latest_value = point_value.double_value + point = time_series[0].points[0] + msg = monitoring_v3.TypedValue.pb(point.value) + + if isinstance(msg, Message): + match msg.WhichOneof("value"): + case "double_value": + latest_value = msg.double_value + case "int64_value": + latest_value = msg.int64_value + case field_type: + raise ValueError(f"Unexpected metric value type: {field_type}") else: - latest_value = float(point_value.int64_value) + raise TypeError("Failed to parse point value as a Protobuf Message") logging.info( f"Metric {name_str} for JobSet {jobset_name}: "