Add dump and inference-engine-checksum comparison helpers for FT tests

miles/utils/test_utils/comparisons/comparators.py

@@ -0,0 +1,55 @@
+import subprocess
+import sys
+from pathlib import Path
+
+
+def run_comparator(
+    *,
+    baseline_path: Path,
+    target_path: Path,
+    diff_thresholds: list[tuple[str, str]],
+    allow_skipped_pattern: str,
+    allow_failed_pattern: str,
+    grouping_skip_keys: list[str] | None = None,
+    extra_args: list[str] | None,
+) -> subprocess.CompletedProcess[str]:
+    # Skip 'rank' when grouping bundles: under FT (target) and non-FT (baseline) the same
+    # logical (pp_rank, cp_rank, ep_rank, tp_rank) coordinate maps to a different absolute
+    # rank ID (e.g. baseline rank=4 vs target cell0 rank=2 for PP=1, CP=0). Without skipping
+    # 'rank' the comparator gets `baseline_load_failed` for every tensor and fails with rc=1.
+    # Callers may pass extra keys (e.g. no_failure skips 'dp'/'edp' too). (Grouping is a
+    # comparator-matching detail, not a pass/fail threshold.)
+    skip_keys: list[str] = list(grouping_skip_keys) if grouping_skip_keys is not None else ["rank"]
+    assert "rank" in skip_keys, f"grouping_skip_keys must include 'rank', got {skip_keys}"
+
+    cmd: list[str] = [
+        sys.executable,
+        "-m",
+        "sglang.srt.debug_utils.comparator",
+        "--baseline-path",
+        str(baseline_path),
+        "--target-path",
+        str(target_path),
+        "--output-format",
+        "json",
+        "--grouping-skip-keys",
+        *skip_keys,
+        "--allow-skipped-pattern",
+        allow_skipped_pattern,
+        "--allow-failed-pattern",
+        allow_failed_pattern,
+    ]
+    if extra_args:
+        cmd.extend(extra_args)
+    # Keep --diff-threshold strictly last: its nargs="*" greedily consumes every
+    # following token, so no flag with a bare value may come after it.
+    cmd.append("--diff-threshold")
+    for pattern, predicate in diff_thresholds:
+        cmd.extend([pattern, predicate])
+
+    result: subprocess.CompletedProcess[str] = subprocess.run(
+        cmd,
+        text=True,
+    )
+
+    return result

miles/utils/test_utils/comparisons/dumps.py

@@ -0,0 +1,67 @@
+from pathlib import Path
+
+from miles.utils.test_utils.comparisons.comparators import run_comparator
+
+# Shared regexes for model-input / metadata tensors that are not weights or grads to
+# compare. Exposed as named constants (not as defaults) so each test passes them
+# explicitly -- every pass/fail knob is visible at the call site, nothing is implicit.
+INPUT_TENSORS_SKIP_PATTERN: str = "input_ids|positions|cu_seqlens_q|cu_seqlens_kv|qkv_format|.*witness.*"
+INPUT_TENSORS_ALLOW_FAILED_PATTERN: str = "input_ids|positions|cu_seqlens_q|cu_seqlens_kv|qkv_format"
+
+
+def compare_dumps(
+    baseline_dir: str,
+    target_dir: str,
+    *,
+    diff_thresholds: list[tuple[str, str]],
+    allow_skipped_pattern: str,
+    allow_failed_pattern: str,
+    phase_subdir: str | None = None,
+    grouping_skip_keys: list[str] | None = None,
+    extra_args: list[str] | None = None,
+) -> None:
+    subdir = phase_subdir or ""
+    baseline_root = Path(baseline_dir) / "dumps" / subdir
+    target_root = Path(target_dir) / "dumps" / subdir
+
+    assert baseline_root.exists(), f"Baseline dump dir does not exist: {baseline_root}"
+    assert target_root.exists(), f"Target dump dir does not exist: {target_root}"
+
+    # Dumps are segmented into leaf dirs (e.g. fwd_bwd/rollout_<id>), each a flat set of
+    # .pt files with its own per-leaf step numbering. The sglang comparator compares one
+    # flat dir at a time, so compare each matching leaf pair independently.
+    baseline_leaves = _find_leaf_dump_dirs(baseline_root)
+    target_leaves = _find_leaf_dump_dirs(target_root)
+
+    assert baseline_leaves, f"No .pt dump files found under {baseline_root}"
+    assert baseline_leaves == target_leaves, (
+        f"Dump leaf-dir mismatch: baseline={baseline_leaves} vs target={target_leaves} "
+        f"(under {baseline_root} vs {target_root})"
+    )
+
+    failed_leaves: list[str] = []
+    for leaf in baseline_leaves:
+        result = run_comparator(
+            baseline_path=baseline_root / leaf,
+            target_path=target_root / leaf,
+            diff_thresholds=diff_thresholds,
+            allow_skipped_pattern=allow_skipped_pattern,
+            allow_failed_pattern=allow_failed_pattern,
+            grouping_skip_keys=grouping_skip_keys,
+            extra_args=extra_args,
+        )
+        if result.returncode != 0:
+            failed_leaves.append(leaf)
+
+    assert not failed_leaves, (
+        f"Dump comparator failed (rc!=0) for {len(failed_leaves)}/{len(baseline_leaves)} leaf dir(s): "
+        f"{failed_leaves} (baseline {baseline_root} vs target {target_root}). The comparator applies the "
+        f"per-tensor predicates ({diff_thresholds}) and the allow/skip patterns itself; see "
+        f"comparator_report.jsonl under {target_root}/<leaf> for the offending tensors."
+    )
+    print(f"Dump comparison passed: {len(baseline_leaves)} leaf dir(s) under {baseline_root} vs {target_root}")
+
+
+def _find_leaf_dump_dirs(root: Path) -> list[str]:
+    leaves: set[str] = {str(p.parent.relative_to(root)) for p in root.rglob("*.pt")}
+    return sorted(leaves)

