130 lines
3.2 KiB
Python
130 lines
3.2 KiB
Python
"""Shared test helpers for comparator tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from io import StringIO
|
|
from typing import Optional
|
|
|
|
from rich.console import Console
|
|
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(
|
|
est_time=0,
|
|
suite="stage-a-test-cpu",
|
|
nightly=True,
|
|
disabled="helper module, no tests",
|
|
)
|
|
|
|
from sglang.srt.debug_utils.comparator.tensor_comparator.types import (
|
|
DiffInfo,
|
|
TensorInfo,
|
|
TensorStats,
|
|
)
|
|
|
|
DEFAULT_PERCENTILES: dict[int, float] = {
|
|
1: -1.8,
|
|
5: -1.5,
|
|
50: 0.0,
|
|
95: 1.5,
|
|
99: 1.8,
|
|
}
|
|
|
|
DEFAULT_ABS_DIFF_PERCENTILES: dict[int, float] = {
|
|
1: 0.0001,
|
|
5: 0.0001,
|
|
50: 0.0002,
|
|
95: 0.0004,
|
|
99: 0.0005,
|
|
}
|
|
|
|
|
|
def make_stats(
|
|
mean: float = 0.0,
|
|
abs_mean: float = 0.8,
|
|
std: float = 1.0,
|
|
min: float = -2.0,
|
|
max: float = 2.0,
|
|
percentiles: Optional[dict[int, float]] = None,
|
|
) -> TensorStats:
|
|
return TensorStats(
|
|
mean=mean,
|
|
abs_mean=abs_mean,
|
|
std=std,
|
|
min=min,
|
|
max=max,
|
|
percentiles=percentiles if percentiles is not None else DEFAULT_PERCENTILES,
|
|
)
|
|
|
|
|
|
def make_diff(
|
|
rel_diff: float = 0.0001,
|
|
max_abs_diff: float = 0.0005,
|
|
mean_abs_diff: float = 0.0002,
|
|
abs_diff_percentiles: Optional[dict[int, float]] = None,
|
|
diff_threshold: float = 1e-3,
|
|
passed: bool = True,
|
|
) -> DiffInfo:
|
|
return DiffInfo(
|
|
rel_diff=rel_diff,
|
|
max_abs_diff=max_abs_diff,
|
|
mean_abs_diff=mean_abs_diff,
|
|
abs_diff_percentiles=(
|
|
abs_diff_percentiles
|
|
if abs_diff_percentiles is not None
|
|
else DEFAULT_ABS_DIFF_PERCENTILES
|
|
),
|
|
max_diff_coord=[2, 3],
|
|
baseline_at_max=1.0,
|
|
target_at_max=1.0005,
|
|
diff_threshold=diff_threshold,
|
|
passed=passed,
|
|
)
|
|
|
|
|
|
_ANSI_ESCAPE_RE = re.compile(r"\033\[([0-9;]*)m")
|
|
|
|
|
|
def assert_rich_tags_balanced(markup: str) -> None:
|
|
"""Render Rich markup to ANSI and verify no styles are active at the end.
|
|
|
|
Tracks ANSI style state through the output. A ``\\033[0m`` (reset)
|
|
clears all active styles; any other ``\\033[Nm`` sets a style.
|
|
At the end of the output, no style should remain active.
|
|
"""
|
|
buf = StringIO()
|
|
console = Console(file=buf, force_terminal=True, width=10000, highlight=False)
|
|
console.print(markup, end="")
|
|
ansi_output: str = buf.getvalue()
|
|
|
|
if "\033[" not in ansi_output:
|
|
return
|
|
|
|
styled = False
|
|
for match in _ANSI_ESCAPE_RE.finditer(ansi_output):
|
|
params: str = match.group(1)
|
|
if params == "0" or params == "":
|
|
styled = False
|
|
else:
|
|
styled = True
|
|
|
|
assert not styled, (
|
|
f"ANSI styles still active at end of output — likely unclosed Rich tag.\n"
|
|
f"Last 200 chars of ANSI output: {ansi_output[-200:]!r}"
|
|
)
|
|
|
|
|
|
def make_tensor_info(
|
|
shape: Optional[list[int]] = None,
|
|
dtype: str = "torch.float32",
|
|
stats: Optional[TensorStats] = None,
|
|
sample: Optional[str] = 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,
|
|
)
|