Handle warnings via sink for structured output and add pair in dump comparator (#19373)

python/sglang/srt/debug_utils/comparator/aligner/unshard/executor.py

@@ -8,7 +8,7 @@ from sglang.srt.debug_utils.comparator.aligner.unshard.types import (
 )
 from sglang.srt.debug_utils.comparator.dims import ParallelAxis
 from sglang.srt.debug_utils.comparator.output_types import (
-    AlignWarning,
+    AnyWarning,
     ReplicatedMismatchWarning,
 )
 
@@ -16,8 +16,8 @@ from sglang.srt.debug_utils.comparator.output_types import (
 def execute_unshard_plan(
     plan: UnshardPlan,
     tensors: list[torch.Tensor],
-) -> tuple[list[torch.Tensor], list[AlignWarning]]:
-    all_warnings: list[AlignWarning] = []
+) -> tuple[list[torch.Tensor], list[AnyWarning]]:
+    all_warnings: list[AnyWarning] = []
     result: list[torch.Tensor] = []
 
     for group_idx, group in enumerate(plan.groups):
@@ -40,7 +40,7 @@ def _apply_unshard(
     *,
     axis: ParallelAxis,
     group_index: int,
-) -> tuple[torch.Tensor, list[AlignWarning]]:
+) -> tuple[torch.Tensor, list[AnyWarning]]:
     if isinstance(params, PickParams):
         warnings = _verify_replicated_group(
             ordered_tensors,

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

@@ -31,13 +31,7 @@ def run(args: argparse.Namespace) -> None:
     assert all(c in df_target.columns for c in ["rank", "step", "dump_index", "name"])
 
     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,
-        ),
+        ConfigRecord.from_args(args),
         output_format=args.output_format,
     )
 

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

@@ -1,7 +1,7 @@
 from abc import abstractmethod
-from typing import Annotated, Literal, Union
+from typing import Annotated, Any, Literal, Union
 
-from pydantic import Discriminator, Field, TypeAdapter
+from pydantic import Discriminator, Field, TypeAdapter, model_validator
 
 from sglang.srt.debug_utils.comparator.tensor_comparison.formatter import (
     format_comparison,
@@ -28,38 +28,45 @@ class ReplicatedMismatchWarning(_StrictBase):
         )
 
 
-AlignWarning = (
-    ReplicatedMismatchWarning  # future: Annotated[Union[...], Discriminator("kind")]
-)
+class GeneralWarning(_StrictBase):
+    kind: Literal["general"] = "general"
+    category: str
+    message: str
+
+    def to_text(self) -> str:
+        return self.message
+
+
+AnyWarning = Annotated[
+    Union[ReplicatedMismatchWarning, GeneralWarning],
+    Discriminator("kind"),
+]
 
 
 class _OutputRecord(_StrictBase):
-    align_warnings: list[AlignWarning] = Field(default_factory=list)
+    warnings: list[AnyWarning] = Field(default_factory=list)
 
     @abstractmethod
     def _format_body(self) -> str: ...
 
     def to_text(self) -> str:
         body = self._format_body()
-        if self.align_warnings:
-            body += "\n" + "\n".join(f"  ⚠ {w.to_text()}" for w in self.align_warnings)
+        if self.warnings:
+            body += "\n" + "\n".join(f"  ⚠ {w.to_text()}" for w in self.warnings)
         return body
 
 
 class ConfigRecord(_OutputRecord):
     type: Literal["config"] = "config"
-    baseline_path: str
-    target_path: str
-    diff_threshold: float
-    start_step: int
-    end_step: int
+    config: dict[str, Any]
+
+    @classmethod
+    def from_args(cls, args) -> "ConfigRecord":
+        """Create ConfigRecord from argparse.Namespace."""
+        return cls(config=vars(args))
 
     def _format_body(self) -> str:
-        return (
-            f"Config: baseline={self.baseline_path} target={self.target_path}\n"
-            f"diff_threshold={self.diff_threshold} "
-            f"steps=[{self.start_step}, {self.end_step}]"
-        )
+        return f"Config: {self.config}"
 
 
 class SkipRecord(_OutputRecord):
@@ -69,7 +76,7 @@ class SkipRecord(_OutputRecord):
 
     @property
     def category(self) -> str:
-        if self.align_warnings:
+        if self.warnings:
             return "failed"
         return "skipped"
 
@@ -82,7 +89,7 @@ class ComparisonRecord(TensorComparisonInfo, _OutputRecord):
 
     @property
     def category(self) -> str:
-        if self.align_warnings:
+        if self.warnings:
             return "failed"
         return "passed" if self.diff is not None and self.diff.passed else "failed"
 
@@ -97,6 +104,15 @@ class SummaryRecord(_OutputRecord):
     failed: int
     skipped: int
 
+    @model_validator(mode="after")
+    def _validate_totals(self) -> "SummaryRecord":
+        expected: int = self.passed + self.failed + self.skipped
+        if self.total != expected:
+            raise ValueError(
+                f"total={self.total} != passed({self.passed}) + failed({self.failed}) + skipped({self.skipped}) = {expected}"
+            )
+        return self
+
     def _format_body(self) -> str:
         return (
             f"Summary: {self.passed} passed, {self.failed} failed, "
@@ -104,8 +120,15 @@ class SummaryRecord(_OutputRecord):
         )
 
 
