485 lines
16 KiB
Python
485 lines
16 KiB
Python
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from sglang.srt.debug_utils.comparator.output_types import SummaryRecord
|
|
from sglang.srt.debug_utils.comparator.utils import (
|
|
Pair,
|
|
argmax_coord,
|
|
auto_descend_dir,
|
|
calc_per_token_rel_diff,
|
|
calc_rel_diff,
|
|
compute_exit_code,
|
|
compute_smaller_dtype,
|
|
try_unify_shape,
|
|
)
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=10, suite="stage-a-test-cpu", nightly=True)
|
|
|
|
|
|
class TestCalcRelDiff:
|
|
def test_identical_tensors(self):
|
|
x = torch.randn(10, 10)
|
|
assert calc_rel_diff(x, x).item() == pytest.approx(0.0, abs=1e-5)
|
|
|
|
def test_orthogonal_tensors(self):
|
|
result = calc_rel_diff(
|
|
torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0])
|
|
).item()
|
|
assert result == pytest.approx(1.0, abs=1e-5)
|
|
|
|
def test_similar_tensors(self):
|
|
x = torch.tensor([1.0, 2.0, 3.0])
|
|
y = torch.tensor([1.01, 2.01, 3.01])
|
|
result = calc_rel_diff(x, y).item()
|
|
assert 0.0 < result < 0.01
|
|
|
|
def test_negated_tensors(self):
|
|
x = torch.tensor([1.0, 2.0])
|
|
result = calc_rel_diff(x, -x).item()
|
|
assert result == pytest.approx(2.0, abs=1e-5)
|
|
|
|
|
|
class TestCalcPerTokenRelDiff:
|
|
def test_identical_tensors(self) -> None:
|
|
"""Identical tensors → per-token diff all zero."""
|
|
x: torch.Tensor = torch.randn(8, 16)
|
|
result: torch.Tensor = calc_per_token_rel_diff(x, x, seq_dim=0)
|
|
|
|
assert result.shape == (8,)
|
|
assert torch.allclose(result, torch.zeros(8), atol=1e-6)
|
|
|
|
def test_different_tensors(self) -> None:
|
|
"""Single token position differs → that position has higher diff."""
|
|
torch.manual_seed(42)
|
|
x: torch.Tensor = torch.randn(8, 16)
|
|
y: torch.Tensor = x.clone()
|
|
y[3, :] += 10.0
|
|
|
|
result: torch.Tensor = calc_per_token_rel_diff(x, y, seq_dim=0)
|
|
|
|
assert result.shape == (8,)
|
|
assert result[3] > result[0]
|
|
assert result[3] > result[7]
|
|
for i in [0, 1, 2, 4, 5, 6, 7]:
|
|
assert result[i] < 1e-6
|
|
|
|
def test_seq_dim_selection(self) -> None:
|
|
"""Different seq_dim values produce correct output shapes."""
|
|
x: torch.Tensor = torch.randn(4, 8, 16)
|
|
y: torch.Tensor = x + torch.randn_like(x) * 0.01
|
|
|
|
assert calc_per_token_rel_diff(x, y, seq_dim=0).shape == (4,)
|
|
assert calc_per_token_rel_diff(x, y, seq_dim=1).shape == (8,)
|
|
assert calc_per_token_rel_diff(x, y, seq_dim=2).shape == (16,)
|
|
|
|
def test_1d_tensor(self) -> None:
|
|
"""1D tensor with seq_dim=0 returns per-element diff."""
|
|
x: torch.Tensor = torch.tensor([1.0, 2.0, 3.0])
|
|
y: torch.Tensor = torch.tensor([1.0, 2.0, 4.0])
|
|
|
|
result: torch.Tensor = calc_per_token_rel_diff(x, y, seq_dim=0)
|
|
|
|
assert result.shape == (3,)
|
|
assert result[0] < 1e-6
|
|
assert result[1] < 1e-6
|
|
assert result[2] > 0.01
|
|
|
|
|
|
class TestArgmaxCoord:
|
|
def test_1d_tensor(self):
|
|
x = torch.tensor([0.0, 0.0, 5.0, 0.0])
|
|
assert argmax_coord(x) == (2,)
|
|
|
|
def test_2d_tensor(self):
|
|
x = torch.zeros(3, 4)
|
|
x[1, 2] = 10.0
|
|
assert argmax_coord(x) == (1, 2)
|
|
|
|
def test_3d_tensor(self):
|
|
x = torch.zeros(2, 3, 4)
|
|
x[1, 2, 3] = 10.0
|
|
assert argmax_coord(x) == (1, 2, 3)
|
|
|
|
|
|
class TestTryUnifyShape:
|
|
def test_squeeze_leading_ones(self):
|
|
target = torch.Size([3, 4])
|
|
assert try_unify_shape(torch.randn(1, 1, 3, 4), target).shape == target
|
|
|
|
def test_no_squeeze_when_leading_dim_not_one(self):
|
|
target = torch.Size([3, 4])
|
|
assert try_unify_shape(torch.randn(2, 3, 4), target).shape == (2, 3, 4)
|
|
|
|
def test_same_shape_noop(self):
|
|
target = torch.Size([3, 4])
|
|
x = torch.randn(3, 4)
|
|
result = try_unify_shape(x, target)
|
|
assert result.shape == target
|
|
assert result.data_ptr() == x.data_ptr()
|
|
|
|
def test_trailing_dims_mismatch(self):
|
|
target = torch.Size([5, 6])
|
|
x = torch.randn(1, 3, 4)
|
|
result = try_unify_shape(x, target)
|
|
assert result.shape == (1, 3, 4)
|
|
|
|
|
|
class TestComputeSmallerDtype:
|
|
def test_float32_bfloat16(self):
|
|
assert (
|
|
compute_smaller_dtype(Pair(x=torch.float32, y=torch.bfloat16))
|
|
== torch.bfloat16
|
|
)
|
|
|
|
def test_reverse_order(self):
|
|
assert (
|
|
compute_smaller_dtype(Pair(x=torch.bfloat16, y=torch.float32))
|
|
== torch.bfloat16
|
|
)
|
|
|
|
def test_same_dtype_returns_none(self):
|
|
assert compute_smaller_dtype(Pair(x=torch.float32, y=torch.float32)) is None
|
|
|
|
def test_unknown_pair_returns_none(self):
|
|
assert compute_smaller_dtype(Pair(x=torch.int32, y=torch.int64)) is None
|
|
|
|
|
|
class TestPairMap:
|
|
def test_map_basic(self):
|
|
pair = Pair(x=[1, 2, 3], y=[4, 5, 6])
|
|
result = pair.map(lambda lst: sum(lst))
|
|
assert result.x == 6
|
|
assert result.y == 15
|
|
|
|
def test_map_type_change(self):
|
|
pair = Pair(x=[1, 2, 3], y=[10, 20])
|
|
result = pair.map(len)
|
|
assert result.x == 3
|
|
assert result.y == 2
|
|
|
|
def test_map_returns_new_pair(self):
|
|
pair = Pair(x="hello", y="world")
|
|
result = pair.map(str.upper)
|
|
assert result.x == "HELLO"
|
|
assert result.y == "WORLD"
|
|
assert result is not pair
|
|
|
|
|
|
class TestComputeExitCode:
|
|
"""Unit tests for compute_exit_code logic."""
|
|
|
|
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_skipped_pattern=".*",
|
|
skipped_names=[],
|
|
allow_failed_pattern=None,
|
|
failed_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_skipped_pattern=".*",
|
|
skipped_names=[],
|
|
allow_failed_pattern=None,
|
|
failed_names=["a", "b"],
|
|
)
|
|
== 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_skipped_pattern=".*",
|
|
skipped_names=[],
|
|
allow_failed_pattern=None,
|
|
failed_names=["a", "b", "c"],
|
|
)
|
|
== 1
|
|
)
|
|
|
|
def test_all_skipped_allow_all(self):
|
|
"""All skipped + allow_skipped_pattern='.*' → exit 1 (nothing passed)."""
|
|
summary = SummaryRecord(total=2, passed=0, failed=0, skipped=2)
|
|
assert (
|
|
compute_exit_code(
|
|
summary,
|
|
allow_skipped_pattern=".*",
|
|
skipped_names=["a", "b"],
|
|
allow_failed_pattern=None,
|
|
failed_names=[],
|
|
)
|
|
== 1
|
|
)
|
|
|
|
def test_all_skipped_forbid_all(self):
|
|
"""All skipped + allow_skipped_pattern='^$' → exit 1."""
|
|
summary = SummaryRecord(total=2, passed=0, failed=0, skipped=2)
|
|
assert (
|
|
compute_exit_code(
|
|
summary,
|
|
allow_skipped_pattern="^$",
|
|
skipped_names=["a", "b"],
|
|
allow_failed_pattern=None,
|
|
failed_names=[],
|
|
)
|
|
== 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_skipped_pattern=".*",
|
|
skipped_names=["a"],
|
|
allow_failed_pattern=None,
|
|
failed_names=[],
|
|
)
|
|
== 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_skipped_pattern="^$",
|
|
skipped_names=["a"],
|
|
allow_failed_pattern=None,
|
|
failed_names=[],
|
|
)
|
|
== 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_skipped_pattern="positions|seq_lens",
|
|
skipped_names=["positions", "seq_lens"],
|
|
allow_failed_pattern=None,
|
|
failed_names=[],
|
|
)
|
|
== 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_skipped_pattern="positions|seq_lens",
|
|
skipped_names=["positions", "seq_lens", "hidden_states"],
|
|
allow_failed_pattern=None,
|
|
failed_names=[],
|
|
)
|
|
== 1
|
|
)
|
|
|
|
def test_allow_failed_pattern_matches_all(self):
|
|
"""allow_failed_pattern='.*' tolerates all failures → exit 0."""
|
|
summary = SummaryRecord(total=3, passed=1, failed=2, skipped=0)
|
|
assert (
|
|
compute_exit_code(
|
|
summary,
|
|
allow_skipped_pattern=".*",
|
|
skipped_names=[],
|
|
allow_failed_pattern=".*",
|
|
failed_names=["a", "b"],
|
|
)
|
|
== 0
|
|
)
|
|
|
|
def test_allow_failed_pattern_matches_specific(self):
|
|
"""Pattern matches all failed names → exit 0."""
|
|
summary = SummaryRecord(total=3, passed=1, failed=2, skipped=0)
|
|
assert (
|
|
compute_exit_code(
|
|
summary,
|
|
allow_skipped_pattern=".*",
|
|
skipped_names=[],
|
|
allow_failed_pattern="hidden_states|logits",
|
|
failed_names=["hidden_states", "logits"],
|
|
)
|
|
== 0
|
|
)
|
|
|
|
def test_allow_failed_pattern_partial_match(self):
|
|
"""Pattern matches some but not all failures → exit 1."""
|
|
summary = SummaryRecord(total=3, passed=0, failed=3, skipped=0)
|
|
assert (
|
|
compute_exit_code(
|
|
summary,
|
|
allow_skipped_pattern=".*",
|
|
skipped_names=[],
|
|
allow_failed_pattern="hidden_states",
|
|
failed_names=["hidden_states", "logits", "attn"],
|
|
)
|
|
== 1
|
|
)
|
|
|
|
def test_allow_failed_pattern_no_failures(self):
|
|
"""Pattern set but no failures → exit 0."""
|
|
summary = SummaryRecord(total=2, passed=2, failed=0, skipped=0)
|
|
assert (
|
|
compute_exit_code(
|
|
summary,
|
|
allow_skipped_pattern=".*",
|
|
skipped_names=[],
|
|
allow_failed_pattern=".*",
|
|
failed_names=[],
|
|
)
|
|
== 0
|
|
)
|
|
|
|
def test_both_failed_and_skipped_patterns(self):
|
|
"""Both patterns set, both satisfied → exit 0."""
|
|
summary = SummaryRecord(total=4, passed=1, failed=1, skipped=2)
|
|
assert (
|
|
compute_exit_code(
|
|
summary,
|
|
allow_skipped_pattern="positions|seq_lens",
|
|
skipped_names=["positions", "seq_lens"],
|
|
allow_failed_pattern="logits",
|
|
failed_names=["logits"],
|
|
)
|
|
== 0
|
|
)
|
|
|
|
def test_failed_pattern_satisfied_but_skipped_not(self):
|
|
"""Failed pattern OK but skipped pattern fails → exit 1."""
|
|
summary = SummaryRecord(total=3, passed=1, failed=1, skipped=1)
|
|
assert (
|
|
compute_exit_code(
|
|
summary,
|
|
allow_skipped_pattern="^$",
|
|
skipped_names=["a"],
|
|
allow_failed_pattern=".*",
|
|
failed_names=["b"],
|
|
)
|
|
== 1
|
|
)
|
|
|
|
def test_zero_passed_exits_one(self):
|
|
"""No tensors passed → exit 1, even when all failures are allowed."""
|
|
summary = SummaryRecord(total=2, passed=0, failed=2, skipped=0)
|
|
assert (
|
|
compute_exit_code(
|
|
summary,
|
|
allow_skipped_pattern=".*",
|
|
skipped_names=[],
|
|
allow_failed_pattern=".*",
|
|
failed_names=["a", "b"],
|
|
)
|
|
== 1
|
|
)
|
|
|
|
def test_zero_passed_all_skipped_exits_one(self):
|
|
"""All skipped, nothing passed → exit 1."""
|
|
summary = SummaryRecord(total=3, passed=0, failed=0, skipped=3)
|
|
assert (
|
|
compute_exit_code(
|
|
summary,
|
|
allow_skipped_pattern=".*",
|
|
skipped_names=["a", "b", "c"],
|
|
allow_failed_pattern=None,
|
|
failed_names=[],
|
|
)
|
|
== 1
|
|
)
|
|
|
|
def test_errored_with_passed_exits_one(self):
|
|
"""Has errored bundle even with passed → exit 1."""
|
|
summary = SummaryRecord(total=3, passed=2, failed=0, skipped=0, errored=1)
|
|
assert (
|
|
compute_exit_code(
|
|
summary,
|
|
allow_skipped_pattern=".*",
|
|
skipped_names=[],
|
|
allow_failed_pattern=None,
|
|
failed_names=[],
|
|
errored_names=["broken_tensor"],
|
|
)
|
|
== 1
|
|
)
|
|
|
|
def test_errored_only_exits_one(self):
|
|
"""All errored → exit 1 (passed==0 already exits 1, but errored also independently triggers)."""
|
|
summary = SummaryRecord(total=1, passed=0, failed=0, skipped=0, errored=1)
|
|
assert (
|
|
compute_exit_code(
|
|
summary,
|
|
allow_skipped_pattern=".*",
|
|
skipped_names=[],
|
|
allow_failed_pattern=None,
|
|
failed_names=[],
|
|
errored_names=["broken_tensor"],
|
|
)
|
|
== 1
|
|
)
|
|
|
|
|
|
def _make_pt(directory: Path) -> None:
|
|
directory.mkdir(parents=True, exist_ok=True)
|
|
torch.save(torch.tensor([1.0]), directory / "dummy.pt")
|
|
|
|
|
|
class TestAutoDescendDir:
|
|
def test_no_descend_when_pt_at_root(self, tmp_path: Path) -> None:
|
|
"""Directory with .pt files directly is returned as-is."""
|
|
_make_pt(tmp_path)
|
|
_make_pt(tmp_path / "child_a")
|
|
assert auto_descend_dir(tmp_path, label="test") == tmp_path
|
|
|
|
def test_descend_into_single_child(self, tmp_path: Path) -> None:
|
|
"""Single child with .pt triggers descend."""
|
|
child: Path = tmp_path / "engine_0"
|
|
_make_pt(child)
|
|
assert auto_descend_dir(tmp_path, label="test") == child
|
|
|
|
def test_descend_single_nonempty_child_among_empty(self, tmp_path: Path) -> None:
|
|
"""Two subdirs but only one has .pt — descend into that one."""
|
|
nonempty: Path = tmp_path / "engine_0"
|
|
_make_pt(nonempty)
|
|
(tmp_path / "empty_child").mkdir()
|
|
assert auto_descend_dir(tmp_path, label="test") == nonempty
|
|
|
|
def test_error_with_multiple_nonempty_children(self, tmp_path: Path) -> None:
|
|
"""Two children with .pt files — ambiguous, raises ValueError."""
|
|
_make_pt(tmp_path / "engine_0")
|
|
_make_pt(tmp_path / "engine_1")
|
|
with pytest.raises(ValueError, match="multiple subdirectories contain data"):
|
|
auto_descend_dir(tmp_path, label="test")
|
|
|
|
def test_error_when_no_data_found(self, tmp_path: Path) -> None:
|
|
"""No .pt files anywhere — raises ValueError."""
|
|
(tmp_path / "empty_child").mkdir()
|
|
with pytest.raises(ValueError, match="no .pt files found"):
|
|
auto_descend_dir(tmp_path, label="test")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main([__file__]))
|