Detect missing GPU health snapshots

Author: lzy-edu

flagscale/runner/heartbeat/config.py

@@ -63,6 +63,7 @@ class HeartbeatLaunchConfig:
     scan_interval_s: float = 1.0
     monitor_nice: int = 10
     expected_world_size: int = 0
+    expected_node_count: int = 0
     hardware_health_enabled: bool = False
     hardware_health_interval_s: float = 60.0
     hardware_health_command_timeout_s: float = 10.0
@@ -193,6 +194,8 @@ def shell_setup_lines(self, node_rank: int) -> list[str]:
                 monitor_cmd.extend(
                     [
                         "--hardware-health-enabled",
+                        "--expected-node-count",
+                        str(self.expected_node_count),
                         "--hardware-health-stale-after",
                         f"{self.hardware_health_stale_after_s:g}",
                     ]
@@ -241,6 +244,14 @@ def _infer_world_size(runner: Any) -> int:
     return nnodes * nproc_per_node if nnodes > 0 and nproc_per_node > 0 else 0
 
 
+def _infer_node_count(runner: Any) -> int:
+    try:
+        nnodes = int(runner.get("nnodes", 1))
+    except (TypeError, ValueError):
+        return 0
+    return nnodes if nnodes > 0 else 0
+
+
 def prepare_heartbeat_launch_config(config: DictConfig, run_id: str) -> HeartbeatLaunchConfig:
     """Resolve ``experiment.runner.heartbeat`` without leaking it to torchrun."""
 
@@ -333,6 +344,7 @@ def prepare_heartbeat_launch_config(config: DictConfig, run_id: str) -> Heartbea
         raise ValueError("heartbeat.expected_world_size must be a non-negative integer") from exc
     if expected_world_size < 0:
         raise ValueError("heartbeat.expected_world_size must be a non-negative integer")
+    expected_node_count = _infer_node_count(config.experiment.runner)
 
     hardware_health = raw_dict.get("hardware_health", {})
     if not isinstance(hardware_health, dict):
@@ -370,6 +382,7 @@ def prepare_heartbeat_launch_config(config: DictConfig, run_id: str) -> Heartbea
         scan_interval_s=_positive_float(raw_dict.get("scan_interval_s", 1.0), "scan_interval_s"),
         monitor_nice=monitor_nice,
         expected_world_size=expected_world_size,
+        expected_node_count=expected_node_count,
         hardware_health_enabled=hardware_health_enabled,
         hardware_health_interval_s=hardware_health_interval_s,
         hardware_health_command_timeout_s=hardware_health_command_timeout_s,

flagscale/runner/heartbeat/health_reader.py

@@ -43,12 +43,14 @@ class HardwareHealthIndex:
     now_unix_ns: int
     stale_after_s: float
     snapshots: tuple[dict[str, Any], ...] = ()
+    missing_node_ranks: tuple[int, ...] = ()
 
     @property
     def overall_status(self) -> str:
         if not self.enabled:
             return "not_collected"
         statuses = [self._snapshot_status(snapshot) for snapshot in self.snapshots]
+        statuses.extend("unavailable" for _ in self.missing_node_ranks)
         return _worst_status(statuses, "unavailable")
 
     def _snapshot_age_s(self, snapshot: dict[str, Any]) -> float | None:
@@ -80,6 +82,18 @@ def summary(self) -> dict[str, Any]:
                     "gpus": snapshot.get("gpus", []),
                 }
             )
+        for node_rank in self.missing_node_ranks:
+            nodes.append(
+                {
+                    "node_rank": node_rank,
+                    "hostname": None,
+                    "status": "unavailable",
+                    "sample_age_s": None,
+                    "source": None,
+                    "error": "gpu_health_snapshot_not_generated",
+                    "gpus": [],
+                }
+            )
         return {"enabled": True, "status": self.overall_status, "nodes": nodes}
 
     def for_rank(self, rank_event: dict[str, Any]) -> dict[str, Any]:
