Use process group to manage the trainer process in AXLearnFT

GitOrigin-RevId: 1994c53

Author: Liang (AIML-Compute) He

axlearn/ft/agent.py

@@ -14,6 +14,7 @@
   --max_restarts=5
 """
 
+import enum
 import signal
 import subprocess
 import sys
@@ -25,11 +26,58 @@
 from axlearn.ft.monitor import StatusMonitor
 from axlearn.ft.utils import TrainerProcessController
 
+
+class TerminationAction(enum.Enum):
+    """Action to take after a termination request."""
+
+    EXIT = "exit"  # Pod shutdown - exit gracefully
+    RESTART = "restart"  # Coordinated restart - restart without counting against max_restarts
+
+
+# Constant for pod shutdown signal reason prefix
+_POD_SHUTDOWN_REASON_PREFIX = "Pod shutdown signal"
+
 flags.DEFINE_integer(
     "max_restarts", 3, "Maximum number of times to restart the trainer subprocess on failure"
 )
 
 
+def _handle_termination_request(
+    process_controller: TrainerProcessController,
+) -> TerminationAction | None:
+    """Check if termination was requested and determine the appropriate action.
+
+    Args:
+        process_controller: The trainer process controller to check for termination requests.
+
+    Returns:
+        TerminationAction.EXIT if pod shutdown was requested (caller should return/exit
+            gracefully),
+        TerminationAction.RESTART if coordinated restart was requested (caller should continue
+            without incrementing restart count),
+        None if no termination was requested (caller should handle returncode normally).
+    """
+    termination_requested, reason = process_controller.check_termination_requested()
+    if not termination_requested:
+        return None
+
+    # Distinguish between pod shutdown (exit) and coordinated restart (restart)
+    if reason.startswith(_POD_SHUTDOWN_REASON_PREFIX):
+        # SIGTERM received - pod is being killed, exit gracefully
+        logging.info(
+            "FT Agent: Pod shutdown signal received (%s), exiting gracefully",
+            reason,
+        )
+        return TerminationAction.EXIT
+
+    # Coordinated restart (JAX re-init) - restart trainer without counting against max_restarts
+    logging.info(
+        "FT Agent: Coordinated restart requested (%s), restarting trainer",
+        reason,
+    )
+    return TerminationAction.RESTART
+
+
 def run_ft_agent():
     """The agent launches trainer as a subprocess with fault tolerance."""
     logging.info("Starting fault tolerant trainer agent...")
@@ -52,14 +100,14 @@ def signal_handler(signum, *_):
 
         # Report pod shutdown to global manager for coordination
         try:
-            monitor.manager.report_pod_shutdown(f"Pod shutdown signal {signum}")
+            monitor.manager.report_pod_shutdown(f"{_POD_SHUTDOWN_REASON_PREFIX} {signum}")
             logging.info("FT Agent: Pod shutdown reported to global manager")
         except Exception as e:  # pylint: disable=broad-exception-caught
             logging.error("FT Agent: Failed to report pod shutdown: %s", e)
 
         # Terminate the trainer process
         logging.info("FT Agent: Terminating trainer due to signal %d", signum)
-        process_controller.terminate_training(f"Pod shutdown signal {signum}")
+        process_controller.terminate_training(f"{_POD_SHUTDOWN_REASON_PREFIX} {signum}")
 
         logging.warning("FT Agent: exit with non-zero code due to signal %d recieved.", signum)
         sys.exit(1)
@@ -77,13 +125,26 @@ def signal_handler(signum, *_):
             logging.info("FT Agent: Starting trainer: %s", " ".join(entrypoint_cmd))
 
             try:
+                # Start trainer in its own process group (start_new_session=True)
+                # This allows os.killpg() to terminate all threads cleanly
                 with subprocess.Popen(
-                    entrypoint_cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True
+                    entrypoint_cmd,
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.STDOUT,
+                    text=True,
+                    start_new_session=True,
                 ) as process:
                     process_controller.set_process(process)
                     returncode = monitor.monitor_training_process(process)
                     process_controller.clear_process()
 
+                    # Handle termination requests (SIGTERM or coordinated restart)
+                    action = _handle_termination_request(process_controller)
+                    if action == TerminationAction.EXIT:
+                        return
+                    if action == TerminationAction.RESTART:
+                        continue  # Restart without incrementing restart_count
+
                     if returncode == 0:
                         logging.info("FT Agent: Training completed successfully")
                         return

axlearn/ft/utils.py

@@ -4,10 +4,10 @@
 
 Environment Variables Used:
     HOSTNAME: Current worker's hostname
-    REPLICA_ID: Current worker's replica ID (0-based)
+    MEGASCALE_SLICE_ID: Current worker's slice/replica ID (0-based)
     TPU_WORKER_ID: Current worker's ID within the replica (0-based)
     TPU_WORKER_HOSTNAMES: Comma-separated list of all worker hostnames in current replica
-    NUM_REPLICAS: Total number of replicas in the training job
+    NUM_TPU_SLICES: Total number of slices/replicas in the training job
 
 Hostname Format:
     TPU worker hostnames follow the pattern: {job_name}-"job"-{replica_id}-{worker_id}.{job_name}
@@ -18,6 +18,7 @@
 import os
 import pathlib
 import re
+import signal
 import subprocess
 import threading
 import time
@@ -35,9 +36,6 @@
 # Client timeout in seconds
 DEFAULT_CLIENT_TIMEOUT = 10.0
 
-# Timeout for graceful training process termination in seconds
-GRACEFUL_TRAINING_TERMINATION_TIMEOUT = 1
-
 # Compiled regex patterns for hostname parsing
 _WORKER_ID_PATTERN = re.compile(r"^.+?-job-\d+-(\d+)\..+")
 _JOB_NAME_PATTERN = re.compile(r"^(.+?)-job-\d+-\d+\.(\1)(?:\.|$)")
@@ -110,26 +108,43 @@ def _cleanup_tpu_lock_files(self) -> None:
     def _do_terminate(self) -> bool:
         """Execute the actual process termination sequence.
 
