Support method decorator for tagging and add minimalistic comparator in dumper (#19559)

test/registered/debug_utils/comparator/test_entrypoint.py

@@ -332,7 +332,6 @@ class TestEntrypointGroupingRaw:
                         enable=True,
                         dir=str(side_dir),
                         exp_name=_FIXED_EXP_NAME,
-                        enable_http_server=False,
                     )
                 )
                 dumper.__dict__["_static_meta"] = {"world_rank": 0, "world_size": 1}
@@ -1110,7 +1109,6 @@ class TestEntrypointAlignment:
                     enable=True,
                     dir=str(d),
                     exp_name=_FIXED_EXP_NAME,
-                    enable_http_server=False,
                 )
             )
 
@@ -1167,7 +1165,6 @@ class TestEntrypointAlignment:
                 enable=True,
                 dir=str(sglang_dir),
                 exp_name=_FIXED_EXP_NAME,
-                enable_http_server=False,
             )
         )
 
@@ -1226,7 +1223,6 @@ class TestEntrypointAlignment:
                 enable=True,
                 dir=str(megatron_dir),
                 exp_name=_FIXED_EXP_NAME,
-                enable_http_server=False,
             )
         )
 
@@ -1468,9 +1464,7 @@ def _assert_single_comparison_passed(records: list[AnyRecord]) -> ComparisonReco
 
 
 def _make_dumper(directory: Path) -> _Dumper:
-    return _Dumper(
-        config=DumperConfig(enable=True, dir=str(directory), enable_http_server=False)
-    )
+    return _Dumper(config=DumperConfig(enable=True, dir=str(directory)))
 
 
 def _create_dumps(
@@ -1528,7 +1522,6 @@ def _create_non_tensor_rank_dump(
                 enable=True,
                 dir=str(directory),
                 exp_name=_FIXED_EXP_NAME,
-                enable_http_server=False,
             )
         )
         dumper.__dict__["_static_meta"] = {"world_rank": rank, "world_size": 1}
@@ -1611,7 +1604,6 @@ def _create_rank_dump(
                 enable=True,
                 dir=str(directory),
                 exp_name=_FIXED_EXP_NAME,
-                enable_http_server=False,
             )
         )
 
@@ -1951,7 +1943,6 @@ class TestEntrypointThdCpZigzag:
                 enable=True,
                 dir=str(sglang_dir),
                 exp_name=_FIXED_EXP_NAME,
-                enable_http_server=False,
             )
         )
 

test/registered/debug_utils/comparator/test_model_validation.py

@@ -403,79 +403,5 @@ class TestAlignerPlanInComparisonRecord:
         assert "unsharder" in text
 
 
-def _make_aligner_plan() -> AlignerPlan:
-    unsharder = UnsharderPlan(
-        axis=ParallelAxis.TP,
-        params=ConcatParams(dim_name="h"),
-        groups=[[0, 1]],
-    )
-    return AlignerPlan(
-        per_step_plans=Pair(
-            x=[
-                AlignerPerStepPlan(
-                    step=0, input_object_indices=[0, 1], sub_plans=[unsharder]
-                )
-            ],
-            y=[
-                AlignerPerStepPlan(
-                    step=0, input_object_indices=[0, 1], sub_plans=[unsharder]
-                )
-            ],
-        ),
-    )
-
-
-class TestAlignerPlanInComparisonRecord:
-    def test_comparison_record_with_aligner_plan(self) -> None:
-        plan: AlignerPlan = _make_aligner_plan()
-        record: ComparisonRecord = _make_comparison_record(
-            diff=_make_diff_info(passed=True),
-        )
-        record_with_plan = record.model_copy(update={"aligner_plan": plan})
-        assert record_with_plan.aligner_plan is not None
-        assert record_with_plan.aligner_plan.per_step_plans.x[0].step == 0
-
-    def test_aligner_plan_json_roundtrip(self) -> None:
-        plan: AlignerPlan = _make_aligner_plan()
-        record: ComparisonRecord = _make_comparison_record(
-            diff=_make_diff_info(passed=True),
-        )
-        record_with_plan = record.model_copy(update={"aligner_plan": plan})
-
-        json_str: str = record_with_plan.model_dump_json()
-        parsed = json.loads(json_str)
-        assert "aligner_plan" in parsed
-        assert (
-            parsed["aligner_plan"]["per_step_plans"]["x"][0]["sub_plans"][0]["type"]
-            == "unsharder"
-        )
-
-        roundtripped: ComparisonRecord = parse_record_json(json_str)
-        assert roundtripped.aligner_plan is not None
-        assert (
-            roundtripped.aligner_plan.per_step_plans.x[0].sub_plans[0].type
-            == "unsharder"
-        )
-
-    def test_comparison_record_without_aligner_plan(self) -> None:
-        record: ComparisonRecord = _make_comparison_record(
-            diff=_make_diff_info(passed=True),
-        )
-        json_str: str = record.model_dump_json()
-        roundtripped: ComparisonRecord = parse_record_json(json_str)
-        assert roundtripped.aligner_plan is None
-
-    def test_aligner_plan_text_format(self) -> None:
-        plan: AlignerPlan = _make_aligner_plan()
-        record: ComparisonRecord = _make_comparison_record(
-            diff=_make_diff_info(passed=True),
-        )
-        record_with_plan = record.model_copy(update={"aligner_plan": plan})
-
-        text: str = record_with_plan.to_text()
-        assert "Aligner Plan:" in text
-        assert "unsharder" in text
-
-
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))

