sgl-project--sglang/python/sglang/test/server_fixtures/streaming_session_fixture.py

"""Streaming-session test fixture.

`TestStreamingSession` is the base class for all streaming-session tests
(default config — Llama-3.1-8B, no spec). Variants in
test_streaming_session.py and test_streaming_session_extra.py inherit
it and only override `setUpClass`.

Also exports:
- ABORT_REPRO_* constants used by the basic file's abort-leak repro.
- _abort_repro_run_all coroutine reused by the basic file.

Lives under sglang.test.server_fixtures so siblings under test/registered
can `import` it without sys.path hacks.
"""

import asyncio
import json
from typing import Any, Optional

import aiohttp
import requests

from sglang.srt.environ import envs
from sglang.srt.utils import kill_process_tree
from sglang.srt.utils.hf_transformers_utils import get_tokenizer
from sglang.test.test_utils import (
    DEFAULT_SMALL_MODEL_NAME_FOR_TEST,
    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
    DEFAULT_URL_FOR_TEST,
    CustomTestCase,
    popen_launch_server,
)

LOGPROB_PROMPTS = [
    "The quick brown fox jumps over the lazy dog.",
    "Pack my box with five dozen liquor jugs.",
    "How vexingly quick daft zebras jump.",
    "Sphinx of black quartz judge my vow.",
    "The five boxing wizards jump quickly.",
]

# Long enough to trigger chunked prefill at 200+ tokens per slice.
LEAK_FILLER = (
    "The quick brown fox jumps over the lazy dog. "
    "Pack my box with five dozen liquor jugs. "
    "How vexingly quick daft zebras jump. "
    "Sphinx of black quartz, judge my vow. "
    "The five boxing wizards jump quickly. "
    "Jackdaws love my big sphinx of quartz. "
    "A wizard's job is to vex chumps quickly in fog. "
    "We promptly judged antique ivory buckles for the next prize. "
) * 20

ABORT_REPRO_CONTEXT_LEN = 512
ABORT_REPRO_PAGE_SIZE = 256
ABORT_REPRO_GEN_LEN = 4
ABORT_REPRO_SESSIONS = 4
ABORT_REPRO_WARMUP_TURNS = 1
ABORT_REPRO_ROUNDS = 8
ABORT_REPRO_STREAM_TOKENS = 16
ABORT_REPRO_ABORT_TOKENS = 600
ABORT_REPRO_NON_STREAMING_TOKENS = 16
ABORT_REPRO_CHUNKED_PREFILL_SIZE = 4096

CONCURRENT_LOGPROB_SESSIONS = 6
CONCURRENT_LOGPROB_TURNS = 5
CONCURRENT_LOGPROB_ROUNDS = 10

STRESS_NUM_SESSIONS = 8
STRESS_NUM_NON_STREAMING = 4
STRESS_NUM_TURNS = 6
STRESS_GEN_LEN = 16


def _make_token_sized_ids(
    tokenizer: Any, prefix: str, min_tokens: int, max_tokens: Optional[int] = None
) -> list[int]:
    text = prefix
    chunk = " pack quartz wizard sphinx zebra fox " * 16
    token_ids = tokenizer.encode(text)
    while len(token_ids) < min_tokens:
        text += chunk
        token_ids = tokenizer.encode(text)
    if max_tokens is not None:
        token_ids = token_ids[:max_tokens]
    return token_ids