@@ -145,13 +159,22 @@ def _match_rank_gpu(
 
 class HardwareHealthReader:
     def __init__(
-        self, heartbeat_dir: Path, run_id: str, enabled: bool, stale_after_s: float
+        self,
+        heartbeat_dir: Path,
+        run_id: str,
+        enabled: bool,
+        stale_after_s: float,
+        expected_node_count: int = 0,
+        monitor_started_unix_ns: int = 0,
     ) -> None:
         self.heartbeat_dir = heartbeat_dir
         self.run_id = run_id
         self.enabled = enabled
         self.stale_after_s = stale_after_s
+        self.expected_node_count = expected_node_count
+        self.monitor_started_unix_ns = monitor_started_unix_ns
         self._reported: set[tuple[str, str, tuple[str, ...]]] = set()
+        self._reported_missing_nodes: set[int] = set()
 
     def poll(self, now_unix_ns: int) -> tuple[HardwareHealthIndex, list[dict[str, Any]]]:
         if not self.enabled:
@@ -164,8 +187,67 @@ def poll(self, now_unix_ns: int) -> tuple[HardwareHealthIndex, list[dict[str, An
                 continue
             if isinstance(payload, dict) and payload.get("run_id") == self.run_id:
                 snapshots.append(payload)
-        index = HardwareHealthIndex(True, now_unix_ns, self.stale_after_s, tuple(snapshots))
-        return index, self._new_findings(index, now_unix_ns)
+        missing_node_ranks = self._missing_node_ranks(snapshots, now_unix_ns)
+        index = HardwareHealthIndex(
+            True,
+            now_unix_ns,
+            self.stale_after_s,
+            tuple(snapshots),
+            missing_node_ranks,
+        )
+        findings = self._new_findings(index, now_unix_ns)
+        findings.extend(self._new_missing_findings(index, now_unix_ns))
+        return index, findings
+
+    def _missing_node_ranks(
+        self, snapshots: list[dict[str, Any]], now_unix_ns: int
+    ) -> tuple[int, ...]:
+        if self.expected_node_count <= 0:
+            return ()
+        monitor_age_s = max(
+            0.0,
+            (now_unix_ns - self.monitor_started_unix_ns) / 1_000_000_000,
+        )
+        if monitor_age_s <= self.stale_after_s:
+            return ()
+        observed_node_ranks: set[int] = set()
+        for snapshot in snapshots:
+            try:
+                node_rank = int(snapshot.get("node_rank"))
+            except (TypeError, ValueError):
+                continue
+            if 0 <= node_rank < self.expected_node_count:
+                observed_node_ranks.add(node_rank)
+        return tuple(
+            node_rank
+            for node_rank in range(self.expected_node_count)
+            if node_rank not in observed_node_ranks
+        )
+
+    def _new_missing_findings(
+        self, index: HardwareHealthIndex, now_unix_ns: int
+    ) -> list[dict[str, Any]]:
+        missing_node_ranks = set(index.missing_node_ranks)
+        self._reported_missing_nodes.intersection_update(missing_node_ranks)
+        findings: list[dict[str, Any]] = []
+        for node_rank in index.missing_node_ranks:
+            if node_rank in self._reported_missing_nodes:
+                continue
+            self._reported_missing_nodes.add(node_rank)
+            findings.append(
+                {
+                    "finding_type": "gpu_health_snapshot_missing",
+                    "run_id": self.run_id,
+                    "detected_at_unix_ns": now_unix_ns,
+                    "node_rank": node_rank,
+                    "hostname": None,
+                    "gpu_device_health": "unavailable",
+                    "expected_file": f"gpu_health_node_{node_rank}.json",
+                    "reason": "gpu_health_snapshot_not_generated",
+                    "confidence": "observed",
+                }
+            )
+        return findings
 
     def _new_findings(self, index: HardwareHealthIndex, now_unix_ns: int) -> list[dict[str, Any]]:
         findings: list[dict[str, Any]] = []

flagscale/runner/heartbeat/monitor.py

@@ -392,6 +392,8 @@ def run_monitor(args: argparse.Namespace) -> int:
         except (AttributeError, OSError):
             logger.debug("Could not adjust heartbeat monitor niceness", exc_info=True)
 
+    monitor_started_s = time.monotonic()
+    monitor_started_unix_ns = time.time_ns()
     analyzer = HeartbeatAnalyzer(
         args.run_id,
         args.initial_process_timeout,
@@ -400,14 +402,16 @@ def run_monitor(args: argparse.Namespace) -> int:
         args.progress_timeout,
         args.checkpoint_timeout,
         expected_world_size=args.expected_world_size,
-        monitor_started_s=time.monotonic(),
+        monitor_started_s=monitor_started_s,
     )
     tailer = JsonlTailer(heartbeat_dir)
     hardware_reader = HardwareHealthReader(
         heartbeat_dir,
         args.run_id,
         args.hardware_health_enabled,
         args.hardware_health_stale_after,
+        expected_node_count=args.expected_node_count,
+        monitor_started_unix_ns=monitor_started_unix_ns,
     )
     stopping = False
     failed_seen_s: float | None = None
@@ -446,6 +450,7 @@ def build_parser() -> argparse.ArgumentParser:
     parser.add_argument("--heartbeat-dir", required=True)
     parser.add_argument("--run-id", required=True)
     parser.add_argument("--expected-world-size", type=int, default=0)
+    parser.add_argument("--expected-node-count", type=int, default=0)
     parser.add_argument("--initial-process-timeout", type=float, default=30.0)
     parser.add_argument("--process-timeout", type=float, default=30.0)
     parser.add_argument("--initial-progress-timeout", type=float, default=600.0)
@@ -469,6 +474,8 @@ def main() -> int:
     args = build_parser().parse_args()
     if args.expected_world_size < 0:
         raise SystemExit("--expected-world-size must be non-negative")
+    if args.expected_node_count < 0:
+        raise SystemExit("--expected-node-count must be non-negative")
     for name in (
         "initial_process_timeout",
         "process_timeout",

tests/unit_tests/runner/heartbeat/test_config.py

@@ -76,27 +76,28 @@ def test_enabled_gpu_progress_heartbeat_has_no_preload_or_nccl_dependency(tmp_pa
 
 
 def test_optional_hardware_health_starts_one_node_local_cpu_collector(tmp_path):
-    resolved = prepare_heartbeat_launch_config(
-        _config(
-            tmp_path,
-            {
+    config = _config(
+        tmp_path,
+        {
+            "enabled": True,
+            "hardware_health": {
                 "enabled": True,
-                "hardware_health": {
-                    "enabled": True,
-                    "interval_s": 60,
-                    "command_timeout_s": 10,
-                    "stale_after_s": 180,
-                },
+                "interval_s": 60,
+                "command_timeout_s": 10,
+                "stale_after_s": 180,
             },
-        ),
-        "run",
+        },
     )
+    config.experiment.runner.nnodes = 2
+    resolved = prepare_heartbeat_launch_config(config, "run")
     node_zero = "\n".join(resolved.shell_setup_lines(0))
     node_one = "\n".join(resolved.shell_setup_lines(1))
 
     assert "flagscale.runner.heartbeat.gpu_health" in node_zero
     assert "gpu_health_node_0.json" in node_zero
     assert "--hardware-health-enabled" in node_zero
+    assert "--expected-node-count 2" in node_zero
+    assert resolved.expected_node_count == 2
     assert "flagscale.runner.heartbeat.gpu_health" in node_one
     assert "gpu_health_node_1.json" in node_one
     assert "flagscale.runner.heartbeat.monitor" not in node_one

tests/unit_tests/runner/heartbeat/test_health_reader.py

@@ -17,8 +17,8 @@
 from flagscale.runner.heartbeat.health_reader import HardwareHealthReader
 
 
-def _snapshot(status="warning"):
-    return {
+def _snapshot(status="warning", **overrides):
+    payload = {
         "run_id": "run",
         "node_rank": 0,
         "hostname": "host-a",
@@ -35,6 +35,8 @@ def _snapshot(status="warning"):
             },
         ],
     }
+    payload.update(overrides)
+    return payload
 
 
 def test_reader_correlates_cuda_visible_devices_and_deduplicates_findings(tmp_path):
@@ -74,3 +76,63 @@ def test_stale_snapshot_is_not_reported_as_hardware_failure(tmp_path):
 
     assert index.overall_status == "stale"
     assert findings == []
+
+
+def test_missing_node_snapshot_is_reported_after_startup_grace(tmp_path):
+    node_zero = tmp_path / "gpu_health_node_0.json"
+    node_zero.write_text(json.dumps(_snapshot("healthy")), encoding="utf-8")
+    reader = HardwareHealthReader(
+        tmp_path,
+        "run",
+        True,
+        stale_after_s=30,
+        expected_node_count=2,
+        monitor_started_unix_ns=0,
+    )
+
+    initial_index, initial_findings = reader.poll(now_unix_ns=20_000_000_000)
+    assert initial_index.overall_status == "healthy"
+    assert initial_index.missing_node_ranks == ()
+    assert initial_findings == []
+
+    missing_index, missing_findings = reader.poll(now_unix_ns=31_000_000_000)
+    assert missing_index.overall_status == "unavailable"
+    assert missing_index.missing_node_ranks == (1,)
+    assert missing_index.summary()["nodes"][1] == {
+        "node_rank": 1,
+        "hostname": None,
+        "status": "unavailable",
+        "sample_age_s": None,
+        "source": None,
+        "error": "gpu_health_snapshot_not_generated",
+        "gpus": [],
+    }
+    assert [finding["finding_type"] for finding in missing_findings] == [
+        "gpu_health_snapshot_missing"
+    ]
+    assert missing_findings[0]["node_rank"] == 1
+    assert missing_findings[0]["expected_file"] == "gpu_health_node_1.json"
+    assert reader.poll(now_unix_ns=32_000_000_000)[1] == []
+
+    node_one = tmp_path / "gpu_health_node_1.json"
+    node_one.write_text(
+        json.dumps(
+            _snapshot(
+                "healthy",
+                node_rank=1,
+                hostname="host-b",
+                collected_at_unix_ns=32_000_000_000,
+            )
+        ),
+        encoding="utf-8",
+    )
+    recovered_index, recovered_findings = reader.poll(now_unix_ns=33_000_000_000)
+    assert recovered_index.overall_status == "healthy"
+    assert recovered_index.missing_node_ranks == ()
+    assert recovered_findings == []
+
+    node_one.unlink()
+    repeated_findings = reader.poll(now_unix_ns=34_000_000_000)[1]
+    assert [finding["finding_type"] for finding in repeated_findings] == [
+        "gpu_health_snapshot_missing"
+    ]