Add CPU-side NCCL H1 probe

Author: lzy-edu

flagscale/runner/backend/backend_megatron.py

@@ -18,11 +18,13 @@
 from omegaconf import DictConfig, OmegaConf
 
 from flagscale.runner.backend.backend_base import BackendBase
+from flagscale.runner.diagnostics import diagnostic_command_body
 from flagscale.runner.heartbeat.config import prepare_heartbeat_launch_config
 from flagscale.runner.runner_train import (
     _get_args_megatron,
     _update_config_train,
 )
+from flagscale.runner.tracing.config import prepare_trace_launch_config
 from flagscale.runner.utils import get_pkg_dir, logger, parse_hostfile, resolve_path
 
 PERF_MONITOR_RUNNER_KEYS = (
@@ -51,6 +53,9 @@ def _prepare(self):
         self.user_args = _get_args_megatron(self.config)
         self.rdzv_id = datetime.now().strftime("%Y%m%d_%H%M%S.%f")
         self.heartbeat_config = prepare_heartbeat_launch_config(self.config, self.rdzv_id)
+        self.trace_config = prepare_trace_launch_config(
+            self.config, self.rdzv_id, self.heartbeat_config
+        )
         self.user_envs = self.config.experiment.get("envs", {})
         self.user_script = self.config.experiment.task.entrypoint
         self.resources = parse_hostfile(self.config.experiment.runner.get("hostfile", None))
@@ -137,6 +142,10 @@ def generate_run_script(
                 f.write(f"{line}\n")
             if self.heartbeat_config.enabled:
                 f.write("\n")
+            for line in self.trace_config.shell_setup_lines(node_rank):
+                f.write(f"{line}\n")
+            if self.trace_config.enabled:
+                f.write("\n")
             f.write(f'cmd="{cmd}"\n')
             f.write("\n")
             if enable_monitoring:
@@ -161,7 +170,9 @@ def generate_run_script(
                 )
             f.write("\n")
 
-            command_body = self.heartbeat_config.training_command_body(node_rank)
+            command_body = diagnostic_command_body(
+                node_rank, self.heartbeat_config, self.trace_config
+            )
             if background:
                 f.write(
                     f'nohup bash -c "{command_body}" >> {host_output_file} 2>&1 & echo $! > {host_pid_file}\n'
@@ -206,6 +217,8 @@ def generate_stop_script(self, host, node_rank):
             f.write("fi\n")
             for line in self.heartbeat_config.stop_shell_lines(node_rank):
                 f.write(f"{line}\n")
+            for line in self.trace_config.stop_shell_lines(node_rank):
+                f.write(f"{line}\n")
             f.write(f"{after_stop}\n")
             f.flush()
             os.fsync(f.fileno())

flagscale/runner/backend/backend_native_train.py

@@ -18,7 +18,10 @@
 from omegaconf import DictConfig, OmegaConf
 
 from flagscale.runner.backend.backend_base import BackendBase
+from flagscale.runner.diagnostics import diagnostic_command_body
+from flagscale.runner.heartbeat.config import prepare_heartbeat_launch_config
 from flagscale.runner.runner_train import _get_args_native, _update_config_train
+from flagscale.runner.tracing.config import prepare_trace_launch_config
 from flagscale.runner.utils import get_pkg_dir, logger, parse_hostfile, resolve_path
 
 
@@ -33,6 +36,10 @@ def _prepare(self):
         _update_config_train(self.config)
         self.user_args = _get_args_native(self.config)
         self.rdzv_id = datetime.now().strftime("%Y%m%d_%H%M%S.%f")
+        self.heartbeat_config = prepare_heartbeat_launch_config(self.config, self.rdzv_id)
+        self.trace_config = prepare_trace_launch_config(
+            self.config, self.rdzv_id, self.heartbeat_config
+        )
         self.user_envs = self.config.experiment.get("envs", {})
         self.user_script = self.config.experiment.task.entrypoint
         self.resources = parse_hostfile(self.config.experiment.runner.get("hostfile", None))
@@ -94,6 +101,14 @@ def generate_run_script(
             f.write("\n")
             f.write(f"export PYTHONPATH={pkg_dir}:{megatron_dir}:${{PYTHONPATH}}\n")
             f.write("\n")
+            for line in self.heartbeat_config.shell_setup_lines(node_rank):
+                f.write(f"{line}\n")
+            if self.heartbeat_config.enabled:
+                f.write("\n")
+            for line in self.trace_config.shell_setup_lines(node_rank):
+                f.write(f"{line}\n")
+            if self.trace_config.enabled:
+                f.write("\n")
             f.write(f'cmd="{cmd}"\n')
             f.write("\n")
             if enable_monitoring:
@@ -118,13 +133,16 @@ def generate_run_script(
                 )
             f.write("\n")
 
+            command_body = diagnostic_command_body(
+                node_rank, self.heartbeat_config, self.trace_config
+            )
             if background:
                 f.write(
-                    f'nohup bash -c "$cmd; sync" >> {host_output_file} 2>&1 & echo $! > {host_pid_file}\n'
+                    f'nohup bash -c "{command_body}" >> {host_output_file} 2>&1 & echo $! > {host_pid_file}\n'
                 )
             else:
                 f.write("set -o pipefail\n")
-                f.write(f'bash -c "$cmd; sync" 2>&1 | tee -a {host_output_file}\n')
+                f.write(f'bash -c "{command_body}" 2>&1 | tee -a {host_output_file}\n')
             f.write("\n")
             f.flush()
             os.fsync(f.fileno())
@@ -160,6 +178,10 @@ def generate_stop_script(self, host, node_rank):
             # TODO: This is a temporary fix. We need to find a better way to stop the job.
             f.write("    pkill -f 'torchrun'\n")
             f.write("fi\n")
+            for line in self.heartbeat_config.stop_shell_lines(node_rank):
+                f.write(f"{line}\n")
+            for line in self.trace_config.stop_shell_lines(node_rank):
+                f.write(f"{line}\n")
             f.write(f"{after_stop}\n")
             f.flush()
             os.fsync(f.fileno())

flagscale/runner/diagnostics.py

@@ -0,0 +1,39 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Shared shell lifecycle helpers for opt-in runner diagnostics."""
+
+from __future__ import annotations
+
+from typing import Protocol
+
+
+class DiagnosticLaunchConfig(Protocol):
+    enabled: bool
+
+    def command_exit_actions(self, node_rank: int) -> list[str]: ...
+
+
+def diagnostic_command_body(node_rank: int, *configs: DiagnosticLaunchConfig) -> str:
+    """Write every enabled monitor's completion marker without nesting shells."""
+
+    actions = [
+        action
+        for config in configs
+        if config.enabled
+        for action in config.command_exit_actions(node_rank)
+    ]
+    if not actions:
+        return "$cmd; sync"
+    return f"$cmd; rc=\\$?; {'; '.join(dict.fromkeys(actions))}; sync; exit \\$rc"

flagscale/runner/heartbeat/config.py

@@ -99,8 +99,15 @@ def hardware_health_log_file(self, node_rank: int) -> str:
 
     def training_command_body(self, node_rank: int) -> str:
         """Run training and notify the node-zero heartbeat monitor on exit."""
-        if not self.enabled:
+        exit_actions = self.command_exit_actions(node_rank)
+        if not exit_actions:
             return "$cmd; sync"
+        return "$cmd; rc=\\$?; " + "; ".join(exit_actions) + "; sync; exit \\$rc"
+
+    def command_exit_actions(self, node_rank: int) -> list[str]:
+        """Return cleanup actions for a shared diagnostic command wrapper."""
+        if not self.enabled:
+            return []
         exit_actions: list[str] = []
         if node_rank == 0:
             completion_file = shlex.quote(self.completion_file)
@@ -111,9 +118,7 @@ def training_command_body(self, node_rank: int) -> str:
                 f"if [ -f {health_pid_file} ]; then "
                 f'kill \\"\\$(cat {health_pid_file})\\" 2>/dev/null || true; fi'
             )
-        if not exit_actions:
-            return "$cmd; sync"
-        return "$cmd; rc=\\$?; " + "; ".join(exit_actions) + "; sync; exit \\$rc"
+        return exit_actions
 
     def shell_setup_lines(self, node_rank: int) -> list[str]:
         if not self.enabled:

flagscale/runner/launcher/launcher_ssh.py

@@ -148,9 +148,11 @@ def _get_runner_cmd_train(
         del runner_args["perf_max_log_files"]
     if "perf_model_type" in runner_args:
         del runner_args["perf_model_type"]
-    # Heartbeat is consumed by FlagScale and is not a torchrun option.
+    # Diagnostics are consumed by FlagScale and are not torchrun options.
     if "heartbeat" in runner_args:
         del runner_args["heartbeat"]
+    if "tracing" in runner_args:
+        del runner_args["tracing"]
     runner_args["rdzv_id"] = rdzv_id
     # runner_args["master_addr"] = master_addr
     # runner_args["master_port"] = master_port

flagscale/runner/runner_train.py

@@ -21,12 +21,14 @@
 from hydra.core.hydra_config import HydraConfig
 from omegaconf import DictConfig, OmegaConf
 
+from flagscale.runner.diagnostics import diagnostic_command_body
 from flagscale.runner.elastic.monitor_service import MonitorService
 from flagscale.runner.heartbeat.config import (
     HeartbeatLaunchConfig,
     prepare_heartbeat_launch_config,
 )
 from flagscale.runner.runner_base_legacy import JobStatus, RunnerBase
+from flagscale.runner.tracing.config import TraceLaunchConfig, prepare_trace_launch_config
 from flagscale.runner.utils import (
     find_latest_stdout_log,
     flatten_dict_to_args,
@@ -212,6 +214,8 @@ def _get_runner_cmd_train(
         del runner_args["enable_monitoring"]
     if "heartbeat" in runner_args:
         del runner_args["heartbeat"]
+    if "tracing" in runner_args:
+        del runner_args["tracing"]
     runner_args["rdzv_id"] = rdzv_id
     # runner_args["master_addr"] = master_addr
     # runner_args["master_port"] = master_port
@@ -245,8 +249,10 @@ def _generate_run_script_train(
     pkg_dir=None,
     enable_monitoring=False,
     heartbeat_config=None,
+    trace_config=None,
 ):
     heartbeat_config = heartbeat_config or HeartbeatLaunchConfig(enabled=False)
+    trace_config = trace_config or TraceLaunchConfig(enabled=False)
     system_config = config.train.system
     logging_config = config.train.system.logging
 
@@ -293,6 +299,10 @@ def _generate_run_script_train(
             f.write(f"{line}\n")
         if heartbeat_config.enabled:
             f.write("\n")
+        for line in trace_config.shell_setup_lines(node_rank):
+            f.write(f"{line}\n")
+        if trace_config.enabled:
+            f.write("\n")
         f.write(f'cmd="{cmd}"\n')
         f.write("\n")
         if enable_monitoring:
@@ -315,7 +325,7 @@ def _generate_run_script_train(
             f.write(f'echo "Monitor service started in background for {host} (node {node_rank})"\n')
         f.write("\n")
 
-        command_body = heartbeat_config.training_command_body(node_rank)
+        command_body = diagnostic_command_body(node_rank, heartbeat_config, trace_config)
         if background:
             f.write(
                 f'nohup bash -c "{command_body}" >> {host_output_file} 2>&1 & echo $! > {host_pid_file}\n'
@@ -331,8 +341,9 @@ def _generate_run_script_train(
     return host_run_script_file
 
 
-def _generate_stop_script_train(config, host, node_rank, heartbeat_config=None):
+def _generate_stop_script_train(config, host, node_rank, heartbeat_config=None, trace_config=None):
     heartbeat_config = heartbeat_config or HeartbeatLaunchConfig(enabled=False)
+    trace_config = trace_config or TraceLaunchConfig(enabled=False)
     if getattr(config, "train", None):
         logging_config = config.train.system.logging
     else:
@@ -362,6 +373,8 @@ def _generate_stop_script_train(config, host, node_rank, heartbeat_config=None):
         f.write("fi\n")
         for line in heartbeat_config.stop_shell_lines(node_rank):
             f.write(f"{line}\n")
+        for line in trace_config.stop_shell_lines(node_rank):
+            f.write(f"{line}\n")
         f.write(f"{after_stop}\n")
         f.flush()
         os.fsync(f.fileno())
@@ -426,6 +439,9 @@ def _prepare(self):
             raise ValueError(f"Unsupported backend: {self.config.experiment.task.backend}")
         self.rdzv_id = datetime.now().strftime("%Y%m%d_%H%M%S.%f")
         self.heartbeat_config = prepare_heartbeat_launch_config(self.config, self.rdzv_id)
+        self.trace_config = prepare_trace_launch_config(
+            self.config, self.rdzv_id, self.heartbeat_config
+        )
         self.user_envs = self.config.experiment.get("envs", {})
         self.user_script = self.config.experiment.task.entrypoint
         self.resources = parse_hostfile(self.config.experiment.runner.get("hostfile", None))
@@ -484,6 +500,7 @@ def _run_each(
             pkg_dir=node_specific_config.get("build_dir", None),
             enable_monitoring=enable_monitoring,
             heartbeat_config=self.heartbeat_config,
+            trace_config=self.trace_config,
         )
 
         if host != "localhost":
@@ -616,6 +633,7 @@ def _stop_each(self, host, node_rank):
             host,
             node_rank,
             self.heartbeat_config,
+            self.trace_config,
         )
         logging_config = self.config.train.system.logging
 
@@ -889,6 +907,9 @@ def _prepare(self):
         _update_config_train(self.config)
         self.rdzv_id = datetime.now().strftime("%Y%m%d_%H%M%S.%f")
         self.heartbeat_config = prepare_heartbeat_launch_config(self.config, self.rdzv_id)
+        self.trace_config = prepare_trace_launch_config(
+            self.config, self.rdzv_id, self.heartbeat_config
+        )
         if self.config.experiment.task.backend == "megatron":
             self.user_args = _get_args_megatron(self.config)
         logger.info("\n************** configuration ***********")
@@ -922,6 +943,7 @@ def _run_each(
             cmd,
             background=background,
             heartbeat_config=self.heartbeat_config,
+            trace_config=self.trace_config,
         )
 
         run_local_command(f"bash {host_run_script_file}", dryrun)

flagscale/runner/tracing/__init__.py

@@ -0,0 +1,22 @@
+# Copyright 2026 FlagOS Contributors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CPU-side NCCL tracing and hang detection utilities."""
+
+from .analyzer import Finding, TraceAnalyzer
+
+__all__ = [
+    "Finding",
+    "TraceAnalyzer",
+]