feat: Add validation DAG for tpu-info CLI tool functionality in TPU worker pods

Author: chengpinglin

dags/tpu_observability/tpu_info_format_validation_dags.py

@@ -13,17 +13,21 @@
 # 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
 import os
 import re
 import subprocess
 import tempfile
-from dataclasses import replace
+from typing import List
 
 from airflow import models
 from airflow.decorators import task
@@ -286,6 +290,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
@@ -525,3 +593,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