chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Beautify text output in dump comparator (#19683)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-02 18:47:01 +08:00
committed by GitHub
parent 5bf3deb4bc
commit 3dd4649b42
9 changed files with 1460 additions and 193 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
python/sglang/srt/debug_utils/comparator/display.py
View File

@@ -3,10 +3,13 @@ from __future__ import annotations
from collections import defaultdict
from io import StringIO
from pathlib import Path
from typing import Any, Optional
from typing import TYPE_CHECKING, Any, Optional
import polars as pl
if TYPE_CHECKING:
from rich.table import Table
from sglang.srt.debug_utils.comparator.output_types import (
InputIdsRecord,
RankInfoRecord,
@@ -39,6 +42,21 @@ def emit_display_records(
def _render_polars_as_text(df: pl.DataFrame, *, title: Optional[str] = None) -> str:
from rich.console import Console
table = _build_rich_table(df, title=title)
buf = StringIO()
Console(file=buf, force_terminal=False, width=200).print(table)
return buf.getvalue().rstrip("\n")
def _render_polars_as_rich_table(
df: pl.DataFrame, *, title: Optional[str] = None
) -> "Table":
return _build_rich_table(df, title=title)
def _build_rich_table(df: pl.DataFrame, *, title: Optional[str] = None) -> "Table":
from rich.table import Table
table = Table(title=title)
@@ -47,9 +65,7 @@ def _render_polars_as_text(df: pl.DataFrame, *, title: Optional[str] = None) ->
for row in df.iter_rows():
table.add_row(*[str(v) for v in row])
buf = StringIO()
Console(file=buf, force_terminal=False, width=200).print(table)
return buf.getvalue().rstrip("\n")
return table
def _collect_rank_info(

287
python/sglang/srt/debug_utils/comparator/output_formatter.py Normal file
View File

@@ -0,0 +1,287 @@
"""Formatting functions for comparator output records.
Extracted from output_types.py to separate data-structure definitions
from rendering / formatting logic.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from rich.console import Group
from rich.markup import escape
from rich.panel import Panel
from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import (
format_comparison,
format_replicated_checks,
)
if TYPE_CHECKING:
from rich.console import RenderableType
from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
TracedAlignerPlan,
TracedSubPlan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan
from sglang.srt.debug_utils.comparator.output_types import (
ConfigRecord,
ErrorLog,
InfoLog,
LogRecord,
NonTensorComparisonRecord,
SkipComparisonRecord,
SummaryRecord,
TensorComparisonRecord,
_OutputRecord,
_TableRecord,
)
# ── Record-level rendering (body + logs) ─────────────────────────────
def _render_record_rich(record: _OutputRecord) -> RenderableType:
body: RenderableType = record._format_rich_body()
log_lines: list[str] = _format_log_lines_rich(
errors=record.errors, infos=record.infos
)
if not log_lines:
return body
log_block: str = "\n".join(log_lines)
if isinstance(body, str):
return body + "\n" + log_block
return Group(body, log_block)
def _render_record_text(record: _OutputRecord) -> str:
body: str = record._format_body()
log_suffix: str = _format_log_lines_text(errors=record.errors, infos=record.infos)
if log_suffix:
body += "\n" + log_suffix
return body
def _format_log_lines_rich(
*, errors: list[ErrorLog], infos: list[InfoLog]
) -> list[str]:
lines: list[str] = []
if errors:
lines.extend(f" [red]✗ {e.to_text()}[/]" for e in errors)
if infos:
lines.extend(f" [dim]ℹ {i.to_text()}[/]" for i in infos)
return lines
def _format_log_lines_text(*, errors: list[ErrorLog], infos: list[InfoLog]) -> str:
lines: list[str] = []
if errors:
lines.extend(f" ✗ {e.to_text()}" for e in errors)
if infos:
lines.extend(f" ℹ {i.to_text()}" for i in infos)
return "\n".join(lines)
# ── ConfigRecord ──────────────────────────────────────────────────────
def _format_config_body(record: ConfigRecord) -> str:
return f"Config: {record.config}"
def _format_config_rich_body(record: ConfigRecord) -> 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")
# ── SkipComparisonRecord ─────────────────────────────────────────────
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:
suffix: str = record._format_location_suffix()
return (
f"[dim]⊘ {escape(record.name)}{suffix} ── skipped ({escape(record.reason)})[/]"
)
# ── _TableRecord ─────────────────────────────────────────────────────
def _format_table_body(record: _TableRecord) -> str:
import polars as pl
from sglang.srt.debug_utils.comparator.display import _render_polars_as_text
return _render_polars_as_text(
pl.DataFrame(record.rows), title=record._table_title()
)
def _format_table_rich_body(record: _TableRecord) -> RenderableType:
import polars as pl
from sglang.srt.debug_utils.comparator.display import (
_render_polars_as_rich_table,
)
return _render_polars_as_rich_table(
pl.DataFrame(record.rows), title=record._table_title()
)
# ── TensorComparisonRecord ───────────────────────────────────────────
def _format_tensor_comparison_body(record: TensorComparisonRecord) -> str:
body: str = record._format_location_prefix() + format_comparison(record)
if record.replicated_checks:
body += "\n" + format_replicated_checks(record.replicated_checks)
if record.traced_plan is not None:
body += "\n" + _format_aligner_plan(record.traced_plan)
return body
def _format_tensor_comparison_rich_body(
record: TensorComparisonRecord,
) -> 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)
# ── NonTensorComparisonRecord ────────────────────────────────────────
def _format_non_tensor_body(record: NonTensorComparisonRecord) -> str:
suffix: str = record._format_location_suffix()
if record.values_equal:
return f"NonTensor: {record.name}{suffix} = {record.baseline_value} ({record.baseline_type}) [equal]"
return (
f"NonTensor: {record.name}{suffix}\n"
f" baseline = {record.baseline_value} ({record.baseline_type})\n"
f" target = {record.target_value} ({record.target_type})"
)
def _format_non_tensor_rich_body(
record: NonTensorComparisonRecord,
) -> RenderableType:
suffix: str = record._format_location_suffix()
name: str = escape(record.name)
baseline_val: str = escape(record.baseline_value)
target_val: str = escape(record.target_value)
if record.values_equal:
return (
f"═ {name}{suffix} = {baseline_val} "
f"({record.baseline_type}) [green]✓[/]"
)
return (
f"═ [bold red]{name}{suffix}[/]\n"
f" baseline = {baseline_val} ({record.baseline_type})\n"
f" target = {target_val} ({record.target_type})"
)
# ── SummaryRecord ────────────────────────────────────────────────────
def _format_summary_body(record: SummaryRecord) -> str:
return (
f"Summary: {record.passed} passed, {record.failed} failed, "
f"{record.skipped} skipped (total {record.total})"
)
def _format_summary_rich_body(record: SummaryRecord) -> RenderableType:
text: str = (
f"[bold green]{record.passed} passed[/] │ "
f"[bold red]{record.failed} failed[/] │ "
f"[yellow]{record.skipped} skipped[/] │ "
f"{record.total} total"
)
return Panel(text, title="SUMMARY", border_style="bold")
# ── LogRecord ────────────────────────────────────────────────────────
def _format_log_body(record: LogRecord) -> str:
return ""
# ── Standalone helpers ───────────────────────────────────────────────
def _format_aligner_plan(traced_plan: TracedAlignerPlan) -> str:
lines: list[str] = ["Aligner Plan:"]
for side_label, traced_side in [
("baseline", traced_plan.per_side.x),
("target", traced_plan.per_side.y),
]:
if not traced_side.step_plans:
lines.append(f" {side_label}: (no steps)")
continue
step_summaries: list[str] = []
for traced_step in traced_side.step_plans:
sub_strs: list[str] = [
_format_sub_plan_text(traced_sub)
for traced_sub in traced_step.sub_plans
]
summary: str = ", ".join(sub_strs) if sub_strs else "passthrough"
step_summaries.append(f"step={traced_step.step}: {summary}")
lines.append(f" {side_label}: [{'; '.join(step_summaries)}]")
lines.extend(_format_cross_side_plan_text(traced_plan.plan))
return "\n".join(lines)
def _format_sub_plan_text(traced_sub: TracedSubPlan) -> str:
sub_desc: str = f"{traced_sub.plan.type}"
if traced_sub.snapshot is not None:
snap = traced_sub.snapshot
in_count: int = len(snap.input_shapes)
out_count: int = len(snap.output_shapes)
in_shape: str = str(snap.input_shapes[0]) if snap.input_shapes else "?"
out_shape: str = str(snap.output_shapes[0]) if snap.output_shapes else "?"
sub_desc += f" {in_count}x{in_shape} -> {out_count}x{out_shape}"
return sub_desc
def _format_cross_side_plan_text(plan: AlignerPlan) -> list[str]:
lines: list[str] = []
if plan.token_aligner_plan is not None:
num_tokens: int = len(plan.token_aligner_plan.locators.x.steps)
lines.append(f" token_aligner: {num_tokens} tokens aligned")
if plan.axis_aligner_plan is not None:
parts: list[str] = []
if plan.axis_aligner_plan.pattern.x:
parts.append(f"x: {plan.axis_aligner_plan.pattern.x}")
if plan.axis_aligner_plan.pattern.y:
parts.append(f"y: {plan.axis_aligner_plan.pattern.y}")
lines.append(f" axis_aligner: {', '.join(parts)}")
return lines

205
python/sglang/srt/debug_utils/comparator/output_types.py
View File

@@ -3,15 +3,10 @@ from __future__ import annotations
from abc import abstractmethod
from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional, Union
import polars as pl
from pydantic import ConfigDict, Discriminator, Field, TypeAdapter, model_validator
from rich.console import RenderableType
from rich.markup import escape
from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import (
format_comparison,
format_replicated_checks,
)
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorComparisonInfo,
@@ -21,7 +16,6 @@ from sglang.srt.debug_utils.comparator.utils import Pair, _StrictBase
if TYPE_CHECKING:
from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
TracedAlignerPlan,
TracedSubPlan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan
from sglang.srt.debug_utils.comparator.report_sink import Verbosity
@@ -89,19 +83,45 @@ class _OutputRecord(_StrictBase):
@abstractmethod
def _format_body(self) -> str: ...
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
def _format_rich_body(self) -> RenderableType:
return self._format_body()
def to_rich(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_text(self) -> str:
body = self._format_body()
if self.errors:
body += "\n" + "\n".join(f" ✗ {e.to_text()}" for e in self.errors)
if self.infos:
body += "\n" + "\n".join(f" ℹ {i.to_text()}" for i in self.infos)
return body
from sglang.srt.debug_utils.comparator.output_formatter import (
_render_record_text,
)
return _render_record_text(self)
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 RecordLocation(_StrictBase):
@@ -132,7 +152,18 @@ class ConfigRecord(_OutputRecord):
config: dict[str, Any]
def _format_body(self) -> str:
return f"Config: {self.config}"
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)
class SkipComparisonRecord(_BaseComparisonRecord):
@@ -147,7 +178,18 @@ class SkipComparisonRecord(_BaseComparisonRecord):
return "skipped"
def _format_body(self) -> str:
return f"Skip: {self.name}{self._format_location_suffix()} ({self.reason})"
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)
class _TableRecord(_OutputRecord):
@@ -158,12 +200,19 @@ class _TableRecord(_OutputRecord):
def _table_title(self) -> str: ...
def _format_body(self) -> str:
from sglang.srt.debug_utils.comparator.display import _render_polars_as_text
return _render_polars_as_text(
pl.DataFrame(self.rows), title=self._table_title()
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_table_body,
)
return _format_table_body(self)
def _format_rich_body(self) -> RenderableType:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_table_rich_body,
)
return _format_table_rich_body(self)
class RankInfoRecord(_TableRecord):
type: Literal["rank_info"] = "rank_info"
@@ -196,12 +245,18 @@ class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord):
return "passed" if self.diff is not None and self.diff.passed else "failed"
def _format_body(self) -> str:
body: str = self._format_location_prefix() + format_comparison(self)
if self.replicated_checks:
body += "\n" + format_replicated_checks(self.replicated_checks)
if self.traced_plan is not None:
body += "\n" + _format_aligner_plan(self.traced_plan)
return body
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)
class NonTensorComparisonRecord(_BaseComparisonRecord):
@@ -220,15 +275,19 @@ class NonTensorComparisonRecord(_BaseComparisonRecord):
return "passed" if self.values_equal else "failed"
def _format_body(self) -> str:
suffix: str = self._format_location_suffix()
if self.values_equal:
return f"NonTensor: {self.name}{suffix} = {self.baseline_value} ({self.baseline_type}) [equal]"
return (
f"NonTensor: {self.name}{suffix}\n"
f" baseline = {self.baseline_value} ({self.baseline_type})\n"
f" target = {self.target_value} ({self.target_type})"
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_non_tensor_body,
)
return _format_non_tensor_body(self)
def _format_rich_body(self) -> RenderableType:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_non_tensor_rich_body,
)
return _format_non_tensor_rich_body(self)
class SummaryRecord(_OutputRecord):
type: Literal["summary"] = "summary"
@@ -247,74 +306,38 @@ class SummaryRecord(_OutputRecord):
return self
def _format_body(self) -> str:
return (
f"Summary: {self.passed} passed, {self.failed} failed, "
f"{self.skipped} skipped (total {self.total})"
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)
class LogRecord(_OutputRecord):
type: Literal["log"] = "log"
def _format_body(self) -> str:
return ""
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:
lines: list[str] = ["Aligner Plan:"]
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_aligner_plan as _impl,
)
for side_label, traced_side in [
("baseline", traced_plan.per_side.x),
("target", traced_plan.per_side.y),
]:
if not traced_side.step_plans:
lines.append(f" {side_label}: (no steps)")
continue
step_summaries: list[str] = []
for traced_step in traced_side.step_plans:
sub_strs: list[str] = [
_format_sub_plan_text(traced_sub)
for traced_sub in traced_step.sub_plans
]
summary: str = ", ".join(sub_strs) if sub_strs else "passthrough"
step_summaries.append(f"step={traced_step.step}: {summary}")
lines.append(f" {side_label}: [{'; '.join(step_summaries)}]")
lines.extend(_format_cross_side_plan_text(traced_plan.plan))
return "\n".join(lines)
def _format_sub_plan_text(traced_sub: TracedSubPlan) -> str:
sub_desc: str = f"{traced_sub.plan.type}"
if traced_sub.snapshot is not None:
snap = traced_sub.snapshot
in_count: int = len(snap.input_shapes)
out_count: int = len(snap.output_shapes)
in_shape: str = str(snap.input_shapes[0]) if snap.input_shapes else "?"
out_shape: str = str(snap.output_shapes[0]) if snap.output_shapes else "?"
sub_desc += f" {in_count}x{in_shape} -> {out_count}x{out_shape}"
return sub_desc
def _format_cross_side_plan_text(plan: AlignerPlan) -> list[str]:
lines: list[str] = []
if plan.token_aligner_plan is not None:
num_tokens: int = len(plan.token_aligner_plan.locators.x.steps)
lines.append(f" token_aligner: {num_tokens} tokens aligned")
if plan.axis_aligner_plan is not None:
parts: list[str] = []
if plan.axis_aligner_plan.pattern.x:
parts.append(f"x: {plan.axis_aligner_plan.pattern.x}")
if plan.axis_aligner_plan.pattern.y:
parts.append(f"y: {plan.axis_aligner_plan.pattern.y}")
lines.append(f" axis_aligner: {', '.join(parts)}")
return lines
return _impl(traced_plan)
AnyRecord = Annotated[

2
python/sglang/srt/debug_utils/comparator/report_sink.py
View File

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

306
python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py
View File

@@ -1,15 +1,83 @@
from __future__ import annotations
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Optional
from rich.markup import escape
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorComparisonInfo,
TensorInfo,
TensorStats,
)
if TYPE_CHECKING:
from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
TracedAlignerPlan,
TracedSubPlan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan
from sglang.srt.debug_utils.comparator.output_types import (
BundleSideInfo,
ReplicatedCheckResult,
ShapeSnapshot,
TensorComparisonRecord,
)
from sglang.srt.debug_utils.comparator.utils import Pair
def _esc_shape(shape: Optional[list[int]]) -> str:
return escape(str(shape))
def _strip_torch_prefix(dtype: str) -> str:
return dtype.replace("torch.", "")
# ---------------------------------------------------------------------------
# Number formatting
# ---------------------------------------------------------------------------
def _fmt_val(value: float) -> str:
return f"{value:.2e}"
def _fmt_diff_colored(diff: float, *, threshold: float = 1e-2) -> str:
formatted: str = f"{diff:+.2e}"
if abs(diff) >= threshold:
return f"[yellow]{formatted}[/]"
return f"[dim]{formatted}[/]"
# ---------------------------------------------------------------------------
# Passed / color / marker helper
# ---------------------------------------------------------------------------
def _category_marker(category: str) -> tuple[bool, str, str]:
passed: bool = category == "passed"
color: str = "green" if passed else "red"
marker: str = f"[{color}]✅[/]" if passed else f"[{color}]❌[/]"
return passed, color, marker
# ---------------------------------------------------------------------------
# Stats formatting helpers
# ---------------------------------------------------------------------------
def _format_stat_line(stat_name: str, val_b: float, val_t: float, diff: float) -> str:
return (
f" [blue]{stat_name:10s}[/] {val_b:>10.4f} vs {val_t:>10.4f}"
f" Δ {_fmt_diff_colored(diff)}"
)
# ---------------------------------------------------------------------------
# Old text-only formatters (kept for to_text() backward compatibility)
# ---------------------------------------------------------------------------
def format_comparison(info: TensorComparisonInfo) -> str:
@@ -129,3 +197,237 @@ def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]:
lines.append("[abs_diff] " + " ".join(quantile_parts))
return lines
# ---------------------------------------------------------------------------
# New Rich markup formatters
# ---------------------------------------------------------------------------
def format_comparison_rich(record: TensorComparisonRecord) -> str:
return _format_comparison_normal(record=record)
def _format_comparison_normal(*, record: TensorComparisonRecord) -> str:
passed, color, marker = _category_marker(record.category)
baseline: TensorInfo = record.baseline
target: TensorInfo = record.target
aligned_shape: str = _esc_shape(record.unified_shape)
dtype_str: str = _strip_torch_prefix(baseline.dtype)
lines: list[str] = []
# L0: Header
lines.append(
f"{marker} [bold {color}]{escape(record.name)}[/] "
f"[dim cyan]── {dtype_str} {aligned_shape}[/]"
)
# L1: Key metrics
if record.diff is not None:
diff: DiffInfo = record.diff
rel_style: str = f"bold {color}" if not passed else color
lines.append(
f" [{rel_style}]rel_diff={_fmt_val(diff.rel_diff)}[/]"
f" max_abs={_fmt_val(diff.max_abs_diff)}"
f" mean_abs={_fmt_val(diff.mean_abs_diff)}"
)
if not passed:
lines.append(
f" max_abs @ {_esc_shape(diff.max_diff_coord)}: "
f"baseline={diff.baseline_at_max} target={diff.target_at_max}"
)
elif record.shape_mismatch:
lines.append(" [yellow]⚠ Shape mismatch[/]")
# Downcast info
if record.diff_downcast is not None and record.downcast_dtype is not None:
dc: DiffInfo = record.diff_downcast
dc_marker: str = "[green]✅[/]" if dc.passed else "[red]❌[/]"
lines.append(
f" {dc_marker} downcast to {record.downcast_dtype}: "
f"rel_diff={_fmt_val(dc.rel_diff)}"
)
# 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))
# Plan section
if record.traced_plan is not None:
lines.append(" [dim]Plan[/]")
lines.extend(_format_plan_section_rich(traced_plan=record.traced_plan))
# Aligned section
lines.append(" [dim]Aligned[/]")
lines.append(
f" {_esc_shape(record.unified_shape)} vs {_esc_shape(target.shape)}"
f" {baseline.dtype} vs {target.dtype}"
)
# Stats section
lines.append(" [dim]Stats[/]")
lines.extend(_format_stats_rich(baseline=baseline.stats, target=target.stats))
# Abs diff percentiles (show when failed)
if not passed 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:
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:
lines.append(" [dim]Replicated Checks[/]")
for check in record.replicated_checks:
chk_marker: str = "[green]✅[/]" if check.passed else "[red]❌[/]"
if check.diff is not None:
lines.append(
f" {chk_marker} axis={check.axis} group={check.group_index}"
f" idx={check.compared_index} vs {check.baseline_index}"
f" rel_diff={_fmt_val(check.diff.rel_diff)}"
f" max_abs={_fmt_val(check.diff.max_abs_diff)}"
)
else:
lines.append(
f" {chk_marker} axis={check.axis} group={check.group_index}"
f" idx={check.compared_index} vs {check.baseline_index}: n/a"
)
return "\n".join(lines)
def _format_bundle_section(bundle_info: Pair[BundleSideInfo]) -> list[str]:
lines: list[str] = []
for label, side in [("baseline", bundle_info.x), ("target", bundle_info.y)]:
if not side.files:
lines.append(f" {label} [dim](no files)[/]")
continue
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"
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}"
)
return lines
def _format_plan_section_rich(
*,
traced_plan: TracedAlignerPlan,
) -> list[str]:
lines: list[str] = []
for side_label, traced_side in [
("baseline", traced_plan.per_side.x),
("target", traced_plan.per_side.y),
]:
if not traced_side.step_plans:
lines.append(f" {side_label} [dim](passthrough)[/]")
continue
parts: list[str] = [
_format_sub_plan_rich(traced_sub)
for traced_step in traced_side.step_plans
for traced_sub in traced_step.sub_plans
]
lines.append(f" {side_label} " + " → ".join(parts))
lines.extend(_format_cross_side_plan_rich(traced_plan.plan))
return lines
def _format_sub_plan_rich(traced_sub: TracedSubPlan) -> str:
sub = traced_sub.plan
snapshot: Optional[ShapeSnapshot] = traced_sub.snapshot
op_name: str = sub.type
axis_str: str = ""
if isinstance(sub, UnsharderPlan):
axis_str = f"({sub.axis})"
shape_change: str = ""
if snapshot:
in_count: int = len(snapshot.input_shapes)
out_count: int = len(snapshot.output_shapes)
in_shape: str = (
_esc_shape(snapshot.input_shapes[0]) if snapshot.input_shapes else "?"
)
out_shape: str = (
_esc_shape(snapshot.output_shapes[0]) if snapshot.output_shapes else "?"
)
shape_change = f" {in_count}×{in_shape} → {out_count}×{out_shape}"
return f"[magenta]{op_name}{axis_str}[/]{shape_change}"
def _format_cross_side_plan_rich(plan: AlignerPlan) -> list[str]:
lines: list[str] = []
if plan.token_aligner_plan is not None:
num_tokens: int = len(plan.token_aligner_plan.locators.x.steps)
lines.append(f" token_aligner [dim]{num_tokens} tokens[/]")
if plan.axis_aligner_plan is not None:
parts: list[str] = []
if plan.axis_aligner_plan.pattern.x:
parts.append(f"x={plan.axis_aligner_plan.pattern.x}")
if plan.axis_aligner_plan.pattern.y:
parts.append(f"y={plan.axis_aligner_plan.pattern.y}")
if parts:
lines.append(f" axis_aligner [dim]{', '.join(parts)}[/]")
else:
lines.append(" axis_aligner [dim](no-op)[/]")
return lines
def _format_stats_rich(
*,
baseline: TensorStats,
target: TensorStats,
) -> 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))
# 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
def _format_abs_diff_percentiles_rich(diff: DiffInfo) -> str:
parts: list[str] = []
for p, value in sorted(diff.abs_diff_percentiles.items()):
formatted: str = f"p{p}={_fmt_val(value)}"
if p >= 99 and value > 0.1:
formatted = f"[yellow]{formatted}[/]"
parts.append(formatted)
return " ".join(parts)