diff --git a/miles/utils/test_utils/comparisons/metrics.py b/miles/utils/test_utils/comparisons/metrics.py new file mode 100644 index 00000000000..70721fbf57e --- /dev/null +++ b/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)] diff --git a/tests/fast/utils/test_utils/comparisons/test_metrics.py b/tests/fast/utils/test_utils/comparisons/test_metrics.py new file mode 100644 index 00000000000..8b1f7a2910d --- /dev/null +++ b/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]