432 lines
14 KiB
Python
432 lines
14 KiB
Python
import sys
|
|
from argparse import Namespace
|
|
from pathlib import Path
|
|
|
|
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.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
|
|
|
|
register_cpu_ci(est_time=30, suite="default", nightly=True)
|
|
|
|
_FIXED_EXP_NAME = "my_exp_name"
|
|
|
|
# Each test has a one-line docstring describing the scenario it covers.
|
|
|
|
|
|
class TestEntrypointGroupingRaw:
|
|
"""Test `--grouping raw` scenarios"""
|
|
|
|
def test_run_basic(self, tmp_path, capsys):
|
|
"""Two matching tensors produce ConfigRecord, 2 ComparisonRecords, and SummaryRecord."""
|
|
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)
|
|
assert isinstance(records[0], ConfigRecord)
|
|
|
|
assert len(_get_comparisons(records)) == 2
|
|
|
|
summary = records[-1]
|
|
assert isinstance(summary, SummaryRecord)
|
|
assert summary.total == 2
|
|
assert summary.skipped == 0
|
|
|
|
def test_filter(self, tmp_path, capsys):
|
|
"""--filter selects only the matching tensor, producing 1 ComparisonRecord."""
|
|
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)
|
|
assert len(_get_comparisons(records)) == 1
|
|
|
|
def test_no_baseline_skip(self, tmp_path, capsys):
|
|
"""Target tensor missing from baseline emits a SkipRecord with reason baseline_load_failed."""
|
|
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, grouping="raw")
|
|
|
|
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"
|
|
|
|
summary = records[-1]
|
|
assert isinstance(summary, SummaryRecord)
|
|
assert summary.skipped == 1
|
|
|
|
def test_step_range(self, tmp_path, capsys):
|
|
"""--start_step/--end_step restricts comparison to a single step out of three."""
|
|
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, grouping="raw"
|
|
)
|
|
|
|
records = _run_and_parse(args, capsys)
|
|
summary = records[-1]
|
|
assert isinstance(summary, SummaryRecord)
|
|
assert summary.total == 1
|
|
|
|
def test_all_valid_records(self, tmp_path, capsys):
|
|
"""Every emitted JSON record is a valid _OutputRecord subclass."""
|
|
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)
|
|
assert all(isinstance(r, _OutputRecord) for r in records)
|
|
|
|
def test_text_output_smoke(self, tmp_path, capsys):
|
|
"""Text output format renders without errors and contains Config/Summary sections."""
|
|
baseline_path, target_path = _create_dumps(tmp_path, ["tensor_a"])
|
|
args = _make_args(
|
|
baseline_path, target_path, output_format="text", grouping="raw"
|
|
)
|
|
capsys.readouterr()
|
|
|
|
run(args)
|
|
|
|
output = capsys.readouterr().out
|
|
assert "Config:" in output
|
|
assert "Summary:" in output
|
|
|
|
|
|
class TestEntrypointGroupingLogical:
|
|
"""Test `--grouping logical` scenarios"""
|
|
|
|
def test_no_dims_single_rank(self, tmp_path, capsys):
|
|
"""Single-rank dumps without dims fall back to raw loading."""
|
|
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)
|
|
assert len(_get_comparisons(records)) == 2
|
|
summary = records[-1]
|
|
assert isinstance(summary, SummaryRecord)
|
|
assert summary.total == 2
|
|
assert summary.skipped == 0
|
|
|
|
def test_tp_unshard_same_size(self, tmp_path, capsys):
|
|
"""Both sides TP=2: shards are concatenated before comparison."""
|
|
torch.manual_seed(42)
|
|
full_baseline = torch.randn(4, 8)
|
|
full_target = full_baseline + torch.randn(4, 8) * 0.001
|
|
|
|
baseline_dir = tmp_path / "baseline"
|
|
target_dir = tmp_path / "target"
|
|
|
|
baseline_path = _create_tp_sharded_dumps(
|
|
baseline_dir,
|
|
full_tensor=full_baseline,
|
|
name="hidden",
|
|
tp_size=2,
|
|
shard_dim=1,
|
|
dims_str="b h(tp)",
|
|
)
|
|
target_path = _create_tp_sharded_dumps(
|
|
target_dir,
|
|
full_tensor=full_target,
|
|
name="hidden",
|
|
tp_size=2,
|
|
shard_dim=1,
|
|
dims_str="b h(tp)",
|
|
)
|
|
|
|
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
|
|
|
|
records = _run_and_parse(args, capsys)
|
|
comp = _assert_single_comparison_passed(records)
|
|
assert comp.name == "hidden"
|
|
|
|
summary = records[-1]
|
|
assert isinstance(summary, SummaryRecord)
|
|
assert summary.total == 1
|
|
assert summary.passed == 1
|
|
|
|
def test_tp_unshard_different_sizes(self, tmp_path, capsys):
|
|
"""Baseline TP=4 vs target TP=2: different shard counts are handled correctly."""
|
|
torch.manual_seed(42)
|
|
full_baseline = torch.randn(4, 8)
|
|
full_target = full_baseline + torch.randn(4, 8) * 0.001
|
|
|
|
baseline_dir = tmp_path / "baseline"
|
|
target_dir = tmp_path / "target"
|
|
|
|
baseline_path = _create_tp_sharded_dumps(
|
|
baseline_dir,
|
|
full_tensor=full_baseline,
|
|
name="hidden",
|
|
tp_size=4,
|
|
shard_dim=1,
|
|
dims_str="b h(tp)",
|
|
)
|
|
target_path = _create_tp_sharded_dumps(
|
|
target_dir,
|
|
full_tensor=full_target,
|
|
name="hidden",
|
|
tp_size=2,
|
|
shard_dim=1,
|
|
dims_str="b h(tp)",
|
|
)
|
|
|
|
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
|
|
|
|
records = _run_and_parse(args, capsys)
|
|
_assert_single_comparison_passed(records)
|
|
|
|
def test_one_side_dims_single_baseline(self, tmp_path, capsys):
|
|
"""Baseline has no dims (single rank), target has TP shards: unshard target only."""
|
|
torch.manual_seed(42)
|
|
full_tensor = torch.randn(4, 8)
|
|
target_full = full_tensor + torch.randn(4, 8) * 0.001
|
|
|
|
baseline_dir = tmp_path / "baseline"
|
|
target_dir = tmp_path / "target"
|
|
|
|
baseline_path = _create_rank_dump(
|
|
baseline_dir, rank=0, name="hidden", tensor=full_tensor
|
|
)
|
|
|
|
target_path = _create_tp_sharded_dumps(
|
|
target_dir,
|
|
full_tensor=target_full,
|
|
name="hidden",
|
|
tp_size=2,
|
|
shard_dim=1,
|
|
dims_str="b h(tp)",
|
|
)
|
|
|
|
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
|
|
|
|
records = _run_and_parse(args, capsys)
|
|
_assert_single_comparison_passed(records)
|
|
|
|
def test_ambiguous_baseline_no_dims(self, tmp_path, capsys):
|
|
"""Multi-rank baseline without dims cannot be unsharded, so it is skipped."""
|
|
torch.manual_seed(42)
|
|
full_tensor = torch.randn(4, 8)
|
|
|
|
baseline_dir = tmp_path / "baseline"
|
|
target_dir = tmp_path / "target"
|
|
|
|
for rank, shard in [(0, full_tensor[:, :4]), (1, full_tensor[:, 4:])]:
|
|
baseline_path = _create_rank_dump(
|
|
baseline_dir, rank=rank, name="hidden", tensor=shard
|
|
)
|
|
|
|
target_path = _create_tp_sharded_dumps(
|
|
target_dir,
|
|
full_tensor=full_tensor,
|
|
name="hidden",
|
|
tp_size=2,
|
|
shard_dim=1,
|
|
dims_str="b h(tp)",
|
|
)
|
|
|
|
args = _make_args(baseline_path, target_path)
|
|
|
|
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"
|
|
|
|
def test_summary_counts_unshard(self, tmp_path, capsys):
|
|
"""Two TP-sharded tensors: summary counts total=2, passed=2, skipped=0."""
|
|
torch.manual_seed(42)
|
|
full_a = torch.randn(4, 8)
|
|
full_b = torch.randn(4, 8)
|
|
|
|
baseline_dir = tmp_path / "baseline"
|
|
target_dir = tmp_path / "target"
|
|
|
|
for tensor_name, tensor in [("t_a", full_a), ("t_b", full_b)]:
|
|
baseline_path = _create_tp_sharded_dumps(
|
|
baseline_dir,
|
|
full_tensor=tensor,
|
|
name=tensor_name,
|
|
tp_size=2,
|
|
shard_dim=1,
|
|
dims_str="b h(tp)",
|
|
)
|
|
target_tensor = tensor + torch.randn_like(tensor) * 0.0001
|
|
target_path = _create_tp_sharded_dumps(
|
|
target_dir,
|
|
full_tensor=target_tensor,
|
|
name=tensor_name,
|
|
tp_size=2,
|
|
shard_dim=1,
|
|
dims_str="b h(tp)",
|
|
)
|
|
|
|
args = _make_args(baseline_path, target_path, diff_threshold=0.01)
|
|
|
|
records = _run_and_parse(args, capsys)
|
|
summary = records[-1]
|
|
assert isinstance(summary, SummaryRecord)
|
|
assert summary.total == 2
|
|
assert summary.passed == 2
|
|
assert summary.failed == 0
|
|
assert summary.skipped == 0
|
|
|
|
|
|
# --------------------------- Assertion helpers -------------------
|
|
|
|
|
|
def _get_comparisons(records: list[AnyRecord]) -> list[ComparisonRecord]:
|
|
return [r for r in records if isinstance(r, ComparisonRecord)]
|
|
|
|
|
|
def _assert_single_comparison_passed(records: list[AnyRecord]) -> ComparisonRecord:
|
|
comparisons = _get_comparisons(records)
|
|
assert len(comparisons) == 1
|
|
assert comparisons[0].diff is not None
|
|
assert comparisons[0].diff.passed
|
|
return comparisons[0]
|
|
|
|
|
|
# --------------------------- Utils ------------------------------
|
|
|
|
|
|
def _make_dumper(directory: Path) -> _Dumper:
|
|
return _Dumper(
|
|
config=DumperConfig(enable=True, dir=str(directory), enable_http_server=False)
|
|
)
|
|
|
|
|
|
def _create_dumps(
|
|
tmp_path: Path,
|
|
tensor_names: list[str],
|
|
*,
|
|
baseline_names: list[str] | None = None,
|
|
num_steps: int = 1,
|
|
) -> tuple[Path, Path]:
|
|
"""Create baseline and target dump directories with given tensor names.
|
|
|
|
If baseline_names is None, uses the same names as tensor_names.
|
|
Each step dumps all names with the same tensor (different per baseline/target).
|
|
"""
|
|
if baseline_names is None:
|
|
baseline_names = tensor_names
|
|
|
|
d_baseline = tmp_path / "baseline"
|
|
d_target = tmp_path / "target"
|
|
d_baseline.mkdir()
|
|
d_target.mkdir()
|
|
|
|
torch.manual_seed(42)
|
|
baseline_tensor = torch.randn(10, 10)
|
|
target_tensor = baseline_tensor + torch.randn(10, 10) * 0.01
|
|
|
|
exp_paths: list[Path] = []
|
|
for d, names, tensor in [
|
|
(d_baseline, baseline_names, baseline_tensor),
|
|
(d_target, tensor_names, target_tensor),
|
|
]:
|
|
dumper = _make_dumper(d)
|
|
for _ in range(num_steps):
|
|
for name in names:
|
|
dumper.dump(name, tensor)
|
|
dumper.step()
|
|
exp_paths.append(d / dumper._config.exp_name)
|
|
|
|
return exp_paths[0], exp_paths[1]
|
|
|
|
|
|
def _make_args(baseline_path: Path, target_path: Path, **overrides) -> Namespace:
|
|
defaults = dict(
|
|
baseline_path=str(baseline_path),
|
|
target_path=str(target_path),
|
|
start_step=0,
|
|
end_step=1000000,
|
|
diff_threshold=1e-3,
|
|
filter=None,
|
|
output_format="json",
|
|
grouping="logical",
|
|
)
|
|
defaults.update(overrides)
|
|
return Namespace(**defaults)
|
|
|
|
|
|
def _run_and_parse(args: Namespace, capsys: pytest.CaptureFixture) -> list[AnyRecord]:
|
|
capsys.readouterr()
|
|
run(args)
|
|
return _parse_jsonl(capsys.readouterr().out)
|
|
|
|
|
|
def _parse_jsonl(output: str) -> list[AnyRecord]:
|
|
return [parse_record_json(line) for line in output.strip().splitlines()]
|
|
|
|
|
|
def _create_rank_dump(
|
|
directory: Path,
|
|
*,
|
|
rank: int,
|
|
name: str,
|
|
tensor: torch.Tensor,
|
|
dims: str | None = None,
|
|
parallel_info: dict | None = None,
|
|
) -> Path:
|
|
"""Create a dump file via the real dumper, as if running on the given rank."""
|
|
with pytest.MonkeyPatch.context() as mp:
|
|
mp.setattr(_dumper_module, "_get_rank", lambda: rank)
|
|
|
|
dumper = _Dumper(
|
|
config=DumperConfig(
|
|
enable=True,
|
|
dir=str(directory),
|
|
exp_name=_FIXED_EXP_NAME,
|
|
enable_http_server=False,
|
|
)
|
|
)
|
|
|
|
static_meta: dict = {"world_rank": rank, "world_size": 1}
|
|
if parallel_info is not None:
|
|
static_meta["sglang_parallel_info"] = parallel_info
|
|
dumper.__dict__["_static_meta"] = static_meta
|
|
|
|
dumper.dump(name, tensor, dims=dims)
|
|
dumper.step()
|
|
|
|
return directory / _FIXED_EXP_NAME
|
|
|
|
|
|
def _create_tp_sharded_dumps(
|
|
directory: Path,
|
|
*,
|
|
full_tensor: torch.Tensor,
|
|
name: str,
|
|
tp_size: int,
|
|
shard_dim: int,
|
|
dims_str: str,
|
|
) -> Path:
|
|
"""Create TP-sharded dump files from a full tensor via the real dumper."""
|
|
shards = list(full_tensor.chunk(tp_size, dim=shard_dim))
|
|
for tp_rank in range(tp_size):
|
|
_create_rank_dump(
|
|
directory,
|
|
rank=tp_rank,
|
|
name=name,
|
|
tensor=shards[tp_rank],
|
|
dims=dims_str,
|
|
parallel_info={"tp_rank": tp_rank, "tp_size": tp_size},
|
|
)
|
|
return directory / _FIXED_EXP_NAME
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main([__file__]))
|