minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix streaming session busy check double-counting; add compat CI tests (#22213)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-04-12 01:48:16 -07:00
committed by GitHub
parent 3f60df8012
commit f1eb4ca90c
2 changed files with 111 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/sessions/test_streaming_session.py
+95 -27
View File
@@ -17,6 +17,7 @@ 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.ci.ci_register import register_cuda_ci
@@ -205,36 +206,28 @@ async def _leak_run_all(base_url: str, tokenizer: Any) -> None:
assert resp.status == 200
# ===================================================================
# Test class
# ===================================================================
class TestStreamingSession(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-streaming-session",
"--chunked-prefill-size",
"512",
],
)
with envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(2):
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-streaming-session",
"--chunked-prefill-size",
"512",
],
)
cls.tokenizer = get_tokenizer(cls.model)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
# ------------------------------------------------------------------
# KV cache mechanics
# ------------------------------------------------------------------
def test_kv_cache_inheritance(self, gen_len=12):
"""Verify KV inheritance, radix cache insertion, and flush reclamation."""
chunks = [
@@ -346,10 +339,6 @@ class TestStreamingSession(CustomTestCase):
"After session close + flush, cache should be fully reclaimed",
)
# ------------------------------------------------------------------
# Logprob leak tests
# ------------------------------------------------------------------
def test_leak_logprob_none(self) -> None:
"""Streaming sessions without logprobs must not leak tokens."""
_logprob_assert_no_leak(self.base_url, self.tokenizer)
@@ -367,10 +356,6 @@ class TestStreamingSession(CustomTestCase):
logprob_start_len=0,
)
# ------------------------------------------------------------------
# Chunked prefill leak test
# ------------------------------------------------------------------
def test_leak_chunked_prefill(self) -> None:
"""Concurrent multi-turn streaming sessions then idle health check."""
requests.post(self.base_url + "/flush_cache")
@@ -399,5 +384,88 @@ class TestStreamingSession(CustomTestCase):
)
class TestStreamingSessionMixedChunk(TestStreamingSession):
"""Streaming session with --enable-mixed-chunk."""
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
with envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(2):
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-streaming-session",
"--chunked-prefill-size",
"512",
"--enable-mixed-chunk",
],
)
cls.tokenizer = get_tokenizer(cls.model)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
@unittest.skip("streaming session + retract has a token leak — tracked separately")
class TestStreamingSessionRetract(TestStreamingSession):
"""Streaming session under retract decode pressure."""
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
with envs.SGLANG_TEST_RETRACT.override(
True
), envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(2):
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-streaming-session",
"--chunked-prefill-size",
"128",
],
)
cls.tokenizer = get_tokenizer(cls.model)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
@unittest.skip("streaming session + retract has a token leak — tracked separately")
class TestStreamingSessionRetractMixedChunk(TestStreamingSession):
"""Streaming session under retract decode with --enable-mixed-chunk."""
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
with envs.SGLANG_TEST_RETRACT.override(
True
), envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.override(2):
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-streaming-session",
"--chunked-prefill-size",
"128",
"--enable-mixed-chunk",
],
)
cls.tokenizer = get_tokenizer(cls.model)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
if __name__ == "__main__":
unittest.main()