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/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, 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..f1893856b --- /dev/null +++ b/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, + ) 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."