From 6ad78b008b1215deff83dc61823a80df8215bc9a Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 19 Mar 2026 19:13:28 -0700 Subject: [PATCH] =?UTF-8?q?Consolidate=20streaming=20session=20tests=20(4?= =?UTF-8?q?=20files=20=E2=86=92=203,=20fewer=20server=20starts)=20(#20961)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../sessions/test_session_control.py | 159 ------- .../sessions/test_session_latency.py | 8 +- .../sessions/test_session_logprob_leak.py | 138 ------ .../sessions/test_streaming_session.py | 403 ++++++++++++++++++ .../sessions/test_streaming_session_leak.py | 173 -------- 5 files changed, 407 insertions(+), 474 deletions(-) delete mode 100644 test/registered/sessions/test_session_logprob_leak.py create mode 100644 test/registered/sessions/test_streaming_session.py delete mode 100644 test/registered/sessions/test_streaming_session_leak.py diff --git a/test/registered/sessions/test_session_control.py b/test/registered/sessions/test_session_control.py index 2b495be04..b9c4a4459 100644 --- a/test/registered/sessions/test_session_control.py +++ b/test/registered/sessions/test_session_control.py @@ -3,13 +3,11 @@ Usage: python3 -m unittest test_session_control.TestSessionControl.test_session_control python3 -m unittest test_session_control.TestSessionControl.test_session_control_with_branching python3 -m unittest test_session_control.TestSessionControl.test_session_control_backtrack_with_abort -python3 -m unittest test_session_control.TestSessionControl.test_streaming_session python3 -m unittest test_session_control.TestSessionControlVision.test_session_control """ import asyncio import json -import time import unittest import aiohttp @@ -45,7 +43,6 @@ class TestSessionControl(unittest.TestCase): other_args=[ "--attention-backend", "triton", - "--enable-streaming-session", ], ) @@ -435,162 +432,6 @@ class TestSessionControl(unittest.TestCase): asyncio.run(self.run_session_control_backtrack_with_abort(replace=True)) asyncio.run(self.run_session_control_backtrack_with_abort(replace=False)) - def test_streaming_session(self, gen_len=12): - chunks = [ - "Let me tell you something about France.", - "The capital of France is", - "The population of the city is", - ] - tokenizer = get_tokenizer(self.model) - chunks_ids = [tokenizer.encode(x) for x in chunks] - for i in range(1, len(chunks_ids)): - if chunks_ids[i][0] == tokenizer.bos_token_id: - chunks_ids[i] = chunks_ids[i][1:] - - # === Part 1: streaming session === - requests.post(self.base_url + "/flush_cache") - session_id = requests.post( - self.base_url + "/open_session", - json={"capacity_of_str_len": 1000, "streaming": True}, - ).json() - rid = None - outputs_from_session = [] - - prev_kv_len = 0 - for turn_idx, chunk_ids in enumerate(chunks_ids): - response = requests.post( - self.base_url + "/generate", - json={ - "input_ids": chunk_ids, - "session_params": {"id": session_id, "rid": rid}, - "sampling_params": { - "temperature": 0, - "max_new_tokens": gen_len, - "no_stop_trim": True, - "skip_special_tokens": False, - }, - }, - ).json() - rid = response["meta_info"]["id"] - outputs_from_session.append(response["text"]) - cached = response["meta_info"]["cached_tokens"] - prompt_tokens = response["meta_info"]["prompt_tokens"] - completion_tokens = response["meta_info"]["completion_tokens"] - - if turn_idx == 0: - # Turn 1 should have no cache hit (cache was flushed). - self.assertEqual( - cached, 0, "Turn 1 should have 0 cached tokens (clean start)" - ) - else: - # Turns 2+ inherit KV from the previous turn (via inherit_kv_states, - # not radix tree matching). cached_tokens reflects the inherited prefix. - self.assertEqual( - cached, - prev_kv_len, - f"Turn {turn_idx + 1}: should inherit {prev_kv_len} KV tokens from previous turn", - ) - prev_kv_len = prompt_tokens + completion_tokens - - # Close the session before checking cache/memory state. - ret = requests.post( - self.base_url + "/close_session", - json={"session_id": session_id}, - ) - self.assertEqual(ret.status_code, 200) - - # === Cache verification (after close, before flush) === - - # Assertion 2: turn 1's prompt was inserted to the cache. - verify_resp = requests.post( - self.base_url + "/generate", - json={ - "input_ids": chunks_ids[0], - "sampling_params": {"temperature": 0, "max_new_tokens": 1}, - }, - ).json() - self.assertGreater( - verify_resp["meta_info"]["cached_tokens"], - 0, - "Turn 1's prompt should be cached in the radix tree", - ) - - # Assertion 3 (insertion): turn 2's prompt tokens should NOT be in cache. - # The tree should only contain turn 1's extent (prompt + output from - # cache_unfinished_req during decode). Turn 2's prompt starts fresh tokens - # that were never inserted. - verify_resp2 = requests.post( - self.base_url + "/generate", - json={ - "input_ids": chunks_ids[1], - "sampling_params": {"temperature": 0, "max_new_tokens": 1}, - }, - ).json() - self.assertEqual( - verify_resp2["meta_info"]["cached_tokens"], - 0, - "Turn 2's prompt should not be in cache (no insertion for turns 2+)", - ) - - # === Memory verification === - - # Assertion 4 & 5: KV is released properly and no memory leak. - # SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE is True by default; - # the scheduler will crash if it detects a leak during idle. - time.sleep(2) - health_resp = requests.get(self.base_url + "/health") - self.assertEqual( - health_resp.status_code, - 200, - "Server should be healthy after session close (no memory leak)", - ) - - # After flush, all cache should be reclaimed. - requests.post(self.base_url + "/flush_cache") - verify_resp3 = requests.post( - self.base_url + "/generate", - json={ - "input_ids": chunks_ids[0], - "sampling_params": {"temperature": 0, "max_new_tokens": 1}, - }, - ).json() - self.assertEqual( - verify_resp3["meta_info"]["cached_tokens"], - 0, - "After session close + flush, cache should be fully reclaimed", - ) - - # === Part 2: non-session baseline for output comparison === - requests.post(self.base_url + "/flush_cache") - - outputs_normal = [] - input_ids = chunks_ids[0][:] - for i in range(len(chunks_ids)): - response = requests.post( - self.base_url + "/generate", - json={ - "input_ids": input_ids, - "sampling_params": { - "temperature": 0, - "max_new_tokens": gen_len, - "no_stop_trim": True, - "skip_special_tokens": False, - }, - }, - ).json() - outputs_normal.append(response["text"]) - if i + 1 < len(chunks_ids): - out_ids = tokenizer.encode(response["text"]) - if out_ids and out_ids[0] == tokenizer.bos_token_id: - out_ids = out_ids[1:] - input_ids = input_ids + out_ids + chunks_ids[i + 1] - - print("outputs from streaming session:") - print(outputs_from_session) - print("outputs from normal queries:") - print(outputs_normal) - self.assertEqual(outputs_from_session, outputs_normal) - def run_session_control_with_branching( self, root_prompt, chunks_per_step, gen_len=16 ): diff --git a/test/registered/sessions/test_session_latency.py b/test/registered/sessions/test_session_latency.py index 774ae8a85..c4ec4d6a9 100644 --- a/test/registered/sessions/test_session_latency.py +++ b/test/registered/sessions/test_session_latency.py @@ -32,11 +32,11 @@ from sglang.test.test_utils import ( ) register_cuda_ci( - est_time=100, + est_time=122, suite="stage-b-test-large-1-gpu", ) -NUM_TURNS = 300 +NUM_TURNS = 150 INPUT_LEN = 16 GEN_LEN = 8 NUM_CONCURRENT = 8 @@ -417,8 +417,8 @@ class TestSessionLatency(CustomTestCase): speedup = reg_tail / stm_tail if stm_tail > 0 else float("inf") self.assertGreaterEqual( speedup, - 2.0, - f"streaming should be >=2x faster on last {TAIL_TURNS} turns " + 1.4, + f"streaming should be >=1.4x faster on last {TAIL_TURNS} turns " f"(regular={reg_tail:.1f}ms, streaming={stm_tail:.1f}ms, speedup={speedup:.2f}x)", ) diff --git a/test/registered/sessions/test_session_logprob_leak.py b/test/registered/sessions/test_session_logprob_leak.py deleted file mode 100644 index ebe2c1bd2..000000000 --- a/test/registered/sessions/test_session_logprob_leak.py +++ /dev/null @@ -1,138 +0,0 @@ -""" -Test for token leak in streaming sessions with return_logprob enabled. - -When logprob_start_len=0, init_next_round_input truncates the prefix match -key to length 0, which bypasses the session slot's committed KV and orphans -allocated tokens. This test verifies that sessions with logprobs enabled -do not leak tokens. - -Usage: - python3 -m pytest test_session_logprob_leak.py -xvs -""" - -import time -import unittest -from typing import Any - -import requests - -from sglang.srt.utils import kill_process_tree -from sglang.test.ci.ci_register import register_cuda_ci -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, -) - -register_cuda_ci(est_time=120, suite="stage-b-test-large-1-gpu") - -NUM_TURNS = 5 -NUM_ROUNDS = 30 - -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.", -] - - -def _generate(base_url, input_ids, **kwargs) -> dict: - payload: dict[str, Any] = { - "input_ids": input_ids, - "sampling_params": { - "temperature": 0, - "max_new_tokens": kwargs.get("max_new_tokens", 8), - "no_stop_trim": True, - "skip_special_tokens": False, - }, - } - if kwargs.get("return_logprob"): - payload["return_logprob"] = True - if kwargs.get("logprob_start_len") is not None: - payload["logprob_start_len"] = kwargs["logprob_start_len"] - if kwargs.get("session_params"): - payload["session_params"] = kwargs["session_params"] - resp = requests.post(base_url + "/generate", json=payload, timeout=120) - assert resp.status_code == 200, f"Generate failed: {resp.text}" - return resp.json() - - -def _run_one_session(base_url, tokenizer, **gen_kwargs): - """Open session → N turns → close.""" - resp = requests.post( - base_url + "/open_session", - json={"capacity_of_str_len": 50000, "streaming": True}, - ) - assert resp.status_code == 200 - session_id = resp.json() - - rid = None - for turn in range(NUM_TURNS): - turn_ids = tokenizer.encode(f"Turn {turn}: {PROMPTS[turn % len(PROMPTS)]}") - result = _generate( - base_url, - turn_ids, - session_params={"id": session_id, "rid": rid}, - **gen_kwargs, - ) - rid = result["meta_info"]["id"] - - requests.post(base_url + "/close_session", json={"session_id": session_id}) - - -def _assert_no_leak(base_url, tokenizer, **gen_kwargs): - """Run many session rounds and verify server stays healthy.""" - requests.post(base_url + "/flush_cache") - for _ in range(NUM_ROUNDS): - _run_one_session(base_url, tokenizer, **gen_kwargs) - time.sleep(3) - assert ( - requests.get(base_url + "/health").status_code == 200 - ), "Server unhealthy — likely a token memory leak." - - -class TestSessionLogprobLeak(CustomTestCase): - @classmethod - def setUpClass(cls) -> None: - 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"], - ) - - @classmethod - def tearDownClass(cls) -> None: - kill_process_tree(cls.process.pid) - - def _tokenizer(self): - from sglang.srt.utils.hf_transformers_utils import get_tokenizer - - return get_tokenizer(self.model) - - def test_session_without_logprob(self) -> None: - """Streaming sessions without logprobs must not leak tokens.""" - _assert_no_leak(self.base_url, self._tokenizer()) - - def test_session_with_output_logprob(self) -> None: - """Streaming sessions with output logprobs must not leak tokens.""" - _assert_no_leak(self.base_url, self._tokenizer(), return_logprob=True) - - def test_session_with_input_logprob(self) -> None: - """Streaming sessions with logprob_start_len=0 must not leak tokens.""" - _assert_no_leak( - self.base_url, - self._tokenizer(), - return_logprob=True, - logprob_start_len=0, - ) - - -if __name__ == "__main__": - unittest.main() diff --git a/test/registered/sessions/test_streaming_session.py b/test/registered/sessions/test_streaming_session.py new file mode 100644 index 000000000..14668b6da --- /dev/null +++ b/test/registered/sessions/test_streaming_session.py @@ -0,0 +1,403 @@ +""" +Streaming session tests: KV cache mechanics, logprob leak, chunked prefill leak. + +All tests share a single server (DEFAULT_SMALL_MODEL) with streaming sessions +and chunked prefill enabled. + +Usage: + python -m pytest test_streaming_session.py -xvs + python -m unittest test_streaming_session.TestStreamingSession +""" + +import asyncio +import time +import unittest +from typing import Any, Optional + +import aiohttp +import requests + +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 +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, +) + +register_cuda_ci(est_time=70, suite="stage-b-test-large-1-gpu") + +# --------------------------------------------------------------------------- +# Logprob leak constants +# --------------------------------------------------------------------------- + +LOGPROB_NUM_TURNS = 5 +LOGPROB_NUM_ROUNDS = 30 + +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.", +] + +# --------------------------------------------------------------------------- +# Chunked prefill leak constants +# --------------------------------------------------------------------------- + +LEAK_NUM_SESSIONS = 4 +LEAK_NUM_TURNS = 5 +LEAK_GEN_LEN = 16 + +# Filler text to trigger chunked prefill (200+ tokens per turn) +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 + + +# --------------------------------------------------------------------------- +# Logprob leak helpers +# --------------------------------------------------------------------------- + + +def _logprob_generate(base_url, input_ids, **kwargs) -> dict: + payload: dict[str, Any] = { + "input_ids": input_ids, + "sampling_params": { + "temperature": 0, + "max_new_tokens": kwargs.get("max_new_tokens", 8), + "no_stop_trim": True, + "skip_special_tokens": False, + }, + } + if kwargs.get("return_logprob"): + payload["return_logprob"] = True + if kwargs.get("logprob_start_len") is not None: + payload["logprob_start_len"] = kwargs["logprob_start_len"] + if kwargs.get("session_params"): + payload["session_params"] = kwargs["session_params"] + resp = requests.post(base_url + "/generate", json=payload, timeout=120) + assert resp.status_code == 200, f"Generate failed: {resp.text}" + return resp.json() + + +def _logprob_run_one_session(base_url, tokenizer, **gen_kwargs): + """Open session -> N turns -> close.""" + resp = requests.post( + base_url + "/open_session", + json={"capacity_of_str_len": 50000, "streaming": True}, + ) + assert resp.status_code == 200 + session_id = resp.json() + + rid = None + for turn in range(LOGPROB_NUM_TURNS): + turn_ids = tokenizer.encode( + f"Turn {turn}: {LOGPROB_PROMPTS[turn % len(LOGPROB_PROMPTS)]}" + ) + result = _logprob_generate( + base_url, + turn_ids, + session_params={"id": session_id, "rid": rid}, + **gen_kwargs, + ) + rid = result["meta_info"]["id"] + + requests.post(base_url + "/close_session", json={"session_id": session_id}) + + +def _logprob_assert_no_leak(base_url, tokenizer, **gen_kwargs): + """Run many session rounds and verify server stays healthy.""" + requests.post(base_url + "/flush_cache") + for _ in range(LOGPROB_NUM_ROUNDS): + _logprob_run_one_session(base_url, tokenizer, **gen_kwargs) + time.sleep(3) + assert ( + requests.get(base_url + "/health").status_code == 200 + ), "Server unhealthy — likely a token memory leak." + + +# --------------------------------------------------------------------------- +# Chunked prefill leak helpers +# --------------------------------------------------------------------------- + + +async def _leak_async_generate( + base_url: str, + session: aiohttp.ClientSession, + input_ids: list[int], + session_params: Optional[dict[str, Any]] = None, +) -> Any: + payload: dict[str, Any] = { + "input_ids": input_ids, + "sampling_params": { + "temperature": 0, + "max_new_tokens": LEAK_GEN_LEN, + "no_stop_trim": True, + "skip_special_tokens": False, + }, + } + if session_params: + payload["session_params"] = session_params + 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 _leak_run_all(base_url: str, tokenizer: Any) -> None: + """Fire all requests per turn simultaneously to create mixed batches.""" + timeout = aiohttp.ClientTimeout(total=300) + async with aiohttp.ClientSession(timeout=timeout) as http: + # Open all sessions + sids = [] + for s in range(LEAK_NUM_SESSIONS): + async with http.post( + base_url + "/open_session", + json={"capacity_of_str_len": 50000, "streaming": True}, + ) as resp: + sids.append(await resp.json()) + + # For each turn, fire ALL streaming + non-streaming requests at once + for turn in range(LEAK_NUM_TURNS): + tasks = [] + # Streaming requests for all sessions + for s in range(LEAK_NUM_SESSIONS): + offset = (s * LEAK_NUM_TURNS + turn) * 200 + text = f"Session {s} turn {turn}: {LEAK_FILLER[offset : offset + 1500]}" + ids = tokenizer.encode(text) + tasks.append( + _leak_async_generate( + base_url, + http, + ids, + session_params={"id": sids[s], "rid": None}, + ) + ) + + # Non-streaming requests interleaved + for ns in range(LEAK_NUM_SESSIONS // 2): + text = f"Non-streaming {ns} turn {turn}: {LEAK_FILLER[ns * 100 : ns * 100 + 500]}" + ids = tokenizer.encode(text) + tasks.append(_leak_async_generate(base_url, http, ids)) + + # Fire all at once — creates mixed batch of streaming + non-streaming + await asyncio.gather(*tasks) + + # Close all sessions + for sid in sids: + async with http.post( + base_url + "/close_session", json={"session_id": sid} + ) as resp: + 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", + ], + ) + 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 = [ + "Let me tell you something about France.", + "The capital of France is", + "The population of the city is", + ] + chunks_ids = [self.tokenizer.encode(x) for x in chunks] + for i in range(1, len(chunks_ids)): + if chunks_ids[i][0] == self.tokenizer.bos_token_id: + chunks_ids[i] = chunks_ids[i][1:] + + # === Part 1: streaming session — check KV inheritance === + requests.post(self.base_url + "/flush_cache") + session_id = requests.post( + self.base_url + "/open_session", + json={"capacity_of_str_len": 1000, "streaming": True}, + ).json() + rid = None + + prev_kv_len = 0 + for turn_idx, chunk_ids in enumerate(chunks_ids): + response = requests.post( + self.base_url + "/generate", + json={ + "input_ids": chunk_ids, + "session_params": {"id": session_id, "rid": rid}, + "sampling_params": { + "temperature": 0, + "max_new_tokens": gen_len, + "no_stop_trim": True, + "skip_special_tokens": False, + }, + }, + ).json() + rid = response["meta_info"]["id"] + cached = response["meta_info"]["cached_tokens"] + prompt_tokens = response["meta_info"]["prompt_tokens"] + completion_tokens = response["meta_info"]["completion_tokens"] + + if turn_idx == 0: + # Turn 1 should have no cache hit (cache was flushed). + self.assertEqual( + cached, 0, "Turn 1 should have 0 cached tokens (clean start)" + ) + else: + # Turns 2+ inherit KV from the previous turn (via inherit_kv_states, + # not radix tree matching). cached_tokens reflects the inherited prefix. + self.assertEqual( + cached, + prev_kv_len, + f"Turn {turn_idx + 1}: should inherit {prev_kv_len} KV tokens from previous turn", + ) + prev_kv_len = prompt_tokens + completion_tokens + + # Close the session before checking cache/memory state. + ret = requests.post( + self.base_url + "/close_session", + json={"session_id": session_id}, + ) + self.assertEqual(ret.status_code, 200) + + # === Cache verification (after close, before flush) === + + # Turn 1's prompt was inserted to the cache. + verify_resp = requests.post( + self.base_url + "/generate", + json={ + "input_ids": chunks_ids[0], + "sampling_params": {"temperature": 0, "max_new_tokens": 1}, + }, + ).json() + self.assertGreater( + verify_resp["meta_info"]["cached_tokens"], + 0, + "Turn 1's prompt should be cached in the radix tree", + ) + + # Turn 2's prompt tokens should NOT be in cache. + # The tree should only contain turn 1's extent (prompt + output from + # cache_unfinished_req during decode). Turn 2's prompt starts fresh tokens + # that were never inserted. + verify_resp2 = requests.post( + self.base_url + "/generate", + json={ + "input_ids": chunks_ids[1], + "sampling_params": {"temperature": 0, "max_new_tokens": 1}, + }, + ).json() + self.assertEqual( + verify_resp2["meta_info"]["cached_tokens"], + 0, + "Turn 2's prompt should not be in cache (no insertion for turns 2+)", + ) + + # === Flush reclamation === + + requests.post(self.base_url + "/flush_cache") + verify_resp3 = requests.post( + self.base_url + "/generate", + json={ + "input_ids": chunks_ids[0], + "sampling_params": {"temperature": 0, "max_new_tokens": 1}, + }, + ).json() + self.assertEqual( + verify_resp3["meta_info"]["cached_tokens"], + 0, + "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) + + def test_leak_logprob_output(self) -> None: + """Streaming sessions with output logprobs must not leak tokens.""" + _logprob_assert_no_leak(self.base_url, self.tokenizer, return_logprob=True) + + def test_leak_logprob_input(self) -> None: + """Streaming sessions with logprob_start_len=0 must not leak tokens.""" + _logprob_assert_no_leak( + self.base_url, + self.tokenizer, + return_logprob=True, + 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") + + asyncio.run(_leak_run_all(self.base_url, self.tokenizer)) + + # Run a few non-streaming requests to flush state + for i in range(3): + ids = self.tokenizer.encode(f"Flush request {i}: final cleanup.") + requests.post( + self.base_url + "/generate", + json={ + "input_ids": ids, + "sampling_params": {"temperature": 0, "max_new_tokens": 4}, + }, + ) + + # Wait for server to go idle and run memory check + time.sleep(5) + health = requests.get(self.base_url + "/health") + self.assertEqual( + health.status_code, + 200, + "Server unhealthy after streaming session close — " + "likely a token memory leak from streaming session lifecycle.", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/sessions/test_streaming_session_leak.py b/test/registered/sessions/test_streaming_session_leak.py deleted file mode 100644 index 65c4517a9..000000000 --- a/test/registered/sessions/test_streaming_session_leak.py +++ /dev/null @@ -1,173 +0,0 @@ -""" -Test for token leak in streaming sessions with chunked prefill. - -Runs concurrent multi-turn streaming sessions interleaved with non-streaming -requests (to create mixed batches), closes all sessions, waits for idle, -and checks server health. - -Usage: - python3 -m pytest test_streaming_session_leak.py -xvs -""" - -import asyncio -import time -import unittest -from typing import Any, Optional - -import aiohttp -import requests - -from sglang.srt.utils import kill_process_tree -from sglang.test.ci.ci_register import register_cuda_ci -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, -) - -register_cuda_ci(est_time=120, suite="stage-b-test-large-1-gpu") - -NUM_SESSIONS = 4 -NUM_TURNS = 5 -GEN_LEN = 16 - -# Filler text to trigger chunked prefill (200+ tokens per turn) -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 - - -async def _async_generate( - base_url: str, - session: aiohttp.ClientSession, - input_ids: list[int], - session_params: Optional[dict[str, Any]] = None, -) -> Any: - payload: dict[str, Any] = { - "input_ids": input_ids, - "sampling_params": { - "temperature": 0, - "max_new_tokens": GEN_LEN, - "no_stop_trim": True, - "skip_special_tokens": False, - }, - } - if session_params: - payload["session_params"] = session_params - 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 _run_all(base_url: str, tokenizer: Any) -> None: - """Fire all requests per turn simultaneously to create mixed batches.""" - timeout = aiohttp.ClientTimeout(total=300) - async with aiohttp.ClientSession(timeout=timeout) as http: - # Open all sessions - sids = [] - for s in range(NUM_SESSIONS): - async with http.post( - base_url + "/open_session", - json={"capacity_of_str_len": 50000, "streaming": True}, - ) as resp: - sids.append(await resp.json()) - - # For each turn, fire ALL streaming + non-streaming requests at once - for turn in range(NUM_TURNS): - tasks = [] - # Streaming requests for all sessions - for s in range(NUM_SESSIONS): - offset = (s * NUM_TURNS + turn) * 200 - text = f"Session {s} turn {turn}: {FILLER[offset : offset + 1500]}" - ids = tokenizer.encode(text) - tasks.append( - _async_generate( - base_url, - http, - ids, - session_params={"id": sids[s], "rid": None}, - ) - ) - - # Non-streaming requests interleaved - for ns in range(NUM_SESSIONS // 2): - text = f"Non-streaming {ns} turn {turn}: {FILLER[ns * 100 : ns * 100 + 500]}" - ids = tokenizer.encode(text) - tasks.append(_async_generate(base_url, http, ids)) - - # Fire all at once — creates mixed batch of streaming + non-streaming - await asyncio.gather(*tasks) - - # Close all sessions - for sid in sids: - async with http.post( - base_url + "/close_session", json={"session_id": sid} - ) as resp: - assert resp.status == 200 - - -class TestStreamingSessionLeak(CustomTestCase): - @classmethod - def setUpClass(cls) -> None: - 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", - ], - ) - - @classmethod - def tearDownClass(cls) -> None: - kill_process_tree(cls.process.pid) - - def test_streaming_session_no_leak(self) -> None: - """Concurrent multi-turn streaming sessions then idle health check.""" - from sglang.srt.utils.hf_transformers_utils import get_tokenizer - - tokenizer = get_tokenizer(self.model) - requests.post(self.base_url + "/flush_cache") - - asyncio.run(_run_all(self.base_url, tokenizer)) - - # Run a few non-streaming requests to flush state - for i in range(3): - ids = tokenizer.encode(f"Flush request {i}: final cleanup.") - requests.post( - self.base_url + "/generate", - json={ - "input_ids": ids, - "sampling_params": {"temperature": 0, "max_new_tokens": 4}, - }, - ) - - # Wait for server to go idle and run memory check - time.sleep(5) - health = requests.get(self.base_url + "/health") - self.assertEqual( - health.status_code, - 200, - "Server unhealthy after streaming session close — " - "likely a token memory leak from streaming session lifecycle.", - ) - - -if __name__ == "__main__": - unittest.main()