Merge pull request #1404 from AI-Hypercomputer:sujinesh/add_disruption_manager_sigterm

PiperOrigin-RevId: 742898218

Author: maxtext authors

benchmarks/disruption_management/disruption_handler.py

@@ -15,12 +15,14 @@
 import dataclasses
 import enum
 import os
-import subprocess
 import sys
 
 parent_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
 sys.path.append(parent_dir)
 
+
+from disruption_management.disruption_utils import execute_command_as_subprocess
+from disruption_management.disruption_utils import get_pod_name_from_regex
 from xpk_configs import XpkClusterConfig
 
 
@@ -29,8 +31,6 @@
 PATHWAYS_STANDARD_TARGET_POD_REGEX_SUFFIX = ".*worker-0-0.*"
 PATHWAYS_STANDARD_STEP_POD_REGEX_SUFFIX = ".*main-0-0.*"
 
-STANDARD_STEP_LOG_REGEX = "completed step: (\\d+)"
-
 PATHWAYS_WORKER_CONTAINER_NAME = "pathways-worker"
 MCJAX_WORKER_CONTAINER_NAME = "jax-tpu"
 
@@ -71,6 +71,9 @@ class DisruptionConfig:
   # Target pod regex needed for triggering disruption.
   target_pod_regex: str = MCJAX_STANDARD_TARGET_POD_REGEX_SUFFIX
 
+  # Step pod regex needed for step-based disruption.
+  step_pod_regex: str = MCJAX_STANDARD_STEP_POD_REGEX_SUFFIX
+
 
 class DisruptionHandler(abc.ABC):
   """Abstract interface for disruption handlers."""
@@ -103,73 +106,50 @@ def trigger_disruption(
         f"🔥🔥🔥 Beginning SIGILL for workload: {workload_name} with pod regex:"
         f" {target_pod_regex} 🔥🔥🔥"
     )
-    pod_name_command = [
-        "kubectl",
-        "get",
-        "pods",
-        "-o=name",
-        "--no-headers",
-        f"| grep -E '{target_pod_regex}'",
-    ]
-    pod_name_command_str = " ".join(pod_name_command)
     container_name = disruption_config.worker_container_name
 
-    try:
-      pod_name_process = subprocess.run(
-          pod_name_command_str,
-          shell=True,
-          check=True,
-          capture_output=True,
-          text=True,
-      )
-      pod_name = pod_name_process.stdout.strip()
-      if not pod_name:
-        print(
-            f"Warning: No pod found matching regex: '{target_pod_regex}' for"
-            f" workload '{workload_name}'"
-        )
-        return
-
-      print(f"🔍 Found pod: {pod_name}")
-      kill_command = [
-          "kubectl",
-          "exec",
-          "-it",
-          pod_name,
-          "-c",
-          container_name,
-          "--",
-          "/bin/bash",
-          "-c",
-          "\"kill -s SIGILL 1\"",
-      ]
-      kill_command_str = " ".join(kill_command)
-      print(f"🔥🔥🔥 Executing command in pod: {kill_command_str} 🔥🔥🔥")
-      subprocess.run(
-          kill_command_str,
-          shell=True,
-          check=True,
-          capture_output=True,
-          text=True,
-      )
-      print(
-          f"✅ Successfully sent SIGILL to pod: {pod_name} in container:"
-          f" {container_name}"
-      )
-
-    except subprocess.CalledProcessError as e:
-      print(
-          "❌ Error sending SIGILL to pod(s) matching regex"
-          f" '{target_pod_regex}' for workload '{workload_name}'"
-      )
-      print(f"Return code: {e.returncode}")
-      print(f"error: {e}")
+    pod_name = get_pod_name_from_regex(workload_name, target_pod_regex)
+    if not pod_name:
+      return
+
+    kill_command = (
+        f"kubectl exec -it {pod_name} -c {container_name} -- /bin/sh -c "
+        f'"kill -s SIGILL 1"'
+    )
+    print(f"🔥🔥🔥 Executing command in pod: {kill_command} 🔥🔥🔥")
+    execute_command_as_subprocess(kill_command)
+
+
+class SIGTERMHandler(DisruptionHandler):
+  """Handles SIGTERM disruption by sending a SIGTERM signal to the pod."""
+
+  def trigger_disruption(
+      self, workload_name: str, cluster_config: XpkClusterConfig,
+      disruption_config, target_pod_regex: str
+  ) -> None:
+    """Triggers SIGTERM disruption by executing kill -s SIGTERM 1 in the pod."""
+    print(
+        f"🔥🔥🔥 Beginning SIGTERM for workload: {workload_name} with pod regex:"
+        f" {target_pod_regex} 🔥🔥🔥"
+    )
+    container_name = disruption_config.worker_container_name
+
+    pod_name = get_pod_name_from_regex(workload_name, target_pod_regex)
+    if not pod_name:
+      return
+
+    kill_command = (
+        f"kubectl exec -it {pod_name} -c {container_name} -- /bin/sh -c "
+        f'"kill -s SIGTERM 1"'
+    )
+    print(f"🔥🔥🔥 Executing command in pod: {kill_command} 🔥🔥🔥")
+    execute_command_as_subprocess(kill_command)
 
 
 def create_disruption_handler(disruption_config):
   """Factory function to create the appropriate disruption handler."""
   if disruption_config.disruption_method == DisruptionMethod.SIGTERM:
