Refactor TPU Observability DAGs to Remove Dynamic Task Mapping

.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/jobset_ttr_kill_process.py

@@ -24,6 +24,7 @@
 
 from airflow import models
 from airflow.decorators import task
+from airflow.exceptions import AirflowFailException
 from airflow.models.baseoperator import chain
 from airflow.utils.trigger_rule import TriggerRule
 from airflow.utils.task_group import TaskGroup
@@ -48,29 +49,41 @@
 
 
 @task
-def kill_tpu_pod_workload(info: node_pool.Info, pod_name: str) -> None:
+def kill_tpu_pod_workloads(info: node_pool.Info, pod_names: list[str]) -> None:
   """
-  Kills the python process on a single pod.
+  Kills the python process on a list of pods.
 
   This task retrieves cluster credentials, then attempts to kill the JAX
-  python process inside the specified pod. It ignores errors if the pod
+  python process inside each specified pod. It ignores errors if the pod
   has already been deleted to ensure pipeline continuity.
   """
   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 -- pkill -9 -f python",
-    ])
-
-    try:
-      subprocess.run_exec(cmd, env=env)
-    except subprocess.ProcessKilledException:
-      logging.info("Process was terminated with SIGKILL")
-    except Exception as e:
-      raise e
+    failed_pods = []
+    for pod_name in pod_names:
+      cmd = " && ".join([
+          jobset.Command.get_credentials_command(info),
+          f"kubectl exec {pod_name} -n default -- pkill -9 -f python",
+      ])
+
+      try:
+        subprocess.run_exec(cmd, env=env)
+        logging.info("Execution succeeded for pod: %s", pod_name)
+      except subprocess.ProcessKilledException:
+        logging.info(f"Process in pod {pod_name} was terminated with SIGKILL")
+        logging.info("Execution failed for pod: %s", pod_name)
+      except Exception as e:
+        logging.error("Execution failed for pod: %s. Error: %s", pod_name, e)
+        failed_pods.append((pod_name, e))
+
+    if failed_pods:
+      failed_pod_names = [name for name, _ in failed_pods]
+      logging.error("The following pods failed execution: %s", failed_pod_names)
+      raise AirflowFailException(
+          f"Task failed because execution failed on pods: {failed_pod_names}"
+      )
 
 
 # Keyword arguments are generated dynamically at runtime (pylint does not
@@ -156,10 +169,11 @@ def kill_tpu_pod_workload(info: node_pool.Info, pod_name: str) -> None:
           workload_type=Workload.JAX_TPU_BENCHMARK,
       )
 
-      kill_tasks = (
-          kill_tpu_pod_workload.override(task_id="kill_tpu_pod_workload")
-          .partial(info=cluster_info)
-          .expand(pod_name=startup.running_pods)
+      kill_tasks = kill_tpu_pod_workloads.override(
+          task_id="kill_tpu_pod_workloads"
+      )(
+          info=cluster_info,
+          pod_names=startup.running_pods,
       )
 
       wait_for_metric_upload = jobset.wait_for_jobset_ttr_to_be_found.override(

dags/tpu_observability/tpu_info_format_validation_dags.py

@@ -23,6 +23,7 @@
 import re
 import subprocess
 import tempfile
+import logging
 
 from airflow import models
 from airflow.decorators import task
@@ -43,10 +44,10 @@
 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_info_util import parse_tpu_info_output
 from dags.tpu_observability.utils.jobset_util import Workload
 from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
 
-
 DAG_ID = "tpu_info_format_validation_dag"
 DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)
 SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID)
@@ -81,6 +82,52 @@ def get_tpu_info_from_pod(info: node_pool.Info, pod_name: str) -> str:
 
 
 @task
