Support agent-friendly output formats in dump comparator (#19275)

This commit is contained in:
fzyzcjy
2026-02-25 09:40:33 +08:00
committed by GitHub
parent d7578ce279
commit 8b1ab4aaf9
11 changed files with 730 additions and 71 deletions
@@ -53,6 +53,7 @@ class TestComputeDiff:
assert diff.rel_diff == pytest.approx(0.0, abs=1e-5)
assert diff.max_abs_diff == pytest.approx(0.0, abs=1e-5)
assert diff.mean_abs_diff == pytest.approx(0.0, abs=1e-5)
assert diff.passed is True
def test_known_offset(self):
x = torch.ones(10, 10)
@@ -62,10 +63,11 @@ class TestComputeDiff:
diff = _compute_diff(x_baseline=x, x_target=y)
assert diff.max_abs_diff == pytest.approx(0.5, abs=1e-4)
assert diff.max_diff_coord == (3, 7)
assert diff.max_diff_coord == [3, 7]
assert diff.baseline_at_max == pytest.approx(1.0, abs=1e-4)
assert diff.target_at_max == pytest.approx(1.5, abs=1e-4)
assert diff.mean_abs_diff == pytest.approx(0.5 / 100, abs=1e-4)
assert diff.passed is False
def test_rel_diff_value(self):
x = torch.tensor([1.0, 0.0])
@@ -73,6 +75,7 @@ class TestComputeDiff:
diff = _compute_diff(x_baseline=x, x_target=y)
assert diff.rel_diff == pytest.approx(1.0, abs=1e-5)
assert diff.passed is False
class TestCompareTensors:
@@ -83,8 +86,8 @@ class TestCompareTensors:
info = compare_tensors(x_baseline=x, x_target=y, name="test")
assert info.name == "test"
assert info.baseline.shape == torch.Size([5, 5])
assert info.target.shape == torch.Size([5, 5])
assert info.baseline.shape == [5, 5]
assert info.target.shape == [5, 5]
assert info.shape_mismatch is False
assert info.diff is not None
assert info.diff_downcast is None
@@ -107,7 +110,7 @@ class TestCompareTensors:
assert info.shape_mismatch is False
assert info.diff is not None
assert info.diff_downcast is not None
assert info.downcast_dtype == torch.bfloat16
assert info.downcast_dtype == "torch.bfloat16"
def test_shape_unification(self):
torch.manual_seed(0)
@@ -117,8 +120,8 @@ class TestCompareTensors:
info = compare_tensors(x_baseline=x, x_target=y, name="unify")
assert info.baseline.shape == torch.Size([1, 1, 4, 8])
assert info.unified_shape == torch.Size([4, 8])
assert info.baseline.shape == [1, 1, 4, 8]
assert info.unified_shape == [4, 8]
assert info.shape_mismatch is False
assert info.diff is not None
assert info.diff.max_abs_diff == pytest.approx(0.0, abs=1e-5)
@@ -0,0 +1,253 @@
import sys
import pytest
from sglang.srt.debug_utils.comparator.tensor_comparison.formatter import (
format_comparison,
)
from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
DiffInfo,
TensorComparisonInfo,
TensorInfo,
TensorStats,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
def _make_stats(
mean: float = 0.0,
std: float = 1.0,
min: float = -2.0,
max: float = 2.0,
p1: float | None = -1.8,
p5: float | None = -1.5,
p95: float | None = 1.5,
p99: float | None = 1.8,
) -> TensorStats:
return TensorStats(
mean=mean, std=std, min=min, max=max, p1=p1, p5=p5, p95=p95, p99=p99
)
def _make_diff(
rel_diff: float = 0.0001,
max_abs_diff: float = 0.0005,
mean_abs_diff: float = 0.0002,
passed: bool = True,
) -> DiffInfo:
return DiffInfo(
rel_diff=rel_diff,
max_abs_diff=max_abs_diff,
mean_abs_diff=mean_abs_diff,
max_diff_coord=[2, 3],
baseline_at_max=1.0,
target_at_max=1.0005,
passed=passed,
)
def _make_tensor_info(
shape: list[int] | None = None,
dtype: str = "torch.float32",
stats: TensorStats | None = None,
sample: str | None = None,
) -> TensorInfo:
return TensorInfo(
shape=shape if shape is not None else [4, 8],
dtype=dtype,
stats=stats if stats is not None else _make_stats(),
sample=sample,
)
# Snapshot strings below are intentionally spelled out in full per test.
# The shared skeleton (stats block, diff block) looks duplicated, but keeping
# each test self-contained makes failures immediately readable without chasing
# helper functions. Do not extract common fragments.
class TestFormatComparison:
def test_normal(self):
info = TensorComparisonInfo(
name="test",
baseline=_make_tensor_info(
stats=_make_stats(mean=0.1, std=1.0, min=-2.0, max=2.0),
),
target=_make_tensor_info(
stats=_make_stats(mean=0.1001, std=1.0001, min=-2.0001, max=2.0001),
),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
assert format_comparison(info) == (
"Raw [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"[mean] 0.1000 vs 0.1001 (diff: 0.0001)\n"
"[std] 1.0000 vs 1.0001 (diff: 0.0001)\n"
"[min] -2.0000 vs -2.0001 (diff: -0.0001)\n"
"[max] 2.0000 vs 2.0001 (diff: 0.0001)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
"max_abs_diff happens at coord=[2, 3] with "
"baseline=1.0 target=1.0005"
)
def test_shape_mismatch(self):
info = TensorComparisonInfo(
name="mismatch",
baseline=_make_tensor_info(shape=[3, 4]),
target=_make_tensor_info(shape=[5, 6]),
unified_shape=[3, 4],
shape_mismatch=True,
)
assert format_comparison(info) == (
"Raw [shape] [3, 4] vs [5, 6]\t"
"[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] [3, 4] vs [5, 6]\t"
"[dtype] torch.float32 vs torch.float32\n"
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"⚠️ Shape mismatch"
)
def test_with_downcast(self):
info = TensorComparisonInfo(
name="downcast",
baseline=_make_tensor_info(),
target=_make_tensor_info(dtype="torch.bfloat16"),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(
rel_diff=0.002, max_abs_diff=0.005, mean_abs_diff=0.001, passed=False
),
diff_downcast=_make_diff(
rel_diff=0.0001, max_abs_diff=0.0005, mean_abs_diff=0.0002, passed=True
),
downcast_dtype="torch.bfloat16",
)
assert format_comparison(info) == (
"Raw [shape] [4, 8] vs [4, 8]\t"
"[🟠dtype] torch.float32 vs torch.bfloat16\n"
"After unify [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.bfloat16\n"
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"❌ rel_diff=0.002\tmax_abs_diff=0.005\tmean_abs_diff=0.001\n"
"max_abs_diff happens at coord=[2, 3] with "
"baseline=1.0 target=1.0005\n"
"When downcast to torch.bfloat16: "
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
"max_abs_diff happens at coord=[2, 3] with "
"baseline=1.0 target=1.0005"
)
def test_with_shape_unification(self):
info = TensorComparisonInfo(
name="unify",
baseline=_make_tensor_info(shape=[1, 1, 4, 8]),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
assert format_comparison(info) == (
"Raw [shape] [1, 1, 4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"Unify shape: [1, 1, 4, 8] -> [4, 8] "
"(to match [4, 8])\n"
"After unify [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
"max_abs_diff happens at coord=[2, 3] with "
"baseline=1.0 target=1.0005"
)
def test_with_samples(self):
info = TensorComparisonInfo(
name="samples",
baseline=_make_tensor_info(sample="tensor([0.1, 0.2, ...])"),
target=_make_tensor_info(sample="tensor([0.1, 0.3, ...])"),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
assert format_comparison(info) == (
"Raw [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"[p1] -1.8000 vs -1.8000 (diff: 0.0000)\n"
"[p5] -1.5000 vs -1.5000 (diff: 0.0000)\n"
"[p95] 1.5000 vs 1.5000 (diff: 0.0000)\n"
"[p99] 1.8000 vs 1.8000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
"max_abs_diff happens at coord=[2, 3] with "
"baseline=1.0 target=1.0005\n"
"x_baseline(sample)=tensor([0.1, 0.2, ...])\n"
"x_target(sample)=tensor([0.1, 0.3, ...])"
)
def test_none_quantiles(self):
stats_no_quantiles = _make_stats(p1=None, p5=None, p95=None, p99=None)
info = TensorComparisonInfo(
name="no_quantiles",
baseline=_make_tensor_info(stats=stats_no_quantiles),
target=_make_tensor_info(stats=stats_no_quantiles),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
assert format_comparison(info) == (
"Raw [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"After unify [shape] [4, 8] vs [4, 8]\t"
"[dtype] torch.float32 vs torch.float32\n"
"[mean] 0.0000 vs 0.0000 (diff: 0.0000)\n"
"[std] 1.0000 vs 1.0000 (diff: 0.0000)\n"
"[min] -2.0000 vs -2.0000 (diff: 0.0000)\n"
"[max] 2.0000 vs 2.0000 (diff: 0.0000)\n"
"✅ rel_diff=0.0001\tmax_abs_diff=0.0005\tmean_abs_diff=0.0002\n"
"max_abs_diff happens at coord=[2, 3] with "
"baseline=1.0 target=1.0005"
)
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
@@ -0,0 +1,122 @@
import json
import sys
import pytest
from sglang.srt.debug_utils.comparator.output_types import (
ComparisonRecord,
ConfigRecord,
SkipRecord,
SummaryRecord,
parse_record_json,
)
from sglang.srt.debug_utils.comparator.tensor_comparison.types import (
DiffInfo,
TensorInfo,
TensorStats,
)
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=10, suite="default", nightly=True)
def _make_stats(**overrides: float) -> TensorStats:
defaults = dict(
mean=0.5,
std=1.0,
min=-2.0,
max=3.0,
p1=-1.8,
p5=-1.5,
p95=2.5,
p99=2.8,
)
defaults.update(overrides)
return TensorStats(**defaults)
def _make_diff(**overrides) -> DiffInfo:
defaults = dict(
rel_diff=1e-4,
max_abs_diff=5e-4,
mean_abs_diff=2e-4,
max_diff_coord=[2, 3],
baseline_at_max=1.0,
target_at_max=1.0005,
passed=True,
)
defaults.update(overrides)
return DiffInfo(**defaults)
def _make_tensor_info(**overrides) -> TensorInfo:
defaults = dict(
shape=[4, 8],
dtype="torch.float32",
stats=_make_stats(),
)
defaults.update(overrides)
return TensorInfo(**defaults)
class TestStrictBase:
def test_rejects_extra_fields(self):
with pytest.raises(Exception):
TensorStats(mean=0.0, std=1.0, min=-1.0, max=1.0, bogus=42)
def test_rejects_extra_fields_on_diff(self):
with pytest.raises(Exception):
DiffInfo(
rel_diff=0.0,
max_abs_diff=0.0,
mean_abs_diff=0.0,
max_diff_coord=[0],
baseline_at_max=0.0,
target_at_max=0.0,
passed=True,
extra_field=123,
)
class TestRecordTypes:
def test_comparison_record_inherits_tensor_fields(self):
record = ComparisonRecord(
name="hidden_states",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
diff=_make_diff(),
)
parsed = json.loads(record.model_dump_json())
assert parsed["type"] == "comparison"
assert parsed["name"] == "hidden_states"
assert "baseline" in parsed
assert "diff" in parsed
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,
),
SkipRecord(name="attn", reason="no_baseline"),
ComparisonRecord(
name="mlp",
baseline=_make_tensor_info(),
target=_make_tensor_info(),
unified_shape=[4, 8],
shape_mismatch=False,
),
SummaryRecord(total=10, passed=8, failed=1, skipped=1),
]:
restored = parse_record_json(record.model_dump_json())
assert type(restored) is type(record)
assert restored == record
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
@@ -6,6 +6,15 @@ import pytest
import torch
from sglang.srt.debug_utils.comparator.entrypoint import run
from sglang.srt.debug_utils.comparator.output_types import (
AnyRecord,
ComparisonRecord,
ConfigRecord,
SkipRecord,
SummaryRecord,
_OutputRecord,
parse_record_json,
)
from sglang.srt.debug_utils.dumper import DumperConfig, _Dumper
from sglang.test.ci.ci_register import register_cpu_ci
@@ -65,36 +74,39 @@ def _make_args(baseline_path: Path, target_path: Path, **overrides) -> Namespace
end_step=1000000,
diff_threshold=1e-3,
filter=None,
output_format="text",
)
defaults.update(overrides)
return Namespace(**defaults)
def _parse_jsonl(output: str) -> list[AnyRecord]:
return [parse_record_json(line) for line in output.strip().splitlines()]
class TestEntrypoint:
def test_run_basic(self, tmp_path, capsys):
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
args = _make_args(baseline_path, target_path)
capsys.readouterr()
run(args)
output = capsys.readouterr().out
assert "df_target" in output
assert "df_baseline" in output
assert output.count("Check:") == 2
assert "tensor_a" in output
assert "tensor_b" in output
assert "Config:" in output
assert "rel_diff" in output
assert "Summary:" in output
assert "Skip" not in output
def test_filter(self, tmp_path, capsys):
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
args = _make_args(baseline_path, target_path, filter="tensor_a")
capsys.readouterr()
run(args)
output = capsys.readouterr().out
assert output.count("Check:") == 1
assert "tensor_a" in output
assert "rel_diff" in output
def test_no_baseline_skip(self, tmp_path, capsys):
baseline_path, target_path = _create_dumps(
@@ -103,22 +115,73 @@ class TestEntrypoint:
baseline_names=["tensor_a"],
)
args = _make_args(baseline_path, target_path)
capsys.readouterr()
run(args)
output = capsys.readouterr().out
assert output.count("Check:") == 1
assert "Skip:" in output
assert "since no baseline" in output
assert "no_baseline" in output
def test_step_range(self, tmp_path, capsys):
baseline_path, target_path = _create_dumps(tmp_path, ["t"], num_steps=3)
args = _make_args(baseline_path, target_path, start_step=1, end_step=1)
capsys.readouterr()
run(args)
output = capsys.readouterr().out
assert output.count("Check:") == 1
assert "Summary:" in output
class TestEntrypointJsonl:
def test_jsonl_basic(self, tmp_path, capsys):
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a", "tensor_b"])
args = _make_args(baseline_path, target_path, output_format="json")
capsys.readouterr()
run(args)
records = _parse_jsonl(capsys.readouterr().out)
assert isinstance(records[0], ConfigRecord)
comparisons = [r for r in records if isinstance(r, ComparisonRecord)]
assert len(comparisons) == 2
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.total == 2
assert summary.skipped == 0
def test_jsonl_skip(self, tmp_path, capsys):
baseline_path, target_path = _create_dumps(
tmp_path,
tensor_names=["tensor_a", "tensor_extra"],
baseline_names=["tensor_a"],
)
args = _make_args(baseline_path, target_path, output_format="json")
capsys.readouterr()
run(args)
records = _parse_jsonl(capsys.readouterr().out)
skips = [r for r in records if isinstance(r, SkipRecord)]
assert len(skips) == 1
assert skips[0].reason == "no_baseline"
summary = records[-1]
assert isinstance(summary, SummaryRecord)
assert summary.skipped == 1
def test_jsonl_all_valid_records(self, tmp_path, capsys):
baseline_path, target_path = _create_dumps(tmp_path, ["t"], num_steps=2)
args = _make_args(baseline_path, target_path, output_format="json")
capsys.readouterr()
run(args)
records = _parse_jsonl(capsys.readouterr().out)
assert all(isinstance(r, _OutputRecord) for r in records)
if __name__ == "__main__":