test/registered/debug_utils/test_dump_comparator.py

@@ -0,0 +1,159 @@
+from argparse import Namespace
+from pathlib import Path
+
+import pytest
+import torch
+
+from sglang.srt.debug_utils.dump_comparator import (
+    _argmax_coord,
+    _calc_rel_diff,
+    _compute_smaller_dtype,
+    _try_unify_shape,
+    main,
+)
+from sglang.srt.debug_utils.dumper import DumperConfig, _Dumper
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=30, suite="default", nightly=True)
+
+
+# ----------------------------- Unit tests -----------------------------
+
+
+class TestCalcRelDiff:
+    def test_identical_vectors(self) -> None:
+        x: torch.Tensor = torch.randn(10, 10)
+        assert _calc_rel_diff(x, x).item() == pytest.approx(0.0, abs=1e-5)
+
+    def test_zero_vectors(self) -> None:
+        z: torch.Tensor = torch.zeros(5)
+        result = _calc_rel_diff(z, z)
+        assert not torch.isnan(result) or True  # should not crash
+
+
+class TestArgmaxCoord:
+    def test_known_position(self) -> None:
+        x: torch.Tensor = torch.zeros(2, 3, 4)
+        x[1, 2, 3] = 10.0
+        assert _argmax_coord(x) == (1, 2, 3)
+
+
+class TestTryUnifyShape:
+    def test_squeeze_leading_ones(self) -> None:
+        target_shape: torch.Size = torch.Size([3, 4])
+        result: torch.Tensor = _try_unify_shape(torch.randn(1, 1, 3, 4), target_shape)
+        assert result.shape == target_shape
+
+    def test_no_op_when_no_leading_ones(self) -> None:
+        target_shape: torch.Size = torch.Size([3, 4])
+        result: torch.Tensor = _try_unify_shape(torch.randn(2, 3, 4), target_shape)
+        assert result.shape == (2, 3, 4)
+
+
+class TestComputeSmallerDtype:
+    def test_known_pair(self) -> None:
+        assert _compute_smaller_dtype(torch.float32, torch.bfloat16) == torch.bfloat16
+        assert _compute_smaller_dtype(torch.bfloat16, torch.float32) == torch.bfloat16
+
+    def test_none_for_same_dtype(self) -> None:
+        assert _compute_smaller_dtype(torch.float32, torch.float32) is None
+
+
+# ----------------------------- Integration tests -----------------------------
+
+
+def _make_dumper(directory: Path) -> _Dumper:
+    return _Dumper(
+        config=DumperConfig(
+            enable=True,
+            dir=str(directory),
+        )
+    )
+
+
+def _create_dumps(
+    tmp_path: Path,
+    tensor_names: list[str],
+    *,
+    baseline_names: list[str] | None = None,
+) -> tuple[Path, Path]:
+    if baseline_names is None:
+        baseline_names = tensor_names
+
+    d_baseline: Path = tmp_path / "baseline"
+    d_target: Path = tmp_path / "target"
+    d_baseline.mkdir()
+    d_target.mkdir()
+
+    torch.manual_seed(42)
+    baseline_tensor: torch.Tensor = torch.randn(10, 10)
+    target_tensor: torch.Tensor = baseline_tensor + torch.randn(10, 10) * 0.01
+
+    exp_paths: list[Path] = []
+    for d, names, tensor in [
+        (d_baseline, baseline_names, baseline_tensor),
+        (d_target, tensor_names, target_tensor),
+    ]:
+        dumper: _Dumper = _make_dumper(d)
+        for name in names:
+            dumper.dump(name, tensor)
+        dumper.step()
+        exp_paths.append(d / dumper._config.exp_name)
+
+    return exp_paths[0], exp_paths[1]
+
+
+def _make_args(
+    baseline_path: Path,
+    target_path: Path,
+    *,
+    filter_pattern: str | None = None,
+) -> Namespace:
+    return Namespace(
+        baseline_path=str(baseline_path),
+        target_path=str(target_path),
+        start_step=0,
+        end_step=1000000,
+        diff_threshold=1e-3,
+        filter=filter_pattern,
+    )
+
+
+class TestMainBasic:
+    def test_matching_tensors(
+        self, tmp_path: Path, capsys: pytest.CaptureFixture
+    ) -> None:
+        baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
+        args: Namespace = _make_args(baseline_path, target_path)
+
+        main(args)
+
+        captured: str = capsys.readouterr().out
+        assert "✅" in captured
+
+    def test_with_filter(self, tmp_path: Path, capsys: pytest.CaptureFixture) -> None:
+        baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
+        args: Namespace = _make_args(
+            baseline_path, target_path, filter_pattern="tensor_a"
+        )
+
+        main(args)
+
+        captured: str = capsys.readouterr().out
+        assert "tensor_a" in captured
+        assert "Check:" in captured
+
+    def test_no_match_skips(
+        self, tmp_path: Path, capsys: pytest.CaptureFixture
+    ) -> None:
+        baseline_path, target_path = _create_dumps(
+            tmp_path,
+            ["only_in_target"],
+            baseline_names=["only_in_baseline"],
+        )
+        args: Namespace = _make_args(baseline_path, target_path)
+
+        main(args)
+
+        captured: str = capsys.readouterr().out
+        assert "Skip" in captured

