298 lines
10 KiB
Python
298 lines
10 KiB
Python
"""Tests for meta_overrider — unit tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import textwrap
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from sglang.srt.debug_utils.comparator.meta_overrider import (
|
|
MetaOverrider,
|
|
MetaOverrideRule,
|
|
_load_yaml_rules,
|
|
_parse_cli_override_arg,
|
|
)
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=10, suite="stage-a-test-cpu", nightly=True)
|
|
|
|
|
|
# ───────────────────── Unit: MetaOverrideRule ─────────────────────
|
|
|
|
|
|
class TestMetaOverrideRule:
|
|
"""Pydantic validation for MetaOverrideRule."""
|
|
|
|
def test_shared_dims_both(self) -> None:
|
|
"""Default side='both' applies dims to both sides."""
|
|
rule = MetaOverrideRule(match="hidden", dims="b s h d")
|
|
assert rule.dims == "b s h d"
|
|
assert rule.side == "both"
|
|
|
|
def test_side_baseline(self) -> None:
|
|
"""side='baseline' is accepted."""
|
|
rule = MetaOverrideRule(match="logits", dims="b s v[tp]", side="baseline")
|
|
assert rule.dims == "b s v[tp]"
|
|
assert rule.side == "baseline"
|
|
|
|
def test_side_target(self) -> None:
|
|
"""side='target' is accepted."""
|
|
rule = MetaOverrideRule(match="logits", dims="b s v[ep]", side="target")
|
|
assert rule.dims == "b s v[ep]"
|
|
assert rule.side == "target"
|
|
|
|
def test_invalid_side_rejected(self) -> None:
|
|
"""Invalid side value is rejected."""
|
|
with pytest.raises(Exception):
|
|
MetaOverrideRule(match="x", dims="b s", side="invalid")
|
|
|
|
def test_dims_required(self) -> None:
|
|
"""Must specify dims."""
|
|
with pytest.raises(Exception):
|
|
MetaOverrideRule(match="x")
|
|
|
|
def test_extra_field_rejected(self) -> None:
|
|
"""Extra fields are rejected by _StrictBase."""
|
|
with pytest.raises(Exception):
|
|
MetaOverrideRule(match="x", dims="b s", bogus="y")
|
|
|
|
|
|
# ──────────────────── Unit: _parse_cli_override_arg ────────────────────
|
|
|
|
|
|
class TestParseCLIOverrideArg:
|
|
"""CLI arg parsing for 'name:dims_string' format."""
|
|
|
|
def test_basic(self) -> None:
|
|
"""Standard 'name:dims' parsing."""
|
|
name, dims_str = _parse_cli_override_arg("hidden_states:b s h d")
|
|
assert name == "hidden_states"
|
|
assert dims_str == "b s h d"
|
|
|
|
def test_colon_in_dims(self) -> None:
|
|
"""Extra colons in dims are kept (maxsplit=1)."""
|
|
name, dims_str = _parse_cli_override_arg("x:a:b")
|
|
assert name == "x"
|
|
assert dims_str == "a:b"
|
|
|
|
def test_whitespace_trimmed(self) -> None:
|
|
"""Leading/trailing whitespace around name and dims is stripped."""
|
|
name, dims_str = _parse_cli_override_arg(" foo : b s ")
|
|
assert name == "foo"
|
|
assert dims_str == "b s"
|
|
|
|
def test_missing_colon(self) -> None:
|
|
"""No colon raises ValueError."""
|
|
with pytest.raises(ValueError, match="Invalid override format"):
|
|
_parse_cli_override_arg("no_colon_here")
|
|
|
|
def test_empty_name(self) -> None:
|
|
"""Empty name raises ValueError."""
|
|
with pytest.raises(ValueError, match="Invalid override format"):
|
|
_parse_cli_override_arg(":b s h")
|
|
|
|
def test_empty_dims(self) -> None:
|
|
"""Empty dims raises ValueError."""
|
|
with pytest.raises(ValueError, match="Invalid override format"):
|
|
_parse_cli_override_arg("foo:")
|
|
|
|
|
|
# ──────────────────── Unit: MetaOverrider ────────────────────
|
|
|
|
|
|
class TestMetaOverrider:
|
|
"""MetaOverrider logic: matching, priority, apply_to_meta."""
|
|
|
|
def test_first_match_wins(self) -> None:
|
|
"""First matching rule takes effect; later rules ignored."""
|
|
overrider = MetaOverrider(
|
|
rules=[
|
|
MetaOverrideRule(match="hidden", dims="FIRST"),
|
|
MetaOverrideRule(match="hidden", dims="SECOND"),
|
|
]
|
|
)
|
|
result: dict = overrider.apply_to_meta(
|
|
name="hidden_states",
|
|
meta={"dims": "old"},
|
|
side="baseline",
|
|
)
|
|
assert result["dims"] == "FIRST"
|
|
|
|
def test_regex_contains_match(self) -> None:
|
|
"""match is a regex contains search, not exact match."""
|
|
overrider = MetaOverrider(
|
|
rules=[MetaOverrideRule(match=r"\.q_proj\.", dims="h d")]
|
|
)
|
|
result: dict = overrider.apply_to_meta(
|
|
name="layers.0.q_proj.weight",
|
|
meta={"dims": "old"},
|
|
side="baseline",
|
|
)
|
|
assert result["dims"] == "h d"
|
|
|
|
def test_no_match_preserves_original(self) -> None:
|
|
"""No matching rule leaves meta untouched."""
|
|
overrider = MetaOverrider(
|
|
rules=[MetaOverrideRule(match="logits", dims="b s v")]
|
|
)
|
|
result: dict = overrider.apply_to_meta(
|
|
name="hidden_states",
|
|
meta={"dims": "original"},
|
|
side="baseline",
|
|
)
|
|
assert result["dims"] == "original"
|
|
|
|
@pytest.mark.parametrize(
|
|
"rule_side,apply_side,should_match",
|
|
[
|
|
("baseline", "baseline", True),
|
|
("baseline", "target", False),
|
|
("target", "target", True),
|
|
("target", "baseline", False),
|
|
("both", "baseline", True),
|
|
("both", "target", True),
|
|
],
|
|
)
|
|
def test_side_filtering(
|
|
self, rule_side: str, apply_side: str, should_match: bool
|
|
) -> None:
|
|
"""Rule only applies when its side matches the apply side."""
|
|
overrider = MetaOverrider(
|
|
rules=[MetaOverrideRule(match="logits", dims="NEW", side=rule_side)]
|
|
)
|
|
result: dict = overrider.apply_to_meta(
|
|
name="logits",
|
|
meta={"dims": "old"},
|
|
side=apply_side,
|
|
)
|
|
assert result["dims"] == ("NEW" if should_match else "old")
|
|
|
|
def test_is_empty(self) -> None:
|
|
"""Empty overrider reports is_empty=True."""
|
|
assert MetaOverrider(rules=[]).is_empty
|
|
assert not MetaOverrider(rules=[MetaOverrideRule(match="x", dims="d")]).is_empty
|
|
|
|
def test_meta_without_dims_key(self) -> None:
|
|
"""Override adds 'dims' even if original meta lacks it."""
|
|
overrider = MetaOverrider(rules=[MetaOverrideRule(match="hidden", dims="NEW")])
|
|
result: dict = overrider.apply_to_meta(
|
|
name="hidden",
|
|
meta={"other": "val"},
|
|
side="baseline",
|
|
)
|
|
assert result["dims"] == "NEW"
|
|
|
|
|
|
# ──────────────────── Unit: from_args_and_config ────────────────────
|
|
|
|
|
|
class TestFromArgsAndConfig:
|
|
"""MetaOverrider.from_args_and_config merges CLI + YAML rules."""
|
|
|
|
def test_cli_before_yaml(self, tmp_path: Path) -> None:
|
|
"""CLI rules are ordered before YAML rules (CLI wins on conflict)."""
|
|
yaml_path = tmp_path / "override.yaml"
|
|
yaml_path.write_text(textwrap.dedent("""\
|
|
overrides:
|
|
- match: "hidden"
|
|
dims: "FROM_YAML"
|
|
"""))
|
|
|
|
overrider = MetaOverrider.from_args_and_config(
|
|
override_dims=["hidden:FROM_CLI"],
|
|
override_baseline_dims=[],
|
|
override_target_dims=[],
|
|
override_config=yaml_path,
|
|
)
|
|
|
|
result: dict = overrider.apply_to_meta(
|
|
name="hidden",
|
|
meta={"dims": "old"},
|
|
side="baseline",
|
|
)
|
|
assert result["dims"] == "FROM_CLI"
|
|
|
|
def test_no_config_no_cli(self) -> None:
|
|
"""Empty CLI + no YAML yields empty overrider."""
|
|
overrider = MetaOverrider.from_args_and_config(
|
|
override_dims=[],
|
|
override_baseline_dims=[],
|
|
override_target_dims=[],
|
|
override_config=None,
|
|
)
|
|
assert overrider.is_empty
|
|
|
|
def test_per_side_cli_produces_separate_rules(self) -> None:
|
|
"""--override-baseline-dims and --override-target-dims produce separate rules with side field."""
|
|
overrider = MetaOverrider.from_args_and_config(
|
|
override_dims=[],
|
|
override_baseline_dims=["hidden:b s h[tp]"],
|
|
override_target_dims=["hidden:b s h[ep]"],
|
|
override_config=None,
|
|
)
|
|
|
|
baseline: dict = overrider.apply_to_meta(
|
|
name="hidden",
|
|
meta={"dims": "old"},
|
|
side="baseline",
|
|
)
|
|
target: dict = overrider.apply_to_meta(
|
|
name="hidden",
|
|
meta={"dims": "old"},
|
|
side="target",
|
|
)
|
|
assert baseline["dims"] == "b s h[tp]"
|
|
assert target["dims"] == "b s h[ep]"
|
|
|
|
|
|
# ──────────────────── Unit: _load_yaml_rules ────────────────────
|
|
|
|
|
|
class TestLoadYamlRules:
|
|
"""YAML loading and validation."""
|
|
|
|
def test_valid_yaml(self, tmp_path: Path) -> None:
|
|
"""Valid YAML with override rules loads correctly."""
|
|
yaml_path = tmp_path / "override.yaml"
|
|
yaml_path.write_text(textwrap.dedent("""\
|
|
overrides:
|
|
- match: "hidden"
|
|
dims: "b s h d"
|
|
- match: "logits"
|
|
dims: "b s v[tp]"
|
|
side: baseline
|
|
"""))
|
|
rules = _load_yaml_rules(yaml_path)
|
|
assert len(rules) == 2
|
|
assert rules[0].dims == "b s h d"
|
|
assert rules[0].side == "both"
|
|
assert rules[1].dims == "b s v[tp]"
|
|
assert rules[1].side == "baseline"
|
|
|
|
def test_empty_yaml(self, tmp_path: Path) -> None:
|
|
"""Empty YAML file returns no rules."""
|
|
yaml_path = tmp_path / "empty.yaml"
|
|
yaml_path.write_text("")
|
|
rules = _load_yaml_rules(yaml_path)
|
|
assert rules == []
|
|
|
|
def test_unknown_top_key_rejected(self, tmp_path: Path) -> None:
|
|
"""Unknown top-level key is rejected by OverrideConfig."""
|
|
yaml_path = tmp_path / "bad.yaml"
|
|
yaml_path.write_text("unknown_key: 42\n")
|
|
with pytest.raises(Exception):
|
|
_load_yaml_rules(yaml_path)
|
|
|
|
def test_overrides_empty_list(self, tmp_path: Path) -> None:
|
|
"""Only 'overrides' key with no entries returns empty list."""
|
|
yaml_path = tmp_path / "minimal.yaml"
|
|
yaml_path.write_text("overrides: []\n")
|
|
rules = _load_yaml_rules(yaml_path)
|
|
assert rules == []
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(pytest.main([__file__]))
|