Merge pull request NVIDIA#246 from hexinw-nvidia/signal

fix: Handle signals during rendezvous

Author: hexinw-nvidia

src/nvidia_resiliency_ext/fault_tolerance/config.py

@@ -91,6 +91,10 @@ class FaultToleranceConfig:
     * `install_exception_hook` [bool] if True, installs sys.excepthook to capture uncaught exceptions
       in training worker processes, format and log the traceback, and use os._exit() to exit the
       process reliably. Default: False.
+    * `num_warmup_iterations` [int] number of warmup iterations before monitoring step section and
+      out-of-section timeouts. The first N iterations (relative to cycle start) are excluded from
+      timeout monitoring as they can be significantly slower than steady-state iterations.
+      Default: 5. Can be overridden by workload (e.g., Megatron-LM via init_workload_monitoring).
     * Attribution service (optional):
       - `attrsvc_host` [str] hostname/IP of the attribution service
       - `attrsvc_port` [int] port of the attribution service
@@ -128,6 +132,9 @@ class FaultToleranceConfig:
     min_progress_iterations: int = 200
     progress_update_interval: float = 30.0  # Seconds between sending progress updates to launcher
     install_exception_hook: bool = False
+    num_warmup_iterations: int = (
+        5  # Number of warmup iterations before monitoring step section and out-of-section timeouts
+    )
     # Attribution service configuration (optional)
     attrsvc_host: Optional[str] = None
     attrsvc_port: Optional[int] = None

src/nvidia_resiliency_ext/fault_tolerance/data.py

@@ -56,6 +56,7 @@ def get_for_current_rank():
                 "Could not find the rank of the current process. "
                 "Is it a part of a distributed workload?"
             )
+        global_rank = int(global_rank)
         local_rank = int(os.environ.get("LOCAL_RANK", -1))
         host = socket.gethostname()
         pid = os.getpid()
@@ -159,11 +160,19 @@ class InitMsg:
         iteration: Current training iteration if available from workload framework.
                   If None, indicates that the workload cannot report iterations,
                   and progress tracking should remain disabled.
