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

fix: streaming session race condition + some metrics (#21875)

Browse Source 
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
This commit is contained in:
ishandhanani
2026-04-12 18:05:23 -07:00
committed by GitHub
co-authored by Claude Opus 4.6 hnyls2002 Liangsheng Yin
parent 37fc47c645
commit c1ab68b45e
11 changed files with 966 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/sessions/test_streaming_session.py
+300
View File
@@ -10,8 +10,12 @@ Usage:
"""
import asyncio
import json
import os
import tempfile
import time
import unittest
from pathlib import Path
from typing import Any, Optional
import aiohttp
@@ -66,6 +70,21 @@ LEAK_FILLER = (
"We promptly judged antique ivory buckles for the next prize. "
) * 20
# ---------------------------------------------------------------------------
# Abort-heavy chunked prefill leak repro constants
# ---------------------------------------------------------------------------
ABORT_REPRO_CONTEXT_LEN = 512
ABORT_REPRO_PAGE_SIZE = 16
ABORT_REPRO_GEN_LEN = 8
ABORT_REPRO_SESSIONS = 4
ABORT_REPRO_WARMUP_TURNS = 2
ABORT_REPRO_ROUNDS = 8
ABORT_REPRO_STREAM_TOKENS = 150
ABORT_REPRO_ABORT_TOKENS = 320
ABORT_REPRO_NON_STREAMING_TOKENS = 96
ABORT_REPRO_CHUNKED_PREFILL_SIZE = 128
# ---------------------------------------------------------------------------
# Logprob leak helpers
@@ -206,6 +225,198 @@ async def _leak_run_all(base_url: str, tokenizer: Any) -> None:
assert resp.status == 200
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", ""
), text
return data
assert resp.status == 400, text
assert "maximum allowed 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
def _read_tail(path: str, num_lines: int = 120) -> str:
if not path or not os.path.exists(path):
return "<missing log>"
lines = Path(path).read_text(errors="replace").splitlines()
return "\n".join(lines[-num_lines:])
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()
# ===================================================================
# Test class
# ===================================================================
class TestStreamingSession(CustomTestCase):
@classmethod
def setUpClass(cls):
@@ -467,5 +678,94 @@ class TestStreamingSessionRetractMixedChunk(TestStreamingSession):
kill_process_tree(cls.process.pid)
class TestStreamingSessionAbortLeakRepro(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEFAULT_SMALL_MODEL_NAME_FOR_TEST
cls.base_url = DEFAULT_URL_FOR_TEST
cls.stdout = tempfile.NamedTemporaryFile(
prefix="streaming-session-abort-repro.",
suffix=".stdout.log",
delete=False,
mode="w+",
encoding="utf-8",
)
cls.stderr = tempfile.NamedTemporaryFile(
prefix="streaming-session-abort-repro.",
suffix=".stderr.log",
delete=False,
mode="w+",
encoding="utf-8",
)
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=[
"--enable-streaming-session",
"--chunked-prefill-size",
str(ABORT_REPRO_CHUNKED_PREFILL_SIZE),
"--context-length",
str(ABORT_REPRO_CONTEXT_LEN),
"--page-size",
str(ABORT_REPRO_PAGE_SIZE),
"--max-running-requests",
"32",
"--log-level",
"info",
],
return_stdout_stderr=(cls.stdout, cls.stderr),
)
cls.tokenizer = get_tokenizer(cls.model)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
for handle in (cls.stdout, cls.stderr):
path = handle.name
handle.close()
if os.path.exists(path):
os.remove(path)
def test_abort_heavy_chunked_prefill_does_not_leak(self) -> None:
requests.post(self.base_url + "/flush_cache")
asyncio.run(_abort_repro_run_all(self.base_url, self.tokenizer))
for i in range(3):
ids = self.tokenizer.encode(f"Post-session cleanup request {i}.")
response = requests.post(
self.base_url + "/generate",
json={
"input_ids": ids,
"sampling_params": {"temperature": 0, "max_new_tokens": 4},
},
timeout=30,
)
self.assertEqual(response.status_code, 200, response.text)
time.sleep(5)
self.assertIsNone(
self.process.poll(),
"Server crashed during abort-heavy streaming session repro.\n"
f"---- stderr tail ----\n{_read_tail(self.stderr.name)}",
)
health = requests.get(self.base_url + "/health", timeout=10)
self.assertEqual(
health.status_code,
200,
"Server unhealthy after abort-heavy streaming session cleanup.\n"
f"---- stderr tail ----\n{_read_tail(self.stderr.name)}",
)
stderr_tail = _read_tail(self.stderr.name)
self.assertNotIn(
"token_to_kv_pool_allocator memory leak detected",
stderr_tail,
stderr_tail,
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_streaming_session_unit.py
+358
View File
@@ -0,0 +1,358 @@
from types import SimpleNamespace
import torch
from sglang.srt.managers.schedule_batch import FINISH_ABORT
from sglang.srt.mem_cache.base_prefix_cache import MatchResult
from sglang.srt.mem_cache.common import release_kv_cache
from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache, SessionSlot
from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=8, suite="stage-a-test-cpu")
class _FakeAllocator:
def __init__(self):
self.freed = []
def free(self, free_index: torch.Tensor):
self.freed.append(free_index.clone())
class _FakeInnerCache:
def __init__(self, req_to_token_pool, allocator, page_size, match_results=None):
self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = allocator
self.page_size = page_size
self.match_results = list(match_results or [])
self.dec_lock_ref_calls = []
def cache_finished_req(self, *args, **kwargs):
raise AssertionError("Streaming requests should not delegate to inner cache")
def match_prefix(self, *args, **kwargs):
if not self.match_results:
raise AssertionError("Unexpected match_prefix call")
return self.match_results.pop(0)
def dec_lock_ref(self, node, *args, **kwargs):
self.dec_lock_ref_calls.append(node)
def supports_mamba(self):
return False
def sanity_check(self):
return None
class _FakeReq:
def __init__(
self, session_id: str, req_pool_idx: int, committed: int, allocated: int
):
self.session = SimpleNamespace(session_id=session_id, streaming=True)
self.req_pool_idx = req_pool_idx
self.kv_committed_len = committed
self.kv_allocated_len = allocated
self.kv_committed_freed = False
self.kv_overallocated_freed = False
self.origin_input_ids = list(range(committed))
self.output_ids = []
self.extra_key = None
self.swa_evicted_seqlen = 0
self.last_node = None
self.cache_protected_len = 0
self.swa_uuid_for_lock = None
self.mamba_pool_idx = None
self.mamba_ping_pong_track_buffer = None
self.mamba_next_track_idx = None
self.mamba_last_track_seqlen = None
self.mamba_branching_seqlen = None
self.pop_overallocated_calls = 0
self.to_finish = None
self.finished_reason = None
def pop_committed_kv_cache(self):
assert not self.kv_committed_freed
self.kv_committed_freed = True
return self.kv_committed_len
def pop_overallocated_kv_cache(self):
assert not self.kv_overallocated_freed
self.pop_overallocated_calls += 1
self.kv_overallocated_freed = True
return self.kv_committed_len, self.kv_allocated_len
def test_streaming_release_kv_cache_trims_overallocated_tail(monkeypatch):
page_size = 16
req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128)
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
allocator = _FakeAllocator()
tree_cache = SessionAwareCache(
_FakeInnerCache(req_to_token_pool, allocator, page_size)
)
req = _FakeReq("session-a", req_pool_idx=0, committed=17, allocated=40)
monkeypatch.setattr(
"sglang.srt.mem_cache.common.get_global_server_args",
lambda: SimpleNamespace(page_size=page_size, speculative_algorithm="eagle"),
)
release_kv_cache(req, tree_cache)
slot = tree_cache.slots["session-a"]
assert req.pop_overallocated_calls == 1
assert req.kv_committed_freed is True
assert req.kv_overallocated_freed is True
assert req.req_pool_idx is None
assert slot.kv_committed_len == 17
assert slot.kv_allocated_len == 17
assert len(allocator.freed) == 1
assert allocator.freed[0].tolist() == list(range(32, 40))
def test_release_session_recomputes_current_tree_owned_prefix():
page_size = 16
req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128)
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
allocator = _FakeAllocator()
full_match = MatchResult(
device_indices=torch.tensor(list(range(16)) + list(range(64, 96))),
last_device_node="stale-expanded",
last_host_node="stale-expanded",
)
protected_match = MatchResult(
device_indices=torch.tensor(list(range(16))),
last_device_node="current-protected",
last_host_node="current-protected",
)
inner = _FakeInnerCache(
req_to_token_pool,
allocator,
page_size,
match_results=[full_match, protected_match],
)
tree_cache = SessionAwareCache(inner)
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv_committed_len=48,
kv_allocated_len=48,
last_node="outdated-node",
cache_protected_len=32,
)
req = _FakeReq("session-a", req_pool_idx=0, committed=48, allocated=48)
tree_cache.release_session("session-a", req)
assert inner.dec_lock_ref_calls == ["current-protected"]
assert req_to_token_pool.free_slots == [0]
assert len(allocator.freed) == 1
assert allocator.freed[0].tolist() == list(range(16, 48))
def test_release_session_never_grows_tree_owned_prefix():
page_size = 16
req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128)
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
allocator = _FakeAllocator()
overmatched = MatchResult(
device_indices=torch.tensor(list(range(48))),
last_device_node="overmatched-node",
last_host_node="overmatched-node",
)
capped_match = MatchResult(
device_indices=torch.tensor(list(range(16))),
last_device_node="original-lock-node",
last_host_node="original-lock-node",
)
inner = _FakeInnerCache(
req_to_token_pool,
allocator,
page_size,
match_results=[overmatched, capped_match],
)
tree_cache = SessionAwareCache(inner)
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv_committed_len=48,
kv_allocated_len=48,
last_node="outdated-node",
cache_protected_len=16,
)
req = _FakeReq("session-a", req_pool_idx=0, committed=48, allocated=48)
tree_cache.release_session("session-a", req)
assert inner.dec_lock_ref_calls == ["original-lock-node"]
assert req_to_token_pool.free_slots == [0]
assert len(allocator.freed) == 1
assert allocator.freed[0].tolist() == list(range(16, 48))
def test_match_prefix_abort_does_not_restore_live_session_slot():
req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
allocator = _FakeAllocator()
inner = _FakeInnerCache(
req_to_token_pool,
allocator,
page_size=16,
match_results=[
MatchResult(
device_indices=torch.tensor([], dtype=torch.int64),
last_device_node=None,
last_host_node=None,
)
],
)
tree_cache = SessionAwareCache(inner)
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv_committed_len=48,
kv_allocated_len=48,
cache_protected_len=16,
)
req = _FakeReq("session-a", req_pool_idx=1, committed=1, allocated=1)
req.to_finish = FINISH_ABORT("too long")
result = tree_cache.match_prefix(
SimpleNamespace(
req=req,
key=SimpleNamespace(token_ids=list(range(64))),
)
)
slot = tree_cache.slots["session-a"]
assert req.req_pool_idx == 1
assert req.kv_committed_len == 1
assert req.kv_allocated_len == 1
assert slot.req_pool_idx == 0
assert slot.kv_committed_len == 48
assert slot.kv_allocated_len == 48
assert len(result.device_indices) == 0
def test_aborted_streaming_turn_preserves_slot_and_accounting(monkeypatch):
page_size = 16
req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
allocator = _FakeAllocator()
tree_cache = SessionAwareCache(
_FakeInnerCache(req_to_token_pool, allocator, page_size)
)
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv_committed_len=48,
kv_allocated_len=48,
cache_protected_len=16,
swa_evicted_seqlen=8,
last_node="lock-node",
)
req = _FakeReq("session-a", req_pool_idx=1, committed=5, allocated=23)
req.finished_reason = FINISH_ABORT("too long")
monkeypatch.setattr(
"sglang.srt.mem_cache.common.get_global_server_args",
lambda: SimpleNamespace(page_size=page_size, speculative_algorithm="eagle"),
)
release_kv_cache(req, tree_cache)
slot = tree_cache.slots["session-a"]
assert slot.req_pool_idx == 0
assert slot.kv_committed_len == 48
assert slot.kv_allocated_len == 48
assert req.kv_committed_freed is True
assert req.kv_overallocated_freed is True
assert req.req_pool_idx is None
assert req.pop_overallocated_calls == 1
assert tree_cache.session_held_tokens() == 32
assert tree_cache.session_held_full_tokens() == 32
assert tree_cache.session_held_swa_tokens() == 32
assert tree_cache.session_held_req_count() == 1
assert req_to_token_pool.free_slots == [1]
assert len(allocator.freed) == 1
assert allocator.freed[0].tolist() == list(range(128, 151))
tree_cache.release_session("session-a")
assert tree_cache.session_held_tokens() == 0
assert tree_cache.session_held_swa_tokens() == 0
assert tree_cache.session_held_req_count() == 0
assert req_to_token_pool.free_slots == [1, 0]
assert len(allocator.freed) == 2
assert allocator.freed[1].tolist() == list(range(16, 48))
def test_session_shrink_frees_orphaned_tail():
"""When a session's KV shrinks (client retried with shorter prompt),
the orphaned tail pages must be freed before save_from_req overwrites
the slot."""
page_size = 16
pool_size = 256
req_to_token = torch.arange(pool_size, dtype=torch.int32).reshape(1, pool_size)
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
allocator = _FakeAllocator()
inner = _FakeInnerCache(req_to_token_pool, allocator, page_size)
tree_cache = SessionAwareCache(inner)
# Session slot has 128 tokens committed
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv_committed_len=128,
kv_allocated_len=128,
last_node="lock-node",
cache_protected_len=16,
)
# New request finished with only 48 tokens (client truncated)
req = _FakeReq("session-a", req_pool_idx=0, committed=48, allocated=48)
tree_cache.cache_finished_req(req)
slot = tree_cache.slots["session-a"]
# Slot should now reflect the shrunk state
assert slot.kv_committed_len == 48
assert slot.kv_allocated_len == 48
# The tail [48:128] should have been freed (page-aligned: [48:128])
assert len(allocator.freed) == 1
assert allocator.freed[0].tolist() == list(range(48, 128))
def test_session_shrink_page_aligns_free_start():
"""The shrink free should page-align the start to avoid freeing
tokens that are still part of the new committed prefix."""
page_size = 16
pool_size = 256
req_to_token = torch.arange(pool_size, dtype=torch.int32).reshape(1, pool_size)
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
allocator = _FakeAllocator()
inner = _FakeInnerCache(req_to_token_pool, allocator, page_size)
tree_cache = SessionAwareCache(inner)
# Session slot has 128 tokens
tree_cache.slots["session-a"] = SessionSlot(
req_pool_idx=0,
kv_committed_len=128,
kv_allocated_len=128,
last_node="lock-node",
cache_protected_len=16,
)
# New request committed 50 tokens (not page-aligned)
req = _FakeReq("session-a", req_pool_idx=0, committed=50, allocated=50)
tree_cache.cache_finished_req(req)
slot = tree_cache.slots["session-a"]
assert slot.kv_committed_len == 50
# Free start should be ceil_align(50, 16) = 64, not 50
# So freed range is [64:128]
assert len(allocator.freed) == 1
assert allocator.freed[0].tolist() == list(range(64, 128))