async def _abort_repro_generate(
    base_url: str,
    session: aiohttp.ClientSession,
    input_ids: list[int],
    max_new_tokens: int,
    session_params: Optional[dict[str, Any]] = None,
    expect_abort: bool = False,
) -> Optional[dict[str, Any]]:
    payload: dict[str, Any] = {
        "input_ids": input_ids,
        "sampling_params": {
            "temperature": 0,
            "max_new_tokens": max_new_tokens,
            "no_stop_trim": True,
            "skip_special_tokens": False,
        },
    }
    if session_params:
        payload["session_params"] = session_params

    async with session.post(base_url + "/generate", json=payload) as resp:
        text = await resp.text()
        if expect_abort:
            if resp.status == 200:
                data = json.loads(text)
                finish_reason = data.get("meta_info", {}).get("finish_reason", {})
                assert finish_reason.get("type") == "abort", text
                assert "maximum allowed length" in finish_reason.get(
                    "message", ""
                ) or "context length" in finish_reason.get("message", ""), text
                return data
            assert resp.status == 400, text
            assert "maximum allowed length" in text or "context length" in text, text
            return None

        assert resp.status == 200, text
        data = json.loads(text)
        finish_reason = data.get("meta_info", {}).get("finish_reason", {})
        assert finish_reason.get("type") != "abort", text
        return data


async def _abort_repro_run_all(base_url: str, tokenizer: Any) -> None:
    timeout = aiohttp.ClientTimeout(total=300)
    async with aiohttp.ClientSession(timeout=timeout) as http:
        session_ids = []
        for _ in range(ABORT_REPRO_SESSIONS):
            async with http.post(
                base_url + "/open_session",
                json={"capacity_of_str_len": 50000, "streaming": True},
            ) as resp:
                assert resp.status == 200, await resp.text()
                session_ids.append(await resp.json())

        try:
            for warmup_turn in range(ABORT_REPRO_WARMUP_TURNS):
                warmup_tasks = []
                for session_idx, session_id in enumerate(session_ids):
                    input_ids = _make_token_sized_ids(
                        tokenizer,
                        prefix=f"[warmup={warmup_turn} session={session_idx}]",
                        min_tokens=ABORT_REPRO_STREAM_TOKENS,
                        max_tokens=ABORT_REPRO_STREAM_TOKENS + 8,
                    )
                    warmup_tasks.append(
                        _abort_repro_generate(
                            base_url,
                            http,
                            input_ids,
                            ABORT_REPRO_GEN_LEN,
                            session_params={"id": session_id, "rid": None},
                        )
                    )
                await asyncio.gather(*warmup_tasks)

            for round_idx in range(ABORT_REPRO_ROUNDS):
                mixed_tasks = []
                for session_idx, session_id in enumerate(session_ids):
                    input_ids = _make_token_sized_ids(
                        tokenizer,
                        prefix=f"[round={round_idx} ok session={session_idx}]",
                        min_tokens=ABORT_REPRO_STREAM_TOKENS,
                        max_tokens=ABORT_REPRO_STREAM_TOKENS + 8,
                    )
                    mixed_tasks.append(
                        _abort_repro_generate(
                            base_url,
                            http,
                            input_ids,
                            ABORT_REPRO_GEN_LEN,
                            session_params={"id": session_id, "rid": None},
                        )
                    )

                for ns_idx in range(2):
                    input_ids = _make_token_sized_ids(
                        tokenizer,
                        prefix=f"[round={round_idx} ns={ns_idx}]",
                        min_tokens=ABORT_REPRO_NON_STREAMING_TOKENS,
                        max_tokens=ABORT_REPRO_NON_STREAMING_TOKENS + 8,
                    )
                    mixed_tasks.append(
                        _abort_repro_generate(
                            base_url,
                            http,
                            input_ids,
                            ABORT_REPRO_GEN_LEN,
                        )
                    )
                await asyncio.gather(*mixed_tasks)

                abort_tasks = []
                for session_idx, session_id in enumerate(session_ids):
                    input_ids = _make_token_sized_ids(
                        tokenizer,
                        prefix=f"[round={round_idx} abort session={session_idx}]",
                        min_tokens=ABORT_REPRO_ABORT_TOKENS,
                    )
                    abort_tasks.append(
                        _abort_repro_generate(
                            base_url,
                            http,
                            input_ids,
                            ABORT_REPRO_GEN_LEN,
                            session_params={"id": session_id, "rid": None},
                            expect_abort=True,
                        )
                    )
                await asyncio.gather(*abort_tasks)

                recovery_tasks = []
                for session_idx, session_id in enumerate(session_ids):
                    input_ids = _make_token_sized_ids(
                        tokenizer,
                        prefix=f"[round={round_idx} recover session={session_idx}]",
                        min_tokens=ABORT_REPRO_NON_STREAMING_TOKENS,
                        max_tokens=ABORT_REPRO_NON_STREAMING_TOKENS + 8,
                    )
                    recovery_tasks.append(
                        _abort_repro_generate(
                            base_url,
                            http,
                            input_ids,
                            ABORT_REPRO_GEN_LEN,
                            session_params={"id": session_id, "rid": None},
                        )
                    )
                recovery_results = await asyncio.gather(*recovery_tasks)
                for result in recovery_results:
                    assert result is not None
                    assert result["meta_info"]["cached_tokens"] > 0, result

                health = requests.get(base_url + "/health", timeout=10)
                if health.status_code != 200:
                    raise RuntimeError(
                        f"server unhealthy after round={round_idx}: "
                        f"{health.status_code} {health.text}"
                    )
        finally:
            for session_id in session_ids:
                async with http.post(
                    base_url + "/close_session", json={"session_id": session_id}
                ) as resp:
                    assert resp.status == 200, await resp.text()


