Add aligner entrypoint and bundle handler in dump comparator (#19375)

test/registered/debug_utils/comparator/test_bundle_matcher.py

@@ -0,0 +1,150 @@
+import sys
+from typing import Any
+
+import polars as pl
+import pytest
+
+from sglang.srt.debug_utils.comparator.bundle_matcher import (
+    TensorBundleInfo,
+    TensorFileInfo,
+    _rows_to_tensor_infos,
+    match_bundles,
+)
+from sglang.srt.debug_utils.comparator.utils import Pair
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=15, suite="default", nightly=True)
+
+
+def _make_row(
+    *, name: str, step: int = 0, rank: int = 0, filename: str | None = None
+) -> dict[str, Any]:
+    if filename is None:
+        filename = f"name={name}___step={step}___rank={rank}.pt"
+    return {"name": name, "step": step, "rank": rank, "filename": filename}
+
+
+def _make_df(rows: list[dict[str, Any]]) -> pl.DataFrame:
+    return pl.DataFrame(rows)
+
+
+class TestMatchBundles:
+    def test_single_tensor_single_step(self) -> None:
+        target_df: pl.DataFrame = _make_df([_make_row(name="t_a")])
+        baseline_df: pl.DataFrame = _make_df([_make_row(name="t_a")])
+
+        results: list[Pair[TensorBundleInfo]] = match_bundles(
+            dfs=Pair(x=baseline_df, y=target_df),
+            skip_keys={"filename"},
+        )
+
+        assert len(results) == 1
+        assert len(results[0].x) == 1
+        assert len(results[0].y) == 1
+        assert results[0].y[0].name == "t_a"
+
+    def test_multiple_names_separate_bundles(self) -> None:
+        target_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a"),
+            _make_row(name="t_b"),
+        ])
+        baseline_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a"),
+            _make_row(name="t_b"),
+        ])
+
+        results: list[Pair[TensorBundleInfo]] = match_bundles(
+            dfs=Pair(x=baseline_df, y=target_df),
+            skip_keys={"filename"},
+        )
+
+        assert len(results) == 2
+        result_names: list[str] = [r.y[0].name for r in results]
+        assert "t_a" in result_names
+        assert "t_b" in result_names
+
+    def test_skip_rank_groups_across_ranks(self) -> None:
+        target_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a", rank=0),
+            _make_row(name="t_a", rank=1),
+        ])
+        baseline_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a", rank=0),
+            _make_row(name="t_a", rank=1),
+        ])
+
+        results: list[Pair[TensorBundleInfo]] = match_bundles(
+            dfs=Pair(x=baseline_df, y=target_df),
+            skip_keys={"filename", "rank"},
+        )
+
+        assert len(results) == 1
+        assert len(results[0].y) == 2
+
+    def test_baseline_missing_tensor(self) -> None:
+        target_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a"),
+            _make_row(name="t_extra"),
+        ])
+        baseline_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a"),
+        ])
+
+        results: list[Pair[TensorBundleInfo]] = match_bundles(
+            dfs=Pair(x=baseline_df, y=target_df),
+            skip_keys={"filename"},
+        )
+
+        assert len(results) == 2
+        extra_pair: Pair[TensorBundleInfo] = [
+            r for r in results if r.y[0].name == "t_extra"
+        ][0]
+        assert extra_pair.x == []
+
+    def test_empty_target_returns_empty(self) -> None:
+        target_df: pl.DataFrame = _make_df([])
+        baseline_df: pl.DataFrame = _make_df([_make_row(name="t_a")])
+
+        results: list[Pair[TensorBundleInfo]] = match_bundles(
+            dfs=Pair(x=baseline_df, y=target_df),
+            skip_keys={"filename"},
+        )
+
+        assert results == []
+
+    def test_skip_step_groups_across_steps(self) -> None:
+        target_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a", step=0),
+            _make_row(name="t_a", step=1),
+        ])
+        baseline_df: pl.DataFrame = _make_df([
+            _make_row(name="t_a", step=0),
+            _make_row(name="t_a", step=1),
+        ])
+
+        results: list[Pair[TensorBundleInfo]] = match_bundles(
+            dfs=Pair(x=baseline_df, y=target_df),
+            skip_keys={"filename", "step"},
+        )
+
+        assert len(results) == 1
+        assert len(results[0].y) == 2
+
+
+class TestRowsToTensorInfos:
+    def test_filters_extra_columns(self) -> None:
+        rows: list[dict[str, Any]] = [
+            {"filename": "a.pt", "name": "t_a", "step": 0, "rank": 7}
+        ]
+        infos: list[TensorFileInfo] = _rows_to_tensor_infos(rows)
+
+        assert len(infos) == 1
+        assert infos[0] == TensorFileInfo(filename="a.pt", name="t_a", step=0)
+
+    def test_empty_rows(self) -> None:
+        infos: list[TensorFileInfo] = _rows_to_tensor_infos([])
+        assert infos == []
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))