106 lines
3.0 KiB
Python
106 lines
3.0 KiB
Python
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__]))
|