From 6703cc448450e1165a5e62bcef765a98fd9cab6c Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Fri, 20 Mar 2026 22:04:50 +0800 Subject: [PATCH] Enhance output formatting in dump comparator (#21029) --- .../comparator/bundle_comparator.py | 81 ++++-- .../srt/debug_utils/comparator/display.py | 13 + .../srt/debug_utils/comparator/entrypoint.py | 12 +- .../comparator/output_formatter.py | 75 +++++- .../debug_utils/comparator/output_types.py | 14 +- .../srt/debug_utils/comparator/report_sink.py | 8 +- .../tensor_comparator/comparator.py | 28 +- .../comparator/tensor_comparator/formatter.py | 46 ++-- .../tensor_comparator/test_comparator.py | 74 +++++ .../tensor_comparator/test_formatter.py | 252 +++++++++++------- .../comparator/test_bundle_comparator.py | 122 ++++++++- .../debug_utils/comparator/test_display.py | 66 ++++- .../debug_utils/comparator/test_entrypoint.py | 50 +++- .../comparator/test_model_validation.py | 17 +- .../comparator/test_output_types.py | 75 +++++- .../debug_utils/comparator/testing_helpers.py | 36 +++ 16 files changed, 815 insertions(+), 154 deletions(-) diff --git a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py index 657e457a4..437d42f65 100644 --- a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py +++ b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py @@ -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] = ( diff --git a/python/sglang/srt/debug_utils/comparator/display.py b/python/sglang/srt/debug_utils/comparator/display.py index 27ae0e955..9f99f18da 100644 --- a/python/sglang/srt/debug_utils/comparator/display.py +++ b/python/sglang/srt/debug_utils/comparator/display.py @@ -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]]]: diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index 1ab483e15..99089011b 100644 --- a/python/sglang/srt/debug_utils/comparator/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -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} diff --git a/python/sglang/srt/debug_utils/comparator/output_formatter.py b/python/sglang/srt/debug_utils/comparator/output_formatter.py index f5859ad0c..37b1afe12 100644 --- a/python/sglang/srt/debug_utils/comparator/output_formatter.py +++ b/python/sglang/srt/debug_utils/comparator/output_formatter.py @@ -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 diff --git a/python/sglang/srt/debug_utils/comparator/output_types.py b/python/sglang/srt/debug_utils/comparator/output_types.py index c24615910..d042623ea 100644 --- a/python/sglang/srt/debug_utils/comparator/output_types.py +++ b/python/sglang/srt/debug_utils/comparator/output_types.py @@ -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 diff --git a/python/sglang/srt/debug_utils/comparator/report_sink.py b/python/sglang/srt/debug_utils/comparator/report_sink.py index 61f9e9ac5..b2e73a31c 100644 --- a/python/sglang/srt/debug_utils/comparator/report_sink.py +++ b/python/sglang/srt/debug_utils/comparator/report_sink.py @@ -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() diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py b/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py index 296429f92..4549218c9 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/comparator.py @@ -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) 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 6021d065d..3996b42c2 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py @@ -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 diff --git a/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py b/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py index 76d197d03..e36071044 100644 --- a/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py +++ b/test/registered/debug_utils/comparator/tensor_comparator/test_comparator.py @@ -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]) 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 057cbcfa1..81ed592c8 100644 --- a/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py +++ b/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py @@ -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[/]", ] diff --git a/test/registered/debug_utils/comparator/test_bundle_comparator.py b/test/registered/debug_utils/comparator/test_bundle_comparator.py index 52c19040a..17cdbd81c 100644 --- a/test/registered/debug_utils/comparator/test_bundle_comparator.py +++ b/test/registered/debug_utils/comparator/test_bundle_comparator.py @@ -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__])) diff --git a/test/registered/debug_utils/comparator/test_display.py b/test/registered/debug_utils/comparator/test_display.py index 9a557a5de..0c7bfa44f 100644 --- a/test/registered/debug_utils/comparator/test_display.py +++ b/test/registered/debug_utils/comparator/test_display.py @@ -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__])) diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index e4f81c831..f82a8de9b 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -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__])) diff --git a/test/registered/debug_utils/comparator/test_model_validation.py b/test/registered/debug_utils/comparator/test_model_validation.py index bbf22ce95..1fc372bed 100644 --- a/test/registered/debug_utils/comparator/test_model_validation.py +++ b/test/registered/debug_utils/comparator/test_model_validation.py @@ -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 diff --git a/test/registered/debug_utils/comparator/test_output_types.py b/test/registered/debug_utils/comparator/test_output_types.py index 897c005ac..100e03463 100644 --- a/test/registered/debug_utils/comparator/test_output_types.py +++ b/test/registered/debug_utils/comparator/test_output_types.py @@ -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( diff --git a/test/registered/debug_utils/comparator/testing_helpers.py b/test/registered/debug_utils/comparator/testing_helpers.py index 9c2b6e8ac..ffe282a19 100644 --- a/test/registered/debug_utils/comparator/testing_helpers.py +++ b/test/registered/debug_utils/comparator/testing_helpers.py @@ -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",