chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Add skip patterns, tee to file, tensor load warning in dump comparator (#19600)

Browse Source
This commit is contained in:
fzyzcjy
2026-03-01 10:36:22 +08:00
committed by GitHub
parent b0b26a7ef1
commit 46960e65cf
10 changed files with 684 additions and 172 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
python/sglang/srt/debug_utils/comparator/bundle_comparator.py
View File

@@ -315,8 +315,16 @@ def _apply_dim_names_from_meta(
def _load_all_values(filenames: list[str], base_path: Path) -> list[ValueWithMeta]:
return [
item
for f in filenames
if (item := ValueWithMeta.load(base_path / f)).value is not LOAD_FAILED
]
result: list[ValueWithMeta] = []
for f in filenames:
item: ValueWithMeta = ValueWithMeta.load(base_path / f)
if item.value is LOAD_FAILED:
warning_sink.add(
GeneralWarning(
category="load_failed",
message=f"Failed to load tensor file: {f}",
)
)
continue
result.append(item)
return result

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

@@ -10,7 +10,7 @@ import polars as pl
from sglang.srt.debug_utils.comparator.output_types import (
InputIdsRecord,
RankInfoRecord,
print_record,
report_sink,
)
from sglang.srt.debug_utils.dump_loader import LOAD_FAILED, ValueWithMeta
@@ -23,25 +23,18 @@ def emit_display_records(
dump_dir: Path,
label: str,
tokenizer: Any,
output_format: str,
) -> None:
rank_rows: Optional[list[dict[str, Any]]] = _collect_rank_info(
df, dump_dir=dump_dir
)
if rank_rows is not None:
print_record(
RankInfoRecord(label=label, rows=rank_rows),
output_format=output_format,
)
report_sink.add(RankInfoRecord(label=label, rows=rank_rows))
input_ids_rows: Optional[list[dict[str, Any]]] = _collect_input_ids_and_positions(
df, dump_dir=dump_dir, tokenizer=tokenizer
)
if input_ids_rows is not None:
print_record(
InputIdsRecord(label=label, rows=input_ids_rows),
output_format=output_format,
)
report_sink.add(InputIdsRecord(label=label, rows=input_ids_rows))
def _render_polars_as_text(df: pl.DataFrame, *, title: Optional[str] = None) -> str:

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

@@ -1,6 +1,8 @@
from __future__ import annotations
import argparse
import re
import sys
from pathlib import Path
from typing import Any, Iterator, Optional, Union
@@ -29,88 +31,124 @@ from sglang.srt.debug_utils.comparator.output_types import (
NonTensorRecord,
SkipRecord,
SummaryRecord,
print_record,
report_sink,
)
from sglang.srt.debug_utils.comparator.per_token_visualizer import (
generate_per_token_heatmap,
)
from sglang.srt.debug_utils.comparator.per_token_visualizer import (
generate_per_token_heatmap,
)
from sglang.srt.debug_utils.comparator.utils import Pair
from sglang.srt.debug_utils.comparator.warning_sink import warning_sink
from sglang.srt.debug_utils.dump_loader import read_meta, read_tokenizer_path
def main() -> None:
args = _parse_args()
run(args)
sys.exit(run(args))
def run(args: argparse.Namespace) -> None:
print_record(
ConfigRecord.from_args(args),
def run(args: argparse.Namespace) -> int:
report_path: Optional[Path] = _resolve_report_path(args)
report_sink.configure(
output_format=args.output_format,
report_path=report_path,
)
warning_sink.set_output_format(args.output_format)
try:
report_sink.add(ConfigRecord.from_args(args))
dfs: Pair[pl.DataFrame] = _read_df(args)
dfs: Pair[pl.DataFrame] = _read_df(args)
tokenizer: Any = _maybe_load_tokenizer(args)
for label, df, dump_dir in [
("baseline", dfs.x, Path(args.baseline_path)),
("target", dfs.y, Path(args.target_path)),
]:
emit_display_records(
df=df,
dump_dir=dump_dir,
label=label,
tokenizer=tokenizer,
output_format=args.output_format,
tokenizer: Any = _maybe_load_tokenizer(args)
for label, df, dump_dir in [
("baseline", dfs.x, Path(args.baseline_path)),
("target", dfs.y, Path(args.target_path)),
]:
emit_display_records(
df=df,
dump_dir=dump_dir,
label=label,
tokenizer=tokenizer,
)
ta_result: TokenAlignerResult = compute_maybe_token_aligner_result(args, dfs)
if ta_result.mode == "smart":
dfs = dfs.map(lambda df: df.filter(~pl.col("name").is_in(AUX_NAMES)))
bundle_info_pairs: list[Pair[TensorBundleInfo]] = match_bundles(
dfs=dfs,
skip_keys=_compute_skip_keys(
args, has_token_aligner=ta_result.mode is not None
),
)
ta_result: TokenAlignerResult = compute_maybe_token_aligner_result(args, dfs)
viz_output_dir: Optional[Path] = (
Path(args.viz_output_dir) if args.viz_bundle_details else None
)
if ta_result.mode == "smart":
dfs = dfs.map(lambda df: df.filter(~pl.col("name").is_in(AUX_NAMES)))
visualize_per_token: Optional[Path] = (
Path(args.visualize_per_token) if args.visualize_per_token else None
)
bundle_info_pairs: list[Pair[TensorBundleInfo]] = match_bundles(
dfs=dfs,
skip_keys=_compute_skip_keys(
args, has_token_aligner=ta_result.mode is not None
),
)
meta_overrider: MetaOverrider = MetaOverrider.from_args_and_config(
override_dims=args.override_dims,
override_baseline_dims=args.override_baseline_dims,
override_target_dims=args.override_target_dims,
override_config=(
Path(args.override_config) if args.override_config else None
),
)
viz_output_dir: Optional[Path] = (
Path(args.viz_output_dir) if args.viz_bundle_details else None
)
comparison_records = _compare_bundle_pairs(
bundle_info_pairs=bundle_info_pairs,
baseline_path=Path(args.baseline_path),
target_path=Path(args.target_path),
token_aligner_mode=ta_result.mode,
token_aligner_plan=ta_result.plan,
diff_threshold=args.diff_threshold,
thd_seq_lens_by_step_pair=ta_result.thd_seq_lens_by_step_pair,
viz_output_dir=viz_output_dir,
compute_per_token=visualize_per_token is not None,
meta_overrider=meta_overrider,
)
summary, skipped_names = _consume_comparison_records(
comparison_records=comparison_records,
visualize_per_token=visualize_per_token,
)
return _compute_exit_code(
summary,
allow_skip_pattern=args.allow_skip_pattern,
skipped_names=skipped_names,
)
finally:
report_sink.close()
if report_path is not None:
print(f"Report: {report_path}", file=sys.stderr)
visualize_per_token: Optional[Path] = (
Path(args.visualize_per_token) if args.visualize_per_token else None
)
meta_overrider: MetaOverrider = MetaOverrider.from_args_and_config(
override_dims=args.override_dims,
override_baseline_dims=args.override_baseline_dims,
override_target_dims=args.override_target_dims,
override_config=Path(args.override_config) if args.override_config else None,
)
def _compute_exit_code(
summary: SummaryRecord,
*,
allow_skip_pattern: str,
skipped_names: list[str],
) -> int:
if summary.failed > 0:
return 1
comparison_records = _compare_bundle_pairs(
bundle_info_pairs=bundle_info_pairs,
baseline_path=Path(args.baseline_path),
target_path=Path(args.target_path),
token_aligner_mode=ta_result.mode,
token_aligner_plan=ta_result.plan,
diff_threshold=args.diff_threshold,
thd_seq_lens_by_step_pair=ta_result.thd_seq_lens_by_step_pair,
viz_output_dir=viz_output_dir,
compute_per_token=visualize_per_token is not None,
meta_overrider=meta_overrider,
)
_consume_comparison_records(
comparison_records=comparison_records,
output_format=args.output_format,
visualize_per_token=visualize_per_token,
)
pattern: re.Pattern[str] = re.compile(allow_skip_pattern)
forbidden: list[str] = [n for n in skipped_names if not pattern.fullmatch(n)]
if forbidden:
return 1
return 0
def _resolve_report_path(args: argparse.Namespace) -> Optional[Path]:
if args.report_path is not None:
return Path(args.report_path) if args.report_path else None
return Path(args.target_path) / "comparator_report.jsonl"
def _maybe_load_tokenizer(args: argparse.Namespace) -> Any:
@@ -195,22 +233,30 @@ def _compare_bundle_pairs(
def _consume_comparison_records(
*,
comparison_records: Iterator[Union[ComparisonRecord, SkipRecord, NonTensorRecord]],
output_format: str,
visualize_per_token: Optional[Path] = None,
) -> None:
) -> tuple[SummaryRecord, list[str]]:
counts: dict[str, int] = {"passed": 0, "failed": 0, "skipped": 0}
collected_comparisons: list[ComparisonRecord] = []
skipped_names: list[str] = []
for record in comparison_records:
counts[record.category] += 1
print_record(record, output_format=output_format)
report_sink.add(record)
if isinstance(record, SkipRecord) and record.category == "skipped":
skipped_names.append(record.name)
if visualize_per_token is not None and isinstance(record, ComparisonRecord):
collected_comparisons.append(record)
print_record(
SummaryRecord(total=sum(counts.values()), **counts),
output_format=output_format,
)
summary: SummaryRecord = SummaryRecord(total=sum(counts.values()), **counts)
report_sink.add(summary)
if visualize_per_token is not None and collected_comparisons:
generate_per_token_heatmap(
records=collected_comparisons,
output_path=visualize_per_token,
)
return summary, skipped_names
if visualize_per_token is not None and collected_comparisons:
generate_per_token_heatmap(
@@ -300,5 +346,21 @@ def _parse_args() -> argparse.Namespace:
default=None,
help="Path to YAML override config file (dims overrides, etc.)",
)
parser.add_argument(
"--allow-skip-pattern",
type=str,
default=".*",
help="Regex pattern for tensor names allowed to be skipped. "
"Default '.*' allows all skips. Use '^$' to forbid all skips.",
)
# Report output
parser.add_argument(
"--report-path",
type=str,
default=None,
help="Path for JSONL report (default: <target-path>/comparator_report.jsonl). "
"Pass empty string '' to disable.",
)
return parser.parse_args()

60
python/sglang/srt/debug_utils/comparator/output_types.py
View File

@@ -1,7 +1,9 @@
from __future__ import annotations
import sys
from abc import abstractmethod
from typing import TYPE_CHECKING, Annotated, Any, Literal, Optional, Union
from pathlib import Path
from typing import IO, TYPE_CHECKING, Annotated, Any, Literal, Optional, Union
import polars as pl
from pydantic import ConfigDict, Discriminator, Field, TypeAdapter, model_validator
@@ -253,8 +255,62 @@ def parse_record_json(json_str: str | bytes) -> AnyRecord:
return _get_any_record_adapter().validate_json(json_str)
def print_record(record: _OutputRecord, output_format: str) -> None:
def _print_to_stdout(record: _OutputRecord, *, output_format: str) -> None:
if output_format == "json":
print(record.model_dump_json())
else:
print(record.to_text())
class ReportSink:
"""Unified entry point for all record output."""
def __init__(self) -> None:
self._output_format: str = "text"
self._report_file: Optional[IO[str]] = None
self._report_path: Optional[Path] = None
def configure(
self,
*,
output_format: str = "text",
report_path: Optional[Path] = None,
) -> None:
self._output_format = output_format
if report_path is not None:
try:
report_path.parent.mkdir(parents=True, exist_ok=True)
self._report_file = open(report_path, "w", encoding="utf-8")
self._report_path = report_path
except OSError as exc:
print(
f"Warning: cannot open report file {report_path}: {exc}",
file=sys.stderr,
)
def add(self, record: _OutputRecord) -> None:
_print_to_stdout(record, output_format=self._output_format)
if self._report_file is not None:
self._report_file.write(record.model_dump_json())
self._report_file.write("\n")
self._report_file.flush()
def close(self) -> None:
if self._report_file is not None:
self._report_file.close()
self._report_file = None
@property
def report_path(self) -> Optional[Path]:
return self._report_path
def _reset(self) -> None:
"""Reset state for test isolation."""
self.close()
self._output_format = "text"
self._report_path = None
report_sink = ReportSink()

11
python/sglang/srt/debug_utils/comparator/warning_sink.py
View File

@@ -9,10 +9,6 @@ from sglang.srt.debug_utils.comparator.output_types import AnyWarning
class WarningSink:
def __init__(self) -> None:
self._stack: list[list[AnyWarning]] = []
self._output_format: str = "text"
def set_output_format(self, output_format: str) -> None:
self._output_format = output_format
@contextmanager
def context(self) -> Generator[list[AnyWarning], None, None]:
@@ -30,13 +26,10 @@ class WarningSink:
else:
from sglang.srt.debug_utils.comparator.output_types import (
WarningRecord,
print_record,
report_sink,
)
print_record(
WarningRecord(warnings=[warning]),
output_format=self._output_format,
)
report_sink.add(WarningRecord(warnings=[warning]))
warning_sink = WarningSink()

9
test/registered/debug_utils/comparator/conftest.py
View File

@@ -0,0 +1,9 @@
import pytest
from sglang.srt.debug_utils.comparator.output_types import report_sink
@pytest.fixture(autouse=True)
def _reset_report_sink() -> None:
yield
report_sink._reset()

90
test/registered/debug_utils/comparator/test_bundle_comparator.py Normal file
View File

@@ -0,0 +1,90 @@
import sys
from pathlib import Path
from unittest.mock import patch
import pytest
import torch
from sglang.srt.debug_utils.comparator.bundle_comparator import _load_all_values
from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
from sglang.srt.debug_utils.comparator.warning_sink import WarningSink
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=15, suite="default", nightly=True)
def _save_tensor(
dump_path: Path,
*,
name: str,
step: int = 0,
rank: int = 0,
) -> str:
filename: str = f"step={step}___rank={rank}___dump_index=0___name={name}.pt"
tensor: torch.Tensor = torch.randn(4)
torch.save({"value": tensor, "meta": {}}, dump_path / filename)
return filename
class TestLoadAllValues:
def test_all_success(self, tmp_path: Path) -> None:
"""All files load successfully — no warnings emitted."""
fn0: str = _save_tensor(tmp_path, name="a", rank=0)
fn1: str = _save_tensor(tmp_path, name="a", rank=1)
sink = WarningSink()
with sink.context() as warnings:
with patch(
"sglang.srt.debug_utils.comparator.bundle_comparator.warning_sink",
sink,
):
result = _load_all_values(filenames=[fn0, fn1], base_path=tmp_path)
assert len(result) == 2
assert len(warnings) == 0
def test_one_corrupted_emits_warning(self, tmp_path: Path) -> None:
"""One corrupted file is filtered out and emits a load_failed warning."""
fn_good: str = _save_tensor(tmp_path, name="a", rank=0)
fn_bad: str = "step=0___rank=1___dump_index=0___name=a.pt"
(tmp_path / fn_bad).write_text("not a valid pt file")
sink = WarningSink()
with sink.context() as warnings:
with patch(
"sglang.srt.debug_utils.comparator.bundle_comparator.warning_sink",
sink,
):
result = _load_all_values(
filenames=[fn_good, fn_bad], base_path=tmp_path
)
assert len(result) == 1
assert len(warnings) == 1
assert isinstance(warnings[0], GeneralWarning)
assert warnings[0].category == "load_failed"
assert fn_bad in warnings[0].message
def test_all_corrupted_emits_warnings_returns_empty(self, tmp_path: Path) -> None:
"""All files corrupted — returns empty list and emits one warning per file."""
fn0: str = "step=0___rank=0___dump_index=0___name=a.pt"
fn1: str = "step=0___rank=1___dump_index=0___name=a.pt"
(tmp_path / fn0).write_text("corrupt")
(tmp_path / fn1).write_text("corrupt")
sink = WarningSink()
with sink.context() as warnings:
with patch(
"sglang.srt.debug_utils.comparator.bundle_comparator.warning_sink",
sink,
):
result = _load_all_values(filenames=[fn0, fn1], base_path=tmp_path)
assert len(result) == 0
assert len(warnings) == 2
assert all(w.category == "load_failed" for w in warnings)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

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

@@ -1,3 +1,4 @@
import subprocess
import sys
import textwrap
from argparse import Namespace
@@ -7,7 +8,7 @@ import pytest
import torch
import sglang.srt.debug_utils.dumper as _dumper_module
from sglang.srt.debug_utils.comparator.entrypoint import run
from sglang.srt.debug_utils.comparator.entrypoint import _compute_exit_code, run
from sglang.srt.debug_utils.comparator.output_types import (
AnyRecord,
ComparisonRecord,
@@ -39,7 +40,7 @@ class TestEntrypointGroupingRaw:
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
args = _make_args(baseline_path, target_path, grouping="raw")
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
assert isinstance(records[0], ConfigRecord)
assert len(_get_comparisons(records)) == 2
@@ -54,7 +55,7 @@ class TestEntrypointGroupingRaw:
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
args = _make_args(baseline_path, target_path, filter="tensor_a", grouping="raw")
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
assert len(_get_comparisons(records)) == 1
def test_no_baseline_skip(self, tmp_path, capsys):
@@ -66,7 +67,7 @@ class TestEntrypointGroupingRaw:
)
args = _make_args(baseline_path, target_path, grouping="raw")
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
skips = [r for r in records if isinstance(r, SkipRecord)]
assert len(skips) == 1
assert skips[0].reason == "baseline_load_failed"
@@ -82,7 +83,7 @@ class TestEntrypointGroupingRaw:
baseline_path, target_path, start_step=1, end_step=1, grouping="raw"
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.total == 1
@@ -92,7 +93,7 @@ class TestEntrypointGroupingRaw:
baseline_path, target_path = _create_dumps(tmp_path, ["t"], num_steps=2)
args = _make_args(baseline_path, target_path, grouping="raw")
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
assert all(isinstance(r, _OutputRecord) for r in records)
def test_comparison_failed(self, tmp_path, capsys):
@@ -111,7 +112,7 @@ class TestEntrypointGroupingRaw:
baseline_path, target_path, grouping="raw", diff_threshold=1e-3
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
assert comparisons[0].diff is not None
@@ -133,7 +134,7 @@ class TestEntrypointGroupingRaw:
)
args = _make_args(baseline_path, target_path, grouping="raw")
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
assert comparisons[0].shape_mismatch is True
@@ -159,7 +160,7 @@ class TestEntrypointGroupingRaw:
)
args = _make_args(baseline_path, target_path, grouping="raw")
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
@@ -189,7 +190,7 @@ class TestEntrypointGroupingRaw:
baseline_path, target_path, grouping="raw", diff_threshold=0.01
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
assert comparisons[0].diff_downcast is not None
@@ -227,7 +228,7 @@ class TestEntrypointGroupingRaw:
diff_threshold=1e-3,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.passed == 1
@@ -242,7 +243,7 @@ class TestEntrypointGroupingRaw:
baseline_path, target_path, filter="nonexistent_pattern", grouping="raw"
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.total == 0
@@ -271,7 +272,7 @@ class TestEntrypointGroupingRaw:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 2
@@ -348,7 +349,7 @@ class TestEntrypointGroupingRaw:
grouping="raw",
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 2
assert all(c.diff is not None and c.diff.passed for c in comparisons)
@@ -367,7 +368,7 @@ class TestEntrypointGroupingLogical:
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
args = _make_args(baseline_path, target_path)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
assert len(_get_comparisons(records)) == 2
summary = records[-1]
assert isinstance(summary, SummaryRecord)
@@ -402,7 +403,7 @@ class TestEntrypointGroupingLogical:
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
@@ -439,7 +440,7 @@ class TestEntrypointGroupingLogical:
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
_assert_single_comparison_passed(records)
def test_one_side_dims_single_baseline(self, tmp_path, capsys):
@@ -466,7 +467,7 @@ class TestEntrypointGroupingLogical:
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
_assert_single_comparison_passed(records)
@pytest.mark.parametrize(
@@ -496,7 +497,7 @@ class TestEntrypointGroupingLogical:
target_dir / _FIXED_EXP_NAME,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
skips = [r for r in records if isinstance(r, SkipRecord)]
assert len(skips) == 1
assert skips[0].reason == expected_reason
@@ -535,7 +536,7 @@ class TestEntrypointGroupingLogical:
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.total == 2
@@ -572,7 +573,7 @@ class TestEntrypointGroupingLogical:
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
# concat along dim 0 (fallback, no token dim) → 2 steps × [4, 8] = [8, 8]
@@ -613,7 +614,7 @@ class TestEntrypointGroupingLogical:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "attn_out"
@@ -651,7 +652,7 @@ class TestEntrypointGroupingLogical:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
assert comparisons[0].name == "t_a"
@@ -696,7 +697,7 @@ class TestEntrypointGroupingLogical:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 2
assert all(c.diff is not None and c.diff.passed for c in comparisons)
@@ -737,7 +738,7 @@ class TestEntrypointGroupingLogical:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
@@ -776,7 +777,7 @@ class TestEntrypointGroupingLogical:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
_assert_single_comparison_passed(records)
def test_ep_cp_tp_three_axis_unshard(self, tmp_path, capsys):
@@ -811,7 +812,7 @@ class TestEntrypointGroupingLogical:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
@@ -845,7 +846,7 @@ class TestEntrypointGroupingLogical:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "attn_out"
@@ -879,7 +880,7 @@ class TestEntrypointGroupingLogical:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
@@ -906,7 +907,7 @@ class TestEntrypointGroupingLogical:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
@@ -934,7 +935,7 @@ class TestEntrypointGroupingLogical:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
@@ -970,7 +971,7 @@ class TestEntrypointGroupingLogical:
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "attn_out"
@@ -1001,7 +1002,7 @@ class TestEntrypointGroupingLogical:
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "attn_out"
@@ -1041,7 +1042,7 @@ class TestEntrypointGroupingLogical:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
@@ -1105,7 +1106,8 @@ class TestEntrypointConcatMode:
args: Namespace = _make_args(
baseline_path, target_path, diff_threshold=diff_threshold
)
return _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
return records
def test_concat_multi_step_different_data(self, tmp_path, capsys):
"""Multi-step concat with different data per step + truncation."""
@@ -1168,7 +1170,7 @@ class TestEntrypointConcatMode:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
# 2 steps × [4, 8] concat along dim 0 (fallback) → [8, 8]
@@ -1330,7 +1332,7 @@ class TestEntrypointConcatMode:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
# all 3 tensors should be compared (not filtered out)
names = {c.name for c in comparisons}
@@ -1425,7 +1427,7 @@ class TestEntrypointConcatMode:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
# CP unshard: [4,4,6] × 2 ranks → [4,8,6] per step
@@ -1477,7 +1479,7 @@ class TestEntrypointConcatMode:
token_aligner="concat_steps",
diff_threshold=1e-3,
)
records: list[AnyRecord] = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons: list[ComparisonRecord] = _get_comparisons(records)
hidden_comparisons: list[ComparisonRecord] = [
@@ -1519,7 +1521,7 @@ class TestEntrypointAxisAligner:
diff_threshold=1e-3,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
assert comp.baseline.shape == [4, 16, 8]
@@ -1556,7 +1558,7 @@ class TestEntrypointAxisAligner:
diff_threshold=1e-3,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
@@ -1589,7 +1591,7 @@ class TestEntrypointAxisAligner:
diff_threshold=1e-3,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.name == "hidden"
assert comp.baseline.shape == [4, 8]
@@ -1702,7 +1704,7 @@ class TestEntrypointReplicatedAxis:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comp = _assert_single_comparison_passed(records)
assert comp.warnings == []
assert all(c.passed for c in comp.replicated_checks)
@@ -1741,7 +1743,7 @@ class TestEntrypointReplicatedAxis:
diff_threshold=0.01,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
assert comparisons[0].category == "failed"
@@ -1787,7 +1789,7 @@ class TestEntrypointReplicatedAxis:
diff_threshold=0.5,
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
@@ -1850,7 +1852,7 @@ class TestEntrypointAlignment:
args = _make_args(
exp_paths[0], exp_paths[1], grouping="logical", token_aligner="smart"
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
# AUX_NAMES are filtered out after plan computation → only hidden_states remains
@@ -1964,7 +1966,7 @@ class TestEntrypointAlignment:
token_aligner="smart",
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
warning_records = [r for r in records if isinstance(r, WarningRecord)]
layout_warnings = [
@@ -2031,7 +2033,7 @@ class TestEntrypointNonTensorValues:
tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.125
)
args = _make_args(baseline_path, target_path, grouping="raw")
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
non_tensors = _get_non_tensors(records)
assert len(non_tensors) == 1
@@ -2050,7 +2052,7 @@ class TestEntrypointNonTensorValues:
tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.25
)
args = _make_args(baseline_path, target_path, grouping="raw")
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
non_tensors = _get_non_tensors(records)
assert len(non_tensors) == 1
@@ -2070,7 +2072,7 @@ class TestEntrypointNonTensorValues:
target_value="flash_attn",
)
args = _make_args(baseline_path, target_path, grouping="raw")
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
non_tensors = _get_non_tensors(records)
assert len(non_tensors) == 1
@@ -2100,7 +2102,7 @@ class TestEntrypointNonTensorValues:
target_dir / _FIXED_EXP_NAME,
grouping="raw",
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
non_tensors = _get_non_tensors(records)
@@ -2121,7 +2123,7 @@ class TestEntrypointNonTensorValues:
tmp_path, name="debug_info", baseline_value=value, target_value=value
)
args = _make_args(baseline_path, target_path, grouping="raw")
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
non_tensors = _get_non_tensors(records)
assert len(non_tensors) == 1
@@ -2135,7 +2137,7 @@ class TestEntrypointNonTensorValues:
tmp_path, name="optional_param", baseline_value=None, target_value=None
)
args = _make_args(baseline_path, target_path, grouping="raw")
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
non_tensors = _get_non_tensors(records)
assert len(non_tensors) == 1
@@ -2151,7 +2153,7 @@ class TestEntrypointNonTensorValues:
tmp_path, name="sm_scale", baseline_value=0.125, target_value=0.125
)
args = _make_args(baseline_path, target_path, grouping="raw")
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
non_tensors = _get_non_tensors(records)
assert len(non_tensors) == 1
@@ -2186,7 +2188,7 @@ class TestEntrypointVisualize:
viz_output_dir=str(viz_dir),
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
assert len(_get_comparisons(records)) == 1
png_files = list(viz_dir.glob("*.png"))
@@ -2341,15 +2343,19 @@ def _make_args(baseline_path: Path, target_path: Path, **overrides) -> Namespace
override_baseline_dims=[],
override_target_dims=[],
override_config=None,
allow_skip_pattern=".*",
report_path="",
)
defaults.update(overrides)
return Namespace(**defaults)
def _run_and_parse(args: Namespace, capsys: pytest.CaptureFixture) -> list[AnyRecord]:
def _run_and_parse(
args: Namespace, capsys: pytest.CaptureFixture
) -> tuple[list[AnyRecord], int]:
capsys.readouterr()
run(args)
return _parse_jsonl(capsys.readouterr().out)
exit_code: int = run(args)
return _parse_jsonl(capsys.readouterr().out), exit_code
def _parse_jsonl(output: str) -> list[AnyRecord]:
@@ -2875,7 +2881,7 @@ class TestEntrypointPerTokenVisualization:
grouping="raw",
visualize_per_token=str(output_png),
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 2
@@ -2892,7 +2898,7 @@ class TestEntrypointPerTokenVisualization:
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
args = _make_args(baseline_path, target_path, grouping="raw")
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons = _get_comparisons(records)
assert len(comparisons) == 1
@@ -2992,7 +2998,7 @@ class TestEntrypointThdCpZigzag:
token_aligner="smart",
diff_threshold=1e-3,
)
records: list[AnyRecord] = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparisons: list[ComparisonRecord] = _get_comparisons(records)
hidden_comparisons: list[ComparisonRecord] = [
@@ -3043,7 +3049,7 @@ class TestEntrypointThdCpZigzag:
token_aligner="smart",
diff_threshold=1e-3,
)
records: list[AnyRecord] = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
# hidden_states should pass comparison (after unshard + reorder)
comparisons: list[ComparisonRecord] = _get_comparisons(records)
@@ -3113,7 +3119,7 @@ class TestEntrypointDpFilter:
grouping="logical",
diff_threshold=1e-3,
)
records: list[AnyRecord] = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparison: ComparisonRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden"
@@ -3169,7 +3175,7 @@ class TestEntrypointDpFilter:
grouping="logical",
diff_threshold=1e-3,
)
records: list[AnyRecord] = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparison: ComparisonRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden"
@@ -3218,7 +3224,7 @@ class TestEntrypointDpFilter:
grouping="logical",
diff_threshold=1e-3,
)
records: list[AnyRecord] = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
comparison: ComparisonRecord = _assert_single_comparison_passed(records)
assert comparison.name == "hidden"
@@ -3344,7 +3350,7 @@ class TestEntrypointMetaOverride:
grouping="logical",
override_dims=["hidden:t h(tp)"],
)
self._assert_all_passed(_run_and_parse(args, capsys))
self._assert_all_passed(_run_and_parse(args, capsys)[0])
@pytest.mark.parametrize(
"baseline_dims, target_dims, override_kwarg",
@@ -3371,7 +3377,7 @@ class TestEntrypointMetaOverride:
)
args = _make_args(baseline_path, target_path, grouping="raw", **override_kwarg)
self._assert_all_passed(_run_and_parse(args, capsys))
self._assert_all_passed(_run_and_parse(args, capsys)[0])
def test_override_config_yaml(self, tmp_path: Path, capsys) -> None:
"""--override-config YAML overrides dims."""
@@ -3390,7 +3396,7 @@ class TestEntrypointMetaOverride:
grouping="raw",
override_config=str(yaml_path),
)
self._assert_all_passed(_run_and_parse(args, capsys))
self._assert_all_passed(_run_and_parse(args, capsys)[0])
def test_no_match_uses_original_dims(self, tmp_path: Path, capsys) -> None:
"""When override regex doesn't match, original dims from dump are used."""
@@ -3406,7 +3412,7 @@ class TestEntrypointMetaOverride:
grouping="raw",
override_dims=["no_match_pattern:b s d"],
)
self._assert_all_passed(_run_and_parse(args, capsys))
self._assert_all_passed(_run_and_parse(args, capsys)[0])
def test_selective_match_multi_tensor(self, tmp_path: Path, capsys) -> None:
"""Override matches only 'logits'; 'hidden' uses original dims."""
@@ -3437,7 +3443,7 @@ class TestEntrypointMetaOverride:
grouping="raw",
override_dims=["logits:t v"],
)
self._assert_all_passed(_run_and_parse(args, capsys), expected_count=2)
self._assert_all_passed(_run_and_parse(args, capsys)[0], expected_count=2)
def test_multiple_cli_override_dims(self, tmp_path: Path, capsys) -> None:
"""Multiple --override-dims for different tensors."""
@@ -3470,7 +3476,7 @@ class TestEntrypointMetaOverride:
grouping="raw",
override_dims=["hidden:t h", "logits:t v"],
)
self._assert_all_passed(_run_and_parse(args, capsys), expected_count=2)
self._assert_all_passed(_run_and_parse(args, capsys)[0], expected_count=2)
def test_per_side_dims_different_parallelism(self, tmp_path: Path, capsys) -> None:
"""baseline TP-sharded, target EP-sharded — per-side override fixes both."""
@@ -3512,7 +3518,7 @@ class TestEntrypointMetaOverride:
override_baseline_dims=["hidden:t h(tp)"],
override_target_dims=["hidden:t h(ep)"],
)
self._assert_all_passed(_run_and_parse(args, capsys))
self._assert_all_passed(_run_and_parse(args, capsys)[0])
def test_yaml_first_match_wins_e2e(self, tmp_path: Path, capsys) -> None:
"""YAML with two matching rules: first rule wins in real pipeline."""
@@ -3533,7 +3539,7 @@ class TestEntrypointMetaOverride:
grouping="raw",
override_config=str(yaml_path),
)
self._assert_all_passed(_run_and_parse(args, capsys))
self._assert_all_passed(_run_and_parse(args, capsys)[0])
def test_cli_overrides_yaml_e2e(self, tmp_path: Path, capsys) -> None:
"""CLI --override-dims wins over YAML rule for the same tensor."""
@@ -3553,7 +3559,7 @@ class TestEntrypointMetaOverride:
override_dims=["hidden:t h"],
override_config=str(yaml_path),
)
self._assert_all_passed(_run_and_parse(args, capsys))
self._assert_all_passed(_run_and_parse(args, capsys)[0])
def test_override_injects_dims_when_absent(self, tmp_path: Path, capsys) -> None:
"""Override injects dims into meta even when dump had no dims annotation."""
@@ -3569,7 +3575,7 @@ class TestEntrypointMetaOverride:
grouping="raw",
override_dims=["hidden:t h"],
)
self._assert_all_passed(_run_and_parse(args, capsys))
self._assert_all_passed(_run_and_parse(args, capsys)[0])
def test_non_tensor_unaffected_by_override(self, tmp_path: Path, capsys) -> None:
"""Non-tensor values pass through without error even with active override."""
@@ -3596,7 +3602,7 @@ class TestEntrypointMetaOverride:
grouping="raw",
override_dims=["hidden:x y"],
)
records = _run_and_parse(args, capsys)
records, _ = _run_and_parse(args, capsys)
non_tensors: list[NonTensorRecord] = [
r for r in records if isinstance(r, NonTensorRecord)
@@ -3613,5 +3619,297 @@ class TestEntrypointMetaOverride:
assert summary.failed == 0
class TestExitCode:
"""Tests for exit code behavior based on comparison results."""
def test_all_passed(self):
"""All passed → exit 0."""
summary = SummaryRecord(total=3, passed=3, failed=0, skipped=0)
assert (
_compute_exit_code(summary, allow_skip_pattern=".*", skipped_names=[]) == 0
)
def test_has_failed_and_passed(self):
"""Has failed and passed → exit 1."""
summary = SummaryRecord(total=4, passed=2, failed=2, skipped=0)
assert (
_compute_exit_code(summary, allow_skip_pattern=".*", skipped_names=[]) == 1
)
def test_all_failed(self):
"""All failed (0 passed) → exit 1."""
summary = SummaryRecord(total=3, passed=0, failed=3, skipped=0)
assert (
_compute_exit_code(summary, allow_skip_pattern=".*", skipped_names=[]) == 1
)
def test_all_skipped_allow_all(self):
"""All skipped + allow_skip_pattern='.*' → exit 0."""
summary = SummaryRecord(total=2, passed=0, failed=0, skipped=2)
assert (
_compute_exit_code(
summary, allow_skip_pattern=".*", skipped_names=["a", "b"]
)
== 0
)
def test_all_skipped_forbid_all(self):
"""All skipped + allow_skip_pattern='^$' → exit 1."""
summary = SummaryRecord(total=2, passed=0, failed=0, skipped=2)
assert (
_compute_exit_code(
summary, allow_skip_pattern="^$", skipped_names=["a", "b"]
)
== 1
)
def test_passed_and_skipped_allow_all(self):
"""Passed + skipped, allow all → exit 0."""
summary = SummaryRecord(total=3, passed=2, failed=0, skipped=1)
assert (
_compute_exit_code(summary, allow_skip_pattern=".*", skipped_names=["a"])
== 0
)
def test_passed_and_skipped_forbid_all(self):
"""Passed + skipped + forbid all → exit 1."""
summary = SummaryRecord(total=3, passed=2, failed=0, skipped=1)
assert (
_compute_exit_code(summary, allow_skip_pattern="^$", skipped_names=["a"])
== 1
)
def test_skip_pattern_matches_specific_name(self):
"""Pattern matching specific name allows that skip, forbids others."""
summary = SummaryRecord(total=4, passed=2, failed=0, skipped=2)
assert (
_compute_exit_code(
summary,
allow_skip_pattern="positions|seq_lens",
skipped_names=["positions", "seq_lens"],
)
== 0
)
def test_skip_pattern_partial_match_forbidden(self):
"""Pattern matches some skips but not all → exit 1."""
summary = SummaryRecord(total=4, passed=1, failed=0, skipped=3)
assert (
_compute_exit_code(
summary,
allow_skip_pattern="positions|seq_lens",
skipped_names=["positions", "seq_lens", "hidden_states"],
)
== 1
)
def test_e2e_all_passed_exit_zero(self, tmp_path, capsys):
"""Integration: all comparisons pass → run() returns 0."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
args = _make_args(baseline_path, target_path, grouping="raw")
records, exit_code = _run_and_parse(args, capsys)
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.passed == 2
assert summary.failed == 0
assert exit_code == 0
def test_e2e_has_failed_exit_nonzero(self, tmp_path, capsys):
"""Integration: a failed comparison → run() returns 1."""
torch.manual_seed(42)
baseline_path = _create_rank_dump(
tmp_path / "baseline", rank=0, name="tensor_a", tensor=torch.randn(10, 10)
)
target_path = _create_rank_dump(
tmp_path / "target",
rank=0,
name="tensor_a",
tensor=torch.randn(10, 10) * 100,
)
args = _make_args(
baseline_path, target_path, grouping="raw", diff_threshold=1e-3
)
records, exit_code = _run_and_parse(args, capsys)
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.failed == 1
assert exit_code == 1
class TestExitCodeSubprocess:
"""E2E subprocess tests: invoke comparator as a child process and verify exit code."""
@staticmethod
def _run_comparator(
baseline_path: Path,
target_path: Path,
*,
grouping: str = "raw",
allow_skip_pattern: str = ".*",
) -> subprocess.CompletedProcess[str]:
cmd: list[str] = [
sys.executable,
"-m",
"sglang.srt.debug_utils.comparator",
"--baseline-path",
str(baseline_path),
"--target-path",
str(target_path),
"--grouping",
grouping,
"--output-format",
"json",
"--allow-skip-pattern",
allow_skip_pattern,
]
return subprocess.run(cmd, capture_output=True, text=True)
def test_all_passed_exit_zero(self, tmp_path):
"""Subprocess: all comparisons pass → exit 0."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
result = self._run_comparator(baseline_path, target_path)
assert result.returncode == 0
def test_failed_exit_nonzero(self, tmp_path):
"""Subprocess: failed comparison → exit 1."""
torch.manual_seed(42)
baseline_path = _create_rank_dump(
tmp_path / "baseline", rank=0, name="t", tensor=torch.randn(10, 10)
)
target_path = _create_rank_dump(
tmp_path / "target", rank=0, name="t", tensor=torch.randn(10, 10) * 100
)
result = self._run_comparator(baseline_path, target_path)
assert result.returncode == 1
def test_skipped_allow_all_exit_zero(self, tmp_path):
"""Subprocess: skipped comparison with allow_skip_pattern='.*' → exit 0."""
baseline_path, target_path = _create_dumps(
tmp_path,
tensor_names=["tensor_a", "tensor_extra"],
baseline_names=["tensor_a"],
)
result = self._run_comparator(
baseline_path, target_path, allow_skip_pattern=".*"
)
assert result.returncode == 0
def test_skipped_forbid_all_exit_nonzero(self, tmp_path):
"""Subprocess: skipped comparison with allow_skip_pattern='^$' → exit 1."""
baseline_path, target_path = _create_dumps(
tmp_path,
tensor_names=["tensor_a", "tensor_extra"],
baseline_names=["tensor_a"],
)
result = self._run_comparator(
baseline_path, target_path, allow_skip_pattern="^$"
)
assert result.returncode == 1
class TestReportOutput:
"""Test JSONL report file output via ReportSink."""
def test_default_report_path(self, tmp_path, capsys):
"""Default writes to <target>/comparator_report.jsonl with ConfigRecord + SummaryRecord."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
args = _make_args(baseline_path, target_path, grouping="raw", report_path=None)
exit_code: int = run(args)
report_file: Path = target_path / "comparator_report.jsonl"
assert report_file.exists()
report_records: list[AnyRecord] = _parse_jsonl(report_file.read_text())
assert isinstance(report_records[0], ConfigRecord)
assert isinstance(report_records[-1], SummaryRecord)
assert exit_code == 0
def test_custom_report_path(self, tmp_path, capsys):
"""--report-path writes to the specified location."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
custom_path: Path = tmp_path / "custom" / "report.jsonl"
args = _make_args(
baseline_path,
target_path,
grouping="raw",
report_path=str(custom_path),
)
run(args)
assert custom_path.exists()
report_records: list[AnyRecord] = _parse_jsonl(custom_path.read_text())
assert isinstance(report_records[0], ConfigRecord)
assert isinstance(report_records[-1], SummaryRecord)
def test_disabled_report(self, tmp_path, capsys):
"""--report-path '' disables file generation."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
args = _make_args(baseline_path, target_path, grouping="raw", report_path="")
run(args)
report_file: Path = target_path / "comparator_report.jsonl"
assert not report_file.exists()
def test_report_matches_stdout_json(self, tmp_path, capsys):
"""In json mode, report content matches stdout output."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
report_file: Path = tmp_path / "report.jsonl"
args = _make_args(
baseline_path,
target_path,
grouping="raw",
output_format="json",
report_path=str(report_file),
)
capsys.readouterr()
run(args)
stdout_lines: list[str] = capsys.readouterr().out.strip().splitlines()
report_lines: list[str] = report_file.read_text().strip().splitlines()
assert stdout_lines == report_lines
def test_text_mode_also_writes_report(self, tmp_path, capsys):
"""Text stdout mode still writes JSONL report."""
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
report_file: Path = tmp_path / "report.jsonl"
args = _make_args(
baseline_path,
target_path,
grouping="raw",
output_format="text",
report_path=str(report_file),
)
run(args)
assert report_file.exists()
report_records: list[AnyRecord] = _parse_jsonl(report_file.read_text())
assert isinstance(report_records[0], ConfigRecord)
assert isinstance(report_records[-1], SummaryRecord)
def test_streaming_flush(self, tmp_path, capsys):
"""Report file is flushed after each record (readable before close)."""
from sglang.srt.debug_utils.comparator.output_types import report_sink
report_file: Path = tmp_path / "stream_report.jsonl"
report_sink.configure(
output_format="json",
report_path=report_file,
)
report_sink.add(ConfigRecord(config={"test": True}))
content: str = report_file.read_text()
assert len(content.strip().splitlines()) == 1
parsed: AnyRecord = parse_record_json(content.strip())
assert isinstance(parsed, ConfigRecord)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))

