From c0fcb719426abe276e60d84fadf87d4364528f0e Mon Sep 17 00:00:00 2001 From: Alfred Yu Date: Tue, 12 Aug 2025 10:03:26 +0800 Subject: [PATCH 01/21] [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 96fcb76bf..ac3982d0e 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 a13a9dc36f79b86e147054304b9dd5cf9986b6f0 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Wed, 7 Jan 2026 16:49:37 +0800 Subject: [PATCH 02/21] feat: Add DAG for validating `tpu-info -help` command execution in TPU worker pods --- .../tpu_info_help_validation_dags.py | 279 ++++++++++++++++++ 1 file changed, 279 insertions(+) create mode 100644 dags/tpu_observability/tpu_info_help_validation_dags.py diff --git a/dags/tpu_observability/tpu_info_help_validation_dags.py b/dags/tpu_observability/tpu_info_help_validation_dags.py new file mode 100644 index 000000000..7d4d76db9 --- /dev/null +++ b/dags/tpu_observability/tpu_info_help_validation_dags.py @@ -0,0 +1,279 @@ +# 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 the `tpu-info -help` command execution inside TPU worker pods. +""" + +import datetime +import tempfile +import os + +from airflow import models +from airflow.decorators import task +from airflow.utils.trigger_rule import TriggerRule +from airflow.utils.task_group import TaskGroup + +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 import subprocess_util as subprocess +from dags.tpu_observability.utils.jobset_util import JobSet, Workload +from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH + + +def _get_tpu_info_help_output(info: node_pool.Info, pod_name: str) -> str: + """ + Retrieves the raw output of `tpu-info -help` from a specific pod. + """ + with tempfile.NamedTemporaryFile() as temp_config_file: + env = os.environ.copy() + env["KUBECONFIG"] = temp_config_file.name + + cmd = " && ".join([ + jobset.Command.get_credentials_command(info), + f"kubectl exec {pod_name} -n default -- tpu-info -help", + ]) + + return subprocess.run_exec(cmd, env=env) + + +@task +def validate_tpu_info_help(info: node_pool.Info, pod_name: str) -> str: + """ + Validates that the `tpu-info -help` output contains all required fields. + """ + output = _get_tpu_info_help_output(info, pod_name) + required_patterns = [ + "Display TPU info and metrics.", + "options:", + "-h, --help", + "-v, --version", + "-p, --process", + "--streaming", + "--rate RATE", + "--list_metrics", + ] + + for pattern in required_patterns: + if pattern not in output: + raise AssertionError( + "Validation failed: Missing expected string " + f"'{pattern}' in tpu-info help.\n" + f"Output received:\n{output}" + ) + + return output + + +def _get_tpu_version_output(info: node_pool.Info, pod_name: str) -> str: + """ + Executes the version command in the pod to get libtpu version and accelerator type. + """ + with tempfile.NamedTemporaryFile() as temp_config_file: + env = os.environ.copy() + env["KUBECONFIG"] = temp_config_file.name + cmd = " && ".join([ + jobset.Command.get_credentials_command(info), + f"kubectl exec {pod_name} -n default -- tpu-info --version", + ]) + return subprocess.run_exec(cmd, env=env) + + +@task +def validate_tpu_info_version(info: node_pool.Info, pod_name: str) -> str: + """ + Validates that `tpu-info --version` returns version and accelerator info. + """ + output = _get_tpu_version_output(info, pod_name) + required_patterns = [ + "tpu-info version:", + "libtpu version:", + "accelerator type:", + ] + + for pattern in required_patterns: + if pattern not in output: + raise AssertionError( + f"Validation failed: Missing expected string '{pattern}' " + f"in tpu-info --version output.\nOutput:\n{output}" + ) + + return output + + +def _get_tpu_process_output(info: node_pool.Info, pod_name: str) -> str: + """ + Executes the process command in the pod to get the TPU process table. + """ + with tempfile.NamedTemporaryFile() as temp_config_file: + env = os.environ.copy() + env["KUBECONFIG"] = temp_config_file.name + cmd = " && ".join([ + jobset.append_credentials_command(info), + f"kubectl exec {pod_name} -n default -- tpu-info --process", + ]) + return subprocess.run_exec(cmd, env=env) + + +@task +def validate_tpu_info_process(info: node_pool.Info, pod_name: str) -> str: + """ + Validates that `tpu-info --process` displays the TPU process table. + """ + output = _get_tpu_process_output(info, pod_name) + required_patterns = [ + "TPU Process Info", + "Chip", + "PID", + "Process Name", + "/dev/vfio/", + "python", + ] + + for pattern in required_patterns: + if pattern not in output: + raise AssertionError( + f"Validation failed: Missing expected string '{pattern}' " + f"in tpu-info --process output.\nOutput:\n{output}" + ) + + return output + + +# Keyword arguments are generated dynamically at runtime (pylint does not +# know this signature). +with models.DAG( # pylint: disable=unexpected-keyword-arg + dag_id="tpu_info_help_validation_dags", + start_date=datetime.datetime(2025, 8, 10), + schedule="0 18 * * *" if composer_env.is_prod_env() else None, + catchup=False, + tags=[ + "cloud-ml-auto-solutions", + "jobset", + "time-to-recover", + "tpu-observability", + "TPU", + "v6e-16", + ], + description=( + "This DAG validates the `tpu-info -help` command " + "execution inside TPU worker pods " + ), + doc_md=""" + ### Description + This DAG validates the `tpu-info -help` command execution inside TPU worker pods + across different TPU versions and machine configurations. + It ensures that the command runs successfully and returns the expected help information. + """, +) as dag: + for machine in MachineConfigMap: + config = machine.value + + jobset_config = JobSet( + jobset_name="tpu-info-help-validation-jobset", + namespace="default", + max_restarts=5, + replicated_job_name="tpu-job-slice", + replicas=1, + backoff_limit=0, + completions=4, + parallelism=4, + tpu_accelerator_type="tpu-v6e-slice", + tpu_topology="4x4", + container_name="jax-tpu-worker", + image="asia-northeast1-docker.pkg.dev/cienet-cmcs/" + "yuna-docker/tpu-info:v0.5.1", + tpu_cores_per_pod=4, + ) + + # 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}" + ): + 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="tpu_info_help_validation_dags", + is_prod=composer_env.is_prod_env(), + machine_type=config.machine_version.value, + tpu_topology=config.tpu_topology, + ) + + create_node_pool = node_pool.create( + node_pool=cluster_info, + ) + + apply_time = jobset.run_workload( + node_pool=cluster_info, + yaml_config=jobset_config.generate_yaml( + workload_script=Workload.JAX_TPU_BENCHMARK + ), + namespace=jobset_config.namespace, + ) + + pod_names = jobset.list_pod_names.override(task_id="list_pod_names")( + node_pool=cluster_info, + namespace=jobset_config.namespace, + ) + + wait_for_job_start = jobset.wait_for_jobset_started.override( + task_id="wait_for_job_start", timeout=1200 + )(cluster_info, pod_name_list=pod_names, job_apply_time=apply_time) + + validate_help = validate_tpu_info_help.partial(info=cluster_info).expand( + pod_name=pod_names + ) + + validate_version = validate_tpu_info_version.partial( + info=cluster_info + ).expand(pod_name=pod_names) + + validate_process = validate_tpu_info_process.partial( + info=cluster_info + ).expand(pod_name=pod_names) + + cleanup_workload = jobset.end_workload.override( + task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE + )( + node_pool=cluster_info, + jobset_name=jobset_config.jobset_name, + namespace=jobset_config.namespace, + ).as_teardown( + setups=apply_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, + ) + + # Airflow uses >> for task chaining, which is pointless for pylint. + # pylint: disable=pointless-statement + ( + cluster_info + >> create_node_pool + >> apply_time + >> pod_names + >> wait_for_job_start + >> validate_help + >> validate_process + >> validate_version + >> cleanup_workload + >> cleanup_node_pool + ) + # pylint: enable=pointless-statement From 2622ffb0565f5d1e7e19322b8e903309863386cf Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Wed, 7 Jan 2026 16:51:22 +0800 Subject: [PATCH 03/21] feat: Enhance DAG documentation for `tpu-info` CLI tool validation --- .../tpu_info_help_validation_dags.py | 21 +++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/dags/tpu_observability/tpu_info_help_validation_dags.py b/dags/tpu_observability/tpu_info_help_validation_dags.py index 7d4d76db9..2a8ea6bf4 100644 --- a/dags/tpu_observability/tpu_info_help_validation_dags.py +++ b/dags/tpu_observability/tpu_info_help_validation_dags.py @@ -13,7 +13,8 @@ # limitations under the License. """ -A DAG to validate the `tpu-info -help` command execution inside TPU worker pods. +A DAG to validate the `tpu-info` CLI tool, ensuring help documentation, +version metadata, and process monitoring are functional inside TPU worker pods. """ import datetime @@ -168,14 +169,22 @@ def validate_tpu_info_process(info: node_pool.Info, pod_name: str) -> str: "v6e-16", ], description=( - "This DAG validates the `tpu-info -help` command " - "execution inside TPU worker pods " + "Validates tpu-info CLI tool: help documentation, version metadata, " + "and process monitoring capabilities inside TPU worker pods." ), doc_md=""" ### Description - This DAG validates the `tpu-info -help` command execution inside TPU worker pods - across different TPU versions and machine configurations. - It ensures that the command runs successfully and returns the expected help information. + This DAG performs an end-to-end validation of the `tpu-info` observability tool + within TPU worker pods. It ensures the CLI tool is correctly installed and + functional across different TPU configurations. + + ### Validation Steps: + 1. **Help Menu Validation**: Verifies `tpu-info -help` displays all required + options (streaming, rate, etc.) and specific usage instructions. + 2. **Process Table Validation**: Confirms `tpu-info --process` can successfully + map PIDs to TPU chips. + 3. **Version Validation**: Ensures `tpu-info --version` correctly reports + the tool version, libtpu version, and accelerator type. """, ) as dag: for machine in MachineConfigMap: From c52fc7284b42b804d5ef3945c7762f2cc37abb5a Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Wed, 7 Jan 2026 17:21:23 +0800 Subject: [PATCH 04/21] feat: Refactor TPU info validation DAG to improve task grouping and command execution --- .../tpu_info_help_validation_dags.py | 31 +++++++++++-------- 1 file changed, 18 insertions(+), 13 deletions(-) diff --git a/dags/tpu_observability/tpu_info_help_validation_dags.py b/dags/tpu_observability/tpu_info_help_validation_dags.py index 2a8ea6bf4..e4f0d99d4 100644 --- a/dags/tpu_observability/tpu_info_help_validation_dags.py +++ b/dags/tpu_observability/tpu_info_help_validation_dags.py @@ -85,6 +85,7 @@ def _get_tpu_version_output(info: node_pool.Info, pod_name: str) -> str: with tempfile.NamedTemporaryFile() as temp_config_file: env = os.environ.copy() env["KUBECONFIG"] = temp_config_file.name + cmd = " && ".join([ jobset.Command.get_credentials_command(info), f"kubectl exec {pod_name} -n default -- tpu-info --version", @@ -121,8 +122,9 @@ def _get_tpu_process_output(info: node_pool.Info, pod_name: str) -> str: with tempfile.NamedTemporaryFile() as temp_config_file: env = os.environ.copy() env["KUBECONFIG"] = temp_config_file.name + cmd = " && ".join([ - jobset.append_credentials_command(info), + jobset.Command.get_credentials_command(info), f"kubectl exec {pod_name} -n default -- tpu-info --process", ]) return subprocess.run_exec(cmd, env=env) @@ -243,17 +245,22 @@ def validate_tpu_info_process(info: node_pool.Info, pod_name: str) -> str: task_id="wait_for_job_start", timeout=1200 )(cluster_info, pod_name_list=pod_names, job_apply_time=apply_time) - validate_help = validate_tpu_info_help.partial(info=cluster_info).expand( - pod_name=pod_names - ) + # Keyword arguments are generated dynamically at runtime (pylint does not + # know this signature). + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id="verification_group" + ) as verification_group: + validate_help = validate_tpu_info_help.partial( + info=cluster_info + ).expand(pod_name=pod_names) - validate_version = validate_tpu_info_version.partial( - info=cluster_info - ).expand(pod_name=pod_names) + validate_version = validate_tpu_info_version.partial( + info=cluster_info + ).expand(pod_name=pod_names) - validate_process = validate_tpu_info_process.partial( - info=cluster_info - ).expand(pod_name=pod_names) + validate_process = validate_tpu_info_process.partial( + info=cluster_info + ).expand(pod_name=pod_names) cleanup_workload = jobset.end_workload.override( task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE @@ -279,9 +286,7 @@ def validate_tpu_info_process(info: node_pool.Info, pod_name: str) -> str: >> apply_time >> pod_names >> wait_for_job_start - >> validate_help - >> validate_process - >> validate_version + >> verification_group >> cleanup_workload >> cleanup_node_pool ) From 3d95ce3e65ffcd14202b4f8d450297e3124c8456 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Fri, 9 Jan 2026 15:07:44 +0800 Subject: [PATCH 05/21] feat: Add DAG for validating `tpu-info` CLI tool functionality in TPU worker pods --- ...ags.py => tpu_info_cli_validation_dags.py} | 135 +++++------------- 1 file changed, 39 insertions(+), 96 deletions(-) rename dags/tpu_observability/{tpu_info_help_validation_dags.py => tpu_info_cli_validation_dags.py} (66%) diff --git a/dags/tpu_observability/tpu_info_help_validation_dags.py b/dags/tpu_observability/tpu_info_cli_validation_dags.py similarity index 66% rename from dags/tpu_observability/tpu_info_help_validation_dags.py rename to dags/tpu_observability/tpu_info_cli_validation_dags.py index e4f0d99d4..3c0057b1d 100644 --- a/dags/tpu_observability/tpu_info_help_validation_dags.py +++ b/dags/tpu_observability/tpu_info_cli_validation_dags.py @@ -18,8 +18,10 @@ """ import datetime -import tempfile import os +import subprocess +import tempfile +from typing import List from airflow import models from airflow.decorators import task @@ -34,29 +36,32 @@ from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH -def _get_tpu_info_help_output(info: node_pool.Info, pod_name: str) -> str: - """ - Retrieves the raw output of `tpu-info -help` from a specific pod. - """ +def run_command(info, pod_name: str, tpu_args: str) -> str: + """Helper to handle KUBECONFIG and execute kubectl.""" with tempfile.NamedTemporaryFile() as temp_config_file: env = os.environ.copy() env["KUBECONFIG"] = temp_config_file.name cmd = " && ".join([ jobset.Command.get_credentials_command(info), - f"kubectl exec {pod_name} -n default -- tpu-info -help", + f"kubectl exec {pod_name} -n default -- {tpu_args}", ]) - return subprocess.run_exec(cmd, env=env) +def check_output_contains(output: str, patterns: List[str], context: str): + """Helper to verify expected strings in output.""" + for pattern in patterns: + if pattern not in output: + raise AssertionError( + f"Validation failed for '{context}': Missing '{pattern}'." + ) + + @task -def validate_tpu_info_help(info: node_pool.Info, pod_name: str) -> str: - """ - Validates that the `tpu-info -help` output contains all required fields. - """ - output = _get_tpu_info_help_output(info, pod_name) - required_patterns = [ +def validate_help(info, pod_name: str) -> str: + output = run_command(info, pod_name, "tpu-info -help") + patterns = [ "Display TPU info and metrics.", "options:", "-h, --help", @@ -66,77 +71,22 @@ def validate_tpu_info_help(info: node_pool.Info, pod_name: str) -> str: "--rate RATE", "--list_metrics", ] - - for pattern in required_patterns: - if pattern not in output: - raise AssertionError( - "Validation failed: Missing expected string " - f"'{pattern}' in tpu-info help.\n" - f"Output received:\n{output}" - ) - + check_output_contains(output, patterns, "tpu-info -help") return output -def _get_tpu_version_output(info: node_pool.Info, pod_name: str) -> str: - """ - Executes the version command in the pod to get libtpu version and accelerator type. - """ - with tempfile.NamedTemporaryFile() as temp_config_file: - env = os.environ.copy() - env["KUBECONFIG"] = temp_config_file.name - - cmd = " && ".join([ - jobset.Command.get_credentials_command(info), - f"kubectl exec {pod_name} -n default -- tpu-info --version", - ]) - return subprocess.run_exec(cmd, env=env) - - @task -def validate_tpu_info_version(info: node_pool.Info, pod_name: str) -> str: - """ - Validates that `tpu-info --version` returns version and accelerator info. - """ - output = _get_tpu_version_output(info, pod_name) - required_patterns = [ - "tpu-info version:", - "libtpu version:", - "accelerator type:", - ] - - for pattern in required_patterns: - if pattern not in output: - raise AssertionError( - f"Validation failed: Missing expected string '{pattern}' " - f"in tpu-info --version output.\nOutput:\n{output}" - ) - +def validate_version(info, pod_name: str) -> str: + output = run_command(info, pod_name, "tpu-info --version") + patterns = ["tpu-info version:", "libtpu version:", "accelerator type:"] + check_output_contains(output, patterns, "tpu-info --version") return output -def _get_tpu_process_output(info: node_pool.Info, pod_name: str) -> str: - """ - Executes the process command in the pod to get the TPU process table. - """ - with tempfile.NamedTemporaryFile() as temp_config_file: - env = os.environ.copy() - env["KUBECONFIG"] = temp_config_file.name - - cmd = " && ".join([ - jobset.Command.get_credentials_command(info), - f"kubectl exec {pod_name} -n default -- tpu-info --process", - ]) - return subprocess.run_exec(cmd, env=env) - - @task -def validate_tpu_info_process(info: node_pool.Info, pod_name: str) -> str: - """ - Validates that `tpu-info --process` displays the TPU process table. - """ - output = _get_tpu_process_output(info, pod_name) - required_patterns = [ +def validate_process(info, pod_name: str) -> str: + output = run_command(info, pod_name, "tpu-info --process") + patterns = [ "TPU Process Info", "Chip", "PID", @@ -144,21 +94,14 @@ def validate_tpu_info_process(info: node_pool.Info, pod_name: str) -> str: "/dev/vfio/", "python", ] - - for pattern in required_patterns: - if pattern not in output: - raise AssertionError( - f"Validation failed: Missing expected string '{pattern}' " - f"in tpu-info --process output.\nOutput:\n{output}" - ) - + check_output_contains(output, patterns, "tpu-info --process") return output # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). with models.DAG( # pylint: disable=unexpected-keyword-arg - dag_id="tpu_info_help_validation_dags", + dag_id="tpu_info_cli_validation_dags", start_date=datetime.datetime(2025, 8, 10), schedule="0 18 * * *" if composer_env.is_prod_env() else None, catchup=False, @@ -193,7 +136,7 @@ def validate_tpu_info_process(info: node_pool.Info, pod_name: str) -> str: config = machine.value jobset_config = JobSet( - jobset_name="tpu-info-help-validation-jobset", + jobset_name="tpu-info-cli-validation-jobset", namespace="default", max_restarts=5, replicated_job_name="tpu-job-slice", @@ -218,7 +161,7 @@ def validate_tpu_info_process(info: node_pool.Info, pod_name: str) -> str: task_id="build_node_pool_info_from_gcs_yaml" )( gcs_path=GCS_CONFIG_PATH, - dag_name="tpu_info_help_validation_dags", + dag_name="tpu_info_cli_validation_dags", is_prod=composer_env.is_prod_env(), machine_type=config.machine_version.value, tpu_topology=config.tpu_topology, @@ -250,17 +193,17 @@ def validate_tpu_info_process(info: node_pool.Info, pod_name: str) -> str: with TaskGroup( # pylint: disable=unexpected-keyword-arg group_id="verification_group" ) as verification_group: - validate_help = validate_tpu_info_help.partial( - info=cluster_info - ).expand(pod_name=pod_names) + help_validation = validate_help.partial(info=cluster_info).expand( + pod_name=pod_names + ) - validate_version = validate_tpu_info_version.partial( - info=cluster_info - ).expand(pod_name=pod_names) + version_validation = validate_version.partial(info=cluster_info).expand( + pod_name=pod_names + ) - validate_process = validate_tpu_info_process.partial( - info=cluster_info - ).expand(pod_name=pod_names) + process_validation = validate_process.partial(info=cluster_info).expand( + pod_name=pod_names + ) cleanup_workload = jobset.end_workload.override( task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE From c2d8149af310a5fa57a648dab97f02fa3a5e891f Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Fri, 9 Jan 2026 15:16:10 +0800 Subject: [PATCH 06/21] refactor: Rename command execution functions for clarity and consistency --- .../tpu_info_cli_validation_dags.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/dags/tpu_observability/tpu_info_cli_validation_dags.py b/dags/tpu_observability/tpu_info_cli_validation_dags.py index 3c0057b1d..ba004cadd 100644 --- a/dags/tpu_observability/tpu_info_cli_validation_dags.py +++ b/dags/tpu_observability/tpu_info_cli_validation_dags.py @@ -36,7 +36,7 @@ from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH -def run_command(info, pod_name: str, tpu_args: str) -> str: +def execute_tpu_info_cli_command(info, pod_name: str, tpu_args: str) -> str: """Helper to handle KUBECONFIG and execute kubectl.""" with tempfile.NamedTemporaryFile() as temp_config_file: env = os.environ.copy() @@ -49,7 +49,9 @@ def run_command(info, pod_name: str, tpu_args: str) -> str: return subprocess.run_exec(cmd, env=env) -def check_output_contains(output: str, patterns: List[str], context: str): +def verify_output_contains_patterns( + output: str, patterns: List[str], context: str +): """Helper to verify expected strings in output.""" for pattern in patterns: if pattern not in output: @@ -60,7 +62,7 @@ def check_output_contains(output: str, patterns: List[str], context: str): @task def validate_help(info, pod_name: str) -> str: - output = run_command(info, pod_name, "tpu-info -help") + output = execute_tpu_info_cli_command(info, pod_name, "tpu-info -help") patterns = [ "Display TPU info and metrics.", "options:", @@ -71,21 +73,21 @@ def validate_help(info, pod_name: str) -> str: "--rate RATE", "--list_metrics", ] - check_output_contains(output, patterns, "tpu-info -help") + verify_output_contains_patterns(output, patterns, "tpu-info -help") return output @task def validate_version(info, pod_name: str) -> str: - output = run_command(info, pod_name, "tpu-info --version") + output = execute_tpu_info_cli_command(info, pod_name, "tpu-info --version") patterns = ["tpu-info version:", "libtpu version:", "accelerator type:"] - check_output_contains(output, patterns, "tpu-info --version") + verify_output_contains_patterns(output, patterns, "tpu-info --version") return output @task def validate_process(info, pod_name: str) -> str: - output = run_command(info, pod_name, "tpu-info --process") + output = execute_tpu_info_cli_command(info, pod_name, "tpu-info --process") patterns = [ "TPU Process Info", "Chip", @@ -94,7 +96,7 @@ def validate_process(info, pod_name: str) -> str: "/dev/vfio/", "python", ] - check_output_contains(output, patterns, "tpu-info --process") + verify_output_contains_patterns(output, patterns, "tpu-info --process") return output From 0492b1c0d98b9ee50a8bdea7a9fbefa843296332 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Fri, 9 Jan 2026 15:22:02 +0800 Subject: [PATCH 07/21] feat: Enhance task ID overrides for pod name and validation tasks in `tpu-info` DAG --- .../tpu_info_cli_validation_dags.py | 22 ++++++++++++------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/dags/tpu_observability/tpu_info_cli_validation_dags.py b/dags/tpu_observability/tpu_info_cli_validation_dags.py index ba004cadd..48ca19857 100644 --- a/dags/tpu_observability/tpu_info_cli_validation_dags.py +++ b/dags/tpu_observability/tpu_info_cli_validation_dags.py @@ -169,11 +169,11 @@ def validate_process(info, pod_name: str) -> str: tpu_topology=config.tpu_topology, ) - create_node_pool = node_pool.create( + create_node_pool = node_pool.create.override(task_id="create_node_pool")( node_pool=cluster_info, ) - apply_time = jobset.run_workload( + apply_time = jobset.run_workload.override(task_id="run_workload")( node_pool=cluster_info, yaml_config=jobset_config.generate_yaml( workload_script=Workload.JAX_TPU_BENCHMARK @@ -195,16 +195,22 @@ def validate_process(info, pod_name: str) -> str: with TaskGroup( # pylint: disable=unexpected-keyword-arg group_id="verification_group" ) as verification_group: - help_validation = validate_help.partial(info=cluster_info).expand( - pod_name=pod_names + help_validation = ( + validate_help.override(task_id="validate_help") + .partial(info=cluster_info) + .expand(pod_name=pod_names) ) - version_validation = validate_version.partial(info=cluster_info).expand( - pod_name=pod_names + version_validation = ( + validate_version.override(task_id="validate_version") + .partial(info=cluster_info) + .expand(pod_name=pod_names) ) - process_validation = validate_process.partial(info=cluster_info).expand( - pod_name=pod_names + process_validation = ( + validate_process.override(task_id="validate_process") + .partial(info=cluster_info) + .expand(pod_name=pod_names) ) cleanup_workload = jobset.end_workload.override( From 12f893136ec63b461662d14b737006e0d1b6b04d Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Fri, 9 Jan 2026 15:32:27 +0800 Subject: [PATCH 08/21] fix: Remove timeout parameter from wait_for_job_start task in `tpu-info` DAG --- dags/tpu_observability/tpu_info_cli_validation_dags.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dags/tpu_observability/tpu_info_cli_validation_dags.py b/dags/tpu_observability/tpu_info_cli_validation_dags.py index 48ca19857..994cbb13d 100644 --- a/dags/tpu_observability/tpu_info_cli_validation_dags.py +++ b/dags/tpu_observability/tpu_info_cli_validation_dags.py @@ -187,7 +187,7 @@ def validate_process(info, pod_name: str) -> str: ) wait_for_job_start = jobset.wait_for_jobset_started.override( - task_id="wait_for_job_start", timeout=1200 + task_id="wait_for_job_start" )(cluster_info, pod_name_list=pod_names, job_apply_time=apply_time) # Keyword arguments are generated dynamically at runtime (pylint does not From fe66f48e977fedf38e02d1f1b6a8519f0ad9c7bd Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Tue, 13 Jan 2026 10:40:17 +0800 Subject: [PATCH 09/21] fix: Remove schedule parameter from DAG definition in `tpu_info_cli_validation_dags` --- dags/tpu_observability/tpu_info_cli_validation_dags.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dags/tpu_observability/tpu_info_cli_validation_dags.py b/dags/tpu_observability/tpu_info_cli_validation_dags.py index 994cbb13d..ac2dfaf1a 100644 --- a/dags/tpu_observability/tpu_info_cli_validation_dags.py +++ b/dags/tpu_observability/tpu_info_cli_validation_dags.py @@ -105,7 +105,7 @@ def validate_process(info, pod_name: str) -> str: with models.DAG( # pylint: disable=unexpected-keyword-arg dag_id="tpu_info_cli_validation_dags", start_date=datetime.datetime(2025, 8, 10), - schedule="0 18 * * *" if composer_env.is_prod_env() else None, + schedule=None, catchup=False, tags=[ "cloud-ml-auto-solutions", From 9e400927c39ed26258d229fe1d0111671b46eaa6 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Wed, 14 Jan 2026 17:32:42 +0800 Subject: [PATCH 10/21] feat: Add validation DAG for `tpu-info` CLI tool functionality in TPU worker pods --- .../tpu_info_format_validation_dags.py | 215 +++++++++++++++++- 1 file changed, 214 insertions(+), 1 deletion(-) diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index c57e99460..859abe0c3 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -13,9 +13,13 @@ # limitations under the License. """ +tpu_info_format_validation_dag: A DAG orchestrates the process of verifying TensorCore utilization metrics. - This is done by comparing data from Cloud Logging and Cloud Monitoring. + +tpu_info_cli_validation_dags: +A DAG to validate the `tpu-info` CLI tool, ensuring help documentation, +version metadata, and process monitoring are functional inside TPU worker pods. """ import datetime @@ -292,6 +296,70 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]): ) +def execute_tpu_info_cli_command(info, pod_name: str, tpu_args: str) -> str: + """Helper to handle KUBECONFIG and execute kubectl.""" + with tempfile.NamedTemporaryFile() as temp_config_file: + env = os.environ.copy() + env["KUBECONFIG"] = temp_config_file.name + + cmd = " && ".join([ + jobset.Command.get_credentials_command(info), + f"kubectl exec {pod_name} -n default -- {tpu_args}", + ]) + return subprocess.run_exec(cmd, env=env) + + +def verify_output_contains_patterns( + output: str, patterns: List[str], context: str +): + """Helper to verify expected strings in output.""" + for pattern in patterns: + if pattern not in output: + raise AssertionError( + f"Validation failed for '{context}': Missing '{pattern}'." + ) + + +@task +def validate_help(info, pod_name: str) -> str: + output = execute_tpu_info_cli_command(info, pod_name, "tpu-info -help") + patterns = [ + "Display TPU info and metrics.", + "options:", + "-h, --help", + "-v, --version", + "-p, --process", + "--streaming", + "--rate RATE", + "--list_metrics", + ] + verify_output_contains_patterns(output, patterns, "tpu-info -help") + return output + + +@task +def validate_version(info, pod_name: str) -> str: + output = execute_tpu_info_cli_command(info, pod_name, "tpu-info --version") + patterns = ["tpu-info version:", "libtpu version:", "accelerator type:"] + verify_output_contains_patterns(output, patterns, "tpu-info --version") + return output + + +@task +def validate_process(info, pod_name: str) -> str: + output = execute_tpu_info_cli_command(info, pod_name, "tpu-info --process") + patterns = [ + "TPU Process Info", + "Chip", + "PID", + "Process Name", + "/dev/vfio/", + "python", + ] + verify_output_contains_patterns(output, patterns, "tpu-info --process") + return output + + # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). with models.DAG( # pylint: disable=unexpected-keyword-arg @@ -531,3 +599,148 @@ def generate_second_node_pool_name( clean_up_workload, cleanup_node_pool, ) + # pylint: enable=pointless-statement + + +# Keyword arguments are generated dynamically at runtime (pylint does not +# know this signature). +with models.DAG( # pylint: disable=unexpected-keyword-arg + dag_id="tpu_info_cli_validation_dags", + start_date=datetime.datetime(2025, 8, 10), + schedule=None, + catchup=False, + tags=[ + "cloud-ml-auto-solutions", + "jobset", + "time-to-recover", + "tpu-observability", + "TPU", + "v6e-16", + ], + description=( + "Validates tpu-info CLI tool: help documentation, version metadata, " + "and process monitoring capabilities inside TPU worker pods." + ), + doc_md=""" + ### Description + This DAG performs an end-to-end validation of the `tpu-info` observability tool + within TPU worker pods. It ensures the CLI tool is correctly installed and + functional across different TPU configurations. + + ### Validation Steps: + 1. **Help Menu Validation**: Verifies `tpu-info -help` displays all required + options (streaming, rate, etc.) and specific usage instructions. + 2. **Process Table Validation**: Confirms `tpu-info --process` can successfully + map PIDs to TPU chips. + 3. **Version Validation**: Ensures `tpu-info --version` correctly reports + the tool version, libtpu version, and accelerator type. + """, +) as dag: + for machine in MachineConfigMap: + config = machine.value + + jobset_config = JobSet( + jobset_name="tpu-info-cli-validation-jobset", + namespace="default", + max_restarts=5, + replicated_job_name="tpu-job-slice", + replicas=1, + backoff_limit=0, + completions=4, + parallelism=4, + tpu_accelerator_type="tpu-v6e-slice", + tpu_topology="4x4", + container_name="jax-tpu-worker", + image="asia-northeast1-docker.pkg.dev/cienet-cmcs/" + "yuna-docker/tpu-info:v0.5.1", + tpu_cores_per_pod=4, + ) + + # 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}" + ): + 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="tpu_info_cli_validation_dags", + is_prod=composer_env.is_prod_env(), + machine_type=config.machine_version.value, + tpu_topology=config.tpu_topology, + ) + + create_node_pool = node_pool.create.override(task_id="create_node_pool")( + node_pool=cluster_info, + ) + + apply_time = jobset.run_workload.override(task_id="run_workload")( + node_pool=cluster_info, + yaml_config=jobset_config.generate_yaml( + workload_script=Workload.JAX_TPU_BENCHMARK + ), + namespace=jobset_config.namespace, + ) + + pod_names = jobset.list_pod_names.override(task_id="list_pod_names")( + node_pool=cluster_info, + namespace=jobset_config.namespace, + ) + + wait_for_job_start = jobset.wait_for_jobset_started.override( + task_id="wait_for_job_start" + )(cluster_info, pod_name_list=pod_names, job_apply_time=apply_time) + + # Keyword arguments are generated dynamically at runtime (pylint does not + # know this signature). + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id="verification_group" + ) as verification_group: + help_validation = ( + validate_help.override(task_id="validate_help") + .partial(info=cluster_info) + .expand(pod_name=pod_names) + ) + + version_validation = ( + validate_version.override(task_id="validate_version") + .partial(info=cluster_info) + .expand(pod_name=pod_names) + ) + + process_validation = ( + validate_process.override(task_id="validate_process") + .partial(info=cluster_info) + .expand(pod_name=pod_names) + ) + + cleanup_workload = jobset.end_workload.override( + task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE + )( + node_pool=cluster_info, + jobset_name=jobset_config.jobset_name, + namespace=jobset_config.namespace, + ).as_teardown( + setups=apply_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, + ) + + # Airflow uses >> for task chaining, which is pointless for pylint. + # pylint: disable=pointless-statement + ( + cluster_info + >> create_node_pool + >> apply_time + >> pod_names + >> wait_for_job_start + >> verification_group + >> cleanup_workload + >> cleanup_node_pool + ) + # pylint: enable=pointless-statement From 02df4475698e14df970491c0a417f6689f5952ee Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Thu, 15 Jan 2026 17:36:35 +0800 Subject: [PATCH 11/21] fix: Update type hint for patterns parameter in verify_output_contains_patterns function --- .../tpu_info_cli_validation_dags.py | 244 ------------------ .../tpu_info_format_validation_dags.py | 2 +- 2 files changed, 1 insertion(+), 245 deletions(-) delete mode 100644 dags/tpu_observability/tpu_info_cli_validation_dags.py diff --git a/dags/tpu_observability/tpu_info_cli_validation_dags.py b/dags/tpu_observability/tpu_info_cli_validation_dags.py deleted file mode 100644 index ac2dfaf1a..000000000 --- a/dags/tpu_observability/tpu_info_cli_validation_dags.py +++ /dev/null @@ -1,244 +0,0 @@ -# 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 the `tpu-info` CLI tool, ensuring help documentation, -version metadata, and process monitoring are functional inside TPU worker pods. -""" - -import datetime -import os -import subprocess -import tempfile -from typing import List - -from airflow import models -from airflow.decorators import task -from airflow.utils.trigger_rule import TriggerRule -from airflow.utils.task_group import TaskGroup - -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 import subprocess_util as subprocess -from dags.tpu_observability.utils.jobset_util import JobSet, Workload -from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH - - -def execute_tpu_info_cli_command(info, pod_name: str, tpu_args: str) -> str: - """Helper to handle KUBECONFIG and execute kubectl.""" - with tempfile.NamedTemporaryFile() as temp_config_file: - env = os.environ.copy() - env["KUBECONFIG"] = temp_config_file.name - - cmd = " && ".join([ - jobset.Command.get_credentials_command(info), - f"kubectl exec {pod_name} -n default -- {tpu_args}", - ]) - return subprocess.run_exec(cmd, env=env) - - -def verify_output_contains_patterns( - output: str, patterns: List[str], context: str -): - """Helper to verify expected strings in output.""" - for pattern in patterns: - if pattern not in output: - raise AssertionError( - f"Validation failed for '{context}': Missing '{pattern}'." - ) - - -@task -def validate_help(info, pod_name: str) -> str: - output = execute_tpu_info_cli_command(info, pod_name, "tpu-info -help") - patterns = [ - "Display TPU info and metrics.", - "options:", - "-h, --help", - "-v, --version", - "-p, --process", - "--streaming", - "--rate RATE", - "--list_metrics", - ] - verify_output_contains_patterns(output, patterns, "tpu-info -help") - return output - - -@task -def validate_version(info, pod_name: str) -> str: - output = execute_tpu_info_cli_command(info, pod_name, "tpu-info --version") - patterns = ["tpu-info version:", "libtpu version:", "accelerator type:"] - verify_output_contains_patterns(output, patterns, "tpu-info --version") - return output - - -@task -def validate_process(info, pod_name: str) -> str: - output = execute_tpu_info_cli_command(info, pod_name, "tpu-info --process") - patterns = [ - "TPU Process Info", - "Chip", - "PID", - "Process Name", - "/dev/vfio/", - "python", - ] - verify_output_contains_patterns(output, patterns, "tpu-info --process") - return output - - -# Keyword arguments are generated dynamically at runtime (pylint does not -# know this signature). -with models.DAG( # pylint: disable=unexpected-keyword-arg - dag_id="tpu_info_cli_validation_dags", - start_date=datetime.datetime(2025, 8, 10), - schedule=None, - catchup=False, - tags=[ - "cloud-ml-auto-solutions", - "jobset", - "time-to-recover", - "tpu-observability", - "TPU", - "v6e-16", - ], - description=( - "Validates tpu-info CLI tool: help documentation, version metadata, " - "and process monitoring capabilities inside TPU worker pods." - ), - doc_md=""" - ### Description - This DAG performs an end-to-end validation of the `tpu-info` observability tool - within TPU worker pods. It ensures the CLI tool is correctly installed and - functional across different TPU configurations. - - ### Validation Steps: - 1. **Help Menu Validation**: Verifies `tpu-info -help` displays all required - options (streaming, rate, etc.) and specific usage instructions. - 2. **Process Table Validation**: Confirms `tpu-info --process` can successfully - map PIDs to TPU chips. - 3. **Version Validation**: Ensures `tpu-info --version` correctly reports - the tool version, libtpu version, and accelerator type. - """, -) as dag: - for machine in MachineConfigMap: - config = machine.value - - jobset_config = JobSet( - jobset_name="tpu-info-cli-validation-jobset", - namespace="default", - max_restarts=5, - replicated_job_name="tpu-job-slice", - replicas=1, - backoff_limit=0, - completions=4, - parallelism=4, - tpu_accelerator_type="tpu-v6e-slice", - tpu_topology="4x4", - container_name="jax-tpu-worker", - image="asia-northeast1-docker.pkg.dev/cienet-cmcs/" - "yuna-docker/tpu-info:v0.5.1", - tpu_cores_per_pod=4, - ) - - # 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}" - ): - 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="tpu_info_cli_validation_dags", - is_prod=composer_env.is_prod_env(), - machine_type=config.machine_version.value, - tpu_topology=config.tpu_topology, - ) - - create_node_pool = node_pool.create.override(task_id="create_node_pool")( - node_pool=cluster_info, - ) - - apply_time = jobset.run_workload.override(task_id="run_workload")( - node_pool=cluster_info, - yaml_config=jobset_config.generate_yaml( - workload_script=Workload.JAX_TPU_BENCHMARK - ), - namespace=jobset_config.namespace, - ) - - pod_names = jobset.list_pod_names.override(task_id="list_pod_names")( - node_pool=cluster_info, - namespace=jobset_config.namespace, - ) - - wait_for_job_start = jobset.wait_for_jobset_started.override( - task_id="wait_for_job_start" - )(cluster_info, pod_name_list=pod_names, job_apply_time=apply_time) - - # Keyword arguments are generated dynamically at runtime (pylint does not - # know this signature). - with TaskGroup( # pylint: disable=unexpected-keyword-arg - group_id="verification_group" - ) as verification_group: - help_validation = ( - validate_help.override(task_id="validate_help") - .partial(info=cluster_info) - .expand(pod_name=pod_names) - ) - - version_validation = ( - validate_version.override(task_id="validate_version") - .partial(info=cluster_info) - .expand(pod_name=pod_names) - ) - - process_validation = ( - validate_process.override(task_id="validate_process") - .partial(info=cluster_info) - .expand(pod_name=pod_names) - ) - - cleanup_workload = jobset.end_workload.override( - task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE - )( - node_pool=cluster_info, - jobset_name=jobset_config.jobset_name, - namespace=jobset_config.namespace, - ).as_teardown( - setups=apply_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, - ) - - # Airflow uses >> for task chaining, which is pointless for pylint. - # pylint: disable=pointless-statement - ( - cluster_info - >> create_node_pool - >> apply_time - >> pod_names - >> wait_for_job_start - >> verification_group - >> cleanup_workload - >> cleanup_node_pool - ) - # pylint: enable=pointless-statement diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index 859abe0c3..3e483d392 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -310,7 +310,7 @@ def execute_tpu_info_cli_command(info, pod_name: str, tpu_args: str) -> str: def verify_output_contains_patterns( - output: str, patterns: List[str], context: str + output: str, patterns: list[str], context: str ): """Helper to verify expected strings in output.""" for pattern in patterns: From 7ce3015c0b847fad090e2f3059fe72c016a556da Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Tue, 27 Jan 2026 15:55:01 +0800 Subject: [PATCH 12/21] refactor: Replace task chaining with chain function for improved readability in DAG --- .../tpu_info_format_validation_dags.py | 20 +++++++++---------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index 3e483d392..e9151d915 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -731,16 +731,14 @@ def generate_second_node_pool_name( setups=create_node_pool, ) - # Airflow uses >> for task chaining, which is pointless for pylint. - # pylint: disable=pointless-statement - ( - cluster_info - >> create_node_pool - >> apply_time - >> pod_names - >> wait_for_job_start - >> verification_group - >> cleanup_workload - >> cleanup_node_pool + chain( + cluster_info, + create_node_pool, + apply_time, + pod_names, + wait_for_job_start, + verification_group, + cleanup_workload, + cleanup_node_pool, ) # pylint: enable=pointless-statement From d92820fe7d9cddc746ca6e9f3a5b692e78847a25 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Tue, 3 Feb 2026 11:53:43 +0800 Subject: [PATCH 13/21] fix: Update schedule parameter in DAG definition for production environment --- dags/tpu_observability/tpu_info_format_validation_dags.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index e9151d915..8e5079d53 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -607,7 +607,7 @@ def generate_second_node_pool_name( with models.DAG( # pylint: disable=unexpected-keyword-arg dag_id="tpu_info_cli_validation_dags", start_date=datetime.datetime(2025, 8, 10), - schedule=None, + schedule="30 14 * * *" if composer_env.is_prod_env() else None, catchup=False, tags=[ "cloud-ml-auto-solutions", From 058c2d11df27d6434cce8189c7ab69b6bed3f1e6 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Tue, 3 Feb 2026 15:27:32 +0800 Subject: [PATCH 14/21] refactor: Simplify jobset configuration by building from GCS YAML in TPU info validation DAG --- .../tpu_info_format_validation_dags.py | 34 ++++++------------- 1 file changed, 10 insertions(+), 24 deletions(-) diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index 8e5079d53..4f7da7ab4 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -639,28 +639,16 @@ def generate_second_node_pool_name( for machine in MachineConfigMap: config = machine.value - jobset_config = JobSet( - jobset_name="tpu-info-cli-validation-jobset", - namespace="default", - max_restarts=5, - replicated_job_name="tpu-job-slice", - replicas=1, - backoff_limit=0, - completions=4, - parallelism=4, - tpu_accelerator_type="tpu-v6e-slice", - tpu_topology="4x4", - container_name="jax-tpu-worker", - image="asia-northeast1-docker.pkg.dev/cienet-cmcs/" - "yuna-docker/tpu-info:v0.5.1", - tpu_cores_per_pod=4, - ) - # 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}" ): + jobset_config = jobset.build_jobset_from_gcs_yaml( + gcs_path=GCS_JOBSET_CONFIG_PATH, + dag_name="tpu_info_cli_validation_dags", + ) + cluster_info = node_pool.build_node_pool_info_from_gcs_yaml.override( task_id="build_node_pool_info_from_gcs_yaml" )( @@ -677,15 +665,13 @@ def generate_second_node_pool_name( apply_time = jobset.run_workload.override(task_id="run_workload")( node_pool=cluster_info, - yaml_config=jobset_config.generate_yaml( - workload_script=Workload.JAX_TPU_BENCHMARK - ), - namespace=jobset_config.namespace, + jobset_config=jobset_config, + workload_type=Workload.JAX_TPU_BENCHMARK, ) pod_names = jobset.list_pod_names.override(task_id="list_pod_names")( node_pool=cluster_info, - namespace=jobset_config.namespace, + jobset_config=jobset_config, ) wait_for_job_start = jobset.wait_for_jobset_started.override( @@ -719,8 +705,7 @@ def generate_second_node_pool_name( task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE )( node_pool=cluster_info, - jobset_name=jobset_config.jobset_name, - namespace=jobset_config.namespace, + jobset_config=jobset_config, ).as_teardown( setups=apply_time ) @@ -732,6 +717,7 @@ def generate_second_node_pool_name( ) chain( + jobset_config, cluster_info, create_node_pool, apply_time, From 1fa00e85de5291a8e69ccbabc44fe394a544d4d7 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Tue, 10 Feb 2026 15:27:51 +0800 Subject: [PATCH 15/21] refactor: Replace task chaining with list for improved clarity in DAG --- dags/tpu_observability/tpu_info_format_validation_dags.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index 4f7da7ab4..709ad8926 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -580,7 +580,7 @@ def generate_second_node_pool_name( ], ) - chain(create_first_node_pool, create_second_node_pool) + [create_first_node_pool, create_second_node_pool] chain(cleanup_first_node_pool, cleanup_second_node_pool) From ff98f0f5e6867c1d7445e2e0bd334b948836176d Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Tue, 10 Feb 2026 15:46:03 +0800 Subject: [PATCH 16/21] refactor: Consolidate TPU CLI validation into a single function using CLI validation specification --- .../tpu_info_format_validation_dags.py | 79 ++++--------------- dags/tpu_observability/utils/tpu_cli_util.py | 34 ++++++++ 2 files changed, 51 insertions(+), 62 deletions(-) create mode 100644 dags/tpu_observability/utils/tpu_cli_util.py diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index 709ad8926..4df30a8ff 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -47,6 +47,7 @@ from dags.tpu_observability.utils import node_pool_util as node_pool from dags.tpu_observability.utils import subprocess_util as subprocess from dags.tpu_observability.utils import tpu_info_util as tpu_info +from dags.tpu_observability.utils.tpu_cli_util import CLI_VALIDATION_SPEC from dags.tpu_observability.utils.jobset_util import Workload from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout @@ -309,55 +310,21 @@ def execute_tpu_info_cli_command(info, pod_name: str, tpu_args: str) -> str: return subprocess.run_exec(cmd, env=env) -def verify_output_contains_patterns( - output: str, patterns: list[str], context: str -): - """Helper to verify expected strings in output.""" - for pattern in patterns: - if pattern not in output: - raise AssertionError( - f"Validation failed for '{context}': Missing '{pattern}'." - ) - - @task -def validate_help(info, pod_name: str) -> str: - output = execute_tpu_info_cli_command(info, pod_name, "tpu-info -help") - patterns = [ - "Display TPU info and metrics.", - "options:", - "-h, --help", - "-v, --version", - "-p, --process", - "--streaming", - "--rate RATE", - "--list_metrics", - ] - verify_output_contains_patterns(output, patterns, "tpu-info -help") - return output - - -@task -def validate_version(info, pod_name: str) -> str: - output = execute_tpu_info_cli_command(info, pod_name, "tpu-info --version") - patterns = ["tpu-info version:", "libtpu version:", "accelerator type:"] - verify_output_contains_patterns(output, patterns, "tpu-info --version") - return output - - -@task -def validate_process(info, pod_name: str) -> str: - output = execute_tpu_info_cli_command(info, pod_name, "tpu-info --process") - patterns = [ - "TPU Process Info", - "Chip", - "PID", - "Process Name", - "/dev/vfio/", - "python", - ] - verify_output_contains_patterns(output, patterns, "tpu-info --process") - return output +def validate_tpu_info_cli(info: node_pool.Info, pod_name: str) -> None: + """ + Validates tpu-info CLI commands inside TPU worker pods. + Consolidated version following the SDK validation pattern. + """ + for cmd_enum, patterns in CLI_VALIDATION_SPEC.items(): + output = execute_tpu_info_cli_command(info, pod_name, cmd_enum.value) + + for pattern in patterns: + if pattern not in output: + raise AssertionError( + f"Validation failed for '{cmd_enum.value}': " + f"Missing expected pattern '{pattern}'." + ) # Keyword arguments are generated dynamically at runtime (pylint does not @@ -683,20 +650,8 @@ def generate_second_node_pool_name( with TaskGroup( # pylint: disable=unexpected-keyword-arg group_id="verification_group" ) as verification_group: - help_validation = ( - validate_help.override(task_id="validate_help") - .partial(info=cluster_info) - .expand(pod_name=pod_names) - ) - - version_validation = ( - validate_version.override(task_id="validate_version") - .partial(info=cluster_info) - .expand(pod_name=pod_names) - ) - - process_validation = ( - validate_process.override(task_id="validate_process") + cli_validation = ( + validate_tpu_info_cli.override(task_id="validate_tpu_info_cli") .partial(info=cluster_info) .expand(pod_name=pod_names) ) diff --git a/dags/tpu_observability/utils/tpu_cli_util.py b/dags/tpu_observability/utils/tpu_cli_util.py new file mode 100644 index 000000000..dbbab07e2 --- /dev/null +++ b/dags/tpu_observability/utils/tpu_cli_util.py @@ -0,0 +1,34 @@ +from enum import Enum + + +class TpuInfoCmd(Enum): + HELP = "tpu-info -help" + VERSION = "tpu-info --version" + PROCESS = "tpu-info --process" + + +CLI_VALIDATION_SPEC: dict[TpuInfoCmd, list[str]] = { + TpuInfoCmd.HELP: [ + "Display TPU info and metrics.", + "options:", + "-h, --help", + "-v, --version", + "-p, --process", + "--streaming", + "--rate RATE", + "--list_metrics", + ], + TpuInfoCmd.VERSION: [ + "tpu-info version:", + "libtpu version:", + "accelerator type:", + ], + TpuInfoCmd.PROCESS: [ + "TPU Process Info", + "Chip", + "PID", + "Process Name", + "/dev/vfio/", + "python", + ], +} From c27ebcc41427e30fae1152eae5b7028304fc313c Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Thu, 12 Feb 2026 11:28:28 +0800 Subject: [PATCH 17/21] refactor: Move TPU CLI validation logic to tpu_info_util and remove tpu_cli_util --- .../tpu_info_format_validation_dags.py | 233 +++++------------- dags/tpu_observability/utils/tpu_cli_util.py | 34 --- dags/tpu_observability/utils/tpu_info_util.py | 47 +++- 3 files changed, 99 insertions(+), 215 deletions(-) delete mode 100644 dags/tpu_observability/utils/tpu_cli_util.py diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index 4df30a8ff..74f5b8a7f 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -47,7 +47,6 @@ from dags.tpu_observability.utils import node_pool_util as node_pool from dags.tpu_observability.utils import subprocess_util as subprocess from dags.tpu_observability.utils import tpu_info_util as tpu_info -from dags.tpu_observability.utils.tpu_cli_util import CLI_VALIDATION_SPEC from dags.tpu_observability.utils.jobset_util import Workload from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout @@ -57,34 +56,6 @@ SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID) -@task -def get_tpu_info_from_pod(info: node_pool.Info, pod_name: str) -> str: - """ - Executes the `tpu-info` command in a specified pod and returns its output. - - This task uses kubectl to run the 'tpu-info' command inside the given pod - in the 'default' namespace. The output of the command is captured and - returned. - - Args: - kubeconfig: The path to the kubeconfig file. - pod_name: The name of the pod to execute the command in. - - Returns: - The standard output from the 'tpu-info' command. - """ - with tempfile.NamedTemporaryFile() as temp_config_file: - env = os.environ.copy() - env["KUBECONFIG"] = temp_config_file.name - - cmd = " && ".join([ - jobset.Command.get_credentials_command(info), - f"kubectl exec {pod_name} -n default -- tpu-info", - ]) - - return subprocess.run_exec(cmd, env=env) - - @task def verify_table_amount(tpu_info_output: list[tpu_info.Table]): """ @@ -312,13 +283,35 @@ def execute_tpu_info_cli_command(info, pod_name: str, tpu_args: str) -> str: @task def validate_tpu_info_cli(info: node_pool.Info, pod_name: str) -> None: - """ - Validates tpu-info CLI commands inside TPU worker pods. - Consolidated version following the SDK validation pattern. - """ - for cmd_enum, patterns in CLI_VALIDATION_SPEC.items(): - output = execute_tpu_info_cli_command(info, pod_name, cmd_enum.value) + """Validates tpu-info CLI commands with internal expectation patterns.""" + + validation_spec = { + tpu_info.TpuInfoCmd.HELP: [ + "Display TPU info and metrics.", + "options:", + "-h, --help", + "-v, --version", + "-p, --process", + "--streaming", + "--rate RATE", + ], + tpu_info.TpuInfoCmd.VERSION: [ + "tpu-info version:", + "libtpu version:", + "accelerator type:", + ], + tpu_info.TpuInfoCmd.PROCESS: [ + "TPU Process Info", + "Chip", + "PID", + "Process Name", + "/dev/vfio/", + "python", + ], + } + for cmd_enum, patterns in validation_spec.items(): + output = execute_tpu_info_cli_command(info, pod_name, cmd_enum.value) for pattern in patterns: if pattern not in output: raise AssertionError( @@ -338,37 +331,31 @@ def validate_tpu_info_cli(info: node_pool.Info, pod_name: str) -> None: catchup=False, tags=["gke", "tpu-observability", "tpu-info", "TPU", "v6e-16"], description=( - "This DAG verifies the format of the tables in the tpu-info output " - "using tpu-info CLI tool. It includes 4 tables: TPU Chips, TPU " - "Runtime Utilization, TensorCore Utilization, and TPU Buffer Transfer " - "Latency." + "Validates the tpu-info CLI tool by verifying its command-line options. " + "It ensures that help documentation, version metadata, and process " + "mapping are functional across TPU worker pods." ), doc_md=""" - # Format Validation DAG - # This DAG verifies the format of the tables in the tpu-info output. + # TPU Info CLI Validation DAG + This DAG automates the end-to-end validation of the `tpu-info` observability tool. ### Description - This DAG automates the validation of the tpu-info command-line tool's - output format.It verifies the structure and content of key metric tables, - including "TPU Chips", "TPU Runtime Utilization", "TensorCore - Utilization", and "TPU Buffer Transfer Latency", by running the tool on a - live GKE cluster with TPU node pools. - - ### Prerequisites - This test requires an existing GKE cluster. - A pre-built Docker image containing the necessary jax, libtpu, and - tpu-info packages must also be available in a repository accessible - by the GKE cluster. + The test verifies that the `tpu-info` CLI correctly interprets various **command-line options** (e.g., `-help`, `--version`, `--process`) and returns the expected metadata or + metric structures. This ensures the tool is correctly installed and compatible + with the current TPU runtime environment. + + ### Validation Scope + The DAG executes the following command options inside the TPU pods: + * **Help Option (`-help`)**: Verifies the presence of required flags like `--streaming` and `--rate`. + * **Version Option (`--version`)**: Validates that the tool reports the correct version and `libtpu` metadata. + * **Process Option (`--process`)**: Ensures that active PIDs can be mapped to specific TPU chips. ### Procedures - The DAG begins by creating temporary GKE TPU node pools for the test. - Once the node pools are running, it schedules a Kubernetes JobSet and - waits for the pods to become active. It then executes the tpu-info - command within these pods to capture the raw text output. This output is - parsed into structured tables, and a series of validation tasks check - each table for the correct structure, row counts, and data formats. - Finally, regardless of the test outcome, the DAG cleans up all created - resources, including the JobSet and the temporary node pools. + 1. **Environment Setup**: Dynamically creates GKE TPU node pools. + 2. **Workload Deployment**: Runs a JAX benchmark JobSet to generate active TPU processes. + 3. **CLI Execution**: Iterates through defined `TpuInfoCmd` options within the worker pods. + 4. **Pattern Matching**: Performs regex and string validation on the CLI output to ensure accuracy. + 5. **Teardown**: Automatically cleans up all Kubernetes resources and node pools to minimize costs. """, ) as dag: for machine in MachineConfigMap: @@ -457,7 +444,7 @@ def generate_second_node_pool_name( ) outputs_of_tpu_info = ( - get_tpu_info_from_pod.override(task_id="get_tpu_info") + tpu_info.get_tpu_info_from_pod.override(task_id="get_tpu_info") .partial(info=cluster_info) .expand(pod_name=running_pods) ) @@ -507,6 +494,12 @@ def generate_second_node_pool_name( .expand(tpu_info_output=output_of_tpu_info) ) + cli_validation = ( + validate_tpu_info_cli.override(task_id="validate_tpu_info_cli") + .partial(info=cluster_info) + .expand(pod_name=pod_names) + ) + clean_up_workload = jobset.end_workload.override( task_id="clean_up_workload", trigger_rule=TriggerRule.ALL_DONE )( @@ -544,6 +537,7 @@ def generate_second_node_pool_name( validate_runtime_metric, validate_tensorcore_metric, validate_latency_metric, + cli_validation, ], ) @@ -566,120 +560,3 @@ def generate_second_node_pool_name( clean_up_workload, cleanup_node_pool, ) - # pylint: enable=pointless-statement - - -# Keyword arguments are generated dynamically at runtime (pylint does not -# know this signature). -with models.DAG( # pylint: disable=unexpected-keyword-arg - dag_id="tpu_info_cli_validation_dags", - start_date=datetime.datetime(2025, 8, 10), - schedule="30 14 * * *" if composer_env.is_prod_env() else None, - catchup=False, - tags=[ - "cloud-ml-auto-solutions", - "jobset", - "time-to-recover", - "tpu-observability", - "TPU", - "v6e-16", - ], - description=( - "Validates tpu-info CLI tool: help documentation, version metadata, " - "and process monitoring capabilities inside TPU worker pods." - ), - doc_md=""" - ### Description - This DAG performs an end-to-end validation of the `tpu-info` observability tool - within TPU worker pods. It ensures the CLI tool is correctly installed and - functional across different TPU configurations. - - ### Validation Steps: - 1. **Help Menu Validation**: Verifies `tpu-info -help` displays all required - options (streaming, rate, etc.) and specific usage instructions. - 2. **Process Table Validation**: Confirms `tpu-info --process` can successfully - map PIDs to TPU chips. - 3. **Version Validation**: Ensures `tpu-info --version` correctly reports - the tool version, libtpu version, and accelerator type. - """, -) 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}" - ): - jobset_config = jobset.build_jobset_from_gcs_yaml( - gcs_path=GCS_JOBSET_CONFIG_PATH, - dag_name="tpu_info_cli_validation_dags", - ) - - 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="tpu_info_cli_validation_dags", - is_prod=composer_env.is_prod_env(), - machine_type=config.machine_version.value, - tpu_topology=config.tpu_topology, - ) - - create_node_pool = node_pool.create.override(task_id="create_node_pool")( - node_pool=cluster_info, - ) - - apply_time = jobset.run_workload.override(task_id="run_workload")( - node_pool=cluster_info, - jobset_config=jobset_config, - workload_type=Workload.JAX_TPU_BENCHMARK, - ) - - pod_names = jobset.list_pod_names.override(task_id="list_pod_names")( - node_pool=cluster_info, - jobset_config=jobset_config, - ) - - wait_for_job_start = jobset.wait_for_jobset_started.override( - task_id="wait_for_job_start" - )(cluster_info, pod_name_list=pod_names, job_apply_time=apply_time) - - # Keyword arguments are generated dynamically at runtime (pylint does not - # know this signature). - with TaskGroup( # pylint: disable=unexpected-keyword-arg - group_id="verification_group" - ) as verification_group: - cli_validation = ( - validate_tpu_info_cli.override(task_id="validate_tpu_info_cli") - .partial(info=cluster_info) - .expand(pod_name=pod_names) - ) - - 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=apply_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( - jobset_config, - cluster_info, - create_node_pool, - apply_time, - pod_names, - wait_for_job_start, - verification_group, - cleanup_workload, - cleanup_node_pool, - ) - # pylint: enable=pointless-statement diff --git a/dags/tpu_observability/utils/tpu_cli_util.py b/dags/tpu_observability/utils/tpu_cli_util.py deleted file mode 100644 index dbbab07e2..000000000 --- a/dags/tpu_observability/utils/tpu_cli_util.py +++ /dev/null @@ -1,34 +0,0 @@ -from enum import Enum - - -class TpuInfoCmd(Enum): - HELP = "tpu-info -help" - VERSION = "tpu-info --version" - PROCESS = "tpu-info --process" - - -CLI_VALIDATION_SPEC: dict[TpuInfoCmd, list[str]] = { - TpuInfoCmd.HELP: [ - "Display TPU info and metrics.", - "options:", - "-h, --help", - "-v, --version", - "-p, --process", - "--streaming", - "--rate RATE", - "--list_metrics", - ], - TpuInfoCmd.VERSION: [ - "tpu-info version:", - "libtpu version:", - "accelerator type:", - ], - TpuInfoCmd.PROCESS: [ - "TPU Process Info", - "Chip", - "PID", - "Process Name", - "/dev/vfio/", - "python", - ], -} diff --git a/dags/tpu_observability/utils/tpu_info_util.py b/dags/tpu_observability/utils/tpu_info_util.py index c408463d4..765f1905f 100644 --- a/dags/tpu_observability/utils/tpu_info_util.py +++ b/dags/tpu_observability/utils/tpu_info_util.py @@ -1,11 +1,24 @@ """Utility for parsing the output of the 'tpu-info' command.""" -from dataclasses import dataclass -from enum import auto -from enum import IntEnum +import os import re +import tempfile +from dataclasses import dataclass +from enum import auto, Enum, IntEnum from airflow.decorators import task +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 import subprocess_util as subprocess + + +class TpuInfoCmd(Enum): + """Defines the available tpu-info CLI commands.""" + + HELP = "tpu-info -help" + VERSION = "tpu-info --version" + PROCESS = "tpu-info --process" + # A type alias for a parsed row, mapping column headers to their values. _TableRow = dict[str, str] @@ -96,6 +109,34 @@ def parse_tpu_info_output(output: str) -> list[Table]: return parsed_tables +@task +def get_tpu_info_from_pod(info: node_pool.Info, pod_name: str) -> str: + """ + Executes the `tpu-info` command in a specified pod and returns its output. + + This task uses kubectl to run the 'tpu-info' command inside the given pod + in the 'default' namespace. The output of the command is captured and + returned. + + Args: + kubeconfig: The path to the kubeconfig file. + pod_name: The name of the pod to execute the command in. + + Returns: + The standard output from the 'tpu-info' command. + """ + with tempfile.NamedTemporaryFile() as temp_config_file: + env = os.environ.copy() + env["KUBECONFIG"] = temp_config_file.name + + cmd = " && ".join([ + jobset.Command.get_credentials_command(info), + f"kubectl exec {pod_name} -n default -- tpu-info", + ]) + + return subprocess.run_exec(cmd, env=env) + + if __name__ == "__main__": full_output = """ Libtpu version: 0.0.23 From 5b86b8ecf9f722a9f7c5f793a771659e27531189 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Thu, 12 Feb 2026 14:28:46 +0800 Subject: [PATCH 18/21] refactor: Update DAG documentation for clarity on TPU observability tool validation --- .../tpu_info_format_validation_dags.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index 74f5b8a7f..8051c9c95 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -13,13 +13,18 @@ # limitations under the License. """ -tpu_info_format_validation_dag: -A DAG orchestrates the process of verifying TensorCore utilization metrics. -This is done by comparing data from Cloud Logging and Cloud Monitoring. - -tpu_info_cli_validation_dags: -A DAG to validate the `tpu-info` CLI tool, ensuring help documentation, -version metadata, and process monitoring are functional inside TPU worker pods. +tpu_info_validation_dag: +A comprehensive DAG that orchestrates the end-to-end validation of the +`tpu-info` observability tool. It performs two primary types of +verification: + +1. Format Validation: Parses the tool's raw output into structured tables + (TPU Chips, Runtime Utilization, TensorCore Utilization, and Latency) + and validates row counts and data integrity using regex. + +2. CLI Validation: Verifies command-line options (`-help`, `--version`, + `--process`) to ensure metadata, help documentation, and + process-to-chip mapping are functional inside live TPU worker pods. """ import datetime From 00cf3397b4ad74b615a55f068d484062164d984d Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Thu, 12 Feb 2026 17:36:49 +0800 Subject: [PATCH 19/21] refactor: Consolidate TPU info CLI command execution and enhance validation handling --- .../tpu_info_format_validation_dags.py | 49 ++++--------------- dags/tpu_observability/utils/tpu_info_util.py | 34 +++++++------ 2 files changed, 28 insertions(+), 55 deletions(-) diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index 8051c9c95..31bf8d225 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -273,55 +273,24 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]): ) -def execute_tpu_info_cli_command(info, pod_name: str, tpu_args: str) -> str: - """Helper to handle KUBECONFIG and execute kubectl.""" - with tempfile.NamedTemporaryFile() as temp_config_file: - env = os.environ.copy() - env["KUBECONFIG"] = temp_config_file.name - - cmd = " && ".join([ - jobset.Command.get_credentials_command(info), - f"kubectl exec {pod_name} -n default -- {tpu_args}", - ]) - return subprocess.run_exec(cmd, env=env) - - @task def validate_tpu_info_cli(info: node_pool.Info, pod_name: str) -> None: - """Validates tpu-info CLI commands with internal expectation patterns.""" - + """Validates tpu-info CLI commands using the consolidated utility.""" validation_spec = { - tpu_info.TpuInfoCmd.HELP: [ - "Display TPU info and metrics.", - "options:", - "-h, --help", - "-v, --version", - "-p, --process", - "--streaming", - "--rate RATE", - ], - tpu_info.TpuInfoCmd.VERSION: [ - "tpu-info version:", - "libtpu version:", - "accelerator type:", - ], - tpu_info.TpuInfoCmd.PROCESS: [ - "TPU Process Info", - "Chip", - "PID", - "Process Name", - "/dev/vfio/", - "python", - ], + tpu_info.TpuInfoCmd.HELP: ["--streaming", "--rate RATE"], + tpu_info.TpuInfoCmd.VERSION: ["tpu-info version:", "libtpu version:"], + tpu_info.TpuInfoCmd.PROCESS: ["TPU Process Info", "/dev/vfio/", "python"], } for cmd_enum, patterns in validation_spec.items(): - output = execute_tpu_info_cli_command(info, pod_name, cmd_enum.value) + output = tpu_info.get_tpu_info_from_pod( + info, pod_name, cmd_str=cmd_enum.value + ) + for pattern in patterns: if pattern not in output: raise AssertionError( - f"Validation failed for '{cmd_enum.value}': " - f"Missing expected pattern '{pattern}'." + f"Validation failed for '{cmd_enum.value}': Missing pattern '{pattern}'." ) diff --git a/dags/tpu_observability/utils/tpu_info_util.py b/dags/tpu_observability/utils/tpu_info_util.py index 765f1905f..fa65aea5d 100644 --- a/dags/tpu_observability/utils/tpu_info_util.py +++ b/dags/tpu_observability/utils/tpu_info_util.py @@ -15,6 +15,7 @@ class TpuInfoCmd(Enum): """Defines the available tpu-info CLI commands.""" + TPU_INFO = "tpu-info" HELP = "tpu-info -help" VERSION = "tpu-info --version" PROCESS = "tpu-info --process" @@ -36,7 +37,8 @@ def parse_body(self): """Parses the raw_body string to populate the structured body attribute.""" class TableLineIndex(IntEnum): - """Below is an example of the text returned by tpu-info, formatted as a table. + """Below is an example of the text returned by tpu-info, + formatted as a table. TPU Chips ┏━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━┓ @@ -109,21 +111,15 @@ def parse_tpu_info_output(output: str) -> list[Table]: return parsed_tables -@task -def get_tpu_info_from_pod(info: node_pool.Info, pod_name: str) -> str: +def get_tpu_info_from_pod( + info: node_pool.Info, pod_name: str, cmd_str: str +) -> str: """ - Executes the `tpu-info` command in a specified pod and returns its output. - - This task uses kubectl to run the 'tpu-info' command inside the given pod - in the 'default' namespace. The output of the command is captured and - returned. - - Args: - kubeconfig: The path to the kubeconfig file. - pod_name: The name of the pod to execute the command in. + Executes a command (default: tpu-info) in a specified pod + and returns its output. - Returns: - The standard output from the 'tpu-info' command. + Consolidated version that handles both standard tpu-info calls and + specific CLI flag validation. """ with tempfile.NamedTemporaryFile() as temp_config_file: env = os.environ.copy() @@ -131,12 +127,20 @@ def get_tpu_info_from_pod(info: node_pool.Info, pod_name: str) -> str: cmd = " && ".join([ jobset.Command.get_credentials_command(info), - f"kubectl exec {pod_name} -n default -- tpu-info", + f"kubectl exec {pod_name} -n default -- {cmd_str}", ]) return subprocess.run_exec(cmd, env=env) +@task +def get_tpu_info_from_pod_task( + info: node_pool.Info, pod_name: str, cmd_str: str +) -> str: + """Airflow task wrapper for get_tpu_info_from_pod.""" + return get_tpu_info_from_pod(info, pod_name, cmd_str) + + if __name__ == "__main__": full_output = """ Libtpu version: 0.0.23 From 8d53f3893ef60a1a96f065fa40ad002334763cf0 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Fri, 13 Feb 2026 15:42:27 +0800 Subject: [PATCH 20/21] refactor: Enhance TPU info CLI validation by updating command output handling and pattern matching --- .../tpu_info_format_validation_dags.py | 60 +++++++++++-------- dags/tpu_observability/utils/tpu_info_util.py | 23 ++----- 2 files changed, 41 insertions(+), 42 deletions(-) diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index 31bf8d225..5030508fd 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -274,24 +274,24 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]): @task -def validate_tpu_info_cli(info: node_pool.Info, pod_name: str) -> None: - """Validates tpu-info CLI commands using the consolidated utility.""" - validation_spec = { - tpu_info.TpuInfoCmd.HELP: ["--streaming", "--rate RATE"], - tpu_info.TpuInfoCmd.VERSION: ["tpu-info version:", "libtpu version:"], - tpu_info.TpuInfoCmd.PROCESS: ["TPU Process Info", "/dev/vfio/", "python"], +def validate_tpu_info_patterns(output: str, cmd_name: str): + """Matches output against patterns defined in a local spec.""" + patterns_map = { + tpu_info.TpuInfoCmd.HELP.value: ["--streaming", "--rate RATE"], + tpu_info.TpuInfoCmd.VERSION.value: [ + "tpu-info version:", + "libtpu version:", + ], + tpu_info.TpuInfoCmd.PROCESS.value: [ + "TPU Process Info", + "/dev/vfio/", + "python", + ], } - - for cmd_enum, patterns in validation_spec.items(): - output = tpu_info.get_tpu_info_from_pod( - info, pod_name, cmd_str=cmd_enum.value - ) - - for pattern in patterns: - if pattern not in output: - raise AssertionError( - f"Validation failed for '{cmd_enum.value}': Missing pattern '{pattern}'." - ) + patterns = patterns_map.get(cmd_name, []) + for pattern in patterns: + if pattern not in output: + raise AssertionError(f"Cmd '{cmd_name}' missing pattern: {pattern}") # Keyword arguments are generated dynamically at runtime (pylint does not @@ -428,7 +428,20 @@ def generate_second_node_pool_name( task_id="get_each_metric_table" ) .partial() - .expand(output=outputs_of_tpu_info) + .expand(output=raw_metric_data.map(lambda x: x["output"])) + ) + + cli_raw_results = ( + tpu_info.get_tpu_info_from_pod.override(task_id="get_cli_output") + .partial(info=cluster_info) + .expand( + pod_name=pod_names, + cmd_str=[ + tpu_info.TpuInfoCmd.HELP.value, + tpu_info.TpuInfoCmd.VERSION.value, + tpu_info.TpuInfoCmd.PROCESS.value, + ], + ) ) # Keyword arguments are generated dynamically at runtime (pylint does not @@ -468,11 +481,9 @@ def generate_second_node_pool_name( .expand(tpu_info_output=output_of_tpu_info) ) - cli_validation = ( - validate_tpu_info_cli.override(task_id="validate_tpu_info_cli") - .partial(info=cluster_info) - .expand(pod_name=pod_names) - ) + cli_validation = validate_tpu_info_patterns.override( + task_id="validate_cli_output" + ).expand_kwargs(cli_raw_results) clean_up_workload = jobset.end_workload.override( task_id="clean_up_workload", trigger_rule=TriggerRule.ALL_DONE @@ -528,8 +539,9 @@ def generate_second_node_pool_name( apply_time, running_pods, wait_for_job_start, - outputs_of_tpu_info, + raw_metric_data, output_of_tpu_info, + cli_raw_results, verification_group, clean_up_workload, cleanup_node_pool, diff --git a/dags/tpu_observability/utils/tpu_info_util.py b/dags/tpu_observability/utils/tpu_info_util.py index fa65aea5d..f8a3a2b0c 100644 --- a/dags/tpu_observability/utils/tpu_info_util.py +++ b/dags/tpu_observability/utils/tpu_info_util.py @@ -111,34 +111,21 @@ def parse_tpu_info_output(output: str) -> list[Table]: return parsed_tables +@task def get_tpu_info_from_pod( info: node_pool.Info, pod_name: str, cmd_str: str -) -> str: - """ - Executes a command (default: tpu-info) in a specified pod - and returns its output. - - Consolidated version that handles both standard tpu-info calls and - specific CLI flag validation. - """ +) -> dict: + """Executes command and returns a dict containing metadata.""" with tempfile.NamedTemporaryFile() as temp_config_file: env = os.environ.copy() env["KUBECONFIG"] = temp_config_file.name - cmd = " && ".join([ jobset.Command.get_credentials_command(info), f"kubectl exec {pod_name} -n default -- {cmd_str}", ]) + output = subprocess.run_exec(cmd, env=env) - return subprocess.run_exec(cmd, env=env) - - -@task -def get_tpu_info_from_pod_task( - info: node_pool.Info, pod_name: str, cmd_str: str -) -> str: - """Airflow task wrapper for get_tpu_info_from_pod.""" - return get_tpu_info_from_pod(info, pod_name, cmd_str) + return {"output": output, "cmd_name": cmd_str} if __name__ == "__main__": From 79f7f5cede2e600f2128ff07ef43c27466fb0602 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Fri, 13 Feb 2026 17:46:03 +0800 Subject: [PATCH 21/21] refactor: Update TPU info retrieval to include command string in pod info fetching --- dags/tpu_observability/tpu_info_format_validation_dags.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/dags/tpu_observability/tpu_info_format_validation_dags.py b/dags/tpu_observability/tpu_info_format_validation_dags.py index 5030508fd..a7dbb5c13 100644 --- a/dags/tpu_observability/tpu_info_format_validation_dags.py +++ b/dags/tpu_observability/tpu_info_format_validation_dags.py @@ -417,9 +417,11 @@ def generate_second_node_pool_name( job_apply_time=apply_time, ) - outputs_of_tpu_info = ( + raw_metric_data = ( tpu_info.get_tpu_info_from_pod.override(task_id="get_tpu_info") - .partial(info=cluster_info) + .partial( + info=cluster_info, cmd_str=tpu_info.TpuInfoCmd.TPU_INFO.value + ) .expand(pod_name=running_pods) ) @@ -435,7 +437,7 @@ def generate_second_node_pool_name( tpu_info.get_tpu_info_from_pod.override(task_id="get_cli_output") .partial(info=cluster_info) .expand( - pod_name=pod_names, + pod_name=running_pods, cmd_str=[ tpu_info.TpuInfoCmd.HELP.value, tpu_info.TpuInfoCmd.VERSION.value,