[AXLearnFT] Optimizing the FT restart training gRPC

GitOrigin-RevId: 0abc902

Author: Liang (AIML-Compute) He

axlearn/ft/manager_server.py

@@ -30,6 +30,12 @@
     worker_status_record_to_proto,
 )
 
+# Maximum parallel threads for forwarding restart requests
+MAX_RESTART_FORWARD_THREADS = 100
+
+# Maximum concurrent async termination tasks per worker
+MAX_TERMINATION_EXECUTOR_THREADS = 1  # Each worker has only 1 training process
+
 
 class ManagerServer(manager_pb2_grpc.ManagerServiceServicer):
     """FT trainer manager server.
@@ -58,6 +64,12 @@ def __init__(
         self.process_controller = process_controller
         self._registry_lock = threading.Lock()
 
+        # Thread pool for async termination tasks (bounded to prevent resource exhaustion)
+        self._termination_executor = futures.ThreadPoolExecutor(
+            max_workers=MAX_TERMINATION_EXECUTOR_THREADS,
+            thread_name_prefix="termination",
+        )
+
         # Initialize registries based on role
         if self._worker_identity.is_replica_manager:
             self._status_registry: Dict[str, WorkerStatusRecord] = {}
@@ -136,11 +148,15 @@ def start_server(self):
         return self.server
 
     def stop_server(self, grace_period: int = 10):
-        """Stop the gRPC server."""
+        """Stop the gRPC server and clean up resources."""
         if self.server:
             self.server.stop(grace_period)
             logging.info("Manager server stopped")
 
+        # Shutdown termination executor (wait for pending tasks to complete)
+        self._termination_executor.shutdown(wait=True)
+        logging.info("Termination executor stopped")
+
     def ReportStatus(
         self, request: manager_pb2.StatusUpdate, context
     ) -> manager_pb2.StatusUpdateResponse:
@@ -296,6 +312,26 @@ def RestartReplicaTraining(
             )
             raise
 
+    def _terminate_training_async(self, hostname: str, reason: str, is_shutdown: bool) -> None:
+        """Terminate training process asynchronously in background thread.
+
+        Args:
+            hostname: Worker hostname for logging
+            reason: Reason for termination
+            is_shutdown: Whether this is a shutdown (vs restart) request
+        """
+        try:
+            success = self.process_controller.terminate_training(reason)
+            if success:
+                action = "shutdown" if is_shutdown else "restart"
+                logging.info("Worker %s training %s completed", hostname, action)
+                if is_shutdown:
+                    logging.warning("Immediate shutdown requested, exiting after termination")
+            else:
+                logging.warning("Worker %s has no active training process", hostname)
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            logging.error("Async termination failed: %s", e)
+
     def RestartTraining(
         self, request: manager_pb2.RestartRequest, context
     ) -> manager_pb2.RestartResponse:
@@ -316,32 +352,25 @@ def RestartTraining(
             # Check if this is a shutdown request (not a restart)
             is_shutdown = request.reason.startswith("IMMEDIATE_SHUTDOWN:")
 
-            # Attempt to terminate the trainer process
-            success = False
-            message = ""
-
-            if self.process_controller:
-                success = self.process_controller.terminate_training(request.reason)
-                if success:
-                    action = "shutdown" if is_shutdown else "restart"
-                    message = f"Worker {worker_identity.hostname} training {action} initiated"
-                    logging.info(message)
-
-                    # For shutdown requests, exit immediately (don't restart)
-                    if is_shutdown:
-                        logging.warning("Immediate shutdown requested, exiting after termination")
-                        # The FT agent should not restart after this termination
-
-                else:
-                    message = f"Worker {worker_identity.hostname} has no active training process"
-                    logging.warning(message)
-            else:
+            if not self.process_controller:
                 message = f"Worker {worker_identity.hostname} has no process controller"
                 logging.error(message)
-                success = False
+                return self._create_success_response(
+                    manager_pb2.RestartResponse, acknowledged=False, message=message
+                )
+
+            # Submit termination to thread pool - return immediately
+            self._termination_executor.submit(
+                self._terminate_training_async,
+                worker_identity.hostname,
+                request.reason,
+                is_shutdown,
+            )
 
+            action = "shutdown" if is_shutdown else "restart"
+            message = f"Worker {worker_identity.hostname} training {action} initiated (async)"
             return self._create_success_response(
-                manager_pb2.RestartResponse, acknowledged=success, message=message
+                manager_pb2.RestartResponse, acknowledged=True, message=message
             )
 
         except grpc.RpcError:
@@ -435,29 +464,42 @@ def _forward_restart_to_workers(self, reason: str) -> Dict[str, bool]:
             logging.error("Only replica managers can forward restart requests")
             return {}
 
-        results = {}
         worker_hostnames = get_all_worker_hostnames()
 
-        with ManagerClient() as client:
-            for hostname in worker_hostnames:
-                try:
-                    logging.info("Forwarding restart to worker: %s", hostname)
+        def restart_single_worker(hostname: str, client: ManagerClient) -> tuple:
+            """Send restart request to a single worker."""
+            try:
+                logging.info("Forwarding restart to worker: %s", hostname)
+                worker_id = extract_worker_id_from_hostname(hostname)
 
-                    worker_id = extract_worker_id_from_hostname(hostname)
+                success = client.restart_worker(
+                    hostname, self._worker_identity.replica_id, reason, worker_id
+                )
 
-                    success = client.restart_worker(
-                        hostname, self._worker_identity.replica_id, reason, worker_id
-                    )
-                    results[hostname] = success
+                if success:
+                    logging.info("Worker %s acknowledged restart request", hostname)
+                else:
+                    logging.warning("Worker %s failed to acknowledge restart request", hostname)
+
+                return (hostname, success)
 
-                    if success:
-                        logging.info("Worker %s acknowledged restart request", hostname)
-                    else:
-                        logging.warning("Worker %s failed to acknowledge restart request", hostname)
+            except Exception as e:  # pylint: disable=broad-exception-caught
+                logging.error("Failed to forward restart to worker %s: %s", hostname, e)
+                return (hostname, False)
 
-                except Exception as e:  # pylint: disable=broad-exception-caught
-                    logging.error("Failed to forward restart to worker %s: %s", hostname, e)
-                    results[hostname] = False
+        # Parallelize restart requests to all workers
+        # Use higher parallelism for restart - these are I/O-bound gRPC calls
+        results = {}
+        max_workers = min(len(worker_hostnames), MAX_RESTART_FORWARD_THREADS)
+        with ManagerClient() as client:
+            with futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+                future_to_hostname = {
+                    executor.submit(restart_single_worker, hostname, client): hostname
+                    for hostname in worker_hostnames
+                }
+                for future in futures.as_completed(future_to_hostname):
+                    hostname, success = future.result()
+                    results[hostname] = success
 
         return results
 

axlearn/ft/utils.py

@@ -16,6 +16,7 @@
 
 import logging
 import os
+import pathlib
 import re
 import subprocess
 import threading
@@ -34,6 +35,9 @@
 # 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)(?:\.|$)")
@@ -88,6 +92,48 @@ def set_process(self, process: subprocess.Popen):
             self.termination_requested = False
             self.termination_reason = ""
 
+    def _cleanup_tpu_lock_files(self) -> None:
+        """Clean up TPU lock files.
+
+        See: https://github.com/jax-ml/jax/issues/10192#issuecomment-1509814942
+        """
+        for lock_file in pathlib.Path("/tmp").glob("libtpu_lockfile*"):
+            try:
+                lock_file.unlink()
+                logging.info("Cleaned up stale lock file: %s", lock_file)
+            except FileNotFoundError:
+                pass  # File already removed, this is fine
+            except OSError as e:
+                logging.warning("Failed to remove TPU lock file %s: %s", lock_file, e)
+
+    def _do_terminate(self) -> bool:
+        """Execute the actual process termination sequence.
+
+        Returns:
+            bool: True if termination succeeded, False otherwise
+        """
+        self.current_process.terminate()
+        try:
+            self.current_process.wait(timeout=GRACEFUL_TRAINING_TERMINATION_TIMEOUT)
+            logging.info("Process terminated gracefully")
+            return True
+        except subprocess.TimeoutExpired:
+            pass
+
+        logging.warning("Process did not terminate gracefully, forcing kill")
+        self.current_process.kill()
+        try:
+            self.current_process.wait(timeout=5)
+            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)
+            return False
+
+        self._cleanup_tpu_lock_files()
+        return True
+
     def terminate_training(self, reason: str = "Restart requested"):
         """Terminate the current trainer process.
 
