From 49ba692657122bb4032c7f2179fe81d548b1522f Mon Sep 17 00:00:00 2001 From: Alfred Yu Date: Tue, 12 Aug 2025 10:03:26 +0800 Subject: [PATCH 01/14] [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 e2ba42c08..9458fc834 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 9e3fa66a5007ba2afd86318b5d3e2f3d961dc436 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Fri, 9 Jan 2026 17:01:55 +0800 Subject: [PATCH 02/14] feat: Add DAG to validate tpu-info streaming performance --- .../tpu_info_streaming_rate.py | 255 ++++++++++++++++++ 1 file changed, 255 insertions(+) create mode 100644 dags/tpu_observability/tpu_info_streaming_rate.py diff --git a/dags/tpu_observability/tpu_info_streaming_rate.py b/dags/tpu_observability/tpu_info_streaming_rate.py new file mode 100644 index 000000000..302ddc539 --- /dev/null +++ b/dags/tpu_observability/tpu_info_streaming_rate.py @@ -0,0 +1,255 @@ +# 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 +""" + +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 + + +# --- Helper Methods --- + + +def execute_tpu_info_cli_command(info, pod_name: str, tpu_args: str) -> str: + """Helper to handle KUBECONFIG and execute kubectl exec.""" + with tempfile.NamedTemporaryFile() as temp_config_file: + env = os.environ.copy() + env["KUBECONFIG"] = temp_config_file.name + + # Note: jobset.Command.get_credentials_command(info) must be accessible + cmd = " && ".join([ + jobset.Command.get_credentials_command(info), + f"kubectl exec {pod_name} -n default -- {tpu_args}", + ]) + # Returning the stdout of the command execution + return subprocess.run_exec(cmd, env=env) + + +def verify_output_contains_patterns( + output: str, patterns: List[str], context: str +): + """Verifies that expected strings exist in the output.""" + for pattern in patterns: + if pattern not in output: + raise AssertionError( + f"Validation failed for '{context}': Missing '{pattern}'." + ) + + +def verify_streaming_frequency(output: str, rate: float, duration: int): + """Quantitatively verifies if the tool outputs data at the requested rate.""" + # Count occurrences of a specific header unique to each data burst + sample_count = output.count("TPU Metrics") + expected_samples = duration / rate + + # 50% tolerance for Kubernetes network latency and tool initialization + threshold = expected_samples * 0.5 + + print( + f"Rate: {rate}s | Expected: ~{expected_samples} | Actual: {sample_count}" + ) + + if sample_count < threshold: + raise AssertionError( + f"Frequency too low. Expected at least {threshold} samples, got {sample_count}." + ) + + +# --- Task Definitions --- + + +@task +def validate_streaming_rate(info, pod_name: str, rate: float) -> str: + """ + Mapped Task: Validates tpu-info streaming performance. + Each instance tests one pod at one specific rate. + """ + # Ensure duration is long enough to catch slow rates (e.g., 5s rate needs >10s) + duration = max(10, int(rate * 3)) + + # Use timeout to stop the infinite streaming output + tpu_args = f"timeout {duration}s tpu-info --streaming --rate {rate}" + + print(f"Validating Pod: {pod_name} at Rate: {rate}s for {duration}s") + output = execute_tpu_info_cli_command(info, pod_name, tpu_args) + + # 1. Check for required UI elements/metrics + verify_output_contains_patterns( + output, + ["TPU Metrics", "Duty Cycle", "Memory Usage"], + f"Streaming check on {pod_name}", + ) + + # 2. Check for timing accuracy + verify_streaming_frequency(output, rate, duration) + + return f"Completed: {pod_name} at {rate}s" + + +# 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_verify_streaming_rate_dags", + start_date=datetime.datetime(2025, 8, 10), + schedule="0 19 * * *" 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 + + jobset_config = JobSet( + jobset_name="tpu-info-verify-streaming-rate-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_verify_streaming_rate_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: + # 1. Define the static rates you want to test + test_rates = [0.1, 0.5, 1.0, 5.0] + + # 2. Use Dynamic Task Mapping (.expand) + # This creates a Cross-Product: (number of pods) x (number of rates) + # If pod_names has 8 pods, this will trigger 32 task instances. + streaming_validation_results = ( + validate_streaming_rate.override(task_id="streaming_rate_test") + .partial(info=cluster_info) + .expand( + pod_name=pod_names, # XComArg from a previous task + rate=test_rates, # Standard Python list + ) + ) + + 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 fd610f05efdef25cbad18d1f1a494af1c2406e5a Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Tue, 13 Jan 2026 10:04:49 +0800 Subject: [PATCH 03/14] feat: Enhance streaming rate validation with detailed logging and improved frequency checks --- .../tpu_info_streaming_rate.py | 128 ++++++++++++------ 1 file changed, 85 insertions(+), 43 deletions(-) diff --git a/dags/tpu_observability/tpu_info_streaming_rate.py b/dags/tpu_observability/tpu_info_streaming_rate.py index 302ddc539..25544f122 100644 --- a/dags/tpu_observability/tpu_info_streaming_rate.py +++ b/dags/tpu_observability/tpu_info_streaming_rate.py @@ -21,6 +21,8 @@ import subprocess import tempfile from typing import List +import logging +import re from airflow import models from airflow.decorators import task @@ -64,54 +66,94 @@ def verify_output_contains_patterns( ) -def verify_streaming_frequency(output: str, rate: float, duration: int): - """Quantitatively verifies if the tool outputs data at the requested rate.""" - # Count occurrences of a specific header unique to each data burst - sample_count = output.count("TPU Metrics") - expected_samples = duration / rate +# --- Helper: Verification Logic --- - # 50% tolerance for Kubernetes network latency and tool initialization - threshold = expected_samples * 0.5 - print( - f"Rate: {rate}s | Expected: ~{expected_samples} | Actual: {sample_count}" +def verify_streaming_frequency(output: str, rate: float, duration: int): + """ + Verifies if the sampling rate is effective by distinguishing between + UI Rendering frequency and Actual Data Update frequency. + """ + # 1. Capture UI Rendering Cycles (Equivalent to: grep -o $'\x1b\[H' | wc -l) + # This proves the --rate argument was accepted and the loop is running. + ui_render_count = output.count("\x1b[H") + + # 2. Capture Unique Data Points (Equivalent to: grep "Last update:" | sort -u | wc -l) + # This accounts for the 1-second display precision limit. + timestamp_pattern = r"Last update: ([\d\- :]+) UTC" + found_timestamps = re.findall(timestamp_pattern, output) + unique_data_points = len(set(found_timestamps)) + + # 3. Calculate Expectations + # Theoretical UI cycles (limited by tool's 4Hz floor if rate > 0.25s) + expected_ui_cycles = duration / rate + + # Real data updates (limited by tool's 1Hz sampling ceiling) + effective_data_rate = max(1.0, rate) + expected_data_points = duration / effective_data_rate + + logging.info(f"[Verify] Target Rate: {rate}s | Duration: {duration}s") + logging.info( + f"[Verify] UI Renders: Found {ui_render_count} | Expected ~{expected_ui_cycles}" + ) + logging.info( + f"[Verify] Data Points: Found {unique_data_points} | Expected ~{expected_data_points}" ) - if sample_count < threshold: + # --- 4. Assertions --- + + # Check A: Data Consistency + # We allow a margin of 2 points for initialization lag and clock boundary crossing. + if unique_data_points < (expected_data_points - 2): raise AssertionError( - f"Frequency too low. Expected at least {threshold} samples, got {sample_count}." + f"Data update frequency too low. Rate {rate}s should yield ~{expected_data_points} " + f"points in {duration}s, but only found {unique_data_points}." ) + # Check B: UI High-Frequency Enforcement (Only for rates <= 0.2s) + # This confirms the tool actually speeds up the rendering loop as requested. + if rate <= 0.2: + if ui_render_count < expected_ui_cycles * 0.7: + raise AssertionError( + f"High-frequency UI rendering failed. Requested {rate}s but " + f"rendering count {ui_render_count} is too low." + ) + + logging.info(f"Frequency validation successful for rate {rate}s.") -# --- Task Definitions --- + +# --- The Task Definition --- @task def validate_streaming_rate(info, pod_name: str, rate: float) -> str: """ - Mapped Task: Validates tpu-info streaming performance. - Each instance tests one pod at one specific rate. + Executes tpu-info --streaming and validates frequency using UI and Data metrics. """ - # Ensure duration is long enough to catch slow rates (e.g., 5s rate needs >10s) - duration = max(10, int(rate * 3)) - - # Use timeout to stop the infinite streaming output - tpu_args = f"timeout {duration}s tpu-info --streaming --rate {rate}" + # Duration: 15s provides enough samples to overcome the 'Last update' precision limit + duration = 15 + + # IMPORTANT: We use 'script -q -c' to simulate a TTY environment. + # Without this, kubectl exec may strip the ANSI control codes (\x1b[H) + # used to count UI refreshes. + tpu_args = ( + f"sh -c \"script -q -c 'timeout {duration}s tpu-info --streaming --rate {rate}' /dev/null\" " + f"|| [ $? -eq 124 ]" + ) - print(f"Validating Pod: {pod_name} at Rate: {rate}s for {duration}s") + logging.info(f"Validating Pod: {pod_name} | Requested Rate: {rate}s") output = execute_tpu_info_cli_command(info, pod_name, tpu_args) - # 1. Check for required UI elements/metrics + # 1. Verify basic output content + patterns = ["TPU Runtime Utilization", "HBM Usage (GiB)", "Last update:"] verify_output_contains_patterns( - output, - ["TPU Metrics", "Duty Cycle", "Memory Usage"], - f"Streaming check on {pod_name}", + output, patterns, f"Content check on {pod_name}" ) - # 2. Check for timing accuracy + # 2. Verify frequency logic (UI vs Data) verify_streaming_frequency(output, rate, duration) - return f"Completed: {pod_name} at {rate}s" + return f"Validated {pod_name} at {rate}s" # Keyword arguments are generated dynamically at runtime (pylint does not @@ -224,21 +266,21 @@ def validate_streaming_rate(info, pod_name: str, rate: float) -> str: ) ) - 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, - ) + # 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 @@ -249,7 +291,7 @@ def validate_streaming_rate(info, pod_name: str, rate: float) -> str: >> pod_names >> wait_for_job_start >> verification_group - >> cleanup_workload - >> cleanup_node_pool + # >> cleanup_workload + # >> cleanup_node_pool ) # pylint: enable=pointless-statement From 65af17e589f84e4e672f7e70f7ccd96ac550eb55 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Tue, 13 Jan 2026 10:28:30 +0800 Subject: [PATCH 04/14] fix: Remove schedule from DAG definition for streaming rate validation --- dags/tpu_observability/tpu_info_streaming_rate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dags/tpu_observability/tpu_info_streaming_rate.py b/dags/tpu_observability/tpu_info_streaming_rate.py index 25544f122..d8307b3b2 100644 --- a/dags/tpu_observability/tpu_info_streaming_rate.py +++ b/dags/tpu_observability/tpu_info_streaming_rate.py @@ -161,7 +161,7 @@ def validate_streaming_rate(info, pod_name: str, rate: float) -> str: with models.DAG( # pylint: disable=unexpected-keyword-arg dag_id="tpu_info_verify_streaming_rate_dags", start_date=datetime.datetime(2025, 8, 10), - schedule="0 19 * * *" if composer_env.is_prod_env() else None, + schedule=None, catchup=False, tags=[ "cloud-ml-auto-solutions", From 995c79c3446ce4cc67e0d54c8d226d1a40261b92 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Tue, 13 Jan 2026 14:19:50 +0800 Subject: [PATCH 05/14] feat: Refactor DAG for tpu-info streaming rate verification with improved documentation and streamlined logic --- .../tpu_info_streaming_rate.py | 143 ++++-------------- 1 file changed, 28 insertions(+), 115 deletions(-) diff --git a/dags/tpu_observability/tpu_info_streaming_rate.py b/dags/tpu_observability/tpu_info_streaming_rate.py index d8307b3b2..d54a34257 100644 --- a/dags/tpu_observability/tpu_info_streaming_rate.py +++ b/dags/tpu_observability/tpu_info_streaming_rate.py @@ -13,7 +13,7 @@ # limitations under the License. """ -A +DAG to verify tpu-info streaming rate functionality on TPU v6e slices. """ import datetime @@ -37,21 +37,16 @@ from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH -# --- Helper Methods --- - - def execute_tpu_info_cli_command(info, pod_name: str, tpu_args: str) -> str: """Helper to handle KUBECONFIG and execute kubectl exec.""" with tempfile.NamedTemporaryFile() as temp_config_file: env = os.environ.copy() env["KUBECONFIG"] = temp_config_file.name - # Note: jobset.Command.get_credentials_command(info) must be accessible cmd = " && ".join([ jobset.Command.get_credentials_command(info), f"kubectl exec {pod_name} -n default -- {tpu_args}", ]) - # Returning the stdout of the command execution return subprocess.run_exec(cmd, env=env) @@ -66,93 +61,23 @@ def verify_output_contains_patterns( ) -# --- Helper: Verification Logic --- - - -def verify_streaming_frequency(output: str, rate: float, duration: int): - """ - Verifies if the sampling rate is effective by distinguishing between - UI Rendering frequency and Actual Data Update frequency. - """ - # 1. Capture UI Rendering Cycles (Equivalent to: grep -o $'\x1b\[H' | wc -l) - # This proves the --rate argument was accepted and the loop is running. - ui_render_count = output.count("\x1b[H") - - # 2. Capture Unique Data Points (Equivalent to: grep "Last update:" | sort -u | wc -l) - # This accounts for the 1-second display precision limit. - timestamp_pattern = r"Last update: ([\d\- :]+) UTC" - found_timestamps = re.findall(timestamp_pattern, output) - unique_data_points = len(set(found_timestamps)) - - # 3. Calculate Expectations - # Theoretical UI cycles (limited by tool's 4Hz floor if rate > 0.25s) - expected_ui_cycles = duration / rate - - # Real data updates (limited by tool's 1Hz sampling ceiling) - effective_data_rate = max(1.0, rate) - expected_data_points = duration / effective_data_rate - - logging.info(f"[Verify] Target Rate: {rate}s | Duration: {duration}s") - logging.info( - f"[Verify] UI Renders: Found {ui_render_count} | Expected ~{expected_ui_cycles}" - ) - logging.info( - f"[Verify] Data Points: Found {unique_data_points} | Expected ~{expected_data_points}" - ) - - # --- 4. Assertions --- - - # Check A: Data Consistency - # We allow a margin of 2 points for initialization lag and clock boundary crossing. - if unique_data_points < (expected_data_points - 2): - raise AssertionError( - f"Data update frequency too low. Rate {rate}s should yield ~{expected_data_points} " - f"points in {duration}s, but only found {unique_data_points}." - ) - - # Check B: UI High-Frequency Enforcement (Only for rates <= 0.2s) - # This confirms the tool actually speeds up the rendering loop as requested. - if rate <= 0.2: - if ui_render_count < expected_ui_cycles * 0.7: - raise AssertionError( - f"High-frequency UI rendering failed. Requested {rate}s but " - f"rendering count {ui_render_count} is too low." - ) - - logging.info(f"Frequency validation successful for rate {rate}s.") - - -# --- The Task Definition --- - - @task def validate_streaming_rate(info, pod_name: str, rate: float) -> str: """ Executes tpu-info --streaming and validates frequency using UI and Data metrics. """ - # Duration: 15s provides enough samples to overcome the 'Last update' precision limit duration = 15 - # IMPORTANT: We use 'script -q -c' to simulate a TTY environment. - # Without this, kubectl exec may strip the ANSI control codes (\x1b[H) - # used to count UI refreshes. tpu_args = ( f"sh -c \"script -q -c 'timeout {duration}s tpu-info --streaming --rate {rate}' /dev/null\" " f"|| [ $? -eq 124 ]" ) - - logging.info(f"Validating Pod: {pod_name} | Requested Rate: {rate}s") output = execute_tpu_info_cli_command(info, pod_name, tpu_args) - # 1. Verify basic output content - patterns = ["TPU Runtime Utilization", "HBM Usage (GiB)", "Last update:"] + patterns = ["Refresh rate:", f"{rate}s"] verify_output_contains_patterns( output, patterns, f"Content check on {pod_name}" ) - - # 2. Verify frequency logic (UI vs Data) - verify_streaming_frequency(output, rate, duration) - return f"Validated {pod_name} at {rate}s" @@ -166,29 +91,21 @@ def validate_streaming_rate(info, pod_name: str, rate: float) -> str: tags=[ "cloud-ml-auto-solutions", "jobset", - "time-to-recover", "tpu-observability", "TPU", "v6e-16", + "tpu-info", + "streaming-rate", ], description=( - "Validates tpu-info CLI tool: help documentation, version metadata, " - "and process monitoring capabilities inside TPU worker pods." + "DAG to verify tpu-info streaming rate functionality on TPU v6e slices." ), 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. - """, + ## TPU Info Streaming Rate Verification DAG + This DAG validates the `tpu-info` CLI tool's ability to stream TPU metrics + at specified rates inside TPU worker pods. It ensures that the tool adheres to + the requested streaming frequency and accurately reflects TPU status updates. + """, ) as dag: for machine in MachineConfigMap: config = machine.value @@ -251,36 +168,32 @@ def validate_streaming_rate(info, pod_name: str, rate: float) -> str: with TaskGroup( # pylint: disable=unexpected-keyword-arg group_id="verification_group" ) as verification_group: - # 1. Define the static rates you want to test test_rates = [0.1, 0.5, 1.0, 5.0] - # 2. Use Dynamic Task Mapping (.expand) - # This creates a Cross-Product: (number of pods) x (number of rates) - # If pod_names has 8 pods, this will trigger 32 task instances. streaming_validation_results = ( validate_streaming_rate.override(task_id="streaming_rate_test") .partial(info=cluster_info) .expand( - pod_name=pod_names, # XComArg from a previous task - rate=test_rates, # Standard Python list + pod_name=pod_names, + rate=test_rates, ) ) - # 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_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, - # ) + 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 @@ -291,7 +204,7 @@ def validate_streaming_rate(info, pod_name: str, rate: float) -> str: >> pod_names >> wait_for_job_start >> verification_group - # >> cleanup_workload - # >> cleanup_node_pool + >> cleanup_workload + >> cleanup_node_pool ) # pylint: enable=pointless-statement From bf5aa1101877cf389cc1d1b468342de398bc8d25 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Tue, 13 Jan 2026 14:34:52 +0800 Subject: [PATCH 06/14] feat: Update DAG description and documentation for tpu-info streaming rate verification --- .../tpu_observability/tpu_info_streaming_rate.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/dags/tpu_observability/tpu_info_streaming_rate.py b/dags/tpu_observability/tpu_info_streaming_rate.py index d54a34257..b785cdeaa 100644 --- a/dags/tpu_observability/tpu_info_streaming_rate.py +++ b/dags/tpu_observability/tpu_info_streaming_rate.py @@ -9,7 +9,7 @@ # 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 +# See the License foar the specific language governing permissions and # limitations under the License. """ @@ -21,8 +21,6 @@ import subprocess import tempfile from typing import List -import logging -import re from airflow import models from airflow.decorators import task @@ -98,13 +96,15 @@ def validate_streaming_rate(info, pod_name: str, rate: float) -> str: "streaming-rate", ], description=( - "DAG to verify tpu-info streaming rate functionality on TPU v6e slices." + "Automated validation of the tpu-info CLI's --rate parameter, " + "ensuring accurate metric streaming frequencies on TPU v6e-16 slices." ), doc_md=""" - ## TPU Info Streaming Rate Verification DAG - This DAG validates the `tpu-info` CLI tool's ability to stream TPU metrics - at specified rates inside TPU worker pods. It ensures that the tool adheres to - the requested streaming frequency and accurately reflects TPU status updates. + ## TPU Info Streaming Rate Verification DAG + + This DAG automates the functional testing of the `tpu-info` CLI tool, specifically focusing on the + `--streaming` and `--rate` flags. It verifies that the tool correctly honors requested update + intervals within a Kubernetes-managed TPU environment. """, ) as dag: for machine in MachineConfigMap: From 5250c9c4805f91a686590bafa29838f4452dc6c9 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Wed, 14 Jan 2026 10:41:27 +0800 Subject: [PATCH 07/14] feat: Enhance node pool management in tpu-info streaming rate DAG with additional node pool creation and cleanup tasks --- .../tpu_info_streaming_rate.py | 55 ++++++++++++++++--- 1 file changed, 47 insertions(+), 8 deletions(-) diff --git a/dags/tpu_observability/tpu_info_streaming_rate.py b/dags/tpu_observability/tpu_info_streaming_rate.py index b785cdeaa..5fdae4fe4 100644 --- a/dags/tpu_observability/tpu_info_streaming_rate.py +++ b/dags/tpu_observability/tpu_info_streaming_rate.py @@ -110,12 +110,19 @@ def validate_streaming_rate(info, pod_name: str, rate: float) -> str: for machine in MachineConfigMap: config = machine.value + @task + def generate_second_node_pool_name( + node_pool_info: node_pool.Info, + ) -> str: + """Generates a second node pool name.""" + return f"{node_pool_info.node_pool_name}-2" + jobset_config = JobSet( jobset_name="tpu-info-verify-streaming-rate-jobset", namespace="default", max_restarts=5, replicated_job_name="tpu-job-slice", - replicas=1, + replicas=2, backoff_limit=0, completions=4, parallelism=4, @@ -142,10 +149,27 @@ def validate_streaming_rate(info, pod_name: str, rate: float) -> str: tpu_topology=config.tpu_topology, ) - create_node_pool = node_pool.create.override(task_id="create_node_pool")( - node_pool=cluster_info, + cluster_info_2 = node_pool.copy_node_pool_info_with_override( + info=cluster_info, + node_pool_name=generate_second_node_pool_name(cluster_info), ) + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id="create_node_pool" + ) as create_node_pool: + create_first_node_pool = node_pool.create.override( + task_id="node_pool_1" + )( + node_pool=cluster_info, + ) + + create_second_node_pool = node_pool.create.override( + task_id="node_pool_2", + retries=2, + )( + node_pool=cluster_info_2, + ) + apply_time = jobset.run_workload.override(task_id="run_workload")( node_pool=cluster_info, yaml_config=jobset_config.generate_yaml( @@ -189,11 +213,26 @@ def validate_streaming_rate(info, pod_name: str, rate: float) -> str: 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, - ) + # Keyword arguments are generated dynamically at runtime (pylint does not + # know this signature). + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id="cleanup_node_pool" + ) as cleanup_node_pool: + cleanup_first_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool_1", + trigger_rule=TriggerRule.ALL_DONE, + retries=2, + )(node_pool=cluster_info).as_teardown( + setups=create_node_pool, + ) + + cleanup_second_node_pool = node_pool.delete.override( + task_id="cleanup_node_pool_2", + trigger_rule=TriggerRule.ALL_DONE, + retries=2, + )(node_pool=cluster_info_2).as_teardown( + setups=create_node_pool, + ) # Airflow uses >> for task chaining, which is pointless for pylint. # pylint: disable=pointless-statement From dd73c88d251fff0862baee2044adb3a3b149ca22 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Wed, 14 Jan 2026 15:30:30 +0800 Subject: [PATCH 08/14] feat: Refactor streaming rate validation to create dynamic task groups for each rate --- .../tpu_info_streaming_rate.py | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/dags/tpu_observability/tpu_info_streaming_rate.py b/dags/tpu_observability/tpu_info_streaming_rate.py index 5fdae4fe4..f17317685 100644 --- a/dags/tpu_observability/tpu_info_streaming_rate.py +++ b/dags/tpu_observability/tpu_info_streaming_rate.py @@ -67,7 +67,8 @@ def validate_streaming_rate(info, pod_name: str, rate: float) -> str: duration = 15 tpu_args = ( - f"sh -c \"script -q -c 'timeout {duration}s tpu-info --streaming --rate {rate}' /dev/null\" " + f"sh -c \"script -q -c 'timeout {duration}s " + f"tpu-info --streaming --rate {rate}' /dev/null\" " f"|| [ $? -eq 124 ]" ) output = execute_tpu_info_cli_command(info, pod_name, tpu_args) @@ -187,21 +188,20 @@ def generate_second_node_pool_name( 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: - test_rates = [0.1, 0.5, 1.0, 5.0] - - streaming_validation_results = ( - validate_streaming_rate.override(task_id="streaming_rate_test") - .partial(info=cluster_info) - .expand( - pod_name=pod_names, - rate=test_rates, - ) - ) + test_rates = [0.1, 0.5, 1.0, 5.0] + for rate in test_rates: + formatted_rate = str(rate).replace(".", "_") + + # Keyword arguments are generated dynamically at runtime (pylint does not + # know this signature). + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id=f"verification_group_rate_{formatted_rate}" + ) as rate_group: + streaming_validation_results = ( + validate_streaming_rate.override(task_id="streaming_rate_test") + .partial(info=cluster_info, rate=rate) + .expand(pod_name=pod_names) + ) cleanup_workload = jobset.end_workload.override( task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE From d91aa606854704f98fc149282f967af84517b3ed Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Thu, 15 Jan 2026 15:41:04 +0800 Subject: [PATCH 09/14] feat: Implement TPUPerformanceAnalyzer for enhanced streaming rate validation and reporting --- .../tpu_info_streaming_rate.py | 248 +++++++++++++++++- 1 file changed, 240 insertions(+), 8 deletions(-) diff --git a/dags/tpu_observability/tpu_info_streaming_rate.py b/dags/tpu_observability/tpu_info_streaming_rate.py index f17317685..eaf9b6804 100644 --- a/dags/tpu_observability/tpu_info_streaming_rate.py +++ b/dags/tpu_observability/tpu_info_streaming_rate.py @@ -18,9 +18,11 @@ import datetime import os +import re import subprocess import tempfile from typing import List +import logging from airflow import models from airflow.decorators import task @@ -35,6 +37,178 @@ from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH +class TPUPerformanceAnalyzer: + + def __init__(self, target_rate: float = 0.1): + """ + Initialize the analyzer with a target sampling rate. + :param target_rate: The expected update interval in seconds (default 0.1s). + """ + self.target_rate = target_rate + self.lower_bound = target_rate * 0.8 + self.upper_bound = target_rate * 1.2 + self.frames = [] + self.update_events = [] + + # Pre-compile regex patterns for optimized performance + self._frame_start_re = re.compile(r"\[(\d{2}:\d{2}:\d{2}\.\d{3})\].*?\[H") + self._chips_re = re.compile( + r"│\s+(/dev/vfio/\d+)\s+│.*?│\s+\d+\s+│\s+(\d+)\s+│" + ) + self._runtime_re = re.compile( + r"│\s+(\d+)\s+│\s+([\d.]+ GiB / [\d.]+ GiB)\s+│\s+([\d.]+)%\s+│" + ) + self._tensor_re = re.compile(r"│\s+(\d+)\s+│\s+([\d.]+)%\s+│") + self._latency_re = re.compile( + r"│\s+([\dMB+]+)\s+│\s+([\d.]+) us\s+│\s+([\d.]+) us\s+│\s+([\d.]+) us\s+│\s+([\d.]+) us\s+│" + ) + + def _init_empty_frame(self): + """Internal helper to initialize a structure for a single log frame.""" + return {"ts": None, "chips": {}, "runtime": {}, "tensor": {}, "latency": {}} + + def _extract_metrics(self, line, frame): + """Internal helper to extract various hardware metrics from a log line.""" + m_c = self._chips_re.search(line) + if m_c: + frame["chips"][m_c.group(1)] = m_c.group(2) + + m_r = self._runtime_re.search(line) + if m_r: + frame["runtime"][m_r.group(1)] = (m_r.group(2), m_r.group(3)) + + m_t = self._tensor_re.search(line) + if m_t: + frame["tensor"][m_t.group(1)] = m_t.group(2) + + m_l = self._latency_re.search(line) + if m_l: + frame["latency"][m_l.group(1)] = ( + m_l.group(2), + m_l.group(3), + m_l.group(4), + m_l.group(5), + ) + + def parse_log(self, log_content: str): + """ + Parse raw log text into structured data frames. + """ + self.frames = [] + lines = log_content.splitlines() + current_frame = self._init_empty_frame() + + for line in lines: + fs_match = self._frame_start_re.search(line) + if fs_match: + if current_frame["ts"]: + self.frames.append(current_frame) + current_frame = self._init_empty_frame() + current_frame["ts"] = datetime.strptime( + fs_match.group(1), "%H:%M:%S.%f" + ) + continue + self._extract_metrics(line, current_frame) + + if current_frame["ts"]: + self.frames.append(current_frame) + + def filter_update_events(self): + """ + Compare consecutive frames and retain only those where hardware data changed. + """ + self.update_events = [] + last_snapshot = None + for f in self.frames: + # Create a snapshot excluding the timestamp for comparison + snapshot = {k: v for k, v in f.items() if k != "ts"} + if last_snapshot is None or snapshot != last_snapshot: + self.update_events.append(f) + last_snapshot = snapshot + + def validate_rate_match(self) -> bool: + """ + Validates if the hardware update frequency aligns with the target rate. + + Logic Rationale: + 1. Eager Hardware Output: To ensure monitoring data does not lag behind the + specified sampling frequency (Target Rate), the hardware driver implements + an eager refresh strategy. This often results in intervals slightly below + or exactly at the target (e.g., 0.08s - 0.10s). + 2. Jitter Tolerance: A +/- 20% buffer (0.08s to 0.12s) is established to + account for system scheduling jitters and network transmission latency. + 3. Sensitivity Verification: A 'True' result confirms the system successfully + captured active hardware state changes at a high frequency, rather than + stale cached data. + """ + if len(self.update_events) < 3: + return False + + for i in range(2, len(self.update_events)): + delta = ( + self.update_events[i]["ts"] - self.update_events[i - 1]["ts"] + ).total_seconds() + if self.lower_bound <= delta <= self.upper_bound: + return True + return False + + def _format_row(self, no, ev, delta): + """Helper to format a single row of report data.""" + pids = ", ".join(ev["chips"].values()) + hbm_str = " | ".join( + [ + f"D{k}:{v[0].split('/')[0].strip()}({v[1]}%)" + for k, v in ev["runtime"].items() + ] + ) + tc_str = ", ".join([f"C{k}:{v}%" for k, v in ev["tensor"].items()]) + lat_str = " || ".join( + [f"{k}: {'|'.join(v)}" for k, v in ev["latency"].items()] + ) + + return ( + f"{no:<3} | {ev['ts'].strftime('%H:%M:%S.%f')[:-3]:<12} | {delta:<7.3f}s | " + f"{pids:<12} | {hbm_str:<60} | {tc_str:<30} | {lat_str}" + ) + + def generate_report(self) -> str: + """ + Generate a complete aligned text report of detected performance updates. + """ + num_events = len(self.update_events) + if num_events < 3: + return "Insufficient data to generate a report (at least 3 unique events required)." + + output = [] + header = ( + f"{'No':<3} | {'Timestamp':<12} | {'Intv (s)':<10} | " + f"{'PIDs':<12} | {'HBM Usage (Device:Used | Duty%)':<60} | " + f"{'TensorCore':<30} | {'Latency Profile (us)'}" + ) + output.append(header) + output.append("-" * 210) + + intervals = [] + for i in range(2, num_events): + ev = self.update_events[i] + prev_ev = self.update_events[i - 1] + delta = (ev["ts"] - prev_ev["ts"]).total_seconds() + intervals.append(delta) + + row = self._format_row(i - 1, ev, delta) + output.append(row) + + if intervals: + avg_intv = sum(intervals) / len(intervals) + is_matched = self.validate_rate_match() + output.append("-" * 210) + output.append( + f"Average Interval (Stable Phase): {avg_intv:.3f} s | Target Rate Match: {is_matched}" + ) + + return "\n".join(output) + + def execute_tpu_info_cli_command(info, pod_name: str, tpu_args: str) -> str: """Helper to handle KUBECONFIG and execute kubectl exec.""" with tempfile.NamedTemporaryFile() as temp_config_file: @@ -80,6 +254,58 @@ def validate_streaming_rate(info, pod_name: str, rate: float) -> str: return f"Validated {pod_name} at {rate}s" +@task +def validate_streaming_rate_iterations( + info, pod_name: str, rate: float, iteration_count: int = 40 +) -> str: + """ + Performs 40 iterations of 30s tests. + The task succeeds if at least 50% of the iterations are valid. + """ + duration = 30 + analyzer = TPUPerformanceAnalyzer(target_rate=rate) + success_count = 0 + pass_threshold = iteration_count / 2 # 50% threshold + + for i in range(1, iteration_count + 1): + # Precise command with Perl microsecond timestamping and terminal line export + tpu_args = ( + f'sh -c "export LINES=50 && ' + f"script -q -c 'timeout {duration}s tpu-info --streaming --rate {rate}' /dev/null\" " + f"| perl -MTime::HiRes=gettimeofday -ne ' " + f"($s, $usec) = gettimeofday; " + f"($sec,$min,$hour) = localtime($s); " + f'printf("[%02d:%02d:%02d.%03d] %s", $hour, $min, $sec, $usec/1000, $_);\' ' + f"|| [ ${{PIPESTATUS[0]}} -eq 124 ]" + ) + + try: + output = execute_tpu_info_cli_command(info, pod_name, tpu_args) + analyzer.parse_log(output) + analyzer.filter_update_events() + logging.info(analyzer.generate_report()) + if analyzer.validate_rate_match(): + success_count += 1 + else: + logging.info( + f"Iteration {i}: Failed validation (Intervals out of range)." + ) + except Exception as e: + logging.error(f"Iteration {i}: Command execution error: {str(e)}") + + # Evaluation logic: Pass if success_count >= 20 + status_msg = f"Pod {pod_name} at {rate}s: {success_count}/{iteration_count} iterations passed." + logging.info(status_msg) + + if success_count < pass_threshold: + raise AssertionError( + f"Validation Failed: Only {success_count}/{iteration_count} passed. " + f"Required at least {pass_threshold}." + ) + + return status_msg + + # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). with models.DAG( # pylint: disable=unexpected-keyword-arg @@ -189,19 +415,25 @@ def generate_second_node_pool_name( )(cluster_info, pod_name_list=pod_names, job_apply_time=apply_time) test_rates = [0.1, 0.5, 1.0, 5.0] + rate_test_groups = [] for rate in test_rates: formatted_rate = str(rate).replace(".", "_") - # Keyword arguments are generated dynamically at runtime (pylint does not - # know this signature). - with TaskGroup( # pylint: disable=unexpected-keyword-arg + with TaskGroup( group_id=f"verification_group_rate_{formatted_rate}" ) as rate_group: - streaming_validation_results = ( - validate_streaming_rate.override(task_id="streaming_rate_test") - .partial(info=cluster_info, rate=rate) - .expand(pod_name=pod_names) + # Fix: Ensure parameters match the task signature exactly + validate_streaming_rate_iterations.override( + task_id="streaming_rate_test", + execution_timeout=datetime.timedelta(minutes=60), + ).partial( + info=cluster_info, + rate=rate, + iteration_count=40, # Pass integer directly to partial + ).expand( + pod_name=pod_names ) + rate_test_groups.append(rate_group) cleanup_workload = jobset.end_workload.override( task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE @@ -242,7 +474,7 @@ def generate_second_node_pool_name( >> apply_time >> pod_names >> wait_for_job_start - >> verification_group + >> rate_test_groups >> cleanup_workload >> cleanup_node_pool ) From ace1cbcaecbd142e51ff8b65e7aef8c290c14f69 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Thu, 15 Jan 2026 16:38:37 +0800 Subject: [PATCH 10/14] feat: Refactor TPUPerformanceAnalyzer to enhance streaming rate verification with dynamic task groups --- .../tpu_info_streaming_rate.py | 89 ++++++++----------- 1 file changed, 38 insertions(+), 51 deletions(-) diff --git a/dags/tpu_observability/tpu_info_streaming_rate.py b/dags/tpu_observability/tpu_info_streaming_rate.py index eaf9b6804..b73dd4007 100644 --- a/dags/tpu_observability/tpu_info_streaming_rate.py +++ b/dags/tpu_observability/tpu_info_streaming_rate.py @@ -21,7 +21,6 @@ import re import subprocess import tempfile -from typing import List import logging from airflow import models @@ -50,7 +49,6 @@ def __init__(self, target_rate: float = 0.1): self.frames = [] self.update_events = [] - # Pre-compile regex patterns for optimized performance self._frame_start_re = re.compile(r"\[(\d{2}:\d{2}:\d{2}\.\d{3})\].*?\[H") self._chips_re = re.compile( r"│\s+(/dev/vfio/\d+)\s+│.*?│\s+\d+\s+│\s+(\d+)\s+│" @@ -104,7 +102,7 @@ def parse_log(self, log_content: str): if current_frame["ts"]: self.frames.append(current_frame) current_frame = self._init_empty_frame() - current_frame["ts"] = datetime.strptime( + current_frame["ts"] = datetime.datetime.strptime( fs_match.group(1), "%H:%M:%S.%f" ) continue @@ -120,7 +118,6 @@ def filter_update_events(self): self.update_events = [] last_snapshot = None for f in self.frames: - # Create a snapshot excluding the timestamp for comparison snapshot = {k: v for k, v in f.items() if k != "ts"} if last_snapshot is None or snapshot != last_snapshot: self.update_events.append(f) @@ -223,7 +220,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 ): """Verifies that expected strings exist in the output.""" for pattern in patterns: @@ -233,36 +230,18 @@ def verify_output_contains_patterns( ) -@task -def validate_streaming_rate(info, pod_name: str, rate: float) -> str: - """ - Executes tpu-info --streaming and validates frequency using UI and Data metrics. - """ - duration = 15 - - tpu_args = ( - f"sh -c \"script -q -c 'timeout {duration}s " - f"tpu-info --streaming --rate {rate}' /dev/null\" " - f"|| [ $? -eq 124 ]" - ) - output = execute_tpu_info_cli_command(info, pod_name, tpu_args) - - patterns = ["Refresh rate:", f"{rate}s"] - verify_output_contains_patterns( - output, patterns, f"Content check on {pod_name}" - ) - return f"Validated {pod_name} at {rate}s" - - @task def validate_streaming_rate_iterations( - info, pod_name: str, rate: float, iteration_count: int = 40 + info, + pod_name: str, + rate: float, + iteration_count: int = 40, + duration: int = 30, ) -> str: """ Performs 40 iterations of 30s tests. The task succeeds if at least 50% of the iterations are valid. """ - duration = 30 analyzer = TPUPerformanceAnalyzer(target_rate=rate) success_count = 0 pass_threshold = iteration_count / 2 # 50% threshold @@ -286,6 +265,7 @@ def validate_streaming_rate_iterations( logging.info(analyzer.generate_report()) if analyzer.validate_rate_match(): success_count += 1 + logging.info(f"Iteration {i}: Passed validation.") else: logging.info( f"Iteration {i}: Failed validation (Intervals out of range)." @@ -293,7 +273,6 @@ def validate_streaming_rate_iterations( except Exception as e: logging.error(f"Iteration {i}: Command execution error: {str(e)}") - # Evaluation logic: Pass if success_count >= 20 status_msg = f"Pod {pod_name} at {rate}s: {success_count}/{iteration_count} iterations passed." logging.info(status_msg) @@ -376,7 +355,9 @@ def generate_second_node_pool_name( tpu_topology=config.tpu_topology, ) - cluster_info_2 = node_pool.copy_node_pool_info_with_override( + cluster_info_2 = node_pool.copy_node_pool_info_with_override.override( + task_id="copy_node_pool_info_with_override" + )( info=cluster_info, node_pool_name=generate_second_node_pool_name(cluster_info), ) @@ -414,26 +395,32 @@ def generate_second_node_pool_name( task_id="wait_for_job_start" )(cluster_info, pod_name_list=pod_names, job_apply_time=apply_time) - test_rates = [0.1, 0.5, 1.0, 5.0] - rate_test_groups = [] - for rate in test_rates: - formatted_rate = str(rate).replace(".", "_") - - with TaskGroup( - group_id=f"verification_group_rate_{formatted_rate}" - ) as rate_group: - # Fix: Ensure parameters match the task signature exactly - validate_streaming_rate_iterations.override( - task_id="streaming_rate_test", - execution_timeout=datetime.timedelta(minutes=60), - ).partial( - info=cluster_info, - rate=rate, - iteration_count=40, # Pass integer directly to partial - ).expand( - pod_name=pod_names - ) - rate_test_groups.append(rate_group) + # Keyword arguments are generated dynamically at runtime (pylint does not + # know this signature). + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id="tpu_streaming_rate_verification" + ) as rate_verification_group: + test_rates = [0.1, 0.5, 1.0, 5.0] + + for rate in test_rates: + formatted_rate = str(rate).replace(".", "_") + + # Keyword arguments are generated dynamically at runtime (pylint does not + # know this signature). + with TaskGroup( # pylint: disable=unexpected-keyword-arg + group_id=f"rate_{formatted_rate}" + ) as rate_iterations_group: + validate_streaming_rate_iterations.override( + task_id="streaming_rate_test", + execution_timeout=datetime.timedelta(minutes=60), + duration=30, # Each iteration runs for 30 seconds + ).partial( + info=cluster_info, + rate=rate, + iteration_count=40, + ).expand( + pod_name=pod_names + ) cleanup_workload = jobset.end_workload.override( task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE @@ -474,7 +461,7 @@ def generate_second_node_pool_name( >> apply_time >> pod_names >> wait_for_job_start - >> rate_test_groups + >> rate_verification_group >> cleanup_workload >> cleanup_node_pool ) From c5d79815a9bbfa7d1ec2d773a147b5c7f0bb1e57 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Fri, 16 Jan 2026 11:37:02 +0800 Subject: [PATCH 11/14] feat: Update tpu-info streaming rate validation to enhance clarity and accuracy in documentation and logic --- .../tpu_info_streaming_rate.py | 38 ++++++++++++------- 1 file changed, 25 insertions(+), 13 deletions(-) diff --git a/dags/tpu_observability/tpu_info_streaming_rate.py b/dags/tpu_observability/tpu_info_streaming_rate.py index b73dd4007..5f5f7cf95 100644 --- a/dags/tpu_observability/tpu_info_streaming_rate.py +++ b/dags/tpu_observability/tpu_info_streaming_rate.py @@ -13,7 +13,8 @@ # limitations under the License. """ -DAG to verify tpu-info streaming rate functionality on TPU v6e slices. +DAG to validate the 'tpu-info' streaming refresh rate by calculating the time interval +between consecutive hardware telemetry updates on TPU v6e slices. """ import datetime @@ -40,7 +41,7 @@ class TPUPerformanceAnalyzer: def __init__(self, target_rate: float = 0.1): """ - Initialize the analyzer with a target sampling rate. + Initialize the analyzer with a target refresh rate. :param target_rate: The expected update interval in seconds (default 0.1s). """ self.target_rate = target_rate @@ -125,11 +126,11 @@ def filter_update_events(self): def validate_rate_match(self) -> bool: """ - Validates if the hardware update frequency aligns with the target rate. + Validates if the hardware update frequency aligns with the target refresh rate. Logic Rationale: 1. Eager Hardware Output: To ensure monitoring data does not lag behind the - specified sampling frequency (Target Rate), the hardware driver implements + specified refresh frequency (Target Refresh Rate e.g., 0.1s), the hardware driver implements an eager refresh strategy. This often results in intervals slightly below or exactly at the target (e.g., 0.08s - 0.10s). 2. Jitter Tolerance: A +/- 20% buffer (0.08s to 0.12s) is established to @@ -200,7 +201,9 @@ def generate_report(self) -> str: is_matched = self.validate_rate_match() output.append("-" * 210) output.append( - f"Average Interval (Stable Phase): {avg_intv:.3f} s | Target Rate Match: {is_matched}" + f"Average Interval (Stable Phase): {avg_intv:.3f} s\n" + f"Target Rate Match: {is_matched}\n" + f"(Verified: Streaming data updated within the target refresh rate window)" ) return "\n".join(output) @@ -237,6 +240,7 @@ def validate_streaming_rate_iterations( rate: float, iteration_count: int = 40, duration: int = 30, + pass_threshold_percent: float = 0.5, ) -> str: """ Performs 40 iterations of 30s tests. @@ -244,7 +248,7 @@ def validate_streaming_rate_iterations( """ analyzer = TPUPerformanceAnalyzer(target_rate=rate) success_count = 0 - pass_threshold = iteration_count / 2 # 50% threshold + pass_threshold = iteration_count * pass_threshold_percent for i in range(1, iteration_count + 1): # Precise command with Perl microsecond timestamping and terminal line export @@ -302,15 +306,20 @@ def validate_streaming_rate_iterations( "streaming-rate", ], description=( - "Automated validation of the tpu-info CLI's --rate parameter, " - "ensuring accurate metric streaming frequencies on TPU v6e-16 slices." + "Validates the tpu-info refresh rate by " + "calculating the time interval " + "between consecutive hardware telemetry updates on TPU v6e-16 slices." ), doc_md=""" - ## TPU Info Streaming Rate Verification DAG + ## TPU Info Streaming Refresh Rate Verification DAG - This DAG automates the functional testing of the `tpu-info` CLI tool, specifically focusing on the - `--streaming` and `--rate` flags. It verifies that the tool correctly honors requested update - intervals within a Kubernetes-managed TPU environment. + This DAG automates the functional testing of the `tpu-info` CLI tool, specifically + validating the accuracy of the streaming **refresh rate**. + + The core verification logic calculates the precise time delta between + consecutive hardware data frames. It ensures that the actual refresh frequency + matches the requested `--rate` parameter within a defined jitter tolerance, + confirming that the system delivers real-time hardware metrics without stale data. """, ) as dag: for machine in MachineConfigMap: @@ -362,6 +371,8 @@ def generate_second_node_pool_name( node_pool_name=generate_second_node_pool_name(cluster_info), ) + # Keyword arguments are generated dynamically at runtime (pylint does not + # know this signature). with TaskGroup( # pylint: disable=unexpected-keyword-arg group_id="create_node_pool" ) as create_node_pool: @@ -398,7 +409,7 @@ def generate_second_node_pool_name( # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). with TaskGroup( # pylint: disable=unexpected-keyword-arg - group_id="tpu_streaming_rate_verification" + group_id="streaming_rate_verification" ) as rate_verification_group: test_rates = [0.1, 0.5, 1.0, 5.0] @@ -418,6 +429,7 @@ def generate_second_node_pool_name( info=cluster_info, rate=rate, iteration_count=40, + pass_threshold_percent=0.5, # At least 50% must pass ).expand( pod_name=pod_names ) From e34d2eb18bc17bf16833915c117c3988f31f2900 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Tue, 20 Jan 2026 13:05:16 +0800 Subject: [PATCH 12/14] feat: Update duration parameter in streaming rate validation to ensure consistent iteration timing --- dags/tpu_observability/tpu_info_streaming_rate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dags/tpu_observability/tpu_info_streaming_rate.py b/dags/tpu_observability/tpu_info_streaming_rate.py index 5f5f7cf95..ddf1ce6e4 100644 --- a/dags/tpu_observability/tpu_info_streaming_rate.py +++ b/dags/tpu_observability/tpu_info_streaming_rate.py @@ -424,11 +424,11 @@ def generate_second_node_pool_name( validate_streaming_rate_iterations.override( task_id="streaming_rate_test", execution_timeout=datetime.timedelta(minutes=60), - duration=30, # Each iteration runs for 30 seconds ).partial( info=cluster_info, rate=rate, iteration_count=40, + duration=30, # Each iteration runs for 30 seconds pass_threshold_percent=0.5, # At least 50% must pass ).expand( pod_name=pod_names From eab54306600482d3f0eb6a5eef5df451789879d4 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Tue, 27 Jan 2026 15:56:47 +0800 Subject: [PATCH 13/14] feat: Refactor task chaining in TPUPerformanceAnalyzer to use chain for improved readability --- .../tpu_info_streaming_rate.py | 21 +++++++++---------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/dags/tpu_observability/tpu_info_streaming_rate.py b/dags/tpu_observability/tpu_info_streaming_rate.py index ddf1ce6e4..0d438db42 100644 --- a/dags/tpu_observability/tpu_info_streaming_rate.py +++ b/dags/tpu_observability/tpu_info_streaming_rate.py @@ -26,6 +26,7 @@ from airflow import models from airflow.decorators import task +from airflow.models.baseoperator import chain from airflow.utils.trigger_rule import TriggerRule from airflow.utils.task_group import TaskGroup @@ -465,16 +466,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 - >> rate_verification_group - >> cleanup_workload - >> cleanup_node_pool + chain( + cluster_info, + create_node_pool, + apply_time, + pod_names, + wait_for_job_start, + rate_verification_group, + cleanup_workload, + cleanup_node_pool, ) # pylint: enable=pointless-statement From b7384fe69e646804c99375d2c76757e837c06033 Mon Sep 17 00:00:00 2001 From: Tony Lin Date: Tue, 3 Feb 2026 13:39:22 +0800 Subject: [PATCH 14/14] feat: Update DAG schedule for streaming rate validation based on environment --- dags/tpu_observability/tpu_info_streaming_rate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dags/tpu_observability/tpu_info_streaming_rate.py b/dags/tpu_observability/tpu_info_streaming_rate.py index 0d438db42..d3dfd2082 100644 --- a/dags/tpu_observability/tpu_info_streaming_rate.py +++ b/dags/tpu_observability/tpu_info_streaming_rate.py @@ -295,7 +295,7 @@ def validate_streaming_rate_iterations( with models.DAG( # pylint: disable=unexpected-keyword-arg dag_id="tpu_info_verify_streaming_rate_dags", start_date=datetime.datetime(2025, 8, 10), - schedule=None, + schedule="0 14 * * *" if composer_env.is_prod_env() else None, catchup=False, tags=[ "cloud-ml-auto-solutions",