Add TPU Observability validation for help, version, and process monitoring

.github/workflows/dag-check.yml

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

.github/workflows/pyink-check.yml

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

.github/workflows/pylint-check.yml

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

.github/workflows/require-checklist.yml

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

.github/workflows/unit-test.yml

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

dags/tpu_observability/tpu_info_format_validation_dags.py

@@ -13,9 +13,18 @@
 # limitations under the License.
 
 """
-A DAG orchestrates the process of verifying TensorCore utilization metrics.
-
-This is done by comparing data from Cloud Logging and Cloud Monitoring.
+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
@@ -52,34 +61,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]):
   """
@@ -292,6 +273,27 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]):
     )
 
 
+@task
+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",
+      ],
+  }
+  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
 # know this signature).
 with models.DAG(  # pylint: disable=unexpected-keyword-arg
@@ -303,37 +305,31 @@ def validate_latency_table(tpu_info_output: list[tpu_info.Table]):
     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:
@@ -421,9 +417,11 @@ def generate_second_node_pool_name(
           job_apply_time=apply_time,
       )
 
-      outputs_of_tpu_info = (
-          get_tpu_info_from_pod.override(task_id="get_tpu_info")
-          .partial(info=cluster_info)
+      raw_metric_data = (
+          tpu_info.get_tpu_info_from_pod.override(task_id="get_tpu_info")
+          .partial(
+              info=cluster_info, cmd_str=tpu_info.TpuInfoCmd.TPU_INFO.value
+          )
           .expand(pod_name=running_pods)
       )
 
@@ -432,7 +430,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=running_pods,
+              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
@@ -472,6 +483,10 @@ def generate_second_node_pool_name(
             .expand(tpu_info_output=output_of_tpu_info)
         )
 
+        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
       )(
@@ -509,10 +524,11 @@ def generate_second_node_pool_name(
               validate_runtime_metric,
               validate_tensorcore_metric,
               validate_latency_metric,
+              cli_validation,
           ],
       )
 
-      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)
 
@@ -525,8 +541,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,

dags/tpu_observability/utils/tpu_info_util.py

@@ -1,11 +1,25 @@
 """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."""
+
+  TPU_INFO = "tpu-info"
+  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]
@@ -23,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
       ┏━━━━━━━━━━━━━┳━━━━━━━━━━━━━━┳━━━━━━━━━┳━━━━━━┓
@@ -96,6 +111,23 @@ 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
+) -> 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 {"output": output, "cmd_name": cmd_str}
+
+
 if __name__ == "__main__":
   full_output = """
                                                                                                                                                    Libtpu version: 0.0.23

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."