+        num_warmup_iters: Number of warmup iterations before monitoring step section
+                         and out-of-section timeouts. If None, server uses default from config.
     """
 
-    def __init__(self, rank_info=None, iteration: Optional[int] = None):
+    def __init__(
+        self,
+        rank_info=None,
+        iteration: Optional[int] = None,
+        num_warmup_iters: Optional[int] = None,
+    ):
         self.rank_info = rank_info
         self.iteration = iteration
+        self.num_warmup_iters = num_warmup_iters
 
 
 class HeartbeatMsg:

src/nvidia_resiliency_ext/fault_tolerance/ft_rendezvous_barrier.py

@@ -14,17 +14,20 @@
 import json
 import logging
 import os
+import signal
 import socket
 import threading
 import time
 from collections import defaultdict
 from dataclasses import dataclass
 from datetime import timedelta
 from enum import Enum
+from types import FrameType
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 from torch.distributed import PrefixStore, Store
 from torch.distributed.elastic.events import NodeState, construct_and_record_rdzv_event
+from torch.distributed.elastic.multiprocessing import SignalException
 from torch.distributed.elastic.rendezvous.api import (
     RendezvousClosedError,
     RendezvousError,
@@ -67,6 +70,24 @@
 log = logging.getLogger(LogConfig.name)
 
 
+def _rdzv_signal_exception_handler(sig: int, frame: Optional[FrameType]) -> None:
+    del frame
+    raise SignalException(f"Received signal {sig} during rendezvous", signal.Signals(sig))
+
+
+def _install_rdzv_signal_handlers() -> Dict[signal.Signals, Any]:
+    prev_handlers: Dict[signal.Signals, Any] = {}
+    for sig_to_handle in (signal.SIGTERM, signal.SIGINT):
+        prev_handlers[sig_to_handle] = signal.getsignal(sig_to_handle)
+        signal.signal(sig_to_handle, _rdzv_signal_exception_handler)
+    return prev_handlers
+
+
+def _restore_rdzv_signal_handlers(prev_handlers: Dict[signal.Signals, Any]) -> None:
+    for sig_to_handle, handler in prev_handlers.items():
+        signal.signal(sig_to_handle, handler)
+
+
 def get_method_name(depth=2):
     if len(inspect.stack()) > depth:
         return inspect.stack()[depth].function
@@ -853,6 +874,13 @@ def perform_rendezvous(
         # Start timing AFTER Step 0 completes, since hot spares may wait indefinitely at Step 0
         self._rendezvous_start_time = time.monotonic()
 
+        # Record rendezvous start event - start profiling AFTER waiting for rendezvous to open
+        # This ensures hot spares waiting at Step 0 don't skew the rendezvous performance measurement
+        rendezvous_start_event_id = record_profiling_event(
+            ProfilingEvent.RENDEZVOUS_STARTED,
+            node_id=node_desc,
+        )
+
         # Step 1: Join the rendezvous and get unique identifier
         self._arrived_count = self.store.add(self.arrived_count_key, 1)
 
@@ -1699,12 +1727,7 @@ def next_rendezvous(self) -> Union[RendezvousInfo, Tuple[Store, int, int]]:
         self._record(message=msg)
         log.info(msg)
 
-        # Record rendezvous start event
-        rendezvous_start_event_id = record_profiling_event(
-            ProfilingEvent.RENDEZVOUS_STARTED,
-            node_id=self._this_node,
-        )
-
+        prev_signal_handlers = _install_rdzv_signal_handlers()
         try:
             # Check node health and control requests before starting rendezvous
             self.ensure_node_is_healthy()
@@ -1738,6 +1761,8 @@ def next_rendezvous(self) -> Union[RendezvousInfo, Tuple[Store, int, int]]:
                 node_state=NodeState.FAILED,
             )
             raise
+        finally:
+            _restore_rdzv_signal_handlers(prev_signal_handlers)
 
         msg = (
             f"The node '{self._this_node}' has joined the rendezvous "

src/nvidia_resiliency_ext/fault_tolerance/launcher.py

@@ -70,14 +70,12 @@
     FT_LAUNCHER_IPC_SOCKET_ENV_VAR,
     FT_RANK_MONITOR_IPC_SOCKET_ENV_VAR,
 )
-from nvidia_resiliency_ext.fault_tolerance.per_cycle_logs import (
-    PerCycleLogsSpecs,
-    PipeBasedLogsSpecs,
-)
+from nvidia_resiliency_ext.fault_tolerance.per_cycle_logs import PipeBasedLogsSpecs
 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 (
     get_processes_by_pgids,
+    is_slurm_job_array,
     patched_method,
     read_obj_from_ipc_stream,
     terminate_mp_processes,
@@ -1657,6 +1655,17 @@ def launch_agent(
         shutdown_rdzv = False
         logger.error(f"Agent .run() raised UnhealthyNodeException: {e}")
         events.record(agent.get_event_failed())
+
+        # Exit behavior depends on deployment mode:
+        # - Job array: raise (exit 1) so replacement job can be launched
+        # - Single job with hot spares: don't raise (instead, exit 0) to avoid killing job
+        #   since --kill-on-bad-exit is the default srun behavior
+        if is_slurm_job_array():
+            logger.info("Job array deployment: exiting with code 1 for replacement.")
+            raise
+        else:
+            logger.info("Single job deployment: exiting with code 0 for hot spare takeover.")
+            # Don't raise - returns None, main() will exit with 0
     except ChildFailedError:
         raise
     except SignalException as e:
@@ -1672,7 +1681,7 @@ def launch_agent(
         else:
             logger.info("All ranks exited gracefully. Launcher exiting without an error.")
     except Exception as e:
-        logger.error(f"Agent .run() raised exception, {e=}", exc_info=True)
+        logger.error(f"Agent .run() raised exception, {e=}")
         events.record(agent.get_event_failed())
         raise
     finally:
@@ -2327,8 +2336,8 @@ def get_args_parser() -> ArgumentParser:
         type=str,
         help="Logging behavior configuration. Options: "
         "(1) None (default): Creates separate log files per rank per restart cycle. "
-        "(2) 'per_cycle': Consolidates all ranks' logs into a single log file per restart cycle. "
-        "(3) Custom entrypoint name from torchrun.logs_specs group for advanced customization.",
+        "(2) Custom entrypoint name from torchrun.logs_specs group for advanced customization. "
+        "Note: For consolidated logging, use --ft-base-logfile instead (automatically uses PipeBasedLogsSpecs).",
     )
 
     #
@@ -2783,14 +2792,14 @@ def _get_logs_specs_class(logs_specs_name: Optional[str]) -> Type[LogsSpecs]:
 
     Built-in options:
     - None (default): Uses DefaultLogsSpecs (per-rank log files per cycle)
-    - 'per_cycle': Uses PerCycleLogsSpecs (single log file per cycle for all ranks)
+
+    Note: The legacy 'per_cycle' option has been removed. Use --ft-base-logfile instead,
+    which automatically uses PipeBasedLogsSpecs for consolidated logging.
     """
     logs_specs_cls = None
 
-    # Handle built-in per_cycle option
-    if logs_specs_name == "per_cycle":
-        logs_specs_cls = PerCycleLogsSpecs
-    elif logs_specs_name is not None:
+    # Try to load from entrypoints
+    if logs_specs_name is not None:
         # Try to load from entrypoints
         eps = metadata.entry_points()
         if hasattr(eps, "select"):  # >= 3.10