minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Enhance error resilience in dump comparator (#19685)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-02 19:08:35 +08:00
committed by GitHub
parent e5ef845cad
commit 7579ab3f33
18 changed files with 495 additions and 303 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/debug_utils/comparator/__init__.py
+2 -2
View File
@@ -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 from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import ( # noqa: F401
AlignerPlan, 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()
python/sglang/srt/debug_utils/comparator/bundle_comparator.py
+14 -14
View File
@@ -31,10 +31,10 @@ from sglang.srt.debug_utils.comparator.meta_overrider import MetaOverrider
from sglang.srt.debug_utils.comparator.output_types import ( from sglang.srt.debug_utils.comparator.output_types import (
BundleFileInfo, BundleFileInfo,
BundleSideInfo, BundleSideInfo,
ComparisonNonTensorRecord,
ComparisonSkipRecord,
ComparisonTensorRecord,
ErrorLog, ErrorLog,
NonTensorComparisonRecord,
SkipComparisonRecord,
TensorComparisonRecord,
_split_logs, _split_logs,
) )
from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import ( from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
@@ -51,7 +51,7 @@ def _collect_bundle_side_info(
metas: list[dict[str, Any]], metas: list[dict[str, Any]],
) -> BundleSideInfo: ) -> BundleSideInfo:
from sglang.srt.debug_utils.comparator.display import ( from sglang.srt.debug_utils.comparator.display import (
PARALLEL_INFO_KEYS, _PARALLEL_INFO_KEYS,
extract_parallel_info, extract_parallel_info,
) )
@@ -61,7 +61,7 @@ def _collect_bundle_side_info(
tensor: torch.Tensor = item.value tensor: torch.Tensor = item.value
parallel_info: dict[str, str] = {} 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, {})) extract_parallel_info(row_data=parallel_info, info=meta.get(key, {}))
files.append( files.append(
@@ -91,7 +91,7 @@ def compare_bundle_pair(
viz_output_dir: Optional[Path] = None, viz_output_dir: Optional[Path] = None,
compute_per_token: bool = False, compute_per_token: bool = False,
meta_overrider: Optional[MetaOverrider] = None, meta_overrider: Optional[MetaOverrider] = None,
) -> Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord]: ) -> Union[ComparisonTensorRecord, ComparisonSkipRecord, ComparisonNonTensorRecord]:
with log_sink.context() as collected_logs: with log_sink.context() as collected_logs:
result = _compare_bundle_pair_inner( result = _compare_bundle_pair_inner(
name=name, name=name,
@@ -124,7 +124,7 @@ def _compare_bundle_pair_inner(
viz_output_dir: Optional[Path] = None, viz_output_dir: Optional[Path] = None,
compute_per_token: bool = False, compute_per_token: bool = False,
meta_overrider: Optional[MetaOverrider] = None, meta_overrider: Optional[MetaOverrider] = None,
) -> Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord]: ) -> Union[ComparisonTensorRecord, ComparisonSkipRecord, ComparisonNonTensorRecord]:
# 1. Load all successfully loaded values # 1. Load all successfully loaded values
all_pair: Pair[list[ValueWithMeta]] = Pair( all_pair: Pair[list[ValueWithMeta]] = Pair(
x=_load_all_values(filenames=filenames_pair.x, base_path=dir_pair.x), 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: if not all_pair.x or not all_pair.y:
reason = "baseline_load_failed" if not all_pair.x else "target_load_failed" 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 # 1b. Dims override: patch meta["dims"] before DP filter reads it
# (--override-dims may add ``# dp:=moe_dp``, so it must run first) # (--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, viz_output_dir: Optional[Path] = None,
compute_per_token: bool = False, compute_per_token: bool = False,
) -> Union[TensorComparisonRecord, SkipComparisonRecord]: ) -> Union[ComparisonTensorRecord, ComparisonSkipRecord]:
if not valid_pair.x or not valid_pair.y: if not valid_pair.x or not valid_pair.y:
reason = "baseline_load_failed" if not valid_pair.x else "target_load_failed" 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) # Plan (meta only, no tensor)
metas_pair: Pair[list[dict[str, Any]]] = valid_pair.map( 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 assert aligner_result.failed_side_xy is not None
side_name: str = _FAILED_SIDE_MAP[aligner_result.failed_side_xy] side_name: str = _FAILED_SIDE_MAP[aligner_result.failed_side_xy]
reason: str = f"{side_name}_load_failed" 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 # Resolve seq_dim for per-token computation
seq_dim: Optional[int] = ( seq_dim: Optional[int] = (
@@ -263,7 +263,7 @@ def _compare_bundle_pair_tensor_type(
diff_threshold=diff_threshold, diff_threshold=diff_threshold,
seq_dim=seq_dim, seq_dim=seq_dim,
) )
record = TensorComparisonRecord( record = ComparisonTensorRecord(
**info.model_dump(), **info.model_dump(),
traced_plan=aligner_result.traced_plan, traced_plan=aligner_result.traced_plan,
replicated_checks=replicated_checks, replicated_checks=replicated_checks,
@@ -331,7 +331,7 @@ def _compare_bundle_pair_non_tensor_type(
*, *,
name: str, name: str,
value_pair: Pair[list[ValueWithMeta]], value_pair: Pair[list[ValueWithMeta]],
) -> NonTensorComparisonRecord: ) -> ComparisonNonTensorRecord:
baseline_value: Any = value_pair.x[0].value baseline_value: Any = value_pair.x[0].value
target_value: Any = value_pair.y[0].value target_value: Any = value_pair.y[0].value
@@ -340,7 +340,7 @@ def _compare_bundle_pair_non_tensor_type(
except Exception: except Exception:
values_equal = False values_equal = False
return NonTensorComparisonRecord( return ComparisonNonTensorRecord(
name=name, name=name,
baseline_value=repr(baseline_value), baseline_value=repr(baseline_value),
target_value=repr(target_value), target_value=repr(target_value),
python/sglang/srt/debug_utils/comparator/display.py
+9 -39
View File
@@ -3,13 +3,9 @@ from __future__ import annotations
from collections import defaultdict from collections import defaultdict
from io import StringIO from io import StringIO
from pathlib import Path from pathlib import Path
from typing import TYPE_CHECKING, Any, Optional from typing import Any, Optional
import polars as pl import polars as pl
import rich.table
if TYPE_CHECKING:
from rich.table import Table
from sglang.srt.debug_utils.comparator.output_types import ( from sglang.srt.debug_utils.comparator.output_types import (
InputIdsRecord, InputIdsRecord,
@@ -43,45 +39,19 @@ def emit_display_records(
def _render_polars_as_text(df: pl.DataFrame, *, title: Optional[str] = None) -> str: def _render_polars_as_text(df: pl.DataFrame, *, title: Optional[str] = None) -> str:
from rich.console import Console 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() buf = StringIO()
Console(file=buf, force_terminal=False, width=200).print(table) Console(file=buf, force_terminal=False, width=200).print(table)
return buf.getvalue().rstrip("\n") 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( def _collect_rank_info(
df: pl.DataFrame, dump_dir: Path df: pl.DataFrame, dump_dir: Path
) -> Optional[list[dict[str, Any]]]: ) -> Optional[list[dict[str, Any]]]:
@@ -99,7 +69,7 @@ def _collect_rank_info(
row_data: dict[str, Any] = {"rank": row["rank"]} row_data: dict[str, Any] = {"rank": row["rank"]}
for key in PARALLEL_INFO_KEYS: 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) table_rows.append(row_data)
return table_rows or None return table_rows or None
@@ -149,7 +119,7 @@ def _collect_input_ids_and_positions(
return table_rows or None 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"): if not info or info.get("error"):
return return
python/sglang/srt/debug_utils/comparator/entrypoint.py
+57 -33
View File
@@ -2,6 +2,7 @@ from __future__ import annotations
import argparse import argparse
import sys import sys
import traceback as _traceback_module
from pathlib import Path from pathlib import Path
from typing import Any, Iterator, Optional, Union 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.display import emit_display_records
from sglang.srt.debug_utils.comparator.meta_overrider import MetaOverrider from sglang.srt.debug_utils.comparator.meta_overrider import MetaOverrider
from sglang.srt.debug_utils.comparator.output_types import ( from sglang.srt.debug_utils.comparator.output_types import (
ComparisonErrorRecord,
ComparisonNonTensorRecord,
ComparisonSkipRecord,
ComparisonTensorRecord,
ConfigRecord, ConfigRecord,
NonTensorComparisonRecord,
RecordLocation, RecordLocation,
SkipComparisonRecord,
SummaryRecord, SummaryRecord,
TensorComparisonRecord,
) )
from sglang.srt.debug_utils.comparator.per_token_visualizer import ( from sglang.srt.debug_utils.comparator.per_token_visualizer import (
generate_per_token_heatmap, generate_per_token_heatmap,
@@ -83,12 +85,6 @@ def run(args: argparse.Namespace) -> int:
verbosity=args.verbosity, 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: try:
report_sink.add(ConfigRecord(config=vars(args))) 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, compute_per_token=visualize_per_token is not None,
meta_overrider=meta_overrider, meta_overrider=meta_overrider,
) )
summary, skipped_names, failed_names = _consume_comparison_records( summary, skipped_names, failed_names, errored_names = (
comparison_records=comparison_records, _consume_comparison_records(
visualize_per_token=visualize_per_token, comparison_records=comparison_records,
visualize_per_token=visualize_per_token,
)
) )
return compute_exit_code( return compute_exit_code(
summary, summary,
@@ -152,6 +150,7 @@ def run(args: argparse.Namespace) -> int:
skipped_names=skipped_names, skipped_names=skipped_names,
allow_failed_pattern=args.allow_failed_pattern, allow_failed_pattern=args.allow_failed_pattern,
failed_names=failed_names, failed_names=failed_names,
errored_names=errored_names,
) )
finally: finally:
report_sink.close() report_sink.close()
@@ -219,7 +218,12 @@ def _compare_bundle_pairs(
compute_per_token: bool = False, compute_per_token: bool = False,
meta_overrider: Optional[MetaOverrider] = None, meta_overrider: Optional[MetaOverrider] = None,
) -> Iterator[ ) -> Iterator[
Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord] Union[
ComparisonTensorRecord,
ComparisonSkipRecord,
ComparisonNonTensorRecord,
ComparisonErrorRecord,
]
]: ]:
for bundle_info_pair in bundle_info_pairs: for bundle_info_pair in bundle_info_pairs:
if not bundle_info_pair.y: if not bundle_info_pair.y:
@@ -229,20 +233,32 @@ def _compare_bundle_pairs(
filenames_pair: Pair[list[str]] = bundle_info_pair.map( filenames_pair: Pair[list[str]] = bundle_info_pair.map(
lambda infos: [info.filename for info in infos] lambda infos: [info.filename for info in infos]
) )
record: Union[ record: Union[
TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord ComparisonTensorRecord,
] = compare_bundle_pair( ComparisonSkipRecord,
name=name, ComparisonNonTensorRecord,
filenames_pair=filenames_pair, ComparisonErrorRecord,
dir_pair=dir_pair, ]
token_aligner_mode=token_aligner_mode, try:
token_aligner_plan=token_aligner_plan, record = compare_bundle_pair(
diff_threshold=diff_threshold, name=name,
thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair, filenames_pair=filenames_pair,
viz_output_dir=viz_output_dir, dir_pair=dir_pair,
compute_per_token=compute_per_token, token_aligner_mode=token_aligner_mode,
meta_overrider=meta_overrider, 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} 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 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( def _consume_comparison_records(
*, *,
comparison_records: Iterator[ comparison_records: Iterator[
Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord] Union[
ComparisonTensorRecord,
ComparisonSkipRecord,
ComparisonNonTensorRecord,
ComparisonErrorRecord,
]
], ],
visualize_per_token: Optional[Path] = None, visualize_per_token: Optional[Path] = None,
) -> tuple[SummaryRecord, list[str], list[str]]: ) -> tuple[SummaryRecord, list[str], list[str], list[str]]:
counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0} counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0, "errored": 0}
collected_comparisons: list[TensorComparisonRecord] = [] collected_comparisons: list[ComparisonTensorRecord] = []
skipped_names: list[str] = [] skipped_names: list[str] = []
failed_names: list[str] = [] failed_names: list[str] = []
errored_names: list[str] = []
for record in comparison_records: for record in comparison_records:
counts[record.category] += 1 counts[record.category] += 1
report_sink.add(record) 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) skipped_names.append(record.name)
if record.category == "failed": if record.category == "failed":
failed_names.append(record.name) failed_names.append(record.name)
if isinstance(record, ComparisonErrorRecord):
errored_names.append(record.name)
if visualize_per_token is not None and isinstance( if visualize_per_token is not None and isinstance(
record, TensorComparisonRecord record, ComparisonTensorRecord
): ):
collected_comparisons.append(record) collected_comparisons.append(record)
@@ -285,7 +309,7 @@ def _consume_comparison_records(
output_path=visualize_per_token, 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: def parse_args(argv: list[str]) -> argparse.Namespace:
python/sglang/srt/debug_utils/comparator/output_formatter.py
+43 -13
View File
@@ -26,14 +26,15 @@ if TYPE_CHECKING:
) )
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan
from sglang.srt.debug_utils.comparator.output_types import ( from sglang.srt.debug_utils.comparator.output_types import (
ComparisonErrorRecord,
ComparisonNonTensorRecord,
ComparisonSkipRecord,
ComparisonTensorRecord,
ConfigRecord, ConfigRecord,
ErrorLog, ErrorLog,
InfoLog, InfoLog,
LogRecord, LogRecord,
NonTensorComparisonRecord,
SkipComparisonRecord,
SummaryRecord, SummaryRecord,
TensorComparisonRecord,
_OutputRecord, _OutputRecord,
_TableRecord, _TableRecord,
) )
@@ -111,15 +112,15 @@ def _format_config_rich_body(
return Panel("\n".join(lines), title="Comparator Config", border_style="cyan") 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})" return f"Skip: {record.name}{record._format_location_suffix()} ({record.reason})"
def _format_skip_rich_body( def _format_skip_rich_body(
record: SkipComparisonRecord, verbosity: Verbosity = "normal" record: ComparisonSkipRecord, verbosity: Verbosity = "normal"
) -> RenderableType: ) -> RenderableType:
suffix: str = record._format_location_suffix() suffix: str = record._format_location_suffix()
return ( 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 ───────────────────────────────────────────────────── # ── _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) body: str = record._format_location_prefix() + format_comparison(record)
if record.replicated_checks: if record.replicated_checks:
body += "\n" + format_replicated_checks(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( def _format_tensor_comparison_rich_body(
record: TensorComparisonRecord, verbosity: Verbosity = "normal" record: ComparisonTensorRecord, verbosity: Verbosity = "normal"
) -> RenderableType: ) -> RenderableType:
from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import ( from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import (
format_comparison_rich, 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() suffix: str = record._format_location_suffix()
if record.values_equal: if record.values_equal:
return f"NonTensor: {record.name}{suffix} = {record.baseline_value} ({record.baseline_type}) [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( def _format_non_tensor_rich_body(
record: NonTensorComparisonRecord, verbosity: Verbosity = "normal" record: ComparisonNonTensorRecord, verbosity: Verbosity = "normal"
) -> RenderableType: ) -> RenderableType:
suffix: str = record._format_location_suffix() suffix: str = record._format_location_suffix()
name: str = escape(record.name) name: str = escape(record.name)
@@ -216,10 +241,13 @@ def _format_non_tensor_rich_body(
def _format_summary_body(record: SummaryRecord) -> str: def _format_summary_body(record: SummaryRecord) -> str:
return ( text: str = (
f"Summary: {record.passed} passed, {record.failed} failed, " f"Summary: {record.passed} passed, {record.failed} failed, "
f"{record.skipped} skipped (total {record.total})" f"{record.skipped} skipped (total {record.total})"
) )
if record.errored > 0:
text += f", {record.errored} errored"
return text
def _format_summary_rich_body( def _format_summary_rich_body(
@@ -231,6 +259,8 @@ def _format_summary_rich_body(
f"[yellow]{record.skipped} skipped[/] │ " f"[yellow]{record.skipped} skipped[/] │ "
f"{record.total} total" f"{record.total} total"
) )
if record.errored > 0:
text += f" │ [bold red]{record.errored} errored[/]"
return Panel(text, title="SUMMARY", border_style="bold") return Panel(text, title="SUMMARY", border_style="bold")
python/sglang/srt/debug_utils/comparator/output_types.py
+33 -30
View File
@@ -13,6 +13,8 @@ from sglang.srt.debug_utils.comparator.output_formatter import ( # noqa: F401
from sglang.srt.debug_utils.comparator.output_formatter import ( from sglang.srt.debug_utils.comparator.output_formatter import (
_format_config_body, _format_config_body,
_format_config_rich_body, _format_config_rich_body,
_format_error_body,
_format_error_rich_body,
_format_log_body, _format_log_body,
_format_non_tensor_body, _format_non_tensor_body,
_format_non_tensor_rich_body, _format_non_tensor_rich_body,
@@ -146,8 +148,8 @@ class ConfigRecord(_OutputRecord):
return _format_config_rich_body(self, verbosity=verbosity) return _format_config_rich_body(self, verbosity=verbosity)
class SkipComparisonRecord(_BaseComparisonRecord): class ComparisonSkipRecord(_BaseComparisonRecord):
type: Literal["skip"] = "skip" type: Literal["comparison_skip"] = "comparison_skip"
name: str name: str
reason: str reason: str
@@ -164,6 +166,23 @@ class SkipComparisonRecord(_BaseComparisonRecord):
return _format_skip_rich_body(self, verbosity=verbosity) 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): class _TableRecord(_OutputRecord):
label: str label: str
rows: list[dict[str, Any]] rows: list[dict[str, Any]]
@@ -177,15 +196,6 @@ class _TableRecord(_OutputRecord):
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType: def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return _format_table_rich_body(self, verbosity=verbosity) 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): class RankInfoRecord(_TableRecord):
type: Literal["rank_info"] = "rank_info" type: Literal["rank_info"] = "rank_info"
@@ -201,10 +211,10 @@ class InputIdsRecord(_TableRecord):
return f"{self.label} input_ids & positions" return f"{self.label} input_ids & positions"
class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord): class ComparisonTensorRecord(TensorComparisonInfo, _BaseComparisonRecord):
model_config = ConfigDict(extra="forbid", defer_build=True) model_config = ConfigDict(extra="forbid", defer_build=True)
type: Literal["comparison"] = "comparison" type: Literal["comparison_tensor"] = "comparison_tensor"
traced_plan: Optional[TracedAlignerPlan] = None traced_plan: Optional[TracedAlignerPlan] = None
replicated_checks: list[ReplicatedCheckResult] = Field(default_factory=list) replicated_checks: list[ReplicatedCheckResult] = Field(default_factory=list)
raw_bundle_info: Optional[Pair[BundleSideInfo]] = None raw_bundle_info: Optional[Pair[BundleSideInfo]] = None
@@ -224,8 +234,8 @@ class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord):
return _format_tensor_comparison_rich_body(self, verbosity=verbosity) return _format_tensor_comparison_rich_body(self, verbosity=verbosity)
class NonTensorComparisonRecord(_BaseComparisonRecord): class ComparisonNonTensorRecord(_BaseComparisonRecord):
type: Literal["non_tensor"] = "non_tensor" type: Literal["comparison_non_tensor"] = "comparison_non_tensor"
name: str name: str
baseline_value: str baseline_value: str
target_value: str target_value: str
@@ -245,15 +255,6 @@ class NonTensorComparisonRecord(_BaseComparisonRecord):
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType: def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return _format_non_tensor_rich_body(self, verbosity=verbosity) 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): class SummaryRecord(_OutputRecord):
type: Literal["summary"] = "summary" type: Literal["summary"] = "summary"
@@ -261,13 +262,15 @@ class SummaryRecord(_OutputRecord):
passed: int passed: int
failed: int failed: int
skipped: int skipped: int
errored: int = 0
@model_validator(mode="after") @model_validator(mode="after")
def _validate_totals(self) -> "SummaryRecord": 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: if self.total != expected:
raise ValueError( 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 return self
@@ -277,7 +280,6 @@ class SummaryRecord(_OutputRecord):
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType: def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return _format_summary_rich_body(self, verbosity=verbosity) return _format_summary_rich_body(self, verbosity=verbosity)
return _format_summary_body(self)
class LogRecord(_OutputRecord): class LogRecord(_OutputRecord):
type: Literal["log"] = "log" type: Literal["log"] = "log"
@@ -291,9 +293,10 @@ AnyRecord = Annotated[
ConfigRecord, ConfigRecord,
RankInfoRecord, RankInfoRecord,
InputIdsRecord, InputIdsRecord,
SkipComparisonRecord, ComparisonSkipRecord,
TensorComparisonRecord, ComparisonErrorRecord,
NonTensorComparisonRecord, ComparisonTensorRecord,
ComparisonNonTensorRecord,
SummaryRecord, SummaryRecord,
LogRecord, LogRecord,
], ],
python/sglang/srt/debug_utils/comparator/per_token_visualizer.py
+3 -3
View File
@@ -9,12 +9,12 @@ from __future__ import annotations
from pathlib import Path from pathlib import Path
from typing import Optional 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( def generate_per_token_heatmap(
*, *,
records: list[TensorComparisonRecord], records: list[ComparisonTensorRecord],
output_path: Path, output_path: Path,
) -> Optional[Path]: ) -> Optional[Path]:
"""Generate a per-token relative difference heatmap PNG. """Generate a per-token relative difference heatmap PNG.
@@ -31,7 +31,7 @@ def generate_per_token_heatmap(
def _collect_per_token_data( def _collect_per_token_data(
*, *,
records: list[TensorComparisonRecord], records: list[ComparisonTensorRecord],
) -> list[tuple[str, list[float]]]: ) -> list[tuple[str, list[float]]]:
rows: list[tuple[str, list[float]]] = [] rows: list[tuple[str, list[float]]] = []
for record in records: for record in records:
python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py
+4 -4
View File
@@ -20,9 +20,9 @@ if TYPE_CHECKING:
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import AlignerPlan
from sglang.srt.debug_utils.comparator.output_types import ( from sglang.srt.debug_utils.comparator.output_types import (
BundleSideInfo, BundleSideInfo,
ComparisonTensorRecord,
ReplicatedCheckResult, ReplicatedCheckResult,
ShapeSnapshot, ShapeSnapshot,
TensorComparisonRecord,
) )
from sglang.srt.debug_utils.comparator.utils import Pair 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( def format_comparison_rich(
record: TensorComparisonRecord, record: ComparisonTensorRecord,
verbosity: Verbosity = "normal", verbosity: Verbosity = "normal",
) -> str: ) -> str:
if verbosity == "minimal": 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) passed, color, marker = _category_marker(record.category)
name_part: str = f"[bold {color}]{escape(record.name):30s}[/]" 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( def _format_comparison_normal_or_verbose(
*, *,
record: TensorComparisonRecord, record: ComparisonTensorRecord,
verbose: bool, verbose: bool,
) -> str: ) -> str:
passed, color, marker = _category_marker(record.category) passed, color, marker = _category_marker(record.category)
python/sglang/srt/debug_utils/comparator/utils.py
+4
View File
@@ -141,10 +141,14 @@ def compute_exit_code(
skipped_names: list[str], skipped_names: list[str],
allow_failed_pattern: Optional[str], allow_failed_pattern: Optional[str],
failed_names: list[str], failed_names: list[str],
errored_names: Optional[list[str]] = None,
) -> int: ) -> int:
if summary.passed == 0: if summary.passed == 0:
return 1 return 1
if errored_names:
return 1
if not _is_all_match_pattern(pattern=allow_failed_pattern, strings=failed_names): if not _is_all_match_pattern(pattern=allow_failed_pattern, strings=failed_names):
return 1 return 1
test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py
+16 -16
View File
@@ -34,9 +34,9 @@ from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, TokenLayou
from sglang.srt.debug_utils.comparator.output_types import ( from sglang.srt.debug_utils.comparator.output_types import (
BundleFileInfo, BundleFileInfo,
BundleSideInfo, BundleSideInfo,
ComparisonTensorRecord,
ReplicatedCheckResult, ReplicatedCheckResult,
ShapeSnapshot, ShapeSnapshot,
TensorComparisonRecord,
) )
from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import ( from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import (
_format_abs_diff_percentiles_rich, _format_abs_diff_percentiles_rich,
@@ -273,9 +273,9 @@ def _make_comparison_record(
replicated_checks: list[ReplicatedCheckResult] | None = None, replicated_checks: list[ReplicatedCheckResult] | None = None,
raw_bundle_info: Pair[BundleSideInfo] | None = None, raw_bundle_info: Pair[BundleSideInfo] | None = None,
traced_plan: TracedAlignerPlan | None = None, traced_plan: TracedAlignerPlan | None = None,
) -> TensorComparisonRecord: ) -> ComparisonTensorRecord:
s: list[int] = shape if shape is not None else [4, 8] s: list[int] = shape if shape is not None else [4, 8]
return TensorComparisonRecord( return ComparisonTensorRecord(
name=name, name=name,
baseline=_make_tensor_info(shape=s, dtype=dtype, sample=sample), baseline=_make_tensor_info(shape=s, dtype=dtype, sample=sample),
target=_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'.""" """format_comparison_rich() with verbosity='minimal'."""
def test_passed(self) -> None: def test_passed(self) -> None:
record: TensorComparisonRecord = _make_comparison_record( record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(rel_diff=1e-4, passed=True), diff=_make_diff(rel_diff=1e-4, passed=True),
) )
result: str = format_comparison_rich(record, verbosity="minimal") result: str = format_comparison_rich(record, verbosity="minimal")
@@ -428,7 +428,7 @@ class TestFormatComparisonRichMinimal:
) )
def test_failed(self) -> None: def test_failed(self) -> None:
record: TensorComparisonRecord = _make_comparison_record( record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(rel_diff=0.5, passed=False), diff=_make_diff(rel_diff=0.5, passed=False),
) )
result: str = format_comparison_rich(record, verbosity="minimal") result: str = format_comparison_rich(record, verbosity="minimal")
@@ -439,7 +439,7 @@ class TestFormatComparisonRichMinimal:
) )
def test_shape_mismatch(self) -> None: def test_shape_mismatch(self) -> None:
record: TensorComparisonRecord = _make_comparison_record( record: ComparisonTensorRecord = _make_comparison_record(
shape_mismatch=True, shape_mismatch=True,
) )
result: str = format_comparison_rich(record, verbosity="minimal") result: str = format_comparison_rich(record, verbosity="minimal")
@@ -450,7 +450,7 @@ class TestFormatComparisonRichMinimal:
) )
def test_no_diff(self) -> None: def test_no_diff(self) -> None:
record: TensorComparisonRecord = _make_comparison_record() record: ComparisonTensorRecord = _make_comparison_record()
result: str = format_comparison_rich(record, verbosity="minimal") result: str = format_comparison_rich(record, verbosity="minimal")
assert result == ("[red]❌[/] [bold red]hidden_states [/]") assert result == ("[red]❌[/] [bold red]hidden_states [/]")
@@ -460,7 +460,7 @@ class TestFormatComparisonRichNormal:
"""format_comparison_rich() with verbosity='normal'.""" """format_comparison_rich() with verbosity='normal'."""
def test_passed(self) -> None: def test_passed(self) -> None:
record: TensorComparisonRecord = _make_comparison_record( record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(rel_diff=1e-4, passed=True), diff=_make_diff(rel_diff=1e-4, passed=True),
) )
result: str = format_comparison_rich(record, verbosity="normal") result: str = format_comparison_rich(record, verbosity="normal")
@@ -477,7 +477,7 @@ class TestFormatComparisonRichNormal:
) )
def test_failed(self) -> None: def test_failed(self) -> None:
record: TensorComparisonRecord = _make_comparison_record( record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff( diff=_make_diff(
rel_diff=0.5, max_abs_diff=1.0, mean_abs_diff=0.3, passed=False 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: def test_shape_mismatch(self) -> None:
record: TensorComparisonRecord = _make_comparison_record( record: ComparisonTensorRecord = _make_comparison_record(
shape_mismatch=True, shape_mismatch=True,
) )
result: str = format_comparison_rich(record, verbosity="normal") result: str = format_comparison_rich(record, verbosity="normal")
@@ -516,7 +516,7 @@ class TestFormatComparisonRichNormal:
) )
def test_with_downcast(self) -> None: 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=_make_diff(rel_diff=0.01, passed=False),
diff_downcast=_make_diff(rel_diff=1e-5, passed=True), diff_downcast=_make_diff(rel_diff=1e-5, passed=True),
downcast_dtype="torch.bfloat16", downcast_dtype="torch.bfloat16",
@@ -543,7 +543,7 @@ class TestFormatComparisonRichNormal:
x=_make_bundle_side_info(num_files=2, dims="b s h(tp) d"), 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"), 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), diff=_make_diff(passed=True),
raw_bundle_info=bundle_info, raw_bundle_info=bundle_info,
) )
@@ -565,7 +565,7 @@ class TestFormatComparisonRichNormal:
def test_with_plan(self) -> None: def test_with_plan(self) -> None:
plan: AlignerPlan = _make_simple_aligner_plan(with_unsharder=True) plan: AlignerPlan = _make_simple_aligner_plan(with_unsharder=True)
record: TensorComparisonRecord = _make_comparison_record( record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(passed=True), diff=_make_diff(passed=True),
traced_plan=_make_traced_plan(plan), traced_plan=_make_traced_plan(plan),
) )
@@ -590,7 +590,7 @@ class TestFormatComparisonRichVerbose:
"""format_comparison_rich() with verbosity='verbose'.""" """format_comparison_rich() with verbosity='verbose'."""
def test_passed_full_detail(self) -> None: 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), diff=_make_diff(rel_diff=1e-4, passed=True),
sample="tensor([0.1, 0.2, ...])", sample="tensor([0.1, 0.2, ...])",
) )
@@ -624,7 +624,7 @@ class TestFormatComparisonRichVerbose:
x=_make_bundle_side_info(num_files=2, with_parallel_info=True), x=_make_bundle_side_info(num_files=2, with_parallel_info=True),
y=_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), diff=_make_diff(passed=True),
raw_bundle_info=bundle_info, raw_bundle_info=bundle_info,
) )
@@ -659,7 +659,7 @@ class TestFormatComparisonRichVerbose:
def test_with_plan_and_traces(self) -> None: def test_with_plan_and_traces(self) -> None:
plan: AlignerPlan = _make_simple_aligner_plan(with_unsharder=True) plan: AlignerPlan = _make_simple_aligner_plan(with_unsharder=True)
record: TensorComparisonRecord = _make_comparison_record( record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(passed=True), diff=_make_diff(passed=True),
traced_plan=_make_traced_plan( traced_plan=_make_traced_plan(
plan, plan,
test/registered/debug_utils/comparator/tensor_comparator/test_types.py
+17 -17
View File
@@ -4,14 +4,14 @@ import sys
import pytest import pytest
from sglang.srt.debug_utils.comparator.output_types import ( from sglang.srt.debug_utils.comparator.output_types import (
ComparisonSkipRecord,
ComparisonTensorRecord,
ConfigRecord, ConfigRecord,
ErrorLog, ErrorLog,
InfoLog, InfoLog,
LogRecord, LogRecord,
ReplicatedCheckResult, ReplicatedCheckResult,
SkipComparisonRecord,
SummaryRecord, SummaryRecord,
TensorComparisonRecord,
parse_record_json, parse_record_json,
) )
from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
@@ -84,7 +84,7 @@ class TestStrictBase:
class TestRecordTypes: class TestRecordTypes:
def test_comparison_record_inherits_tensor_fields(self): def test_comparison_record_inherits_tensor_fields(self):
record = TensorComparisonRecord( record = ComparisonTensorRecord(
name="hidden_states", name="hidden_states",
baseline=_make_tensor_info(), baseline=_make_tensor_info(),
target=_make_tensor_info(), target=_make_tensor_info(),
@@ -93,7 +93,7 @@ class TestRecordTypes:
diff=_make_diff(), diff=_make_diff(),
) )
parsed = json.loads(record.model_dump_json()) parsed = json.loads(record.model_dump_json())
assert parsed["type"] == "comparison" assert parsed["type"] == "comparison_tensor"
assert parsed["name"] == "hidden_states" assert parsed["name"] == "hidden_states"
assert "baseline" in parsed assert "baseline" in parsed
assert "diff" in parsed assert "diff" in parsed
@@ -109,8 +109,8 @@ class TestRecordTypes:
"end_step": 100, "end_step": 100,
} }
), ),
SkipComparisonRecord(name="attn", reason="no_baseline"), ComparisonSkipRecord(name="attn", reason="no_baseline"),
TensorComparisonRecord( ComparisonTensorRecord(
name="mlp", name="mlp",
baseline=_make_tensor_info(), baseline=_make_tensor_info(),
target=_make_tensor_info(), target=_make_tensor_info(),
@@ -149,8 +149,8 @@ def _make_replicated_check(**overrides) -> ReplicatedCheckResult:
class TestWarnings: class TestWarnings:
def test_comparison_record_failed_when_diff_passed_but_errors(self): def test_comparison_record_failed_when_diff_passed_but_errors(self):
"""TensorComparisonRecord with diff.passed=True but errors → category=='failed'.""" """ComparisonTensorRecord with diff.passed=True but errors → category=='failed'."""
record = TensorComparisonRecord( record = ComparisonTensorRecord(
name="hidden", name="hidden",
baseline=_make_tensor_info(), baseline=_make_tensor_info(),
target=_make_tensor_info(), target=_make_tensor_info(),
@@ -162,8 +162,8 @@ class TestWarnings:
assert record.category == "failed" assert record.category == "failed"
def test_skip_record_failed_when_errors(self): def test_skip_record_failed_when_errors(self):
"""SkipComparisonRecord with errors → category=='failed' instead of 'skipped'.""" """ComparisonSkipRecord with errors → category=='failed' instead of 'skipped'."""
record = SkipComparisonRecord( record = ComparisonSkipRecord(
name="x", name="x",
reason="no_baseline", reason="no_baseline",
errors=[ErrorLog(category="test", message="some warning")], errors=[ErrorLog(category="test", message="some warning")],
@@ -171,8 +171,8 @@ class TestWarnings:
assert record.category == "failed" assert record.category == "failed"
def test_replicated_checks_all_passed(self): def test_replicated_checks_all_passed(self):
"""TensorComparisonRecord with all replicated_checks passed → category=='passed'.""" """ComparisonTensorRecord with all replicated_checks passed → category=='passed'."""
record = TensorComparisonRecord( record = ComparisonTensorRecord(
name="hidden", name="hidden",
baseline=_make_tensor_info(), baseline=_make_tensor_info(),
target=_make_tensor_info(), target=_make_tensor_info(),
@@ -184,8 +184,8 @@ class TestWarnings:
assert record.category == "passed" assert record.category == "passed"
def test_replicated_checks_failed_means_record_failed(self): def test_replicated_checks_failed_means_record_failed(self):
"""TensorComparisonRecord with any replicated_check.passed=False → category=='failed'.""" """ComparisonTensorRecord with any replicated_check.passed=False → category=='failed'."""
record = TensorComparisonRecord( record = ComparisonTensorRecord(
name="hidden", name="hidden",
baseline=_make_tensor_info(), baseline=_make_tensor_info(),
target=_make_tensor_info(), target=_make_tensor_info(),
@@ -197,7 +197,7 @@ class TestWarnings:
assert record.category == "failed" assert record.category == "failed"
def test_replicated_check_json_round_trip(self): 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( check = _make_replicated_check(
axis="cp", axis="cp",
group_index=2, group_index=2,
@@ -205,7 +205,7 @@ class TestWarnings:
baseline_index=0, baseline_index=0,
passed=False, passed=False,
) )
record = TensorComparisonRecord( record = ComparisonTensorRecord(
name="mlp", name="mlp",
baseline=_make_tensor_info(), baseline=_make_tensor_info(),
target=_make_tensor_info(), target=_make_tensor_info(),
@@ -216,7 +216,7 @@ class TestWarnings:
) )
restored = parse_record_json(record.model_dump_json()) restored = parse_record_json(record.model_dump_json())
assert isinstance(restored, TensorComparisonRecord) assert isinstance(restored, ComparisonTensorRecord)
assert len(restored.replicated_checks) == 1 assert len(restored.replicated_checks) == 1
restored_check: ReplicatedCheckResult = restored.replicated_checks[0] restored_check: ReplicatedCheckResult = restored.replicated_checks[0]
test/registered/debug_utils/comparator/test_entrypoint.py
+123 -50
View File
@@ -7,6 +7,7 @@ from pathlib import Path
import pytest import pytest
import torch import torch
import sglang.srt.debug_utils.comparator.entrypoint as _entrypoint_module
import sglang.srt.debug_utils.dumper as _dumper_module import sglang.srt.debug_utils.dumper as _dumper_module
from sglang.srt.debug_utils.comparator.entrypoint import ( from sglang.srt.debug_utils.comparator.entrypoint import (
parse_args, parse_args,
@@ -14,14 +15,15 @@ from sglang.srt.debug_utils.comparator.entrypoint import (
) )
from sglang.srt.debug_utils.comparator.output_types import ( from sglang.srt.debug_utils.comparator.output_types import (
AnyRecord, AnyRecord,
ComparisonErrorRecord,
ComparisonNonTensorRecord,
ComparisonSkipRecord,
ComparisonTensorRecord,
ConfigRecord, ConfigRecord,
InfoLog, InfoLog,
LogRecord, LogRecord,
NonTensorComparisonRecord,
ReplicatedCheckResult, ReplicatedCheckResult,
SkipComparisonRecord,
SummaryRecord, SummaryRecord,
TensorComparisonRecord,
_OutputRecord, _OutputRecord,
parse_record_json, parse_record_json,
) )
@@ -39,7 +41,7 @@ class TestEntrypointGroupingRaw:
"""Test `--grouping-skip-keys` empty (raw) scenarios""" """Test `--grouping-skip-keys` empty (raw) scenarios"""
def test_run_basic(self, tmp_path, capsys): 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"]) baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
argv = _make_argv(baseline_path, target_path, preset="raw") argv = _make_argv(baseline_path, target_path, preset="raw")
@@ -54,7 +56,7 @@ class TestEntrypointGroupingRaw:
assert summary.skipped == 0 assert summary.skipped == 0
def test_filter(self, tmp_path, capsys): 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"]) baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
argv = _make_argv(baseline_path, target_path, filter="tensor_a", preset="raw") argv = _make_argv(baseline_path, target_path, filter="tensor_a", preset="raw")
@@ -62,7 +64,7 @@ class TestEntrypointGroupingRaw:
assert len(_get_comparisons(records)) == 1 assert len(_get_comparisons(records)) == 1
def test_no_baseline_skip(self, tmp_path, capsys): 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( baseline_path, target_path = _create_dumps(
tmp_path, tmp_path,
tensor_names=["tensor_a", "tensor_extra"], tensor_names=["tensor_a", "tensor_extra"],
@@ -71,7 +73,7 @@ class TestEntrypointGroupingRaw:
argv = _make_argv(baseline_path, target_path, preset="raw") argv = _make_argv(baseline_path, target_path, preset="raw")
records, _ = _run_and_parse(argv, capsys) 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 len(skips) == 1
assert skips[0].reason == "baseline_load_failed" assert skips[0].reason == "baseline_load_failed"
@@ -100,7 +102,7 @@ class TestEntrypointGroupingRaw:
assert all(isinstance(r, _OutputRecord) for r in records) assert all(isinstance(r, _OutputRecord) for r in records)
def test_comparison_failed(self, tmp_path, capsys): 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) torch.manual_seed(42)
baseline_path = _create_rank_dump( baseline_path = _create_rank_dump(
tmp_path / "baseline", rank=0, name="tensor_a", tensor=torch.randn(10, 10) tmp_path / "baseline", rank=0, name="tensor_a", tensor=torch.randn(10, 10)
@@ -251,7 +253,7 @@ class TestEntrypointGroupingRaw:
assert summary.total == 0 assert summary.total == 0
def test_raw_multi_rank(self, tmp_path, capsys): 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) torch.manual_seed(42)
tensor = torch.randn(4, 4) 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): 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) torch.manual_seed(42)
tensor = torch.randn(4, 8) tensor = torch.randn(4, 8)
@@ -500,7 +502,7 @@ class TestEntrypointGroupingLogical:
) )
records, _ = _run_and_parse(argv, capsys) 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 len(skips) == 1
assert skips[0].reason == expected_reason assert skips[0].reason == expected_reason
@@ -1092,7 +1094,7 @@ class TestEntrypointPerStepMode:
"""Test per-step comparison mode (sglang_dev preset behavior).""" """Test per-step comparison mode (sglang_dev preset behavior)."""
def test_multi_step_per_step_comparison(self, tmp_path, capsys): 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) torch.manual_seed(42)
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"], num_steps=3) baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"], num_steps=3)
argv = _make_argv(baseline_path, target_path, diff_threshold=0.1) 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) assert all(c.baseline.shape == [4, 8] for c in comparisons)
def test_single_step_has_step_field(self, tmp_path, capsys): 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) baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"], num_steps=1)
argv = _make_argv(baseline_path, target_path) argv = _make_argv(baseline_path, target_path)
@@ -1458,7 +1460,7 @@ class TestEntrypointConcatMode:
assert len(comparisons) == 3 assert len(comparisons) == 3
def test_concat_aligner_plan_fields(self, tmp_path, capsys): 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) torch.manual_seed(42)
records = self._run_concat( records = self._run_concat(
@@ -1599,8 +1601,8 @@ class TestEntrypointConcatMode:
) )
records, _ = _run_and_parse(argv, capsys) records, _ = _run_and_parse(argv, capsys)
comparisons: list[TensorComparisonRecord] = _get_comparisons(records) comparisons: list[ComparisonTensorRecord] = _get_comparisons(records)
hidden_comparisons: list[TensorComparisonRecord] = [ hidden_comparisons: list[ComparisonTensorRecord] = [
c for c in comparisons if c.name == "hidden_states" c for c in comparisons if c.name == "hidden_states"
] ]
assert len(hidden_comparisons) >= 1 assert len(hidden_comparisons) >= 1
@@ -2158,7 +2160,7 @@ class TestEntrypointNonTensorValues:
"""Test non-tensor value comparison through the full entrypoint pipeline.""" """Test non-tensor value comparison through the full entrypoint pipeline."""
def test_non_tensor_float_same_value(self, tmp_path: Path, capsys) -> None: 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( baseline_path, target_path = _create_non_tensor_dumps(
tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.125 tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.125
) )
@@ -2177,7 +2179,7 @@ class TestEntrypointNonTensorValues:
assert summary.failed == 0 assert summary.failed == 0
def test_non_tensor_float_different_value(self, tmp_path: Path, capsys) -> None: 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( baseline_path, target_path = _create_non_tensor_dumps(
tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.25 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" assert non_tensors[0].target_type == "dict"
def test_non_tensor_none_value(self, tmp_path: Path, capsys) -> None: 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( baseline_path, target_path = _create_non_tensor_dumps(
tmp_path, name="optional_param", baseline_value=None, target_value=None tmp_path, name="optional_param", baseline_value=None, target_value=None
) )
@@ -2278,7 +2280,7 @@ class TestEntrypointNonTensorValues:
assert non_tensors[0].category == "passed" assert non_tensors[0].category == "passed"
def test_non_tensor_json_roundtrip(self, tmp_path: Path, capsys) -> None: 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( baseline_path, target_path = _create_non_tensor_dumps(
tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.125 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() json_str: str = non_tensors[0].model_dump_json()
roundtripped = parse_record_json(json_str) roundtripped = parse_record_json(json_str)
assert isinstance(roundtripped, NonTensorComparisonRecord) assert isinstance(roundtripped, ComparisonNonTensorRecord)
assert roundtripped.name == "sm_scale" assert roundtripped.name == "sm_scale"
assert roundtripped.values_equal is True assert roundtripped.values_equal is True
@@ -2344,17 +2346,17 @@ class TestEntrypointVisualize:
# --------------------------- Assertion helpers ------------------- # --------------------------- Assertion helpers -------------------
def _get_comparisons(records: list[AnyRecord]) -> list[TensorComparisonRecord]: def _get_comparisons(records: list[AnyRecord]) -> list[ComparisonTensorRecord]:
return [r for r in records if isinstance(r, TensorComparisonRecord)] return [r for r in records if isinstance(r, ComparisonTensorRecord)]
def _get_non_tensors(records: list[AnyRecord]) -> list[NonTensorComparisonRecord]: def _get_non_tensors(records: list[AnyRecord]) -> list[ComparisonNonTensorRecord]:
return [r for r in records if isinstance(r, NonTensorComparisonRecord)] return [r for r in records if isinstance(r, ComparisonNonTensorRecord)]
def _assert_single_comparison_passed( def _assert_single_comparison_passed(
records: list[AnyRecord], records: list[AnyRecord],
) -> TensorComparisonRecord: ) -> ComparisonTensorRecord:
comparisons = _get_comparisons(records) comparisons = _get_comparisons(records)
assert len(comparisons) == 1 assert len(comparisons) == 1
assert comparisons[0].diff is not None assert comparisons[0].diff is not None
@@ -3235,8 +3237,8 @@ class TestEntrypointThdCpZigzag:
) )
records, _ = _run_and_parse(argv, capsys) records, _ = _run_and_parse(argv, capsys)
comparisons: list[TensorComparisonRecord] = _get_comparisons(records) comparisons: list[ComparisonTensorRecord] = _get_comparisons(records)
hidden_comparisons: list[TensorComparisonRecord] = [ hidden_comparisons: list[ComparisonTensorRecord] = [
c for c in comparisons if c.name == "hidden_states" c for c in comparisons if c.name == "hidden_states"
] ]
assert len(hidden_comparisons) >= 1 assert len(hidden_comparisons) >= 1
@@ -3287,8 +3289,8 @@ class TestEntrypointThdCpZigzag:
records, _ = _run_and_parse(argv, capsys) records, _ = _run_and_parse(argv, capsys)
# hidden_states should pass comparison (after unshard + reorder) # hidden_states should pass comparison (after unshard + reorder)
comparisons: list[TensorComparisonRecord] = _get_comparisons(records) comparisons: list[ComparisonTensorRecord] = _get_comparisons(records)
hidden_comparisons: list[TensorComparisonRecord] = [ hidden_comparisons: list[ComparisonTensorRecord] = [
c for c in comparisons if c.name == "hidden_states" c for c in comparisons if c.name == "hidden_states"
] ]
assert len(hidden_comparisons) >= 1 assert len(hidden_comparisons) >= 1
@@ -3355,7 +3357,7 @@ class TestEntrypointDpFilter:
) )
records, _ = _run_and_parse(argv, capsys) records, _ = _run_and_parse(argv, capsys)
comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden" assert comparison.name == "hidden"
def test_dp2_megatron_both_sides(self, tmp_path: Path, capsys) -> None: def test_dp2_megatron_both_sides(self, tmp_path: Path, capsys) -> None:
@@ -3410,7 +3412,7 @@ class TestEntrypointDpFilter:
) )
records, _ = _run_and_parse(argv, capsys) records, _ = _run_and_parse(argv, capsys)
comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden" assert comparison.name == "hidden"
def test_dp2_tp2_sglang(self, tmp_path: Path, capsys) -> None: def test_dp2_tp2_sglang(self, tmp_path: Path, capsys) -> None:
@@ -3458,7 +3460,7 @@ class TestEntrypointDpFilter:
) )
records, _ = _run_and_parse(argv, capsys) records, _ = _run_and_parse(argv, capsys)
comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden" assert comparison.name == "hidden"
def test_dp2_both_nonempty_raises(self, tmp_path: Path, capsys) -> None: def test_dp2_both_nonempty_raises(self, tmp_path: Path, capsys) -> None:
@@ -3496,10 +3498,12 @@ class TestEntrypointDpFilter:
diff_threshold=1e-3, diff_threshold=1e-3,
) )
with pytest.raises( records, exit_code = _run_and_parse(argv, capsys)
AssertionError, match="Expected exactly 1 non-empty dp_rank" errors = [r for r in records if isinstance(r, ComparisonErrorRecord)]
): assert len(errors) == 1
_run_and_parse(argv, capsys) 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: class TestEntrypointDpGroupAlias:
@@ -3542,7 +3546,7 @@ class TestEntrypointDpGroupAlias:
) )
records, _ = _run_and_parse(argv, capsys) records, _ = _run_and_parse(argv, capsys)
comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden" assert comparison.name == "hidden"
def test_dp_alias_via_override_dims(self, tmp_path: Path, capsys) -> None: 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) records, _ = _run_and_parse(argv, capsys)
comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden" assert comparison.name == "hidden"
def test_dp_alias_with_real_alias_group_filters( def test_dp_alias_with_real_alias_group_filters(
@@ -3640,7 +3644,7 @@ class TestEntrypointDpGroupAlias:
) )
records, _ = _run_and_parse(argv, capsys) records, _ = _run_and_parse(argv, capsys)
comparison: TensorComparisonRecord = _assert_single_comparison_passed(records) comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden" assert comparison.name == "hidden"
@@ -3679,7 +3683,7 @@ class TestEntrypointMetaOverride:
records: list[AnyRecord], *, expected_count: int = 1 records: list[AnyRecord], *, expected_count: int = 1
) -> None: ) -> None:
"""Assert that exactly expected_count comparisons exist and all passed.""" """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 len(comparisons) == expected_count
assert all(c.diff is not None and c.diff.passed for c in comparisons) 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) records, _ = _run_and_parse(argv, capsys)
non_tensors: list[NonTensorComparisonRecord] = [ non_tensors: list[ComparisonNonTensorRecord] = [
r for r in records if isinstance(r, NonTensorComparisonRecord) r for r in records if isinstance(r, ComparisonNonTensorRecord)
] ]
assert len(non_tensors) == 1 assert len(non_tensors) == 1
assert non_tensors[0].name == "sm_scale" assert non_tensors[0].name == "sm_scale"
assert non_tensors[0].values_equal assert non_tensors[0].values_equal
comparisons: list[TensorComparisonRecord] = _get_comparisons(records) comparisons: list[ComparisonTensorRecord] = _get_comparisons(records)
assert len(comparisons) == 1 assert len(comparisons) == 1
assert comparisons[0].name == "hidden" assert comparisons[0].name == "hidden"
@@ -4356,12 +4360,9 @@ class TestEntrypointDpAttentionMissingAlias:
assert exit_code == 1 assert exit_code == 1
comparisons: list[TensorComparisonRecord] = _get_comparisons(records) errors = [r for r in records if isinstance(r, ComparisonErrorRecord)]
assert len(comparisons) == 1 assert len(errors) == 1
comparison: TensorComparisonRecord = comparisons[0] assert errors[0].category == "errored"
assert comparison.shape_mismatch is True
assert comparison.diff is None
assert comparison.category == "failed"
class TestEntrypointAutoDescend: class TestEntrypointAutoDescend:
@@ -4460,5 +4461,77 @@ class TestEntrypointAutoDescend:
run(parse_args(argv)) 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__": if __name__ == "__main__":
sys.exit(pytest.main([__file__])) sys.exit(pytest.main([__file__]))
test/registered/debug_utils/comparator/test_manually_verify.py
+6 -6
View File
@@ -208,7 +208,7 @@ class TestPerTokenHeatmapManualVerify:
rows for different tensor names. Colorbar shows log10 scale. rows for different tensor names. Colorbar shows log10 scale.
""" """
from sglang.srt.debug_utils.comparator.output_types import ( from sglang.srt.debug_utils.comparator.output_types import (
TensorComparisonRecord, ComparisonTensorRecord,
) )
from sglang.srt.debug_utils.comparator.per_token_visualizer import ( from sglang.srt.debug_utils.comparator.per_token_visualizer import (
generate_per_token_heatmap, generate_per_token_heatmap,
@@ -222,7 +222,7 @@ class TestPerTokenHeatmapManualVerify:
hidden_dim: int = 128 hidden_dim: int = 128
num_tensors: int = 5 num_tensors: int = 5
records: list[TensorComparisonRecord] = [] records: list[ComparisonTensorRecord] = []
for i in range(num_tensors): for i in range(num_tensors):
baseline: torch.Tensor = torch.randn(seq_len, hidden_dim) baseline: torch.Tensor = torch.randn(seq_len, hidden_dim)
noise_scale: torch.Tensor = torch.linspace( noise_scale: torch.Tensor = torch.linspace(
@@ -237,7 +237,7 @@ class TestPerTokenHeatmapManualVerify:
diff_threshold=1e-3, diff_threshold=1e-3,
seq_dim=0, 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" output_path: Path = tmp_path / "per_token_increasing_diff.png"
result = generate_per_token_heatmap(records=records, output_path=output_path) result = generate_per_token_heatmap(records=records, output_path=output_path)
@@ -253,7 +253,7 @@ class TestPerTokenHeatmapManualVerify:
rest is dark/cold. rest is dark/cold.
""" """
from sglang.srt.debug_utils.comparator.output_types import ( from sglang.srt.debug_utils.comparator.output_types import (
TensorComparisonRecord, ComparisonTensorRecord,
) )
from sglang.srt.debug_utils.comparator.per_token_visualizer import ( from sglang.srt.debug_utils.comparator.per_token_visualizer import (
generate_per_token_heatmap, generate_per_token_heatmap,
@@ -268,7 +268,7 @@ class TestPerTokenHeatmapManualVerify:
spike_pos: int = 32 spike_pos: int = 32
num_tensors: int = 4 num_tensors: int = 4
records: list[TensorComparisonRecord] = [] records: list[ComparisonTensorRecord] = []
for i in range(num_tensors): for i in range(num_tensors):
baseline: torch.Tensor = torch.randn(seq_len, hidden_dim) baseline: torch.Tensor = torch.randn(seq_len, hidden_dim)
target: torch.Tensor = baseline.clone() target: torch.Tensor = baseline.clone()
@@ -281,7 +281,7 @@ class TestPerTokenHeatmapManualVerify:
diff_threshold=1e-3, diff_threshold=1e-3,
seq_dim=0, 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" output_path: Path = tmp_path / "per_token_single_spike.png"
result = generate_per_token_heatmap(records=records, output_path=output_path) result = generate_per_token_heatmap(records=records, output_path=output_path)
test/registered/debug_utils/comparator/test_model_validation.py
+95 -43
View File
@@ -4,6 +4,12 @@ import sys
import pytest import pytest
from pydantic import ValidationError 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 ( from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
AlignerPerStepPlan, AlignerPerStepPlan,
AlignerPlan, 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.dims_spec import ParallelAxis, TokenLayout
from sglang.srt.debug_utils.comparator.output_types import ( from sglang.srt.debug_utils.comparator.output_types import (
ComparisonErrorRecord,
ComparisonNonTensorRecord,
ComparisonSkipRecord,
ComparisonTensorRecord,
ErrorLog, ErrorLog,
NonTensorComparisonRecord,
SkipComparisonRecord,
SummaryRecord, SummaryRecord,
TensorComparisonRecord,
parse_record_json, parse_record_json,
) )
from sglang.srt.debug_utils.comparator.tensor_comparator.types import ( from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
@@ -156,6 +163,14 @@ class TestSummaryRecord:
with pytest.raises(ValidationError, match="total=10"): with pytest.raises(ValidationError, match="total=10"):
SummaryRecord(total=10, passed=5, failed=2, skipped=1) 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: class TestAxisInfo:
def test_valid(self): def test_valid(self):
@@ -208,9 +223,9 @@ def _make_comparison_record(
*, *,
diff: DiffInfo | None, diff: DiffInfo | None,
errors: list | None = None, errors: list | None = None,
) -> TensorComparisonRecord: ) -> ComparisonTensorRecord:
ti: TensorInfo = _make_tensor_info() ti: TensorInfo = _make_tensor_info()
return TensorComparisonRecord( return ComparisonTensorRecord(
name="t", name="t",
baseline=ti, baseline=ti,
target=ti, target=ti,
@@ -223,7 +238,7 @@ def _make_comparison_record(
class TestOutputRecordCategories: class TestOutputRecordCategories:
def test_skip_record_with_errors_is_failed(self) -> None: def test_skip_record_with_errors_is_failed(self) -> None:
record = SkipComparisonRecord( record = ComparisonSkipRecord(
name="t", name="t",
reason="test", reason="test",
errors=[ErrorLog(category="c", message="m")], errors=[ErrorLog(category="c", message="m")],
@@ -231,28 +246,28 @@ class TestOutputRecordCategories:
assert record.category == "failed" assert record.category == "failed"
def test_skip_record_no_warnings_is_skipped(self) -> None: 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" assert record.category == "skipped"
def test_comparison_record_diff_none_is_failed(self) -> None: 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" assert record.category == "failed"
def test_comparison_record_passed_with_errors_is_failed(self) -> None: 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), diff=_make_diff_info(passed=True),
errors=[ErrorLog(category="c", message="m")], errors=[ErrorLog(category="c", message="m")],
) )
assert record.category == "failed" assert record.category == "failed"
def test_comparison_record_passed_no_warnings_is_passed(self) -> None: 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), diff=_make_diff_info(passed=True),
) )
assert record.category == "passed" assert record.category == "passed"
def test_non_tensor_record_equal_is_passed(self) -> None: def test_non_tensor_record_equal_is_passed(self) -> None:
record = NonTensorComparisonRecord( record = ComparisonNonTensorRecord(
name="sm_scale", name="sm_scale",
baseline_value="0.125", baseline_value="0.125",
target_value="0.125", target_value="0.125",
@@ -263,7 +278,7 @@ class TestOutputRecordCategories:
assert record.category == "passed" assert record.category == "passed"
def test_non_tensor_record_different_is_failed(self) -> None: def test_non_tensor_record_different_is_failed(self) -> None:
record = NonTensorComparisonRecord( record = ComparisonNonTensorRecord(
name="sm_scale", name="sm_scale",
baseline_value="0.125", baseline_value="0.125",
target_value="0.25", target_value="0.25",
@@ -274,7 +289,7 @@ class TestOutputRecordCategories:
assert record.category == "failed" assert record.category == "failed"
def test_non_tensor_record_with_errors_is_failed(self) -> None: def test_non_tensor_record_with_errors_is_failed(self) -> None:
record = NonTensorComparisonRecord( record = ComparisonNonTensorRecord(
name="sm_scale", name="sm_scale",
baseline_value="0.125", baseline_value="0.125",
target_value="0.125", target_value="0.125",
@@ -286,7 +301,7 @@ class TestOutputRecordCategories:
assert record.category == "failed" assert record.category == "failed"
def test_non_tensor_record_json_roundtrip(self) -> None: def test_non_tensor_record_json_roundtrip(self) -> None:
record = NonTensorComparisonRecord( record = ComparisonNonTensorRecord(
name="sm_scale", name="sm_scale",
baseline_value="0.125", baseline_value="0.125",
target_value="0.25", target_value="0.25",
@@ -296,14 +311,14 @@ class TestOutputRecordCategories:
) )
json_str: str = record.model_dump_json() json_str: str = record.model_dump_json()
roundtripped = parse_record_json(json_str) roundtripped = parse_record_json(json_str)
assert isinstance(roundtripped, NonTensorComparisonRecord) assert isinstance(roundtripped, ComparisonNonTensorRecord)
assert roundtripped.name == "sm_scale" assert roundtripped.name == "sm_scale"
assert roundtripped.values_equal is False assert roundtripped.values_equal is False
assert roundtripped.baseline_value == "0.125" assert roundtripped.baseline_value == "0.125"
assert roundtripped.target_value == "0.25" assert roundtripped.target_value == "0.25"
def test_non_tensor_record_text_format_equal(self) -> None: def test_non_tensor_record_text_format_equal(self) -> None:
record = NonTensorComparisonRecord( record = ComparisonNonTensorRecord(
name="sm_scale", name="sm_scale",
baseline_value="0.125", baseline_value="0.125",
target_value="0.125", target_value="0.125",
@@ -316,7 +331,7 @@ class TestOutputRecordCategories:
assert "[equal]" in text assert "[equal]" in text
def test_non_tensor_record_text_format_different(self) -> None: def test_non_tensor_record_text_format_different(self) -> None:
record = NonTensorComparisonRecord( record = ComparisonNonTensorRecord(
name="sm_scale", name="sm_scale",
baseline_value="0.125", baseline_value="0.125",
target_value="0.25", target_value="0.25",
@@ -328,14 +343,38 @@ class TestOutputRecordCategories:
assert "baseline" in text assert "baseline" in text
assert "target" 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( unsharder = UnsharderPlan(
axis=ParallelAxis.TP, axis=ParallelAxis.TP,
params=ConcatParams(dim_name="h"), params=ConcatParams(dim_name="h"),
groups=[[0, 1]], groups=[[0, 1]],
) )
return AlignerPlan( plan = AlignerPlan(
per_step_plans=Pair( per_step_plans=Pair(
x=[ x=[
AlignerPerStepPlan( 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: class TestAlignerPlanInComparisonTensorRecord:
def test_comparison_record_with_aligner_plan(self) -> None: def test_comparison_record_with_traced_plan(self) -> None:
plan: AlignerPlan = _make_aligner_plan() traced_plan: TracedAlignerPlan = _make_traced_aligner_plan()
record: TensorComparisonRecord = _make_comparison_record( record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff_info(passed=True), 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})
assert record_with_plan.aligner_plan is not None assert record_with_plan.traced_plan is not None
assert record_with_plan.aligner_plan.per_step_plans.x[0].step == 0 assert record_with_plan.traced_plan.per_side.x.step_plans[0].step == 0
def test_aligner_plan_json_roundtrip(self) -> None: def test_traced_plan_json_roundtrip(self) -> None:
plan: AlignerPlan = _make_aligner_plan() traced_plan: TracedAlignerPlan = _make_traced_aligner_plan()
record: TensorComparisonRecord = _make_comparison_record( record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff_info(passed=True), 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() json_str: str = record_with_plan.model_dump_json()
parsed = json.loads(json_str) parsed = json.loads(json_str)
assert "aligner_plan" in parsed assert "traced_plan" in parsed
assert ( 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" == "unsharder"
) )
roundtripped: TensorComparisonRecord = parse_record_json(json_str) roundtripped: ComparisonTensorRecord = parse_record_json(json_str)
assert roundtripped.aligner_plan is not None assert roundtripped.traced_plan is not None
assert ( 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" == "unsharder"
) )
def test_comparison_record_without_aligner_plan(self) -> None: def test_comparison_record_without_traced_plan(self) -> None:
record: TensorComparisonRecord = _make_comparison_record( record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff_info(passed=True), diff=_make_diff_info(passed=True),
) )
json_str: str = record.model_dump_json() json_str: str = record.model_dump_json()
roundtripped: TensorComparisonRecord = parse_record_json(json_str) roundtripped: ComparisonTensorRecord = parse_record_json(json_str)
assert roundtripped.aligner_plan is None assert roundtripped.traced_plan is None
def test_aligner_plan_text_format(self) -> None: def test_traced_plan_text_format(self) -> None:
plan: AlignerPlan = _make_aligner_plan() traced_plan: TracedAlignerPlan = _make_traced_aligner_plan()
record: TensorComparisonRecord = _make_comparison_record( record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff_info(passed=True), 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() text: str = record_with_plan.to_text()
assert "Aligner Plan:" in text assert "Aligner Plan:" in text
test/registered/debug_utils/comparator/test_output_types.py
+26 -26
View File
@@ -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.dims_spec import ParallelAxis, TokenLayout
from sglang.srt.debug_utils.comparator.output_types import ( from sglang.srt.debug_utils.comparator.output_types import (
ComparisonNonTensorRecord,
ComparisonSkipRecord,
ComparisonTensorRecord,
ConfigRecord, ConfigRecord,
ErrorLog, ErrorLog,
InfoLog, InfoLog,
LogRecord, LogRecord,
NonTensorComparisonRecord,
RecordLocation, RecordLocation,
SkipComparisonRecord,
SummaryRecord, SummaryRecord,
TensorComparisonRecord,
_format_aligner_plan, _format_aligner_plan,
_split_logs, _split_logs,
) )
@@ -150,20 +150,20 @@ class TestConfigRecord:
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# SkipComparisonRecord # ComparisonSkipRecord
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
class TestSkipComparisonRecord: class TestComparisonSkipRecord:
def test_format_body_no_step(self) -> None: def test_format_body_no_step(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord( record: ComparisonSkipRecord = ComparisonSkipRecord(
name="layer.weight", name="layer.weight",
reason="zero-dim tensor", reason="zero-dim tensor",
) )
assert record._format_body() == "Skip: layer.weight (zero-dim tensor)" assert record._format_body() == "Skip: layer.weight (zero-dim tensor)"
def test_format_body_with_step(self) -> None: def test_format_body_with_step(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord( record: ComparisonSkipRecord = ComparisonSkipRecord(
name="layer.weight", name="layer.weight",
reason="scalar", reason="scalar",
location=RecordLocation(step=3), location=RecordLocation(step=3),
@@ -171,7 +171,7 @@ class TestSkipComparisonRecord:
assert record._format_body() == "Skip: layer.weight (step=3) (scalar)" assert record._format_body() == "Skip: layer.weight (step=3) (scalar)"
def test_format_rich_body(self) -> None: def test_format_rich_body(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord( record: ComparisonSkipRecord = ComparisonSkipRecord(
name="attn.qkv", name="attn.qkv",
reason="no baseline", reason="no baseline",
) )
@@ -179,14 +179,14 @@ class TestSkipComparisonRecord:
assert body == "[dim]⊘ attn.qkv ── skipped (no baseline)[/]" assert body == "[dim]⊘ attn.qkv ── skipped (no baseline)[/]"
def test_category_skipped(self) -> None: def test_category_skipped(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord( record: ComparisonSkipRecord = ComparisonSkipRecord(
name="x", name="x",
reason="r", reason="r",
) )
assert record.category == "skipped" assert record.category == "skipped"
def test_category_failed(self) -> None: def test_category_failed(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord( record: ComparisonSkipRecord = ComparisonSkipRecord(
name="x", name="x",
reason="r", reason="r",
errors=[ErrorLog(category="e", message="boom")], 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: def test_format_body_equal(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord( record: ComparisonNonTensorRecord = ComparisonNonTensorRecord(
name="config.lr", name="config.lr",
baseline_value="0.001", baseline_value="0.001",
target_value="0.001", target_value="0.001",
@@ -212,7 +212,7 @@ class TestNonTensorComparisonRecord:
assert record._format_body() == "NonTensor: config.lr = 0.001 (float) [equal]" assert record._format_body() == "NonTensor: config.lr = 0.001 (float) [equal]"
def test_format_body_not_equal(self) -> None: def test_format_body_not_equal(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord( record: ComparisonNonTensorRecord = ComparisonNonTensorRecord(
name="config.lr", name="config.lr",
baseline_value="0.001", baseline_value="0.001",
target_value="0.01", target_value="0.01",
@@ -227,7 +227,7 @@ class TestNonTensorComparisonRecord:
) )
def test_format_rich_body_equal(self) -> None: def test_format_rich_body_equal(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord( record: ComparisonNonTensorRecord = ComparisonNonTensorRecord(
name="config.lr", name="config.lr",
baseline_value="0.001", baseline_value="0.001",
target_value="0.001", target_value="0.001",
@@ -238,7 +238,7 @@ class TestNonTensorComparisonRecord:
assert record._format_rich_body() == ("═ config.lr = 0.001 (float) [green]✓[/]") assert record._format_rich_body() == ("═ config.lr = 0.001 (float) [green]✓[/]")
def test_format_rich_body_not_equal(self) -> None: def test_format_rich_body_not_equal(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord( record: ComparisonNonTensorRecord = ComparisonNonTensorRecord(
name="config.lr", name="config.lr",
baseline_value="0.001", baseline_value="0.001",
target_value="0.01", target_value="0.01",
@@ -253,7 +253,7 @@ class TestNonTensorComparisonRecord:
) )
def test_with_step(self) -> None: def test_with_step(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord( record: ComparisonNonTensorRecord = ComparisonNonTensorRecord(
name="bias", name="bias",
baseline_value="True", baseline_value="True",
target_value="True", target_value="True",
@@ -265,7 +265,7 @@ class TestNonTensorComparisonRecord:
assert "(step=5)" in record._format_body() assert "(step=5)" in record._format_body()
def test_category(self) -> None: def test_category(self) -> None:
passed: NonTensorComparisonRecord = NonTensorComparisonRecord( passed: ComparisonNonTensorRecord = ComparisonNonTensorRecord(
name="x", name="x",
baseline_value="1", baseline_value="1",
target_value="1", target_value="1",
@@ -273,7 +273,7 @@ class TestNonTensorComparisonRecord:
target_type="int", target_type="int",
values_equal=True, values_equal=True,
) )
failed: NonTensorComparisonRecord = NonTensorComparisonRecord( failed: ComparisonNonTensorRecord = ComparisonNonTensorRecord(
name="x", name="x",
baseline_value="1", baseline_value="1",
target_value="2", target_value="2",
@@ -328,13 +328,13 @@ class TestSummaryRecord:
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# TensorComparisonRecord._format_body # ComparisonTensorRecord._format_body
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
class TestTensorComparisonRecordFormatBody: class TestComparisonTensorRecordFormatBody:
def test_basic(self) -> None: def test_basic(self) -> None:
record: TensorComparisonRecord = TensorComparisonRecord( record: ComparisonTensorRecord = ComparisonTensorRecord(
name="hidden", name="hidden",
baseline=_make_tensor_info(), baseline=_make_tensor_info(),
target=_make_tensor_info(), target=_make_tensor_info(),
@@ -365,7 +365,7 @@ class TestTensorComparisonRecordFormatBody:
def test_with_replicated_checks(self) -> None: def test_with_replicated_checks(self) -> None:
from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
record: TensorComparisonRecord = TensorComparisonRecord( record: ComparisonTensorRecord = ComparisonTensorRecord(
name="hidden", name="hidden",
baseline=_make_tensor_info(), baseline=_make_tensor_info(),
target=_make_tensor_info(), target=_make_tensor_info(),
@@ -420,7 +420,7 @@ class TestTensorComparisonRecordFormatBody:
y=TracedSidePlan(step_plans=[]), y=TracedSidePlan(step_plans=[]),
), ),
) )
record: TensorComparisonRecord = TensorComparisonRecord( record: ComparisonTensorRecord = ComparisonTensorRecord(
name="hidden", name="hidden",
baseline=_make_tensor_info(), baseline=_make_tensor_info(),
target=_make_tensor_info(), target=_make_tensor_info(),
@@ -453,7 +453,7 @@ class TestTensorComparisonRecordFormatBody:
) )
def test_with_step(self) -> None: def test_with_step(self) -> None:
record: TensorComparisonRecord = TensorComparisonRecord( record: ComparisonTensorRecord = ComparisonTensorRecord(
name="hidden", name="hidden",
baseline=_make_tensor_info(), baseline=_make_tensor_info(),
target=_make_tensor_info(), target=_make_tensor_info(),
@@ -679,7 +679,7 @@ class TestOutputRecordLogAttachment:
assert text == "Config: {'a': 1}\n ✗ err1\n ℹ note1" assert text == "Config: {'a': 1}\n ✗ err1\n ℹ note1"
def test_to_rich_string_body(self) -> None: def test_to_rich_string_body(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord( record: ComparisonSkipRecord = ComparisonSkipRecord(
name="x", name="x",
reason="r", reason="r",
errors=[ErrorLog(category="e", message="oops")], errors=[ErrorLog(category="e", message="oops")],
test/registered/debug_utils/comparator/test_per_token_visualizer.py
+7 -7
View File
@@ -9,7 +9,7 @@ from pathlib import Path
import pytest import pytest
import torch 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 ( from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
compare_tensor_pair, compare_tensor_pair,
) )
@@ -31,8 +31,8 @@ def _make_comparison_record(
baseline: torch.Tensor, baseline: torch.Tensor,
target: torch.Tensor, target: torch.Tensor,
seq_dim: int = 0, seq_dim: int = 0,
) -> TensorComparisonRecord: ) -> ComparisonTensorRecord:
"""Build a TensorComparisonRecord with per-token data from raw tensors.""" """Build a ComparisonTensorRecord with per-token data from raw tensors."""
info = compare_tensor_pair( info = compare_tensor_pair(
x_baseline=baseline, x_baseline=baseline,
x_target=target, x_target=target,
@@ -40,7 +40,7 @@ def _make_comparison_record(
diff_threshold=1e-3, diff_threshold=1e-3,
seq_dim=seq_dim, seq_dim=seq_dim,
) )
return TensorComparisonRecord(**info.model_dump()) return ComparisonTensorRecord(**info.model_dump())
class TestPerTokenVisualizer: class TestPerTokenVisualizer:
@@ -68,7 +68,7 @@ class TestPerTokenVisualizer:
name="no_per_token", name="no_per_token",
diff_threshold=1e-3, diff_threshold=1e-3,
) )
record = TensorComparisonRecord(**info.model_dump()) record = ComparisonTensorRecord(**info.model_dump())
output_path: Path = tmp_path / "no_data.png" output_path: Path = tmp_path / "no_data.png"
result = generate_per_token_heatmap(records=[record], output_path=output_path) result = generate_per_token_heatmap(records=[record], output_path=output_path)
@@ -82,7 +82,7 @@ class TestPerTokenVisualizer:
) )
torch.manual_seed(42) torch.manual_seed(42)
records: list[TensorComparisonRecord] = [ records: list[ComparisonTensorRecord] = [
_make_comparison_record( _make_comparison_record(
name=f"tensor_{i}", name=f"tensor_{i}",
baseline=torch.randn(16, 32), baseline=torch.randn(16, 32),
@@ -108,7 +108,7 @@ class TestPerTokenVisualizer:
) )
torch.manual_seed(42) torch.manual_seed(42)
records: list[TensorComparisonRecord] = [ records: list[ComparisonTensorRecord] = [
_make_comparison_record( _make_comparison_record(
name="short", name="short",
baseline=torch.randn(4, 8), baseline=torch.randn(4, 8),
test/registered/debug_utils/comparator/test_utils.py
+30
View File
@@ -410,6 +410,36 @@ class TestComputeExitCode:
== 1 == 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: def _make_pt(directory: Path) -> None:
directory.mkdir(parents=True, exist_ok=True) directory.mkdir(parents=True, exist_ok=True)
test/registered/debug_utils/comparator/testing_helpers.py
+6
View File
@@ -4,6 +4,12 @@ from __future__ import annotations
from typing import Optional 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 ( from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo, DiffInfo,
TensorInfo, TensorInfo,