feat: Add progress-based early termination for in-job restarts

Implement centralized training progress tracking to automatically
terminate jobs when restarts fail to make meaningful progress, preventing wasted
compute on jobs stuck in restart loops.

Problem:
--------

The InJob restart would continue restarting failed jobs up to
max_restarts regardless of whether restarts were effective. Jobs that
crash repeatedly wast compute time across many restart attempts before
exhausting the restart budget.

Solution:
---------

- Immediate arming: Rank monitors send iteration=1 on init to enable
  tracking
- Periodic update: Rank monitors send iteration updates every 30
  seconds.
- Early termination after 3 consecutive restarts with <200 iteration progress
- Works for both fast-failing (seconds) and slow-failing (minutes/hours)
  jobs

Configuration:
--------------

--ft-max-no-progress-restarts=3      # Max restarts without progress (0 = disabled)
--ft-min-progress-iterations=200     # Min iterations to show progress
--ft-progress-update-interval=30.0   # Update frequency (seconds)

Author: hexinw-nvidia

src/nvidia_resiliency_ext/fault_tolerance/config.py

@@ -85,6 +85,15 @@ class FaultToleranceConfig:
     skip_section_response: bool = True
     use_infra_group_rank: bool = True
     numa_bind_strict: bool = False
+    # Progress tracking configuration (controlled by max_no_progress_restarts)
+    max_no_progress_restarts: int = 3
+    min_progress_iterations: int = 200
+    progress_update_interval: float = 30.0  # Seconds between sending progress updates to launcher
+
+    @property
+    def is_progress_tracking_enabled(self) -> bool:
+        """Check if progress tracking is enabled (controlled by max_no_progress_restarts > 0)."""
+        return self.max_no_progress_restarts > 0
 
     @staticmethod
     def from_kwargs(ignore_not_recognized: bool = True, **kwargs) -> 'FaultToleranceConfig':

src/nvidia_resiliency_ext/fault_tolerance/data.py

@@ -188,10 +188,12 @@ def __init__(
         rank: int,
         section: str,
         action: SectionAction,
+        iteration: Optional[int] = None,
     ):
         self.rank = rank
         self.section = section
         self.action = action
+        self.iteration = iteration
 
 
 class UpdateConfigMsg:

src/nvidia_resiliency_ext/fault_tolerance/launcher.py

@@ -69,9 +69,11 @@
     FT_LAUNCHER_IPC_SOCKET_ENV_VAR,
     FT_RANK_MONITOR_IPC_SOCKET_ENV_VAR,
 )
+from nvidia_resiliency_ext.fault_tolerance.progress_tracker import TrainingProgressTracker
 from nvidia_resiliency_ext.fault_tolerance.rank_monitor_server import RankMonitorServer
 from nvidia_resiliency_ext.fault_tolerance.utils import (
     patched_method,
+    read_obj_from_ipc_stream,
     terminate_mp_processes,
     write_obj_to_ipc_stream,
 )
