sgl-project--sglang/python/sglang/test/kv_canary/e2e_base.py

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.srt.utils import kill_process_tree
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,
)

# 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:
            kill_process_tree(cls.process.pid)
        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

    _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: int = 8,
        *,
        assert_all_success: bool = True,
        max_new_tokens: int = 2048,
        timeout: float = 240.0,
        ignore_eos: Optional[bool] = None,
    ) -> list[dict]:
        """Fan out n parallel /generate requests; return list of response dicts."""
        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)
    ]