+        Uses process group killing to ensure all threads are terminated.
+        The trainer subprocess should be started with start_new_session=True.
+
         Returns:
             bool: True if termination succeeded, False otherwise
         """
-        self.current_process.terminate()
+        pid = self.current_process.pid
+
+        # Try to get process group ID
         try:
-            self.current_process.wait(timeout=GRACEFUL_TRAINING_TERMINATION_TIMEOUT)
-            logging.info("Process terminated gracefully")
+            pgid = os.getpgid(pid)
+        except ProcessLookupError:
+            logging.info("Process %d already dead", pid)
+            self._cleanup_tpu_lock_files()
             return True
-        except subprocess.TimeoutExpired:
-            pass
 
-        logging.warning("Process did not terminate gracefully, forcing kill")
-        self.current_process.kill()
+        # If trainer has its own process group, kill the entire group
+        if pgid == pid:
+            logging.warning("Sending SIGKILL to process group %d", pgid)
+            try:
+                os.killpg(pgid, signal.SIGKILL)
+            except ProcessLookupError:
+                logging.info("Process group %d already dead", pgid)
+            except PermissionError as e:
+                logging.warning("Cannot kill process group %d: %s, killing process only", pgid, e)
+                self.current_process.kill()
+        else:
+            logging.warning("Sending SIGKILL to process %d", pid)
+            self.current_process.kill()
+
         try:
-            self.current_process.wait(timeout=5)
+            self.current_process.wait(timeout=10)
             logging.info(
                 "Process killed successfully, exit code: %s", self.current_process.returncode
             )
         except subprocess.TimeoutExpired:
-            logging.error("Process still alive after SIGKILL, pid=%d", self.current_process.pid)
+            logging.error("Process still alive after SIGKILL, pid=%d", pid)
             return False
 
         self._cleanup_tpu_lock_files()
@@ -220,19 +235,20 @@ def get_replica_id() -> int:
     """Get replica ID from environment variable with validation.
 
     Returns:
-        The replica ID as an integer.
+        The replica ID as an integer. Defaults to 0 for single-slice jobs.
 
     Raises:
-        ValueError: If REPLICA_ID is not set or not a valid integer.
+        ValueError: If MEGASCALE_SLICE_ID is not a valid integer.
     """
-    replica_id_str = os.environ.get("REPLICA_ID")
-    if replica_id_str is None:
-        raise ValueError("REPLICA_ID environment variable is not set")
+    # Default to "0" for single-slice jobs where MEGASCALE_SLICE_ID is not set
+    replica_id_str = os.environ.get("MEGASCALE_SLICE_ID", "0")
 
     try:
         return int(replica_id_str)
     except ValueError as e:
-        raise ValueError(f"REPLICA_ID must be a valid integer, got: '{replica_id_str}'") from e
+        raise ValueError(
+            f"MEGASCALE_SLICE_ID must be a valid integer, got: '{replica_id_str}'"
+        ) from e
 
 
 def get_worker_id() -> int:
@@ -258,23 +274,22 @@ def get_num_replicas() -> int:
     """Get total number of replicas from environment variable with validation.
 
     Returns:
-        The total number of replicas as an integer.
+        The total number of replicas as an integer. Defaults to 1 for single-slice jobs.
 
     Raises:
-        ValueError: If NUM_REPLICAS is not set or not a valid integer.
+        ValueError: If NUM_TPU_SLICES is not a valid positive integer.
     """
-    num_replicas_str = os.environ.get("NUM_REPLICAS")
-    if num_replicas_str is None:
-        raise ValueError("NUM_REPLICAS environment variable is not set")
+    # Default to "1" for single-slice jobs where NUM_TPU_SLICES is not set
+    num_replicas_str = os.environ.get("NUM_TPU_SLICES", "1")
 
     try:
         num_replicas = int(num_replicas_str)
         if num_replicas < 1:
-            raise ValueError(f"NUM_REPLICAS must be at least 1, got: {num_replicas}")
+            raise ValueError(f"NUM_TPU_SLICES must be at least 1, got: {num_replicas}")
         return num_replicas
     except ValueError as e:
         raise ValueError(
-            f"NUM_REPLICAS must be a valid positive integer, got: '{num_replicas_str}'"
+            f"NUM_TPU_SLICES must be a valid positive integer, got: '{num_replicas_str}'"
         ) from e