feat: add OOM fault injection for JobSet TTR validation

chiajunglien · chiajunglien · commit eef2c1c1ecf9 · 2026-05-11T08:10:23.000Z
diff --git a/dags/tpu_observability/jobset_ttr_pod_oom.py b/dags/tpu_observability/jobset_ttr_pod_oom.py
@@ -19,6 +19,7 @@
 import tempfile
 import os
 import random
+import time
 from typing import List
 
 from airflow import models
@@ -31,38 +32,75 @@
 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 JobSet
-from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH
+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):
+def trigger_oom_failure(info, pod_name: str, namespace: str):
   """
-  Triggers an OOM (Out-of-Memory) event on a specified TPU pod.
+  Injects an OOM fault by running a memory-intensive loop inside a pod.
 
-  This task connects to the specified pod via kubectl exec and runs a
-  memory-intensive Python script (oomkill.py). The script is designed
-  to exhaust the container's memory limit, forcing a SIGKILL (Exit Code 137).
-  A connection error is expected and caught when the pod is killed.
+  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: Node pool and cluster information.
-    pod_name (str): The name of the target pod to be OOMKilled.
+    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.
   """
+
+  wait_time = 60
+  print(
+      f"Waiting {wait_time}s for Pod {pod_name} to stabilize before OOM Test..."
+  )
+  time.sleep(wait_time)
+
+  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
 
-    cmd = " && ".join([
-        jobset.Command.get_credentials_command(info),
-        f"kubectl exec {pod_name} -n default -- python3 oomkill.py",
-    ])
+    credentials_cmd = jobset.Command.get_credentials_command(info)
+    exec_cmd = (
+        f'kubectl exec {pod_name} -n {namespace} -- python3 -c "{python_logic}"'
+    )
+
+    full_command = f"{credentials_cmd} && {exec_cmd}"
 
     try:
-      print(f"Executing OOM script in {pod_name}...")
-      subprocess.run_exec(cmd, env=env)
-    except RuntimeError as e:
-      print(f"Expectation: Connection closed due to OOMKilled. Info: {e}")
+      print(f"Blasting {pod_name} memory now...")
+      subprocess.run_exec(full_command, env=env)
+    except Exception as e:
+      if "137" in str(e):
+        print(f"Success: Pod {pod_name} was OOMKilled (Exit Code 137).")
+      else:
+        print(f"Connection lost (Possible OOM): {e}")
 
 
 @task
@@ -88,9 +126,10 @@ def pick_random_pod(active_pods: List[str]) -> str:
 
 
 with models.DAG(
-    dag_id="jobset_ttr_pod_oom",
+    dag_id=DAG_ID,
     start_date=datetime.datetime(2026, 1, 15),
-    schedule="0 18 * * *" if composer_env.is_prod_env() else None,
+    schedule=SCHEDULE if composer_env.is_prod_env() else None,
+    dagrun_timeout=DAGRUN_TIMEOUT,
     catchup=False,
     tags=[
         "cloud-ml-auto-solutions",
@@ -130,103 +169,82 @@ def pick_random_pod(active_pods: List[str]) -> str:
   for machine in MachineConfigMap:
     config = machine.value
 
-    jobset_config = JobSet(
-        jobset_name="jobset-ttr-pod-oom-v6e-workload",
-        namespace="default",
-        max_restarts=5,
-        replicated_job_name="tpu-job-slice",
-        replicas=1,
-        backoff_limit=0,
-        completions=4,
-        parallelism=4,
-        tpu_accelerator_type="tpu-v6e-slice",
-        tpu_topology="4x4",
-        container_name="jax-tpu-worker",
-        image="us-docker.pkg.dev/cienet-cmcs/emma-tpu-test-repo/oom-test:v1",
-        tpu_cores_per_pod=4,
-    )
-    raw_yaml = jobset_config.generate_yaml(workload_script="sleep infinity")
-    custom_yaml = raw_yaml.replace(
-        "google.com/tpu: 4",
-        'google.com/tpu: 4\n                  memory: "4Gi"',
-    )
-  # 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}"
-  ):
-    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="jobset_ttr_pod_oom",
-        is_prod=composer_env.is_prod_env(),
-        machine_type=config.machine_version.value,
-        tpu_topology=config.tpu_topology,
-    )
-
-    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,
-        yaml_config=custom_yaml,
-        namespace=jobset_config.namespace,
-    )
-
-    ensure_all_pods_running = jobset.wait_for_all_pods_running.override(
-        task_id="ensure_all_pods_running"
-    )(
-        num_pods=(jobset_config.replicas * jobset_config.parallelism),
-        node_pool=cluster_info,
-    )
-
-    found_pods = jobset.list_pod_names.override(task_id="list_pod_names")(
-        node_pool=cluster_info,
-        namespace=jobset_config.namespace,
-    )
-
-    select_random_pod = pick_random_pod.override(task_id="select_random_pod")(
-        active_pods=found_pods
-    )
-
-    trigger_oom_killed = trigger_oom_failure.override(
-        task_id="trigger_oom_killed"
-    )(info=cluster_info, pod_name=select_random_pod)
-
-    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_name=jobset_config.jobset_name,
-    )
-
-    cleanup_workload = jobset.end_workload.override(
-        task_id="cleanup_workload", trigger_rule=TriggerRule.ALL_DONE
-    )(
-        node_pool=cluster_info,
-        jobset_name=jobset_config.jobset_name,
-        namespace=jobset_config.namespace,
-    ).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,
-        found_pods,
-        select_random_pod,
-        trigger_oom_killed,
-        wait_for_metric_upload,
-        cleanup_workload,
-        cleanup_node_pool,
-    )
+    # 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,
+      )