13
test/registered/debug_utils/comparator/test_warning_sink.py
View File

@@ -3,7 +3,10 @@ import sys
import pytest
from sglang.srt.debug_utils.comparator.output_types import GeneralWarning
from sglang.srt.debug_utils.comparator.output_types import (
GeneralWarning,
report_sink,
)
from sglang.srt.debug_utils.comparator.warning_sink import WarningSink
from sglang.test.ci.ci_register import register_cpu_ci
@@ -53,7 +56,7 @@ class TestWarningSink:
def test_add_outside_context_prints(self, capsys) -> None:
sink = WarningSink()
sink.set_output_format("text")
report_sink.configure(output_format="text")
sink.add(_make_warning())
@@ -62,7 +65,7 @@ class TestWarningSink:
def test_context_captures_instead_of_printing(self, capsys) -> None:
sink = WarningSink()
sink.set_output_format("text")
report_sink.configure(output_format="text")
with sink.context() as collected:
sink.add(_make_warning())
@@ -73,7 +76,7 @@ class TestWarningSink:
def test_json_output_outside_context(self, capsys) -> None:
sink = WarningSink()
sink.set_output_format("json")
report_sink.configure(output_format="json")
sink.add(_make_warning())
@@ -84,7 +87,7 @@ class TestWarningSink:
def test_exception_in_context_cleans_stack(self, capsys) -> None:
sink = WarningSink()
sink.set_output_format("text")
report_sink.configure(output_format="text")
with pytest.raises(RuntimeError):
with sink.context() as collected:

2
test/registered/debug_utils/test_engine_dumper_comparator_e2e.py
View File

@@ -34,7 +34,7 @@ register_cuda_ci(est_time=120, suite="nightly-2-gpu", nightly=True)
MODEL = "Qwen/Qwen3-0.6B"
EXP_NAME = "e2e_source_patcher"
DUMPER_FILTER = r"layer_id=[012]"
DUMPER_FILTER = "layer_id in [0, 1, 2]"
PATCH_CONFIG_YAML: str = """\
patches: