Make H1 checkpoint phase aware

Author: lzy-edu

flagscale/runner/tracing/analyzer.py

@@ -67,6 +67,7 @@ class _CollectiveRound:
     expected_nranks: int
     first_seen_monotonic_s: float
     enters: dict[int, dict[str, Any]] = field(default_factory=dict)
+    detection_phase: str = "unknown"
 
 
 @dataclass
@@ -93,6 +94,7 @@ def __init__(
         heartbeat_timeout_s: float = 30.0,
         collective_timeout_s: float = 60.0,
         delayed_enter_threshold_s: float = 30.0,
+        checkpoint_timeout_s: float = 1800.0,
         p2p_timeout_s: float = 60.0,
         p2p_match_window_s: float = 30.0,
         detect_heartbeat_timeouts: bool = False,
@@ -106,6 +108,7 @@ def __init__(
         )
         self.collective_timeout_s = float(collective_timeout_s)
         self.delayed_enter_threshold_s = float(delayed_enter_threshold_s)
+        self.checkpoint_timeout_s = float(checkpoint_timeout_s)
         self.p2p_timeout_s = float(p2p_timeout_s)
         self.p2p_match_window_s = float(p2p_match_window_s)
         self.detect_heartbeat_timeouts = detect_heartbeat_timeouts
@@ -243,6 +246,8 @@ def scan(
             )
         completed_rounds: list[tuple[str, str, int]] = []
         for key, collective in list(self._rounds.items()):
+            detection_phase = self._collective_detection_phase(collective)
+            checkpointing = detection_phase == "checkpointing"
             mismatch = self._detect_signature_mismatch(collective, now_unix_ns)
             if mismatch is not None:
                 self._emit_once(("collective_signature_mismatch", *key), mismatch, findings)
@@ -255,19 +260,26 @@ def scan(
                 # never entered. Keep evaluating the incomplete round for H1.
 
             if len(collective.enters) >= collective.expected_nranks:
-                delayed = self._detect_delayed_enter(collective, now_unix_ns)
+                delayed = self._detect_delayed_enter(
+                    collective,
+                    now_unix_ns,
+                    detection_phase=detection_phase,
+                )
                 self._emit_once(("delayed_collective_enter", *key), delayed, findings)
                 completed_rounds.append(key)
                 continue
 
             elapsed_s = now_monotonic_s - collective.first_seen_monotonic_s
-            if elapsed_s < self.collective_timeout_s:
+            timeout_s = self.checkpoint_timeout_s if checkpointing else self.collective_timeout_s
+            if elapsed_s < timeout_s:
                 continue
             missing = self._detect_missing_enter(
                 collective,
                 now_monotonic_s=now_monotonic_s,
                 now_unix_ns=now_unix_ns,
                 elapsed_s=elapsed_s,
+                detection_phase=detection_phase,
+                timeout_s=timeout_s,
             )
             self._emit_once(("collective_missing_enter", *key), missing, findings)
 
@@ -760,7 +772,11 @@ def _detect_signature_mismatch(
         )
 
     def _detect_delayed_enter(
-        self, collective: _CollectiveRound, now_unix_ns: int
+        self,
+        collective: _CollectiveRound,
+        now_unix_ns: int,
+        *,
+        detection_phase: str,
     ) -> Finding | None:
         if not collective.enters:
             return None
@@ -773,12 +789,17 @@ def _detect_delayed_enter(
         earliest = min(timestamps.values())
         latest = max(timestamps.values())
         spread_s = (latest - earliest) / 1_000_000_000
-        if spread_s <= self.delayed_enter_threshold_s:
+        threshold_s = (
+            self.checkpoint_timeout_s
+            if detection_phase == "checkpointing"
+            else self.delayed_enter_threshold_s
+        )
+        if spread_s <= threshold_s:
             return None
         slow_ranks = sorted(
             rank
             for rank, timestamp in timestamps.items()
-            if (timestamp - earliest) / 1_000_000_000 > self.delayed_enter_threshold_s
+            if (timestamp - earliest) / 1_000_000_000 > threshold_s
         )
         return Finding(
             hang_type="delayed_collective_enter",
@@ -791,6 +812,10 @@ def _detect_delayed_enter(
                 "api": next(iter(collective.enters.values())).get("api"),
                 "enter_spread_s": spread_s,
                 "slow_comm_ranks": slow_ranks,
+                "detection_phase": detection_phase,
+                "detection_threshold_s": threshold_s,
+                "threshold_reason": self._threshold_reason(detection_phase),
+                "reason": "collective_enter_spread_exceeded_threshold",
                 "clock_assumption": "hosts have synchronized wall clocks",
                 "confidence": "suspected",
             },
@@ -803,6 +828,8 @@ def _detect_missing_enter(
         now_monotonic_s: float,
         now_unix_ns: int,
         elapsed_s: float,
+        detection_phase: str,
+        timeout_s: float,
     ) -> Finding:
         entered = sorted(collective.enters)
         missing = sorted(set(range(collective.expected_nranks)) - set(entered))
@@ -828,6 +855,7 @@ def _detect_missing_enter(
                     {
                         "heartbeat": state,
                         "heartbeat_age_s": max(0.0, age_s),
+                        "phase": str(heartbeat.event.get("phase") or "unknown"),
                     }
                 )
                 known_states.append(state)
@@ -861,12 +889,49 @@ def _detect_missing_enter(
                 "missing_comm_ranks": missing,
                 "missing_rank_status": rank_status,
                 "waited_s": elapsed_s,
+                "detection_phase": detection_phase,
+                "detection_threshold_s": timeout_s,
+                "threshold_reason": self._threshold_reason(detection_phase),
                 "reason": reason,
                 "confidence": confidence,
                 "trace_event_loss_possible": trace_event_loss_possible,
             },
         )
 
+    def _collective_detection_phase(self, collective: _CollectiveRound) -> str:
+        """Return a stable phase for one collective round.
+
+        A checkpoint can block the heartbeat publisher, so a round that has once
+        been correlated with ``checkpointing`` keeps that phase until it resolves.
+        """
+        if collective.detection_phase == "checkpointing":
+            return collective.detection_phase
+
+        global_ranks = {
+            _as_int(event.get("rank"), default=-1) for event in collective.enters.values()
+        }
+        global_ranks.update(
+            self._comm_members.get((collective.run_id, collective.comm_uid_hash), {}).values()
+        )
+        phases = {
+            str(heartbeat.event.get("phase") or "unknown")
+            for rank in global_ranks
+            if rank >= 0 and (heartbeat := self._heartbeats.get(rank)) is not None
+        }
+        if "checkpointing" in phases:
+            collective.detection_phase = "checkpointing"
+        elif len(phases) == 1:
+            collective.detection_phase = next(iter(phases))
+        elif phases:
+            collective.detection_phase = "mixed"
+        return collective.detection_phase
+
+    @staticmethod
+    def _threshold_reason(detection_phase: str) -> str:
+        if detection_phase == "checkpointing":
+            return "heartbeat_phase_checkpointing"
+        return "normal_collective_phase"
+
     def _emit_once(
         self,
         dedupe_key: tuple[Any, ...],

flagscale/runner/tracing/config.py

@@ -61,6 +61,7 @@ class TraceLaunchConfig:
     heartbeat_timeout_s: float = 30.0
     collective_timeout_s: float = 60.0
     delayed_enter_threshold_s: float = 30.0
+    checkpoint_timeout_s: float = 1800.0
     p2p_timeout_s: float = 60.0
     p2p_match_window_s: float = 30.0
     failure_grace_period_s: float = 60.0
@@ -121,6 +122,8 @@ def shell_setup_lines(self, node_rank: int) -> list[str]:
                 f"{self.collective_timeout_s:g}",
                 "--delayed-enter-threshold",
                 f"{self.delayed_enter_threshold_s:g}",
+                "--checkpoint-timeout",
+                f"{self.checkpoint_timeout_s:g}",
                 "--p2p-timeout",
                 f"{self.p2p_timeout_s:g}",
                 "--p2p-match-window",
@@ -246,6 +249,19 @@ def prepare_trace_launch_config(
     collective_timeout_s = _positive_float(
         raw_dict.get("collective_timeout_s", 60.0), "collective_timeout_s"
     )
+    delayed_enter_threshold_s = _positive_float(
+        raw_dict.get("delayed_enter_threshold_s", 30.0),
+        "delayed_enter_threshold_s",
+    )
+    inherited_checkpoint_timeout_s = (
+        getattr(heartbeat_config, "checkpoint_timeout_s", 1800.0)
+        if heartbeat_config is not None and heartbeat_config.enabled
+        else 1800.0
+    )
+    checkpoint_timeout_s = _positive_float(
+        raw_dict.get("checkpoint_timeout_s", inherited_checkpoint_timeout_s),
+        "checkpoint_timeout_s",
+    )
     p2p_timeout_s = _positive_float(
         raw_dict.get("p2p_timeout_s", collective_timeout_s), "p2p_timeout_s"
     )
@@ -280,10 +296,8 @@ def prepare_trace_launch_config(
             else 30.0
         ),
         collective_timeout_s=collective_timeout_s,
-        delayed_enter_threshold_s=_positive_float(
-            raw_dict.get("delayed_enter_threshold_s", 30.0),
-            "delayed_enter_threshold_s",
-        ),
+        delayed_enter_threshold_s=delayed_enter_threshold_s,
+        checkpoint_timeout_s=checkpoint_timeout_s,
         p2p_timeout_s=p2p_timeout_s,
         p2p_match_window_s=p2p_match_window_s,
         failure_grace_period_s=failure_grace_period_s,
@@ -301,5 +315,13 @@ def prepare_trace_launch_config(
         raise ValueError(
             "tracing.failure_grace_period_s must be at least the largest enabled detection timeout"
         )
+    if resolved.checkpoint_timeout_s < max(
+        resolved.collective_timeout_s,
+        resolved.delayed_enter_threshold_s,
+    ):
+        raise ValueError(
+            "tracing.checkpoint_timeout_s must be greater than or equal to the normal "
+            "collective thresholds"
+        )
 
     return resolved

flagscale/runner/tracing/monitor.py

@@ -110,6 +110,7 @@ def run_monitor(args: argparse.Namespace) -> int:
         heartbeat_timeout_s=args.heartbeat_timeout,
         collective_timeout_s=args.collective_timeout,
         delayed_enter_threshold_s=args.delayed_enter_threshold,
+        checkpoint_timeout_s=args.checkpoint_timeout,
         p2p_timeout_s=args.p2p_timeout,
         p2p_match_window_s=args.p2p_match_window,
     )
@@ -190,6 +191,7 @@ def build_parser() -> argparse.ArgumentParser:
     parser.add_argument("--heartbeat-timeout", type=float, default=30.0)
     parser.add_argument("--collective-timeout", type=float, default=60.0)
     parser.add_argument("--delayed-enter-threshold", type=float, default=30.0)
+    parser.add_argument("--checkpoint-timeout", type=float, default=1800.0)
     parser.add_argument("--p2p-timeout", type=float, default=60.0)
     parser.add_argument("--p2p-match-window", type=float, default=30.0)
     parser.add_argument("--failure-grace-period", type=float, default=60.0)
@@ -210,6 +212,7 @@ def main() -> int:
         "heartbeat_timeout",
         "collective_timeout",
         "delayed_enter_threshold",
+        "checkpoint_timeout",
         "p2p_timeout",
         "p2p_match_window",
         "failure_grace_period",

tests/unit_tests/runner/tracing/test_analyzer.py

@@ -185,6 +185,84 @@ def test_delayed_enter_is_reported_after_every_rank_enters():
     assert findings[0].details["enter_spread_s"] == 7.0
 
 
+def test_checkpoint_missing_enter_uses_checkpoint_timeout_and_still_reports_hang():
+    analyzer = TraceAnalyzer(
+        run_id=RUN_ID,
+        collective_timeout_s=5,
+        checkpoint_timeout_s=20,
+    )
+    _ingest(analyzer, _event("comm_init", rank=0, comm_rank=0))
+    _ingest(analyzer, _event("comm_init", rank=1, comm_rank=1))
+    _ingest(
+        analyzer,
+        _event("heartbeat", rank=1, phase="checkpointing"),
+        observed=1.0,
+    )
+    _ingest(analyzer, _event("nccl_call", rank=0, comm_rank=0), observed=1.0)
+
+    assert analyzer.scan(now_monotonic_s=8.0, now_unix_ns=8_000_000_000) == []
+
+    findings = analyzer.scan(now_monotonic_s=22.0, now_unix_ns=22_000_000_000)
+    assert [finding.hang_type for finding in findings] == ["collective_missing_enter"]
+    assert findings[0].details["detection_phase"] == "checkpointing"
+    assert findings[0].details["detection_threshold_s"] == 20
+    assert findings[0].details["threshold_reason"] == "heartbeat_phase_checkpointing"
+
+
+def test_training_missing_enter_keeps_normal_timeout():
+    analyzer = TraceAnalyzer(
+        run_id=RUN_ID,
+        collective_timeout_s=5,
+        checkpoint_timeout_s=20,
+    )
+    _ingest(analyzer, _event("comm_init", rank=1, comm_rank=1))
+    _ingest(analyzer, _event("heartbeat", rank=1, phase="train"), observed=1.0)
+    _ingest(analyzer, _event("nccl_call", rank=0, comm_rank=0), observed=1.0)
+
+    findings = analyzer.scan(now_monotonic_s=8.0, now_unix_ns=8_000_000_000)
+    assert [finding.hang_type for finding in findings] == ["collective_missing_enter"]
+    assert findings[0].details["detection_phase"] == "train"
+    assert findings[0].details["detection_threshold_s"] == 5
+    assert findings[0].details["threshold_reason"] == "normal_collective_phase"
+
+
+def test_checkpoint_delayed_enter_keeps_phase_after_heartbeat_returns_to_train():
+    analyzer = TraceAnalyzer(
+        run_id=RUN_ID,
+        delayed_enter_threshold_s=5,
+        checkpoint_timeout_s=20,
+    )
+    _ingest(
+        analyzer,
+        _event("heartbeat", rank=0, phase="checkpointing"),
+        observed=1.0,
+    )
+    _ingest(
+        analyzer,
+        _event("nccl_call", rank=0, comm_rank=0, timestamp_unix_ns=1_000_000_000),
+        observed=1.0,
+    )
+    assert analyzer.scan(now_monotonic_s=2.0, now_unix_ns=2_000_000_000) == []
+
+    _ingest(
+        analyzer,
+        _event(
+            "heartbeat",
+            rank=0,
+            phase="train",
+            timestamp_unix_ns=3_000_000_000,
+        ),
+        observed=3.0,
+    )
+    _ingest(
+        analyzer,
+        _event("nccl_call", comm_rank=1, rank=1, timestamp_unix_ns=8_000_000_000),
+        observed=8.0,
+    )
+
+    assert analyzer.scan(now_monotonic_s=9.0, now_unix_ns=9_000_000_000) == []
+
+
 def test_events_from_other_runs_are_ignored():
     analyzer = TraceAnalyzer(run_id=RUN_ID)
     accepted = _ingest(analyzer, _event("nccl_call", run_id="other-run"))

tests/unit_tests/runner/tracing/test_config.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+from types import SimpleNamespace
 
 import pytest
 from omegaconf import OmegaConf
@@ -56,7 +57,13 @@ def test_enabled_tracing_renders_cpu_probe_and_analyzer_shell(tmp_path):
         },
     )
 
-    resolved = prepare_trace_launch_config(config, "run-123")
+    heartbeat = SimpleNamespace(
+        enabled=True,
+        heartbeat_dir=str(tmp_path / "heartbeat"),
+        process_timeout_s=30,
+        checkpoint_timeout_s=600,
+    )
+    resolved = prepare_trace_launch_config(config, "run-123", heartbeat)
     shell = "\n".join(resolved.shell_setup_lines(0))
 
     assert resolved.enabled is True
@@ -67,6 +74,8 @@ def test_enabled_tracing_renders_cpu_probe_and_analyzer_shell(tmp_path):
     assert "flagscale.runner.tracing.monitor" in shell
     assert "--p2p-timeout 12" in shell
     assert "--p2p-match-window 4" in shell
+    assert "--checkpoint-timeout 600" in shell
+    assert resolved.checkpoint_timeout_s == 600
     assert "rc=\\$?" in resolved.command_body(0)
     assert resolved.shell_setup_lines(1)
     assert "flagscale.runner.tracing.monitor" not in "\n".join(resolved.shell_setup_lines(1))
@@ -97,6 +106,20 @@ def test_p2p_match_window_cannot_exceed_timeout(tmp_path):
         prepare_trace_launch_config(config, "run")
 
 
+def test_checkpoint_timeout_cannot_be_shorter_than_normal_thresholds(tmp_path):
+    config = _config(
+        tmp_path,
+        {
+            "enabled": True,
+            "collective_timeout_s": 10,
+            "delayed_enter_threshold_s": 5,
+            "checkpoint_timeout_s": 9,
+        },
+    )
+    with pytest.raises(ValueError, match="checkpoint_timeout_s"):
+        prepare_trace_launch_config(config, "run")
+
+
 def test_no_shared_filesystem_is_rejected_for_cross_rank_analysis(tmp_path):
     config = _config(tmp_path, {"enabled": True})
     config.experiment.runner.no_shared_fs = True