+def validate_tpu_info_format(
+    info: node_pool.Info,
+    tpu_config: TpuConfig,
+    pod_names: list[str],
+) -> None:
+  """Executes tpu-info command and runs all validations sequentially across all pods."""
+
+  failed_pods = []
+  for pod_name in pod_names:
+    logging.info(
+        "Executing tpu-info and performing format validation for pod: %s",
+        pod_name,
+    )
+
+    try:
+      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",
+        ])
+
+        raw_output = subprocess.run_exec(cmd, env=env)
+
+      tpu_info_output = parse_tpu_info_output(raw_output)
+
+      verify_table_amount(tpu_info_output)
+      validate_chips_table(tpu_info_output, tpu_config)
+      validate_runtime_table(tpu_info_output)
+      validate_tensorcore_table(tpu_info_output)
+      validate_latency_table(tpu_info_output)
+      logging.info("Validation succeeded for pod: %s", pod_name)
+    except Exception as e:
+      logging.error("Validation failed for pod: %s. Error: %s", pod_name, e)
+      failed_pods.append((pod_name, e))
+
+  if failed_pods:
+    failed_pod_names = [name for name, _ in failed_pods]
+    logging.error("The following pods failed validation: %s", failed_pod_names)
+    raise AirflowFailException(
+        f"Task failed because validation failed on pods: {failed_pod_names}"
+    )
+
+
 def verify_table_amount(tpu_info_output: list[tpu_info.Table]):
   """
   Verifies if all expected tables are present.
@@ -103,7 +150,6 @@ def verify_table_amount(tpu_info_output: list[tpu_info.Table]):
     )
 
 
-@task
 def validate_chips_table(
     tpu_info_output: list[tpu_info.Table],
     tpu_config: TpuConfig,
@@ -156,7 +202,6 @@ def validate_chips_table(
     )
 
 
-@task
 def validate_runtime_table(tpu_info_output: list[tpu_info.Table]):
   """
   Validates the row count and content of table 'TPU Runtime Utilization'
@@ -211,7 +256,6 @@ def validate_runtime_table(tpu_info_output: list[tpu_info.Table]):
     )
 
 
-@task
 def validate_tensorcore_table(tpu_info_output: list[tpu_info.Table]):
   """
   Validates the row count and content of table 'TensorCore Utilization'
@@ -251,7 +295,6 @@ def validate_tensorcore_table(tpu_info_output: list[tpu_info.Table]):
     )
 
 
-@task
 def validate_latency_table(tpu_info_output: list[tpu_info.Table]):
   """
   Validates the row count and content of table 'TPU Buffer Transfer Latency'
@@ -407,57 +450,14 @@ def generate_second_node_pool_name(
           workload_type=Workload.JAX_TPU_BENCHMARK,
       )
 
-      outputs_of_tpu_info = (
-          get_tpu_info_from_pod.override(task_id="get_tpu_info")
-          .partial(info=cluster_info)
-          .expand(pod_name=startup.running_pods)
-      )
-
-      output_of_tpu_info = (
-          tpu_info.parse_tpu_info_output.override(
-              task_id="get_each_metric_table"
-          )
-          .partial()
-          .expand(output=outputs_of_tpu_info)
+      validate_format = validate_tpu_info_format.override(
+          task_id="validate_tpu_info_format"
+      )(
+          info=cluster_info,
+          tpu_config=config,
+          pod_names=startup.running_pods,
       )
 
-      # 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:
-        verify_table_amount_task = (
-            verify_table_amount.override(task_id="verify_table_amount_task")
-            .partial()
-            .expand(tpu_info_output=output_of_tpu_info)
-        )
-
-        validate_tpu_chips_metric = (
-            validate_chips_table.override(task_id="validate_tpu_chips_metric")
-            .partial(tpu_config=config)
-            .expand(tpu_info_output=output_of_tpu_info)
-        )
-
-        validate_runtime_metric = (
-            validate_runtime_table.override(task_id="validate_runtime_metric")
-            .partial()
-            .expand(tpu_info_output=output_of_tpu_info)
-        )
-
-        validate_tensorcore_metric = (
-            validate_tensorcore_table.override(
-                task_id="validate_tensorcore_metric"
-            )
-            .partial()
-            .expand(tpu_info_output=output_of_tpu_info)
-        )
-
-        validate_latency_metric = (
-            validate_latency_table.override(task_id="validate_latency_metric")
-            .partial()
-            .expand(tpu_info_output=output_of_tpu_info)
-        )
-
       clean_up_workload = jobset.end_workload.override(
           task_id="clean_up_workload", trigger_rule=TriggerRule.ALL_DONE
       )(
@@ -488,16 +488,6 @@ def generate_second_node_pool_name(
             setups=create_node_pool,
         )
 
-      chain(
-          verify_table_amount_task,
-          [
-              validate_tpu_chips_metric,
-              validate_runtime_metric,
-              validate_tensorcore_metric,
-              validate_latency_metric,
-          ],
-      )
-
       chain(create_first_node_pool, create_second_node_pool)
 
       chain(cleanup_first_node_pool, cleanup_second_node_pool)
@@ -509,9 +499,7 @@ def generate_second_node_pool_name(
           cluster_info_2,
           create_node_pool,
           *startup.tasks,
-          outputs_of_tpu_info,
-          output_of_tpu_info,
-          verification_group,
+          validate_format,
           clean_up_workload,
           cleanup_node_pool,
       )