miles/utils/test_utils/comparisons/inference_engine_checksums.py

@@ -0,0 +1,65 @@
+from pathlib import Path
+
+from miles.utils.event_analyzer.rules import inference_engine_weight_checksum_consistency
+from miles.utils.event_analyzer.rules.checksum_compare import ChecksumMismatchIssue, compare_flat_dicts
+from miles.utils.event_logger.logger import read_events
+from miles.utils.event_logger.models import InferenceEngineWeightChecksumEvent
+
+
+def compare_inference_engine_checksums(baseline_dir: str, target_dir: str) -> None:
+    baseline = _read_inference_engine_checksum_events(Path(baseline_dir))
+    target = _read_inference_engine_checksum_events(Path(target_dir))
+    assert baseline, f"No InferenceEngineWeightChecksumEvents found in baseline dir: {baseline_dir}"
+    assert target, f"No InferenceEngineWeightChecksumEvents found in target dir: {target_dir}"
+
+    # Each side's engines must already agree internally (same invariant as the production rule), so
+    # one representative engine per rollout then proves baseline == target regardless of engine count.
+    assert not inference_engine_weight_checksum_consistency.check(
+        baseline
+    ), "Baseline engines disagree with each other"
+    assert not inference_engine_weight_checksum_consistency.check(target), "Target engines disagree with each other"
+
+    baseline_by_rollout = _checksums_by_rollout_id(baseline)
+    target_by_rollout = _checksums_by_rollout_id(target)
+    assert baseline_by_rollout.keys() == target_by_rollout.keys(), (
+        f"Engine checksum rollout_id sets differ: "
+        f"baseline={sorted(baseline_by_rollout)} "
+        f"vs target={sorted(target_by_rollout)}"
+    )
+
+    mismatches: list[ChecksumMismatchIssue] = []
+    for rollout_id in sorted(baseline_by_rollout):
+        mismatches += list(
+            compare_flat_dicts(
+                a=baseline_by_rollout[rollout_id],
+                b=target_by_rollout[rollout_id],
+                label_a=f"baseline/rollout_{rollout_id}",
+                label_b=f"target/rollout_{rollout_id}",
+            )
+        )
+    assert not mismatches, "Engine weight checksum baseline-vs-target mismatch:\n" + "\n".join(
+        f"  - {m.label_a} vs {m.label_b} key {m.key}: {m.value_a} != {m.value_b}" for m in mismatches
+    )
+    print(f"Engine weight checksum comparison passed: {len(baseline_by_rollout)} rollout(s) compared")
+
+
+def _checksums_by_rollout_id(events: list[InferenceEngineWeightChecksumEvent]) -> dict[int, dict[str, str]]:
+    by_rollout: dict[int, dict[str, str]] = {}
+    for event in events:
+        if event.rollout_id is None:
+            continue
+        assert (
+            event.rollout_id not in by_rollout
+        ), f"Duplicate InferenceEngineWeightChecksumEvent for rollout {event.rollout_id}"
+        assert event.engine_checksums, f"No engine checksums for rollout {event.rollout_id}"
+        by_rollout[event.rollout_id] = event.engine_checksums[0]
+    return by_rollout
+
+
+def _read_inference_engine_checksum_events(dump_dir: Path) -> list[InferenceEngineWeightChecksumEvent]:
+    """Read all InferenceEngineWeightChecksumEvents 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, InferenceEngineWeightChecksumEvent)]

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

@@ -0,0 +1,45 @@
+from pathlib import Path
+
+from miles.utils.test_utils.comparisons.dumps import _find_leaf_dump_dirs
+
+
+class TestFindLeafDumpDirs:
+    def test_two_pt_files_in_one_leaf_yield_single_entry(self, tmp_path: Path) -> None:
+        """Multiple .pt files sharing one leaf dir dedup to a single relative entry."""
+        leaf = tmp_path / "fwd_bwd" / "rollout_0"
+        leaf.mkdir(parents=True)
+        (leaf / "step_0.pt").touch()
+        (leaf / "step_1.pt").touch()
+
+        assert _find_leaf_dump_dirs(tmp_path) == ["fwd_bwd/rollout_0"]
+
+    def test_two_leaves_returned_sorted(self, tmp_path: Path) -> None:
+        """Distinct leaf dirs are returned sorted by their relative path string."""
+        leaf_b = tmp_path / "leaf_b"
+        leaf_a = tmp_path / "leaf_a"
+        leaf_b.mkdir()
+        leaf_a.mkdir()
+        (leaf_b / "x.pt").touch()
+        (leaf_a / "y.pt").touch()
+
+        assert _find_leaf_dump_dirs(tmp_path) == ["leaf_a", "leaf_b"]
+
+    def test_pt_file_directly_in_root_yields_dot(self, tmp_path: Path) -> None:
+        """A .pt file directly under root has parent equal to root, reported as '.'."""
+        (tmp_path / "step_0.pt").touch()
+
+        assert _find_leaf_dump_dirs(tmp_path) == ["."]
+
+    def test_no_pt_files_yields_empty_list(self, tmp_path: Path) -> None:
+        """A tree with no .pt files produces an empty list."""
+        (tmp_path / "sub").mkdir()
+
+        assert _find_leaf_dump_dirs(tmp_path) == []
+
+    def test_non_pt_files_are_ignored(self, tmp_path: Path) -> None:
+        """Files not matching *.pt (including *.pth) are ignored by the glob."""
+        (tmp_path / "notes.txt").touch()
+        (tmp_path / "weights.pth").touch()
+        (tmp_path / "data.json").touch()
+
+        assert _find_leaf_dump_dirs(tmp_path) == []

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

@@ -0,0 +1,149 @@
+"""Tests for test_utils.comparisons.inference_engine_checksums.compare_inference_engine_checksums."""
+
+from pathlib import Path
+from typing import Any
+
+import pytest
+
+from miles.utils.event_logger.logger import EventLogger
+from miles.utils.event_logger.models import InferenceEngineWeightChecksumEvent
+from miles.utils.process_identity import MainProcessIdentity
+from miles.utils.test_utils.comparisons.inference_engine_checksums import compare_inference_engine_checksums
+
+
+def _write_inference_engine_events(side_dir: Path, partials: list[dict[str, Any]]) -> None:
+    events_dir = side_dir / "events"
+    event_logger = EventLogger(log_dir=events_dir, source=MainProcessIdentity())
+    for partial in partials:
+        event_logger.log(InferenceEngineWeightChecksumEvent, partial, print_log=False)
+    event_logger.close()
+
+
+def _partial(*, rollout_id: int | None, engine_checksums: list[dict[str, str]]) -> dict[str, Any]:
+    return dict(rollout_id=rollout_id, engine_checksums=engine_checksums)
+
+
+class TestCompareInferenceEngineChecksums:
+    def test_identical_passes(self, tmp_path: Path) -> None:
+        """Internally-consistent sides with equal representative checksums pass."""
+        partials = [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "aaa"}])]
+        _write_inference_engine_events(tmp_path / "baseline", partials)
+        _write_inference_engine_events(tmp_path / "target", partials)
+
+        compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target"))
+
+    def test_differing_engine_counts_still_pass(self, tmp_path: Path) -> None:
+        """Engine count may differ between sides; only internal agreement + representative equality matter."""
+        _write_inference_engine_events(
+            tmp_path / "baseline", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])]
+        )
+        _write_inference_engine_events(
+            tmp_path / "target",
+            [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "aaa"}, {"rank0/w": "aaa"}])],
+        )
+
+        compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target"))
+
+    def test_none_rollout_id_skipped(self, tmp_path: Path) -> None:
+        """The initial out-of-loop sync (rollout_id=None) is not compared: it differs here yet the
+        per-rollout checksums match, so the comparison still passes."""
+        _write_inference_engine_events(
+            tmp_path / "baseline",
+            [
+                _partial(rollout_id=None, engine_checksums=[{"rank0/w": "init_baseline"}]),
+                _partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}]),
+            ],
+        )
+        _write_inference_engine_events(
+            tmp_path / "target",
+            [
+                _partial(rollout_id=None, engine_checksums=[{"rank0/w": "init_target"}]),
+                _partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}]),
+            ],
+        )
+
+        compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target"))
+
+    def test_recurring_none_across_phases_skipped(self, tmp_path: Path) -> None:
+        """A multi-phase resume yields several None events per side; all are skipped, so a side with
+        more None events than the other still passes when the per-rollout checksums match."""
+        _write_inference_engine_events(
+            tmp_path / "baseline",
+            [
+                _partial(rollout_id=None, engine_checksums=[{"rank0/w": "init_a"}]),
+                _partial(rollout_id=2, engine_checksums=[{"rank0/w": "aaa"}]),
+                _partial(rollout_id=None, engine_checksums=[{"rank0/w": "init_b"}]),
+                _partial(rollout_id=5, engine_checksums=[{"rank0/w": "bbb"}]),
+            ],
+        )
+        _write_inference_engine_events(
+            tmp_path / "target",
+            [
+                _partial(rollout_id=2, engine_checksums=[{"rank0/w": "aaa"}]),
+                _partial(rollout_id=5, engine_checksums=[{"rank0/w": "bbb"}]),
+            ],
+        )
+
+        compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target"))
+
+    def test_baseline_engines_disagree_fails(self, tmp_path: Path) -> None:
+        """If baseline's own engines disagree, the comparison fails (caught by the consistency rule)."""
+        _write_inference_engine_events(
+            tmp_path / "baseline", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "zzz"}])]
+        )
+        _write_inference_engine_events(
+            tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])]
+        )
+
+        with pytest.raises(AssertionError, match="Baseline engines disagree"):
+            compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target"))
+
+    def test_target_engines_disagree_fails(self, tmp_path: Path) -> None:
+        """If target's own engines disagree, the comparison fails."""
+        _write_inference_engine_events(
+            tmp_path / "baseline", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])]
+        )
+        _write_inference_engine_events(
+            tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}, {"rank0/w": "zzz"}])]
+        )
+
+        with pytest.raises(AssertionError, match="Target engines disagree"):
+            compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target"))
+
+    def test_representative_mismatch_fails(self, tmp_path: Path) -> None:
+        """Internally-consistent sides whose representatives differ fail and name the tensor."""
+        _write_inference_engine_events(
+            tmp_path / "baseline", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])]
+        )
+        _write_inference_engine_events(
+            tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "zzz"}])]
+        )
+
+        with pytest.raises(AssertionError, match=r"key rank0/w"):
+            compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target"))
+
+    def test_missing_rollout_fails(self, tmp_path: Path) -> None:
+        """A rollout present only on one side fails closed."""
+        _write_inference_engine_events(
+            tmp_path / "baseline",
+            [
+                _partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}]),
+                _partial(rollout_id=2, engine_checksums=[{"rank0/w": "ccc"}]),
+            ],
+        )
+        _write_inference_engine_events(
+            tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])]
+        )
+
+        with pytest.raises(AssertionError, match="rollout_id sets differ"):
+            compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target"))
+
+    def test_empty_baseline_fails(self, tmp_path: Path) -> None:
+        """No baseline events fails closed rather than vacuously passing."""
+        _write_inference_engine_events(tmp_path / "baseline", [])
+        _write_inference_engine_events(
+            tmp_path / "target", [_partial(rollout_id=1, engine_checksums=[{"rank0/w": "aaa"}])]
+        )
+
+        with pytest.raises(AssertionError, match="No InferenceEngineWeightChecksumEvents found in baseline"):
+            compare_inference_engine_checksums(str(tmp_path / "baseline"), str(tmp_path / "target"))