Support agent-friendly output formats in dump comparator (#19275)

test/registered/debug_utils/comparator/test_entrypoint.py

@@ -6,6 +6,15 @@ import pytest
 import torch
 
 from sglang.srt.debug_utils.comparator.entrypoint import run
+from sglang.srt.debug_utils.comparator.output_types import (
+    AnyRecord,
+    ComparisonRecord,
+    ConfigRecord,
+    SkipRecord,
+    SummaryRecord,
+    _OutputRecord,
+    parse_record_json,
+)
 from sglang.srt.debug_utils.dumper import DumperConfig, _Dumper
 from sglang.test.ci.ci_register import register_cpu_ci
 
@@ -65,36 +74,39 @@ def _make_args(baseline_path: Path, target_path: Path, **overrides) -> Namespace
         end_step=1000000,
         diff_threshold=1e-3,
         filter=None,
+        output_format="text",
     )
     defaults.update(overrides)
     return Namespace(**defaults)
 
 
+def _parse_jsonl(output: str) -> list[AnyRecord]:
+    return [parse_record_json(line) for line in output.strip().splitlines()]
+
+
 class TestEntrypoint:
     def test_run_basic(self, tmp_path, capsys):
         baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
         args = _make_args(baseline_path, target_path)
+        capsys.readouterr()
 
         run(args)
 
         output = capsys.readouterr().out
-        assert "df_target" in output
-        assert "df_baseline" in output
-        assert output.count("Check:") == 2
-        assert "tensor_a" in output
-        assert "tensor_b" in output
+        assert "Config:" in output
         assert "rel_diff" in output
+        assert "Summary:" in output
         assert "Skip" not in output
 
     def test_filter(self, tmp_path, capsys):
         baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
         args = _make_args(baseline_path, target_path, filter="tensor_a")
+        capsys.readouterr()
 
         run(args)
 
         output = capsys.readouterr().out
-        assert output.count("Check:") == 1
-        assert "tensor_a" in output
+        assert "rel_diff" in output
 
     def test_no_baseline_skip(self, tmp_path, capsys):
         baseline_path, target_path = _create_dumps(
@@ -103,22 +115,73 @@ class TestEntrypoint:
             baseline_names=["tensor_a"],
         )
         args = _make_args(baseline_path, target_path)
+        capsys.readouterr()
 
         run(args)
 
         output = capsys.readouterr().out
-        assert output.count("Check:") == 1
         assert "Skip:" in output
-        assert "since no baseline" in output
+        assert "no_baseline" in output
 
     def test_step_range(self, tmp_path, capsys):
         baseline_path, target_path = _create_dumps(tmp_path, ["t"], num_steps=3)
         args = _make_args(baseline_path, target_path, start_step=1, end_step=1)
+        capsys.readouterr()
 
         run(args)
 
         output = capsys.readouterr().out
-        assert output.count("Check:") == 1
+        assert "Summary:" in output
+
+
+class TestEntrypointJsonl:
+    def test_jsonl_basic(self, tmp_path, capsys):
+        baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
+        args = _make_args(baseline_path, target_path, output_format="json")
+        capsys.readouterr()
+
+        run(args)
+
+        records = _parse_jsonl(capsys.readouterr().out)
+        assert isinstance(records[0], ConfigRecord)
+
+        comparisons = [r for r in records if isinstance(r, ComparisonRecord)]
+        assert len(comparisons) == 2
+
+        summary = records[-1]
+        assert isinstance(summary, SummaryRecord)
+        assert summary.total == 2
+        assert summary.skipped == 0
+
+    def test_jsonl_skip(self, tmp_path, capsys):
+        baseline_path, target_path = _create_dumps(
+            tmp_path,
+            tensor_names=["tensor_a", "tensor_extra"],
+            baseline_names=["tensor_a"],
+        )
+        args = _make_args(baseline_path, target_path, output_format="json")
+        capsys.readouterr()
+
+        run(args)
+
+        records = _parse_jsonl(capsys.readouterr().out)
+        skips = [r for r in records if isinstance(r, SkipRecord)]
+        assert len(skips) == 1
+        assert skips[0].reason == "no_baseline"
+
+        summary = records[-1]
+        assert isinstance(summary, SummaryRecord)
+        assert summary.skipped == 1
+
+    def test_jsonl_all_valid_records(self, tmp_path, capsys):
+        baseline_path, target_path = _create_dumps(tmp_path, ["t"], num_steps=2)
+        args = _make_args(baseline_path, target_path, output_format="json")
+        capsys.readouterr()
+
+        run(args)
+
+        records = _parse_jsonl(capsys.readouterr().out)
+        assert all(isinstance(r, _OutputRecord) for r in records)
 
 
 if __name__ == "__main__":