chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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

4
python/sglang/srt/debug_utils/comparator/__init__.py
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
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()

28
python/sglang/srt/debug_utils/comparator/bundle_comparator.py
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 (
BundleFileInfo,
BundleSideInfo,
ComparisonNonTensorRecord,
ComparisonSkipRecord,
ComparisonTensorRecord,
ErrorLog,
NonTensorComparisonRecord,
SkipComparisonRecord,
TensorComparisonRecord,
_split_logs,
)
from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
@@ -51,7 +51,7 @@ def _collect_bundle_side_info(
metas: list[dict[str, Any]],
) -> BundleSideInfo:
from sglang.srt.debug_utils.comparator.display import (
PARALLEL_INFO_KEYS,
_PARALLEL_INFO_KEYS,
extract_parallel_info,
)
@@ -61,7 +61,7 @@ def _collect_bundle_side_info(
tensor: torch.Tensor = item.value
parallel_info: dict[str, str] = {}
for key in PARALLEL_INFO_KEYS:
for key in _PARALLEL_INFO_KEYS:
extract_parallel_info(row_data=parallel_info, info=meta.get(key, {}))
files.append(
@@ -91,7 +91,7 @@ def compare_bundle_pair(
viz_output_dir: Optional[Path] = None,
compute_per_token: bool = False,
meta_overrider: Optional[MetaOverrider] = None,
) -> Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord]:
) -> Union[ComparisonTensorRecord, ComparisonSkipRecord, ComparisonNonTensorRecord]:
with log_sink.context() as collected_logs:
result = _compare_bundle_pair_inner(
name=name,
@@ -124,7 +124,7 @@ def _compare_bundle_pair_inner(
viz_output_dir: Optional[Path] = None,
compute_per_token: bool = False,
meta_overrider: Optional[MetaOverrider] = None,
) -> Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord]:
) -> Union[ComparisonTensorRecord, ComparisonSkipRecord, ComparisonNonTensorRecord]:
# 1. Load all successfully loaded values
all_pair: Pair[list[ValueWithMeta]] = Pair(
x=_load_all_values(filenames=filenames_pair.x, base_path=dir_pair.x),
@@ -133,7 +133,7 @@ def _compare_bundle_pair_inner(
if not all_pair.x or not all_pair.y:
reason = "baseline_load_failed" if not all_pair.x else "target_load_failed"
return SkipComparisonRecord(name=name, reason=reason)
return ComparisonSkipRecord(name=name, reason=reason)
# 1b. Dims override: patch meta["dims"] before DP filter reads it
# (--override-dims may add ``# dp:=moe_dp``, so it must run first)
@@ -203,10 +203,10 @@ def _compare_bundle_pair_tensor_type(
),
viz_output_dir: Optional[Path] = None,
compute_per_token: bool = False,
) -> Union[TensorComparisonRecord, SkipComparisonRecord]:
) -> Union[ComparisonTensorRecord, ComparisonSkipRecord]:
if not valid_pair.x or not valid_pair.y:
reason = "baseline_load_failed" if not valid_pair.x else "target_load_failed"
return SkipComparisonRecord(name=name, reason=reason)
return ComparisonSkipRecord(name=name, reason=reason)
# Plan (meta only, no tensor)
metas_pair: Pair[list[dict[str, Any]]] = valid_pair.map(
@@ -245,7 +245,7 @@ def _compare_bundle_pair_tensor_type(
assert aligner_result.failed_side_xy is not None
side_name: str = _FAILED_SIDE_MAP[aligner_result.failed_side_xy]
reason: str = f"{side_name}_load_failed"
return SkipComparisonRecord(name=name, reason=reason)
return ComparisonSkipRecord(name=name, reason=reason)
# Resolve seq_dim for per-token computation
seq_dim: Optional[int] = (
@@ -263,7 +263,7 @@ def _compare_bundle_pair_tensor_type(
diff_threshold=diff_threshold,
seq_dim=seq_dim,
)
record = TensorComparisonRecord(
record = ComparisonTensorRecord(
**info.model_dump(),
traced_plan=aligner_result.traced_plan,
replicated_checks=replicated_checks,
@@ -331,7 +331,7 @@ def _compare_bundle_pair_non_tensor_type(
*,
name: str,
value_pair: Pair[list[ValueWithMeta]],
) -> NonTensorComparisonRecord:
) -> ComparisonNonTensorRecord:
baseline_value: Any = value_pair.x[0].value
target_value: Any = value_pair.y[0].value
@@ -340,7 +340,7 @@ def _compare_bundle_pair_non_tensor_type(
except Exception:
values_equal = False
return NonTensorComparisonRecord(
return ComparisonNonTensorRecord(
name=name,
baseline_value=repr(baseline_value),
target_value=repr(target_value),

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

@@ -3,13 +3,9 @@ from __future__ import annotations
from collections import defaultdict
from io import StringIO
from pathlib import Path
from typing import TYPE_CHECKING, Any, Optional
from typing import Any, Optional
import polars as pl
import rich.table
if TYPE_CHECKING:
from rich.table import Table
from sglang.srt.debug_utils.comparator.output_types import (
InputIdsRecord,
@@ -43,45 +39,19 @@ def emit_display_records(
def _render_polars_as_text(df: pl.DataFrame, *, title: Optional[str] = None) -> str:
from rich.console import Console
from rich.table import Table
table = _build_rich_table(df, title=title)
table = Table(title=title)
for col in df.columns:
table.add_column(col)
for row in df.iter_rows():
table.add_row(*[str(v) for v in row])
buf = StringIO()
Console(file=buf, force_terminal=False, width=200).print(table)
return buf.getvalue().rstrip("\n")
def _render_polars_as_rich_table(
df: pl.DataFrame, *, title: Optional[str] = None
) -> "Table":
return _build_rich_table(df, title=title)
def _build_rich_table(df: pl.DataFrame, *, title: Optional[str] = None) -> "Table":
from rich.table import Table
table = Table(title=title)
for col in df.columns:
table.add_column(col)
for row in df.iter_rows():
table.add_row(*[str(v) for v in row])
return table
def _render_polars_as_rich_table(
df: pl.DataFrame, *, title: Optional[str] = None
) -> "rich.table.Table":
from rich.table import Table
table = Table(title=title)
for col in df.columns:
table.add_column(col)
for row in df.iter_rows():
table.add_row(*[str(v) for v in row])
return table
def _collect_rank_info(
df: pl.DataFrame, dump_dir: Path
) -> Optional[list[dict[str, Any]]]:
@@ -99,7 +69,7 @@ def _collect_rank_info(
row_data: dict[str, Any] = {"rank": row["rank"]}
for key in PARALLEL_INFO_KEYS:
extract_parallel_info(row_data=row_data, info=meta.get(key, {}))
_extract_parallel_info(row_data=row_data, info=meta.get(key, {}))
table_rows.append(row_data)
return table_rows or None
@@ -149,7 +119,7 @@ def _collect_input_ids_and_positions(
return table_rows or None
def extract_parallel_info(row_data: dict[str, Any], info: dict[str, Any]) -> None:
def _extract_parallel_info(row_data: dict[str, Any], info: dict[str, Any]) -> None:
if not info or info.get("error"):
return

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

@@ -2,6 +2,7 @@ from __future__ import annotations
import argparse
import sys
import traceback as _traceback_module
from pathlib import Path
from typing import Any, Iterator, Optional, Union
@@ -25,12 +26,13 @@ from sglang.srt.debug_utils.comparator.bundle_matcher import (
from sglang.srt.debug_utils.comparator.display import emit_display_records
from sglang.srt.debug_utils.comparator.meta_overrider import MetaOverrider
from sglang.srt.debug_utils.comparator.output_types import (
ComparisonErrorRecord,
ComparisonNonTensorRecord,
ComparisonSkipRecord,
ComparisonTensorRecord,
ConfigRecord,
NonTensorComparisonRecord,
RecordLocation,
SkipComparisonRecord,
SummaryRecord,
TensorComparisonRecord,
)
from sglang.srt.debug_utils.comparator.per_token_visualizer import (
generate_per_token_heatmap,
@@ -83,12 +85,6 @@ def run(args: argparse.Namespace) -> int:
verbosity=args.verbosity,
)
report_path: Optional[Path] = _resolve_report_path(
target_path=dir_pair.y,
report_path_arg=args.report_path,
)
report_sink.configure(output_format=args.output_format, report_path=report_path)
try:
report_sink.add(ConfigRecord(config=vars(args)))
@@ -142,9 +138,11 @@ def run(args: argparse.Namespace) -> int:
compute_per_token=visualize_per_token is not None,
meta_overrider=meta_overrider,
)
summary, skipped_names, failed_names = _consume_comparison_records(
comparison_records=comparison_records,
visualize_per_token=visualize_per_token,
summary, skipped_names, failed_names, errored_names = (
_consume_comparison_records(
comparison_records=comparison_records,
visualize_per_token=visualize_per_token,
)
)
return compute_exit_code(
summary,
@@ -152,6 +150,7 @@ def run(args: argparse.Namespace) -> int:
skipped_names=skipped_names,
allow_failed_pattern=args.allow_failed_pattern,
failed_names=failed_names,
errored_names=errored_names,
)
finally:
report_sink.close()
@@ -219,7 +218,12 @@ def _compare_bundle_pairs(
compute_per_token: bool = False,
meta_overrider: Optional[MetaOverrider] = None,
) -> Iterator[
Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord]
Union[
ComparisonTensorRecord,
ComparisonSkipRecord,
ComparisonNonTensorRecord,
ComparisonErrorRecord,
]
]:
for bundle_info_pair in bundle_info_pairs:
if not bundle_info_pair.y:
@@ -229,20 +233,32 @@ def _compare_bundle_pairs(
filenames_pair: Pair[list[str]] = bundle_info_pair.map(
lambda infos: [info.filename for info in infos]
)
record: Union[
TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord
] = compare_bundle_pair(
name=name,
filenames_pair=filenames_pair,
dir_pair=dir_pair,
token_aligner_mode=token_aligner_mode,
token_aligner_plan=token_aligner_plan,
diff_threshold=diff_threshold,
thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
viz_output_dir=viz_output_dir,
compute_per_token=compute_per_token,
meta_overrider=meta_overrider,
)
ComparisonTensorRecord,
ComparisonSkipRecord,
ComparisonNonTensorRecord,
ComparisonErrorRecord,
]
try:
record = compare_bundle_pair(
name=name,
filenames_pair=filenames_pair,
dir_pair=dir_pair,
token_aligner_mode=token_aligner_mode,
token_aligner_plan=token_aligner_plan,
diff_threshold=diff_threshold,
thd_seq_lens_by_step_pair=thd_seq_lens_by_step_pair,
viz_output_dir=viz_output_dir,
compute_per_token=compute_per_token,
meta_overrider=meta_overrider,
)
except Exception as exc:
record = ComparisonErrorRecord(
name=name,
exception_type=type(exc).__name__,
traceback_str=_traceback_module.format_exc(),
)
target_steps: set[int] = {info.step for info in bundle_info_pair.y}
step: Optional[int] = target_steps.pop() if len(target_steps) == 1 else None
@@ -255,24 +271,32 @@ def _compare_bundle_pairs(
def _consume_comparison_records(
*,
comparison_records: Iterator[
Union[TensorComparisonRecord, SkipComparisonRecord, NonTensorComparisonRecord]
Union[
ComparisonTensorRecord,
ComparisonSkipRecord,
ComparisonNonTensorRecord,
ComparisonErrorRecord,
]
],
visualize_per_token: Optional[Path] = None,
) -> tuple[SummaryRecord, list[str], list[str]]:
counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0}
collected_comparisons: list[TensorComparisonRecord] = []
) -> tuple[SummaryRecord, list[str], list[str], list[str]]:
counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0, "errored": 0}
collected_comparisons: list[ComparisonTensorRecord] = []
skipped_names: list[str] = []
failed_names: list[str] = []
errored_names: list[str] = []
for record in comparison_records:
counts[record.category] += 1
report_sink.add(record)
if isinstance(record, SkipComparisonRecord) and record.category == "skipped":
if isinstance(record, ComparisonSkipRecord) and record.category == "skipped":
skipped_names.append(record.name)
if record.category == "failed":
failed_names.append(record.name)
if isinstance(record, ComparisonErrorRecord):
errored_names.append(record.name)
if visualize_per_token is not None and isinstance(
record, TensorComparisonRecord
record, ComparisonTensorRecord
):
collected_comparisons.append(record)
@@ -285,7 +309,7 @@ def _consume_comparison_records(
output_path=visualize_per_token,
)
return summary, skipped_names, failed_names
return summary, skipped_names, failed_names, errored_names
def parse_args(argv: list[str]) -> argparse.Namespace:

56
python/sglang/srt/debug_utils/comparator/output_formatter.py
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.output_types import (
ComparisonErrorRecord,
ComparisonNonTensorRecord,
ComparisonSkipRecord,
ComparisonTensorRecord,
ConfigRecord,
ErrorLog,
InfoLog,
LogRecord,
NonTensorComparisonRecord,
SkipComparisonRecord,
SummaryRecord,
TensorComparisonRecord,
_OutputRecord,
_TableRecord,
)
@@ -111,15 +112,15 @@ def _format_config_rich_body(
return Panel("\n".join(lines), title="Comparator Config", border_style="cyan")
# ── SkipComparisonRecord ─────────────────────────────────────────────
# ── ComparisonSkipRecord ─────────────────────────────────────────────
def _format_skip_body(record: SkipComparisonRecord) -> str:
def _format_skip_body(record: ComparisonSkipRecord) -> str:
return f"Skip: {record.name}{record._format_location_suffix()} ({record.reason})"
def _format_skip_rich_body(
record: SkipComparisonRecord, verbosity: Verbosity = "normal"
record: ComparisonSkipRecord, verbosity: Verbosity = "normal"
) -> RenderableType:
suffix: str = record._format_location_suffix()
return (
@@ -127,6 +128,30 @@ def _format_skip_rich_body(
)
# ── ComparisonErrorRecord ────────────────────────────────────────────
def _format_error_body(record: ComparisonErrorRecord) -> str:
prefix: str = record._format_location_prefix()
return (
f"{prefix}Error: {record.name} ({record.exception_type})\n"
f"{record.traceback_str}"
)
def _format_error_rich_body(
record: ComparisonErrorRecord, verbosity: Verbosity = "normal"
) -> RenderableType:
prefix: str = record._format_location_prefix_rich()
name: str = escape(record.name)
header: str = (
f"{prefix}[bold red]{name} ── errored ({escape(record.exception_type)})[/]"
)
if verbosity == "minimal":
return header
return header + f"\n[dim]{escape(record.traceback_str)}[/]"
# ── _TableRecord ─────────────────────────────────────────────────────
@@ -154,10 +179,10 @@ def _format_table_rich_body(
)
# ── TensorComparisonRecord ───────────────────────────────────────────
# ── ComparisonTensorRecord ───────────────────────────────────────────
def _format_tensor_comparison_body(record: TensorComparisonRecord) -> str:
def _format_tensor_comparison_body(record: ComparisonTensorRecord) -> str:
body: str = record._format_location_prefix() + format_comparison(record)
if record.replicated_checks:
body += "\n" + format_replicated_checks(record.replicated_checks)
@@ -167,7 +192,7 @@ def _format_tensor_comparison_body(record: TensorComparisonRecord) -> str:
def _format_tensor_comparison_rich_body(
record: TensorComparisonRecord, verbosity: Verbosity = "normal"
record: ComparisonTensorRecord, verbosity: Verbosity = "normal"
) -> RenderableType:
from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import (
format_comparison_rich,
@@ -178,10 +203,10 @@ def _format_tensor_comparison_rich_body(
)
# ── NonTensorComparisonRecord ────────────────────────────────────────
# ── ComparisonNonTensorRecord ────────────────────────────────────────
def _format_non_tensor_body(record: NonTensorComparisonRecord) -> str:
def _format_non_tensor_body(record: ComparisonNonTensorRecord) -> str:
suffix: str = record._format_location_suffix()
if record.values_equal:
return f"NonTensor: {record.name}{suffix} = {record.baseline_value} ({record.baseline_type}) [equal]"
@@ -193,7 +218,7 @@ def _format_non_tensor_body(record: NonTensorComparisonRecord) -> str:
def _format_non_tensor_rich_body(
record: NonTensorComparisonRecord, verbosity: Verbosity = "normal"
record: ComparisonNonTensorRecord, verbosity: Verbosity = "normal"
) -> RenderableType:
suffix: str = record._format_location_suffix()
name: str = escape(record.name)
@@ -216,10 +241,13 @@ def _format_non_tensor_rich_body(
def _format_summary_body(record: SummaryRecord) -> str:
return (
text: str = (
f"Summary: {record.passed} passed, {record.failed} failed, "
f"{record.skipped} skipped (total {record.total})"
)
if record.errored > 0:
text += f", {record.errored} errored"
return text
def _format_summary_rich_body(
@@ -231,6 +259,8 @@ def _format_summary_rich_body(
f"[yellow]{record.skipped} skipped[/] │ "
f"{record.total} total"
)
if record.errored > 0:
text += f" │ [bold red]{record.errored} errored[/]"
return Panel(text, title="SUMMARY", border_style="bold")

63
python/sglang/srt/debug_utils/comparator/output_types.py
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 (
_format_config_body,
_format_config_rich_body,
_format_error_body,
_format_error_rich_body,
_format_log_body,
_format_non_tensor_body,
_format_non_tensor_rich_body,
@@ -146,8 +148,8 @@ class ConfigRecord(_OutputRecord):
return _format_config_rich_body(self, verbosity=verbosity)
class SkipComparisonRecord(_BaseComparisonRecord):
type: Literal["skip"] = "skip"
class ComparisonSkipRecord(_BaseComparisonRecord):
type: Literal["comparison_skip"] = "comparison_skip"
name: str
reason: str
@@ -164,6 +166,23 @@ class SkipComparisonRecord(_BaseComparisonRecord):
return _format_skip_rich_body(self, verbosity=verbosity)
class ComparisonErrorRecord(_BaseComparisonRecord):
type: Literal["comparison_error"] = "comparison_error"
name: str
exception_type: str
traceback_str: str
@property
def category(self) -> str:
return "errored"
def _format_body(self) -> str:
return _format_error_body(self)
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return _format_error_rich_body(self, verbosity=verbosity)
class _TableRecord(_OutputRecord):
label: str
rows: list[dict[str, Any]]
@@ -177,15 +196,6 @@ class _TableRecord(_OutputRecord):
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return _format_table_rich_body(self, verbosity=verbosity)
return _format_table_body(self)
def _format_rich_body(self) -> RenderableType:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_table_rich_body,
)
return _format_table_rich_body(self)
class RankInfoRecord(_TableRecord):
type: Literal["rank_info"] = "rank_info"
@@ -201,10 +211,10 @@ class InputIdsRecord(_TableRecord):
return f"{self.label} input_ids & positions"
class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord):
class ComparisonTensorRecord(TensorComparisonInfo, _BaseComparisonRecord):
model_config = ConfigDict(extra="forbid", defer_build=True)
type: Literal["comparison"] = "comparison"
type: Literal["comparison_tensor"] = "comparison_tensor"
traced_plan: Optional[TracedAlignerPlan] = None
replicated_checks: list[ReplicatedCheckResult] = Field(default_factory=list)
raw_bundle_info: Optional[Pair[BundleSideInfo]] = None
@@ -224,8 +234,8 @@ class TensorComparisonRecord(TensorComparisonInfo, _BaseComparisonRecord):
return _format_tensor_comparison_rich_body(self, verbosity=verbosity)
class NonTensorComparisonRecord(_BaseComparisonRecord):
type: Literal["non_tensor"] = "non_tensor"
class ComparisonNonTensorRecord(_BaseComparisonRecord):
type: Literal["comparison_non_tensor"] = "comparison_non_tensor"
name: str
baseline_value: str
target_value: str
@@ -245,15 +255,6 @@ class NonTensorComparisonRecord(_BaseComparisonRecord):
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return _format_non_tensor_rich_body(self, verbosity=verbosity)
return _format_non_tensor_body(self)
def _format_rich_body(self) -> RenderableType:
from sglang.srt.debug_utils.comparator.output_formatter import (
_format_non_tensor_rich_body,
)
return _format_non_tensor_rich_body(self)
class SummaryRecord(_OutputRecord):
type: Literal["summary"] = "summary"
@@ -261,13 +262,15 @@ class SummaryRecord(_OutputRecord):
passed: int
failed: int
skipped: int
errored: int = 0
@model_validator(mode="after")
def _validate_totals(self) -> "SummaryRecord":
expected: int = self.passed + self.failed + self.skipped
expected: int = self.passed + self.failed + self.skipped + self.errored
if self.total != expected:
raise ValueError(
f"total={self.total} != passed({self.passed}) + failed({self.failed}) + skipped({self.skipped}) = {expected}"
f"total={self.total} != passed({self.passed}) + failed({self.failed}) "
f"+ skipped({self.skipped}) + errored({self.errored}) = {expected}"
)
return self
@@ -277,7 +280,6 @@ class SummaryRecord(_OutputRecord):
def _format_rich_body(self, verbosity: Verbosity = "normal") -> RenderableType:
return _format_summary_rich_body(self, verbosity=verbosity)
return _format_summary_body(self)
class LogRecord(_OutputRecord):
type: Literal["log"] = "log"
@@ -291,9 +293,10 @@ AnyRecord = Annotated[
ConfigRecord,
RankInfoRecord,
InputIdsRecord,
SkipComparisonRecord,
TensorComparisonRecord,
NonTensorComparisonRecord,
ComparisonSkipRecord,
ComparisonErrorRecord,
ComparisonTensorRecord,
ComparisonNonTensorRecord,
SummaryRecord,
LogRecord,
],

6
python/sglang/srt/debug_utils/comparator/per_token_visualizer.py
View File

@@ -9,12 +9,12 @@ from __future__ import annotations
from pathlib import Path
from typing import Optional
from sglang.srt.debug_utils.comparator.output_types import TensorComparisonRecord
from sglang.srt.debug_utils.comparator.output_types import ComparisonTensorRecord
def generate_per_token_heatmap(
*,
records: list[TensorComparisonRecord],
records: list[ComparisonTensorRecord],
output_path: Path,
) -> Optional[Path]:
"""Generate a per-token relative difference heatmap PNG.
@@ -31,7 +31,7 @@ def generate_per_token_heatmap(
def _collect_per_token_data(
*,
records: list[TensorComparisonRecord],
records: list[ComparisonTensorRecord],
) -> list[tuple[str, list[float]]]:
rows: list[tuple[str, list[float]]] = []
for record in records:

8
python/sglang/srt/debug_utils/comparator/tensor_comparator/formatter.py
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.output_types import (
BundleSideInfo,
ComparisonTensorRecord,
ReplicatedCheckResult,
ShapeSnapshot,
TensorComparisonRecord,
)
from sglang.srt.debug_utils.comparator.utils import Pair
@@ -207,7 +207,7 @@ def _format_diff(diff: DiffInfo, prefix_text: str = "") -> list[str]:
def format_comparison_rich(
record: TensorComparisonRecord,
record: ComparisonTensorRecord,
verbosity: Verbosity = "normal",
) -> str:
if verbosity == "minimal":
@@ -219,7 +219,7 @@ def format_comparison_rich(
)
def _format_comparison_minimal(record: TensorComparisonRecord) -> str:
def _format_comparison_minimal(record: ComparisonTensorRecord) -> str:
passed, color, marker = _category_marker(record.category)
name_part: str = f"[bold {color}]{escape(record.name):30s}[/]"
@@ -233,7 +233,7 @@ def _format_comparison_minimal(record: TensorComparisonRecord) -> str:
def _format_comparison_normal_or_verbose(
*,
record: TensorComparisonRecord,
record: ComparisonTensorRecord,
verbose: bool,
) -> str:
passed, color, marker = _category_marker(record.category)

4
python/sglang/srt/debug_utils/comparator/utils.py
View File

@@ -141,10 +141,14 @@ def compute_exit_code(
skipped_names: list[str],
allow_failed_pattern: Optional[str],
failed_names: list[str],
errored_names: Optional[list[str]] = None,
) -> int:
if summary.passed == 0:
return 1
if errored_names:
return 1
if not _is_all_match_pattern(pattern=allow_failed_pattern, strings=failed_names):
return 1

32
test/registered/debug_utils/comparator/tensor_comparator/test_formatter.py
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 (
BundleFileInfo,
BundleSideInfo,
ComparisonTensorRecord,
ReplicatedCheckResult,
ShapeSnapshot,
TensorComparisonRecord,
)
from sglang.srt.debug_utils.comparator.tensor_comparator.formatter import (
_format_abs_diff_percentiles_rich,
@@ -273,9 +273,9 @@ def _make_comparison_record(
replicated_checks: list[ReplicatedCheckResult] | None = None,
raw_bundle_info: Pair[BundleSideInfo] | None = None,
traced_plan: TracedAlignerPlan | None = None,
) -> TensorComparisonRecord:
) -> ComparisonTensorRecord:
s: list[int] = shape if shape is not None else [4, 8]
return TensorComparisonRecord(
return ComparisonTensorRecord(
name=name,
baseline=_make_tensor_info(shape=s, dtype=dtype, sample=sample),
target=_make_tensor_info(shape=s, dtype=dtype, sample=sample),
@@ -417,7 +417,7 @@ class TestFormatComparisonRichMinimal:
"""format_comparison_rich() with verbosity='minimal'."""
def test_passed(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(rel_diff=1e-4, passed=True),
)
result: str = format_comparison_rich(record, verbosity="minimal")
@@ -428,7 +428,7 @@ class TestFormatComparisonRichMinimal:
)
def test_failed(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(rel_diff=0.5, passed=False),
)
result: str = format_comparison_rich(record, verbosity="minimal")
@@ -439,7 +439,7 @@ class TestFormatComparisonRichMinimal:
)
def test_shape_mismatch(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
record: ComparisonTensorRecord = _make_comparison_record(
shape_mismatch=True,
)
result: str = format_comparison_rich(record, verbosity="minimal")
@@ -450,7 +450,7 @@ class TestFormatComparisonRichMinimal:
)
def test_no_diff(self) -> None:
record: TensorComparisonRecord = _make_comparison_record()
record: ComparisonTensorRecord = _make_comparison_record()
result: str = format_comparison_rich(record, verbosity="minimal")
assert result == ("[red]❌[/] [bold red]hidden_states [/]")
@@ -460,7 +460,7 @@ class TestFormatComparisonRichNormal:
"""format_comparison_rich() with verbosity='normal'."""
def test_passed(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(rel_diff=1e-4, passed=True),
)
result: str = format_comparison_rich(record, verbosity="normal")
@@ -477,7 +477,7 @@ class TestFormatComparisonRichNormal:
)
def test_failed(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(
rel_diff=0.5, max_abs_diff=1.0, mean_abs_diff=0.3, passed=False
),
@@ -499,7 +499,7 @@ class TestFormatComparisonRichNormal:
)
def test_shape_mismatch(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
record: ComparisonTensorRecord = _make_comparison_record(
shape_mismatch=True,
)
result: str = format_comparison_rich(record, verbosity="normal")
@@ -516,7 +516,7 @@ class TestFormatComparisonRichNormal:
)
def test_with_downcast(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(rel_diff=0.01, passed=False),
diff_downcast=_make_diff(rel_diff=1e-5, passed=True),
downcast_dtype="torch.bfloat16",
@@ -543,7 +543,7 @@ class TestFormatComparisonRichNormal:
x=_make_bundle_side_info(num_files=2, dims="b s h(tp) d"),
y=_make_bundle_side_info(num_files=2, dims="b s h(tp) d"),
)
record: TensorComparisonRecord = _make_comparison_record(
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(passed=True),
raw_bundle_info=bundle_info,
)
@@ -565,7 +565,7 @@ class TestFormatComparisonRichNormal:
def test_with_plan(self) -> None:
plan: AlignerPlan = _make_simple_aligner_plan(with_unsharder=True)
record: TensorComparisonRecord = _make_comparison_record(
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(passed=True),
traced_plan=_make_traced_plan(plan),
)
@@ -590,7 +590,7 @@ class TestFormatComparisonRichVerbose:
"""format_comparison_rich() with verbosity='verbose'."""
def test_passed_full_detail(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(rel_diff=1e-4, passed=True),
sample="tensor([0.1, 0.2, ...])",
)
@@ -624,7 +624,7 @@ class TestFormatComparisonRichVerbose:
x=_make_bundle_side_info(num_files=2, with_parallel_info=True),
y=_make_bundle_side_info(num_files=2, with_parallel_info=True),
)
record: TensorComparisonRecord = _make_comparison_record(
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(passed=True),
raw_bundle_info=bundle_info,
)
@@ -659,7 +659,7 @@ class TestFormatComparisonRichVerbose:
def test_with_plan_and_traces(self) -> None:
plan: AlignerPlan = _make_simple_aligner_plan(with_unsharder=True)
record: TensorComparisonRecord = _make_comparison_record(
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff(passed=True),
traced_plan=_make_traced_plan(
plan,

34
test/registered/debug_utils/comparator/tensor_comparator/test_types.py
View File

@@ -4,14 +4,14 @@ import sys
import pytest
from sglang.srt.debug_utils.comparator.output_types import (
ComparisonSkipRecord,
ComparisonTensorRecord,
ConfigRecord,
ErrorLog,
InfoLog,
LogRecord,
ReplicatedCheckResult,
SkipComparisonRecord,
SummaryRecord,
TensorComparisonRecord,
parse_record_json,
)
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
@@ -84,7 +84,7 @@ class TestStrictBase:
class TestRecordTypes:
def test_comparison_record_inherits_tensor_fields(self):
record = TensorComparisonRecord(
record = ComparisonTensorRecord(
name="hidden_states",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -93,7 +93,7 @@ class TestRecordTypes:
diff=_make_diff(),
)
parsed = json.loads(record.model_dump_json())
assert parsed["type"] == "comparison"
assert parsed["type"] == "comparison_tensor"
assert parsed["name"] == "hidden_states"
assert "baseline" in parsed
assert "diff" in parsed
@@ -109,8 +109,8 @@ class TestRecordTypes:
"end_step": 100,
}
),
SkipComparisonRecord(name="attn", reason="no_baseline"),
TensorComparisonRecord(
ComparisonSkipRecord(name="attn", reason="no_baseline"),
ComparisonTensorRecord(
name="mlp",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -149,8 +149,8 @@ def _make_replicated_check(**overrides) -> ReplicatedCheckResult:
class TestWarnings:
def test_comparison_record_failed_when_diff_passed_but_errors(self):
"""TensorComparisonRecord with diff.passed=True but errors → category=='failed'."""
record = TensorComparisonRecord(
"""ComparisonTensorRecord with diff.passed=True but errors → category=='failed'."""
record = ComparisonTensorRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -162,8 +162,8 @@ class TestWarnings:
assert record.category == "failed"
def test_skip_record_failed_when_errors(self):
"""SkipComparisonRecord with errors → category=='failed' instead of 'skipped'."""
record = SkipComparisonRecord(
"""ComparisonSkipRecord with errors → category=='failed' instead of 'skipped'."""
record = ComparisonSkipRecord(
name="x",
reason="no_baseline",
errors=[ErrorLog(category="test", message="some warning")],
@@ -171,8 +171,8 @@ class TestWarnings:
assert record.category == "failed"
def test_replicated_checks_all_passed(self):
"""TensorComparisonRecord with all replicated_checks passed → category=='passed'."""
record = TensorComparisonRecord(
"""ComparisonTensorRecord with all replicated_checks passed → category=='passed'."""
record = ComparisonTensorRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -184,8 +184,8 @@ class TestWarnings:
assert record.category == "passed"
def test_replicated_checks_failed_means_record_failed(self):
"""TensorComparisonRecord with any replicated_check.passed=False → category=='failed'."""
record = TensorComparisonRecord(
"""ComparisonTensorRecord with any replicated_check.passed=False → category=='failed'."""
record = ComparisonTensorRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -197,7 +197,7 @@ class TestWarnings:
assert record.category == "failed"
def test_replicated_check_json_round_trip(self):
"""ReplicatedCheckResult survives JSON round-trip via TensorComparisonRecord."""
"""ReplicatedCheckResult survives JSON round-trip via ComparisonTensorRecord."""
check = _make_replicated_check(
axis="cp",
group_index=2,
@@ -205,7 +205,7 @@ class TestWarnings:
baseline_index=0,
passed=False,
)
record = TensorComparisonRecord(
record = ComparisonTensorRecord(
name="mlp",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -216,7 +216,7 @@ class TestWarnings:
)
restored = parse_record_json(record.model_dump_json())
assert isinstance(restored, TensorComparisonRecord)
assert isinstance(restored, ComparisonTensorRecord)
assert len(restored.replicated_checks) == 1
restored_check: ReplicatedCheckResult = restored.replicated_checks[0]

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

@@ -7,6 +7,7 @@ from pathlib import Path
import pytest
import torch
import sglang.srt.debug_utils.comparator.entrypoint as _entrypoint_module
import sglang.srt.debug_utils.dumper as _dumper_module
from sglang.srt.debug_utils.comparator.entrypoint import (
parse_args,
@@ -14,14 +15,15 @@ from sglang.srt.debug_utils.comparator.entrypoint import (
)
from sglang.srt.debug_utils.comparator.output_types import (
AnyRecord,
ComparisonErrorRecord,
ComparisonNonTensorRecord,
ComparisonSkipRecord,
ComparisonTensorRecord,
ConfigRecord,
InfoLog,
LogRecord,
NonTensorComparisonRecord,
ReplicatedCheckResult,
SkipComparisonRecord,
SummaryRecord,
TensorComparisonRecord,
_OutputRecord,
parse_record_json,
)
@@ -39,7 +41,7 @@ class TestEntrypointGroupingRaw:
"""Test `--grouping-skip-keys` empty (raw) scenarios"""
def test_run_basic(self, tmp_path, capsys):
"""Two matching tensors produce ConfigRecord, 2 TensorComparisonRecords, and SummaryRecord."""
"""Two matching tensors produce ConfigRecord, 2 ComparisonTensorRecords, and SummaryRecord."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
argv = _make_argv(baseline_path, target_path, preset="raw")
@@ -54,7 +56,7 @@ class TestEntrypointGroupingRaw:
assert summary.skipped == 0
def test_filter(self, tmp_path, capsys):
"""--filter selects only the matching tensor, producing 1 TensorComparisonRecord."""
"""--filter selects only the matching tensor, producing 1 ComparisonTensorRecord."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
argv = _make_argv(baseline_path, target_path, filter="tensor_a", preset="raw")
@@ -62,7 +64,7 @@ class TestEntrypointGroupingRaw:
assert len(_get_comparisons(records)) == 1
def test_no_baseline_skip(self, tmp_path, capsys):
"""Target tensor missing from baseline emits a SkipComparisonRecord with reason baseline_load_failed."""
"""Target tensor missing from baseline emits a ComparisonSkipRecord with reason baseline_load_failed."""
baseline_path, target_path = _create_dumps(
tmp_path,
tensor_names=["tensor_a", "tensor_extra"],
@@ -71,7 +73,7 @@ class TestEntrypointGroupingRaw:
argv = _make_argv(baseline_path, target_path, preset="raw")
records, _ = _run_and_parse(argv, capsys)
skips = [r for r in records if isinstance(r, SkipComparisonRecord)]
skips = [r for r in records if isinstance(r, ComparisonSkipRecord)]
assert len(skips) == 1
assert skips[0].reason == "baseline_load_failed"
@@ -100,7 +102,7 @@ class TestEntrypointGroupingRaw:
assert all(isinstance(r, _OutputRecord) for r in records)
def test_comparison_failed(self, tmp_path, capsys):
"""Completely different tensors produce a failed TensorComparisonRecord."""
"""Completely different tensors produce a failed ComparisonTensorRecord."""
torch.manual_seed(42)
baseline_path = _create_rank_dump(
tmp_path / "baseline", rank=0, name="tensor_a", tensor=torch.randn(10, 10)
@@ -251,7 +253,7 @@ class TestEntrypointGroupingRaw:
assert summary.total == 0
def test_raw_multi_rank(self, tmp_path, capsys):
"""Two ranks in raw grouping produce two TensorComparisonRecords (one per rank)."""
"""Two ranks in raw grouping produce two ComparisonTensorRecords (one per rank)."""
torch.manual_seed(42)
tensor = torch.randn(4, 4)
@@ -480,7 +482,7 @@ class TestEntrypointGroupingLogical:
],
)
def test_ambiguous_no_dims_skip(self, tmp_path, capsys, bad_side, expected_reason):
"""Multi-rank without dims on one side produces a SkipComparisonRecord with the appropriate reason."""
"""Multi-rank without dims on one side produces a ComparisonSkipRecord with the appropriate reason."""
torch.manual_seed(42)
tensor = torch.randn(4, 8)
@@ -500,7 +502,7 @@ class TestEntrypointGroupingLogical:
)
records, _ = _run_and_parse(argv, capsys)
skips = [r for r in records if isinstance(r, SkipComparisonRecord)]
skips = [r for r in records if isinstance(r, ComparisonSkipRecord)]
assert len(skips) == 1
assert skips[0].reason == expected_reason
@@ -1092,7 +1094,7 @@ class TestEntrypointPerStepMode:
"""Test per-step comparison mode (sglang_dev preset behavior)."""
def test_multi_step_per_step_comparison(self, tmp_path, capsys):
"""Multiple steps produce one TensorComparisonRecord per step with step field set."""
"""Multiple steps produce one ComparisonTensorRecord per step with step field set."""
torch.manual_seed(42)
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"], num_steps=3)
argv = _make_argv(baseline_path, target_path, diff_threshold=0.1)
@@ -1149,7 +1151,7 @@ class TestEntrypointPerStepMode:
assert all(c.baseline.shape == [4, 8] for c in comparisons)
def test_single_step_has_step_field(self, tmp_path, capsys):
"""Single step produces TensorComparisonRecord with location.step=0."""
"""Single step produces ComparisonTensorRecord with location.step=0."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"], num_steps=1)
argv = _make_argv(baseline_path, target_path)
@@ -1458,7 +1460,7 @@ class TestEntrypointConcatMode:
assert len(comparisons) == 3
def test_concat_aligner_plan_fields(self, tmp_path, capsys):
"""TensorComparisonRecord.traced_plan reports mode='concat' with plan=None."""
"""ComparisonTensorRecord.traced_plan reports mode='concat' with plan=None."""
torch.manual_seed(42)
records = self._run_concat(
@@ -1599,8 +1601,8 @@ class TestEntrypointConcatMode:
)
records, _ = _run_and_parse(argv, capsys)
comparisons: list[TensorComparisonRecord] = _get_comparisons(records)
hidden_comparisons: list[TensorComparisonRecord] = [
comparisons: list[ComparisonTensorRecord] = _get_comparisons(records)
hidden_comparisons: list[ComparisonTensorRecord] = [
c for c in comparisons if c.name == "hidden_states"
]
assert len(hidden_comparisons) >= 1
@@ -2158,7 +2160,7 @@ class TestEntrypointNonTensorValues:
"""Test non-tensor value comparison through the full entrypoint pipeline."""
def test_non_tensor_float_same_value(self, tmp_path: Path, capsys) -> None:
"""Two sides dump the same float → NonTensorComparisonRecord with values_equal=True, category=passed."""
"""Two sides dump the same float → ComparisonNonTensorRecord with values_equal=True, category=passed."""
baseline_path, target_path = _create_non_tensor_dumps(
tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.125
)
@@ -2177,7 +2179,7 @@ class TestEntrypointNonTensorValues:
assert summary.failed == 0
def test_non_tensor_float_different_value(self, tmp_path: Path, capsys) -> None:
"""Two sides dump different floats → NonTensorComparisonRecord with values_equal=False, category=failed."""
"""Two sides dump different floats → ComparisonNonTensorRecord with values_equal=False, category=failed."""
baseline_path, target_path = _create_non_tensor_dumps(
tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.25
)
@@ -2262,7 +2264,7 @@ class TestEntrypointNonTensorValues:
assert non_tensors[0].target_type == "dict"
def test_non_tensor_none_value(self, tmp_path: Path, capsys) -> None:
"""Dumping None is displayed as NonTensorComparisonRecord, not skipped as load failure."""
"""Dumping None is displayed as ComparisonNonTensorRecord, not skipped as load failure."""
baseline_path, target_path = _create_non_tensor_dumps(
tmp_path, name="optional_param", baseline_value=None, target_value=None
)
@@ -2278,7 +2280,7 @@ class TestEntrypointNonTensorValues:
assert non_tensors[0].category == "passed"
def test_non_tensor_json_roundtrip(self, tmp_path: Path, capsys) -> None:
"""NonTensorComparisonRecord JSON output can be parsed back correctly."""
"""ComparisonNonTensorRecord JSON output can be parsed back correctly."""
baseline_path, target_path = _create_non_tensor_dumps(
tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.125
)
@@ -2290,7 +2292,7 @@ class TestEntrypointNonTensorValues:
json_str: str = non_tensors[0].model_dump_json()
roundtripped = parse_record_json(json_str)
assert isinstance(roundtripped, NonTensorComparisonRecord)
assert isinstance(roundtripped, ComparisonNonTensorRecord)
assert roundtripped.name == "sm_scale"
assert roundtripped.values_equal is True
@@ -2344,17 +2346,17 @@ class TestEntrypointVisualize:
# --------------------------- Assertion helpers -------------------
def _get_comparisons(records: list[AnyRecord]) -> list[TensorComparisonRecord]:
return [r for r in records if isinstance(r, TensorComparisonRecord)]
def _get_comparisons(records: list[AnyRecord]) -> list[ComparisonTensorRecord]:
return [r for r in records if isinstance(r, ComparisonTensorRecord)]
def _get_non_tensors(records: list[AnyRecord]) -> list[NonTensorComparisonRecord]:
return [r for r in records if isinstance(r, NonTensorComparisonRecord)]
def _get_non_tensors(records: list[AnyRecord]) -> list[ComparisonNonTensorRecord]:
return [r for r in records if isinstance(r, ComparisonNonTensorRecord)]
def _assert_single_comparison_passed(
records: list[AnyRecord],
) -> TensorComparisonRecord:
) -> ComparisonTensorRecord:
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
assert comparisons[0].diff is not None
@@ -3235,8 +3237,8 @@ class TestEntrypointThdCpZigzag:
)
records, _ = _run_and_parse(argv, capsys)
comparisons: list[TensorComparisonRecord] = _get_comparisons(records)
hidden_comparisons: list[TensorComparisonRecord] = [
comparisons: list[ComparisonTensorRecord] = _get_comparisons(records)
hidden_comparisons: list[ComparisonTensorRecord] = [
c for c in comparisons if c.name == "hidden_states"
]
assert len(hidden_comparisons) >= 1
@@ -3287,8 +3289,8 @@ class TestEntrypointThdCpZigzag:
records, _ = _run_and_parse(argv, capsys)
# hidden_states should pass comparison (after unshard + reorder)
comparisons: list[TensorComparisonRecord] = _get_comparisons(records)
hidden_comparisons: list[TensorComparisonRecord] = [
comparisons: list[ComparisonTensorRecord] = _get_comparisons(records)
hidden_comparisons: list[ComparisonTensorRecord] = [
c for c in comparisons if c.name == "hidden_states"
]
assert len(hidden_comparisons) >= 1
@@ -3355,7 +3357,7 @@ class TestEntrypointDpFilter:
)
records, _ = _run_and_parse(argv, capsys)
comparison: TensorComparisonRecord = _assert_single_comparison_passed(records)
comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden"
def test_dp2_megatron_both_sides(self, tmp_path: Path, capsys) -> None:
@@ -3410,7 +3412,7 @@ class TestEntrypointDpFilter:
)
records, _ = _run_and_parse(argv, capsys)
comparison: TensorComparisonRecord = _assert_single_comparison_passed(records)
comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden"
def test_dp2_tp2_sglang(self, tmp_path: Path, capsys) -> None:
@@ -3458,7 +3460,7 @@ class TestEntrypointDpFilter:
)
records, _ = _run_and_parse(argv, capsys)
comparison: TensorComparisonRecord = _assert_single_comparison_passed(records)
comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden"
def test_dp2_both_nonempty_raises(self, tmp_path: Path, capsys) -> None:
@@ -3496,10 +3498,12 @@ class TestEntrypointDpFilter:
diff_threshold=1e-3,
)
with pytest.raises(
AssertionError, match="Expected exactly 1 non-empty dp_rank"
):
_run_and_parse(argv, capsys)
records, exit_code = _run_and_parse(argv, capsys)
errors = [r for r in records if isinstance(r, ComparisonErrorRecord)]
assert len(errors) == 1
assert errors[0].exception_type == "AssertionError"
assert "Expected exactly 1 non-empty dp_rank" in errors[0].traceback_str
assert exit_code == 1
class TestEntrypointDpGroupAlias:
@@ -3542,7 +3546,7 @@ class TestEntrypointDpGroupAlias:
)
records, _ = _run_and_parse(argv, capsys)
comparison: TensorComparisonRecord = _assert_single_comparison_passed(records)
comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden"
def test_dp_alias_via_override_dims(self, tmp_path: Path, capsys) -> None:
@@ -3599,7 +3603,7 @@ class TestEntrypointDpGroupAlias:
)
records, _ = _run_and_parse(argv, capsys)
comparison: TensorComparisonRecord = _assert_single_comparison_passed(records)
comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden"
def test_dp_alias_with_real_alias_group_filters(
@@ -3640,7 +3644,7 @@ class TestEntrypointDpGroupAlias:
)
records, _ = _run_and_parse(argv, capsys)
comparison: TensorComparisonRecord = _assert_single_comparison_passed(records)
comparison: ComparisonTensorRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden"
@@ -3679,7 +3683,7 @@ class TestEntrypointMetaOverride:
records: list[AnyRecord], *, expected_count: int = 1
) -> None:
"""Assert that exactly expected_count comparisons exist and all passed."""
comparisons: list[TensorComparisonRecord] = _get_comparisons(records)
comparisons: list[ComparisonTensorRecord] = _get_comparisons(records)
assert len(comparisons) == expected_count
assert all(c.diff is not None and c.diff.passed for c in comparisons)
@@ -3975,14 +3979,14 @@ class TestEntrypointMetaOverride:
)
records, _ = _run_and_parse(argv, capsys)
non_tensors: list[NonTensorComparisonRecord] = [
r for r in records if isinstance(r, NonTensorComparisonRecord)
non_tensors: list[ComparisonNonTensorRecord] = [
r for r in records if isinstance(r, ComparisonNonTensorRecord)
]
assert len(non_tensors) == 1
assert non_tensors[0].name == "sm_scale"
assert non_tensors[0].values_equal
comparisons: list[TensorComparisonRecord] = _get_comparisons(records)
comparisons: list[ComparisonTensorRecord] = _get_comparisons(records)
assert len(comparisons) == 1
assert comparisons[0].name == "hidden"
@@ -4356,12 +4360,9 @@ class TestEntrypointDpAttentionMissingAlias:
assert exit_code == 1
comparisons: list[TensorComparisonRecord] = _get_comparisons(records)
assert len(comparisons) == 1
comparison: TensorComparisonRecord = comparisons[0]
assert comparison.shape_mismatch is True
assert comparison.diff is None
assert comparison.category == "failed"
errors = [r for r in records if isinstance(r, ComparisonErrorRecord)]
assert len(errors) == 1
assert errors[0].category == "errored"
class TestEntrypointAutoDescend:
@@ -4460,5 +4461,77 @@ class TestEntrypointAutoDescend:
run(parse_args(argv))
class TestErrorResilience:
"""Bundle comparison exception → continue with remaining bundles."""
def test_one_bundle_errors_others_continue(self, tmp_path, capsys, monkeypatch):
"""One bundle raises exception → other bundles still compared, summary correct."""
baseline_path, target_path = _create_dumps(
tmp_path, ["tensor_a", "tensor_b", "tensor_c"]
)
argv = _make_argv(baseline_path, target_path, preset="raw")
original = _entrypoint_module.compare_bundle_pair
def _patched(**kwargs):
if kwargs["name"] == "tensor_b":
raise RuntimeError("intentional test error")
return original(**kwargs)
monkeypatch.setattr(_entrypoint_module, "compare_bundle_pair", _patched)
records, exit_code = _run_and_parse(argv, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 2
errors = [r for r in records if isinstance(r, ComparisonErrorRecord)]
assert len(errors) == 1
assert errors[0].name == "tensor_b"
assert errors[0].exception_type == "RuntimeError"
assert "intentional test error" in errors[0].traceback_str
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.errored == 1
assert summary.passed == 2
assert summary.total == 3
assert exit_code == 1
def test_all_bundles_error_exits_one(self, tmp_path, capsys, monkeypatch):
"""All bundles error → exit 1, summary all errored."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
argv = _make_argv(baseline_path, target_path, preset="raw")
def _always_raise(**kwargs):
raise ValueError("always fail")
monkeypatch.setattr(_entrypoint_module, "compare_bundle_pair", _always_raise)
records, exit_code = _run_and_parse(argv, capsys)
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.errored == 1
assert summary.passed == 0
assert exit_code == 1
def test_error_record_json_roundtrip_in_output(self, tmp_path, capsys, monkeypatch):
"""ComparisonErrorRecord correctly serializes and deserializes in output."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
argv = _make_argv(baseline_path, target_path, preset="raw")
def _raise(**kwargs):
raise TypeError("bad type")
monkeypatch.setattr(_entrypoint_module, "compare_bundle_pair", _raise)
records, _ = _run_and_parse(argv, capsys)
errors = [r for r in records if isinstance(r, ComparisonErrorRecord)]
assert len(errors) == 1
assert errors[0].exception_type == "TypeError"
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

12
test/registered/debug_utils/comparator/test_manually_verify.py
View File

@@ -208,7 +208,7 @@ class TestPerTokenHeatmapManualVerify:
rows for different tensor names. Colorbar shows log10 scale.
"""
from sglang.srt.debug_utils.comparator.output_types import (
TensorComparisonRecord,
ComparisonTensorRecord,
)
from sglang.srt.debug_utils.comparator.per_token_visualizer import (
generate_per_token_heatmap,
@@ -222,7 +222,7 @@ class TestPerTokenHeatmapManualVerify:
hidden_dim: int = 128
num_tensors: int = 5
records: list[TensorComparisonRecord] = []
records: list[ComparisonTensorRecord] = []
for i in range(num_tensors):
baseline: torch.Tensor = torch.randn(seq_len, hidden_dim)
noise_scale: torch.Tensor = torch.linspace(
@@ -237,7 +237,7 @@ class TestPerTokenHeatmapManualVerify:
diff_threshold=1e-3,
seq_dim=0,
)
records.append(TensorComparisonRecord(**info.model_dump()))
records.append(ComparisonTensorRecord(**info.model_dump()))
output_path: Path = tmp_path / "per_token_increasing_diff.png"
result = generate_per_token_heatmap(records=records, output_path=output_path)
@@ -253,7 +253,7 @@ class TestPerTokenHeatmapManualVerify:
rest is dark/cold.
"""
from sglang.srt.debug_utils.comparator.output_types import (
TensorComparisonRecord,
ComparisonTensorRecord,
)
from sglang.srt.debug_utils.comparator.per_token_visualizer import (
generate_per_token_heatmap,
@@ -268,7 +268,7 @@ class TestPerTokenHeatmapManualVerify:
spike_pos: int = 32
num_tensors: int = 4
records: list[TensorComparisonRecord] = []
records: list[ComparisonTensorRecord] = []
for i in range(num_tensors):
baseline: torch.Tensor = torch.randn(seq_len, hidden_dim)
target: torch.Tensor = baseline.clone()
@@ -281,7 +281,7 @@ class TestPerTokenHeatmapManualVerify:
diff_threshold=1e-3,
seq_dim=0,
)
records.append(TensorComparisonRecord(**info.model_dump()))
records.append(ComparisonTensorRecord(**info.model_dump()))
output_path: Path = tmp_path / "per_token_single_spike.png"
result = generate_per_token_heatmap(records=records, output_path=output_path)

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

@@ -4,6 +4,12 @@ import sys
import pytest
from pydantic import ValidationError
from sglang.srt.debug_utils.comparator.aligner.entrypoint.traced_types import (
TracedAlignerPlan,
TracedSidePlan,
TracedStepPlan,
TracedSubPlan,
)
from sglang.srt.debug_utils.comparator.aligner.entrypoint.types import (
AlignerPerStepPlan,
AlignerPlan,
@@ -22,11 +28,12 @@ from sglang.srt.debug_utils.comparator.aligner.unsharder.types import (
)
from sglang.srt.debug_utils.comparator.dims_spec import ParallelAxis, TokenLayout
from sglang.srt.debug_utils.comparator.output_types import (
ComparisonErrorRecord,
ComparisonNonTensorRecord,
ComparisonSkipRecord,
ComparisonTensorRecord,
ErrorLog,
NonTensorComparisonRecord,
SkipComparisonRecord,
SummaryRecord,
TensorComparisonRecord,
parse_record_json,
)
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
@@ -156,6 +163,14 @@ class TestSummaryRecord:
with pytest.raises(ValidationError, match="total=10"):
SummaryRecord(total=10, passed=5, failed=2, skipped=1)
def test_valid_with_errored(self):
record = SummaryRecord(total=10, passed=6, failed=2, skipped=1, errored=1)
assert record.errored == 1
def test_total_mismatch_with_errored(self):
with pytest.raises(ValidationError, match="total=10"):
SummaryRecord(total=10, passed=6, failed=2, skipped=1, errored=0)
class TestAxisInfo:
def test_valid(self):
@@ -208,9 +223,9 @@ def _make_comparison_record(
*,
diff: DiffInfo | None,
errors: list | None = None,
) -> TensorComparisonRecord:
) -> ComparisonTensorRecord:
ti: TensorInfo = _make_tensor_info()
return TensorComparisonRecord(
return ComparisonTensorRecord(
name="t",
baseline=ti,
target=ti,
@@ -223,7 +238,7 @@ def _make_comparison_record(
class TestOutputRecordCategories:
def test_skip_record_with_errors_is_failed(self) -> None:
record = SkipComparisonRecord(
record = ComparisonSkipRecord(
name="t",
reason="test",
errors=[ErrorLog(category="c", message="m")],
@@ -231,28 +246,28 @@ class TestOutputRecordCategories:
assert record.category == "failed"
def test_skip_record_no_warnings_is_skipped(self) -> None:
record = SkipComparisonRecord(name="t", reason="test")
record = ComparisonSkipRecord(name="t", reason="test")
assert record.category == "skipped"
def test_comparison_record_diff_none_is_failed(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(diff=None)
record: ComparisonTensorRecord = _make_comparison_record(diff=None)
assert record.category == "failed"
def test_comparison_record_passed_with_errors_is_failed(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
errors=[ErrorLog(category="c", message="m")],
)
assert record.category == "failed"
def test_comparison_record_passed_no_warnings_is_passed(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
)
assert record.category == "passed"
def test_non_tensor_record_equal_is_passed(self) -> None:
record = NonTensorComparisonRecord(
record = ComparisonNonTensorRecord(
name="sm_scale",
baseline_value="0.125",
target_value="0.125",
@@ -263,7 +278,7 @@ class TestOutputRecordCategories:
assert record.category == "passed"
def test_non_tensor_record_different_is_failed(self) -> None:
record = NonTensorComparisonRecord(
record = ComparisonNonTensorRecord(
name="sm_scale",
baseline_value="0.125",
target_value="0.25",
@@ -274,7 +289,7 @@ class TestOutputRecordCategories:
assert record.category == "failed"
def test_non_tensor_record_with_errors_is_failed(self) -> None:
record = NonTensorComparisonRecord(
record = ComparisonNonTensorRecord(
name="sm_scale",
baseline_value="0.125",
target_value="0.125",
@@ -286,7 +301,7 @@ class TestOutputRecordCategories:
assert record.category == "failed"
def test_non_tensor_record_json_roundtrip(self) -> None:
record = NonTensorComparisonRecord(
record = ComparisonNonTensorRecord(
name="sm_scale",
baseline_value="0.125",
target_value="0.25",
@@ -296,14 +311,14 @@ class TestOutputRecordCategories:
)
json_str: str = record.model_dump_json()
roundtripped = parse_record_json(json_str)
assert isinstance(roundtripped, NonTensorComparisonRecord)
assert isinstance(roundtripped, ComparisonNonTensorRecord)
assert roundtripped.name == "sm_scale"
assert roundtripped.values_equal is False
assert roundtripped.baseline_value == "0.125"
assert roundtripped.target_value == "0.25"
def test_non_tensor_record_text_format_equal(self) -> None:
record = NonTensorComparisonRecord(
record = ComparisonNonTensorRecord(
name="sm_scale",
baseline_value="0.125",
target_value="0.125",
@@ -316,7 +331,7 @@ class TestOutputRecordCategories:
assert "[equal]" in text
def test_non_tensor_record_text_format_different(self) -> None:
record = NonTensorComparisonRecord(
record = ComparisonNonTensorRecord(
name="sm_scale",
baseline_value="0.125",
target_value="0.25",
@@ -328,14 +343,38 @@ class TestOutputRecordCategories:
assert "baseline" in text
assert "target" in text
def test_error_record_category_is_errored(self) -> None:
record = ComparisonErrorRecord(
name="t", exception_type="ValueError", traceback_str="..."
)
assert record.category == "errored"
def _make_aligner_plan() -> AlignerPlan:
def test_error_record_json_roundtrip(self) -> None:
record = ComparisonErrorRecord(
name="t", exception_type="ValueError", traceback_str="traceback..."
)
json_str: str = record.model_dump_json()
roundtripped = parse_record_json(json_str)
assert isinstance(roundtripped, ComparisonErrorRecord)
assert roundtripped.name == "t"
assert roundtripped.exception_type == "ValueError"
def test_error_record_text_format(self) -> None:
record = ComparisonErrorRecord(
name="t", exception_type="RuntimeError", traceback_str="Traceback..."
)
text: str = record.to_text()
assert "RuntimeError" in text
assert "Traceback" in text
def _make_traced_aligner_plan() -> TracedAlignerPlan:
unsharder = UnsharderPlan(
axis=ParallelAxis.TP,
params=ConcatParams(dim_name="h"),
groups=[[0, 1]],
)
return AlignerPlan(
plan = AlignerPlan(
per_step_plans=Pair(
x=[
AlignerPerStepPlan(
@@ -349,54 +388,67 @@ def _make_aligner_plan() -> AlignerPlan:
],
),
)
traced_sub = TracedSubPlan(plan=unsharder, snapshot=None)
traced_step = TracedStepPlan(
step=0, input_object_indices=[0, 1], sub_plans=[traced_sub]
)
return TracedAlignerPlan(
plan=plan,
per_side=Pair(
x=TracedSidePlan(step_plans=[traced_step]),
y=TracedSidePlan(step_plans=[traced_step]),
),
)
class TestAlignerPlanInTensorComparisonRecord:
def test_comparison_record_with_aligner_plan(self) -> None:
plan: AlignerPlan = _make_aligner_plan()
record: TensorComparisonRecord = _make_comparison_record(
class TestAlignerPlanInComparisonTensorRecord:
def test_comparison_record_with_traced_plan(self) -> None:
traced_plan: TracedAlignerPlan = _make_traced_aligner_plan()
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
)
record_with_plan = record.model_copy(update={"aligner_plan": plan})
assert record_with_plan.aligner_plan is not None
assert record_with_plan.aligner_plan.per_step_plans.x[0].step == 0
record_with_plan = record.model_copy(update={"traced_plan": traced_plan})
assert record_with_plan.traced_plan is not None
assert record_with_plan.traced_plan.per_side.x.step_plans[0].step == 0
def test_aligner_plan_json_roundtrip(self) -> None:
plan: AlignerPlan = _make_aligner_plan()
record: TensorComparisonRecord = _make_comparison_record(
def test_traced_plan_json_roundtrip(self) -> None:
traced_plan: TracedAlignerPlan = _make_traced_aligner_plan()
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
)
record_with_plan = record.model_copy(update={"aligner_plan": plan})
record_with_plan = record.model_copy(update={"traced_plan": traced_plan})
json_str: str = record_with_plan.model_dump_json()
parsed = json.loads(json_str)
assert "aligner_plan" in parsed
assert "traced_plan" in parsed
assert (
parsed["aligner_plan"]["per_step_plans"]["x"][0]["sub_plans"][0]["type"]
parsed["traced_plan"]["per_side"]["x"]["step_plans"][0]["sub_plans"][0][
"plan"
]["type"]
== "unsharder"
)
roundtripped: TensorComparisonRecord = parse_record_json(json_str)
assert roundtripped.aligner_plan is not None
roundtripped: ComparisonTensorRecord = parse_record_json(json_str)
assert roundtripped.traced_plan is not None
assert (
roundtripped.aligner_plan.per_step_plans.x[0].sub_plans[0].type
roundtripped.traced_plan.per_side.x.step_plans[0].sub_plans[0].plan.type
== "unsharder"
)
def test_comparison_record_without_aligner_plan(self) -> None:
record: TensorComparisonRecord = _make_comparison_record(
def test_comparison_record_without_traced_plan(self) -> None:
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
)
json_str: str = record.model_dump_json()
roundtripped: TensorComparisonRecord = parse_record_json(json_str)
assert roundtripped.aligner_plan is None
roundtripped: ComparisonTensorRecord = parse_record_json(json_str)
assert roundtripped.traced_plan is None
def test_aligner_plan_text_format(self) -> None:
plan: AlignerPlan = _make_aligner_plan()
record: TensorComparisonRecord = _make_comparison_record(
def test_traced_plan_text_format(self) -> None:
traced_plan: TracedAlignerPlan = _make_traced_aligner_plan()
record: ComparisonTensorRecord = _make_comparison_record(
diff=_make_diff_info(passed=True),
)
record_with_plan = record.model_copy(update={"aligner_plan": plan})
record_with_plan = record.model_copy(update={"traced_plan": traced_plan})
text: str = record_with_plan.to_text()
assert "Aligner Plan:" in text

52
test/registered/debug_utils/comparator/test_output_types.py
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.output_types import (
ComparisonNonTensorRecord,
ComparisonSkipRecord,
ComparisonTensorRecord,
ConfigRecord,
ErrorLog,
InfoLog,
LogRecord,
NonTensorComparisonRecord,
RecordLocation,
SkipComparisonRecord,
SummaryRecord,
TensorComparisonRecord,
_format_aligner_plan,
_split_logs,
)
@@ -150,20 +150,20 @@ class TestConfigRecord:
# ---------------------------------------------------------------------------
# SkipComparisonRecord
# ComparisonSkipRecord
# ---------------------------------------------------------------------------
class TestSkipComparisonRecord:
class TestComparisonSkipRecord:
def test_format_body_no_step(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord(
record: ComparisonSkipRecord = ComparisonSkipRecord(
name="layer.weight",
reason="zero-dim tensor",
)
assert record._format_body() == "Skip: layer.weight (zero-dim tensor)"
def test_format_body_with_step(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord(
record: ComparisonSkipRecord = ComparisonSkipRecord(
name="layer.weight",
reason="scalar",
location=RecordLocation(step=3),
@@ -171,7 +171,7 @@ class TestSkipComparisonRecord:
assert record._format_body() == "Skip: layer.weight (step=3) (scalar)"
def test_format_rich_body(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord(
record: ComparisonSkipRecord = ComparisonSkipRecord(
name="attn.qkv",
reason="no baseline",
)
@@ -179,14 +179,14 @@ class TestSkipComparisonRecord:
assert body == "[dim]⊘ attn.qkv ── skipped (no baseline)[/]"
def test_category_skipped(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord(
record: ComparisonSkipRecord = ComparisonSkipRecord(
name="x",
reason="r",
)
assert record.category == "skipped"
def test_category_failed(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord(
record: ComparisonSkipRecord = ComparisonSkipRecord(
name="x",
reason="r",
errors=[ErrorLog(category="e", message="boom")],
@@ -195,13 +195,13 @@ class TestSkipComparisonRecord:
# ---------------------------------------------------------------------------
# NonTensorComparisonRecord
# ComparisonNonTensorRecord
# ---------------------------------------------------------------------------
class TestNonTensorComparisonRecord:
class TestComparisonNonTensorRecord:
def test_format_body_equal(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord(
record: ComparisonNonTensorRecord = ComparisonNonTensorRecord(
name="config.lr",
baseline_value="0.001",
target_value="0.001",
@@ -212,7 +212,7 @@ class TestNonTensorComparisonRecord:
assert record._format_body() == "NonTensor: config.lr = 0.001 (float) [equal]"
def test_format_body_not_equal(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord(
record: ComparisonNonTensorRecord = ComparisonNonTensorRecord(
name="config.lr",
baseline_value="0.001",
target_value="0.01",
@@ -227,7 +227,7 @@ class TestNonTensorComparisonRecord:
)
def test_format_rich_body_equal(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord(
record: ComparisonNonTensorRecord = ComparisonNonTensorRecord(
name="config.lr",
baseline_value="0.001",
target_value="0.001",
@@ -238,7 +238,7 @@ class TestNonTensorComparisonRecord:
assert record._format_rich_body() == ("═ config.lr = 0.001 (float) [green]✓[/]")
def test_format_rich_body_not_equal(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord(
record: ComparisonNonTensorRecord = ComparisonNonTensorRecord(
name="config.lr",
baseline_value="0.001",
target_value="0.01",
@@ -253,7 +253,7 @@ class TestNonTensorComparisonRecord:
)
def test_with_step(self) -> None:
record: NonTensorComparisonRecord = NonTensorComparisonRecord(
record: ComparisonNonTensorRecord = ComparisonNonTensorRecord(
name="bias",
baseline_value="True",
target_value="True",
@@ -265,7 +265,7 @@ class TestNonTensorComparisonRecord:
assert "(step=5)" in record._format_body()
def test_category(self) -> None:
passed: NonTensorComparisonRecord = NonTensorComparisonRecord(
passed: ComparisonNonTensorRecord = ComparisonNonTensorRecord(
name="x",
baseline_value="1",
target_value="1",
@@ -273,7 +273,7 @@ class TestNonTensorComparisonRecord:
target_type="int",
values_equal=True,
)
failed: NonTensorComparisonRecord = NonTensorComparisonRecord(
failed: ComparisonNonTensorRecord = ComparisonNonTensorRecord(
name="x",
baseline_value="1",
target_value="2",
@@ -328,13 +328,13 @@ class TestSummaryRecord:
# ---------------------------------------------------------------------------
# TensorComparisonRecord._format_body
# ComparisonTensorRecord._format_body
# ---------------------------------------------------------------------------
class TestTensorComparisonRecordFormatBody:
class TestComparisonTensorRecordFormatBody:
def test_basic(self) -> None:
record: TensorComparisonRecord = TensorComparisonRecord(
record: ComparisonTensorRecord = ComparisonTensorRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -365,7 +365,7 @@ class TestTensorComparisonRecordFormatBody:
def test_with_replicated_checks(self) -> None:
from sglang.srt.debug_utils.comparator.output_types import ReplicatedCheckResult
record: TensorComparisonRecord = TensorComparisonRecord(
record: ComparisonTensorRecord = ComparisonTensorRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -420,7 +420,7 @@ class TestTensorComparisonRecordFormatBody:
y=TracedSidePlan(step_plans=[]),
),
)
record: TensorComparisonRecord = TensorComparisonRecord(
record: ComparisonTensorRecord = ComparisonTensorRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -453,7 +453,7 @@ class TestTensorComparisonRecordFormatBody:
)
def test_with_step(self) -> None:
record: TensorComparisonRecord = TensorComparisonRecord(
record: ComparisonTensorRecord = ComparisonTensorRecord(
name="hidden",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
@@ -679,7 +679,7 @@ class TestOutputRecordLogAttachment:
assert text == "Config: {'a': 1}\n ✗ err1\n ℹ note1"
def test_to_rich_string_body(self) -> None:
record: SkipComparisonRecord = SkipComparisonRecord(
record: ComparisonSkipRecord = ComparisonSkipRecord(
name="x",
reason="r",
errors=[ErrorLog(category="e", message="oops")],

14
test/registered/debug_utils/comparator/test_per_token_visualizer.py
View File

@@ -9,7 +9,7 @@ from pathlib import Path
import pytest
import torch
from sglang.srt.debug_utils.comparator.output_types import TensorComparisonRecord
from sglang.srt.debug_utils.comparator.output_types import ComparisonTensorRecord
from sglang.srt.debug_utils.comparator.tensor_comparator.comparator import (
compare_tensor_pair,
)
@@ -31,8 +31,8 @@ def _make_comparison_record(
baseline: torch.Tensor,
target: torch.Tensor,
seq_dim: int = 0,
) -> TensorComparisonRecord:
"""Build a TensorComparisonRecord with per-token data from raw tensors."""
) -> ComparisonTensorRecord:
"""Build a ComparisonTensorRecord with per-token data from raw tensors."""
info = compare_tensor_pair(
x_baseline=baseline,
x_target=target,
@@ -40,7 +40,7 @@ def _make_comparison_record(
diff_threshold=1e-3,
seq_dim=seq_dim,
)
return TensorComparisonRecord(**info.model_dump())
return ComparisonTensorRecord(**info.model_dump())
class TestPerTokenVisualizer:
@@ -68,7 +68,7 @@ class TestPerTokenVisualizer:
name="no_per_token",
diff_threshold=1e-3,
)
record = TensorComparisonRecord(**info.model_dump())
record = ComparisonTensorRecord(**info.model_dump())
output_path: Path = tmp_path / "no_data.png"
result = generate_per_token_heatmap(records=[record], output_path=output_path)
@@ -82,7 +82,7 @@ class TestPerTokenVisualizer:
)
torch.manual_seed(42)
records: list[TensorComparisonRecord] = [
records: list[ComparisonTensorRecord] = [
_make_comparison_record(
name=f"tensor_{i}",
baseline=torch.randn(16, 32),
@@ -108,7 +108,7 @@ class TestPerTokenVisualizer:
)
torch.manual_seed(42)
records: list[TensorComparisonRecord] = [
records: list[ComparisonTensorRecord] = [
_make_comparison_record(
name="short",
baseline=torch.randn(4, 8),

30
test/registered/debug_utils/comparator/test_utils.py
View File

@@ -410,6 +410,36 @@ class TestComputeExitCode:
== 1
)
def test_errored_with_passed_exits_one(self):
"""Has errored bundle even with passed → exit 1."""
summary = SummaryRecord(total=3, passed=2, failed=0, skipped=0, errored=1)
assert (
compute_exit_code(
summary,
allow_skipped_pattern=".*",
skipped_names=[],
allow_failed_pattern=None,
failed_names=[],
errored_names=["broken_tensor"],
)
== 1
)
def test_errored_only_exits_one(self):
"""All errored → exit 1 (passed==0 already exits 1, but errored also independently triggers)."""
summary = SummaryRecord(total=1, passed=0, failed=0, skipped=0, errored=1)
assert (
compute_exit_code(
summary,
allow_skipped_pattern=".*",
skipped_names=[],
allow_failed_pattern=None,
failed_names=[],
errored_names=["broken_tensor"],
)
== 1
)
def _make_pt(directory: Path) -> None:
directory.mkdir(parents=True, exist_ok=True)

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

@@ -4,6 +4,12 @@ from __future__ import annotations
from typing import Optional
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(
est_time=0, suite="default", nightly=True, disabled="helper module, no tests"
)
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
DiffInfo,
TensorInfo,