From 7579ab3f3351444f65bb4758554a73fecf110cfb Mon Sep 17 00:00:00 2001 From: fzyzcjy <5236035+fzyzcjy@users.noreply.github.com> Date: Mon, 2 Mar 2026 19:08:35 +0800 Subject: [PATCH] Enhance error resilience in dump comparator (#19685) --- .../srt/debug_utils/comparator/__init__.py | 4 +- .../comparator/bundle_comparator.py | 28 +-- .../srt/debug_utils/comparator/display.py | 48 +---- .../srt/debug_utils/comparator/entrypoint.py | 90 +++++---- .../comparator/output_formatter.py | 56 ++++-- .../debug_utils/comparator/output_types.py | 63 ++++--- .../comparator/per_token_visualizer.py | 6 +- .../comparator/tensor_comparator/formatter.py | 8 +- .../srt/debug_utils/comparator/utils.py | 4 + .../tensor_comparator/test_formatter.py | 32 ++-- .../tensor_comparator/test_types.py | 34 ++-- .../debug_utils/comparator/test_entrypoint.py | 173 +++++++++++++----- .../comparator/test_manually_verify.py | 12 +- .../comparator/test_model_validation.py | 138 +++++++++----- .../comparator/test_output_types.py | 52 +++--- .../comparator/test_per_token_visualizer.py | 14 +- .../debug_utils/comparator/test_utils.py | 30 +++ .../debug_utils/comparator/testing_helpers.py | 6 + 18 files changed, 495 insertions(+), 303 deletions(-) diff --git a/python/sglang/srt/debug_utils/comparator/__init__.py b/python/sglang/srt/debug_utils/comparator/__init__.py index ab7ed6c6b..3aec40062 100644 --- a/python/sglang/srt/debug_utils/comparator/__init__.py +++ b/python/sglang/srt/debug_utils/comparator/__init__.py @@ -4,6 +4,6 @@ from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import ( from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( # noqa: F401 AlignerPlan, ) -from sglang.srt.debug_utils.comparator.output_types import TensorComparisonRecord +from sglang.srt.debug_utils.comparator.output_types import ComparisonTensorRecord -TensorComparisonRecord.model_rebuild() +ComparisonTensorRecord.model_rebuild() diff --git a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py index 35a453cae..657e457a4 100644 --- a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py +++ b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py @@ -31,10 +31,10 @@ from sglang.srt.debug_utils.comparator.meta_overrider import MetaOverrider from sglang.srt.debug_utils.comparator.output_types import ( BundleFileInfo, BundleSideInfo, + ComparisonNonTensorRecord, + ComparisonSkipRecord, + ComparisonTensorRecord, ErrorLog, - NonTensorComparisonRecord, - SkipComparisonRecord, - TensorComparisonRecord, _split_logs, ) from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( @@ -51,7 +51,7 @@ def _collect_bundle_side_info( metas: list[dict[str, Any]], ) -> BundleSideInfo: from sglang.srt.debug_utils.comparator.display import ( - PARALLEL_INFO_KEYS, + _PARALLEL_INFO_KEYS, extract_parallel_info, ) @@ -61,7 +61,7 @@ def _collect_bundle_side_info( tensor: torch.Tensor = item.value parallel_info: dict[str, str] = {} - for key in PARALLEL_INFO_KEYS: + for key in _PARALLEL_INFO_KEYS: extract_parallel_info(row_data=parallel_info, info=meta.get(key, {})) files.append( @@ -91,7 +91,7 @@ def compare_bundle_pair( viz_output_dir: Optional[Path] = None, compute_per_token: bool = False, meta_overrider: Optional[MetaOverrider] = None, -) -> Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord]: +) -> Union[ComparisonTensorRecord, ComparisonSkipRecord, ComparisonNonTensorRecord]: with log_sink.context() as collected_logs: result = _compare_bundle_pair_inner( name=name, @@ -124,7 +124,7 @@ def _compare_bundle_pair_inner( viz_output_dir: Optional[Path] = None, compute_per_token: bool = False, meta_overrider: Optional[MetaOverrider] = None, -) -> Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord]: +) -> Union[ComparisonTensorRecord, ComparisonSkipRecord, ComparisonNonTensorRecord]: # 1. Load all successfully loaded values all_pair: Pair[list[ValueWithMeta]] = Pair( x=_load_all_values(filenames=filenames_pair.x, base_path=dir_pair.x), @@ -133,7 +133,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 SkipComparisonRecord(name=name, reason=reason) + return ComparisonSkipRecord(name=name, reason=reason) # 1b. Dims override: patch meta["dims"] before DP filter reads it # (--override-dims may add ``# dp:=moe_dp``, so it must run first) @@ -203,10 +203,10 @@ def _compare_bundle_pair_tensor_type( ), viz_output_dir: Optional[Path] = None, compute_per_token: bool = False, -) -> Union[TensorComparisonRecord, SkipComparisonRecord]: +) -> 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 SkipComparisonRecord(name=name, reason=reason) + return ComparisonSkipRecord(name=name, reason=reason) # Plan (meta only, no tensor) metas_pair: Pair[list[dict[str, Any]]] = valid_pair.map( @@ -245,7 +245,7 @@ def _compare_bundle_pair_tensor_type( 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 SkipComparisonRecord(name=name, reason=reason) + return ComparisonSkipRecord(name=name, reason=reason) # Resolve seq_dim for per-token computation seq_dim: Optional[int] = ( @@ -263,7 +263,7 @@ def _compare_bundle_pair_tensor_type( diff_threshold=diff_threshold, seq_dim=seq_dim, ) - record = TensorComparisonRecord( + record = ComparisonTensorRecord( **info.model_dump(), traced_plan=aligner_result.traced_plan, replicated_checks=replicated_checks, @@ -331,7 +331,7 @@ def _compare_bundle_pair_non_tensor_type( *, name: str, value_pair: Pair[list[ValueWithMeta]], -) -> NonTensorComparisonRecord: +) -> ComparisonNonTensorRecord: baseline_value: Any = value_pair.x[0].value target_value: Any = value_pair.y[0].value @@ -340,7 +340,7 @@ def _compare_bundle_pair_non_tensor_type( except Exception: values_equal = False - return NonTensorComparisonRecord( + return ComparisonNonTensorRecord( name=name, baseline_value=repr(baseline_value), target_value=repr(target_value), diff --git a/python/sglang/srt/debug_utils/comparator/display.py b/python/sglang/srt/debug_utils/comparator/display.py index 32ae5fbc3..27ae0e955 100644 --- a/python/sglang/srt/debug_utils/comparator/display.py +++ b/python/sglang/srt/debug_utils/comparator/display.py @@ -3,13 +3,9 @@ from __future__ import annotations from collections import defaultdict from io import StringIO from pathlib import Path -from typing import TYPE_CHECKING, Any, Optional +from typing import Any, Optional import polars as pl -import rich.table - -if TYPE_CHECKING: - from rich.table import Table from sglang.srt.debug_utils.comparator.output_types import ( InputIdsRecord, @@ -43,45 +39,19 @@ def emit_display_records( def _render_polars_as_text(df: pl.DataFrame, *, title: Optional[str] = None) -> str: from rich.console import Console + from rich.table import Table - table = _build_rich_table(df, title=title) + 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]) buf = StringIO() Console(file=buf, force_terminal=False, width=200).print(table) return buf.getvalue().rstrip("\n") -def _render_polars_as_rich_table( - df: pl.DataFrame, *, title: Optional[str] = None -) -> "Table": - return _build_rich_table(df, title=title) - - -def _build_rich_table(df: pl.DataFrame, *, title: Optional[str] = None) -> "Table": - from rich.table import Table - - table = Table(title=title) - 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 _render_polars_as_rich_table( - df: pl.DataFrame, *, title: Optional[str] = None -) -> "rich.table.Table": - 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]]]: @@ -99,7 +69,7 @@ def _collect_rank_info( row_data: dict[str, Any] = {"rank": row["rank"]} for key in PARALLEL_INFO_KEYS: - extract_parallel_info(row_data=row_data, info=meta.get(key, {})) + _extract_parallel_info(row_data=row_data, info=meta.get(key, {})) table_rows.append(row_data) return table_rows or None @@ -149,7 +119,7 @@ def _collect_input_ids_and_positions( return table_rows or None -def extract_parallel_info(row_data: dict[str, Any], info: dict[str, Any]) -> None: +def _extract_parallel_info(row_data: dict[str, Any], info: dict[str, Any]) -> None: if not info or info.get("error"): return diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index dfb311cf8..1ab483e15 100644 --- a/python/sglang/srt/debug_utils/comparator/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -2,6 +2,7 @@ from __future__ import annotations import argparse import sys +import traceback as _traceback_module from pathlib import Path from typing import Any, Iterator, Optional, Union @@ -25,12 +26,13 @@ from sglang.srt.debug_utils.comparator.bundle_matcher import ( from sglang.srt.debug_utils.comparator.display import emit_display_records from sglang.srt.debug_utils.comparator.meta_overrider import MetaOverrider from sglang.srt.debug_utils.comparator.output_types import ( + ComparisonErrorRecord, + ComparisonNonTensorRecord, + ComparisonSkipRecord, + ComparisonTensorRecord, ConfigRecord, - NonTensorComparisonRecord, RecordLocation, - SkipComparisonRecord, SummaryRecord, - TensorComparisonRecord, ) from sglang.srt.debug_utils.comparator.per_token_visualizer import ( generate_per_token_heatmap, @@ -83,12 +85,6 @@ def run(args: argparse.Namespace) -> int: verbosity=args.verbosity, ) - report_path: Optional[Path] = _resolve_report_path( - target_path=dir_pair.y, - report_path_arg=args.report_path, - ) - report_sink.configure(output_format=args.output_format, report_path=report_path) - try: report_sink.add(ConfigRecord(config=vars(args))) @@ -142,9 +138,11 @@ def run(args: argparse.Namespace) -> int: compute_per_token=visualize_per_token is not None, meta_overrider=meta_overrider, ) - summary, skipped_names, failed_names = _consume_comparison_records( - comparison_records=comparison_records, - visualize_per_token=visualize_per_token, + summary, skipped_names, failed_names, errored_names = ( + _consume_comparison_records( + comparison_records=comparison_records, + visualize_per_token=visualize_per_token, + ) ) return compute_exit_code( summary, @@ -152,6 +150,7 @@ def run(args: argparse.Namespace) -> int: skipped_names=skipped_names, allow_failed_pattern=args.allow_failed_pattern, failed_names=failed_names, + errored_names=errored_names, ) finally: report_sink.close() @@ -219,7 +218,12 @@ def _compare_bundle_pairs( compute_per_token: bool = False, meta_overrider: Optional[MetaOverrider] = None, ) -> Iterator[ - Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord] + Union[ + ComparisonTensorRecord, + ComparisonSkipRecord, + ComparisonNonTensorRecord, + ComparisonErrorRecord, + ] ]: for bundle_info_pair in bundle_info_pairs: if not bundle_info_pair.y: @@ -229,20 +233,32 @@ def _compare_bundle_pairs( filenames_pair: Pair[list[str]] = bundle_info_pair.map( lambda infos: [info.filename for info in infos] ) + record: Union[ - TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord - ] = compare_bundle_pair( - name=name, - filenames_pair=filenames_pair, - dir_pair=dir_pair, - token_aligner_mode=token_aligner_mode, - token_aligner_plan=token_aligner_plan, - diff_threshold=diff_threshold, - thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair, - viz_output_dir=viz_output_dir, - compute_per_token=compute_per_token, - meta_overrider=meta_overrider, - ) + ComparisonTensorRecord, + ComparisonSkipRecord, + ComparisonNonTensorRecord, + ComparisonErrorRecord, + ] + try: + record = compare_bundle_pair( + name=name, + filenames_pair=filenames_pair, + dir_pair=dir_pair, + token_aligner_mode=token_aligner_mode, + token_aligner_plan=token_aligner_plan, + diff_threshold=diff_threshold, + thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair, + viz_output_dir=viz_output_dir, + compute_per_token=compute_per_token, + meta_overrider=meta_overrider, + ) + except Exception as exc: + record = ComparisonErrorRecord( + name=name, + exception_type=type(exc).__name__, + traceback_str=_traceback_module.format_exc(), + ) target_steps: set[int] = {info.step for info in bundle_info_pair.y} step: Optional[int] = target_steps.pop() if len(target_steps) == 1 else None @@ -255,24 +271,32 @@ def _compare_bundle_pairs( def _consume_comparison_records( *, comparison_records: Iterator[ - Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord] + Union[ + ComparisonTensorRecord, + ComparisonSkipRecord, + ComparisonNonTensorRecord, + ComparisonErrorRecord, + ] ], visualize_per_token: Optional[Path] = None, -) -> tuple[SummaryRecord, list[str], list[str]]: - counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0} - collected_comparisons: list[TensorComparisonRecord] = [] +) -> tuple[SummaryRecord, list[str], list[str], list[str]]: + counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0, "errored": 0} + collected_comparisons: list[ComparisonTensorRecord] = [] skipped_names: list[str] = [] failed_names: list[str] = [] + errored_names: list[str] = [] for record in comparison_records: counts[record.category] += 1 report_sink.add(record) - if isinstance(record, SkipComparisonRecord) and record.category == "skipped": + if isinstance(record, ComparisonSkipRecord) and record.category == "skipped": skipped_names.append(record.name) if record.category == "failed": failed_names.append(record.name) + if isinstance(record, ComparisonErrorRecord): + errored_names.append(record.name) if visualize_per_token is not None and isinstance( - record, TensorComparisonRecord + record, ComparisonTensorRecord ): collected_comparisons.append(record) @@ -285,7 +309,7 @@ def _consume_comparison_records( output_path=visualize_per_token, ) - return summary, skipped_names, failed_names + return summary, skipped_names, failed_names, errored_names def parse_args(argv: list[str]) -> argparse.Namespace: diff --git a/python/sglang/srt/debug_utils/comparator/output_formatter.py b/python/sglang/srt/debug_utils/comparator/output_formatter.py index cd93a3d2c..f5859ad0c 100644 --- a/python/sglang/srt/debug_utils/comparator/output_formatter.py +++ b/python/sglang/srt/debug_utils/comparator/output_formatter.py @@ -26,14 +26,15 @@ if TYPE_CHECKING: ) from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan from sglang.srt.debug_utils.comparator.output_types import ( + ComparisonErrorRecord, + ComparisonNonTensorRecord, + ComparisonSkipRecord, + ComparisonTensorRecord, ConfigRecord, ErrorLog, InfoLog, LogRecord, - NonTensorComparisonRecord, - SkipComparisonRecord, SummaryRecord, - TensorComparisonRecord, _OutputRecord, _TableRecord, ) @@ -111,15 +112,15 @@ def _format_config_rich_body( return Panel("\n".join(lines), title="Comparator Config", border_style="cyan") -# ── SkipComparisonRecord ───────────────────────────────────────────── +# ── ComparisonSkipRecord ───────────────────────────────────────────── -def _format_skip_body(record: SkipComparisonRecord) -> str: +def _format_skip_body(record: ComparisonSkipRecord) -> str: return f"Skip: {record.name}{record._format_location_suffix()} ({record.reason})" def _format_skip_rich_body( - record: SkipComparisonRecord, verbosity: Verbosity = "normal" + record: ComparisonSkipRecord, verbosity: Verbosity = "normal" ) -> RenderableType: suffix: str = record._format_location_suffix() return ( @@ -127,6 +128,30 @@ def _format_skip_rich_body( ) +# ── ComparisonErrorRecord ──────────────────────────────────────────── + + +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.traceback_str}" + ) + + +def _format_error_rich_body( + record: ComparisonErrorRecord, verbosity: Verbosity = "normal" +) -> RenderableType: + prefix: str = record._format_location_prefix_rich() + name: str = escape(record.name) + header: str = ( + f"{prefix}[bold red]{name} ── errored ({escape(record.exception_type)})[/]" + ) + if verbosity == "minimal": + return header + return header + f"\n[dim]{escape(record.traceback_str)}[/]" + + # ── _TableRecord ───────────────────────────────────────────────────── @@ -154,10 +179,10 @@ def _format_table_rich_body( ) -# ── TensorComparisonRecord ─────────────────────────────────────────── +# ── ComparisonTensorRecord ─────────────────────────────────────────── -def _format_tensor_comparison_body(record: TensorComparisonRecord) -> str: +def _format_tensor_comparison_body(record: ComparisonTensorRecord) -> str: body: str = record._format_location_prefix() + format_comparison(record) if record.replicated_checks: body += "\n" + format_replicated_checks(record.replicated_checks) @@ -167,7 +192,7 @@ def _format_tensor_comparison_body(record: TensorComparisonRecord) -> str: def _format_tensor_comparison_rich_body( - record: TensorComparisonRecord, verbosity: Verbosity = "normal" + record: ComparisonTensorRecord, verbosity: Verbosity = "normal" ) -> RenderableType: from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import ( format_comparison_rich, @@ -178,10 +203,10 @@ def _format_tensor_comparison_rich_body( ) -# ── NonTensorComparisonRecord ──────────────────────────────────────── +# ── ComparisonNonTensorRecord ──────────────────────────────────────── -def _format_non_tensor_body(record: NonTensorComparisonRecord) -> str: +def _format_non_tensor_body(record: ComparisonNonTensorRecord) -> str: suffix: str = record._format_location_suffix() if record.values_equal: return f"NonTensor: {record.name}{suffix} = {record.baseline_value} ({record.baseline_type}) [equal]" @@ -193,7 +218,7 @@ def _format_non_tensor_body(record: NonTensorComparisonRecord) -> str: def _format_non_tensor_rich_body( - record: NonTensorComparisonRecord, verbosity: Verbosity = "normal" + record: ComparisonNonTensorRecord, verbosity: Verbosity = "normal" ) -> RenderableType: suffix: str = record._format_location_suffix() name: str = escape(record.name) @@ -216,10 +241,13 @@ def _format_non_tensor_rich_body( def _format_summary_body(record: SummaryRecord) -> str: - return ( + text: str = ( f"Summary: {record.passed} passed, {record.failed} failed, " f"{record.skipped} skipped (total {record.total})" ) + if record.errored > 0: + text += f", {record.errored} errored" + return text def _format_summary_rich_body( @@ -231,6 +259,8 @@ def _format_summary_rich_body( f"[yellow]{record.skipped} skipped[/] │ " f"{record.total} total" ) + if record.errored > 0: + text += f" │ [bold red]{record.errored} errored[/]" return Panel(text, title="SUMMARY", border_style="bold") diff --git a/python/sglang/srt/debug_utils/comparator/output_types.py b/python/sglang/srt/debug_utils/comparator/output_types.py index e07281e7e..c24615910 100644 --- a/python/sglang/srt/debug_utils/comparator/output_types.py +++ b/python/sglang/srt/debug_utils/comparator/output_types.py @@ -13,6 +13,8 @@ from sglang.srt.debug_utils.comparator.output_formatter import ( # noqa: F401 from sglang.srt.debug_utils.comparator.output_formatter import ( _format_config_body, _format_config_rich_body, + _format_error_body, + _format_error_rich_body, _format_log_body, _format_non_tensor_body, _format_non_tensor_rich_body, @@ -146,8 +148,8 @@ class ConfigRecord(_OutputRecord): return _format_config_rich_body(self, verbosity=verbosity) -class SkipComparisonRecord(_BaseComparisonRecord): - type: Literal["skip"] = "skip" +class ComparisonSkipRecord(_BaseComparisonRecord): + type: Literal["comparison_skip"] = "comparison_skip" name: str reason: str @@ -164,6 +166,23 @@ class SkipComparisonRecord(_BaseComparisonRecord): return _format_skip_rich_body(self, verbosity=verbosity) +class ComparisonErrorRecord(_BaseComparisonRecord): + type: Literal["comparison_error"] = "comparison_error" + name: str + exception_type: str + traceback_str: str + + @property + def category(self) -> str: + return "errored" + + def _format_body(self) -> str: + return _format_error_body(self) + + def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType: + return _format_error_rich_body(self, verbosity=verbosity) + + class _TableRecord(_OutputRecord): label: str rows: list[dict[str, Any]] @@ -177,15 +196,6 @@ class _TableRecord(_OutputRecord): def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType: return _format_table_rich_body(self, verbosity=verbosity) - return _format_table_body(self) - - def _format_rich_body(self) -> RenderableType: - from sglang.srt.debug_utils.comparator.output_formatter import ( - _format_table_rich_body, - ) - - return _format_table_rich_body(self) - class RankInfoRecord(_TableRecord): type: Literal["rank_info"] = "rank_info" @@ -201,10 +211,10 @@ class InputIdsRecord(_TableRecord): return f"{self.label} input_ids & positions" -class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord): +class ComparisonTensorRecord(TensorComparisonInfo, _BaseComparisonRecord): model_config = ConfigDict(extra="forbid", defer_build=True) - type: Literal["comparison"] = "comparison" + type: Literal["comparison_tensor"] = "comparison_tensor" traced_plan: Optional[TracedAlignerPlan] = None replicated_checks: list[ReplicatedCheckResult] = Field(default_factory=list) raw_bundle_info: Optional[Pair[BundleSideInfo]] = None @@ -224,8 +234,8 @@ class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord): return _format_tensor_comparison_rich_body(self, verbosity=verbosity) -class NonTensorComparisonRecord(_BaseComparisonRecord): - type: Literal["non_tensor"] = "non_tensor" +class ComparisonNonTensorRecord(_BaseComparisonRecord): + type: Literal["comparison_non_tensor"] = "comparison_non_tensor" name: str baseline_value: str target_value: str @@ -245,15 +255,6 @@ class NonTensorComparisonRecord(_BaseComparisonRecord): def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType: return _format_non_tensor_rich_body(self, verbosity=verbosity) - return _format_non_tensor_body(self) - - def _format_rich_body(self) -> RenderableType: - from sglang.srt.debug_utils.comparator.output_formatter import ( - _format_non_tensor_rich_body, - ) - - return _format_non_tensor_rich_body(self) - class SummaryRecord(_OutputRecord): type: Literal["summary"] = "summary" @@ -261,13 +262,15 @@ class SummaryRecord(_OutputRecord): passed: int failed: int skipped: int + errored: int = 0 @model_validator(mode="after") def _validate_totals(self) -> "SummaryRecord": - expected: int = self.passed + self.failed + self.skipped + expected: int = self.passed + self.failed + self.skipped + self.errored if self.total != expected: raise ValueError( - f"total={self.total} != passed({self.passed}) + failed({self.failed}) + skipped({self.skipped}) = {expected}" + f"total={self.total} != passed({self.passed}) + failed({self.failed}) " + f"+ skipped({self.skipped}) + errored({self.errored}) = {expected}" ) return self @@ -277,7 +280,6 @@ class SummaryRecord(_OutputRecord): def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType: return _format_summary_rich_body(self, verbosity=verbosity) - return _format_summary_body(self) class LogRecord(_OutputRecord): type: Literal["log"] = "log" @@ -291,9 +293,10 @@ AnyRecord = Annotated[ ConfigRecord, RankInfoRecord, InputIdsRecord, - SkipComparisonRecord, - TensorComparisonRecord, - NonTensorComparisonRecord, + ComparisonSkipRecord, + ComparisonErrorRecord, + ComparisonTensorRecord, + ComparisonNonTensorRecord, SummaryRecord, LogRecord, ], diff --git a/python/sglang/srt/debug_utils/comparator/per_token_visualizer.py b/python/sglang/srt/debug_utils/comparator/per_token_visualizer.py index 705fd0490..9f1a30c2c 100644 --- a/python/sglang/srt/debug_utils/comparator/per_token_visualizer.py +++ b/python/sglang/srt/debug_utils/comparator/per_token_visualizer.py @@ -9,12 +9,12 @@ from __future__ import annotations from pathlib import Path from typing import Optional -from sglang.srt.debug_utils.comparator.output_types import TensorComparisonRecord +from sglang.srt.debug_utils.comparator.output_types import ComparisonTensorRecord def generate_per_token_heatmap( *, - records: list[TensorComparisonRecord], + records: list[ComparisonTensorRecord], output_path: Path, ) -> Optional[Path]: """Generate a per-token relative difference heatmap PNG. @@ -31,7 +31,7 @@ def generate_per_token_heatmap( def _collect_per_token_data( *, - records: list[TensorComparisonRecord], + records: list[ComparisonTensorRecord], ) -> list[tuple[str, list[float]]]: rows: list[tuple[str, list[float]]] = [] for record in records: diff --git a/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py b/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py index 47dd09efb..6021d065d 100644 --- a/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py +++ b/python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py @@ -20,9 +20,9 @@ if TYPE_CHECKING: from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan from sglang.srt.debug_utils.comparator.output_types import ( BundleSideInfo, + ComparisonTensorRecord, ReplicatedCheckResult, ShapeSnapshot, - TensorComparisonRecord, ) from sglang.srt.debug_utils.comparator.utils import Pair @@ -207,7 +207,7 @@ def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]: def format_comparison_rich( - record: TensorComparisonRecord, + record: ComparisonTensorRecord, verbosity: Verbosity = "normal", ) -> str: if verbosity == "minimal": @@ -219,7 +219,7 @@ def format_comparison_rich( ) -def _format_comparison_minimal(record: TensorComparisonRecord) -> str: +def _format_comparison_minimal(record: ComparisonTensorRecord) -> str: passed, color, marker = _category_marker(record.category) name_part: str = f"[bold {color}]{escape(record.name):30s}[/]" @@ -233,7 +233,7 @@ def _format_comparison_minimal(record: TensorComparisonRecord) -> str: def _format_comparison_normal_or_verbose( *, - record: TensorComparisonRecord, + record: ComparisonTensorRecord, verbose: bool, ) -> str: passed, color, marker = _category_marker(record.category) diff --git a/python/sglang/srt/debug_utils/comparator/utils.py b/python/sglang/srt/debug_utils/comparator/utils.py index c2823ce85..2e8cedc01 100644 --- a/python/sglang/srt/debug_utils/comparator/utils.py +++ b/python/sglang/srt/debug_utils/comparator/utils.py @@ -141,10 +141,14 @@ def compute_exit_code( skipped_names: list[str], allow_failed_pattern: Optional[str], failed_names: list[str], + errored_names: Optional[list[str]] = None, ) -> int: if summary.passed == 0: return 1 + if errored_names: + return 1 + if not _is_all_match_pattern(pattern=allow_failed_pattern, strings=failed_names): return 1 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 24329bf1c..1fb37ecbb 100644 --- a/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py +++ b/test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py @@ -34,9 +34,9 @@ from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, TokenLayou from sglang.srt.debug_utils.comparator.output_types import ( BundleFileInfo, BundleSideInfo, + ComparisonTensorRecord, ReplicatedCheckResult, ShapeSnapshot, - TensorComparisonRecord, ) from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import ( _format_abs_diff_percentiles_rich, @@ -273,9 +273,9 @@ def _make_comparison_record( replicated_checks: list[ReplicatedCheckResult] | None = None, raw_bundle_info: Pair[BundleSideInfo] | None = None, traced_plan: TracedAlignerPlan | None = None, -) -> TensorComparisonRecord: +) -> ComparisonTensorRecord: s: list[int] = shape if shape is not None else [4, 8] - return TensorComparisonRecord( + return ComparisonTensorRecord( name=name, baseline=_make_tensor_info(shape=s, dtype=dtype, sample=sample), target=_make_tensor_info(shape=s, dtype=dtype, sample=sample), @@ -417,7 +417,7 @@ class TestFormatComparisonRichMinimal: """format_comparison_rich() with verbosity='minimal'.""" def test_passed(self) -> None: - record: TensorComparisonRecord = _make_comparison_record( + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff(rel_diff=1e-4, passed=True), ) result: str = format_comparison_rich(record, verbosity="minimal") @@ -428,7 +428,7 @@ class TestFormatComparisonRichMinimal: ) def test_failed(self) -> None: - record: TensorComparisonRecord = _make_comparison_record( + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff(rel_diff=0.5, passed=False), ) result: str = format_comparison_rich(record, verbosity="minimal") @@ -439,7 +439,7 @@ class TestFormatComparisonRichMinimal: ) def test_shape_mismatch(self) -> None: - record: TensorComparisonRecord = _make_comparison_record( + record: ComparisonTensorRecord = _make_comparison_record( shape_mismatch=True, ) result: str = format_comparison_rich(record, verbosity="minimal") @@ -450,7 +450,7 @@ class TestFormatComparisonRichMinimal: ) def test_no_diff(self) -> None: - record: TensorComparisonRecord = _make_comparison_record() + record: ComparisonTensorRecord = _make_comparison_record() result: str = format_comparison_rich(record, verbosity="minimal") assert result == ("[red]❌[/] [bold red]hidden_states [/]") @@ -460,7 +460,7 @@ class TestFormatComparisonRichNormal: """format_comparison_rich() with verbosity='normal'.""" def test_passed(self) -> None: - record: TensorComparisonRecord = _make_comparison_record( + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff(rel_diff=1e-4, passed=True), ) result: str = format_comparison_rich(record, verbosity="normal") @@ -477,7 +477,7 @@ class TestFormatComparisonRichNormal: ) def test_failed(self) -> None: - record: TensorComparisonRecord = _make_comparison_record( + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff( rel_diff=0.5, max_abs_diff=1.0, mean_abs_diff=0.3, passed=False ), @@ -499,7 +499,7 @@ class TestFormatComparisonRichNormal: ) def test_shape_mismatch(self) -> None: - record: TensorComparisonRecord = _make_comparison_record( + record: ComparisonTensorRecord = _make_comparison_record( shape_mismatch=True, ) result: str = format_comparison_rich(record, verbosity="normal") @@ -516,7 +516,7 @@ class TestFormatComparisonRichNormal: ) def test_with_downcast(self) -> None: - record: TensorComparisonRecord = _make_comparison_record( + record: ComparisonTensorRecord = _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", @@ -543,7 +543,7 @@ class TestFormatComparisonRichNormal: 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( + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff(passed=True), raw_bundle_info=bundle_info, ) @@ -565,7 +565,7 @@ class TestFormatComparisonRichNormal: def test_with_plan(self) -> None: plan: AlignerPlan = _make_simple_aligner_plan(with_unsharder=True) - record: TensorComparisonRecord = _make_comparison_record( + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff(passed=True), traced_plan=_make_traced_plan(plan), ) @@ -590,7 +590,7 @@ class TestFormatComparisonRichVerbose: """format_comparison_rich() with verbosity='verbose'.""" def test_passed_full_detail(self) -> None: - record: TensorComparisonRecord = _make_comparison_record( + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff(rel_diff=1e-4, passed=True), sample="tensor([0.1, 0.2, ...])", ) @@ -624,7 +624,7 @@ class TestFormatComparisonRichVerbose: x=_make_bundle_side_info(num_files=2, with_parallel_info=True), y=_make_bundle_side_info(num_files=2, with_parallel_info=True), ) - record: TensorComparisonRecord = _make_comparison_record( + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff(passed=True), raw_bundle_info=bundle_info, ) @@ -659,7 +659,7 @@ class TestFormatComparisonRichVerbose: def test_with_plan_and_traces(self) -> None: plan: AlignerPlan = _make_simple_aligner_plan(with_unsharder=True) - record: TensorComparisonRecord = _make_comparison_record( + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff(passed=True), traced_plan=_make_traced_plan( plan, diff --git a/test/registered/debug_utils/comparator/tensor_comparator/test_types.py b/test/registered/debug_utils/comparator/tensor_comparator/test_types.py index 560cef972..17ddcedb9 100644 --- a/test/registered/debug_utils/comparator/tensor_comparator/test_types.py +++ b/test/registered/debug_utils/comparator/tensor_comparator/test_types.py @@ -4,14 +4,14 @@ import sys import pytest from sglang.srt.debug_utils.comparator.output_types import ( + ComparisonSkipRecord, + ComparisonTensorRecord, ConfigRecord, ErrorLog, InfoLog, LogRecord, ReplicatedCheckResult, - SkipComparisonRecord, SummaryRecord, - TensorComparisonRecord, parse_record_json, ) from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( @@ -84,7 +84,7 @@ class TestStrictBase: class TestRecordTypes: def test_comparison_record_inherits_tensor_fields(self): - record = TensorComparisonRecord( + record = ComparisonTensorRecord( name="hidden_states", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -93,7 +93,7 @@ class TestRecordTypes: diff=_make_diff(), ) parsed = json.loads(record.model_dump_json()) - assert parsed["type"] == "comparison" + assert parsed["type"] == "comparison_tensor" assert parsed["name"] == "hidden_states" assert "baseline" in parsed assert "diff" in parsed @@ -109,8 +109,8 @@ class TestRecordTypes: "end_step": 100, } ), - SkipComparisonRecord(name="attn", reason="no_baseline"), - TensorComparisonRecord( + ComparisonSkipRecord(name="attn", reason="no_baseline"), + ComparisonTensorRecord( name="mlp", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -149,8 +149,8 @@ def _make_replicated_check(**overrides) -> ReplicatedCheckResult: class TestWarnings: def test_comparison_record_failed_when_diff_passed_but_errors(self): - """TensorComparisonRecord with diff.passed=True but errors → category=='failed'.""" - record = TensorComparisonRecord( + """ComparisonTensorRecord with diff.passed=True but errors → category=='failed'.""" + record = ComparisonTensorRecord( name="hidden", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -162,8 +162,8 @@ class TestWarnings: assert record.category == "failed" def test_skip_record_failed_when_errors(self): - """SkipComparisonRecord with errors → category=='failed' instead of 'skipped'.""" - record = SkipComparisonRecord( + """ComparisonSkipRecord with errors → category=='failed' instead of 'skipped'.""" + record = ComparisonSkipRecord( name="x", reason="no_baseline", errors=[ErrorLog(category="test", message="some warning")], @@ -171,8 +171,8 @@ class TestWarnings: assert record.category == "failed" def test_replicated_checks_all_passed(self): - """TensorComparisonRecord with all replicated_checks passed → category=='passed'.""" - record = TensorComparisonRecord( + """ComparisonTensorRecord with all replicated_checks passed → category=='passed'.""" + record = ComparisonTensorRecord( name="hidden", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -184,8 +184,8 @@ class TestWarnings: assert record.category == "passed" def test_replicated_checks_failed_means_record_failed(self): - """TensorComparisonRecord with any replicated_check.passed=False → category=='failed'.""" - record = TensorComparisonRecord( + """ComparisonTensorRecord with any replicated_check.passed=False → category=='failed'.""" + record = ComparisonTensorRecord( name="hidden", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -197,7 +197,7 @@ class TestWarnings: assert record.category == "failed" def test_replicated_check_json_round_trip(self): - """ReplicatedCheckResult survives JSON round-trip via TensorComparisonRecord.""" + """ReplicatedCheckResult survives JSON round-trip via ComparisonTensorRecord.""" check = _make_replicated_check( axis="cp", group_index=2, @@ -205,7 +205,7 @@ class TestWarnings: baseline_index=0, passed=False, ) - record = TensorComparisonRecord( + record = ComparisonTensorRecord( name="mlp", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -216,7 +216,7 @@ class TestWarnings: ) restored = parse_record_json(record.model_dump_json()) - assert isinstance(restored, TensorComparisonRecord) + assert isinstance(restored, ComparisonTensorRecord) assert len(restored.replicated_checks) == 1 restored_check: ReplicatedCheckResult = restored.replicated_checks[0] diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 41e5d5ae8..2ab49899f 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -7,6 +7,7 @@ from pathlib import Path import pytest import torch +import sglang.srt.debug_utils.comparator.entrypoint as _entrypoint_module import sglang.srt.debug_utils.dumper as _dumper_module from sglang.srt.debug_utils.comparator.entrypoint import ( parse_args, @@ -14,14 +15,15 @@ from sglang.srt.debug_utils.comparator.entrypoint import ( ) from sglang.srt.debug_utils.comparator.output_types import ( AnyRecord, + ComparisonErrorRecord, + ComparisonNonTensorRecord, + ComparisonSkipRecord, + ComparisonTensorRecord, ConfigRecord, InfoLog, LogRecord, - NonTensorComparisonRecord, ReplicatedCheckResult, - SkipComparisonRecord, SummaryRecord, - TensorComparisonRecord, _OutputRecord, parse_record_json, ) @@ -39,7 +41,7 @@ class TestEntrypointGroupingRaw: """Test `--grouping-skip-keys` empty (raw) scenarios""" def test_run_basic(self, tmp_path, capsys): - """Two matching tensors produce ConfigRecord, 2 TensorComparisonRecords, and SummaryRecord.""" + """Two matching tensors produce ConfigRecord, 2 ComparisonTensorRecords, and SummaryRecord.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"]) argv = _make_argv(baseline_path, target_path, preset="raw") @@ -54,7 +56,7 @@ class TestEntrypointGroupingRaw: assert summary.skipped == 0 def test_filter(self, tmp_path, capsys): - """--filter selects only the matching tensor, producing 1 TensorComparisonRecord.""" + """--filter selects only the matching tensor, producing 1 ComparisonTensorRecord.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"]) argv = _make_argv(baseline_path, target_path, filter="tensor_a", preset="raw") @@ -62,7 +64,7 @@ class TestEntrypointGroupingRaw: assert len(_get_comparisons(records)) == 1 def test_no_baseline_skip(self, tmp_path, capsys): - """Target tensor missing from baseline emits a SkipComparisonRecord with reason baseline_load_failed.""" + """Target tensor missing from baseline emits a ComparisonSkipRecord with reason baseline_load_failed.""" baseline_path, target_path = _create_dumps( tmp_path, tensor_names=["tensor_a", "tensor_extra"], @@ -71,7 +73,7 @@ class TestEntrypointGroupingRaw: argv = _make_argv(baseline_path, target_path, preset="raw") records, _ = _run_and_parse(argv, capsys) - skips = [r for r in records if isinstance(r, SkipComparisonRecord)] + skips = [r for r in records if isinstance(r, ComparisonSkipRecord)] assert len(skips) == 1 assert skips[0].reason == "baseline_load_failed" @@ -100,7 +102,7 @@ class TestEntrypointGroupingRaw: assert all(isinstance(r, _OutputRecord) for r in records) def test_comparison_failed(self, tmp_path, capsys): - """Completely different tensors produce a failed TensorComparisonRecord.""" + """Completely different tensors produce a failed ComparisonTensorRecord.""" torch.manual_seed(42) baseline_path = _create_rank_dump( tmp_path / "baseline", rank=0, name="tensor_a", tensor=torch.randn(10, 10) @@ -251,7 +253,7 @@ class TestEntrypointGroupingRaw: assert summary.total == 0 def test_raw_multi_rank(self, tmp_path, capsys): - """Two ranks in raw grouping produce two TensorComparisonRecords (one per rank).""" + """Two ranks in raw grouping produce two ComparisonTensorRecords (one per rank).""" torch.manual_seed(42) tensor = torch.randn(4, 4) @@ -480,7 +482,7 @@ class TestEntrypointGroupingLogical: ], ) def test_ambiguous_no_dims_skip(self, tmp_path, capsys, bad_side, expected_reason): - """Multi-rank without dims on one side produces a SkipComparisonRecord with the appropriate reason.""" + """Multi-rank without dims on one side produces a ComparisonSkipRecord with the appropriate reason.""" torch.manual_seed(42) tensor = torch.randn(4, 8) @@ -500,7 +502,7 @@ class TestEntrypointGroupingLogical: ) records, _ = _run_and_parse(argv, capsys) - skips = [r for r in records if isinstance(r, SkipComparisonRecord)] + skips = [r for r in records if isinstance(r, ComparisonSkipRecord)] assert len(skips) == 1 assert skips[0].reason == expected_reason @@ -1092,7 +1094,7 @@ class TestEntrypointPerStepMode: """Test per-step comparison mode (sglang_dev preset behavior).""" def test_multi_step_per_step_comparison(self, tmp_path, capsys): - """Multiple steps produce one TensorComparisonRecord per step with step field set.""" + """Multiple steps produce one ComparisonTensorRecord per step with step field set.""" torch.manual_seed(42) baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"], num_steps=3) argv = _make_argv(baseline_path, target_path, diff_threshold=0.1) @@ -1149,7 +1151,7 @@ class TestEntrypointPerStepMode: assert all(c.baseline.shape == [4, 8] for c in comparisons) def test_single_step_has_step_field(self, tmp_path, capsys): - """Single step produces TensorComparisonRecord with location.step=0.""" + """Single step produces ComparisonTensorRecord with location.step=0.""" baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"], num_steps=1) argv = _make_argv(baseline_path, target_path) @@ -1458,7 +1460,7 @@ class TestEntrypointConcatMode: assert len(comparisons) == 3 def test_concat_aligner_plan_fields(self, tmp_path, capsys): - """TensorComparisonRecord.traced_plan reports mode='concat' with plan=None.""" + """ComparisonTensorRecord.traced_plan reports mode='concat' with plan=None.""" torch.manual_seed(42) records = self._run_concat( @@ -1599,8 +1601,8 @@ class TestEntrypointConcatMode: ) records, _ = _run_and_parse(argv, capsys) - comparisons: list[TensorComparisonRecord] = _get_comparisons(records) - hidden_comparisons: list[TensorComparisonRecord] = [ + comparisons: list[ComparisonTensorRecord] = _get_comparisons(records) + hidden_comparisons: list[ComparisonTensorRecord] = [ c for c in comparisons if c.name == "hidden_states" ] assert len(hidden_comparisons) >= 1 @@ -2158,7 +2160,7 @@ class TestEntrypointNonTensorValues: """Test non-tensor value comparison through the full entrypoint pipeline.""" def test_non_tensor_float_same_value(self, tmp_path: Path, capsys) -> None: - """Two sides dump the same float → NonTensorComparisonRecord with values_equal=True, category=passed.""" + """Two sides dump the same float → ComparisonNonTensorRecord with values_equal=True, category=passed.""" baseline_path, target_path = _create_non_tensor_dumps( tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.125 ) @@ -2177,7 +2179,7 @@ class TestEntrypointNonTensorValues: assert summary.failed == 0 def test_non_tensor_float_different_value(self, tmp_path: Path, capsys) -> None: - """Two sides dump different floats → NonTensorComparisonRecord with values_equal=False, category=failed.""" + """Two sides dump different floats → ComparisonNonTensorRecord with values_equal=False, category=failed.""" baseline_path, target_path = _create_non_tensor_dumps( tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.25 ) @@ -2262,7 +2264,7 @@ class TestEntrypointNonTensorValues: assert non_tensors[0].target_type == "dict" def test_non_tensor_none_value(self, tmp_path: Path, capsys) -> None: - """Dumping None is displayed as NonTensorComparisonRecord, not skipped as load failure.""" + """Dumping None is displayed as ComparisonNonTensorRecord, not skipped as load failure.""" baseline_path, target_path = _create_non_tensor_dumps( tmp_path, name="optional_param", baseline_value=None, target_value=None ) @@ -2278,7 +2280,7 @@ class TestEntrypointNonTensorValues: assert non_tensors[0].category == "passed" def test_non_tensor_json_roundtrip(self, tmp_path: Path, capsys) -> None: - """NonTensorComparisonRecord JSON output can be parsed back correctly.""" + """ComparisonNonTensorRecord JSON output can be parsed back correctly.""" baseline_path, target_path = _create_non_tensor_dumps( tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.125 ) @@ -2290,7 +2292,7 @@ class TestEntrypointNonTensorValues: json_str: str = non_tensors[0].model_dump_json() roundtripped = parse_record_json(json_str) - assert isinstance(roundtripped, NonTensorComparisonRecord) + assert isinstance(roundtripped, ComparisonNonTensorRecord) assert roundtripped.name == "sm_scale" assert roundtripped.values_equal is True @@ -2344,17 +2346,17 @@ class TestEntrypointVisualize: # --------------------------- Assertion helpers ------------------- -def _get_comparisons(records: list[AnyRecord]) -> list[TensorComparisonRecord]: - return [r for r in records if isinstance(r, TensorComparisonRecord)] +def _get_comparisons(records: list[AnyRecord]) -> list[ComparisonTensorRecord]: + return [r for r in records if isinstance(r, ComparisonTensorRecord)] -def _get_non_tensors(records: list[AnyRecord]) -> list[NonTensorComparisonRecord]: - return [r for r in records if isinstance(r, NonTensorComparisonRecord)] +def _get_non_tensors(records: list[AnyRecord]) -> list[ComparisonNonTensorRecord]: + return [r for r in records if isinstance(r, ComparisonNonTensorRecord)] def _assert_single_comparison_passed( records: list[AnyRecord], -) -> TensorComparisonRecord: +) -> ComparisonTensorRecord: comparisons = _get_comparisons(records) assert len(comparisons) == 1 assert comparisons[0].diff is not None @@ -3235,8 +3237,8 @@ class TestEntrypointThdCpZigzag: ) records, _ = _run_and_parse(argv, capsys) - comparisons: list[TensorComparisonRecord] = _get_comparisons(records) - hidden_comparisons: list[TensorComparisonRecord] = [ + comparisons: list[ComparisonTensorRecord] = _get_comparisons(records) + hidden_comparisons: list[ComparisonTensorRecord] = [ c for c in comparisons if c.name == "hidden_states" ] assert len(hidden_comparisons) >= 1 @@ -3287,8 +3289,8 @@ class TestEntrypointThdCpZigzag: records, _ = _run_and_parse(argv, capsys) # hidden_states should pass comparison (after unshard + reorder) - comparisons: list[TensorComparisonRecord] = _get_comparisons(records) - hidden_comparisons: list[TensorComparisonRecord] = [ + comparisons: list[ComparisonTensorRecord] = _get_comparisons(records) + hidden_comparisons: list[ComparisonTensorRecord] = [ c for c in comparisons if c.name == "hidden_states" ] assert len(hidden_comparisons) >= 1 @@ -3355,7 +3357,7 @@ class TestEntrypointDpFilter: ) records, _ = _run_and_parse(argv, capsys) - comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) + comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records) assert comparison.name == "hidden" def test_dp2_megatron_both_sides(self, tmp_path: Path, capsys) -> None: @@ -3410,7 +3412,7 @@ class TestEntrypointDpFilter: ) records, _ = _run_and_parse(argv, capsys) - comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) + comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records) assert comparison.name == "hidden" def test_dp2_tp2_sglang(self, tmp_path: Path, capsys) -> None: @@ -3458,7 +3460,7 @@ class TestEntrypointDpFilter: ) records, _ = _run_and_parse(argv, capsys) - comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) + comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records) assert comparison.name == "hidden" def test_dp2_both_nonempty_raises(self, tmp_path: Path, capsys) -> None: @@ -3496,10 +3498,12 @@ class TestEntrypointDpFilter: diff_threshold=1e-3, ) - with pytest.raises( - AssertionError, match="Expected exactly 1 non-empty dp_rank" - ): - _run_and_parse(argv, capsys) + records, exit_code = _run_and_parse(argv, capsys) + 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 exit_code == 1 class TestEntrypointDpGroupAlias: @@ -3542,7 +3546,7 @@ class TestEntrypointDpGroupAlias: ) records, _ = _run_and_parse(argv, capsys) - comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) + comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records) assert comparison.name == "hidden" def test_dp_alias_via_override_dims(self, tmp_path: Path, capsys) -> None: @@ -3599,7 +3603,7 @@ class TestEntrypointDpGroupAlias: ) records, _ = _run_and_parse(argv, capsys) - comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) + comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records) assert comparison.name == "hidden" def test_dp_alias_with_real_alias_group_filters( @@ -3640,7 +3644,7 @@ class TestEntrypointDpGroupAlias: ) records, _ = _run_and_parse(argv, capsys) - comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) + comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records) assert comparison.name == "hidden" @@ -3679,7 +3683,7 @@ class TestEntrypointMetaOverride: records: list[AnyRecord], *, expected_count: int = 1 ) -> None: """Assert that exactly expected_count comparisons exist and all passed.""" - comparisons: list[TensorComparisonRecord] = _get_comparisons(records) + comparisons: list[ComparisonTensorRecord] = _get_comparisons(records) assert len(comparisons) == expected_count assert all(c.diff is not None and c.diff.passed for c in comparisons) @@ -3975,14 +3979,14 @@ class TestEntrypointMetaOverride: ) records, _ = _run_and_parse(argv, capsys) - non_tensors: list[NonTensorComparisonRecord] = [ - r for r in records if isinstance(r, NonTensorComparisonRecord) + non_tensors: list[ComparisonNonTensorRecord] = [ + r for r in records if isinstance(r, ComparisonNonTensorRecord) ] assert len(non_tensors) == 1 assert non_tensors[0].name == "sm_scale" assert non_tensors[0].values_equal - comparisons: list[TensorComparisonRecord] = _get_comparisons(records) + comparisons: list[ComparisonTensorRecord] = _get_comparisons(records) assert len(comparisons) == 1 assert comparisons[0].name == "hidden" @@ -4356,12 +4360,9 @@ class TestEntrypointDpAttentionMissingAlias: assert exit_code == 1 - comparisons: list[TensorComparisonRecord] = _get_comparisons(records) - assert len(comparisons) == 1 - comparison: TensorComparisonRecord = comparisons[0] - assert comparison.shape_mismatch is True - assert comparison.diff is None - assert comparison.category == "failed" + errors = [r for r in records if isinstance(r, ComparisonErrorRecord)] + assert len(errors) == 1 + assert errors[0].category == "errored" class TestEntrypointAutoDescend: @@ -4460,5 +4461,77 @@ class TestEntrypointAutoDescend: run(parse_args(argv)) +class TestErrorResilience: + """Bundle comparison exception → continue with remaining bundles.""" + + def test_one_bundle_errors_others_continue(self, tmp_path, capsys, monkeypatch): + """One bundle raises exception → other bundles still compared, summary correct.""" + baseline_path, target_path = _create_dumps( + tmp_path, ["tensor_a", "tensor_b", "tensor_c"] + ) + argv = _make_argv(baseline_path, target_path, preset="raw") + + original = _entrypoint_module.compare_bundle_pair + + def _patched(**kwargs): + if kwargs["name"] == "tensor_b": + raise RuntimeError("intentional test error") + return original(**kwargs) + + monkeypatch.setattr(_entrypoint_module, "compare_bundle_pair", _patched) + + records, exit_code = _run_and_parse(argv, capsys) + + comparisons = _get_comparisons(records) + assert len(comparisons) == 2 + + errors = [r for r in records if isinstance(r, ComparisonErrorRecord)] + 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 + + summary = records[-1] + assert isinstance(summary, SummaryRecord) + assert summary.errored == 1 + assert summary.passed == 2 + assert summary.total == 3 + + assert exit_code == 1 + + def test_all_bundles_error_exits_one(self, tmp_path, capsys, monkeypatch): + """All bundles error → exit 1, summary all errored.""" + baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"]) + argv = _make_argv(baseline_path, target_path, preset="raw") + + def _always_raise(**kwargs): + raise ValueError("always fail") + + monkeypatch.setattr(_entrypoint_module, "compare_bundle_pair", _always_raise) + + records, exit_code = _run_and_parse(argv, capsys) + + summary = records[-1] + assert isinstance(summary, SummaryRecord) + assert summary.errored == 1 + assert summary.passed == 0 + assert exit_code == 1 + + def test_error_record_json_roundtrip_in_output(self, tmp_path, capsys, monkeypatch): + """ComparisonErrorRecord correctly serializes and deserializes in output.""" + baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"]) + argv = _make_argv(baseline_path, target_path, preset="raw") + + def _raise(**kwargs): + raise TypeError("bad type") + + 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 errors[0].exception_type == "TypeError" + + if __name__ == "__main__": sys.exit(pytest.main([__file__])) diff --git a/test/registered/debug_utils/comparator/test_manually_verify.py b/test/registered/debug_utils/comparator/test_manually_verify.py index fd23617bc..439f1cfa3 100644 --- a/test/registered/debug_utils/comparator/test_manually_verify.py +++ b/test/registered/debug_utils/comparator/test_manually_verify.py @@ -208,7 +208,7 @@ class TestPerTokenHeatmapManualVerify: rows for different tensor names. Colorbar shows log10 scale. """ from sglang.srt.debug_utils.comparator.output_types import ( - TensorComparisonRecord, + ComparisonTensorRecord, ) from sglang.srt.debug_utils.comparator.per_token_visualizer import ( generate_per_token_heatmap, @@ -222,7 +222,7 @@ class TestPerTokenHeatmapManualVerify: hidden_dim: int = 128 num_tensors: int = 5 - records: list[TensorComparisonRecord] = [] + records: list[ComparisonTensorRecord] = [] for i in range(num_tensors): baseline: torch.Tensor = torch.randn(seq_len, hidden_dim) noise_scale: torch.Tensor = torch.linspace( @@ -237,7 +237,7 @@ class TestPerTokenHeatmapManualVerify: diff_threshold=1e-3, seq_dim=0, ) - records.append(TensorComparisonRecord(**info.model_dump())) + records.append(ComparisonTensorRecord(**info.model_dump())) output_path: Path = tmp_path / "per_token_increasing_diff.png" result = generate_per_token_heatmap(records=records, output_path=output_path) @@ -253,7 +253,7 @@ class TestPerTokenHeatmapManualVerify: rest is dark/cold. """ from sglang.srt.debug_utils.comparator.output_types import ( - TensorComparisonRecord, + ComparisonTensorRecord, ) from sglang.srt.debug_utils.comparator.per_token_visualizer import ( generate_per_token_heatmap, @@ -268,7 +268,7 @@ class TestPerTokenHeatmapManualVerify: spike_pos: int = 32 num_tensors: int = 4 - records: list[TensorComparisonRecord] = [] + records: list[ComparisonTensorRecord] = [] for i in range(num_tensors): baseline: torch.Tensor = torch.randn(seq_len, hidden_dim) target: torch.Tensor = baseline.clone() @@ -281,7 +281,7 @@ class TestPerTokenHeatmapManualVerify: diff_threshold=1e-3, seq_dim=0, ) - records.append(TensorComparisonRecord(**info.model_dump())) + records.append(ComparisonTensorRecord(**info.model_dump())) output_path: Path = tmp_path / "per_token_single_spike.png" result = generate_per_token_heatmap(records=records, output_path=output_path) diff --git a/test/registered/debug_utils/comparator/test_model_validation.py b/test/registered/debug_utils/comparator/test_model_validation.py index 82306af28..096559eaa 100644 --- a/test/registered/debug_utils/comparator/test_model_validation.py +++ b/test/registered/debug_utils/comparator/test_model_validation.py @@ -4,6 +4,12 @@ import sys import pytest from pydantic import ValidationError +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, @@ -22,11 +28,12 @@ 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 ( + ComparisonErrorRecord, + ComparisonNonTensorRecord, + ComparisonSkipRecord, + ComparisonTensorRecord, ErrorLog, - NonTensorComparisonRecord, - SkipComparisonRecord, SummaryRecord, - TensorComparisonRecord, parse_record_json, ) from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( @@ -156,6 +163,14 @@ class TestSummaryRecord: with pytest.raises(ValidationError, match="total=10"): SummaryRecord(total=10, passed=5, failed=2, skipped=1) + def test_valid_with_errored(self): + record = SummaryRecord(total=10, passed=6, failed=2, skipped=1, errored=1) + assert record.errored == 1 + + def test_total_mismatch_with_errored(self): + with pytest.raises(ValidationError, match="total=10"): + SummaryRecord(total=10, passed=6, failed=2, skipped=1, errored=0) + class TestAxisInfo: def test_valid(self): @@ -208,9 +223,9 @@ def _make_comparison_record( *, diff: DiffInfo | None, errors: list | None = None, -) -> TensorComparisonRecord: +) -> ComparisonTensorRecord: ti: TensorInfo = _make_tensor_info() - return TensorComparisonRecord( + return ComparisonTensorRecord( name="t", baseline=ti, target=ti, @@ -223,7 +238,7 @@ def _make_comparison_record( class TestOutputRecordCategories: def test_skip_record_with_errors_is_failed(self) -> None: - record = SkipComparisonRecord( + record = ComparisonSkipRecord( name="t", reason="test", errors=[ErrorLog(category="c", message="m")], @@ -231,28 +246,28 @@ class TestOutputRecordCategories: assert record.category == "failed" def test_skip_record_no_warnings_is_skipped(self) -> None: - record = SkipComparisonRecord(name="t", reason="test") + record = ComparisonSkipRecord(name="t", reason="test") assert record.category == "skipped" def test_comparison_record_diff_none_is_failed(self) -> None: - record: TensorComparisonRecord = _make_comparison_record(diff=None) + record: ComparisonTensorRecord = _make_comparison_record(diff=None) assert record.category == "failed" def test_comparison_record_passed_with_errors_is_failed(self) -> None: - record: TensorComparisonRecord = _make_comparison_record( + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff_info(passed=True), errors=[ErrorLog(category="c", message="m")], ) assert record.category == "failed" def test_comparison_record_passed_no_warnings_is_passed(self) -> None: - record: TensorComparisonRecord = _make_comparison_record( + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) assert record.category == "passed" def test_non_tensor_record_equal_is_passed(self) -> None: - record = NonTensorComparisonRecord( + record = ComparisonNonTensorRecord( name="sm_scale", baseline_value="0.125", target_value="0.125", @@ -263,7 +278,7 @@ class TestOutputRecordCategories: assert record.category == "passed" def test_non_tensor_record_different_is_failed(self) -> None: - record = NonTensorComparisonRecord( + record = ComparisonNonTensorRecord( name="sm_scale", baseline_value="0.125", target_value="0.25", @@ -274,7 +289,7 @@ class TestOutputRecordCategories: assert record.category == "failed" def test_non_tensor_record_with_errors_is_failed(self) -> None: - record = NonTensorComparisonRecord( + record = ComparisonNonTensorRecord( name="sm_scale", baseline_value="0.125", target_value="0.125", @@ -286,7 +301,7 @@ class TestOutputRecordCategories: assert record.category == "failed" def test_non_tensor_record_json_roundtrip(self) -> None: - record = NonTensorComparisonRecord( + record = ComparisonNonTensorRecord( name="sm_scale", baseline_value="0.125", target_value="0.25", @@ -296,14 +311,14 @@ class TestOutputRecordCategories: ) json_str: str = record.model_dump_json() roundtripped = parse_record_json(json_str) - assert isinstance(roundtripped, NonTensorComparisonRecord) + assert isinstance(roundtripped, ComparisonNonTensorRecord) assert roundtripped.name == "sm_scale" assert roundtripped.values_equal is False assert roundtripped.baseline_value == "0.125" assert roundtripped.target_value == "0.25" def test_non_tensor_record_text_format_equal(self) -> None: - record = NonTensorComparisonRecord( + record = ComparisonNonTensorRecord( name="sm_scale", baseline_value="0.125", target_value="0.125", @@ -316,7 +331,7 @@ class TestOutputRecordCategories: assert "[equal]" in text def test_non_tensor_record_text_format_different(self) -> None: - record = NonTensorComparisonRecord( + record = ComparisonNonTensorRecord( name="sm_scale", baseline_value="0.125", target_value="0.25", @@ -328,14 +343,38 @@ class TestOutputRecordCategories: assert "baseline" in text assert "target" in text + def test_error_record_category_is_errored(self) -> None: + record = ComparisonErrorRecord( + name="t", exception_type="ValueError", traceback_str="..." + ) + assert record.category == "errored" -def _make_aligner_plan() -> AlignerPlan: + def test_error_record_json_roundtrip(self) -> None: + record = ComparisonErrorRecord( + name="t", exception_type="ValueError", 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" + + def test_error_record_text_format(self) -> None: + record = ComparisonErrorRecord( + name="t", exception_type="RuntimeError", traceback_str="Traceback..." + ) + text: str = record.to_text() + assert "RuntimeError" in text + assert "Traceback" in text + + +def _make_traced_aligner_plan() -> TracedAlignerPlan: unsharder = UnsharderPlan( axis=ParallelAxis.TP, params=ConcatParams(dim_name="h"), groups=[[0, 1]], ) - return AlignerPlan( + plan = AlignerPlan( per_step_plans=Pair( x=[ AlignerPerStepPlan( @@ -349,54 +388,67 @@ def _make_aligner_plan() -> AlignerPlan: ], ), ) + traced_sub = TracedSubPlan(plan=unsharder, snapshot=None) + traced_step = TracedStepPlan( + step=0, input_object_indices=[0, 1], sub_plans=[traced_sub] + ) + return TracedAlignerPlan( + plan=plan, + per_side=Pair( + x=TracedSidePlan(step_plans=[traced_step]), + y=TracedSidePlan(step_plans=[traced_step]), + ), + ) -class TestAlignerPlanInTensorComparisonRecord: - def test_comparison_record_with_aligner_plan(self) -> None: - plan: AlignerPlan = _make_aligner_plan() - record: TensorComparisonRecord = _make_comparison_record( +class TestAlignerPlanInComparisonTensorRecord: + def test_comparison_record_with_traced_plan(self) -> None: + traced_plan: TracedAlignerPlan = _make_traced_aligner_plan() + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) - record_with_plan = record.model_copy(update={"aligner_plan": plan}) - assert record_with_plan.aligner_plan is not None - assert record_with_plan.aligner_plan.per_step_plans.x[0].step == 0 + record_with_plan = record.model_copy(update={"traced_plan": traced_plan}) + assert record_with_plan.traced_plan is not None + assert record_with_plan.traced_plan.per_side.x.step_plans[0].step == 0 - def test_aligner_plan_json_roundtrip(self) -> None: - plan: AlignerPlan = _make_aligner_plan() - record: TensorComparisonRecord = _make_comparison_record( + def test_traced_plan_json_roundtrip(self) -> None: + traced_plan: TracedAlignerPlan = _make_traced_aligner_plan() + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) - record_with_plan = record.model_copy(update={"aligner_plan": plan}) + record_with_plan = record.model_copy(update={"traced_plan": traced_plan}) json_str: str = record_with_plan.model_dump_json() parsed = json.loads(json_str) - assert "aligner_plan" in parsed + assert "traced_plan" in parsed assert ( - parsed["aligner_plan"]["per_step_plans"]["x"][0]["sub_plans"][0]["type"] + parsed["traced_plan"]["per_side"]["x"]["step_plans"][0]["sub_plans"][0][ + "plan" + ]["type"] == "unsharder" ) - roundtripped: TensorComparisonRecord = parse_record_json(json_str) - assert roundtripped.aligner_plan is not None + roundtripped: ComparisonTensorRecord = parse_record_json(json_str) + assert roundtripped.traced_plan is not None assert ( - roundtripped.aligner_plan.per_step_plans.x[0].sub_plans[0].type + roundtripped.traced_plan.per_side.x.step_plans[0].sub_plans[0].plan.type == "unsharder" ) - def test_comparison_record_without_aligner_plan(self) -> None: - record: TensorComparisonRecord = _make_comparison_record( + def test_comparison_record_without_traced_plan(self) -> None: + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) json_str: str = record.model_dump_json() - roundtripped: TensorComparisonRecord = parse_record_json(json_str) - assert roundtripped.aligner_plan is None + roundtripped: ComparisonTensorRecord = parse_record_json(json_str) + assert roundtripped.traced_plan is None - def test_aligner_plan_text_format(self) -> None: - plan: AlignerPlan = _make_aligner_plan() - record: TensorComparisonRecord = _make_comparison_record( + def test_traced_plan_text_format(self) -> None: + traced_plan: TracedAlignerPlan = _make_traced_aligner_plan() + record: ComparisonTensorRecord = _make_comparison_record( diff=_make_diff_info(passed=True), ) - record_with_plan = record.model_copy(update={"aligner_plan": plan}) + record_with_plan = record.model_copy(update={"traced_plan": traced_plan}) text: str = record_with_plan.to_text() assert "Aligner Plan:" 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 59149040b..5eb316015 100644 --- a/test/registered/debug_utils/comparator/test_output_types.py +++ b/test/registered/debug_utils/comparator/test_output_types.py @@ -34,15 +34,15 @@ 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 ( + ComparisonNonTensorRecord, + ComparisonSkipRecord, + ComparisonTensorRecord, ConfigRecord, ErrorLog, InfoLog, LogRecord, - NonTensorComparisonRecord, RecordLocation, - SkipComparisonRecord, SummaryRecord, - TensorComparisonRecord, _format_aligner_plan, _split_logs, ) @@ -150,20 +150,20 @@ class TestConfigRecord: # --------------------------------------------------------------------------- -# SkipComparisonRecord +# ComparisonSkipRecord # --------------------------------------------------------------------------- -class TestSkipComparisonRecord: +class TestComparisonSkipRecord: def test_format_body_no_step(self) -> None: - record: SkipComparisonRecord = SkipComparisonRecord( + record: ComparisonSkipRecord = ComparisonSkipRecord( name="layer.weight", reason="zero-dim tensor", ) assert record._format_body() == "Skip: layer.weight (zero-dim tensor)" def test_format_body_with_step(self) -> None: - record: SkipComparisonRecord = SkipComparisonRecord( + record: ComparisonSkipRecord = ComparisonSkipRecord( name="layer.weight", reason="scalar", location=RecordLocation(step=3), @@ -171,7 +171,7 @@ class TestSkipComparisonRecord: assert record._format_body() == "Skip: layer.weight (step=3) (scalar)" def test_format_rich_body(self) -> None: - record: SkipComparisonRecord = SkipComparisonRecord( + record: ComparisonSkipRecord = ComparisonSkipRecord( name="attn.qkv", reason="no baseline", ) @@ -179,14 +179,14 @@ class TestSkipComparisonRecord: assert body == "[dim]⊘ attn.qkv ── skipped (no baseline)[/]" def test_category_skipped(self) -> None: - record: SkipComparisonRecord = SkipComparisonRecord( + record: ComparisonSkipRecord = ComparisonSkipRecord( name="x", reason="r", ) assert record.category == "skipped" def test_category_failed(self) -> None: - record: SkipComparisonRecord = SkipComparisonRecord( + record: ComparisonSkipRecord = ComparisonSkipRecord( name="x", reason="r", errors=[ErrorLog(category="e", message="boom")], @@ -195,13 +195,13 @@ class TestSkipComparisonRecord: # --------------------------------------------------------------------------- -# NonTensorComparisonRecord +# ComparisonNonTensorRecord # --------------------------------------------------------------------------- -class TestNonTensorComparisonRecord: +class TestComparisonNonTensorRecord: def test_format_body_equal(self) -> None: - record: NonTensorComparisonRecord = NonTensorComparisonRecord( + record: ComparisonNonTensorRecord = ComparisonNonTensorRecord( name="config.lr", baseline_value="0.001", target_value="0.001", @@ -212,7 +212,7 @@ class TestNonTensorComparisonRecord: assert record._format_body() == "NonTensor: config.lr = 0.001 (float) [equal]" def test_format_body_not_equal(self) -> None: - record: NonTensorComparisonRecord = NonTensorComparisonRecord( + record: ComparisonNonTensorRecord = ComparisonNonTensorRecord( name="config.lr", baseline_value="0.001", target_value="0.01", @@ -227,7 +227,7 @@ class TestNonTensorComparisonRecord: ) def test_format_rich_body_equal(self) -> None: - record: NonTensorComparisonRecord = NonTensorComparisonRecord( + record: ComparisonNonTensorRecord = ComparisonNonTensorRecord( name="config.lr", baseline_value="0.001", target_value="0.001", @@ -238,7 +238,7 @@ class TestNonTensorComparisonRecord: assert record._format_rich_body() == ("═ config.lr = 0.001 (float) [green]✓[/]") def test_format_rich_body_not_equal(self) -> None: - record: NonTensorComparisonRecord = NonTensorComparisonRecord( + record: ComparisonNonTensorRecord = ComparisonNonTensorRecord( name="config.lr", baseline_value="0.001", target_value="0.01", @@ -253,7 +253,7 @@ class TestNonTensorComparisonRecord: ) def test_with_step(self) -> None: - record: NonTensorComparisonRecord = NonTensorComparisonRecord( + record: ComparisonNonTensorRecord = ComparisonNonTensorRecord( name="bias", baseline_value="True", target_value="True", @@ -265,7 +265,7 @@ class TestNonTensorComparisonRecord: assert "(step=5)" in record._format_body() def test_category(self) -> None: - passed: NonTensorComparisonRecord = NonTensorComparisonRecord( + passed: ComparisonNonTensorRecord = ComparisonNonTensorRecord( name="x", baseline_value="1", target_value="1", @@ -273,7 +273,7 @@ class TestNonTensorComparisonRecord: target_type="int", values_equal=True, ) - failed: NonTensorComparisonRecord = NonTensorComparisonRecord( + failed: ComparisonNonTensorRecord = ComparisonNonTensorRecord( name="x", baseline_value="1", target_value="2", @@ -328,13 +328,13 @@ class TestSummaryRecord: # --------------------------------------------------------------------------- -# TensorComparisonRecord._format_body +# ComparisonTensorRecord._format_body # --------------------------------------------------------------------------- -class TestTensorComparisonRecordFormatBody: +class TestComparisonTensorRecordFormatBody: def test_basic(self) -> None: - record: TensorComparisonRecord = TensorComparisonRecord( + record: ComparisonTensorRecord = ComparisonTensorRecord( name="hidden", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -365,7 +365,7 @@ class TestTensorComparisonRecordFormatBody: def test_with_replicated_checks(self) -> None: from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult - record: TensorComparisonRecord = TensorComparisonRecord( + record: ComparisonTensorRecord = ComparisonTensorRecord( name="hidden", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -420,7 +420,7 @@ class TestTensorComparisonRecordFormatBody: y=TracedSidePlan(step_plans=[]), ), ) - record: TensorComparisonRecord = TensorComparisonRecord( + record: ComparisonTensorRecord = ComparisonTensorRecord( name="hidden", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -453,7 +453,7 @@ class TestTensorComparisonRecordFormatBody: ) def test_with_step(self) -> None: - record: TensorComparisonRecord = TensorComparisonRecord( + record: ComparisonTensorRecord = ComparisonTensorRecord( name="hidden", baseline=_make_tensor_info(), target=_make_tensor_info(), @@ -679,7 +679,7 @@ class TestOutputRecordLogAttachment: assert text == "Config: {'a': 1}\n ✗ err1\n ℹ note1" def test_to_rich_string_body(self) -> None: - record: SkipComparisonRecord = SkipComparisonRecord( + record: ComparisonSkipRecord = ComparisonSkipRecord( name="x", reason="r", errors=[ErrorLog(category="e", message="oops")], diff --git a/test/registered/debug_utils/comparator/test_per_token_visualizer.py b/test/registered/debug_utils/comparator/test_per_token_visualizer.py index b99f95aba..a9c9c1733 100644 --- a/test/registered/debug_utils/comparator/test_per_token_visualizer.py +++ b/test/registered/debug_utils/comparator/test_per_token_visualizer.py @@ -9,7 +9,7 @@ from pathlib import Path import pytest import torch -from sglang.srt.debug_utils.comparator.output_types import TensorComparisonRecord +from sglang.srt.debug_utils.comparator.output_types import ComparisonTensorRecord from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( compare_tensor_pair, ) @@ -31,8 +31,8 @@ def _make_comparison_record( baseline: torch.Tensor, target: torch.Tensor, seq_dim: int = 0, -) -> TensorComparisonRecord: - """Build a TensorComparisonRecord with per-token data from raw tensors.""" +) -> ComparisonTensorRecord: + """Build a ComparisonTensorRecord with per-token data from raw tensors.""" info = compare_tensor_pair( x_baseline=baseline, x_target=target, @@ -40,7 +40,7 @@ def _make_comparison_record( diff_threshold=1e-3, seq_dim=seq_dim, ) - return TensorComparisonRecord(**info.model_dump()) + return ComparisonTensorRecord(**info.model_dump()) class TestPerTokenVisualizer: @@ -68,7 +68,7 @@ class TestPerTokenVisualizer: name="no_per_token", diff_threshold=1e-3, ) - record = TensorComparisonRecord(**info.model_dump()) + record = ComparisonTensorRecord(**info.model_dump()) output_path: Path = tmp_path / "no_data.png" result = generate_per_token_heatmap(records=[record], output_path=output_path) @@ -82,7 +82,7 @@ class TestPerTokenVisualizer: ) torch.manual_seed(42) - records: list[TensorComparisonRecord] = [ + records: list[ComparisonTensorRecord] = [ _make_comparison_record( name=f"tensor_{i}", baseline=torch.randn(16, 32), @@ -108,7 +108,7 @@ class TestPerTokenVisualizer: ) torch.manual_seed(42) - records: list[TensorComparisonRecord] = [ + records: list[ComparisonTensorRecord] = [ _make_comparison_record( name="short", baseline=torch.randn(4, 8), diff --git a/test/registered/debug_utils/comparator/test_utils.py b/test/registered/debug_utils/comparator/test_utils.py index 954cce00a..2b74a2462 100644 --- a/test/registered/debug_utils/comparator/test_utils.py +++ b/test/registered/debug_utils/comparator/test_utils.py @@ -410,6 +410,36 @@ class TestComputeExitCode: == 1 ) + def test_errored_with_passed_exits_one(self): + """Has errored bundle even with passed → exit 1.""" + summary = SummaryRecord(total=3, passed=2, failed=0, skipped=0, errored=1) + assert ( + compute_exit_code( + summary, + allow_skipped_pattern=".*", + skipped_names=[], + allow_failed_pattern=None, + failed_names=[], + errored_names=["broken_tensor"], + ) + == 1 + ) + + def test_errored_only_exits_one(self): + """All errored → exit 1 (passed==0 already exits 1, but errored also independently triggers).""" + summary = SummaryRecord(total=1, passed=0, failed=0, skipped=0, errored=1) + assert ( + compute_exit_code( + summary, + allow_skipped_pattern=".*", + skipped_names=[], + allow_failed_pattern=None, + failed_names=[], + errored_names=["broken_tensor"], + ) + == 1 + ) + def _make_pt(directory: Path) -> None: directory.mkdir(parents=True, exist_ok=True) diff --git a/test/registered/debug_utils/comparator/testing_helpers.py b/test/registered/debug_utils/comparator/testing_helpers.py index 1106e3a1a..9af13d034 100644 --- a/test/registered/debug_utils/comparator/testing_helpers.py +++ b/test/registered/debug_utils/comparator/testing_helpers.py @@ -4,6 +4,12 @@ from __future__ import annotations from typing import Optional +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci( + est_time=0, suite="default", nightly=True, disabled="helper module, no tests" +) + from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( DiffInfo, TensorInfo,