Add metric comparison helpers for FT tests

miles/utils/test_utils/comparisons/metrics.py

@@ -0,0 +1,206 @@
+import logging
+import math
+from collections import defaultdict
+from pathlib import Path
+
+import polars as pl
+from sglang.srt.debug_utils.comparator.display import _render_polars_as_text
+
+from miles.utils.event_logger.logger import read_events
+from miles.utils.event_logger.models import MetricEvent
+
+logger = logging.getLogger(__name__)
+
+_REQUIRED_METRIC_KEYS: list[str] = ["train/grad_norm", "train/loss"]
+
+
+def compare_metrics(
+    baseline_dir: str,
+    target_dir: str,
+    *,
+    rtol: float,
+    atol: float,
+    key_prefixes: list[str],
+    exclude_keys: list[str],
+) -> None:
+    baseline_events = _read_metric_events(Path(baseline_dir))
+    target_events = _read_metric_events(Path(target_dir))
+
+    # FT retries (healing path) leave events from earlier failed attempts. Only
+    # the highest-attempt events per rollout_id reflect the successful run.
+    baseline_events = _keep_only_final_attempt(baseline_events)
+    target_events = _keep_only_final_attempt(target_events)
+
+    issues: list[str] = []
+    issues += _check_event_counts(baseline_events, target_events, baseline_dir, target_dir)
+
+    if not issues:
+        for step_idx, (b_event, t_event) in enumerate(zip(baseline_events, target_events, strict=True)):
+            _print_step_comparison_table(step_idx, b_event, t_event, key_prefixes, exclude_keys=exclude_keys)
+            issues += _check_step_metrics(
+                step_idx, b_event, t_event, key_prefixes, rtol, atol=atol, exclude_keys=exclude_keys
+            )
+
+    issues += _check_required_keys_exist(baseline_events)
+
+    assert not issues, f"MetricEvent comparison found {len(issues)} issue(s):\n" + "\n".join(
+        f"  - {i}" for i in issues
+    )
+    print(f"MetricEvent comparison passed: {len(baseline_events)} steps compared")
+
+
+def _keep_only_final_attempt(events: list[MetricEvent]) -> list[MetricEvent]:
+    """Keep only events from the highest-attempt for each rollout_id.
+
+    During FT healing, a crashed rollout is retried at attempt+1; events from
+    the failed attempt are partial and should be discarded for comparison.
+
+    Rollout-side metrics (e.g. RolloutManager log_rollout_metrics) have
+    attempt=None — they are not part of the FT retry stream, so we treat them
+    as a single attempt (normalized to 0).
+    """
+
+    def _attempt(e: MetricEvent) -> int:
+        return e.attempt if e.attempt is not None else 0
+
+    max_attempt_by_rollout: dict[int, int] = defaultdict(int)
+    for e in events:
+        max_attempt_by_rollout[e.rollout_id] = max(max_attempt_by_rollout[e.rollout_id], _attempt(e))
+    return [e for e in events if _attempt(e) == max_attempt_by_rollout[e.rollout_id]]
+
+
+def _check_event_counts(
+    baseline: list[MetricEvent],
+    target: list[MetricEvent],
+    baseline_dir: str,
+    target_dir: str,
+) -> list[str]:
+    issues: list[str] = []
+    if len(baseline) == 0:
+        issues.append(f"No MetricEvents found in baseline dir: {baseline_dir}")
+    if len(target) == 0:
+        issues.append(f"No MetricEvents found in target dir: {target_dir}")
+    if len(baseline) > 0 and len(target) > 0 and len(baseline) != len(target):
+        issues.append(f"MetricEvent count mismatch: baseline={len(baseline)}, target={len(target)}")
+    return issues
+
+
+def _check_step_metrics(
+    step_idx: int,
+    baseline_event: MetricEvent,
+    target_event: MetricEvent,
+    key_prefixes: list[str],
+    rtol: float,
+    *,
+    atol: float,
+    exclude_keys: list[str] | None = None,
+) -> list[str]:
+    issues: list[str] = []
+    for key in baseline_event.metrics:
+        if not any(key.startswith(prefix) for prefix in key_prefixes):
+            continue
+        if exclude_keys and key in exclude_keys:
+            continue
+
+        if key not in target_event.metrics:
+            issues.append(f"Step {step_idx}: metric '{key}' present in baseline but missing in target")
+            continue
+
+        issues += _check_single_metric(
+            step_idx, key, baseline_event.metrics[key], target_event.metrics[key], rtol, atol=atol
+        )
+    return issues
+
+
+def _check_single_metric(
+    step_idx: int,
+    key: str,
+    baseline_val: object,
+    target_val: object,
+    rtol: float,
+    atol: float,
+) -> list[str]:
+    if not isinstance(baseline_val, (int, float)) or not isinstance(target_val, (int, float)):
+        return []
+
+    if math.isnan(baseline_val) or math.isnan(target_val):
+        return [f"Step {step_idx}, metric '{key}': NaN detected (baseline={baseline_val}, target={target_val})"]
+    if math.isinf(baseline_val) or math.isinf(target_val):
+        if baseline_val != target_val:
+            return [f"Step {step_idx}, metric '{key}': inf mismatch (baseline={baseline_val}, target={target_val})"]
+        return []
+
+    if baseline_val == 0.0 and target_val == 0.0:
+        return []
+
+    abs_diff = abs(baseline_val - target_val)
+    if abs_diff <= atol:
+        return []
+
+    rel_diff = abs_diff / max(abs(baseline_val), abs(target_val), 1e-12)
+    if rel_diff > rtol:
+        return [
+            f"Step {step_idx}, metric '{key}': baseline={baseline_val}, target={target_val}, "
+            f"rel_diff={rel_diff:.6f} > rtol={rtol}"
+        ]
+    return []
+
+
+def _print_step_comparison_table(
+    step_idx: int,
+    baseline_event: MetricEvent,
+    target_event: MetricEvent,
+    key_prefixes: list[str],
+    *,
+    exclude_keys: list[str] | None = None,
+) -> None:
+    rows: list[dict[str, str]] = []
+    for key in sorted(baseline_event.metrics):
+        if not any(key.startswith(p) for p in key_prefixes):
+            continue
+        b_val = baseline_event.metrics[key]
+        t_val = target_event.metrics.get(key)
+        if not isinstance(b_val, (int, float)) or t_val is None or not isinstance(t_val, (int, float)):
+            continue
+        excluded = "(excluded)" if exclude_keys and key in exclude_keys else ""
+        abs_diff = abs(b_val - t_val)
+        denom = max(abs(b_val), abs(t_val), 1e-12)
+        rel_diff = abs_diff / denom
+        rows.append(
+            {
+                "metric": key,
+                "baseline": f"{b_val:.6e}",
+                "target": f"{t_val:.6e}",
+                "abs_diff": f"{abs_diff:.2e}",
+                "rel_diff": f"{rel_diff:.4%}{excluded}",
+            }
+        )
+
+    if not rows:
+        return
+    df = pl.DataFrame(rows)
+    print(_render_polars_as_text(df, title=f"Step {step_idx} metric comparison"))
+
+
+def _check_required_keys_exist(events: list[MetricEvent]) -> list[str]:
+    all_keys: set[str] = set()
+    for event in events:
+        all_keys.update(event.metrics.keys())
+
+    issues: list[str] = []
+    for required in _REQUIRED_METRIC_KEYS:
+        if required not in all_keys:
+            issues.append(
+                f"Required metric '{required}' not found in any baseline MetricEvent. "
+                f"Available keys: {sorted(all_keys)}"
+            )
+    return issues
+
+
+def _read_metric_events(dump_dir: Path) -> list[MetricEvent]:
+    """Read all MetricEvents from the events directory."""
+    events_dir: Path = dump_dir / "events"
+    if not events_dir.exists():
+        return []
+    all_events = read_events(events_dir)
+    return [e for e in all_events if isinstance(e, MetricEvent)]

tests/fast/utils/test_utils/comparisons/test_metrics.py

@@ -0,0 +1,113 @@
+from datetime import datetime, timezone
+from typing import Any
+
+from miles.utils.event_logger.models import MetricEvent
+from miles.utils.process_identity import MainProcessIdentity
+from miles.utils.test_utils.comparisons.metrics import _check_single_metric, _keep_only_final_attempt
+
+_FIXED_TS = datetime(2026, 1, 1, tzinfo=timezone.utc)
+_FIXED_SOURCE = MainProcessIdentity()
+
+
+def _metric_event(
+    *, rollout_id: int | None, attempt: int | None, metrics: dict[str, Any] | None = None
+) -> MetricEvent:
+    return MetricEvent(
+        timestamp=_FIXED_TS,
+        source=_FIXED_SOURCE,
+        rollout_id=rollout_id,
+        attempt=attempt,
+        metrics=metrics if metrics is not None else {},
+    )
+
+
+class TestKeepOnlyFinalAttempt:
+    def test_keeps_highest_attempt_for_single_rollout(self) -> None:
+        """Among attempts 0,1,2 for one rollout_id, only the attempt=2 event survives."""
+        events = [
+            _metric_event(rollout_id=1, attempt=0),
+            _metric_event(rollout_id=1, attempt=1),
+            _metric_event(rollout_id=1, attempt=2),
+        ]
+        kept = _keep_only_final_attempt(events)
+        assert [e.attempt for e in kept] == [2]
+
+    def test_highest_attempt_resolved_independently_per_rollout(self) -> None:
+        """Each rollout_id keeps its own max attempt; different maxima coexist."""
+        events = [
+            _metric_event(rollout_id=1, attempt=0),
+            _metric_event(rollout_id=1, attempt=1),
+            _metric_event(rollout_id=2, attempt=0),
+        ]
+        kept = _keep_only_final_attempt(events)
+        assert {(e.rollout_id, e.attempt) for e in kept} == {(1, 1), (2, 0)}
+
+    def test_none_attempt_normalized_to_zero_and_dropped_when_mixed(self) -> None:
+        """attempt=None normalizes to 0, so it is dropped when an attempt=1 event shares the rollout_id."""
+        events = [
+            _metric_event(rollout_id=1, attempt=None),
+            _metric_event(rollout_id=1, attempt=1),
+        ]
+        kept = _keep_only_final_attempt(events)
+        assert [e.attempt for e in kept] == [1]
+
+    def test_empty_input_returns_empty(self) -> None:
+        """An empty event list yields an empty result."""
+        assert _keep_only_final_attempt([]) == []
+
+    def test_ties_on_max_attempt_all_kept(self) -> None:
+        """Multiple events tied at the max attempt for a rollout_id are all retained."""
+        events = [
+            _metric_event(rollout_id=1, attempt=2, metrics={"a": 1}),
+            _metric_event(rollout_id=1, attempt=2, metrics={"b": 2}),
+        ]
+        kept = _keep_only_final_attempt(events)
+        assert len(kept) == 2
+        assert [e.metrics for e in kept] == [{"a": 1}, {"b": 2}]
+
+
+class TestCheckSingleMetric:
+    def test_equal_values_no_issue(self) -> None:
+        """Exactly equal numeric values produce no issue."""
+        assert _check_single_metric(0, "k", 1.5, 1.5, rtol=0.01, atol=0.0) == []
+
+    def test_within_atol_no_issue(self) -> None:
+        """A difference within atol is accepted even if relative difference would exceed rtol."""
+        assert _check_single_metric(0, "k", 1.0, 1.0 + 1e-9, rtol=0.0, atol=1e-6) == []
+
+    def test_relative_difference_above_rtol_reports_issue(self) -> None:
+        """A relative difference above rtol (and above atol) yields exactly one issue."""
+        issues = _check_single_metric(3, "train/loss", 1.0, 2.0, rtol=0.1, atol=0.0)
+        assert len(issues) == 1
+        assert "train/loss" in issues[0]
+        assert "rel_diff" in issues[0]
+
+    def test_nan_detected(self) -> None:
+        """A NaN on either side produces a 'NaN detected' issue."""
+        issues = _check_single_metric(0, "k", float("nan"), 1.0, rtol=0.1, atol=0.0)
+        assert len(issues) == 1
+        assert "NaN detected" in issues[0]
+
+    def test_matching_inf_no_issue(self) -> None:
+        """inf == inf compares equal and produces no issue."""
+        assert _check_single_metric(0, "k", float("inf"), float("inf"), rtol=0.1, atol=0.0) == []
+
+    def test_inf_vs_finite_reports_mismatch(self) -> None:
+        """inf versus a finite value produces an 'inf mismatch' issue."""
+        issues = _check_single_metric(0, "k", float("inf"), 1.0, rtol=0.1, atol=0.0)
+        assert len(issues) == 1
+        assert "inf mismatch" in issues[0]
+
+    def test_both_zero_no_issue(self) -> None:
+        """Two exact zeros short-circuit to no issue."""
+        assert _check_single_metric(0, "k", 0.0, 0.0, rtol=0.0, atol=0.0) == []
+
+    def test_non_numeric_skipped(self) -> None:
+        """A non-numeric value on either side is skipped (no issue)."""
+        assert _check_single_metric(0, "k", "abc", 1.0, rtol=0.0, atol=0.0) == []
+
+    def test_tiny_baseline_uses_relative_floor(self) -> None:
+        """A near-zero baseline uses the 1e-12 denominator floor, making a tiny abs diff a large rel diff."""
+        issues = _check_single_metric(0, "k", 0.0, 5e-13, rtol=0.1, atol=0.0)
+        assert len(issues) == 1
+        assert "rel_diff" in issues[0]