From badfd1391efa08c1a19a2de3f7d6848190e0223d Mon Sep 17 00:00:00 2001 From: Alfred Yu Date: Tue, 12 Aug 2025 10:03:26 +0800 Subject: [PATCH 01/15] [cienet-private] Enable workflow checks for the primary working branches --- .github/workflows/dag-check.yml | 1 - .github/workflows/pyink-check.yml | 3 +- .github/workflows/pylint-check.yml | 3 +- .github/workflows/require-checklist.yml | 5 ++- .github/workflows/unit-test.yml | 1 - scripts/code-style.sh | 41 +++++++++++++++++++------ 6 files changed, 40 insertions(+), 14 deletions(-) diff --git a/.github/workflows/dag-check.yml b/.github/workflows/dag-check.yml index 96fcb76bf..ac3982d0e 100644 --- a/.github/workflows/dag-check.yml +++ b/.github/workflows/dag-check.yml @@ -3,7 +3,6 @@ name: DAG Check on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: diff --git a/.github/workflows/pyink-check.yml b/.github/workflows/pyink-check.yml index ea7a8f250..c5cb1cd83 100644 --- a/.github/workflows/pyink-check.yml +++ b/.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 diff --git a/.github/workflows/pylint-check.yml b/.github/workflows/pylint-check.yml index 5e23d2812..02f4cb31e 100644 --- a/.github/workflows/pylint-check.yml +++ b/.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 diff --git a/.github/workflows/require-checklist.yml b/.github/workflows/require-checklist.yml index d15d19d99..4da288575 100644 --- a/.github/workflows/require-checklist.yml +++ b/.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 \ No newline at end of file + requireChecklist: false # If this is true and there are no checklists detected, the action will fail diff --git a/.github/workflows/unit-test.yml b/.github/workflows/unit-test.yml index 4f1723cb2..b771ca7e9 100644 --- a/.github/workflows/unit-test.yml +++ b/.github/workflows/unit-test.yml @@ -3,7 +3,6 @@ name: Unit Test on: pull_request: - branches: [master] types: [opened, synchronize, edited] push: diff --git a/scripts/code-style.sh b/scripts/code-style.sh index 36cfa13e9..ea50709ac 100755 --- a/scripts/code-style.sh +++ b/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." From f7ddca96bd5bbdb88cd9e1248f1597f2b1846f38 Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Thu, 15 Jan 2026 16:30:35 +0800 Subject: [PATCH 02/15] feat: add jobset_ttr_pod_oom DAG for v6e recovery validation --- dags/tpu_observability/jobset_ttr_pod_oom.py | 234 +++++++++++++++++++ 1 file changed, 234 insertions(+) create mode 100644 dags/tpu_observability/jobset_ttr_pod_oom.py diff --git a/dags/tpu_observability/jobset_ttr_pod_oom.py b/dags/tpu_observability/jobset_ttr_pod_oom.py new file mode 100644 index 000000000..3e2c6fd2b --- /dev/null +++ b/dags/tpu_observability/jobset_ttr_pod_oom.py @@ -0,0 +1,234 @@ +# 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 the `tpumonitoring` SDK, ensuring help() and +list_supported_metrics() are functional inside TPU worker pods.""" + +import datetime +import tempfile +import os +import random +from typing import List + +from airflow import models +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 JobSet, Workload +from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH + + +@task +def trigger_oom_failure(info, pod_name: str): + """ + Triggers an OOM (Out-of-Memory) event on a specified TPU 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. + + Args: + info: Node pool and cluster information. + pod_name (str): The name of the target pod to be OOMKilled. + """ + 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", + ]) + + try: + print(f"Executing OOM script in {pod_name}...") + subprocess.run_exec(cmd, env=env) + except Exception as e: + print(f"Expectation: Connection closed due to OOMKilled. Info: {e}") + +@task +def pick_random_pod(pod_names: 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 pod_names: + raise ValueError("No pods found to attack!") + chosen_pod = random.choice(pod_names) + print(f"Randomly selected pod for OOM test: {chosen_pod}") + return chosen_pod + +with models.DAG( + dag_id="jobset_ttr_pod_oom", + start_date=datetime.datetime(2026, 1, 15), + schedule="0 18 * * *" if composer_env.is_prod_env() else None, + catchup=False, + tags=[ + "cloud-ml-auto-solutions", + "jobset", + "tpu-observability", + "pod_oom", + "TPU", + "v6e-16", + "SDK", + "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 with memory limits, + injects a memory-intensive workload into a running pod, and uses a sensor + to confirm that the JobSet controller triggers a recovery and reports the + recovery duration. + ### Prerequisites + This test requires an existing cluster and a container image containing + the `oomkill.py` script to run. + ### Procedures + First, the node pool is created. A JobSet YAML with specific memory constraints + is then launched on the cluster and given time for all pods to reach a + `Running` state. After this, a random pod is selected and an OOM event is + triggered via `kubectl exec` to interrupt the JobSet (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 + + 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, + ) + + pod_names = 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')( + pod_names=pod_names + ) + + 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, + ) + + # Airflow uses >> for task chaining, which is pointless for pylint. + # pylint: disable=pointless-statement + ( + cluster_info + >> create_node_pool + >> start_workload + >> ensure_all_pods_running + >> pod_names + >> select_random_pod + >> trigger_oom_killed + >> wait_for_metric_upload + >> cleanup_workload + >> cleanup_node_pool + ) + # pylint: enable=pointless-statement From bd853dc5e62200e4afd6598941d19ccacb026486 Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Thu, 15 Jan 2026 16:31:43 +0800 Subject: [PATCH 03/15] format --- dags/tpu_observability/jobset_ttr_pod_oom.py | 107 ++++++++++--------- 1 file changed, 54 insertions(+), 53 deletions(-) diff --git a/dags/tpu_observability/jobset_ttr_pod_oom.py b/dags/tpu_observability/jobset_ttr_pod_oom.py index 3e2c6fd2b..9aa801b3d 100644 --- a/dags/tpu_observability/jobset_ttr_pod_oom.py +++ b/dags/tpu_observability/jobset_ttr_pod_oom.py @@ -36,53 +36,55 @@ @task def trigger_oom_failure(info, pod_name: str): - """ - Triggers an OOM (Out-of-Memory) event on a specified TPU 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. - - Args: - info: Node pool and cluster information. - pod_name (str): The name of the target pod to be OOMKilled. - """ - 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", - ]) - - try: - print(f"Executing OOM script in {pod_name}...") - subprocess.run_exec(cmd, env=env) - except Exception as e: - print(f"Expectation: Connection closed due to OOMKilled. Info: {e}") + """ + Triggers an OOM (Out-of-Memory) event on a specified TPU 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. + + Args: + info: Node pool and cluster information. + pod_name (str): The name of the target pod to be OOMKilled. + """ + 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", + ]) + + try: + print(f"Executing OOM script in {pod_name}...") + subprocess.run_exec(cmd, env=env) + except Exception as e: + print(f"Expectation: Connection closed due to OOMKilled. Info: {e}") + @task def pick_random_pod(pod_names: List[str]) -> str: - """ - Randomly selects one pod from a list of available JobSet pods. + """ + 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. - 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. - Args: - pod_names (List[str]): List of active pod names in the JobSet. + Returns: + str: The name of the randomly selected pod. + """ + if not pod_names: + raise ValueError("No pods found to attack!") + chosen_pod = random.choice(pod_names) + print(f"Randomly selected pod for OOM test: {chosen_pod}") + return chosen_pod - Returns: - str: The name of the randomly selected pod. - """ - if not pod_names: - raise ValueError("No pods found to attack!") - chosen_pod = random.choice(pod_names) - print(f"Randomly selected pod for OOM test: {chosen_pod}") - return chosen_pod with models.DAG( dag_id="jobset_ttr_pod_oom", @@ -146,7 +148,7 @@ def pick_random_pod(pod_names: List[str]) -> str: 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\"" + 'google.com/tpu: 4\n memory: "4Gi"', ) # Keyword arguments are generated dynamically at runtime (pylint does not # know this signature). @@ -185,21 +187,20 @@ def pick_random_pod(pod_names: List[str]) -> str: namespace=jobset_config.namespace, ) - select_random_pod = pick_random_pod.override(task_id='select_random_pod')( - pod_names=pod_names + select_random_pod = pick_random_pod.override(task_id="select_random_pod")( + pod_names=pod_names ) - trigger_oom_killed = trigger_oom_failure.override(task_id='trigger_oom_killed')( - info=cluster_info, - pod_name=select_random_pod - ) + 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, - ) + 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 From 0e334d5204550869a6cfa10ff790572f7e42c878 Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Thu, 15 Jan 2026 16:58:12 +0800 Subject: [PATCH 04/15] fix --- dags/tpu_observability/utils/jobset_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dags/tpu_observability/utils/jobset_util.py b/dags/tpu_observability/utils/jobset_util.py index 085de83eb..f015c8a89 100644 --- a/dags/tpu_observability/utils/jobset_util.py +++ b/dags/tpu_observability/utils/jobset_util.py @@ -749,7 +749,7 @@ def wait_for_jobset_started( return all(p > threshold_value for p in last_n_data_points) -@task.sensor(poke_interval=60, timeout=3600, mode="poke") +@task.sensor(poke_interval=60, timeout=3600, mode="reschedule") def wait_for_jobset_ttr_to_be_found( node_pool: node_pool_info, jobset_config: JobSet, From e26b0c02b37a202059f2eb76fdf14a7bbe52f5be Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Thu, 22 Jan 2026 10:45:48 +0800 Subject: [PATCH 05/15] fix --- dags/tpu_observability/utils/jobset_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dags/tpu_observability/utils/jobset_util.py b/dags/tpu_observability/utils/jobset_util.py index f015c8a89..085de83eb 100644 --- a/dags/tpu_observability/utils/jobset_util.py +++ b/dags/tpu_observability/utils/jobset_util.py @@ -749,7 +749,7 @@ def wait_for_jobset_started( return all(p > threshold_value for p in last_n_data_points) -@task.sensor(poke_interval=60, timeout=3600, mode="reschedule") +@task.sensor(poke_interval=60, timeout=3600, mode="poke") def wait_for_jobset_ttr_to_be_found( node_pool: node_pool_info, jobset_config: JobSet, From de9cb9b64aee16af6b744650d8fe1d1448d42183 Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Fri, 23 Jan 2026 16:31:01 +0800 Subject: [PATCH 06/15] fix --- dags/tpu_observability/jobset_ttr_pod_oom.py | 47 +++++++++----------- 1 file changed, 22 insertions(+), 25 deletions(-) diff --git a/dags/tpu_observability/jobset_ttr_pod_oom.py b/dags/tpu_observability/jobset_ttr_pod_oom.py index 9aa801b3d..bf3b71bf9 100644 --- a/dags/tpu_observability/jobset_ttr_pod_oom.py +++ b/dags/tpu_observability/jobset_ttr_pod_oom.py @@ -22,6 +22,7 @@ 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 @@ -30,7 +31,7 @@ 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, Workload +from dags.tpu_observability.utils.jobset_util import JobSet from dags.tpu_observability.configs.common import MachineConfigMap, GCS_CONFIG_PATH @@ -60,12 +61,12 @@ def trigger_oom_failure(info, pod_name: str): try: print(f"Executing OOM script in {pod_name}...") subprocess.run_exec(cmd, env=env) - except Exception as e: + except RuntimeError as e: print(f"Expectation: Connection closed due to OOMKilled. Info: {e}") @task -def pick_random_pod(pod_names: List[str]) -> str: +def pick_random_pod(active_pods: List[str]) -> str: """ Randomly selects one pod from a list of available JobSet pods. @@ -79,9 +80,9 @@ def pick_random_pod(pod_names: List[str]) -> str: Returns: str: The name of the randomly selected pod. """ - if not pod_names: + if not active_pods: raise ValueError("No pods found to attack!") - chosen_pod = random.choice(pod_names) + chosen_pod = random.choice(active_pods) print(f"Randomly selected pod for OOM test: {chosen_pod}") return chosen_pod @@ -98,13 +99,12 @@ def pick_random_pod(pod_names: List[str]) -> str: "pod_oom", "TPU", "v6e-16", - "SDK", "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." + "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 @@ -182,13 +182,13 @@ def pick_random_pod(pod_names: List[str]) -> str: node_pool=cluster_info, ) - pod_names = jobset.list_pod_names.override(task_id="list_pod_names")( + 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")( - pod_names=pod_names + active_pods=found_pods ) trigger_oom_killed = trigger_oom_failure.override( @@ -218,18 +218,15 @@ def pick_random_pod(pod_names: List[str]) -> str: setups=create_node_pool, ) - # Airflow uses >> for task chaining, which is pointless for pylint. - # pylint: disable=pointless-statement - ( - cluster_info - >> create_node_pool - >> start_workload - >> ensure_all_pods_running - >> pod_names - >> select_random_pod - >> trigger_oom_killed - >> wait_for_metric_upload - >> cleanup_workload - >> cleanup_node_pool + chain( + 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, ) - # pylint: enable=pointless-statement + From acc2591b936c251d55eb5875d0b7a58939c69152 Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Fri, 23 Jan 2026 16:31:52 +0800 Subject: [PATCH 07/15] format --- dags/tpu_observability/jobset_ttr_pod_oom.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/dags/tpu_observability/jobset_ttr_pod_oom.py b/dags/tpu_observability/jobset_ttr_pod_oom.py index bf3b71bf9..2ee5c9bc1 100644 --- a/dags/tpu_observability/jobset_ttr_pod_oom.py +++ b/dags/tpu_observability/jobset_ttr_pod_oom.py @@ -219,14 +219,13 @@ def pick_random_pod(active_pods: List[str]) -> str: ) chain( - 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, + 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, ) - From 3ae182a7f6d359a67b94972aa1378ff250c56775 Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Fri, 23 Jan 2026 16:33:41 +0800 Subject: [PATCH 08/15] fix --- dags/tpu_observability/jobset_ttr_pod_oom.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dags/tpu_observability/jobset_ttr_pod_oom.py b/dags/tpu_observability/jobset_ttr_pod_oom.py index 2ee5c9bc1..527eb7e3e 100644 --- a/dags/tpu_observability/jobset_ttr_pod_oom.py +++ b/dags/tpu_observability/jobset_ttr_pod_oom.py @@ -219,6 +219,7 @@ def pick_random_pod(active_pods: List[str]) -> str: ) chain( + cluster_info, create_node_pool, start_workload, ensure_all_pods_running, From a3c0efaad3ef90c84665963231f49eabbe1fa2a2 Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Fri, 23 Jan 2026 16:39:55 +0800 Subject: [PATCH 09/15] format --- dags/tpu_observability/jobset_ttr_pod_oom.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/dags/tpu_observability/jobset_ttr_pod_oom.py b/dags/tpu_observability/jobset_ttr_pod_oom.py index 527eb7e3e..625e6ebbe 100644 --- a/dags/tpu_observability/jobset_ttr_pod_oom.py +++ b/dags/tpu_observability/jobset_ttr_pod_oom.py @@ -46,8 +46,8 @@ def trigger_oom_failure(info, pod_name: str): A connection error is expected and caught when the pod is killed. Args: - info: Node pool and cluster information. - pod_name (str): The name of the target pod to be OOMKilled. + info: Node pool and cluster information. + pod_name (str): The name of the target pod to be OOMKilled. """ with tempfile.NamedTemporaryFile() as temp_config_file: env = os.environ.copy() @@ -75,10 +75,10 @@ def pick_random_pod(active_pods: List[str]) -> str: JobSet controller handles partial failures within a replica. Args: - pod_names (List[str]): List of active pod names in the JobSet. + pod_names (List[str]): List of active pod names in the JobSet. Returns: - str: The name of the randomly selected pod. + str: The name of the randomly selected pod. """ if not active_pods: raise ValueError("No pods found to attack!") From 3d45d34b734debc98523bea95ace7b0f4ace70d0 Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Fri, 23 Jan 2026 16:54:56 +0800 Subject: [PATCH 10/15] fix --- dags/tpu_observability/jobset_ttr_pod_oom.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dags/tpu_observability/jobset_ttr_pod_oom.py b/dags/tpu_observability/jobset_ttr_pod_oom.py index 625e6ebbe..91cc0fce9 100644 --- a/dags/tpu_observability/jobset_ttr_pod_oom.py +++ b/dags/tpu_observability/jobset_ttr_pod_oom.py @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""A DAG to validate the `tpumonitoring` SDK, ensuring help() and -list_supported_metrics() are functional inside TPU worker pods.""" +"""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 From 718efc0e08e616af0096abe598318797742d534d Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Mon, 30 Mar 2026 14:37:35 +0800 Subject: [PATCH 11/15] feat: add OOM fault injection for JobSet TTR validation --- dags/tpu_observability/jobset_ttr_pod_oom.py | 258 ++++++++++--------- 1 file changed, 138 insertions(+), 120 deletions(-) diff --git a/dags/tpu_observability/jobset_ttr_pod_oom.py b/dags/tpu_observability/jobset_ttr_pod_oom.py index 91cc0fce9..9aacb9c2e 100644 --- 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, + ) From e26dd092591a6ffc531808a2d056eb15faf30d69 Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Mon, 30 Mar 2026 14:40:25 +0800 Subject: [PATCH 12/15] fix docstring --- dags/tpu_observability/jobset_ttr_pod_oom.py | 24 ++++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/dags/tpu_observability/jobset_ttr_pod_oom.py b/dags/tpu_observability/jobset_ttr_pod_oom.py index 9aacb9c2e..97177dde6 100644 --- a/dags/tpu_observability/jobset_ttr_pod_oom.py +++ b/dags/tpu_observability/jobset_ttr_pod_oom.py @@ -149,20 +149,20 @@ def pick_random_pod(active_pods: List[str]) -> str: ### 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 with memory limits, - injects a memory-intensive workload into a running pod, and uses a sensor - to confirm that the JobSet controller triggers a recovery and reports the - recovery duration. + 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 a container image containing - the `oomkill.py` script to run. + 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 with specific memory constraints - is then launched on the cluster and given time for all pods to reach a - `Running` state. After this, a random pod is selected and an OOM event is - triggered via `kubectl exec` to interrupt the JobSet (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, + 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: From 62e8841a66b88464b5659e5ade36de1f0bef2129 Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Mon, 30 Mar 2026 17:33:37 +0800 Subject: [PATCH 13/15] fix --- dags/tpu_observability/jobset_ttr_pod_oom.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/dags/tpu_observability/jobset_ttr_pod_oom.py b/dags/tpu_observability/jobset_ttr_pod_oom.py index 97177dde6..80838f5f2 100644 --- a/dags/tpu_observability/jobset_ttr_pod_oom.py +++ b/dags/tpu_observability/jobset_ttr_pod_oom.py @@ -96,11 +96,10 @@ def trigger_oom_failure(info, pod_name: str, namespace: str): try: print(f"Blasting {pod_name} memory now...") subprocess.run_exec(full_command, env=env) + except subprocess.ProcessKilledException: + print(f"Success: Pod {pod_name} was OOMKilled.") 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}") + print(f"Connection lost or other error: {e}") @task From 348248d6e75bdd0408f6a4cbb8c10fa92f0a589c Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Tue, 7 Apr 2026 14:24:08 +0800 Subject: [PATCH 14/15] fix: update TTR DAG to use logging and active pod state polling --- dags/tpu_observability/jobset_ttr_pod_oom.py | 25 +++++++++++--------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/dags/tpu_observability/jobset_ttr_pod_oom.py b/dags/tpu_observability/jobset_ttr_pod_oom.py index 80838f5f2..f1893856b 100644 --- a/dags/tpu_observability/jobset_ttr_pod_oom.py +++ b/dags/tpu_observability/jobset_ttr_pod_oom.py @@ -19,7 +19,7 @@ import tempfile import os import random -import time +import logging from typing import List from airflow import models @@ -66,11 +66,7 @@ def trigger_oom_failure(info, pod_name: str, namespace: str): 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) + logging.info(f"Checking status for Pod {pod_name}...") python_logic = ( "import time\n" @@ -87,19 +83,26 @@ def trigger_oom_failure(info, pod_name: str, namespace: str): 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} && {exec_cmd}" + full_command = f"{credentials_cmd} && {wait_cmd} && {exec_cmd}" try: - print(f"Blasting {pod_name} memory now...") + logging.info(f"Blasting {pod_name} memory now...") subprocess.run_exec(full_command, env=env) except subprocess.ProcessKilledException: - print(f"Success: Pod {pod_name} was OOMKilled.") + logging.info(f"Success: Pod {pod_name} was OOMKilled.") except Exception as e: - print(f"Connection lost or other error: {e}") + logging.error(f"Connection lost or other error: {e}") + raise e @task @@ -120,7 +123,7 @@ def pick_random_pod(active_pods: List[str]) -> str: if not active_pods: raise ValueError("No pods found to attack!") chosen_pod = random.choice(active_pods) - print(f"Randomly selected pod for OOM test: {chosen_pod}") + logging.info(f"Randomly selected pod for OOM test: {chosen_pod}") return chosen_pod From 934886f7aed34bc883d9d24a48489f8f18597b59 Mon Sep 17 00:00:00 2001 From: EmmaLien Date: Mon, 11 May 2026 08:19:34 +0000 Subject: [PATCH 15/15] fix --- dags/common/scheduling_helper/scheduling_helper.py | 1 + 1 file changed, 1 insertion(+) diff --git a/dags/common/scheduling_helper/scheduling_helper.py b/dags/common/scheduling_helper/scheduling_helper.py index e0cff02dc..e3eaa7814 100644 --- a/dags/common/scheduling_helper/scheduling_helper.py +++ b/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,