async def _async_generate(
    base_url: str,
    session: aiohttp.ClientSession,
    input_ids: list[int],
    max_new_tokens: int = 8,
    session_params: Optional[dict[str, Any]] = None,
    return_logprob: bool = False,
    logprob_start_len: Optional[int] = None,
) -> dict[str, Any]:
    payload: dict[str, Any] = {
        "input_ids": input_ids,
        "sampling_params": {
            "temperature": 0,
            "max_new_tokens": max_new_tokens,
            "no_stop_trim": True,
            "skip_special_tokens": False,
        },
    }
    if session_params:
        payload["session_params"] = session_params
    if return_logprob:
        payload["return_logprob"] = True
    if logprob_start_len is not None:
        payload["logprob_start_len"] = logprob_start_len
    timeout = aiohttp.ClientTimeout(total=300)
    async with session.post(
        base_url + "/generate", json=payload, timeout=timeout
    ) as resp:
        assert resp.status == 200, f"Generate failed: {await resp.text()}"
        return await resp.json()


async def _concurrent_logprob_run(base_url: str, tokenizer: Any, **gen_kwargs) -> None:
    """N sessions per round, all requests fired simultaneously per turn so
    the running batch has real concurrency (retract can actually kick one).
    """
    timeout = aiohttp.ClientTimeout(total=300)
    async with aiohttp.ClientSession(timeout=timeout) as http:
        for _ in range(CONCURRENT_LOGPROB_ROUNDS):
            sids: list[str] = []
            for _ in range(CONCURRENT_LOGPROB_SESSIONS):
                async with http.post(
                    base_url + "/open_session",
                    json={"capacity_of_str_len": 50000, "streaming": True},
                ) as resp:
                    assert resp.status == 200
                    sids.append(await resp.json())

            rids: list[Optional[str]] = [None] * CONCURRENT_LOGPROB_SESSIONS
            for turn in range(CONCURRENT_LOGPROB_TURNS):
                tasks = []
                for s in range(CONCURRENT_LOGPROB_SESSIONS):
                    text = (
                        f"S{s} T{turn}: "
                        f"{LOGPROB_PROMPTS[turn % len(LOGPROB_PROMPTS)]}"
                    )
                    ids = tokenizer.encode(text)
                    tasks.append(
                        _async_generate(
                            base_url,
                            http,
                            ids,
                            session_params={"id": sids[s], "rid": rids[s]},
                            **gen_kwargs,
                        )
                    )
                results = await asyncio.gather(*tasks)
                for s in range(CONCURRENT_LOGPROB_SESSIONS):
                    rids[s] = results[s]["meta_info"]["id"]

            for sid in sids:
                async with http.post(
                    base_url + "/close_session", json={"session_id": sid}
                ) as resp:
                    assert resp.status == 200


