diff --git a/python/sglang/srt/debug_utils/comparator/display.py b/python/sglang/srt/debug_utils/comparator/display.py index a20141436..ba4179c0b 100644 --- a/python/sglang/srt/debug_utils/comparator/display.py +++ b/python/sglang/srt/debug_utils/comparator/display.py @@ -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( diff --git a/python/sglang/srt/debug_utils/comparator/output_formatter.py b/python/sglang/srt/debug_utils/comparator/output_formatter.py new file mode 100644 index 000000000..dfa25368f --- /dev/null +++ b/python/sglang/srt/debug_utils/comparator/output_formatter.py @@ -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 diff --git a/python/sglang/srt/debug_utils/comparator/output_types.py b/python/sglang/srt/debug_utils/comparator/output_types.py index dea5a9acc..db3722986 100644 --- a/python/sglang/srt/debug_utils/comparator/output_types.py +++ b/python/sglang/srt/debug_utils/comparator/output_types.py @@ -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[ diff --git a/python/sglang/srt/debug_utils/comparator/report_sink.py b/python/sglang/srt/debug_utils/comparator/report_sink.py index 61f9e9ac5..cab8f5cd8 100644 --- a/python/sglang/srt/debug_utils/comparator/report_sink.py +++ b/python/sglang/srt/debug_utils/comparator/report_sink.py @@ -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 diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py b/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py index 1946aa56a..ece8fd214 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py @@ -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) diff --git a/test/registered/debug_utils/comparator/conftest.py b/test/registered/debug_utils/comparator/conftest.py index 3b26f0af8..b2a2a67d3 100644 --- a/test/registered/debug_utils/comparator/conftest.py +++ b/test/registered/debug_utils/comparator/conftest.py @@ -19,23 +19,6 @@ from sglang.srt.debug_utils.comparator.report_sink import report_sink collect_ignore_glob: list[str] = [] -def pytest_configure(config: pytest.Config) -> None: - config.addinivalue_line( - "filterwarnings", - "ignore:Unknown config option. asyncio_mode:pytest.PytestConfigWarning", - ) - config.addinivalue_line( - "filterwarnings", - "ignore:builtin type Swig.*:DeprecationWarning", - ) - config.addinivalue_line( - "filterwarnings", - "ignore:Named tensors and all their associated APIs:UserWarning", - ) - -collect_ignore_glob: list[str] = [] - - def pytest_configure(config: pytest.Config) -> None: config.addinivalue_line( "filterwarnings", diff --git a/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py b/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py index b216fba33..a9f8a87af 100644 --- a/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py +++ b/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py @@ -1,96 +1,63 @@ import sys import pytest +from registered.debug_utils.comparator.testing_helpers import make_diff as _make_diff +from registered.debug_utils.comparator.testing_helpers import make_stats as _make_stats +from registered.debug_utils.comparator.testing_helpers import ( + make_tensor_info as _make_tensor_info, +) +from sglang.srt.debug_utils.comparator.aligner.axis_aligner import AxisAlignerPlan +from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import ( + TracedAlignerPlan, + TracedSidePlan, + TracedStepPlan, + TracedSubPlan, +) +from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( + AlignerPerStepPlan, + AlignerPlan, +) +from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ( + ReordererPlan, + ZigzagToNaturalParams, +) +from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import ( + TokenAlignerPlan, + TokenLocator, +) +from sglang.srt.debug_utils.comparator.aligner.unsharder.types import ( + ConcatParams, + UnsharderPlan, +) +from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, TokenLayout +from sglang.srt.debug_utils.comparator.output_types import ( + BundleFileInfo, + BundleSideInfo, + ReplicatedCheckResult, + ShapeSnapshot, + TensorComparisonRecord, +) from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import ( + _format_abs_diff_percentiles_rich, + _format_bundle_section, + _format_plan_section_rich, + _format_stats_rich, format_comparison, + format_comparison_rich, + format_replicated_checks, ) from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( DiffInfo, TensorComparisonInfo, - TensorInfo, TensorStats, ) +from sglang.srt.debug_utils.comparator.utils import Pair from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="default", nightly=True) -_DEFAULT_PERCENTILES: dict[int, float] = { - 1: -1.8, - 5: -1.5, - 50: 0.0, - 95: 1.5, - 99: 1.8, -} - - -def _make_stats( - mean: float = 0.0, - abs_mean: float = 0.8, - std: float = 1.0, - min: float = -2.0, - max: float = 2.0, - percentiles: dict[int, float] | None = None, -) -> TensorStats: - return TensorStats( - mean=mean, - abs_mean=abs_mean, - std=std, - min=min, - max=max, - percentiles=percentiles if percentiles is not None else _DEFAULT_PERCENTILES, - ) - - -_DEFAULT_ABS_DIFF_PERCENTILES: dict[int, float] = { - 1: 0.0001, - 5: 0.0001, - 50: 0.0002, - 95: 0.0004, - 99: 0.0005, -} - - -def _make_diff( - rel_diff: float = 0.0001, - max_abs_diff: float = 0.0005, - mean_abs_diff: float = 0.0002, - abs_diff_percentiles: dict[int, float] | None = None, - diff_threshold: float = 1e-3, - passed: bool = True, -) -> DiffInfo: - return DiffInfo( - rel_diff=rel_diff, - max_abs_diff=max_abs_diff, - mean_abs_diff=mean_abs_diff, - abs_diff_percentiles=( - abs_diff_percentiles - if abs_diff_percentiles is not None - else _DEFAULT_ABS_DIFF_PERCENTILES - ), - max_diff_coord=[2, 3], - baseline_at_max=1.0, - target_at_max=1.0005, - diff_threshold=diff_threshold, - passed=passed, - ) - - -def _make_tensor_info( - shape: list[int] | None = None, - dtype: str = "torch.float32", - stats: TensorStats | None = None, - sample: str | None = None, -) -> TensorInfo: - return TensorInfo( - shape=shape if shape is not None else [4, 8], - dtype=dtype, - stats=stats if stats is not None else _make_stats(), - sample=sample, - ) - - # Snapshot strings below are intentionally spelled out in full per test. # The shared skeleton (stats block, diff block) looks duplicated, but keeping # each test self-contained makes failures immediately readable without chasing @@ -294,5 +261,552 @@ class TestFormatComparison: ) +def _make_comparison_record( + name: str = "hidden_states", + shape: list[int] | None = None, + dtype: str = "torch.float32", + diff: DiffInfo | None = None, + shape_mismatch: bool = False, + sample: str | None = None, + diff_downcast: DiffInfo | None = None, + downcast_dtype: str | None = None, + replicated_checks: list[ReplicatedCheckResult] | None = None, + raw_bundle_info: Pair[BundleSideInfo] | None = None, + traced_plan: TracedAlignerPlan | None = None, +) -> TensorComparisonRecord: + s: list[int] = shape if shape is not None else [4, 8] + return TensorComparisonRecord( + name=name, + baseline=_make_tensor_info(shape=s, dtype=dtype, sample=sample), + target=_make_tensor_info(shape=s, dtype=dtype, sample=sample), + unified_shape=s, + shape_mismatch=shape_mismatch, + diff=diff, + diff_downcast=diff_downcast, + downcast_dtype=downcast_dtype, + replicated_checks=replicated_checks or [], + raw_bundle_info=raw_bundle_info, + traced_plan=traced_plan, + ) + + +def _make_bundle_side_info( + num_files: int = 2, + shape: list[int] | None = None, + dtype: str = "torch.float32", + dims: str | None = None, + with_parallel_info: bool = False, +) -> BundleSideInfo: + s: list[int] = shape if shape is not None else [2, 4096] + files: list[BundleFileInfo] = [] + for i in range(num_files): + par: dict[str, str] | None = ( + {"tp": f"{i}/{num_files}"} if with_parallel_info else None + ) + files.append(BundleFileInfo(shape=s, dtype=dtype, rank=i, parallel_info=par)) + return BundleSideInfo(num_files=num_files, files=files, dims=dims) + + +def _make_simple_aligner_plan( + *, + with_unsharder: bool = False, + with_reorderer: bool = False, + with_token_aligner: bool = False, + with_axis_aligner: bool = False, + axis_aligner_noop: bool = False, +) -> AlignerPlan: + baseline_plans: list[AlignerPerStepPlan] = [] + target_plans: list[AlignerPerStepPlan] = [] + + if with_unsharder: + unsharder: UnsharderPlan = UnsharderPlan( + axis=ParallelAxis.TP, + params=ConcatParams(dim_name="h"), + groups=[[0, 1]], + ) + target_plans.append( + AlignerPerStepPlan( + step=0, input_object_indices=[0, 1], sub_plans=[unsharder] + ) + ) + + if with_reorderer: + reorderer: ReordererPlan = ReordererPlan( + params=ZigzagToNaturalParams(dim_name="s", cp_size=2), + ) + target_plans.append( + AlignerPerStepPlan(step=0, input_object_indices=[0], sub_plans=[reorderer]) + ) + + token_aligner_plan: TokenAlignerPlan | None = None + if with_token_aligner: + token_aligner_plan = TokenAlignerPlan( + locators=Pair( + x=TokenLocator(steps=[0, 0, 0], token_index_in_step=[0, 1, 2]), + y=TokenLocator(steps=[0, 0, 0], token_index_in_step=[0, 1, 2]), + ), + layouts=Pair(x=TokenLayout.T, y=TokenLayout.T), + ) + + axis_aligner_plan: AxisAlignerPlan | None = None + if with_axis_aligner: + if axis_aligner_noop: + axis_aligner_plan = AxisAlignerPlan(pattern=Pair(x=None, y=None)) + else: + axis_aligner_plan = AxisAlignerPlan( + pattern=Pair(x="b s d -> s b d", y=None) + ) + + return AlignerPlan( + per_step_plans=Pair(x=baseline_plans, y=target_plans), + token_aligner_plan=token_aligner_plan, + axis_aligner_plan=axis_aligner_plan, + ) + + +def _make_traced_plan( + plan: AlignerPlan, + *, + target_input_shapes: list[list[int]] | None = None, + target_output_shapes: list[list[int]] | None = None, +) -> TracedAlignerPlan: + """Build a TracedAlignerPlan by attaching snapshots to target sub_plans.""" + baseline_traced_steps: list[TracedStepPlan] = [ + TracedStepPlan( + step=sp.step, + input_object_indices=sp.input_object_indices, + sub_plans=[TracedSubPlan(plan=sub) for sub in sp.sub_plans], + ) + for sp in plan.per_step_plans.x + ] + + target_traced_steps: list[TracedStepPlan] = [] + for sp in plan.per_step_plans.y: + traced_subs: list[TracedSubPlan] = [] + for sub in sp.sub_plans: + snapshot: ShapeSnapshot | None = None + if target_input_shapes is not None or target_output_shapes is not None: + snapshot = ShapeSnapshot( + input_shapes=target_input_shapes or [[2, 4096], [2, 4096]], + output_shapes=target_output_shapes or [[4, 4096]], + ) + traced_subs.append(TracedSubPlan(plan=sub, snapshot=snapshot)) + target_traced_steps.append( + TracedStepPlan( + step=sp.step, + input_object_indices=sp.input_object_indices, + sub_plans=traced_subs, + ) + ) + + return TracedAlignerPlan( + plan=plan, + per_side=Pair( + x=TracedSidePlan(step_plans=baseline_traced_steps), + y=TracedSidePlan(step_plans=target_traced_steps), + ), + ) + + +# --------------------------------------------------------------------------- +# Rich format snapshot tests (normal mode only) +# --------------------------------------------------------------------------- + + +class TestFormatComparisonRichNormal: + """format_comparison_rich() snapshot tests.""" + + def test_passed(self) -> None: + record: TensorComparisonRecord = _make_comparison_record( + diff=_make_diff(rel_diff=1e-4, passed=True), + ) + result: str = format_comparison_rich(record) + + assert result == ( + "[green]✅[/] [bold green]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n" + " [green]rel_diff=1.00e-04[/] max_abs=5.00e-04 mean_abs=2.00e-04\n" + " [dim]Aligned[/]\n" + " [4, 8] vs [4, 8] torch.float32 vs torch.float32\n" + " [dim]Stats[/]\n" + " [blue]mean [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n" + " [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]\n" + " [blue]range [/] [-2.0000, 2.0000] vs [-2.0000, 2.0000]" + ) + + def test_failed(self) -> None: + record: TensorComparisonRecord = _make_comparison_record( + diff=_make_diff( + rel_diff=0.5, max_abs_diff=1.0, mean_abs_diff=0.3, passed=False + ), + ) + result: str = format_comparison_rich(record) + + assert result == ( + "[red]❌[/] [bold red]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n" + " [bold red]rel_diff=5.00e-01[/] max_abs=1.00e+00 mean_abs=3.00e-01\n" + " max_abs @ [2, 3]: baseline=1.0 target=1.0005\n" + " [dim]Aligned[/]\n" + " [4, 8] vs [4, 8] torch.float32 vs torch.float32\n" + " [dim]Stats[/]\n" + " [blue]mean [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n" + " [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]\n" + " [blue]range [/] [-2.0000, 2.0000] vs [-2.0000, 2.0000]\n" + " [dim]Abs Diff Percentiles[/]\n" + " p1=1.00e-04 p5=1.00e-04 p50=2.00e-04 p95=4.00e-04 p99=5.00e-04" + ) + + def test_shape_mismatch(self) -> None: + record: TensorComparisonRecord = _make_comparison_record( + shape_mismatch=True, + ) + result: str = format_comparison_rich(record) + + assert result == ( + "[red]❌[/] [bold red]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n" + " [yellow]⚠ Shape mismatch[/]\n" + " [dim]Aligned[/]\n" + " [4, 8] vs [4, 8] torch.float32 vs torch.float32\n" + " [dim]Stats[/]\n" + " [blue]mean [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n" + " [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]\n" + " [blue]range [/] [-2.0000, 2.0000] vs [-2.0000, 2.0000]" + ) + + def test_with_downcast(self) -> None: + record: TensorComparisonRecord = _make_comparison_record( + diff=_make_diff(rel_diff=0.01, passed=False), + diff_downcast=_make_diff(rel_diff=1e-5, passed=True), + downcast_dtype="torch.bfloat16", + ) + result: str = format_comparison_rich(record) + + assert result == ( + "[red]❌[/] [bold red]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n" + " [bold red]rel_diff=1.00e-02[/] max_abs=5.00e-04 mean_abs=2.00e-04\n" + " max_abs @ [2, 3]: baseline=1.0 target=1.0005\n" + " [green]✅[/] downcast to torch.bfloat16: rel_diff=1.00e-05\n" + " [dim]Aligned[/]\n" + " [4, 8] vs [4, 8] torch.float32 vs torch.float32\n" + " [dim]Stats[/]\n" + " [blue]mean [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n" + " [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]\n" + " [blue]range [/] [-2.0000, 2.0000] vs [-2.0000, 2.0000]\n" + " [dim]Abs Diff Percentiles[/]\n" + " p1=1.00e-04 p5=1.00e-04 p50=2.00e-04 p95=4.00e-04 p99=5.00e-04" + ) + + def test_with_bundle_info(self) -> None: + bundle_info: Pair[BundleSideInfo] = Pair( + x=_make_bundle_side_info(num_files=2, dims="b s h(tp) d"), + y=_make_bundle_side_info(num_files=2, dims="b s h(tp) d"), + ) + record: TensorComparisonRecord = _make_comparison_record( + diff=_make_diff(passed=True), + raw_bundle_info=bundle_info, + ) + result: str = format_comparison_rich(record) + + assert result == ( + "[green]✅[/] [bold green]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n" + " [green]rel_diff=1.00e-04[/] max_abs=5.00e-04 mean_abs=2.00e-04\n" + " [dim]Bundle[/]\n" + " baseline [cyan]2 files[/] × [2, 4096] float32 [dim]dims: b s h(tp) d[/]\n" + " target [cyan]2 files[/] × [2, 4096] float32 [dim]dims: b s h(tp) d[/]\n" + " [dim]Aligned[/]\n" + " [4, 8] vs [4, 8] torch.float32 vs torch.float32\n" + " [dim]Stats[/]\n" + " [blue]mean [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n" + " [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]\n" + " [blue]range [/] [-2.0000, 2.0000] vs [-2.0000, 2.0000]" + ) + + def test_with_plan(self) -> None: + plan: AlignerPlan = _make_simple_aligner_plan(with_unsharder=True) + record: TensorComparisonRecord = _make_comparison_record( + diff=_make_diff(passed=True), + traced_plan=_make_traced_plan(plan), + ) + result: str = format_comparison_rich(record) + + assert result == ( + "[green]✅[/] [bold green]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n" + " [green]rel_diff=1.00e-04[/] max_abs=5.00e-04 mean_abs=2.00e-04\n" + " [dim]Plan[/]\n" + " baseline [dim](passthrough)[/]\n" + " target [magenta]unsharder(ParallelAxis.TP)[/]\n" + " [dim]Aligned[/]\n" + " [4, 8] vs [4, 8] torch.float32 vs torch.float32\n" + " [dim]Stats[/]\n" + " [blue]mean [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n" + " [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]\n" + " [blue]range [/] [-2.0000, 2.0000] vs [-2.0000, 2.0000]" + ) + + +class TestFormatBundleSection: + """_format_bundle_section() snapshot tests.""" + + def test_single_shape(self) -> None: + bundle: Pair[BundleSideInfo] = Pair( + x=_make_bundle_side_info(num_files=2, shape=[2, 4096]), + y=_make_bundle_side_info(num_files=2, shape=[2, 4096]), + ) + lines: list[str] = _format_bundle_section(bundle) + + assert lines == [ + " baseline [cyan]2 files[/] × [2, 4096] float32", + " target [cyan]2 files[/] × [2, 4096] float32", + ] + + def test_mixed_shapes(self) -> None: + side: BundleSideInfo = BundleSideInfo( + num_files=2, + files=[ + BundleFileInfo(shape=[2, 4096], dtype="torch.float32", rank=0), + BundleFileInfo(shape=[3, 4096], dtype="torch.float32", rank=1), + ], + ) + bundle: Pair[BundleSideInfo] = Pair(x=side, y=side) + lines: list[str] = _format_bundle_section(bundle) + + assert lines == [ + " baseline [cyan]2 files[/] × mixed shapes float32", + " target [cyan]2 files[/] × mixed shapes float32", + ] + + def test_no_files(self) -> None: + empty: BundleSideInfo = BundleSideInfo(num_files=0, files=[]) + bundle: Pair[BundleSideInfo] = Pair(x=empty, y=empty) + lines: list[str] = _format_bundle_section(bundle) + + assert lines == [ + " baseline [dim](no files)[/]", + " target [dim](no files)[/]", + ] + + def test_with_dims(self) -> None: + bundle: Pair[BundleSideInfo] = Pair( + x=_make_bundle_side_info(num_files=1, dims="b s h(tp) d"), + y=_make_bundle_side_info(num_files=1, dims="b s h(tp) d"), + ) + lines: list[str] = _format_bundle_section(bundle) + + assert lines == [ + " baseline [cyan]1 files[/] × [2, 4096] float32 [dim]dims: b s h(tp) d[/]", + " target [cyan]1 files[/] × [2, 4096] float32 [dim]dims: b s h(tp) d[/]", + ] + + +class TestFormatPlanSectionRich: + """_format_plan_section_rich() snapshot tests.""" + + def test_passthrough(self) -> None: + plan: AlignerPlan = _make_simple_aligner_plan() + traced: TracedAlignerPlan = _make_traced_plan(plan) + lines: list[str] = _format_plan_section_rich(traced_plan=traced) + + assert lines == [ + " baseline [dim](passthrough)[/]", + " target [dim](passthrough)[/]", + ] + + def test_unsharder_op(self) -> None: + plan: AlignerPlan = _make_simple_aligner_plan(with_unsharder=True) + traced: TracedAlignerPlan = _make_traced_plan(plan) + lines: list[str] = _format_plan_section_rich(traced_plan=traced) + + assert lines == [ + " baseline [dim](passthrough)[/]", + " target [magenta]unsharder(ParallelAxis.TP)[/]", + ] + + def test_reorderer_op(self) -> None: + plan: AlignerPlan = _make_simple_aligner_plan(with_reorderer=True) + traced: TracedAlignerPlan = _make_traced_plan(plan) + lines: list[str] = _format_plan_section_rich(traced_plan=traced) + + assert lines == [ + " baseline [dim](passthrough)[/]", + " target [magenta]reorderer[/]", + ] + + def test_with_shape_traces(self) -> None: + plan: AlignerPlan = _make_simple_aligner_plan(with_unsharder=True) + traced: TracedAlignerPlan = _make_traced_plan( + plan, + target_input_shapes=[[2, 4096], [2, 4096]], + target_output_shapes=[[4, 4096]], + ) + lines: list[str] = _format_plan_section_rich(traced_plan=traced) + + assert lines == [ + " baseline [dim](passthrough)[/]", + " target [magenta]unsharder(ParallelAxis.TP)[/] 2×[2, 4096] → 1×[4, 4096]", + ] + + def test_with_token_aligner(self) -> None: + plan: AlignerPlan = _make_simple_aligner_plan(with_token_aligner=True) + traced: TracedAlignerPlan = _make_traced_plan(plan) + lines: list[str] = _format_plan_section_rich(traced_plan=traced) + + assert lines == [ + " baseline [dim](passthrough)[/]", + " target [dim](passthrough)[/]", + " token_aligner [dim]3 tokens[/]", + ] + + def test_with_axis_aligner(self) -> None: + plan: AlignerPlan = _make_simple_aligner_plan(with_axis_aligner=True) + traced: TracedAlignerPlan = _make_traced_plan(plan) + lines: list[str] = _format_plan_section_rich(traced_plan=traced) + + assert lines == [ + " baseline [dim](passthrough)[/]", + " target [dim](passthrough)[/]", + " axis_aligner [dim]x=b s d -> s b d[/]", + ] + + def test_axis_aligner_noop(self) -> None: + plan: AlignerPlan = _make_simple_aligner_plan( + with_axis_aligner=True, axis_aligner_noop=True + ) + traced: TracedAlignerPlan = _make_traced_plan(plan) + lines: list[str] = _format_plan_section_rich(traced_plan=traced) + + assert lines == [ + " baseline [dim](passthrough)[/]", + " target [dim](passthrough)[/]", + " axis_aligner [dim](no-op)[/]", + ] + + +class TestFormatStatsRich: + """_format_stats_rich() snapshot tests.""" + + def test_basic(self) -> None: + baseline: TensorStats = _make_stats(mean=0.0, std=1.0, min=-2.0, max=2.0) + target: TensorStats = _make_stats( + mean=0.0001, std=1.0001, min=-2.0001, max=2.0001 + ) + lines: list[str] = _format_stats_rich(baseline=baseline, target=target) + + assert lines == [ + " [blue]mean [/] 0.0000 vs 0.0001 Δ [dim]+1.00e-04[/]", + " [blue]std [/] 1.0000 vs 1.0001 Δ [dim]+1.00e-04[/]", + " [blue]range [/] [-2.0000, 2.0000] vs [-2.0001, 2.0001]", + ] + + def test_large_delta(self) -> None: + baseline: TensorStats = _make_stats(mean=0.0) + target: TensorStats = _make_stats(mean=1.0) + lines: list[str] = _format_stats_rich(baseline=baseline, target=target) + + assert lines == [ + " [blue]mean [/] 0.0000 vs 1.0000 Δ [yellow]+1.00e+00[/]", + " [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]", + " [blue]range [/] [-2.0000, 2.0000] vs [-2.0000, 2.0000]", + ] + + def test_small_delta(self) -> None: + baseline: TensorStats = _make_stats(mean=0.0) + target: TensorStats = _make_stats(mean=0.001) + lines: list[str] = _format_stats_rich(baseline=baseline, target=target) + + assert lines == [ + " [blue]mean [/] 0.0000 vs 0.0010 Δ [dim]+1.00e-03[/]", + " [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]", + " [blue]range [/] [-2.0000, 2.0000] vs [-2.0000, 2.0000]", + ] + + +class TestFormatAbsDiffPercentilesRich: + """_format_abs_diff_percentiles_rich() snapshot tests.""" + + def test_normal_values(self) -> None: + diff: DiffInfo = _make_diff() + result: str = _format_abs_diff_percentiles_rich(diff) + + assert result == ( + "p1=1.00e-04 p5=1.00e-04 p50=2.00e-04 " "p95=4.00e-04 p99=5.00e-04" + ) + + def test_high_p99_coloring(self) -> None: + diff: DiffInfo = _make_diff( + abs_diff_percentiles={99: 0.5}, + ) + result: str = _format_abs_diff_percentiles_rich(diff) + + assert result == "[yellow]p99=5.00e-01[/]" + + def test_low_p99_no_coloring(self) -> None: + diff: DiffInfo = _make_diff( + abs_diff_percentiles={99: 0.01}, + ) + result: str = _format_abs_diff_percentiles_rich(diff) + + assert result == "p99=1.00e-02" + + +class TestFormatReplicatedChecks: + """format_replicated_checks() snapshot tests.""" + + def test_all_passed(self) -> None: + checks: list[ReplicatedCheckResult] = [ + ReplicatedCheckResult( + axis="tp", + group_index=0, + compared_index=1, + baseline_index=0, + passed=True, + atol=1e-3, + diff=_make_diff(rel_diff=1e-6, max_abs_diff=1e-5, mean_abs_diff=1e-6), + ), + ] + result: str = format_replicated_checks(checks) + + assert result == ( + "Replicated checks:\n" + " ✅ axis=tp group=0 idx=1 vs 0: " + "rel_diff=1.000000e-06 max_abs_diff=1.000000e-05 mean_abs_diff=1.000000e-06" + ) + + def test_one_failed(self) -> None: + checks: list[ReplicatedCheckResult] = [ + ReplicatedCheckResult( + axis="tp", + group_index=0, + compared_index=1, + baseline_index=0, + passed=False, + atol=1e-3, + diff=_make_diff(rel_diff=0.5, max_abs_diff=1.0, mean_abs_diff=0.3), + ), + ] + result: str = format_replicated_checks(checks) + + assert result == ( + "Replicated checks:\n" + " ❌ axis=tp group=0 idx=1 vs 0: " + "rel_diff=5.000000e-01 max_abs_diff=1.000000e+00 mean_abs_diff=3.000000e-01" + ) + + def test_no_diff(self) -> None: + checks: list[ReplicatedCheckResult] = [ + ReplicatedCheckResult( + axis="tp", + group_index=0, + compared_index=1, + baseline_index=0, + passed=True, + atol=1e-3, + ), + ] + result: str = format_replicated_checks(checks) + + assert result == ( + "Replicated checks:\n" " ✅ axis=tp group=0 idx=1 vs 0: n/a diff" + ) + + if __name__ == "__main__": sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/test_display.py b/test/registered/debug_utils/comparator/test_display.py index 451ac56b8..ebbf7bb0f 100644 --- a/test/registered/debug_utils/comparator/test_display.py +++ b/test/registered/debug_utils/comparator/test_display.py @@ -1,10 +1,12 @@ import sys +from io import StringIO from pathlib import Path from typing import Any, Optional import polars as pl import pytest import torch +from rich.console import Console from sglang.srt.debug_utils.comparator.display import ( _collect_input_ids_and_positions, @@ -21,6 +23,12 @@ from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=10, suite="default", nightly=True) +def _render_rich(renderable: object) -> str: + buf: StringIO = StringIO() + Console(file=buf, force_terminal=False, width=120).print(renderable) + return buf.getvalue().rstrip("\n") + + def _save_dump_file( directory: Path, *, @@ -276,6 +284,24 @@ class TestRankInfoRecordSnapshot: assert "1/2" in text assert "0/1" in text + def test_to_rich_snapshot(self) -> None: + from rich.table import Table + + record = RankInfoRecord( + label="baseline", + rows=[ + {"rank": 0, "tp": "0/2", "pp": "0/1"}, + {"rank": 1, "tp": "1/2", "pp": "0/1"}, + ], + ) + body = record._format_rich_body() + + assert isinstance(body, Table) + rendered: str = _render_rich(body) + assert "baseline ranks" in rendered + assert "0/2" in rendered + assert "1/2" in rendered + def test_json_roundtrip(self) -> None: record = RankInfoRecord( label="target", @@ -310,6 +336,29 @@ class TestInputIdsRecordSnapshot: assert "10, 20, 30" in text assert "0, 1, 2" in text + def test_to_rich_snapshot(self) -> None: + from rich.table import Table + + record = InputIdsRecord( + label="target", + rows=[ + { + "step": 0, + "rank": 0, + "num_tokens": 3, + "input_ids": "[10, 20, 30]", + "positions": "[0, 1, 2]", + }, + ], + ) + body = record._format_rich_body() + + assert isinstance(body, Table) + rendered: str = _render_rich(body) + assert "target input_ids & positions" in rendered + assert "10, 20, 30" in rendered + assert "0, 1, 2" in rendered + def test_json_roundtrip(self) -> None: record = InputIdsRecord( label="baseline", diff --git a/test/registered/debug_utils/comparator/test_output_types.py b/test/registered/debug_utils/comparator/test_output_types.py index e7b5f85f3..59149040b 100644 --- a/test/registered/debug_utils/comparator/test_output_types.py +++ b/test/registered/debug_utils/comparator/test_output_types.py @@ -6,7 +6,8 @@ from registered.debug_utils.comparator.testing_helpers import make_diff as _make from registered.debug_utils.comparator.testing_helpers import ( make_tensor_info as _make_tensor_info, ) -from rich.console import Console +from rich.console import Console, Group +from rich.panel import Panel from sglang.srt.debug_utils.comparator.aligner.axis_aligner import AxisAlignerPlan from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import ( @@ -121,6 +122,23 @@ class TestConfigRecord: record: ConfigRecord = ConfigRecord(config={"a": 1, "b": "two"}) assert record._format_body() == "Config: {'a': 1, 'b': 'two'}" + def test_format_rich_body(self) -> None: + record: ConfigRecord = ConfigRecord(config={"threshold": 0.001, "mode": "fast"}) + body = record._format_rich_body() + + assert isinstance(body, Panel) + rendered: str = _render_rich(body) + assert rendered == ( + "╭───────────────────────────────────────────────── Comparator Config " + "──────────────────────────────────────────────────╮\n" + "│ threshold : 0.001" + " │\n" + "│ mode : fast" + " │\n" + "╰──────────────────────────────────────────────────────────────────────" + "────────────────────────────────────────────────╯" + ) + def test_to_text_with_errors(self) -> None: record: ConfigRecord = ConfigRecord( config={"x": 1}, @@ -152,6 +170,14 @@ class TestSkipComparisonRecord: ) assert record._format_body() == "Skip: layer.weight (step=3) (scalar)" + def test_format_rich_body(self) -> None: + record: SkipComparisonRecord = SkipComparisonRecord( + name="attn.qkv", + reason="no baseline", + ) + body: str = record._format_rich_body() + assert body == "[dim]⊘ attn.qkv ── skipped (no baseline)[/]" + def test_category_skipped(self) -> None: record: SkipComparisonRecord = SkipComparisonRecord( name="x", @@ -200,6 +226,32 @@ class TestNonTensorComparisonRecord: " target = 0.01 (float)" ) + def test_format_rich_body_equal(self) -> None: + record: NonTensorComparisonRecord = NonTensorComparisonRecord( + name="config.lr", + baseline_value="0.001", + target_value="0.001", + baseline_type="float", + target_type="float", + values_equal=True, + ) + assert record._format_rich_body() == ("═ config.lr = 0.001 (float) [green]✓[/]") + + def test_format_rich_body_not_equal(self) -> None: + record: NonTensorComparisonRecord = NonTensorComparisonRecord( + name="config.lr", + baseline_value="0.001", + target_value="0.01", + baseline_type="float", + target_type="float", + values_equal=False, + ) + assert record._format_rich_body() == ( + "═ [bold red]config.lr[/]\n" + " baseline = 0.001 (float)\n" + " target = 0.01 (float)" + ) + def test_with_step(self) -> None: record: NonTensorComparisonRecord = NonTensorComparisonRecord( name="bias", @@ -250,6 +302,26 @@ class TestSummaryRecord: "Summary: 7 passed, 2 failed, 1 skipped (total 10)" ) + def test_format_rich_body(self) -> None: + record: SummaryRecord = SummaryRecord( + total=10, + passed=7, + failed=2, + skipped=1, + ) + body = record._format_rich_body() + assert isinstance(body, Panel) + + rendered: str = _render_rich(body) + assert rendered == ( + "╭────────────────────────────────────────────────────── SUMMARY " + "───────────────────────────────────────────────────────╮\n" + "│ 7 passed │ 2 failed │ 1 skipped │ 10 total" + " │\n" + "╰──────────────────────────────────────────────────────────────────────" + "────────────────────────────────────────────────╯" + ) + def test_validation_error(self) -> None: with pytest.raises(ValueError, match="total=5 !="): SummaryRecord(total=5, passed=1, failed=1, skipped=1) @@ -567,7 +639,7 @@ class TestFormatAlignerPlan: # --------------------------------------------------------------------------- -# _OutputRecord log attachment (to_text) +# _OutputRecord log attachment (to_text / to_rich) # --------------------------------------------------------------------------- @@ -606,6 +678,27 @@ class TestOutputRecordLogAttachment: assert text == "Config: {'a': 1}\n ✗ err1\n ℹ note1" + def test_to_rich_string_body(self) -> None: + record: SkipComparisonRecord = SkipComparisonRecord( + name="x", + reason="r", + errors=[ErrorLog(category="e", message="oops")], + ) + body = record.to_rich() + + assert isinstance(body, str) + assert body == "[dim]⊘ x ── skipped (r)[/]\n [red]✗ oops[/]" + + def test_to_rich_group_body(self) -> None: + record: ConfigRecord = ConfigRecord( + config={"a": 1}, + errors=[ErrorLog(category="e", message="oops")], + ) + body = record.to_rich() + + # Panel body + log block → Group + assert isinstance(body, Group) + if __name__ == "__main__": sys.exit(pytest.main([__file__]))