Support multiple verbosity in dump comparator (#19684)

This commit is contained in:
fzyzcjy
2026-03-02 18:47:30 +08:00
committed by GitHub
parent 3dd4649b42
commit e5ef845cad
7 changed files with 519 additions and 249 deletions

View File

@@ -6,6 +6,7 @@ from pathlib import Path
from typing import TYPE_CHECKING, Any, Optional
import polars as pl
import rich.table
if TYPE_CHECKING:
from rich.table import Table
@@ -68,6 +69,19 @@ def _build_rich_table(df: pl.DataFrame, *, title: Optional[str] = None) -> "Tabl
return table
def _render_polars_as_rich_table(
df: pl.DataFrame, *, title: Optional[str] = None
) -> "rich.table.Table":
from rich.table import Table
table = Table(title=title)
for col in df.columns:
table.add_column(col)
for row in df.iter_rows():
table.add_row(*[str(v) for v in row])
return table
def _collect_rank_info(
df: pl.DataFrame, dump_dir: Path
) -> Optional[list[dict[str, Any]]]:

View File

@@ -6,7 +6,7 @@ from rendering / formatting logic.
from __future__ import annotations
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Literal
from rich.console import Group
from rich.markup import escape
@@ -38,12 +38,16 @@ if TYPE_CHECKING:
_TableRecord,
)
Verbosity = Literal["minimal", "normal", "verbose"]
# ── Record-level rendering (body + logs) ─────────────────────────────
def _render_record_rich(record: _OutputRecord) -> RenderableType:
body: RenderableType = record._format_rich_body()
def _render_record_rich(
record: _OutputRecord, *, verbosity: Verbosity = "normal"
) -> RenderableType:
body: RenderableType = record._format_rich_body(verbosity=verbosity)
log_lines: list[str] = _format_log_lines_rich(
errors=record.errors, infos=record.infos
@@ -100,7 +104,9 @@ def _format_config_body(record: ConfigRecord) -> str:
return f"Config: {record.config}"
def _format_config_rich_body(record: ConfigRecord) -> RenderableType:
def _format_config_rich_body(
record: ConfigRecord, verbosity: Verbosity = "normal"
) -> RenderableType:
lines: list[str] = [f" [bold]{k}[/] : {v}" for k, v in record.config.items()]
return Panel("\n".join(lines), title="Comparator Config", border_style="cyan")
@@ -112,7 +118,9 @@ def _format_skip_body(record: SkipComparisonRecord) -> str:
return f"Skip: {record.name}{record._format_location_suffix()} ({record.reason})"
def _format_skip_rich_body(record: SkipComparisonRecord) -> RenderableType:
def _format_skip_rich_body(
record: SkipComparisonRecord, verbosity: Verbosity = "normal"
) -> RenderableType:
suffix: str = record._format_location_suffix()
return (
f"[dim]⊘ {escape(record.name)}{suffix} ── skipped ({escape(record.reason)})[/]"
@@ -132,7 +140,9 @@ def _format_table_body(record: _TableRecord) -> str:
)
def _format_table_rich_body(record: _TableRecord) -> RenderableType:
def _format_table_rich_body(
record: _TableRecord, verbosity: Verbosity = "normal"
) -> RenderableType:
import polars as pl
from sglang.srt.debug_utils.comparator.display import (
@@ -157,13 +167,15 @@ def _format_tensor_comparison_body(record: TensorComparisonRecord) -> str:
def _format_tensor_comparison_rich_body(
record: TensorComparisonRecord,
record: TensorComparisonRecord, verbosity: Verbosity = "normal"
) -> RenderableType:
from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import (
format_comparison_rich,
)
return record._format_location_prefix_rich() + format_comparison_rich(record=record)
return record._format_location_prefix_rich() + format_comparison_rich(
record=record, verbosity=verbosity
)
# ── NonTensorComparisonRecord ────────────────────────────────────────
@@ -181,7 +193,7 @@ def _format_non_tensor_body(record: NonTensorComparisonRecord) -> str:
def _format_non_tensor_rich_body(
record: NonTensorComparisonRecord,
record: NonTensorComparisonRecord, verbosity: Verbosity = "normal"
) -> RenderableType:
suffix: str = record._format_location_suffix()
name: str = escape(record.name)
@@ -210,7 +222,9 @@ def _format_summary_body(record: SummaryRecord) -> str:
)
def _format_summary_rich_body(record: SummaryRecord) -> RenderableType:
def _format_summary_rich_body(
record: SummaryRecord, verbosity: Verbosity = "normal"
) -> RenderableType:
text: str = (
f"[bold green]{record.passed} passed[/] │ "
f"[bold red]{record.failed} failed[/] │ "

View File

@@ -7,6 +7,26 @@ from pydantic import ConfigDict, Discriminator, Field, TypeAdapter, model_valida
from rich.console import RenderableType
from rich.markup import escape
from sglang.srt.debug_utils.comparator.output_formatter import ( # noqa: F401 — re-export
_format_aligner_plan as _format_aligner_plan,
)
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_config_body,
_format_config_rich_body,
_format_log_body,
_format_non_tensor_body,
_format_non_tensor_rich_body,
_format_skip_body,
_format_skip_rich_body,
_format_summary_body,
_format_summary_rich_body,
_format_table_body,
_format_table_rich_body,
_format_tensor_comparison_body,
_format_tensor_comparison_rich_body,
_render_record_rich,
_render_record_text,
)
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorComparisonInfo,
@@ -17,7 +37,6 @@ if TYPE_CHECKING:
from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
TracedAlignerPlan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan
from sglang.srt.debug_utils.comparator.report_sink import Verbosity
@@ -83,21 +102,13 @@ class _OutputRecord(_StrictBase):
@abstractmethod
def _format_body(self) -> str: ...
def _format_rich_body(self) -> RenderableType:
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return self._format_body()
def to_rich(self) -> RenderableType:
from sglang.srt.debug_utils.comparator.output_formatter import (
_render_record_rich,
)
return _render_record_rich(self)
def to_rich(self, verbosity: Verbosity = "normal") -> RenderableType:
return _render_record_rich(self, verbosity=verbosity)
def to_text(self) -> str:
from sglang.srt.debug_utils.comparator.output_formatter import (
_render_record_text,
)
return _render_record_text(self)
@@ -124,46 +135,15 @@ class _BaseComparisonRecord(_OutputRecord):
return ""
class RecordLocation(_StrictBase):
step: Optional[int] = None
class _BaseComparisonRecord(_OutputRecord):
location: RecordLocation = Field(default_factory=RecordLocation)
def _format_location_prefix(self) -> str:
if self.location.step is not None:
return f"[step={self.location.step}] "
return ""
def _format_location_prefix_rich(self) -> str:
if self.location.step is not None:
return escape(f"[step={self.location.step}]") + " "
return ""
def _format_location_suffix(self) -> str:
if self.location.step is not None:
return f" (step={self.location.step})"
return ""
class ConfigRecord(_OutputRecord):
type: Literal["config"] = "config"
config: dict[str, Any]
def _format_body(self) -> str:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_config_body,
)
return _format_config_body(self)
def _format_rich_body(self) -> RenderableType:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_config_rich_body,
)
return _format_config_rich_body(self)
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return _format_config_rich_body(self, verbosity=verbosity)
class SkipComparisonRecord(_BaseComparisonRecord):
@@ -178,18 +158,10 @@ class SkipComparisonRecord(_BaseComparisonRecord):
return "skipped"
def _format_body(self) -> str:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_skip_body,
)
return _format_skip_body(self)
def _format_rich_body(self) -> RenderableType:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_skip_rich_body,
)
return _format_skip_rich_body(self)
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return _format_skip_rich_body(self, verbosity=verbosity)
class _TableRecord(_OutputRecord):
@@ -200,9 +172,10 @@ class _TableRecord(_OutputRecord):
def _table_title(self) -> str: ...
def _format_body(self) -> str:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_table_body,
)
return _format_table_body(self)
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return _format_table_rich_body(self, verbosity=verbosity)
return _format_table_body(self)
@@ -245,18 +218,10 @@ class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord):
return "passed" if self.diff is not None and self.diff.passed else "failed"
def _format_body(self) -> str:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_tensor_comparison_body,
)
return _format_tensor_comparison_body(self)
def _format_rich_body(self) -> RenderableType:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_tensor_comparison_rich_body,
)
return _format_tensor_comparison_rich_body(self)
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return _format_tensor_comparison_rich_body(self, verbosity=verbosity)
class NonTensorComparisonRecord(_BaseComparisonRecord):
@@ -275,9 +240,10 @@ class NonTensorComparisonRecord(_BaseComparisonRecord):
return "passed" if self.values_equal else "failed"
def _format_body(self) -> str:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_non_tensor_body,
)
return _format_non_tensor_body(self)
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return _format_non_tensor_rich_body(self, verbosity=verbosity)
return _format_non_tensor_body(self)
@@ -306,40 +272,20 @@ class SummaryRecord(_OutputRecord):
return self
def _format_body(self) -> str:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_summary_body,
)
return _format_summary_body(self)
def _format_rich_body(self) -> RenderableType:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_summary_rich_body,
)
return _format_summary_rich_body(self)
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return _format_summary_rich_body(self, verbosity=verbosity)
return _format_summary_body(self)
class LogRecord(_OutputRecord):
type: Literal["log"] = "log"
def _format_body(self) -> str:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_log_body,
)
return _format_log_body(self)
# Re-export _format_aligner_plan for backward compatibility (used by tests)
def _format_aligner_plan(traced_plan: TracedAlignerPlan) -> str:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_aligner_plan as _impl,
)
return _impl(traced_plan)
AnyRecord = Annotated[
Union[
ConfigRecord,

View File

@@ -80,7 +80,7 @@ class ReportSink:
print(record.model_dump_json())
else:
console: Console = self._get_console()
console.print(record.to_rich())
console.print(record.to_rich(verbosity=self._verbosity))
console.print() # blank line between records

View File

@@ -1,6 +1,6 @@
from __future__ import annotations
from typing import TYPE_CHECKING, Optional
from typing import TYPE_CHECKING, Literal, Optional
from rich.markup import escape
@@ -26,6 +26,8 @@ if TYPE_CHECKING:
)
from sglang.srt.debug_utils.comparator.utils import Pair
Verbosity = Literal["minimal", "normal", "verbose"]
def _esc_shape(shape: Optional[list[int]]) -> str:
return escape(str(shape))
@@ -64,7 +66,7 @@ def _category_marker(category: str) -> tuple[bool, str, str]:
# ---------------------------------------------------------------------------
# Stats formatting helpers
# Stats formatting helpers (shared between compact / verbose)
# ---------------------------------------------------------------------------
@@ -204,11 +206,36 @@ def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]:
# ---------------------------------------------------------------------------
def format_comparison_rich(record: TensorComparisonRecord) -> str:
return _format_comparison_normal(record=record)
def format_comparison_rich(
record: TensorComparisonRecord,
verbosity: Verbosity = "normal",
) -> str:
if verbosity == "minimal":
return _format_comparison_minimal(record)
return _format_comparison_normal_or_verbose(
record=record,
verbose=(verbosity == "verbose"),
)
def _format_comparison_normal(*, record: TensorComparisonRecord) -> str:
def _format_comparison_minimal(record: TensorComparisonRecord) -> str:
passed, color, marker = _category_marker(record.category)
name_part: str = f"[bold {color}]{escape(record.name):30s}[/]"
if record.diff is not None:
return f"{marker} {name_part} rel_diff={_fmt_val(record.diff.rel_diff)}"
elif record.shape_mismatch:
return f"{marker} {name_part} [yellow]shape mismatch[/]"
else:
return f"{marker} {name_part}"
def _format_comparison_normal_or_verbose(
*,
record: TensorComparisonRecord,
verbose: bool,
) -> str:
passed, color, marker = _category_marker(record.category)
baseline: TensorInfo = record.baseline
@@ -254,12 +281,19 @@ def _format_comparison_normal(*, record: TensorComparisonRecord) -> str:
# Bundle section
if record.raw_bundle_info is not None:
lines.append(" [dim]Bundle[/]")
lines.extend(_format_bundle_section(bundle_info=record.raw_bundle_info))
lines.extend(
_format_bundle_section(bundle_info=record.raw_bundle_info, verbose=verbose)
)
# Plan section
if record.traced_plan is not None:
lines.append(" [dim]Plan[/]")
lines.extend(_format_plan_section_rich(traced_plan=record.traced_plan))
lines.extend(
_format_plan_section_rich(
traced_plan=record.traced_plan,
verbose=verbose,
)
)
# Aligned section
lines.append(" [dim]Aligned[/]")
@@ -270,22 +304,28 @@ def _format_comparison_normal(*, record: TensorComparisonRecord) -> str:
# Stats section
lines.append(" [dim]Stats[/]")
lines.extend(_format_stats_rich(baseline=baseline.stats, target=target.stats))
lines.extend(
_format_stats_rich(
baseline=baseline.stats, target=target.stats, verbose=verbose
)
)
# Abs diff percentiles (show when failed)
if not passed and record.diff is not None and record.diff.abs_diff_percentiles:
show_detail: bool = verbose or not passed
# Abs diff percentiles
if show_detail and record.diff is not None and record.diff.abs_diff_percentiles:
lines.append(" [dim]Abs Diff Percentiles[/]")
lines.append(" " + _format_abs_diff_percentiles_rich(record.diff))
# Samples (show when failed)
if not passed and baseline.sample is not None:
# Samples
if show_detail and baseline.sample is not None:
lines.append(" [dim]Samples[/]")
lines.append(f" baseline {escape(baseline.sample)}")
if target.sample is not None:
lines.append(f" target {escape(target.sample)}")
# Replicated checks (show when failed)
if not passed and record.replicated_checks:
# Replicated checks
if show_detail and record.replicated_checks:
lines.append(" [dim]Replicated Checks[/]")
for check in record.replicated_checks:
chk_marker: str = "[green]✅[/]" if check.passed else "[red]❌[/]"
@@ -305,7 +345,9 @@ def _format_comparison_normal(*, record: TensorComparisonRecord) -> str:
return "\n".join(lines)
def _format_bundle_section(bundle_info: Pair[BundleSideInfo]) -> list[str]:
def _format_bundle_section(
bundle_info: Pair[BundleSideInfo], *, verbose: bool = False
) -> list[str]:
lines: list[str] = []
for label, side in [("baseline", bundle_info.x), ("target", bundle_info.y)]:
@@ -315,19 +357,37 @@ def _format_bundle_section(bundle_info: Pair[BundleSideInfo]) -> list[str]:
dtype_desc: str = _strip_torch_prefix(side.files[0].dtype)
shapes: list[list[int]] = [f.shape for f in side.files]
unique_shapes: set[str] = {str(s) for s in shapes}
shape_desc: str
if len(unique_shapes) == 1:
shape_desc = _esc_shape(shapes[0])
else:
shape_desc = "mixed shapes"
if verbose:
dims_part: str = f" dims: {side.dims}" if side.dims else ""
lines.append(
f" {label} [cyan]{side.num_files} files[/]"
f" {dtype_desc}{dims_part}"
)
dims_part: str = f" [dim]dims: {side.dims}[/]" if side.dims else ""
lines.append(
f" {label} [cyan]{side.num_files} files[/]"
f" × {shape_desc} {dtype_desc}{dims_part}"
)
for idx, f in enumerate(side.files):
rank_part: str = f"rank={f.rank}" if f.rank is not None else ""
par_part: str = ""
if f.parallel_info:
par_part = " " + " ".join(
f"{k}={v}" for k, v in f.parallel_info.items()
)
lines.append(
f" [{idx}] {_esc_shape(f.shape)} {rank_part}{par_part}"
)
else:
shapes: list[list[int]] = [f.shape for f in side.files]
unique_shapes: set[str] = {str(s) for s in shapes}
shape_desc: str
if len(unique_shapes) == 1:
shape_desc = _esc_shape(shapes[0])
else:
shape_desc = "mixed shapes"
dims_part = f" [dim]dims: {side.dims}[/]" if side.dims else ""
lines.append(
f" {label} [cyan]{side.num_files} files[/]"
f" × {shape_desc} {dtype_desc}{dims_part}"
)
return lines
@@ -335,6 +395,7 @@ def _format_bundle_section(bundle_info: Pair[BundleSideInfo]) -> list[str]:
def _format_plan_section_rich(
*,
traced_plan: TracedAlignerPlan,
verbose: bool = False,
) -> list[str]:
lines: list[str] = []
@@ -406,19 +467,35 @@ def _format_stats_rich(
*,
baseline: TensorStats,
target: TensorStats,
verbose: bool = False,
) -> list[str]:
lines: list[str] = []
# Compact: mean, std, range (min/max combined)
for stat_name in ("mean", "std"):
val_b: float = getattr(baseline, stat_name)
val_t: float = getattr(target, stat_name)
lines.append(_format_stat_line(stat_name, val_b, val_t, val_t - val_b))
if verbose:
# All stat fields
for stat_name in TensorStats.model_fields:
if stat_name == "percentiles":
continue
val_b: float = getattr(baseline, stat_name)
val_t: float = getattr(target, stat_name)
lines.append(_format_stat_line(stat_name, val_b, val_t, val_t - val_b))
# Range line: combine min/max (escape brackets to avoid Rich markup)
range_baseline: str = escape(f"[{baseline.min:.4f}, {baseline.max:.4f}]")
range_target: str = escape(f"[{target.min:.4f}, {target.max:.4f}]")
lines.append(f" [blue]{'range':10s}[/] {range_baseline} vs {range_target}")
# Percentiles
for p in sorted(set(baseline.percentiles) & set(target.percentiles)):
val_b = baseline.percentiles[p]
val_t = target.percentiles[p]
lines.append(_format_stat_line(f"p{p}", val_b, val_t, val_t - val_b))
else:
# Compact: mean, std, range (min/max combined)
for stat_name in ("mean", "std"):
val_b = getattr(baseline, stat_name)
val_t = getattr(target, stat_name)
lines.append(_format_stat_line(stat_name, val_b, val_t, val_t - val_b))
# Range line: combine min/max (escape brackets to avoid Rich markup)
range_baseline: str = escape(f"[{baseline.min:.4f}, {baseline.max:.4f}]")
range_target: str = escape(f"[{target.min:.4f}, {target.max:.4f}]")
lines.append(f" [blue]{'range':10s}[/] {range_baseline} vs {range_target}")
return lines