Feat: add jobset_ttr_pod_oom DAG for v6e recovery validation

.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/common/scheduling_helper/scheduling_helper.py

@@ -64,6 +64,7 @@ class DayOfWeek(enum.Enum):
         "jobset_uptime_validation": dt.timedelta(minutes=90),
         "jobset_ttr_drain_restart": DefaultTimeout,
         "tpu_info_metrics_verification": DefaultTimeout,
+        "jobset_ttr_pod_oom": dt.timedelta(minutes=90),
     },
     TPU_INTERRUPTION_MOCK_CLUSTER.name: {
         "validate_interruption_count_gce_bare_metal_preemption": DefaultTimeout,

dags/tpu_observability/jobset_ttr_pod_oom.py

@@ -0,0 +1,252 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A DAG to validate JobSet Time-to-Recover (TTR) metrics
+by injecting Out-of-Memory (OOM) faults into TPU worker pods."""
+
+import datetime
+import tempfile
+import os
+import random
+import logging
+from typing import List
+
+from airflow import models
+from airflow.models.baseoperator import chain
+from airflow.utils.trigger_rule import TriggerRule
+from airflow.utils.task_group import TaskGroup
+from airflow.decorators import task
+
+from dags import composer_env
+from dags.tpu_observability.utils import jobset_util as jobset
+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.jobset_util import Workload
+from dags.tpu_observability.configs.common import (
+    MachineConfigMap,
+    GCS_CONFIG_PATH,
+    GCS_JOBSET_CONFIG_PATH,
+)
+from dags.common.scheduling_helper.scheduling_helper import SchedulingHelper, get_dag_timeout
+
+DAG_ID = "jobset_ttr_pod_oom"
+DAGRUN_TIMEOUT = get_dag_timeout(DAG_ID)
+SCHEDULE = SchedulingHelper.arrange_schedule_time(DAG_ID)
+
+
+@task
+def trigger_oom_failure(info, pod_name: str, namespace: str):
+  """
+  Injects an OOM fault by running a memory-intensive loop inside a pod.
+
+  This task waits for the target pod to stabilize, then executes a Python
+    one-liner via 'kubectl exec' that rapidly allocates memory. The script
+    uses bytearray allocation combined with an immediate write operation
+    (a[-1][0] = 1) to ensure physical RAM is committed by the OS, effectively
+    triggering a SIGKILL from the OOM Killer (Exit Code 137).
+
+  Args:
+    info: An object containing cluster and node pool credentials/metadata.
+    pod_name: The name of the target TPU worker pod to inject the fault.
+    namespace: The Kubernetes namespace where the pod is running.
+
+  Raises:
+    Exception: If the subprocess execution fails for reasons other than
+    the expected connection loss during an OOM event.
+  """
+
+  logging.info(f"Checking status for Pod {pod_name}...")
+
+  python_logic = (
+      "import time\n"
+      "a = []\n"
+      "print('Starting memory stress test...')\n"
+      "while True:\n"
+      "    a.append(bytearray(1024**3))\n"
+      "    a[-1][0] = 1\n"
+      "    time.sleep(0.01)"
+  )
+
+  with tempfile.NamedTemporaryFile() as temp_config_file:
+    env = os.environ.copy()
+    env["KUBECONFIG"] = temp_config_file.name
+
+    credentials_cmd = jobset.Command.get_credentials_command(info)
+
+    wait_cmd = (
+        f"kubectl wait --for=condition=Ready pod/{pod_name} "
+        f"-n {namespace} --timeout=60s"
+    )
+
+    exec_cmd = (
+        f'kubectl exec {pod_name} -n {namespace} -- python3 -c "{python_logic}"'
+    )
+
+    full_command = f"{credentials_cmd} && {wait_cmd} && {exec_cmd}"
+
+    try:
+      logging.info(f"Blasting {pod_name} memory now...")
+      subprocess.run_exec(full_command, env=env)
+    except subprocess.ProcessKilledException:
+      logging.info(f"Success: Pod {pod_name} was OOMKilled.")
+    except Exception as e:
+      logging.error(f"Connection lost or other error: {e}")
+      raise e
+
+
+@task
+def pick_random_pod(active_pods: List[str]) -> str:
+  """
+  Randomly selects one pod from a list of available JobSet pods.
+
+  This ensures that the fault injection is performed on a single
+  unit of the TPU slice, allowing the test to validate how the
+  JobSet controller handles partial failures within a replica.
+
+  Args:
+    pod_names (List[str]): List of active pod names in the JobSet.
+
+  Returns:
+    str: The name of the randomly selected pod.
+  """
+  if not active_pods:
+    raise ValueError("No pods found to attack!")
+  chosen_pod = random.choice(active_pods)
+  logging.info(f"Randomly selected pod for OOM test: {chosen_pod}")
+  return chosen_pod
+
+
+with models.DAG(
+    dag_id=DAG_ID,
+    start_date=datetime.datetime(2026, 1, 15),
+    schedule=SCHEDULE if composer_env.is_prod_env() else None,
+    dagrun_timeout=DAGRUN_TIMEOUT,
+    catchup=False,
+    tags=[
+        "cloud-ml-auto-solutions",
+        "jobset",
+        "tpu-observability",
+        "pod_oom",
+        "TPU",
+        "v6e-16",
+        "Validation",
+    ],
+    description=(
+        "This DAG tests the JobSet time-to-recover metric by injecting "
+        "an OOM event into a random pod to trigger a recovery, "
+        "then polls Cloud Monitoring to verify the metric is updated."
+    ),
+    doc_md="""
+        ### JobSet Time-To-Recover (TTR) Test Using Random Pod OOM Injection
+        ### Description
+        This DAG verifies that JobSet can recover from a single pod failure caused by
+        an Out-Of-Memory (OOM) event. It launches a JobSet, injects a memory-intensive
+        Python stressor into a running pod, and uses a sensor to confirm that the
+        JobSet controller triggers a recovery and reports the recovery duration (TTR).
+        ### Prerequisites
+        This test requires an existing cluster and the ability to execute commands
+        within the pod via `kubectl exec`.
+        ### Procedures
+        First, the node pool is created. A JobSet YAML is then launched on the cluster
+        and given time for all pods to reach a `Running` state. After stabilization,
+        a random pod is selected and an OOM event is triggered via `kubectl exec`
+        using a Python-based memory allocator that forces physical RAM commitment
+        until the process is terminated (expecting Exit Code 137). A sensor is
+        finally run which will poll Cloud Monitoring to detect that the JobSet
+        Time-To-Recover (TTR) metric has been updated, resulting in a success,
+        or timeout, and fail.
+      """,
+) as dag:
+  for machine in MachineConfigMap:
+    config = machine.value
+
+    # 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}"
+    ):
+      selector = jobset.generate_node_pool_selector("jobset-ttr-pod-oom")
+
+      jobset_config = jobset.build_jobset_from_gcs_yaml(
+          gcs_path=GCS_JOBSET_CONFIG_PATH,
+          dag_name=DAG_ID,
+          node_pool_selector=selector,
+      )
+
+      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=DAG_ID,
+          is_prod=composer_env.is_prod_env(),
+          machine_type=config.machine_version.value,
+          tpu_topology=config.tpu_topology,
+          node_pool_selector=selector,
+      )
+
+      create_node_pool = node_pool.create.override(task_id="create_node_pool")(
+          node_pool=cluster_info,
+      )
+
+      start_workload = jobset.run_workload.override(task_id="start_workload")(
+          node_pool=cluster_info,
+          jobset_config=jobset_config,
+          workload_type=Workload.JAX_TPU_BENCHMARK,
+      )
+
+      ensure_all_pods_running = jobset.wait_for_all_pods_running.override(
+          task_id="ensure_all_pods_running"
+      )(
+          node_pool=cluster_info,
+          jobset_config=jobset_config,
+      )
+
+      select_random_pod = pick_random_pod.override(task_id="select_random_pod")(
+          active_pods=ensure_all_pods_running
+      )
+
+      trigger_oom_killed = trigger_oom_failure.override(
+          task_id="trigger_oom_killed"
+      )(info=cluster_info, pod_name=select_random_pod, namespace="default")
+
+      wait_for_metric_upload = jobset.wait_for_jobset_ttr_to_be_found.override(
+          task_id="wait_for_jobset_ttr_to_be_found"
+      )(
+          node_pool=cluster_info,
+          jobset_config=jobset_config,
+      )
+
+      cleanup_workload = jobset.end_workload.override(
+          task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE
+      )(node_pool=cluster_info, jobset_config=jobset_config).as_teardown(
+          setups=start_workload
+      )
+
+      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,
+      )
+
+      chain(
+          cluster_info,
+          create_node_pool,
+          start_workload,
+          ensure_all_pods_running,
+          select_random_pod,
+          trigger_oom_killed,
+          wait_for_metric_upload,
+          cleanup_workload,
+          cleanup_node_pool,
+      )

scripts/code-style.sh

@@ -19,14 +19,37 @@ set -e
 
 FOLDERS_TO_FORMAT=("dags" "xlml")
 
-for folder in "${FOLDERS_TO_FORMAT[@]}"
-do
-  pyink "$folder" --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80 --check --diff
-done
-
-for folder in "${FOLDERS_TO_FORMAT[@]}"
-do
-  pylint "./$folder" --fail-under=9.6
-done
+HEAD_SHA="$(git rev-parse HEAD)"
+BASE_BRANCH="dev"
+
+if ! git rev-parse --verify "$BASE_BRANCH" >/dev/null 2>&1; then
+  git fetch origin "$BASE_BRANCH":"$BASE_BRANCH" || {
+    echo "[code-style] base branch '$BASE_BRANCH' not found, skip diff-based check."
+    exit 0
+  }
+fi
+
+CHANGED_PY_FILES="$(
+  git diff --name-only --diff-filter=ACM "${BASE_BRANCH}" "${HEAD_SHA}" \
+    | grep '\.py$' \
+    | while read -r f; do
+        for folder in "${FOLDERS_TO_FORMAT[@]}"; do
+          if [[ "$f" == "$folder/"* ]]; then
+            echo "$f"
+            break
+          fi
+        done
+      done \
+    | sort -u
+)"
+
+if [[ -z "${CHANGED_PY_FILES}" ]]; then
+  echo "[pre-push hook] no changed files detected between ${HEAD_SHA} and ${BASE_BRANCH}"
+  exit 1
+fi
+
+pyink ${CHANGED_PY_FILES} --pyink-indentation=2 --pyink-use-majority-quotes --line-length=80 --check --diff
+
+pylint ${CHANGED_PY_FILES} --fail-under=9.6 --disable=E1123
 
 echo "Successfully clean up all codes."