@@ -192,6 +194,15 @@ def _wrap_entrypoint_with_numactl(
 # https://github.com/pytorch/pytorch/blob/release/2.3/torch/distributed/elastic/agent/server/local_elastic_agent.py
 
 
+@dataclass
+class RankMonitorState:
+    """State for a single rank monitor process and its IPC connections."""
+    process: Any  # multiprocessing.Process
+    reader: Optional[asyncio.StreamReader] = None
+    writer: Optional[asyncio.StreamWriter] = None
+    listener_task: Optional[asyncio.Task] = None
+
+
 class LocalElasticAgent(SimpleElasticAgent):
     """An implementation of :py:class:`torchelastic.agent.server.ElasticAgent` that handles host-local workers.
 
@@ -317,8 +328,15 @@ def __init__(
         self._term_timeout = term_timeout
         self._workers_stop_timeout = workers_stop_timeout
         self._is_store_host = is_store_host
-        self._local_rank_to_rmon: Dict[int, Any] = dict()
+        # Rank monitor state (process, IPC connections, listener tasks) per local rank
+        self._rank_monitors: Dict[int, RankMonitorState] = dict()
         self._ft_cfg = fault_tol_cfg
+        # Centralized progress tracking (always instantiated, active only if configured)
+        self._progress_tracker = TrainingProgressTracker(
+            min_progress_iterations=fault_tol_cfg.min_progress_iterations,
+            max_no_progress_restarts=fault_tol_cfg.max_no_progress_restarts,
+        )
+        self._rank_iterations: Dict[int, int] = dict()  # Track max iteration per rank
         self._children_pgids: Set[int] = set()
         self._restart_policy = restart_policy
         self._node_id = self._get_fq_hostname()
@@ -367,7 +385,7 @@ def _open_rendezvous_for_restart(self):
                     self._worker_group.group_rank if self._worker_group else "N/A"
                 )
             except Exception as e:
-                logger.warning(f"Failed to open rendezvous: {e}")
+                logger.error(f"Failed to open rendezvous: {e}")
         # For legacy rendezvous, no action needed - it uses different mechanism
 
     def _invoke_run(self, role: str = DEFAULT_ROLE) -> RunResult:
@@ -420,7 +438,16 @@ def _invoke_run_with_any_failed_policy(self, role: str = DEFAULT_ROLE) -> RunRes
                     rank=self._worker_group.group_rank,
                 )
 
-                if self._remaining_restarts > 0:
+                self._progress_tracker.analyze_previous_cycle()
+                should_terminate_early = self._progress_tracker.should_terminate_early()
+
+                if should_terminate_early:
+                    logger.error(
+                        "[%s] Progress tracker detected no progress across restarts. "
+                        "No more restarts will be attempted.",
+                        role
+                    )
+                elif self._remaining_restarts > 0:
                     logger.info(
                         "[%s] Worker group %s. "
                         "%s/%s attempts left;"
@@ -434,14 +461,13 @@ def _invoke_run_with_any_failed_policy(self, role: str = DEFAULT_ROLE) -> RunRes
                     # Open rendezvous before restarting (for barrier-based rendezvous)
                     self._open_rendezvous_for_restart()
                     self._restart_workers(self._worker_group)
-                else:
-                    self._stop_workers(self._worker_group)
-                    self._worker_group.state = WorkerState.FAILED
-                    # to preserve torchrun's behaviour, should not return WorkerState.UNHEALTHY.
-                    # we use WorkerState.UNHEALTHY to denote a worker group that is still
-                    # running but has some failed workers. torchrun does not use WorkerState.UNHEALTHY
-                    run_result = self._monitor_workers(self._worker_group)
-                    return run_result
+                    continue  # Continue monitoring after restart
+
+                # No more restarts (either exhausted or early termination)
+                self._stop_workers(self._worker_group)
+                self._worker_group.state = WorkerState.FAILED
+                run_result = self._monitor_workers(self._worker_group)
+                return run_result
             elif state == WorkerState.HEALTHY:
                 # Check for cluster-wide issues: unhealthy nodes or new nodes waiting
                 unhealthy_count = self._check_cluster_unhealthy_count()
@@ -579,31 +605,113 @@ def get_rank_mon_socket_path(self, local_rank):
 
     def setup_rank_monitors(self, envs: Dict[int, Dict[str, str]]) -> None:
         fork_mp_ctx = torch.multiprocessing.get_context("fork")
+        new_monitors = []  # Track newly started monitors
+
         for worker_env in envs.values():
             # Start rank monitors if not already started
             # Each rank (re)connects to its rank monitor when it starts
             # Monitor of the local rank0 on the store hosting node is the restarter logger
             local_rank = int(worker_env['LOCAL_RANK'])
             is_restarter_logger = self._is_store_host and local_rank == 0
             rmon_ipc_socket = worker_env[FT_RANK_MONITOR_IPC_SOCKET_ENV_VAR]
-            if local_rank not in self._local_rank_to_rmon:
-                self._local_rank_to_rmon[local_rank] = RankMonitorServer.run_in_subprocess(
+            if local_rank not in self._rank_monitors:
+                rmon_proc = RankMonitorServer.run_in_subprocess(
                     cfg=self._ft_cfg,
                     ipc_socket_path=rmon_ipc_socket,
                     is_restarter_logger=is_restarter_logger,
                     mp_ctx=fork_mp_ctx,
                     env=worker_env,
                 )
+                self._rank_monitors[local_rank] = RankMonitorState(process=rmon_proc)
+                new_monitors.append((local_rank, rmon_proc))
+
+        # Establish bidirectional IPC connections to new rank monitors
+        if new_monitors:
+            async def connect_all():
+                await asyncio.gather(
+                    *[self._connect_to_rank_monitor(lr, rmon) for lr, rmon in new_monitors]
+                )
+            asyncio.run(connect_all())
 
     def shutdown_rank_monitors(self):
-        for local_rank, rmon_proc in self._local_rank_to_rmon.items():
+        # Stop listener tasks, close connections, and send shutdown messages
+        for local_rank, state in self._rank_monitors.items():
+            # Cancel listener task
+            if state.listener_task and not state.listener_task.done():
+                state.listener_task.cancel()
+
+            # Close connection with shutdown message
+            if state.writer:
+                try:
+                    async def send_shutdown():
+                        await write_obj_to_ipc_stream("shutdown", state.writer)
+                        state.writer.close()
+                        await state.writer.wait_closed()
+                    asyncio.run(send_shutdown())
+                except Exception as e:
+                    logger.debug(f"Error closing rank monitor connection for rank {local_rank}: {e}")
+
+        # Terminate rank monitor processes
+        for local_rank, state in self._rank_monitors.items():
             with contextlib.suppress(Exception):
-                rmon_proc.terminate()
+                state.process.terminate()
             with contextlib.suppress(Exception):
-                rmon_proc.join()
+                state.process.join()
             with contextlib.suppress(Exception):
                 os.unlink(self.get_rank_mon_socket_path(local_rank))
 
+    async def _connect_to_rank_monitor(self, local_rank: int, rmon_proc) -> None:
+        """Establish persistent connection to rank monitor for bidirectional IPC.
+
+        Note: This is called after rank_monitor_ready_event is set, which guarantees
+        the socket file already exists.
+        """
+        launcher_to_rmon_socket = f"{tempfile.gettempdir()}/_ft_launcher{rmon_proc.pid}_to_rmon.socket"
+
+        reader, writer = await asyncio.open_unix_connection(launcher_to_rmon_socket)
+        state = self._rank_monitors[local_rank]
+        state.reader = reader
+        state.writer = writer
+        logger.debug(f"Connected to rank monitor {local_rank} at {launcher_to_rmon_socket}")
+
+        # Start listener task for this connection
+        state.listener_task = asyncio.create_task(self._listen_to_rank_monitor(local_rank, reader))
+
+    def _update_progress_iteration(self, local_rank: int, iteration: int):
+        """Update iteration for a specific rank and aggregate using MIN strategy."""
+        # Update this rank's max iteration
+        self._rank_iterations[local_rank] = max(
+            self._rank_iterations.get(local_rank, 0), iteration
+        )
+
+        # Use minimum across all ranks (most conservative - slowest rank determines progress)
+        min_iteration = min(self._rank_iterations.values()) if self._rank_iterations else 0
+        self._progress_tracker.update_iteration(min_iteration)
+
+        logger.debug(
+            f"Updated iteration for rank {local_rank}={iteration}, "
+            f"cluster min={min_iteration}, all_ranks={self._rank_iterations}"
+        )
+
+    async def _listen_to_rank_monitor(self, local_rank: int, reader) -> None:
+        """Listen for messages from rank monitor."""
+        try:
+            while True:
+                msg = await read_obj_from_ipc_stream(reader)
+                if isinstance(msg, dict) and msg.get("type") == "iteration_update":
+                    # Handle iteration update from rank monitor
+                    iteration = msg["iteration"]
+                    self._update_progress_iteration(local_rank, iteration)
+                    logger.debug(f"[Rank {local_rank}] Received iteration update: {iteration}")
+                else:
+                    logger.debug(f"Received message from rank monitor {local_rank}: {msg}")
+        except (asyncio.IncompleteReadError, ConnectionResetError, BrokenPipeError, EOFError):
+            logger.debug(f"Rank monitor {local_rank} connection closed")
+        except asyncio.CancelledError:
+            logger.debug(f"Listener for rank monitor {local_rank} cancelled")
+        except Exception as e:
+            logger.error(f"Error listening to rank monitor {local_rank}: {e}")
+
     def _setup_local_watchdog(self, envs: Dict[int, Dict[str, str]]) -> None:
         enable_watchdog_env_name = TORCHELASTIC_ENABLE_FILE_TIMER
         watchdog_enabled = os.getenv(enable_watchdog_env_name)
@@ -682,22 +790,27 @@ def _stop_workers(self, worker_group: WorkerGroup, *args, **kwargs) -> None:
 
         logger.info(f"Stopping workers... Timeout = {self._workers_stop_timeout} sec.")
 
-        # Send close message to rank monitors
-        for local_rank, rmon_proc in self._local_rank_to_rmon.items():
-            try:
-                launcher_to_rmon_socket = f"{tempfile.gettempdir()}/_ft_launcher{rmon_proc.pid}_to_rmon.socket"
-                if os.path.exists(launcher_to_rmon_socket):
-                    async def send_close_msg():
-                        reader, writer = await asyncio.open_unix_connection(launcher_to_rmon_socket)
-                        try:
-                            await write_obj_to_ipc_stream("close_worker_ipc_connection", writer)
-                        finally:
-                            writer.close()
-                            await writer.wait_closed()
-
-                    asyncio.run(send_close_msg())
-            except Exception as e:
-                logger.warning(f"Failed to send close message to rank monitor {local_rank}: {e}")
+        # Send close message to rank monitors through persistent connections
+        async def send_close_messages():
+            tasks = []
+            for local_rank, state in self._rank_monitors.items():
+                if state.writer:
+                    async def send_msg(writer, local_rank):
+                        await write_obj_to_ipc_stream("close_worker_ipc_connection", writer)
+                    tasks.append(send_msg(state.writer, local_rank))
+            if tasks:
+                # return_exceptions=True catches exceptions from send_msg, no need for try-except inside
+                results = await asyncio.gather(*tasks, return_exceptions=True)
+                for local_rank, result in zip([lr for lr, s in self._rank_monitors.items() if s.writer], results):
+                    if isinstance(result, Exception):
+                        # Connection errors during shutdown are expected (rank monitor may be dead)
+                        if isinstance(result, (ConnectionError, BrokenPipeError, OSError)):
+                            logger.debug(f"Rank monitor {local_rank} already disconnected: {result}")
+                        else:
+                            logger.warning(f"Unexpected error sending close message to rank monitor {local_rank}: {result}")
+
+        if self._rank_monitors:
+            asyncio.run(send_close_messages())
 
         self._shutdown(timeout=self._workers_stop_timeout)
 
@@ -2074,6 +2187,32 @@ def get_args_parser() -> ArgumentParser:
         help="Do not raise an error if there is no Fault Tolerance pkg config provided, just use default settings.",
     )
 
+    #
+    # Progress tracking arguments
+    #
+
+    parser.add_argument(
+        "--ft-max-no-progress-restarts",
+        "--ft-max_no_progress_restarts",
+        type=int,
+        default=3,
+        dest="ft_max_no_progress_restarts",
+        help="Maximum consecutive restarts without progress before early termination. "
+        "Progress tracking is enabled when this value > 0. "
+        "Set to 0 or -1 to disable progress tracking. "
+        "Default: 3 (progress tracking enabled).",
+    )
+
+    parser.add_argument(
+        "--ft-min-progress-iterations",
+        "--ft-min_progress_iterations",
+        type=int,
+        default=200,
+        dest="ft_min_progress_iterations",
+        help="Minimum iterations required to consider a restart as making progress. "
+        "Default: 200.",
+    )
+
     #
     # Positional arguments.
     #