diff --git a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py index 683d28f0d..021990b51 100644 --- a/python/sglang/srt/debug_utils/comparator/bundle_comparator.py +++ b/python/sglang/srt/debug_utils/comparator/bundle_comparator.py @@ -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 diff --git a/python/sglang/srt/debug_utils/comparator/display.py b/python/sglang/srt/debug_utils/comparator/display.py index 56a9f46b7..0ef7e5ab6 100644 --- a/python/sglang/srt/debug_utils/comparator/display.py +++ b/python/sglang/srt/debug_utils/comparator/display.py @@ -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: diff --git a/python/sglang/srt/debug_utils/comparator/entrypoint.py b/python/sglang/srt/debug_utils/comparator/entrypoint.py index ced54e0bc..ce52aa57a 100644 --- a/python/sglang/srt/debug_utils/comparator/entrypoint.py +++ b/python/sglang/srt/debug_utils/comparator/entrypoint.py @@ -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: /comparator_report.jsonl). " + "Pass empty string '' to disable.", + ) return parser.parse_args() diff --git a/python/sglang/srt/debug_utils/comparator/output_types.py b/python/sglang/srt/debug_utils/comparator/output_types.py index c83db102e..f75179276 100644 --- a/python/sglang/srt/debug_utils/comparator/output_types.py +++ b/python/sglang/srt/debug_utils/comparator/output_types.py @@ -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() diff --git a/python/sglang/srt/debug_utils/comparator/warning_sink.py b/python/sglang/srt/debug_utils/comparator/warning_sink.py index 1c443f8a6..a5e106f42 100644 --- a/python/sglang/srt/debug_utils/comparator/warning_sink.py +++ b/python/sglang/srt/debug_utils/comparator/warning_sink.py @@ -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() diff --git a/test/registered/debug_utils/comparator/conftest.py b/test/registered/debug_utils/comparator/conftest.py index e69de29bb..f356008e3 100644 --- a/test/registered/debug_utils/comparator/conftest.py +++ b/test/registered/debug_utils/comparator/conftest.py @@ -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() diff --git a/test/registered/debug_utils/comparator/test_bundle_comparator.py b/test/registered/debug_utils/comparator/test_bundle_comparator.py new file mode 100644 index 000000000..f9491b9c8 --- /dev/null +++ b/test/registered/debug_utils/comparator/test_bundle_comparator.py @@ -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__])) diff --git a/test/registered/debug_utils/comparator/test_entrypoint.py b/test/registered/debug_utils/comparator/test_entrypoint.py index 7d705816e..add958808 100644 --- a/test/registered/debug_utils/comparator/test_entrypoint.py +++ b/test/registered/debug_utils/comparator/test_entrypoint.py @@ -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 /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__])) diff --git a/test/registered/debug_utils/comparator/test_warning_sink.py b/test/registered/debug_utils/comparator/test_warning_sink.py index 7a41dcf4c..efc66cc32 100644 --- a/test/registered/debug_utils/comparator/test_warning_sink.py +++ b/test/registered/debug_utils/comparator/test_warning_sink.py @@ -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: diff --git a/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py b/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py index 59477f910..05d46761d 100644 --- a/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py +++ b/test/registered/debug_utils/test_engine_dumper_comparator_e2e.py @@ -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: