Refactor TPU Observability DAGs to Remove Dynamic Task Mapping

[yuna-tzeng]

Description

This PR refactors the TPU observability DAGs to remove Airflow dynamic task mapping (the .expand() pattern). Instead, validations and operations across multiple pods are now executed sequentially inside a single Airflow task. This change simplifies the DAG structure, improves task-level logging, and streamlines error handling.

Removal of Dynamic Task Mapping (.expand())

Converted tasks that previously targeted a single pod using dynamic task mapping to accept a list of pod names (list[str]) and process them sequentially.
Removed complex nested TaskGroup structures and the usage of .zip() and chain(), resulting in cleaner and more maintainable DAG graphs.
Removed the @task decorator from parsing and helper validation functions (such as parse_tpu_info_output, validate_chips_table, validate_runtime_table, validate_tensorcore_table, and validate_latency_table), converting them back to standard Python functions to reduce Airflow task instantiation overhead.

Modified Files & Implementation Details

1. dags/tpu_observability/jobset_ttr_kill_process.py

   • Refactored kill_tpu_pod_workload to kill_tpu_pod_workloads, taking pod_names: list[str].
   • Replaced .partial(...).expand(pod_name=startup.running_pods) in the DAG definition with a direct single-task invocation passing the list of running pods.

2. dags/tpu_observability/tpu_info_format_validation_dags.py

   • Removed @task decorators from validate_chips_table, validate_runtime_table, validate_tensorcore_table, and validate_latency_table.
   • Introduced a new validate_tpu_info_format task to coordinate the tpu-info execution, parsing, and format checks for all pods.
   • Removed the verification_group TaskGroup and its dynamic mapping steps, replacing them with a single task.

3. dags/tpu_observability/tpu_info_metrics_dag.py

   • Removed get_tpu_info_metric_from_pod and run_metric_verification tasks.
   • Created a new verify_metric_for_all_pods task to sequentially execute tpu-info --metric, parse outputs, query Cloud Monitoring, and run validation strategies for all pods.
   • Removed the TaskGroup and .zip mapping from the main DAG workflow.

4. dags/tpu_observability/tpu_sdk_monitoring_validation_dag.py

   • Refactored validate_monitoring_sdk to accept pod_names: list[str].
   • Removed .expand(pod_name=startup.running_pods).
   • Added sequential pod validation loop and error aggregation logic.

5. dags/tpu_observability/utils/tpu_info_util.py

   • Removed @task decorator from parse_tpu_info_output to keep it a standard utility function.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run one-shot tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed.

[alfredyu-cienet]

FYI, Kim's change has been merged, so let's rebase this PR and resolve conflicts

[alfredyu-cienet]

maybe we can apply TaskGroupWithTimeout for those DAGs in this PR as well?

[alfredyu-cienet]

Suggested change

	all_verification_tasks,
	*all_verification_tasks,