Feat: add jobset_ttr_pod_oom DAG for v6e recovery validation

[chiajunglien]

Description

The primary goal of this PR is to implement the jobset_ttr_pod_oom DAG, designed to validate the self-healing capability of TPU JobSets under memory pressure. Unlike simple pod deletion, this DAG automates OOM (Out-Of-Memory) fault injection to ensure JAX workloads and the JobSet controller correctly handle and recover from container-level failures.

Technical Implementation

• Dynamic Provisioning: Builds node pool info from GCS configurations and provisions TPU resources based on MachineConfigMap (e.g., v6e-16).
• Resource Constraints: Dynamically injects memory limits (e.g., 4Gi) into the JobSet YAML to facilitate rapid OOM triggering.
• Fault Injection: Implements @task trigger_oom_killed which executes oomkill.py inside a random pod via kubectl exec.
• Detection Phase: Monitors the JobSet's reaction. It expects the JobSet to detect the missing pod and trigger a full-slice restart (Killing event) to maintain JAX synchronization.
• Performance Metric: Measures the Time To Recovery (TTR).

Airflow/Composer

• GCP Composer name: emmalien-test (under GCP project: cienet-cmcs)
• GCP Composer version: 2.13.1
• GCP Airflow version: 2.10.5
• Test results: https://35b13c750f364c56b55b9367bb681dee-dot-us-central1.composer.googleusercontent.com/dags/jobset_ttr_pod_oom/grid

Required Variables

• Cluster Information (This DAG requires an existing cluster)

  • PROJECT_ID: The GCP project ID where the cluster resides. (Default to cienet-cmcs)
  • CLUSTER_NAME: The name of the target GKE cluster. (Default to tpu-observability-automation-dev)
  • LOCATION: The region of the GKE cluster. (Default to us-central1)

• Node Pool Configurations

  • NODE_POOL_NAME: The base name for the new node pool. (Default to jobset-ttr-pod-oom-v6e)
  • NODE_LOCATIONS: The zone for the nodes in the normal test path. (Default to us-central1-b)
  • NUM_NODES: The number of nodes to create in the pool. (Default to 4)
  • MACHINE_TYPE: The machine type for the GKE nodes. (Default to ct6e-standard-4t)
  • TPU_TOPOLOGY: The TPU topology for the node pool. (Default to 4x4)

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]

what's the content of this script, can we specify the raw commands from DAG directly?

[chiajunglien]

Done.