Beautify text output in dump comparator (#19683)

This commit is contained in:
fzyzcjy
2026-03-02 18:47:01 +08:00
committed by GitHub
parent 5bf3deb4bc
commit 3dd4649b42
9 changed files with 1460 additions and 193 deletions
@@ -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",
@@ -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__]))
@@ -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",
@@ -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__]))