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

Enhance output formatting in dump comparator (#21029)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-20 22:04:50 +08:00
committed by GitHub
parent fdbcb8156e
commit 6703cc4484
16 changed files with 815 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files

81
python/sglang/srt/debug_utils/comparator/bundle_comparator.py
View File

@@ -21,6 +21,7 @@ from sglang.srt.debug_utils.comparator.aligner.token_aligner.smart.types import
from sglang.srt.debug_utils.comparator.dims_spec import (
SEQ_DIM_NAME,
TOKEN_DIM_NAME,
ParallelAxis,
apply_dim_names,
parse_dims,
resolve_dim_names,
@@ -39,11 +40,53 @@ from sglang.srt.debug_utils.comparator.output_types import (
)
from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
compare_tensor_pair,
compute_tensor_info,
)
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.srt.debug_utils.dump_loader import LOAD_FAILED, ValueWithMeta
_FAILED_SIDE_MAP: dict[str, str] = {"x": "baseline", "y": "target"}
def _build_skip_from_one_empty_side(
*, name: str, pair: Pair[list[ValueWithMeta]]
) -> ComparisonSkipRecord:
"""Build a skip record when one side of *pair* is empty.
The non-empty side's tensor info is attached to the record.
"""
assert not pair.x or not pair.y
if not pair.x:
reason, available_side, available_items = (
"baseline_load_failed",
"target",
pair.y,
)
else:
reason, available_side, available_items = (
"target_load_failed",
"baseline",
pair.x,
)
tensor_items: list[ValueWithMeta] = [
it for it in available_items if isinstance(it.value, torch.Tensor)
]
if not tensor_items:
return ComparisonSkipRecord(name=name, reason=reason)
first_tensor: torch.Tensor = tensor_items[0].value
tensor_info = compute_tensor_info(first_tensor, include_sample=True)
metas: list[dict[str, Any]] = [it.meta for it in tensor_items]
bundle_info: BundleSideInfo = _collect_bundle_side_info(
items=tensor_items, metas=metas
)
return ComparisonSkipRecord(
name=name,
reason=reason,
available_side=available_side, # type: ignore[arg-type]
available_tensor_info=tensor_info,
available_bundle_info=bundle_info,
)
def _collect_bundle_side_info(
@@ -51,8 +94,8 @@ def _collect_bundle_side_info(
metas: list[dict[str, Any]],
) -> BundleSideInfo:
from sglang.srt.debug_utils.comparator.display import (
_PARALLEL_INFO_KEYS,
extract_parallel_info,
PARALLEL_INFO_KEYS,
_extract_parallel_info,
)
files: list[BundleFileInfo] = []
@@ -61,8 +104,8 @@ def _collect_bundle_side_info(
tensor: torch.Tensor = item.value
parallel_info: dict[str, str] = {}
for key in _PARALLEL_INFO_KEYS:
extract_parallel_info(row_data=parallel_info, info=meta.get(key, {}))
for key in PARALLEL_INFO_KEYS:
_extract_parallel_info(row_data=parallel_info, info=meta.get(key, {}))
files.append(
BundleFileInfo(
@@ -70,6 +113,7 @@ def _collect_bundle_side_info(
dtype=str(tensor.dtype),
rank=meta.get("rank"),
parallel_info=parallel_info if parallel_info else None,
filename=meta.get("filename"),
)
)
@@ -132,8 +176,7 @@ def _compare_bundle_pair_inner(
)
if not all_pair.x or not all_pair.y:
reason = "baseline_load_failed" if not all_pair.x else "target_load_failed"
return ComparisonSkipRecord(name=name, reason=reason)
return _build_skip_from_one_empty_side(name=name, pair=all_pair)
# 1b. Dims override: patch meta["dims"] before DP filter reads it
# (--override-dims may add ``# dp:=moe_dp``, so it must run first)
@@ -157,7 +200,7 @@ def _compare_bundle_pair_inner(
# 1c. DP filter: keep only the non-empty dp_rank
all_pair = all_pair.map(
lambda items: filter_to_non_empty_dp_rank(
items, dp_group_alias=_extract_dp_alias_from_items(items)
items, dp_axis=_extract_dp_axis_from_items(items)
)
)
@@ -181,14 +224,14 @@ def _compare_bundle_pair_inner(
)
def _extract_dp_alias_from_items(items: list[ValueWithMeta]) -> Optional[str]:
"""Extract dp group alias from the first item's ``meta["dims"]``."""
def _extract_dp_axis_from_items(items: list[ValueWithMeta]) -> ParallelAxis:
"""Extract dp axis from the first item's ``meta["dims"]``."""
if not items:
return None
return ParallelAxis.DP
dims_str: Optional[str] = items[0].meta.get("dims")
if dims_str is None:
return None
return parse_dims(dims_str).dp_group_alias
return ParallelAxis.DP
return parse_dims(dims_str).dp_axis
def _compare_bundle_pair_tensor_type(
@@ -205,8 +248,7 @@ def _compare_bundle_pair_tensor_type(
compute_per_token: bool = False,
) -> Union[ComparisonTensorRecord, ComparisonSkipRecord]:
if not valid_pair.x or not valid_pair.y:
reason = "baseline_load_failed" if not valid_pair.x else "target_load_failed"
return ComparisonSkipRecord(name=name, reason=reason)
return _build_skip_from_one_empty_side(name=name, pair=valid_pair)
# Plan (meta only, no tensor)
metas_pair: Pair[list[dict[str, Any]]] = valid_pair.map(
@@ -243,9 +285,12 @@ def _compare_bundle_pair_tensor_type(
if aligner_result.tensors is None:
assert aligner_result.failed_side_xy is not None
side_name: str = _FAILED_SIDE_MAP[aligner_result.failed_side_xy]
reason: str = f"{side_name}_load_failed"
return ComparisonSkipRecord(name=name, reason=reason)
failed_xy: str = aligner_result.failed_side_xy
pair_with_failed_emptied: Pair[list[ValueWithMeta]] = Pair(
x=[] if failed_xy == "x" else valid_pair.x,
y=[] if failed_xy == "y" else valid_pair.y,
)
return _build_skip_from_one_empty_side(name=name, pair=pair_with_failed_emptied)
# Resolve seq_dim for per-token computation
seq_dim: Optional[int] = (

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

@@ -52,6 +52,19 @@ def _render_polars_as_text(df: pl.DataFrame, *, title: Optional[str] = None) ->
return buf.getvalue().rstrip("\n")
def _render_polars_as_rich_table(
df: pl.DataFrame, *, title: Optional[str] = None
) -> Any:
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]]]:

12
python/sglang/srt/debug_utils/comparator/entrypoint.py
View File

@@ -48,6 +48,14 @@ from sglang.srt.debug_utils.dump_loader import read_meta, read_tokenizer_path
_DEFAULT_SKIP_KEYS: set[str] = {"dump_index", "filename"}
_DIMS_DEBUG_HINT: str = (
"\nHint: If this is a dims annotation issue, do NOT re-run expensive dumps.\n"
"Use --override-dims at comparison time, e.g.:\n"
' python -m sglang.srt.debug_utils.comparator --override-dims "tensor_name:b s h[tp] d"\n'
"(Use --override-baseline-dims / --override-target-dims for per-side overrides.\n"
" Use --override-config for bulk overrides via YAML file.)"
)
def main() -> None:
args = parse_args(sys.argv[1:])
@@ -254,10 +262,12 @@ def _compare_bundle_pairs(
meta_overrider=meta_overrider,
)
except Exception as exc:
tb = _traceback_module.format_exc()
record = ComparisonErrorRecord(
name=name,
exception_type=type(exc).__name__,
traceback_str=_traceback_module.format_exc(),
exception_message=str(exc),
traceback_str=f"{_DIMS_DEBUG_HINT}\n\n{tb}",
)
target_steps: set[int] = {info.step for info in bundle_info_pair.y}

75
python/sglang/srt/debug_utils/comparator/output_formatter.py
View File

@@ -116,17 +116,72 @@ def _format_config_rich_body(
def _format_skip_body(record: ComparisonSkipRecord) -> str:
return f"Skip: {record.name}{record._format_location_suffix()} ({record.reason})"
text: str = (
f"Skip: {record.name}{record._format_location_suffix()} ({record.reason})"
)
if record.available_side is not None and record.available_tensor_info is not None:
info = record.available_tensor_info
text += f"\n {record.available_side}: shape={info.shape} dtype={info.dtype}"
text += (
f" mean={info.stats.mean:.4f} std={info.stats.std:.4f}"
f" range=[{info.stats.min:.4f}, {info.stats.max:.4f}]"
)
if info.sample is not None:
text += f"\n sample: {info.sample}"
return text
def _format_skip_rich_body(
record: ComparisonSkipRecord, verbosity: Verbosity = "normal"
) -> RenderableType:
suffix: str = record._format_location_suffix()
return (
header: str = (
f"[dim]⊘ {escape(record.name)}{suffix} ── skipped ({escape(record.reason)})[/]"
)
if (
verbosity == "minimal"
or record.available_side is None
or record.available_tensor_info is None
):
return header
info = record.available_tensor_info
side: str = record.available_side
dtype_str: str = info.dtype.replace("torch.", "")
lines: list[str] = [header]
# Bundle info line
if record.available_bundle_info is not None:
bi = record.available_bundle_info
shapes: list[list[int]] = [f.shape for f in bi.files]
unique_shapes: set[str] = {str(s) for s in shapes}
shape_desc: str = (
escape(str(shapes[0])) if len(unique_shapes) == 1 else "mixed shapes"
)
dims_part: str = f" [dim]dims: {bi.dims}[/]" if bi.dims else ""
lines.append(
f" {side:8s} [cyan]{bi.num_files} files[/]"
f" × {shape_desc} {dtype_str}{dims_part}"
)
else:
lines.append(f" {side:8s} {escape(str(info.shape))} {dtype_str}")
# Stats line (compact single-side)
stats = info.stats
range_str: str = escape(f"[{stats.min:.4f}, {stats.max:.4f}]")
lines.append(
f" [dim]stats[/] mean={stats.mean:.4f} std={stats.std:.4f}"
f" range={range_str}"
)
# Sample
if info.sample is not None:
lines.append(f" [dim]sample[/] {escape(info.sample)}")
return "\n".join(lines)
# ── ComparisonErrorRecord ────────────────────────────────────────────
@@ -135,6 +190,7 @@ def _format_error_body(record: ComparisonErrorRecord) -> str:
prefix: str = record._format_location_prefix()
return (
f"{prefix}Error: {record.name} ({record.exception_type})\n"
f"{record.exception_message}\n"
f"{record.traceback_str}"
)
@@ -145,7 +201,8 @@ def _format_error_rich_body(
prefix: str = record._format_location_prefix_rich()
name: str = escape(record.name)
header: str = (
f"{prefix}[bold red]{name} ── errored ({escape(record.exception_type)})[/]"
f"{prefix}[bold red]{name} ── errored ({escape(record.exception_type)}): "
f"{escape(record.exception_message)}[/]"
)
if verbosity == "minimal":
return header
@@ -300,7 +357,17 @@ def _format_aligner_plan(traced_plan: TracedAlignerPlan) -> str:
def _format_sub_plan_text(traced_sub: TracedSubPlan) -> str:
sub_desc: str = f"{traced_sub.plan.type}"
from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
sub = traced_sub.plan
qualifier: str = ""
if isinstance(sub, UnsharderPlan):
qualifier = f"({sub.axis.value})"
elif isinstance(sub, ReordererPlan):
qualifier = f"({sub.params.op})"
sub_desc: str = f"{sub.type}{qualifier}"
if traced_sub.snapshot is not None:
snap = traced_sub.snapshot

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

@@ -4,7 +4,7 @@ from abc import abstractmethod
from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional, Union
from pydantic import ConfigDict, Discriminator, Field, TypeAdapter, model_validator
from rich.console import RenderableType
from rich.console import Group, RenderableType
from rich.markup import escape
from sglang.srt.debug_utils.comparator.output_formatter import ( # noqa: F401 — re-export
@@ -32,6 +32,7 @@ from sglang.srt.debug_utils.comparator.output_formatter import (
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorComparisonInfo,
TensorInfo,
)
from sglang.srt.debug_utils.comparator.utils import Pair, _StrictBase
@@ -84,6 +85,7 @@ class BundleFileInfo(_StrictBase):
dtype: str
rank: Optional[int] = None
parallel_info: Optional[dict[str, str]] = None # e.g. {"tp": "0/4", "ep": "1/2"}
filename: Optional[str] = None
class BundleSideInfo(_StrictBase):
@@ -121,6 +123,12 @@ class RecordLocation(_StrictBase):
class _BaseComparisonRecord(_OutputRecord):
location: RecordLocation = Field(default_factory=RecordLocation)
def to_rich(self, verbosity: Verbosity = "normal") -> RenderableType:
result = _render_record_rich(self, verbosity=verbosity)
if isinstance(result, str):
return result + "\n"
return Group(result, "")
def _format_location_prefix(self) -> str:
if self.location.step is not None:
return f"[step={self.location.step}] "
@@ -152,6 +160,9 @@ class ComparisonSkipRecord(_BaseComparisonRecord):
type: Literal["comparison_skip"] = "comparison_skip"
name: str
reason: str
available_side: Optional[Literal["baseline", "target"]] = None
available_tensor_info: Optional[TensorInfo] = None
available_bundle_info: Optional[BundleSideInfo] = None
@property
def category(self) -> str:
@@ -170,6 +181,7 @@ class ComparisonErrorRecord(_BaseComparisonRecord):
type: Literal["comparison_error"] = "comparison_error"
name: str
exception_type: str
exception_message: str
traceback_str: str
@property

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

@@ -1,5 +1,6 @@
from __future__ import annotations
import os
import sys
from pathlib import Path
from typing import IO, Literal, Optional
@@ -72,7 +73,11 @@ class ReportSink:
def _get_console(self) -> Console:
if self._console is None:
self._console = Console()
try:
width = os.get_terminal_size().columns
except OSError:
width = 200
self._console = Console(force_terminal=True, width=width)
return self._console
def _print_to_stdout(self, record: _OutputRecord) -> None:
@@ -81,7 +86,6 @@ class ReportSink:
else:
console: Console = self._get_console()
console.print(record.to_rich(verbosity=self._verbosity))
console.print() # blank line between records
report_sink = ReportSink()

28
python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py
View File

@@ -23,6 +23,22 @@ QUANTILE_NUMEL_THRESHOLD = 10_000_000
SAMPLE_DIFF_THRESHOLD = 1e-3
def compute_tensor_info(
tensor: torch.Tensor, *, include_sample: bool = False
) -> TensorInfo:
"""Compute TensorInfo (shape, dtype, stats, optional sample) for a single tensor."""
stats: TensorStats = _compute_tensor_stats(tensor.float())
sample: Optional[str] = (
str(get_truncated_value(tensor.float())) if include_sample else None
)
return TensorInfo(
shape=list(tensor.shape),
dtype=str(tensor.dtype),
stats=stats,
sample=sample,
)
def compare_tensor_pair(
x_baseline: torch.Tensor,
x_target: torch.Tensor,
@@ -30,16 +46,8 @@ def compare_tensor_pair(
diff_threshold: float = 1e-3,
seq_dim: Optional[int] = None,
) -> TensorComparisonInfo:
baseline_info = TensorInfo(
shape=list(x_baseline.shape),
dtype=str(x_baseline.dtype),
stats=_compute_tensor_stats(x_baseline.float()),
)
target_info = TensorInfo(
shape=list(x_target.shape),
dtype=str(x_target.dtype),
stats=_compute_tensor_stats(x_target.float()),
)
baseline_info: TensorInfo = compute_tensor_info(x_baseline)
target_info: TensorInfo = compute_tensor_info(x_target)
x_baseline = try_unify_shape(x_baseline, target_shape=x_target.shape)
unified_shape = list(x_baseline.shape)

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

@@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Literal, Optional
from rich.markup import escape
from sglang.srt.debug_utils.comparator.aligner.reorderer.types import ReordererPlan
from sglang.srt.debug_utils.comparator.aligner.unsharder.types import UnsharderPlan
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
@@ -70,11 +71,13 @@ def _category_marker(category: str) -> tuple[bool, str, str]:
# ---------------------------------------------------------------------------
_STAT_HEADER = (
f" [dim]{'':10s} {'baseline':>10s} {'target':>10s} {'Δ':s}[/]"
)
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)}"
)
return f" [blue]{stat_name:10s}[/] {val_b:>10.4f} {val_t:>10.4f} {_fmt_diff_colored(diff)}"
# ---------------------------------------------------------------------------
@@ -352,7 +355,7 @@ def _format_bundle_section(
for label, side in [("baseline", bundle_info.x), ("target", bundle_info.y)]:
if not side.files:
lines.append(f" {label} [dim](no files)[/]")
lines.append(f" {label:8s} [dim](no files)[/]")
continue
dtype_desc: str = _strip_torch_prefix(side.files[0].dtype)
@@ -360,7 +363,7 @@ def _format_bundle_section(
if verbose:
dims_part: str = f" dims: {side.dims}" if side.dims else ""
lines.append(
f" {label} [cyan]{side.num_files} files[/]"
f" {label:8s} [cyan]{side.num_files} files[/]"
f" {dtype_desc}{dims_part}"
)
@@ -371,8 +374,9 @@ def _format_bundle_section(
par_part = " " + " ".join(
f"{k}={v}" for k, v in f.parallel_info.items()
)
file_part: str = f" [dim]{escape(f.filename)}[/]" if f.filename else ""
lines.append(
f" [{idx}] {_esc_shape(f.shape)} {rank_part}{par_part}"
f" [{idx}] {_esc_shape(f.shape)} {rank_part}{par_part}{file_part}"
)
else:
shapes: list[list[int]] = [f.shape for f in side.files]
@@ -385,7 +389,7 @@ def _format_bundle_section(
dims_part = f" [dim]dims: {side.dims}[/]" if side.dims else ""
lines.append(
f" {label} [cyan]{side.num_files} files[/]"
f" {label:8s} [cyan]{side.num_files} files[/]"
f" × {shape_desc} {dtype_desc}{dims_part}"
)
@@ -404,7 +408,7 @@ def _format_plan_section_rich(
("target", traced_plan.per_side.y),
]:
if not traced_side.step_plans:
lines.append(f" {side_label} [dim](passthrough)[/]")
lines.append(f" {side_label:8s} [dim](passthrough)[/]")
continue
parts: list[str] = [
@@ -412,7 +416,7 @@ def _format_plan_section_rich(
for traced_step in traced_side.step_plans
for traced_sub in traced_step.sub_plans
]
lines.append(f" {side_label} " + " → ".join(parts))
lines.append(f" {side_label:8s} " + " → ".join(parts))
lines.extend(_format_cross_side_plan_rich(traced_plan.plan))
return lines
@@ -423,9 +427,11 @@ def _format_sub_plan_rich(traced_sub: TracedSubPlan) -> str:
snapshot: Optional[ShapeSnapshot] = traced_sub.snapshot
op_name: str = sub.type
axis_str: str = ""
qualifier: str = ""
if isinstance(sub, UnsharderPlan):
axis_str = f"({sub.axis})"
qualifier = f"({sub.axis.value})"
elif isinstance(sub, ReordererPlan):
qualifier = f"({sub.params.op})"
shape_change: str = ""
if snapshot:
@@ -437,9 +443,9 @@ def _format_sub_plan_rich(traced_sub: TracedSubPlan) -> str:
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}"
shape_change = f" ({in_count}×{in_shape} → {out_count}×{out_shape})"
return f"[magenta]{op_name}{axis_str}[/]{shape_change}"
return f"[magenta]{op_name}{qualifier}[/]{shape_change}"
def _format_cross_side_plan_rich(plan: AlignerPlan) -> list[str]:
@@ -469,7 +475,7 @@ def _format_stats_rich(
target: TensorStats,
verbose: bool = False,
) -> list[str]:
lines: list[str] = []
lines: list[str] = [_STAT_HEADER]
if verbose:
# All stat fields
@@ -486,7 +492,7 @@ def _format_stats_rich(
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)
# Compact: mean, std, range, then percentiles
for stat_name in ("mean", "std"):
val_b = getattr(baseline, stat_name)
val_t = getattr(target, stat_name)
@@ -495,7 +501,13 @@ def _format_stats_rich(
# 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}")
lines.append(f" [blue]{'range':10s}[/] {range_baseline} {range_target}")
# Percentiles (compact: same as verbose)
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))
return lines

74
test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py
View File

@@ -9,6 +9,7 @@ from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
_compute_tensor_stats,
compare_tensor_pair,
compute_diff,
compute_tensor_info,
)
from sglang.srt.debug_utils.comparator.tensor_comparator.types import DiffInfo
from sglang.test.ci.ci_register import register_cpu_ci
@@ -16,6 +17,79 @@ from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=20, suite="stage-a-cpu-only", nightly=True)
class TestComputeTensorInfo:
def test_basic_tensor_returns_correct_shape_and_dtype(self) -> None:
tensor = torch.randn(2, 3)
info = compute_tensor_info(tensor)
assert info.shape == [2, 3]
assert info.dtype == "torch.float32"
assert info.stats.mean == pytest.approx(tensor.float().mean().item(), abs=1e-4)
def test_include_sample_false_returns_none_sample(self) -> None:
tensor = torch.randn(2, 3)
info = compute_tensor_info(tensor, include_sample=False)
assert info.sample is None
def test_include_sample_true_returns_string_sample(self) -> None:
tensor = torch.randn(2, 3)
info = compute_tensor_info(tensor, include_sample=True)
assert info.sample is not None
assert isinstance(info.sample, str)
def test_empty_tensor_stats_are_zero(self) -> None:
tensor = torch.tensor([])
info = compute_tensor_info(tensor)
assert info.stats.mean == 0.0
assert info.stats.std == 0.0
assert info.shape == [0]
def test_integer_tensor_converted_to_float_for_stats(self) -> None:
"""Integer tensors should be cast to float internally for stats computation."""
tensor = torch.tensor([1, 2, 3, 4], dtype=torch.int32)
info = compute_tensor_info(tensor)
assert info.dtype == "torch.int32"
assert info.stats.mean == pytest.approx(2.5, abs=1e-4)
assert info.stats.min == pytest.approx(1.0, abs=1e-4)
assert info.stats.max == pytest.approx(4.0, abs=1e-4)
def test_bfloat16_tensor_shape_and_stats(self) -> None:
"""bfloat16 tensors produce correct shape and dtype string."""
tensor = torch.ones(3, 4, dtype=torch.bfloat16)
info = compute_tensor_info(tensor)
assert info.shape == [3, 4]
assert info.dtype == "torch.bfloat16"
assert info.stats.mean == pytest.approx(1.0, abs=1e-2)
def test_multidimensional_shape(self) -> None:
"""Shape is preserved for high-rank tensors."""
tensor = torch.randn(2, 3, 4, 5)
info = compute_tensor_info(tensor)
assert info.shape == [2, 3, 4, 5]
def test_scalar_tensor(self) -> None:
"""Scalar (0-dim) tensor produces empty shape list."""
tensor = torch.tensor(3.14)
info = compute_tensor_info(tensor)
assert info.shape == []
assert info.stats.mean == pytest.approx(3.14, abs=1e-4)
assert info.stats.min == pytest.approx(3.14, abs=1e-4)
assert info.stats.max == pytest.approx(3.14, abs=1e-4)
def test_include_sample_true_contains_tensor_representation(self) -> None:
"""Sample string should contain some recognizable tensor content."""
tensor = torch.tensor([1.0, 2.0])
info = compute_tensor_info(tensor, include_sample=True)
assert info.sample is not None
assert "1." in info.sample or "2." in info.sample
def test_percentiles_present_for_small_tensor(self) -> None:
"""Small tensors (< threshold) should have percentile data."""
tensor = torch.randn(100)
info = compute_tensor_info(tensor)
assert len(info.stats.percentiles) > 0
assert 50 in info.stats.percentiles
class TestComputeTensorStats:
def test_basic_stats(self):
x = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])

252
test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py
View File

@@ -1,6 +1,9 @@
import sys
import pytest
from registered.debug_utils.comparator.testing_helpers import (
assert_rich_tags_balanced,
)
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 (
@@ -57,6 +60,14 @@ from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="stage-a-cpu-only", nightly=True)
_DEFAULT_PERCENTILE_LINES: list[str] = [
" [blue]p1 [/] -1.8000 -1.8000 [dim]+0.00e+00[/]",
" [blue]p5 [/] -1.5000 -1.5000 [dim]+0.00e+00[/]",
" [blue]p50 [/] 0.0000 0.0000 [dim]+0.00e+00[/]",
" [blue]p95 [/] 1.5000 1.5000 [dim]+0.00e+00[/]",
" [blue]p99 [/] 1.8000 1.8000 [dim]+0.00e+00[/]",
]
# Snapshot strings below are intentionally spelled out in full per test.
# The shared skeleton (stats block, diff block) looks duplicated, but keeping
@@ -421,6 +432,7 @@ class TestFormatComparisonRichMinimal:
diff=_make_diff(rel_diff=1e-4, passed=True),
)
result: str = format_comparison_rich(record, verbosity="minimal")
assert_rich_tags_balanced(result)
assert result == (
"[green]✅[/] [bold green]hidden_states [/] "
@@ -432,6 +444,7 @@ class TestFormatComparisonRichMinimal:
diff=_make_diff(rel_diff=0.5, passed=False),
)
result: str = format_comparison_rich(record, verbosity="minimal")
assert_rich_tags_balanced(result)
assert result == (
"[red]❌[/] [bold red]hidden_states [/] "
@@ -443,6 +456,7 @@ class TestFormatComparisonRichMinimal:
shape_mismatch=True,
)
result: str = format_comparison_rich(record, verbosity="minimal")
assert_rich_tags_balanced(result)
assert result == (
"[red]❌[/] [bold red]hidden_states [/] "
@@ -452,6 +466,7 @@ class TestFormatComparisonRichMinimal:
def test_no_diff(self) -> None:
record: ComparisonTensorRecord = _make_comparison_record()
result: str = format_comparison_rich(record, verbosity="minimal")
assert_rich_tags_balanced(result)
assert result == ("[red]❌[/] [bold red]hidden_states [/]")
@@ -464,6 +479,7 @@ class TestFormatComparisonRichNormal:
diff=_make_diff(rel_diff=1e-4, passed=True),
)
result: str = format_comparison_rich(record, verbosity="normal")
assert_rich_tags_balanced(result)
assert result == (
"[green]✅[/] [bold green]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
@@ -471,9 +487,15 @@ class TestFormatComparisonRichNormal:
" [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]"
" [dim] baseline target Δ[/]\n"
" [blue]mean [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 1.0000 [dim]+0.00e+00[/]\n"
" [blue]range [/] [-2.0000, 2.0000] [-2.0000, 2.0000]\n"
" [blue]p1 [/] -1.8000 -1.8000 [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 -1.5000 [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 1.5000 [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 1.8000 [dim]+0.00e+00[/]"
)
def test_failed(self) -> None:
@@ -483,6 +505,7 @@ class TestFormatComparisonRichNormal:
),
)
result: str = format_comparison_rich(record, verbosity="normal")
assert_rich_tags_balanced(result)
assert result == (
"[red]❌[/] [bold red]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
@@ -491,9 +514,15 @@ class TestFormatComparisonRichNormal:
" [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] baseline target Δ[/]\n"
" [blue]mean [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 1.0000 [dim]+0.00e+00[/]\n"
" [blue]range [/] [-2.0000, 2.0000] [-2.0000, 2.0000]\n"
" [blue]p1 [/] -1.8000 -1.8000 [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 -1.5000 [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 1.5000 [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 1.8000 [dim]+0.00e+00[/]\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"
)
@@ -503,6 +532,7 @@ class TestFormatComparisonRichNormal:
shape_mismatch=True,
)
result: str = format_comparison_rich(record, verbosity="normal")
assert_rich_tags_balanced(result)
assert result == (
"[red]❌[/] [bold red]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
@@ -510,9 +540,15 @@ class TestFormatComparisonRichNormal:
" [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]"
" [dim] baseline target Δ[/]\n"
" [blue]mean [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 1.0000 [dim]+0.00e+00[/]\n"
" [blue]range [/] [-2.0000, 2.0000] [-2.0000, 2.0000]\n"
" [blue]p1 [/] -1.8000 -1.8000 [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 -1.5000 [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 1.5000 [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 1.8000 [dim]+0.00e+00[/]"
)
def test_with_downcast(self) -> None:
@@ -522,6 +558,7 @@ class TestFormatComparisonRichNormal:
downcast_dtype="torch.bfloat16",
)
result: str = format_comparison_rich(record, verbosity="normal")
assert_rich_tags_balanced(result)
assert result == (
"[red]❌[/] [bold red]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
@@ -531,9 +568,15 @@ class TestFormatComparisonRichNormal:
" [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] baseline target Δ[/]\n"
" [blue]mean [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 1.0000 [dim]+0.00e+00[/]\n"
" [blue]range [/] [-2.0000, 2.0000] [-2.0000, 2.0000]\n"
" [blue]p1 [/] -1.8000 -1.8000 [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 -1.5000 [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 1.5000 [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 1.8000 [dim]+0.00e+00[/]\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"
)
@@ -548,19 +591,26 @@ class TestFormatComparisonRichNormal:
raw_bundle_info=bundle_info,
)
result: str = format_comparison_rich(record, verbosity="normal")
assert_rich_tags_balanced(result)
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"
" 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]"
" [dim] baseline target Δ[/]\n"
" [blue]mean [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 1.0000 [dim]+0.00e+00[/]\n"
" [blue]range [/] [-2.0000, 2.0000] [-2.0000, 2.0000]\n"
" [blue]p1 [/] -1.8000 -1.8000 [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 -1.5000 [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 1.5000 [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 1.8000 [dim]+0.00e+00[/]"
)
def test_with_plan(self) -> None:
@@ -570,19 +620,26 @@ class TestFormatComparisonRichNormal:
traced_plan=_make_traced_plan(plan),
)
result: str = format_comparison_rich(record, verbosity="normal")
assert_rich_tags_balanced(result)
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"
" target [magenta]unsharder(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]"
" [dim] baseline target Δ[/]\n"
" [blue]mean [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 1.0000 [dim]+0.00e+00[/]\n"
" [blue]range [/] [-2.0000, 2.0000] [-2.0000, 2.0000]\n"
" [blue]p1 [/] -1.8000 -1.8000 [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 -1.5000 [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 1.5000 [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 1.8000 [dim]+0.00e+00[/]"
)
@@ -595,6 +652,7 @@ class TestFormatComparisonRichVerbose:
sample="tensor([0.1, 0.2, ...])",
)
result: str = format_comparison_rich(record, verbosity="verbose")
assert_rich_tags_balanced(result)
assert result == (
"[green]✅[/] [bold green]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
@@ -602,16 +660,17 @@ class TestFormatComparisonRichVerbose:
" [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]abs_mean [/] 0.8000 vs 0.8000 Δ [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]min [/] -2.0000 vs -2.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]max [/] 2.0000 vs 2.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]p1 [/] -1.8000 vs -1.8000 Δ [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 vs -1.5000 Δ [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 vs 1.5000 Δ [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 vs 1.8000 Δ [dim]+0.00e+00[/]\n"
" [dim] baseline target Δ[/]\n"
" [blue]mean [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]abs_mean [/] 0.8000 0.8000 [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 1.0000 [dim]+0.00e+00[/]\n"
" [blue]min [/] -2.0000 -2.0000 [dim]+0.00e+00[/]\n"
" [blue]max [/] 2.0000 2.0000 [dim]+0.00e+00[/]\n"
" [blue]p1 [/] -1.8000 -1.8000 [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 -1.5000 [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 1.5000 [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 1.8000 [dim]+0.00e+00[/]\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\n"
" [dim]Samples[/]\n"
@@ -629,6 +688,7 @@ class TestFormatComparisonRichVerbose:
raw_bundle_info=bundle_info,
)
result: str = format_comparison_rich(record, verbosity="verbose")
assert_rich_tags_balanced(result)
assert result == (
"[green]✅[/] [bold green]hidden_states[/] [dim cyan]── float32 [4, 8][/]\n"
@@ -637,22 +697,23 @@ class TestFormatComparisonRichVerbose:
" baseline [cyan]2 files[/] float32\n"
" [0] [2, 4096] rank=0 tp=0/2\n"
" [1] [2, 4096] rank=1 tp=1/2\n"
" target [cyan]2 files[/] float32\n"
" target [cyan]2 files[/] float32\n"
" [0] [2, 4096] rank=0 tp=0/2\n"
" [1] [2, 4096] rank=1 tp=1/2\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]abs_mean [/] 0.8000 vs 0.8000 Δ [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]min [/] -2.0000 vs -2.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]max [/] 2.0000 vs 2.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]p1 [/] -1.8000 vs -1.8000 Δ [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 vs -1.5000 Δ [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 vs 1.5000 Δ [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 vs 1.8000 Δ [dim]+0.00e+00[/]\n"
" [dim] baseline target Δ[/]\n"
" [blue]mean [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]abs_mean [/] 0.8000 0.8000 [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 1.0000 [dim]+0.00e+00[/]\n"
" [blue]min [/] -2.0000 -2.0000 [dim]+0.00e+00[/]\n"
" [blue]max [/] 2.0000 2.0000 [dim]+0.00e+00[/]\n"
" [blue]p1 [/] -1.8000 -1.8000 [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 -1.5000 [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 1.5000 [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 1.8000 [dim]+0.00e+00[/]\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"
)
@@ -668,26 +729,28 @@ class TestFormatComparisonRichVerbose:
),
)
result: str = format_comparison_rich(record, verbosity="verbose")
assert_rich_tags_balanced(result)
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)[/] 2×[2, 4096] → 1×[4, 4096]\n"
" target [magenta]unsharder(tp)[/] (2×[2, 4096] → 1×[4, 4096])\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]abs_mean [/] 0.8000 vs 0.8000 Δ [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]min [/] -2.0000 vs -2.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]max [/] 2.0000 vs 2.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]p1 [/] -1.8000 vs -1.8000 Δ [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 vs -1.5000 Δ [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 vs 1.5000 Δ [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 vs 1.8000 Δ [dim]+0.00e+00[/]\n"
" [dim] baseline target Δ[/]\n"
" [blue]mean [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]abs_mean [/] 0.8000 0.8000 [dim]+0.00e+00[/]\n"
" [blue]std [/] 1.0000 1.0000 [dim]+0.00e+00[/]\n"
" [blue]min [/] -2.0000 -2.0000 [dim]+0.00e+00[/]\n"
" [blue]max [/] 2.0000 2.0000 [dim]+0.00e+00[/]\n"
" [blue]p1 [/] -1.8000 -1.8000 [dim]+0.00e+00[/]\n"
" [blue]p5 [/] -1.5000 -1.5000 [dim]+0.00e+00[/]\n"
" [blue]p50 [/] 0.0000 0.0000 [dim]+0.00e+00[/]\n"
" [blue]p95 [/] 1.5000 1.5000 [dim]+0.00e+00[/]\n"
" [blue]p99 [/] 1.8000 1.8000 [dim]+0.00e+00[/]\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"
)
@@ -705,7 +768,7 @@ class TestFormatBundleSection:
assert lines == [
" baseline [cyan]2 files[/] × [2, 4096] float32",
" target [cyan]2 files[/] × [2, 4096] float32",
" target [cyan]2 files[/] × [2, 4096] float32",
]
def test_mixed_shapes(self) -> None:
@@ -721,7 +784,7 @@ class TestFormatBundleSection:
assert lines == [
" baseline [cyan]2 files[/] × mixed shapes float32",
" target [cyan]2 files[/] × mixed shapes float32",
" target [cyan]2 files[/] × mixed shapes float32",
]
def test_no_files(self) -> None:
@@ -731,7 +794,7 @@ class TestFormatBundleSection:
assert lines == [
" baseline [dim](no files)[/]",
" target [dim](no files)[/]",
" target [dim](no files)[/]",
]
def test_with_dims(self) -> None:
@@ -743,7 +806,7 @@ class TestFormatBundleSection:
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[/]",
" target [cyan]1 files[/] × [2, 4096] float32 [dim]dims: b s h(tp) d[/]",
]
@@ -761,7 +824,7 @@ class TestFormatBundleSectionVerbose:
" baseline [cyan]2 files[/] float32",
" [0] [2, 4096] rank=0 tp=0/2",
" [1] [2, 4096] rank=1 tp=1/2",
" target [cyan]2 files[/] float32",
" target [cyan]2 files[/] float32",
" [0] [2, 4096] rank=0 tp=0/2",
" [1] [2, 4096] rank=1 tp=1/2",
]
@@ -773,7 +836,7 @@ class TestFormatBundleSectionVerbose:
assert lines == [
" baseline [dim](no files)[/]",
" target [dim](no files)[/]",
" target [dim](no files)[/]",
]
@@ -787,7 +850,7 @@ class TestFormatPlanSectionRich:
assert lines == [
" baseline [dim](passthrough)[/]",
" target [dim](passthrough)[/]",
" target [dim](passthrough)[/]",
]
def test_unsharder_op(self) -> None:
@@ -797,7 +860,7 @@ class TestFormatPlanSectionRich:
assert lines == [
" baseline [dim](passthrough)[/]",
" target [magenta]unsharder(ParallelAxis.TP)[/]",
" target [magenta]unsharder(tp)[/]",
]
def test_reorderer_op(self) -> None:
@@ -807,7 +870,7 @@ class TestFormatPlanSectionRich:
assert lines == [
" baseline [dim](passthrough)[/]",
" target [magenta]reorderer[/]",
" target [magenta]reorderer(zigzag_to_natural)[/]",
]
def test_with_shape_traces(self) -> None:
@@ -821,7 +884,7 @@ class TestFormatPlanSectionRich:
assert lines == [
" baseline [dim](passthrough)[/]",
" target [magenta]unsharder(ParallelAxis.TP)[/] 2×[2, 4096] → 1×[4, 4096]",
" target [magenta]unsharder(tp)[/] (2×[2, 4096] → 1×[4, 4096])",
]
def test_with_token_aligner(self) -> None:
@@ -831,7 +894,7 @@ class TestFormatPlanSectionRich:
assert lines == [
" baseline [dim](passthrough)[/]",
" target [dim](passthrough)[/]",
" target [dim](passthrough)[/]",
" token_aligner [dim]3 tokens[/]",
]
@@ -842,7 +905,7 @@ class TestFormatPlanSectionRich:
assert lines == [
" baseline [dim](passthrough)[/]",
" target [dim](passthrough)[/]",
" target [dim](passthrough)[/]",
" axis_aligner [dim]x=b s d -> s b d[/]",
]
@@ -855,7 +918,7 @@ class TestFormatPlanSectionRich:
assert lines == [
" baseline [dim](passthrough)[/]",
" target [dim](passthrough)[/]",
" target [dim](passthrough)[/]",
" axis_aligner [dim](no-op)[/]",
]
@@ -869,33 +932,42 @@ class TestFormatStatsRich:
mean=0.0001, std=1.0001, min=-2.0001, max=2.0001
)
lines: list[str] = _format_stats_rich(baseline=baseline, target=target)
assert_rich_tags_balanced("\n".join(lines))
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]",
" [dim] baseline target Δ[/]",
" [blue]mean [/] 0.0000 0.0001 [dim]+1.00e-04[/]",
" [blue]std [/] 1.0000 1.0001 [dim]+1.00e-04[/]",
" [blue]range [/] [-2.0000, 2.0000] [-2.0001, 2.0001]",
*_DEFAULT_PERCENTILE_LINES,
]
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_rich_tags_balanced("\n".join(lines))
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]",
" [dim] baseline target Δ[/]",
" [blue]mean [/] 0.0000 1.0000 [yellow]+1.00e+00[/]",
" [blue]std [/] 1.0000 1.0000 [dim]+0.00e+00[/]",
" [blue]range [/] [-2.0000, 2.0000] [-2.0000, 2.0000]",
*_DEFAULT_PERCENTILE_LINES,
]
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_rich_tags_balanced("\n".join(lines))
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]",
" [dim] baseline target Δ[/]",
" [blue]mean [/] 0.0000 0.0010 [dim]+1.00e-03[/]",
" [blue]std [/] 1.0000 1.0000 [dim]+0.00e+00[/]",
" [blue]range [/] [-2.0000, 2.0000] [-2.0000, 2.0000]",
*_DEFAULT_PERCENTILE_LINES,
]
@@ -910,16 +982,17 @@ class TestFormatStatsRichVerbose:
)
assert lines == [
" [blue]mean [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]",
" [blue]abs_mean [/] 0.8000 vs 0.8000 Δ [dim]+0.00e+00[/]",
" [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]",
" [blue]min [/] -2.0000 vs -2.0000 Δ [dim]+0.00e+00[/]",
" [blue]max [/] 2.0000 vs 2.0000 Δ [dim]+0.00e+00[/]",
" [blue]p1 [/] -1.8000 vs -1.8000 Δ [dim]+0.00e+00[/]",
" [blue]p5 [/] -1.5000 vs -1.5000 Δ [dim]+0.00e+00[/]",
" [blue]p50 [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]",
" [blue]p95 [/] 1.5000 vs 1.5000 Δ [dim]+0.00e+00[/]",
" [blue]p99 [/] 1.8000 vs 1.8000 Δ [dim]+0.00e+00[/]",
" [dim] baseline target Δ[/]",
" [blue]mean [/] 0.0000 0.0000 [dim]+0.00e+00[/]",
" [blue]abs_mean [/] 0.8000 0.8000 [dim]+0.00e+00[/]",
" [blue]std [/] 1.0000 1.0000 [dim]+0.00e+00[/]",
" [blue]min [/] -2.0000 -2.0000 [dim]+0.00e+00[/]",
" [blue]max [/] 2.0000 2.0000 [dim]+0.00e+00[/]",
" [blue]p1 [/] -1.8000 -1.8000 [dim]+0.00e+00[/]",
" [blue]p5 [/] -1.5000 -1.5000 [dim]+0.00e+00[/]",
" [blue]p50 [/] 0.0000 0.0000 [dim]+0.00e+00[/]",
" [blue]p95 [/] 1.5000 1.5000 [dim]+0.00e+00[/]",
" [blue]p99 [/] 1.8000 1.8000 [dim]+0.00e+00[/]",
]
def test_no_percentiles(self) -> None:
@@ -930,11 +1003,12 @@ class TestFormatStatsRichVerbose:
)
assert lines == [
" [blue]mean [/] 0.0000 vs 0.0000 Δ [dim]+0.00e+00[/]",
" [blue]abs_mean [/] 0.8000 vs 0.8000 Δ [dim]+0.00e+00[/]",
" [blue]std [/] 1.0000 vs 1.0000 Δ [dim]+0.00e+00[/]",
" [blue]min [/] -2.0000 vs -2.0000 Δ [dim]+0.00e+00[/]",
" [blue]max [/] 2.0000 vs 2.0000 Δ [dim]+0.00e+00[/]",
" [dim] baseline target Δ[/]",
" [blue]mean [/] 0.0000 0.0000 [dim]+0.00e+00[/]",
" [blue]abs_mean [/] 0.8000 0.8000 [dim]+0.00e+00[/]",
" [blue]std [/] 1.0000 1.0000 [dim]+0.00e+00[/]",
" [blue]min [/] -2.0000 -2.0000 [dim]+0.00e+00[/]",
" [blue]max [/] 2.0000 2.0000 [dim]+0.00e+00[/]",
]