async def _stress_run_all(base_url: str, tokenizer: Any) -> None:
    """Streaming + non-streaming mixed batches under retract pressure.
    Long prompts (~200+ tokens) trigger chunked prefill so retract can
    interrupt mid-extend.
    """
    timeout = aiohttp.ClientTimeout(total=300)
    async with aiohttp.ClientSession(timeout=timeout) as http:
        sids: list[str] = []
        for _ in range(STRESS_NUM_SESSIONS):
            async with http.post(
                base_url + "/open_session",
                json={"capacity_of_str_len": 50000, "streaming": True},
            ) as resp:
                assert resp.status == 200
                sids.append(await resp.json())

        rids: list[Optional[str]] = [None] * STRESS_NUM_SESSIONS
        for turn in range(STRESS_NUM_TURNS):
            tasks = []
            # Streaming requests — long prompts to trigger chunked prefill.
            for s in range(STRESS_NUM_SESSIONS):
                offset = (s * STRESS_NUM_TURNS + turn) * 200
                text = (
                    f"Session {s} turn {turn}: " f"{LEAK_FILLER[offset : offset + 800]}"
                )
                ids = tokenizer.encode(text)
                tasks.append(
                    _async_generate(
                        base_url,
                        http,
                        ids,
                        max_new_tokens=STRESS_GEN_LEN,
                        session_params={"id": sids[s], "rid": rids[s]},
                    )
                )

            # Non-streaming requests interleaved.
            for ns in range(STRESS_NUM_NON_STREAMING):
                text = (
                    f"Non-streaming {ns} turn {turn}: "
                    f"{LEAK_FILLER[ns * 100 : ns * 100 + 400]}"
                )
                ids = tokenizer.encode(text)
                tasks.append(
                    _async_generate(
                        base_url,
                        http,
                        ids,
                        max_new_tokens=STRESS_GEN_LEN,
                    )
                )

            results = await asyncio.gather(*tasks)
            for s in range(STRESS_NUM_SESSIONS):
                rids[s] = results[s]["meta_info"]["id"]

        for sid in sids:
            async with http.post(
                base_url + "/close_session", json={"session_id": sid}
            ) as resp:
                assert resp.status == 200


class StreamingSessionServerBase(CustomTestCase):
    """Minimal streaming-session server fixture.

    Subclasses override class attrs to customize launch:
    - `model`: defaults to the small model.
    - `extra_args`: appended after `--enable-streaming-session` (set
      `--chunked-prefill-size`, `--page-size`, spec args, etc. here).
    - `env_overrides`: list of `(env_attr_name, value)` tuples; each is
      pushed onto the `setUpClass` context stack so the env override is
      live during `popen_launch_server` and torn down on
      `tearDownClass`-time. `SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY=1`
      is always applied on top of these.
    """

    model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
    base_url = DEFAULT_URL_FOR_TEST
    extra_args: list = []
    env_overrides: list = []

    @classmethod
    def setUpClass(cls):
        import contextlib

        with contextlib.ExitStack() as stack:
            stack.enter_context(
                envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(1)
            )
            stack.enter_context(envs.SGLANG_CHECK_KV_PAGE_INVARIANTS.override(True))
            for name, val in cls.env_overrides:
                stack.enter_context(getattr(envs, name).override(val))
            cls.process = popen_launch_server(
                cls.model,
                cls.base_url,
                timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
                other_args=["--enable-streaming-session"] + list(cls.extra_args),
            )
        cls.tokenizer = get_tokenizer(cls.model)

    @classmethod
    def tearDownClass(cls):
        kill_process_tree(cls.process.pid)