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

python/sglang/srt/debug_utils/comparator/entrypoint.py

@@ -3,13 +3,16 @@ from pathlib import Path
 
 import polars as pl
 
-from sglang.srt.debug_utils.comparator.tensor_comparison import (
-    compare_tensors,
-    print_comparison,
+from sglang.srt.debug_utils.comparator.output_types import (
+    ComparisonRecord,
+    ConfigRecord,
+    SkipRecord,
+    SummaryRecord,
+    print_record,
 )
+from sglang.srt.debug_utils.comparator.tensor_comparison import compare_tensors
 from sglang.srt.debug_utils.comparator.utils import load_object
 from sglang.srt.debug_utils.dump_loader import find_row, read_meta
-from sglang.srt.debug_utils.dumper import get_truncated_value
 
 
 def main() -> None:
@@ -27,8 +30,19 @@ def run(args: argparse.Namespace) -> None:
     assert all(c in df_target.columns for c in ["rank", "step", "dump_index", "name"])
 
     df_baseline = read_meta(args.baseline_path)
-    print("df_target", df_target)
-    print("df_baseline", df_baseline)
+
+    print_record(
+        ConfigRecord(
+            baseline_path=args.baseline_path,
+            target_path=args.target_path,
+            diff_threshold=args.diff_threshold,
+            start_step=args.start_step,
+            end_step=args.end_step,
+        ),
+        output_format=args.output_format,
+    )
+
+    counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0}
 
     for row in df_target.iter_rows(named=True):
         path_target = Path(args.target_path) / row["filename"]
@@ -47,26 +61,23 @@ def run(args: argparse.Namespace) -> None:
         )
 
         if row_baseline is None:
-            print(f"Skip: target={str(path_target)} since no baseline")
-            x_target = load_object(path_target)
-            if x_target is not None:
-                print(f"x_target(sample)={get_truncated_value(x_target)}")
+            counts["skipped"] += 1
+            print_record(
+                SkipRecord(name=row["name"], reason="no_baseline"),
+                output_format=args.output_format,
+            )
             continue
 
         path_baseline = Path(args.baseline_path) / row_baseline["filename"]
-        print(
-            f"Check:\n"
-            f"target={str(path_target)} (duplicate_index={row['duplicate_index']})\n"
-            f"baseline={str(path_baseline)} (duplicate_index={row_baseline['duplicate_index']})"
-        )
 
         x_baseline = load_object(path_baseline)
         x_target = load_object(path_target)
 
         if x_baseline is None or x_target is None:
-            print(
-                f"Skip comparison because of None: "
-                f"x_baseline={x_baseline}, x_target={x_target}"
+            counts["skipped"] += 1
+            print_record(
+                SkipRecord(name=row["name"], reason="load_failed"),
+                output_format=args.output_format,
             )
             continue
 
@@ -74,9 +85,23 @@ def run(args: argparse.Namespace) -> None:
             x_baseline=x_baseline,
             x_target=x_target,
             name=row["name"],
+            diff_threshold=args.diff_threshold,
         )
-        print_comparison(info=info, diff_threshold=args.diff_threshold)
-        print()
+
+        if info.diff is not None and info.diff.passed:
+            counts["passed"] += 1
+        else:
+            counts["failed"] += 1
+
+        print_record(
+            ComparisonRecord(**info.model_dump()),
+            output_format=args.output_format,
+        )
+
+    print_record(
+        SummaryRecord(total=sum(counts.values()), **counts),
+        output_format=args.output_format,
+    )
 
 
 def _parse_args() -> argparse.Namespace:
@@ -90,4 +115,11 @@ def _parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--filter", type=str, default=None, help="Regex to filter filenames"
     )
+    parser.add_argument(
+        "--output-format",
+        type=str,
+        choices=["text", "json"],
+        default="text",
+        help="Output format: text (default) or json (JSONL, one JSON object per line)",
+    )
     return parser.parse_args()