Handle warnings via sink for structured output and add pair in dump comparator (#19373)
This commit is contained in:
fzyzcjy
@@ -98,11 +98,13 @@ class TestRecordTypes:
|
||||
def test_discriminated_union_parsing(self):
|
||||
for record in [
|
||||
ConfigRecord(
|
||||
baseline_path="/a",
|
||||
target_path="/b",
|
||||
diff_threshold=1e-3,
|
||||
start_step=0,
|
||||
end_step=100,
|
||||
config={
|
||||
"baseline_path": "/a",
|
||||
"target_path": "/b",
|
||||
"diff_threshold": 1e-3,
|
||||
"start_step": 0,
|
||||
"end_step": 100,
|
||||
},
|
||||
),
|
||||
SkipRecord(name="attn", reason="no_baseline"),
|
||||
ComparisonRecord(
|
||||
@@ -133,7 +135,7 @@ def _make_warning(**overrides) -> ReplicatedMismatchWarning:
|
||||
|
||||
class TestAlignWarnings:
|
||||
def test_comparison_record_failed_when_diff_passed_but_warnings(self):
|
||||
"""ComparisonRecord with diff.passed=True but align_warnings → category=='failed'."""
|
||||
"""ComparisonRecord with diff.passed=True but warnings → category=='failed'."""
|
||||
record = ComparisonRecord(
|
||||
name="hidden",
|
||||
baseline=_make_tensor_info(),
|
||||
@@ -141,21 +143,21 @@ class TestAlignWarnings:
|
||||
unified_shape=[4, 8],
|
||||
shape_mismatch=False,
|
||||
diff=_make_diff(passed=True),
|
||||
align_warnings=[_make_warning()],
|
||||
warnings=[_make_warning()],
|
||||
)
|
||||
assert record.category == "failed"
|
||||
|
||||
def test_skip_record_failed_when_warnings(self):
|
||||
"""SkipRecord with align_warnings → category=='failed' instead of 'skipped'."""
|
||||
"""SkipRecord with warnings → category=='failed' instead of 'skipped'."""
|
||||
record = SkipRecord(
|
||||
name="x",
|
||||
reason="no_baseline",
|
||||
align_warnings=[_make_warning()],
|
||||
warnings=[_make_warning()],
|
||||
)
|
||||
assert record.category == "failed"
|
||||
|
||||
def test_align_warnings_json_round_trip(self):
|
||||
"""align_warnings survive model_dump_json → parse_record_json round-trip."""
|
||||
def test_warnings_json_round_trip(self):
|
||||
"""warnings survive model_dump_json → parse_record_json round-trip."""
|
||||
warning = _make_warning(
|
||||
axis="cp",
|
||||
group_index=2,
|
||||
@@ -170,14 +172,14 @@ class TestAlignWarnings:
|
||||
unified_shape=[4, 8],
|
||||
shape_mismatch=False,
|
||||
diff=_make_diff(),
|
||||
align_warnings=[warning],
|
||||
warnings=[warning],
|
||||
)
|
||||
|
||||
restored = parse_record_json(record.model_dump_json())
|
||||
assert isinstance(restored, ComparisonRecord)
|
||||
assert len(restored.align_warnings) == 1
|
||||
assert len(restored.warnings) == 1
|
||||
|
||||
restored_warning = restored.align_warnings[0]
|
||||
restored_warning = restored.warnings[0]
|
||||
assert restored_warning.axis == "cp"
|
||||
assert restored_warning.group_index == 2
|
||||
assert restored_warning.differing_index == 3
|
||||
|
||||
@@ -540,7 +540,7 @@ class TestEntrypointGroupingLogical:
|
||||
assert summary.skipped == 0
|
||||
|
||||
def test_multi_step_tp(self, tmp_path, capsys):
|
||||
"""Two steps with TP=2 shards produce two logical groups (one per step)."""
|
||||
"""Two steps with TP=2 shards produce two per-step comparisons (no aux → no alignment)."""
|
||||
torch.manual_seed(42)
|
||||
full_tensor = torch.randn(4, 8)
|
||||
|
||||
@@ -571,6 +571,8 @@ class TestEntrypointGroupingLogical:
|
||||
records = _run_and_parse(args, capsys)
|
||||
comparisons = _get_comparisons(records)
|
||||
assert len(comparisons) == 2
|
||||
assert comparisons[0].baseline.shape == [4, 8]
|
||||
assert comparisons[1].baseline.shape == [4, 8]
|
||||
|
||||
summary = records[-1]
|
||||
assert isinstance(summary, SummaryRecord)
|
||||
@@ -612,7 +614,7 @@ class TestEntrypointGroupingLogical:
|
||||
assert comp.name == "attn_out"
|
||||
|
||||
def test_filter_logical(self, tmp_path, capsys):
|
||||
"""--filter in logical grouping selects only matching tensor groups."""
|
||||
"""--filter in logical grouping selects only matching tensor bundles."""
|
||||
torch.manual_seed(42)
|
||||
full_a = torch.randn(4, 8)
|
||||
full_b = torch.randn(4, 8)
|
||||
@@ -736,7 +738,7 @@ class TestEntrypointGroupingLogical:
|
||||
assert comp.name == "hidden"
|
||||
|
||||
def test_cp_tp_different_sizes(self, tmp_path, capsys):
|
||||
"""Baseline CP=2+TP=2 vs target CP=1+TP=4: both sides independently unshard."""
|
||||
"""Baseline CP=2+TP=2 vs target CP=1+TP=4: both sides independently unsharder."""
|
||||
torch.manual_seed(42)
|
||||
full_baseline = torch.randn(4, 8, 16)
|
||||
full_target = full_baseline + torch.randn(4, 8, 16) * 0.001
|
||||
@@ -882,7 +884,7 @@ class TestEntrypointReplicatedAxis:
|
||||
"""Test replicated-axis scenarios through the full entrypoint pipeline."""
|
||||
|
||||
def test_replicated_axis_identical_replicas_passed(self, tmp_path, capsys):
|
||||
"""CP2 TP2, TP replicated and identical → passed, no align_warnings."""
|
||||
"""CP2 TP2, TP replicated and identical → passed, no warnings."""
|
||||
torch.manual_seed(42)
|
||||
full_baseline = torch.randn(4, 8, 6)
|
||||
full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
|
||||
@@ -912,14 +914,14 @@ class TestEntrypointReplicatedAxis:
|
||||
|
||||
records = _run_and_parse(args, capsys)
|
||||
comp = _assert_single_comparison_passed(records)
|
||||
assert comp.align_warnings == []
|
||||
assert comp.warnings == []
|
||||
|
||||
summary = records[-1]
|
||||
assert isinstance(summary, SummaryRecord)
|
||||
assert summary.passed == 1
|
||||
|
||||
def test_replicated_mismatch_fails(self, tmp_path, capsys):
|
||||
"""CP2 TP2, TP replicas differ (> atol) → failed with align_warnings."""
|
||||
"""CP2 TP2, TP replicas differ (> atol) → failed with warnings."""
|
||||
torch.manual_seed(42)
|
||||
full_baseline = torch.randn(4, 8, 6)
|
||||
full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
|
||||
@@ -952,14 +954,14 @@ class TestEntrypointReplicatedAxis:
|
||||
comparisons = _get_comparisons(records)
|
||||
assert len(comparisons) == 1
|
||||
assert comparisons[0].category == "failed"
|
||||
assert len(comparisons[0].align_warnings) > 0
|
||||
assert len(comparisons[0].warnings) > 0
|
||||
|
||||
summary = records[-1]
|
||||
assert isinstance(summary, SummaryRecord)
|
||||
assert summary.failed == 1
|
||||
|
||||
def test_summary_counts_failed_from_align_warnings_only(self, tmp_path, capsys):
|
||||
"""Diff itself passes but TP replicas differ → summary.failed=1 from align_warnings."""
|
||||
def test_summary_counts_failed_from_warnings_only(self, tmp_path, capsys):
|
||||
"""Diff itself passes but TP replicas differ → summary.failed=1 from warnings."""
|
||||
torch.manual_seed(42)
|
||||
full_baseline = torch.randn(4, 8, 6)
|
||||
full_target = full_baseline + torch.randn(4, 8, 6) * 0.0001
|
||||
@@ -1001,7 +1003,7 @@ class TestEntrypointReplicatedAxis:
|
||||
comp = comparisons[0]
|
||||
assert comp.diff is not None
|
||||
assert comp.diff.passed
|
||||
assert len(comp.align_warnings) > 0
|
||||
assert len(comp.warnings) > 0
|
||||
assert comp.category == "failed"
|
||||
|
||||
summary = records[-1]
|
||||
|
||||
@@ -0,0 +1,114 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from sglang.srt.debug_utils.comparator.output_types import (
|
||||
ComparisonRecord,
|
||||
GeneralWarning,
|
||||
SkipRecord,
|
||||
SummaryRecord,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
|
||||
DiffInfo,
|
||||
TensorInfo,
|
||||
TensorStats,
|
||||
)
|
||||
from sglang.srt.debug_utils.comparator.utils import _check_equal_lengths
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=10, suite="default", nightly=True)
|
||||
|
||||
|
||||
class TestCheckEqualLengths:
|
||||
def test_all_equal(self):
|
||||
_check_equal_lengths(a=[1, 2], b=[3, 4])
|
||||
|
||||
def test_empty_lists(self):
|
||||
_check_equal_lengths(a=[], b=[])
|
||||
|
||||
def test_mismatch_raises(self):
|
||||
with pytest.raises(ValueError, match="Length mismatch"):
|
||||
_check_equal_lengths(a=[1, 2], b=[3])
|
||||
|
||||
|
||||
class TestSummaryRecord:
|
||||
def test_valid(self):
|
||||
record = SummaryRecord(total=10, passed=7, failed=2, skipped=1)
|
||||
assert record.total == 10
|
||||
|
||||
def test_total_mismatch(self):
|
||||
with pytest.raises(ValidationError, match="total=10"):
|
||||
SummaryRecord(total=10, passed=5, failed=2, skipped=1)
|
||||
|
||||
|
||||
def _make_tensor_info() -> TensorInfo:
|
||||
return TensorInfo(
|
||||
shape=[4, 4],
|
||||
dtype="float32",
|
||||
stats=TensorStats(mean=0.0, std=1.0, min=-2.0, max=2.0),
|
||||
)
|
||||
|
||||
|
||||
def _make_diff_info(*, passed: bool) -> DiffInfo:
|
||||
return DiffInfo(
|
||||
rel_diff=0.001,
|
||||
max_abs_diff=0.01,
|
||||
mean_abs_diff=0.005,
|
||||
max_diff_coord=[0, 0],
|
||||
baseline_at_max=1.0,
|
||||
target_at_max=1.01,
|
||||
passed=passed,
|
||||
)
|
||||
|
||||
|
||||
def _make_comparison_record(
|
||||
*,
|
||||
diff: DiffInfo | None,
|
||||
warnings: list | None = None,
|
||||
) -> ComparisonRecord:
|
||||
ti: TensorInfo = _make_tensor_info()
|
||||
return ComparisonRecord(
|
||||
name="t",
|
||||
baseline=ti,
|
||||
target=ti,
|
||||
unified_shape=[4, 4],
|
||||
shape_mismatch=False,
|
||||
diff=diff,
|
||||
warnings=warnings or [],
|
||||
)
|
||||
|
||||
|
||||
class TestOutputRecordCategories:
|
||||
def test_skip_record_with_warnings_is_failed(self) -> None:
|
||||
record = SkipRecord(
|
||||
name="t",
|
||||
reason="test",
|
||||
warnings=[GeneralWarning(category="c", message="m")],
|
||||
)
|
||||
assert record.category == "failed"
|
||||
|
||||
def test_skip_record_no_warnings_is_skipped(self) -> None:
|
||||
record = SkipRecord(name="t", reason="test")
|
||||
assert record.category == "skipped"
|
||||
|
||||
def test_comparison_record_diff_none_is_failed(self) -> None:
|
||||
record: ComparisonRecord = _make_comparison_record(diff=None)
|
||||
assert record.category == "failed"
|
||||
|
||||
def test_comparison_record_passed_with_warnings_is_failed(self) -> None:
|
||||
record: ComparisonRecord = _make_comparison_record(
|
||||
diff=_make_diff_info(passed=True),
|
||||
warnings=[GeneralWarning(category="c", message="m")],
|
||||
)
|
||||
assert record.category == "failed"
|
||||
|
||||
def test_comparison_record_passed_no_warnings_is_passed(self) -> None:
|
||||
record: ComparisonRecord = _make_comparison_record(
|
||||
diff=_make_diff_info(passed=True),
|
||||
)
|
||||
assert record.category == "passed"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -0,0 +1,105 @@
|
||||
import json
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
from sglang.srt.debug_utils.comparator.output_types import ReplicatedMismatchWarning
|
||||
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=10, suite="default", nightly=True)
|
||||
|
||||
|
||||
def _make_warning(**overrides) -> ReplicatedMismatchWarning:
|
||||
defaults: dict = dict(
|
||||
axis="tp",
|
||||
group_index=0,
|
||||
differing_index=1,
|
||||
baseline_index=0,
|
||||
max_abs_diff=0.1,
|
||||
)
|
||||
defaults.update(overrides)
|
||||
return ReplicatedMismatchWarning(**defaults)
|
||||
|
||||
|
||||
class TestWarningSink:
|
||||
def test_basic_collection(self) -> None:
|
||||
sink = WarningSink()
|
||||
warning = _make_warning()
|
||||
|
||||
with sink.context() as collected:
|
||||
sink.add(warning)
|
||||
|
||||
assert len(collected) == 1
|
||||
assert collected[0] is warning
|
||||
|
||||
def test_nested_contexts(self) -> None:
|
||||
sink = WarningSink()
|
||||
outer_warning = _make_warning(group_index=0)
|
||||
inner_warning = _make_warning(group_index=1)
|
||||
|
||||
with sink.context() as outer:
|
||||
sink.add(outer_warning)
|
||||
with sink.context() as inner:
|
||||
sink.add(inner_warning)
|
||||
assert len(inner) == 1
|
||||
assert inner[0] is inner_warning
|
||||
|
||||
assert len(outer) == 1
|
||||
assert outer[0] is outer_warning
|
||||
|
||||
def test_empty_context(self) -> None:
|
||||
sink = WarningSink()
|
||||
with sink.context() as collected:
|
||||
pass
|
||||
assert collected == []
|
||||
|
||||
def test_add_outside_context_prints(self, capsys) -> None:
|
||||
sink = WarningSink()
|
||||
sink.set_output_format("text")
|
||||
|
||||
sink.add(_make_warning())
|
||||
|
||||
captured = capsys.readouterr()
|
||||
assert "Replicated along tp" in captured.out
|
||||
|
||||
def test_context_captures_instead_of_printing(self, capsys) -> None:
|
||||
sink = WarningSink()
|
||||
sink.set_output_format("text")
|
||||
|
||||
with sink.context() as collected:
|
||||
sink.add(_make_warning())
|
||||
|
||||
assert len(collected) == 1
|
||||
captured = capsys.readouterr()
|
||||
assert captured.out == ""
|
||||
|
||||
def test_json_output_outside_context(self, capsys) -> None:
|
||||
sink = WarningSink()
|
||||
sink.set_output_format("json")
|
||||
|
||||
sink.add(_make_warning())
|
||||
|
||||
captured = capsys.readouterr()
|
||||
parsed: dict = json.loads(captured.out.strip())
|
||||
assert "warnings" in parsed
|
||||
assert len(parsed["warnings"]) == 1
|
||||
|
||||
def test_exception_in_context_cleans_stack(self, capsys) -> None:
|
||||
sink = WarningSink()
|
||||
sink.set_output_format("text")
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
with sink.context() as collected:
|
||||
sink.add(_make_warning())
|
||||
raise RuntimeError("boom")
|
||||
|
||||
assert len(collected) == 1
|
||||
|
||||
sink.add(_make_warning(group_index=99))
|
||||
captured = capsys.readouterr()
|
||||
assert "Replicated along tp" in captured.out
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
Reference in New Issue
Block a user