122
test/registered/debug_utils/comparator/test_bundle_comparator.py
View File

@@ -5,9 +5,14 @@ from unittest.mock import patch
import pytest
import torch
from sglang.srt.debug_utils.comparator.bundle_comparator import _load_all_values
from sglang.srt.debug_utils.comparator.bundle_comparator import (
_build_skip_from_one_empty_side,
_load_all_values,
)
from sglang.srt.debug_utils.comparator.log_sink import LogSink
from sglang.srt.debug_utils.comparator.output_types import ErrorLog
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.srt.debug_utils.dump_loader import ValueWithMeta
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=15, suite="stage-a-cpu-only", nightly=True)
@@ -86,5 +91,120 @@ class TestLoadAllValues:
assert all(w.category == "load_failed" for w in warnings)
def _tensor_item(value: torch.Tensor, rank: int = 0) -> ValueWithMeta:
return ValueWithMeta(
value=value,
meta={
"rank": rank,
"dims": "b s",
"sglang_parallel_info": {},
"megatron_parallel_info": {},
"filename": f"rank_{rank}.pt",
},
)
class TestBuildSkipFromOneEmptySide:
def test_baseline_empty_sets_reason_and_side(self) -> None:
"""Empty baseline → reason='baseline_load_failed', available_side='target'."""
item = _tensor_item(torch.randn(2, 3))
record = _build_skip_from_one_empty_side(
name="test_tensor",
pair=Pair(x=[], y=[item]),
)
assert record.reason == "baseline_load_failed"
assert record.available_side == "target"
assert record.available_tensor_info is not None
def test_target_empty_sets_reason_and_side(self) -> None:
"""Empty target → reason='target_load_failed', available_side='baseline'."""
item = _tensor_item(torch.randn(2, 3))
record = _build_skip_from_one_empty_side(
name="test_tensor",
pair=Pair(x=[item], y=[]),
)
assert record.reason == "target_load_failed"
assert record.available_side == "baseline"
assert record.available_tensor_info is not None
def test_no_tensor_items_returns_minimal_skip(self) -> None:
"""All items are non-tensor → skip record with no tensor info."""
non_tensor_item = ValueWithMeta(value="not_a_tensor", meta={"rank": 0})
record = _build_skip_from_one_empty_side(
name="test_tensor",
pair=Pair(x=[], y=[non_tensor_item]),
)
assert record.reason == "baseline_load_failed"
assert record.available_tensor_info is None
assert record.available_bundle_info is None
def test_with_tensor_items_populates_info(self) -> None:
"""Tensor items present → tensor_info and bundle_info are populated."""
item = _tensor_item(torch.randn(2, 3))
record = _build_skip_from_one_empty_side(
name="test_tensor",
pair=Pair(x=[], y=[item]),
)
assert record.available_tensor_info is not None
assert record.available_tensor_info.shape == [2, 3]
assert record.available_bundle_info is not None
assert record.available_bundle_info.num_files >= 1
def test_multiple_tensor_items_uses_first_for_info(self) -> None:
"""When multiple tensor items exist, tensor_info comes from the first."""
item1 = _tensor_item(torch.randn(2, 3), rank=0)
item2 = _tensor_item(torch.randn(4, 5), rank=1)
record = _build_skip_from_one_empty_side(
name="multi",
pair=Pair(x=[], y=[item1, item2]),
)
assert record.available_tensor_info is not None
assert record.available_tensor_info.shape == [2, 3]
assert record.available_bundle_info is not None
assert record.available_bundle_info.num_files == 2
def test_mixed_tensor_and_non_tensor_filters_non_tensor(self) -> None:
"""Non-tensor items are filtered; tensor_info comes from tensor items only."""
non_tensor = ValueWithMeta(value="string_value", meta={"rank": 0})
tensor_item = _tensor_item(torch.randn(5, 6), rank=1)
record = _build_skip_from_one_empty_side(
name="mixed",
pair=Pair(x=[], y=[non_tensor, tensor_item]),
)
assert record.available_tensor_info is not None
assert record.available_tensor_info.shape == [5, 6]
assert record.available_bundle_info is not None
assert record.available_bundle_info.num_files == 1
def test_tensor_info_includes_sample(self) -> None:
"""Tensor info should include a sample string for skip records."""
item = _tensor_item(torch.tensor([1.0, 2.0, 3.0]))
record = _build_skip_from_one_empty_side(
name="sample_check",
pair=Pair(x=[item], y=[]),
)
assert record.available_tensor_info is not None
assert record.available_tensor_info.sample is not None
def test_name_preserved_in_record(self) -> None:
"""The tensor name is preserved in the skip record."""
item = _tensor_item(torch.randn(2, 3))
record = _build_skip_from_one_empty_side(
name="my_layer.weight",
pair=Pair(x=[], y=[item]),
)
assert record.name == "my_layer.weight"
def test_bundle_info_has_dims_from_meta(self) -> None:
"""Bundle info dims field should come from the meta."""
item = _tensor_item(torch.randn(2, 3))
record = _build_skip_from_one_empty_side(
name="dims_check",
pair=Pair(x=[], y=[item]),
)
assert record.available_bundle_info is not None
assert record.available_bundle_info.dims == "b s"
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

