252 lines
10 KiB
Python
252 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import io
|
|
import os
|
|
import string
|
|
import time
|
|
from typing import ClassVar, Literal, Optional
|
|
|
|
from sglang.srt.kv_canary.config import CanaryMode
|
|
from sglang.srt.kv_canary.runner.swa_divergence import SwaDivergenceLog
|
|
from sglang.test.kv_canary.mode_config import _MODE_CONFIGS, _ModeConfig
|
|
from sglang.test.kv_canary.utils import build_canary_server_args, post_parallel_generate
|
|
from sglang.test.kv_canary.violation_assert_mixin import CanaryViolationAssertMixin
|
|
from sglang.test.test_utils import (
|
|
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
DEFAULT_URL_FOR_TEST,
|
|
CustomTestCase,
|
|
popen_launch_server,
|
|
terminate_and_kill_process_tree,
|
|
)
|
|
|
|
# Long prompt body shared by all canary e2e tests. The repetition count is chosen
|
|
# so the tokenised prompt comfortably exceeds the SWA sliding window of swa-mode
|
|
# fixtures (gemma-4-E2B); short prompts would never exercise the SWA-windowed
|
|
# verify path. Token count is roughly 7k after BPE.
|
|
_LONG_PROMPT_BODY = ("The quick brown fox jumps over the lazy dog. " * 700).strip()
|
|
_UNIQUE_PROMPT_FIRST_CHARS = string.ascii_letters + string.digits
|
|
|
|
|
|
class CapturedServerE2EBase(CanaryViolationAssertMixin, CustomTestCase):
|
|
process: ClassVar[Optional[object]] = None
|
|
base_url: ClassVar[str] = DEFAULT_URL_FOR_TEST
|
|
_stdout_buf: ClassVar[Optional[io.StringIO]] = None
|
|
_stderr_buf: ClassVar[Optional[io.StringIO]] = None
|
|
|
|
@classmethod
|
|
def tearDownClass(cls) -> None:
|
|
if cls.process is not None:
|
|
terminate_and_kill_process_tree(cls.process, wait_timeout=60)
|
|
for buf in (cls._stdout_buf, cls._stderr_buf):
|
|
if buf is not None:
|
|
buf.close()
|
|
cls._stdout_buf = None
|
|
cls._stderr_buf = None
|
|
|
|
def _captured_log_text(
|
|
self, side: Optional[Literal["prefill", "decode"]] = None
|
|
) -> str:
|
|
stdout_text = (
|
|
self._stdout_buf.getvalue() if self._stdout_buf is not None else ""
|
|
)
|
|
stderr_text = (
|
|
self._stderr_buf.getvalue() if self._stderr_buf is not None else ""
|
|
)
|
|
return stdout_text + stderr_text
|
|
|
|
def assert_log_contains(self, substring: str) -> None:
|
|
log_text = self._captured_log_text()
|
|
if substring not in log_text:
|
|
raise AssertionError(
|
|
f"Expected substring {substring!r} not found in captured log. "
|
|
f"Log tail:\n{log_text[-2000:]}"
|
|
)
|
|
|
|
|
|
class CanaryE2EBase(CapturedServerE2EBase):
|
|
model_mode: ClassVar[Literal["mha", "swa", "dsv4"]]
|
|
kv_canary_mode: ClassVar[CanaryMode]
|
|
extra_env: ClassVar[dict[str, str]] = {}
|
|
extra_server_args: ClassVar[tuple[str, ...]] = ()
|
|
use_unique_prompts: ClassVar[bool] = False
|
|
# SWA divergence assertions need slot recycling across batches; setting > 1 makes the
|
|
# test methods send N sequential batches so the SWA allocator's full→swa index mapping
|
|
# diverges from identity. Default 1 keeps MHA tests fast.
|
|
workload_n_batches: ClassVar[int] = 1
|
|
# Default workload for send_parallel_requests, tuned for the CUDA-kernel canary
|
|
# path; a slower backend retunes its own subclass here instead of passing sizes at
|
|
# every call site.
|
|
default_parallel_n: ClassVar[int] = 8
|
|
default_max_new_tokens: ClassVar[int] = 2048
|
|
default_request_timeout: ClassVar[float] = 240.0
|
|
|
|
_cfg: ClassVar[Optional[_ModeConfig]] = None
|
|
|
|
@classmethod
|
|
def setUpClass(cls) -> None:
|
|
cls._cfg = _MODE_CONFIGS[cls.model_mode]
|
|
server_env = os.environ.copy()
|
|
server_env.setdefault("SGLANG_KV_CANARY_ENABLE_VERIFY_TOKEN_ASSERT", "1")
|
|
server_env.update(cls.extra_env)
|
|
if cls.model_mode == "swa":
|
|
server_env.setdefault(
|
|
"SGLANG_KV_CANARY_SWA_DIVERGENCE_STATS_INTERVAL", "20"
|
|
)
|
|
# SWA mode uses google/gemma-4-E2B-it, whose forward does a
|
|
# ``positions += 1`` in-place. canary's WRITE/VERIFY require
|
|
# forward_batch.positions to stay 0-indexed, so flip the gemma
|
|
# path to out-of-place shift for these tests.
|
|
server_env.setdefault("SGLANG_GEMMA_OUT_OF_PLACE_POSITION_MUTATION", "1")
|
|
|
|
cls._stdout_buf = io.StringIO()
|
|
cls._stderr_buf = io.StringIO()
|
|
|
|
server_args = build_canary_server_args(
|
|
kv_canary_mode=cls.kv_canary_mode,
|
|
mode_cfg=cls._cfg,
|
|
extra_server_args=(
|
|
"--max-total-tokens",
|
|
"65536",
|
|
"--skip-server-warmup",
|
|
*cls.extra_server_args,
|
|
),
|
|
)
|
|
cls.process = popen_launch_server(
|
|
cls._cfg.model_path,
|
|
cls.base_url,
|
|
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
|
other_args=server_args,
|
|
env=server_env,
|
|
return_stdout_stderr=(cls._stdout_buf, cls._stderr_buf),
|
|
)
|
|
|
|
def make_prompts(self, n: int) -> list[str]:
|
|
if self.use_unique_prompts:
|
|
return _make_unique_prompts(n)
|
|
return [_LONG_PROMPT_BODY] * n
|
|
|
|
def send_parallel_requests(
|
|
self,
|
|
n: Optional[int] = None,
|
|
*,
|
|
assert_all_success: bool = True,
|
|
max_new_tokens: Optional[int] = None,
|
|
timeout: Optional[float] = None,
|
|
ignore_eos: Optional[bool] = None,
|
|
) -> list[dict]:
|
|
"""Fan out n parallel /generate requests; return list of response dicts.
|
|
|
|
Unset sizes fall back to the ``default_*`` class attributes, so a subclass can
|
|
retune the whole workload for its backend in one place.
|
|
"""
|
|
if n is None:
|
|
n = self.default_parallel_n
|
|
if max_new_tokens is None:
|
|
max_new_tokens = self.default_max_new_tokens
|
|
if timeout is None:
|
|
timeout = self.default_request_timeout
|
|
if ignore_eos is None:
|
|
ignore_eos = self.model_mode == "swa"
|
|
results = post_parallel_generate(
|
|
url=self.base_url + "/generate",
|
|
prompts=self.make_prompts(n),
|
|
max_new_tokens=max_new_tokens,
|
|
timeout=timeout,
|
|
ignore_eos=ignore_eos,
|
|
)
|
|
if assert_all_success:
|
|
for result in results:
|
|
self.assertEqual(result.get("status_code"), 200, result)
|
|
return results
|
|
|
|
def maybe_assert_swa_divergence_observed(self) -> None:
|
|
if self.model_mode == "swa":
|
|
self.assert_swa_divergence_observed()
|
|
|
|
def assert_swa_divergence_observed(
|
|
self,
|
|
*,
|
|
min_swa_out_of_window_tokens: int = 1,
|
|
min_swa_full_idx_divergence: int = 1,
|
|
require_verify_lag: bool = True,
|
|
flush_wait_seconds: float = 3.0,
|
|
max_retries: int = 10,
|
|
) -> None:
|
|
"""Assert that the SWA path was genuinely exercised.
|
|
|
|
Three signals must all hold:
|
|
- ``swa_out_of_window_tokens >= 1``: at least one token has slid out of the
|
|
sliding window (its SWA mapping is 0). This only appears once a request decodes
|
|
past the window, so the window evicts — proves the SWA window slide actually ran.
|
|
- ``swa_full_idx_divergence >= 1``: SWA pool has actually remapped at least one
|
|
slot to a non-identity index (i.e. real slot reuse / eviction occurred). The
|
|
workload must drive SWA pool pressure for this to fire — required because the
|
|
"pool reuse" path is the one production hits under sustained long-context
|
|
traffic, and we must keep it covered.
|
|
- ``verify_swa < verify_full``: SWA verify kernel processed fewer tokens than
|
|
FULL — proves both kernel groups ran and the window short-circuited SWA.
|
|
|
|
The first two signals are checked as the *peak* across all sampled forwards, not
|
|
only the last sample. The divergence reporter snapshots one live forward batch per
|
|
interval; under PP it snapshots a single micro-batch, which may hold only in-window
|
|
requests even when another micro-batch diverged. "Was the SWA path ever exercised?"
|
|
is a max-over-samples question, so a trailing in-window sample must not mask an
|
|
earlier diverging one. ``verify_swa``/``verify_full`` are monotonic running totals,
|
|
so the lag check reads the last sample.
|
|
"""
|
|
samples: list[tuple[SwaDivergenceLog, str]] = []
|
|
for _ in range(max_retries):
|
|
time.sleep(flush_wait_seconds)
|
|
samples = SwaDivergenceLog.find_all(self._captured_log_text())
|
|
if samples:
|
|
break
|
|
|
|
if not samples:
|
|
raise AssertionError(
|
|
"No kv_canary swa_divergence line found in server log after "
|
|
f"{max_retries} retries (wait={flush_wait_seconds}s each). "
|
|
f"Log tail:\n{self._captured_log_text()[-2000:]}"
|
|
)
|
|
|
|
peak_out_of_window = max(p.swa_out_of_window_tokens for p, _ in samples)
|
|
peak_full_idx_divergence = max(p.swa_full_idx_divergence for p, _ in samples)
|
|
last_parsed, last_line = samples[-1]
|
|
|
|
if peak_out_of_window < min_swa_out_of_window_tokens:
|
|
raise AssertionError(
|
|
f"SWA path not exercised: peak swa_out_of_window_tokens={peak_out_of_window} "
|
|
f"< min={min_swa_out_of_window_tokens} across {len(samples)} samples. "
|
|
f"Last line: {last_line}"
|
|
)
|
|
if peak_full_idx_divergence < min_swa_full_idx_divergence:
|
|
raise AssertionError(
|
|
f"SWA pool reuse not exercised: peak swa_full_idx_divergence={peak_full_idx_divergence} "
|
|
f"< min={min_swa_full_idx_divergence} across {len(samples)} samples. The workload "
|
|
f"did not drive enough SWA pool pressure to force slot remap. Last line: {last_line}"
|
|
)
|
|
if require_verify_lag and not (
|
|
last_parsed.verify_swa < last_parsed.verify_full
|
|
):
|
|
raise AssertionError(
|
|
f"SWA path not exercised: verify_swa={last_parsed.verify_swa} "
|
|
f"not strictly less than verify_full={last_parsed.verify_full}. "
|
|
f"Line: {last_line}"
|
|
)
|
|
|
|
|
|
def _make_unique_prompts(n: int) -> list[str]:
|
|
if n > len(_UNIQUE_PROMPT_FIRST_CHARS):
|
|
raise ValueError(
|
|
f"unique prompt count {n} exceeds supported count "
|
|
f"{len(_UNIQUE_PROMPT_FIRST_CHARS)}"
|
|
)
|
|
|
|
return [
|
|
(
|
|
f"{_UNIQUE_PROMPT_FIRST_CHARS[i]}"
|
|
f"{hex(i * 0x9E3779B1 & 0xFFFFFFFF)[2:]} "
|
|
f"{_LONG_PROMPT_BODY}"
|
|
)
|
|
for i in range(n)
|
|
]
|