-    raise NotImplementedError("SIGTERM Disruption Handler not implemented yet.")
+    return SIGTERMHandler()
   elif disruption_config.disruption_method == DisruptionMethod.SIGILL:
     return SIGILLHandler()
   else:

benchmarks/disruption_management/disruption_manager.py

@@ -113,10 +113,11 @@ def _monitor_and_disrupt_workload(
   ) -> None:
     """Monitors workload progress, triggers disruptions, and recoveries."""
     target_pod_regex = f"{workload_name}{disruption_config.target_pod_regex}"
+    step_pod_regex = f"{workload_name}{disruption_config.step_pod_regex}"
 
     # Create Monitor based on trigger type
     monitor: Monitor = create_monitor(
-        workload_name, disruption_config, target_pod_regex
+        workload_name, disruption_config, step_pod_regex
     )
     disruption_handler: DisruptionHandler = create_disruption_handler(
         disruption_config

benchmarks/disruption_management/disruption_utils.py

@@ -0,0 +1,137 @@
+"""
+ 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
+      https://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.
+ """
+
+"""Utility functions for disrupting workloads."""
+
+import subprocess
+import time
+
+POLLING_INTERVAL_SECONDS = 10
+
+
+def execute_command_as_subprocess(
+    command: str,
+) -> None:
+  """Executes the command in the pod."""
+  print(f"Executing command: {command}")
+  try:
+    process = subprocess.run(
+        command,
+        shell=True,
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    if process.stdout:
+      print(process.stdout)
+    if process.stderr:
+      print(f"stderr: {process.stderr}")
+    print(f"✅ Successfully executed command: {command}")
+  except subprocess.CalledProcessError as e:
+    print(f"❌ Error executing command: {command}")
+    print(f"Return code: {e.returncode}")
+    print(f"Error: {e}")
+
+
+def get_pod_name_from_regex(workload_name: str, pod_regex: str) -> str | None:
+  """Returns the name of the first pod matching the regex."""
+  print(
+      f"Workload '{workload_name}': Getting pod name matching"
+      f" '{pod_regex}'..."
+  )
+  pod_name_command = [
+      "kubectl",
+      "get",
+      "pods",
+      "-o=custom-columns=NAME:.metadata.name",
+      "--no-headers",
+      f"| grep -E '{pod_regex}'",
+  ]
+  pod_name_command_str = " ".join(pod_name_command)
+  try:
+    process = subprocess.run(
+        pod_name_command_str,
+        shell=True,
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    pod_names = process.stdout.strip().splitlines()
+    if pod_names:
+      # Assuming there's only one step pod.
+      pod_name = pod_names[0]
+      print(f"Workload '{workload_name}': Found pod: {pod_name}")
+      return pod_name
+    else:
+      print(
+          f"Workload '{workload_name}': No pod found matching"
+          f" regex '{pod_regex}'."
+      )
+  except subprocess.CalledProcessError as e:
+    print(
+        f"Workload '{workload_name}': Error getting pod information:"
+        f" {e}"
+    )
+
+
+def get_pod_status(workload_name: str, pod_name: str) -> str | None:
+  """Returns the status of the pod."""
+  print(
+      f"Workload '{workload_name}': Getting status of pod '{pod_name}'..."
+  )
+  pod_status_command = [
+      "kubectl",
+      "get",
+      "pod",
+      pod_name,
+      "-o=jsonpath='{.status.phase}'",
+  ]
+  pod_status_command_str = " ".join(pod_status_command)
+  status_process = subprocess.run(
+      pod_status_command_str,
+      shell=True,
+      check=True,
+      capture_output=True,
+      text=True,
+  )
+  pod_status = status_process.stdout.strip()
+  print(
+      f"Workload '{workload_name}': Pod '{pod_name}' is in '{pod_status}'"
+      " state."
+  )
+  return pod_status
+
+
+def wait_for_pod_to_start(workload_name: str, pod_regex: str) -> str | None:
+  """Waits for the step pod to be in 'Running' state and returns its name."""
+  print(
+      f"Workload '{workload_name}': Waiting for pod matching"
+      f" '{pod_regex}' to be in 'Running' state..."
+  )
+  while True:
+    pod_name = get_pod_name_from_regex(workload_name, pod_regex)
+    if pod_name:
+      pod_status = get_pod_status(workload_name, pod_name)
+      if pod_status == "Running":
+        print(
+            f"Workload '{workload_name}': Step pod '{pod_name}'"
+            f" is now in 'Running' state."
+        )
+        return pod_name
+    time.sleep(POLLING_INTERVAL_SECONDS)
+
+  print(
+      f"Workload '{workload_name}': Timed out waiting for step pod"
+      f" matching '{pod_regex}' to reach 'Running' state."
+  )
+  return None