66
test/registered/debug_utils/comparator/test_display.py
View File

@@ -12,8 +12,8 @@ from sglang.srt.debug_utils.comparator.display import (
_collect_input_ids_and_positions,
_collect_rank_info,
_extract_parallel_info,
_render_polars_as_rich_table,
_render_polars_as_text,
_extract_parallel_info,
)
from sglang.srt.debug_utils.comparator.output_types import (
InputIdsRecord,
@@ -432,5 +432,69 @@ class TestExtractParallelInfo:
assert "tp" not in row_data
class TestRenderPolarsAsRichTable:
def test_basic_dataframe_renders_table(self) -> None:
df = pl.DataFrame({"a": [1, 2], "b": ["x", "y"]})
table = _render_polars_as_rich_table(df)
assert len(table.columns) == 2
assert table.row_count == 2
def test_empty_dataframe_returns_table_with_no_rows(self) -> None:
df = pl.DataFrame(
{"a": pl.Series([], dtype=pl.Int64), "b": pl.Series([], dtype=pl.Utf8)}
)
table = _render_polars_as_rich_table(df)
assert len(table.columns) == 2
assert table.row_count == 0
def test_title_passed_to_table(self) -> None:
df = pl.DataFrame({"a": [1]})
table = _render_polars_as_rich_table(df, title="My Title")
assert table.title == "My Title"
def test_no_title_defaults_to_none(self) -> None:
df = pl.DataFrame({"x": [1]})
table = _render_polars_as_rich_table(df)
assert table.title is None
def test_column_names_match_dataframe(self) -> None:
df = pl.DataFrame({"alpha": [1], "beta": [2], "gamma": [3]})
table = _render_polars_as_rich_table(df)
column_headers: list[str] = [col.header for col in table.columns]
assert column_headers == ["alpha", "beta", "gamma"]
def test_values_converted_to_strings(self) -> None:
"""Numeric and None values should be stringified in the rendered output."""
df = pl.DataFrame({"num": [42], "text": ["hello"]})
table = _render_polars_as_rich_table(df)
rendered: str = _render_rich(table)
assert "42" in rendered
assert "hello" in rendered
def test_single_column_dataframe(self) -> None:
df = pl.DataFrame({"only_col": [10, 20, 30]})
table = _render_polars_as_rich_table(df)
assert len(table.columns) == 1
assert table.row_count == 3
def test_many_rows_all_present(self) -> None:
"""All rows from the dataframe appear in the rich table."""
df = pl.DataFrame({"val": list(range(50))})
table = _render_polars_as_rich_table(df)
assert table.row_count == 50
def test_null_values_rendered_as_string(self) -> None:
"""Null values should be converted to their string representation."""
df = pl.DataFrame({"a": [1, None, 3]})
table = _render_polars_as_rich_table(df)
assert table.row_count == 3
rendered: str = _render_rich(table)
assert (
"null" in rendered.lower()
or "none" in rendered.lower()
or "None" in rendered
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

50
test/registered/debug_utils/comparator/test_entrypoint.py
View File

@@ -2121,7 +2121,7 @@ class TestEntrypointReplicatedAxis:
errors = [r for r in records if isinstance(r, ComparisonErrorRecord)]
assert len(errors) == 1
assert "not orthogonal" in errors[0].traceback_str
assert "not orthogonal" in errors[0].exception_message
summary = records[-1]
assert isinstance(summary, SummaryRecord)
@@ -3745,7 +3745,7 @@ class TestEntrypointDpFilter:
errors = [r for r in records if isinstance(r, ComparisonErrorRecord)]
assert len(errors) == 1
assert errors[0].exception_type == "AssertionError"
assert "Expected exactly 1 non-empty dp_rank" in errors[0].traceback_str
assert "Expected exactly 1 non-empty dp_rank" in errors[0].exception_message
assert exit_code == 1
@@ -4905,7 +4905,8 @@ class TestErrorResilience:
assert len(errors) == 1
assert errors[0].name == "tensor_b"
assert errors[0].exception_type == "RuntimeError"
assert "intentional test error" in errors[0].traceback_str
assert "intentional test error" in errors[0].exception_message
assert "--override-dims" in errors[0].traceback_str
summary = records[-1]
assert isinstance(summary, SummaryRecord)
@@ -4948,6 +4949,49 @@ class TestErrorResilience:
assert len(errors) == 1
assert errors[0].exception_type == "TypeError"
def test_error_record_contains_dims_hint(self, tmp_path, capsys, monkeypatch):
"""Error record includes --override-dims hint with all variant flags."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
argv = _make_argv(baseline_path, target_path, preset="raw")
def _raise(**kwargs):
raise ValueError("Invalid dim token: 'zzz'")
monkeypatch.setattr(_entrypoint_module, "compare_bundle_pair", _raise)
records, _ = _run_and_parse(argv, capsys)
errors = [r for r in records if isinstance(r, ComparisonErrorRecord)]
assert len(errors) == 1
assert "Invalid dim token: 'zzz'" in errors[0].exception_message
tb = errors[0].traceback_str
assert "--override-dims" in tb
assert "--override-baseline-dims" in tb
assert "--override-target-dims" in tb
assert "--override-config" in tb
assert "do NOT re-run expensive dumps" in tb
def test_error_record_hint_appears_before_traceback(
self, tmp_path, capsys, monkeypatch
):
"""Hint appears before the full stack trace in traceback_str."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
argv = _make_argv(baseline_path, target_path, preset="raw")
def _raise(**kwargs):
raise RuntimeError("some dims problem")
monkeypatch.setattr(_entrypoint_module, "compare_bundle_pair", _raise)
records, _ = _run_and_parse(argv, capsys)
errors = [r for r in records if isinstance(r, ComparisonErrorRecord)]
assert len(errors) == 1
tb = errors[0].traceback_str
hint_pos = tb.index("--override-dims")
traceback_pos = tb.index("Traceback (most recent call last)")
assert hint_pos < traceback_pos
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

17
test/registered/debug_utils/comparator/test_model_validation.py
View File

@@ -345,26 +345,37 @@ class TestOutputRecordCategories:
def test_error_record_category_is_errored(self) -> None:
record = ComparisonErrorRecord(
name="t", exception_type="ValueError", traceback_str="..."
name="t",
exception_type="ValueError",
exception_message="bad",
traceback_str="...",
)
assert record.category == "errored"
def test_error_record_json_roundtrip(self) -> None:
record = ComparisonErrorRecord(
name="t", exception_type="ValueError", traceback_str="traceback..."
name="t",
exception_type="ValueError",
exception_message="bad",
traceback_str="traceback...",
)
json_str: str = record.model_dump_json()
roundtripped = parse_record_json(json_str)
assert isinstance(roundtripped, ComparisonErrorRecord)
assert roundtripped.name == "t"
assert roundtripped.exception_type == "ValueError"
assert roundtripped.exception_message == "bad"
def test_error_record_text_format(self) -> None:
record = ComparisonErrorRecord(
name="t", exception_type="RuntimeError", traceback_str="Traceback..."
name="t",
exception_type="RuntimeError",
exception_message="oops",
traceback_str="Traceback...",
)
text: str = record.to_text()
assert "RuntimeError" in text
assert "oops" in text
assert "Traceback" in text

75
test/registered/debug_utils/comparator/test_output_types.py
View File

@@ -2,7 +2,11 @@ import sys
from io import StringIO
import pytest
from registered.debug_utils.comparator.testing_helpers import (
assert_rich_tags_balanced,
)
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,
)
@@ -34,6 +38,8 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
)
from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, TokenLayout
from sglang.srt.debug_utils.comparator.output_types import (
BundleFileInfo,
BundleSideInfo,
ComparisonNonTensorRecord,
ComparisonSkipRecord,
ComparisonTensorRecord,
@@ -193,6 +199,63 @@ class TestComparisonSkipRecord:
)
assert record.category == "failed"
def test_format_body_with_available_side(self) -> None:
record: ComparisonSkipRecord = ComparisonSkipRecord(
name="layer.weight",
reason="baseline_load_failed",
available_side="target",
available_tensor_info=_make_tensor_info(
shape=[4, 8],
dtype="torch.float32",
stats=_make_stats(mean=0.5, std=1.2, min=-2.0, max=3.0),
sample="tensor([0.1, 0.2, ...])",
),
)
body: str = record._format_body()
assert "baseline_load_failed" in body
assert "target: shape=[4, 8]" in body
assert "mean=0.5000" in body
assert "sample: tensor([0.1, 0.2, ...])" in body
def test_format_rich_body_with_available_side(self) -> None:
record: ComparisonSkipRecord = ComparisonSkipRecord(
name="attn.qkv",
reason="baseline_load_failed",
available_side="target",
available_tensor_info=_make_tensor_info(
shape=[4, 8],
dtype="torch.float32",
stats=_make_stats(mean=0.5, std=1.2, min=-2.0, max=3.0),
sample="tensor([0.1, 0.2, ...])",
),
available_bundle_info=BundleSideInfo(
num_files=2,
files=[
BundleFileInfo(shape=[4, 8], dtype="torch.float32"),
BundleFileInfo(shape=[4, 8], dtype="torch.float32"),
],
),
)
body: str = record._format_rich_body()
assert "skipped (baseline_load_failed)" in body
assert "target" in body
assert "2 files" in body
assert "mean=0.5000" in body
assert "tensor(" in body
assert_rich_tags_balanced(body)
def test_format_rich_body_minimal_hides_available_side(self) -> None:
record: ComparisonSkipRecord = ComparisonSkipRecord(
name="x",
reason="target_load_failed",
available_side="baseline",
available_tensor_info=_make_tensor_info(),
)
body: str = record._format_rich_body(verbosity="minimal")
assert "skipped" in body
assert "stats" not in body
assert_rich_tags_balanced(body)
# ---------------------------------------------------------------------------
# ComparisonNonTensorRecord
@@ -546,7 +609,9 @@ class TestFormatAlignerPlan:
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n" " baseline: (no steps)\n" " target: [step=0: unsharder]"
"Aligner Plan:\n"
" baseline: (no steps)\n"
" target: [step=0: unsharder(tp)]"
)
def test_reorderer(self) -> None:
@@ -566,7 +631,9 @@ class TestFormatAlignerPlan:
result: str = _format_aligner_plan(_wrap_plan(plan))
assert result == (
"Aligner Plan:\n" " baseline: (no steps)\n" " target: [step=0: reorderer]"
"Aligner Plan:\n"
" baseline: (no steps)\n"
" target: [step=0: reorderer(zigzag_to_natural)]"
)
def test_multi_step(self) -> None:
@@ -596,7 +663,7 @@ class TestFormatAlignerPlan:
assert result == (
"Aligner Plan:\n"
" baseline: (no steps)\n"
" target: [step=0: unsharder; step=1: reorderer]"
" target: [step=0: unsharder(tp); step=1: reorderer(zigzag_to_natural)]"
)
def test_with_token_aligner(self) -> None:
@@ -687,7 +754,7 @@ class TestOutputRecordLogAttachment:
body = record.to_rich()
assert isinstance(body, str)
assert body == "[dim]⊘ x ── skipped (r)[/]\n [red]✗ oops[/]"
assert body == "[dim]⊘ x ── skipped (r)[/]\n [red]✗ oops[/]\n"
def test_to_rich_group_body(self) -> None:
record: ConfigRecord = ConfigRecord(

36
test/registered/debug_utils/comparator/testing_helpers.py
View File

@@ -2,8 +2,12 @@
from __future__ import annotations
import re
from io import StringIO
from typing import Optional
from rich.console import Console
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(
@@ -79,6 +83,38 @@ def make_diff(
)
_ANSI_ESCAPE_RE = re.compile(r"\033\[([0-9;]*)m")
def assert_rich_tags_balanced(markup: str) -> None:
"""Render Rich markup to ANSI and verify no styles are active at the end.
Tracks ANSI style state through the output. A ``\\033[0m`` (reset)
clears all active styles; any other ``\\033[Nm`` sets a style.
At the end of the output, no style should remain active.
"""
buf = StringIO()
console = Console(file=buf, force_terminal=True, width=10000, highlight=False)
console.print(markup, end="")
ansi_output: str = buf.getvalue()
if "\033[" not in ansi_output:
return
styled = False
for match in _ANSI_ESCAPE_RE.finditer(ansi_output):
params: str = match.group(1)
if params == "0" or params == "":
styled = False
else:
styled = True
assert not styled, (
f"ANSI styles still active at end of output — likely unclosed Rich tag.\n"
f"Last 200 chars of ANSI output: {ansi_output[-200:]!r}"
)
def make_tensor_info(
shape: Optional[list[int]] = None,
dtype: str = "torch.float32",