+class WarningRecord(_OutputRecord):
+    type: Literal["warning"] = "warning"
+
+    def _format_body(self) -> str:
+        return ""
+
+
 AnyRecord = Annotated[
-    Union[ConfigRecord, SkipRecord, ComparisonRecord, SummaryRecord],
+    Union[ConfigRecord, SkipRecord, ComparisonRecord, SummaryRecord, WarningRecord],
     Discriminator("type"),
 ]
 

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

@@ -20,7 +20,7 @@ from sglang.srt.debug_utils.comparator.aligner.unshard.planner import (
 from sglang.srt.debug_utils.comparator.aligner.unshard.types import UnshardPlan
 from sglang.srt.debug_utils.comparator.dims import parse_dims
 from sglang.srt.debug_utils.comparator.output_types import (
-    AlignWarning,
+    AnyWarning,
     ComparisonRecord,
     SkipRecord,
 )
@@ -53,11 +53,11 @@ def process_tensor_group(
 
     b_tensor, b_warns = _execute_plans(b_extracted, b_plans)
     t_tensor, t_warns = _execute_plans(t_extracted, t_plans)
-    all_warnings: list[AlignWarning] = b_warns + t_warns
+    all_warnings: list[AnyWarning] = b_warns + t_warns
 
     if b_tensor is None or t_tensor is None:
         reason = "baseline_load_failed" if b_tensor is None else "target_load_failed"
-        return SkipRecord(name=name, reason=reason, align_warnings=all_warnings)
+        return SkipRecord(name=name, reason=reason, warnings=all_warnings)
 
     info = compare_tensors(
         x_baseline=b_tensor,
@@ -66,7 +66,7 @@ def process_tensor_group(
         diff_threshold=diff_threshold,
     )
 
-    return ComparisonRecord(**info.model_dump(), align_warnings=all_warnings)
+    return ComparisonRecord(**info.model_dump(), warnings=all_warnings)
 
 
 def _load_tensors(filenames: list[str], base_path: Path) -> list[ValueWithMeta]:
@@ -114,7 +114,7 @@ def _extract_tensors(
 def _execute_plans(
     tensors: list[torch.Tensor],
     plans: list[Plan],
-) -> tuple[Optional[torch.Tensor], list[AlignWarning]]:
+) -> tuple[Optional[torch.Tensor], list[AnyWarning]]:
     if not tensors:
         return None, []
 
@@ -123,7 +123,7 @@ def _execute_plans(
             return None, []
         return tensors[0], []
 
-    warnings: list[AlignWarning] = []
+    warnings: list[AnyWarning] = []
     current = tensors
     for plan in plans:
         current, new_warnings = _execute_plan(current, plan)
@@ -136,7 +136,7 @@ def _execute_plans(
 def _execute_plan(
     tensors: list[torch.Tensor],
     plan: Plan,
-) -> tuple[list[torch.Tensor], list[AlignWarning]]:
+) -> tuple[list[torch.Tensor], list[AnyWarning]]:
     if isinstance(plan, UnshardPlan):
         return execute_unshard_plan(plan, tensors)
     elif isinstance(plan, ReorderPlan):

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

@@ -1,9 +1,22 @@
+from __future__ import annotations
+
 import functools
-from typing import Optional, Tuple
+from typing import Callable, Generic, Optional, Tuple, TypeVar
 
 import torch
 from pydantic import BaseModel, ConfigDict
 
+_T = TypeVar("_T")
+_U = TypeVar("_U")
+
+
+def _check_equal_lengths(**named_lists: list) -> None:
+    lengths: dict[str, int] = {name: len(lst) for name, lst in named_lists.items()}
+    unique: set[int] = set(lengths.values())
+    if len(unique) > 1:
+        details: str = ", ".join(f"{name}={length}" for name, length in lengths.items())
+        raise ValueError(f"Length mismatch: {details}")
+
 
 class _StrictBase(BaseModel):
     model_config = ConfigDict(extra="forbid")
@@ -13,6 +26,14 @@ class _FrozenBase(BaseModel):
     model_config = ConfigDict(frozen=True, extra="forbid")
 
 
+class Pair(_FrozenBase, Generic[_T]):
+    x: _T
+    y: _T
+
+    def map(self, fn: Callable[[_T], _U]) -> Pair[_U]:
+        return Pair(x=fn(self.x), y=fn(self.y))
+
+
 def argmax_coord(x: torch.Tensor) -> Tuple[int, ...]:
     flat_idx = x.argmax()
     return tuple(idx.item() for idx in torch.unravel_index(flat_idx, x.shape))

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

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Generator
+
+from sglang.srt.debug_utils.comparator.output_types import AnyWarning
+
+
+class WarningSink:
+    def __init__(self) -> None:
+        self._stack: list[list[AnyWarning]] = []
+        self._output_format: str = "text"
+
+    def set_output_format(self, output_format: str) -> None:
+        self._output_format = output_format
+
+    @contextmanager
+    def context(self) -> Generator[list[AnyWarning], None, None]:
+        bucket: list[AnyWarning] = []
+        self._stack.append(bucket)
+        try:
+            yield bucket
+        finally:
+            popped = self._stack.pop()
+            assert popped is bucket
+
+    def add(self, warning: AnyWarning) -> None:
+        if self._stack:
+            self._stack[-1].append(warning)
+        else:
+            from sglang.srt.debug_utils.comparator.output_types import (
+                WarningRecord,
+                print_record,
+            )
+
+            print_record(
+                WarningRecord(warnings=[warning]),
+                output_format=self._output_format,
+            )
+
+
+warning_sink = WarningSink()