Merge pull request #1759 from NVIDIA-NeMo/ko3n1g/chore/test-ft

Author: ko3n1g

scripts/performance/resiliency_plugins.py

@@ -0,0 +1,88 @@
+import logging
+from dataclasses import dataclass
+from typing import Callable, List, Optional, Union
+
+import nemo_run as run
+from nemo_run import Plugin
+
+
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class FaultTolerancePluginScriptArgs:
+    """Arguments for FaultTolerancePlugin to pass to run.Script."""
+
+    enable_ft_package: bool
+    calc_ft_timeouts: bool
+
+
+def _default_fault_tolerance_converter(args: FaultTolerancePluginScriptArgs) -> List[str]:
+    """Default converter for FaultTolerancePlugin that generates hydra-style overrides."""
+    return [
+        f"ft.enable_ft_package={str(args.enable_ft_package).lower()}",
+        f"ft.calc_ft_timeouts={str(args.calc_ft_timeouts).lower()}",
+    ]
+
+
+@dataclass(kw_only=True)
+class FaultTolerancePlugin(Plugin):
+    """
+    A plugin for setting up fault tolerance configuration.
+    This plugin enables workload hang detection, automatic calculation of timeouts used for hang detection,
+    detection of rank(s) terminated due to an error and workload respawning in case of a failure.
+
+
+    Args:
+        enable_ft_package (bool): Enable the fault tolerance package. Default is True.
+        calc_ft_timeouts (bool): Automatically compute timeouts. Default is True.
+        num_in_job_restarts (int): Max number of restarts on failure, within the same job. Default is 3.
+        num_job_retries_on_failure (int): Max number of new job restarts on failure. Default is 2.
+        initial_rank_heartbeat_timeout (int): Timeouts are time intervals used by a rank monitor to detect
+            that a rank is not alive. This is the max timeout for the initial heartbeat. Default is 1800.
+        rank_heartbeat_timeout (int): This is the timeout for subsequent hearbeats after the initial heartbeat.
+            Default is 300.
+        script_args_converter_fn (Optional[Callable]): A function that takes FaultTolerancePluginScriptArgs
+                                                        and returns a list of CLI arguments. If not provided,
+                                                        uses the default hydra-style converter.
+
+    Note:
+        This plugin is incompatible with NsysPlugin. Nsys profiling cannot be used when fault tolerance
+        is enabled.
+    """
+
+    enable_ft_package: bool = True
+    calc_ft_timeouts: bool = True
+    num_in_job_restarts: int = 3
+    num_job_retries_on_failure: int = 2
+    initial_rank_heartbeat_timeout: int = 1800
+    rank_heartbeat_timeout: int = 300
+    script_args_converter_fn: Optional[Callable[[FaultTolerancePluginScriptArgs], List[str]]] = None
+
+    def setup(self, task: Union["run.Partial", "run.Script"], executor: "run.Executor"):
+        """Set up the fault tolerance plugin."""
+        # Set up fault tolerance launcher for both task types
+        executor.launcher = run.FaultTolerance(
+            max_restarts=self.num_in_job_restarts,
+            initial_rank_heartbeat_timeout=self.initial_rank_heartbeat_timeout,
+            rank_heartbeat_timeout=self.rank_heartbeat_timeout,
+        )
+        executor.retries = self.num_job_retries_on_failure
+
+        if isinstance(task, run.Script):
+            # For run.Script, append CLI overrides to the script arguments
+            # Create args dataclass
+            script_args = FaultTolerancePluginScriptArgs(
+                enable_ft_package=self.enable_ft_package,
+                calc_ft_timeouts=self.calc_ft_timeouts,
+            )
+
+            # Use custom converter or default
+            converter = self.script_args_converter_fn or _default_fault_tolerance_converter
+            cli_overrides = converter(script_args)
+
+            task.args.extend(cli_overrides)
+            logger.info(f"{self.__class__.__name__} added CLI overrides: {', '.join(cli_overrides)}")
+        else:
+            raise NotImplementedError("FaultTolerancePlugin is only supported for run.Script tasks")

scripts/performance/setup_experiment.py

@@ -22,6 +22,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional
 
+import nemo_run as run
 from nemo_run.config import get_nemorun_home
 
 
@@ -34,9 +35,6 @@
     from .utils.evaluate import calc_convergence_and_performance
     from .utils.executors import dgxc_executor, slurm_executor
 
-import nemo_run as run
-
-
 try:
     import wandb
 
@@ -46,8 +44,10 @@
 
 try:
     from perf_plugins import NsysPlugin, PerfEnvPlugin
+    from resiliency_plugins import FaultTolerancePlugin
 except (ImportError, ModuleNotFoundError):
     from .perf_plugins import NsysPlugin, PerfEnvPlugin
+    from .resiliency_plugins import FaultTolerancePlugin
 
 import logging
 
@@ -327,6 +327,18 @@ def main(
             )
         )
 
+    if use_recipes:
+        plugins.append(
+            FaultTolerancePlugin(
+                enable_ft_package=True,
+                calc_ft_timeouts=True,
+                num_in_job_restarts=10,
+                num_job_retries_on_failure=10,
+                initial_rank_heartbeat_timeout=1800,
+                rank_heartbeat_timeout=300,
+            )
+        )
+
     nemorun_script = run.Script(
         path=str(run_script_path),
         entrypoint="python",

scripts/performance/utils/executors.py

@@ -223,6 +223,6 @@ def dgxc_executor(
             else {}
         ),
         env_vars=env_vars,
-        launcher="ft",
+        launcher="torchrun",
     )
     return executor