test/registered/debug_utils/test_dumper.py

@@ -451,7 +451,6 @@ class TestDumperDistributed:
             config=DumperConfig(
                 enable=True,
                 collective_timeout=3,
-                enable_http_server=False,
             ),
         )
 
@@ -663,12 +662,11 @@ class TestDumpDictFormat:
 
 
 def _make_test_dumper(tmp_path, **overrides) -> _Dumper:
-    """Create a _Dumper for CPU testing without HTTP server or distributed."""
+    """Create a _Dumper for CPU testing without distributed."""
     defaults = dict(
         enable=True,
         dir=str(tmp_path),
         exp_name="test",
-        enable_http_server=False,
     )
     defaults.update(overrides)
     config = DumperConfig(**defaults)
@@ -2293,7 +2291,6 @@ class TestDumperDims:
             config=DumperConfig(
                 enable=True,
                 dir=str(tmp_path),
-                enable_http_server=False,
                 enable_grad=True,
             )
         )
@@ -2323,7 +2320,6 @@ class TestDumperDims:
             config=DumperConfig(
                 enable=True,
                 dir=str(tmp_path),
-                enable_http_server=False,
                 enable_grad=True,
             )
         )
@@ -2340,5 +2336,62 @@ class TestDumperDims:
         assert grad_data["meta"]["dims"] == "b h(tp)"
 
 
+class TestCtxDecorator:
+    def test_ctx_dynamic_lambda(self, tmp_path: Path) -> None:
+        d = _make_test_dumper(tmp_path)
+
+        class FakeLayer:
+            def __init__(self, layer_id: int) -> None:
+                self.layer_id = layer_id
+
+            @d.ctx(lambda self: dict(layer_id=self.layer_id))
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                d.dump("hidden", x)
+                return x
+
+        layer = FakeLayer(layer_id=42)
+        layer.forward(torch.randn(3))
+
+        filenames = _get_filenames(tmp_path)
+        _assert_files(filenames, exist=["layer_id=42"])
+
+    def test_ctx_static_kwargs(self, tmp_path: Path) -> None:
+        d = _make_test_dumper(tmp_path)
+
+        @d.ctx(phase="decode")
+        def decode_step(x: torch.Tensor) -> torch.Tensor:
+            d.dump("step_out", x)
+            return x
+
+        decode_step(torch.randn(3))
+
+        filenames = _get_filenames(tmp_path)
+        _assert_files(filenames, exist=["phase=decode"])
+
+    def test_ctx_clears_on_exception(self, tmp_path: Path) -> None:
+        d = _make_test_dumper(tmp_path)
+
+        @d.ctx(phase="train")
+        def buggy_fn() -> None:
+            raise RuntimeError("boom")
+
+        with pytest.raises(RuntimeError, match="boom"):
+            buggy_fn()
+
+        assert d._state.global_ctx == {}
+
+    def test_ctx_rejects_mixed_args(self) -> None:
+        d = _make_test_dumper("/tmp")
+
+        with pytest.raises(ValueError, match="cannot mix"):
+            d.ctx(lambda self: dict(a=1), phase="x")
+
+    def test_ctx_rejects_empty_args(self) -> None:
+        d = _make_test_dumper("/tmp")
+
+        with pytest.raises(ValueError, match="must provide"):
+            d.ctx()
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__]))