refactor: Consolidate TPU info CLI command execution and enhance validation handling

Author: chengpinglin

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}'."
         )
 
 

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,34 +111,36 @@ 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()
     env["KUBECONFIG"] = temp_config_file.name
 
     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