@@ -98,27 +144,20 @@ def terminate_training(self, reason: str = "Restart requested"):
             bool: True if termination was initiated, False if no process to terminate
         """
         with self.lock:
-            if self.current_process:
-                logging.warning("Terminating trainer process: %s", reason)
-                self.termination_requested = True
-                self.termination_reason = reason
-
-                try:
-                    self.current_process.terminate()  # Send SIGTERM first
-                    try:
-                        self.current_process.wait(timeout=5)  # Wait for graceful shutdown
-                        logging.info("Process terminated gracefully")
-                    except subprocess.TimeoutExpired:
-                        logging.warning("Process did not terminate gracefully, forcing kill")
-                        self.current_process.kill()  # Force kill if timeout
-                    return True
-                except Exception as e:  # pylint: disable=broad-exception-caught
-                    logging.error("Failed to terminate process: %s", e)
-                    return False
-            else:
+            if not self.current_process:
                 logging.warning("No trainer process to terminate")
                 return False
 
+            logging.warning("Terminating trainer process: %s", reason)
+            self.termination_requested = True
+            self.termination_reason = reason
+
+            try:
+                return self._do_terminate()
+            except Exception as e:  # pylint: disable=broad-exception-caught
+                logging.error("Failed to terminate process: %s", e)
+                return False
+
     def check_termination_requested(self):
         """Check if termination was requested and get reason.