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

Refactor streaming session abort handling (#22790)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-04-15 00:13:05 -07:00
committed by GitHub
parent 45a83ffbe3
commit aa78564e1a
4 changed files with 423 additions and 176 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/sessions/test_streaming_session.py
+275 -7
View File
@@ -73,13 +73,13 @@ LEAK_FILLER = (
ABORT_REPRO_CONTEXT_LEN = 512
ABORT_REPRO_PAGE_SIZE = 16
ABORT_REPRO_GEN_LEN = 8
ABORT_REPRO_GEN_LEN = 4
ABORT_REPRO_SESSIONS = 4
ABORT_REPRO_WARMUP_TURNS = 2
ABORT_REPRO_WARMUP_TURNS = 1
ABORT_REPRO_ROUNDS = 8
ABORT_REPRO_STREAM_TOKENS = 150
ABORT_REPRO_ABORT_TOKENS = 320
ABORT_REPRO_NON_STREAMING_TOKENS = 96
ABORT_REPRO_STREAM_TOKENS = 16
ABORT_REPRO_ABORT_TOKENS = 600
ABORT_REPRO_NON_STREAMING_TOKENS = 16
ABORT_REPRO_CHUNKED_PREFILL_SIZE = 128
@@ -265,10 +265,10 @@ async def _abort_repro_generate(
assert finish_reason.get("type") == "abort", text
assert "maximum allowed length" in finish_reason.get(
"message", ""
), text
) or "context length" in finish_reason.get("message", ""), text
return data
assert resp.status == 400, text
assert "maximum allowed length" in text, text
assert "maximum allowed length" in text or "context length" in text, text
return None
assert resp.status == 200, text
@@ -584,6 +584,274 @@ class TestStreamingSession(CustomTestCase):
"likely a token memory leak from streaming session lifecycle.",
)
def test_nth_mid_abort_recovery(self) -> None:
"""Abort a running streaming session request (nth turn) via the
abort API. Session rolls back to last successful turn."""
requests.post(self.base_url + "/flush_cache")
resp = requests.post(
self.base_url + "/open_session",
json={"capacity_of_str_len": 50000, "streaming": True},
)
self.assertEqual(resp.status_code, 200)
session_id = resp.json()
try:
# Turn 1: normal generate to create slot.
ids_1 = self.tokenizer.encode("Tell me a very long story about a wizard.")
resp_1 = requests.post(
self.base_url + "/generate",
json={
"input_ids": ids_1,
"sampling_params": {"temperature": 0, "max_new_tokens": 16},
"session_params": {"id": session_id, "rid": None},
},
timeout=30,
)
self.assertEqual(resp_1.status_code, 200, resp_1.text)
data_1 = resp_1.json()
turn_1_total = (
data_1["meta_info"]["prompt_tokens"]
+ data_1["meta_info"]["completion_tokens"]
)
# Turn 2: long generate, then abort mid-decode.
ids_2 = self.tokenizer.encode(" Continue the story in great detail.")
import threading
result = [None]
def do_generate():
r = requests.post(
self.base_url + "/generate",
json={
"input_ids": ids_2,
"sampling_params": {
"temperature": 0,
"max_new_tokens": 100000,
},
"session_params": {"id": session_id, "rid": None},
},
timeout=60,
)
result[0] = r
t = threading.Thread(target=do_generate)
t.start()
time.sleep(0.5)
abort_resp = requests.post(
self.base_url + "/abort_request",
json={"rid": "", "abort_all": True},
timeout=10,
)
self.assertEqual(abort_resp.status_code, 200, abort_resp.text)
t.join(timeout=30)
self.assertIsNotNone(result[0], "Turn 2 should have returned")
data_2 = result[0].json()
self.assertEqual(
data_2["meta_info"]["finish_reason"]["type"],
"abort",
"Turn 2 should be aborted, not finished normally",
)
# Turn 3: recovery. Rolls back to turn 1.
ids_3 = self.tokenizer.encode(" What happens next?")
for attempt in range(20):
resp_3 = requests.post(
self.base_url + "/generate",
json={
"input_ids": ids_3,
"sampling_params": {"temperature": 0, "max_new_tokens": 8},
"session_params": {"id": session_id, "rid": None},
},
timeout=30,
)
if resp_3.status_code == 200:
break
time.sleep(0.5)
self.assertEqual(resp_3.status_code, 200, resp_3.text)
data_3 = resp_3.json()
# prompt_tokens = turn_1_total + append (BOS stripped).
bos = 1 if ids_3[0] == self.tokenizer.bos_token_id else 0
expected_prompt_3 = turn_1_total + len(ids_3) - bos
self.assertEqual(
data_3["meta_info"]["prompt_tokens"],
expected_prompt_3,
"prompt_tokens must equal turn_1_total + append (no stale abort context)",
)
finally:
requests.post(
self.base_url + "/close_session",
json={"session_id": session_id},
)
health = requests.get(self.base_url + "/health", timeout=10)
self.assertEqual(health.status_code, 200)
def test_first_mid_abort_recovery(self) -> None:
"""Abort the very first request on a streaming session mid-decode.
No slot exists yet (ephemeral slot created and nuked).
Verify the session is still usable afterward."""
requests.post(self.base_url + "/flush_cache")
resp = requests.post(
self.base_url + "/open_session",
json={"capacity_of_str_len": 50000, "streaming": True},
)
self.assertEqual(resp.status_code, 200)
session_id = resp.json()
try:
ids_1 = self.tokenizer.encode("Tell me a very long story about a wizard.")
import threading
result = [None]
def do_generate():
r = requests.post(
self.base_url + "/generate",
json={
"input_ids": ids_1,
"sampling_params": {
"temperature": 0,
"max_new_tokens": 100000,
},
"session_params": {"id": session_id, "rid": None},
},
timeout=60,
)
result[0] = r
t = threading.Thread(target=do_generate)
t.start()
time.sleep(0.5)
abort_resp = requests.post(
self.base_url + "/abort_request",
json={"rid": "", "abort_all": True},
timeout=10,
)
self.assertEqual(abort_resp.status_code, 200, abort_resp.text)
t.join(timeout=30)
self.assertIsNotNone(result[0], "Turn 1 should have returned")
data_1 = result[0].json()
self.assertEqual(
data_1["meta_info"]["finish_reason"]["type"],
"abort",
"Turn 1 should be aborted, not finished normally",
)
# Turn 2: recovery. No inherited context (req_nodes empty).
ids_2 = self.tokenizer.encode("Tell me a short joke.")
for attempt in range(20):
resp_2 = requests.post(
self.base_url + "/generate",
json={
"input_ids": ids_2,
"sampling_params": {"temperature": 0, "max_new_tokens": 8},
"session_params": {"id": session_id, "rid": None},
},
timeout=30,
)
if resp_2.status_code == 200:
break
time.sleep(0.5)
self.assertEqual(resp_2.status_code, 200, resp_2.text)
data_2 = resp_2.json()
self.assertEqual(
data_2["meta_info"]["prompt_tokens"],
len(ids_2),
"prompt_tokens must equal turn 2 input only (no inherited context)",
)
finally:
requests.post(
self.base_url + "/close_session",
json={"session_id": session_id},
)
health = requests.get(self.base_url + "/health", timeout=10)
self.assertEqual(health.status_code, 200)
def test_preabort_recovery(self) -> None:
"""Pre-aborted request (unsupported offset) does not corrupt session.
The slot is preserved, and the next turn inherits correctly."""
requests.post(self.base_url + "/flush_cache")
resp = requests.post(
self.base_url + "/open_session",
json={"capacity_of_str_len": 50000, "streaming": True},
)
self.assertEqual(resp.status_code, 200)
session_id = resp.json()
try:
# Turn 1: normal generate to create slot.
ids_1 = self.tokenizer.encode("Tell me a very long story about a wizard.")
resp_1 = requests.post(
self.base_url + "/generate",
json={
"input_ids": ids_1,
"sampling_params": {"temperature": 0, "max_new_tokens": 16},
"session_params": {"id": session_id, "rid": None},
},
timeout=30,
)
self.assertEqual(resp_1.status_code, 200, resp_1.text)
data_1 = resp_1.json()
turn_1_total = (
data_1["meta_info"]["prompt_tokens"]
+ data_1["meta_info"]["completion_tokens"]
)
# Turn 2: pre-aborted via unsupported offset parameter.
ids_2 = self.tokenizer.encode(" This should be rejected.")
resp_2 = requests.post(
self.base_url + "/generate",
json={
"input_ids": ids_2,
"sampling_params": {"temperature": 0, "max_new_tokens": 8},
"session_params": {
"id": session_id,
"rid": None,
"offset": 1,
},
},
timeout=30,
)
self.assertIn(resp_2.status_code, (200, 400), resp_2.text)
# Turn 3: normal append. Slot should be intact from turn 1.
ids_3 = self.tokenizer.encode(" What happens next?")
resp_3 = requests.post(
self.base_url + "/generate",
json={
"input_ids": ids_3,
"sampling_params": {"temperature": 0, "max_new_tokens": 8},
"session_params": {"id": session_id, "rid": None},
},
timeout=30,
)
self.assertEqual(resp_3.status_code, 200, resp_3.text)
data_3 = resp_3.json()
bos = 1 if ids_3[0] == self.tokenizer.bos_token_id else 0
expected_prompt_3 = turn_1_total + len(ids_3) - bos
self.assertEqual(
data_3["meta_info"]["prompt_tokens"],
expected_prompt_3,
"prompt_tokens must equal turn_1_total + append (slot preserved)",
)
finally:
requests.post(
self.base_url + "/close_session",
json={"session_id": session_id},
)
health = requests.get(self.base_url + "/health", timeout=10)
self.assertEqual(health.status_code, 200)
class TestStreamingSessionMixedChunk(TestStreamingSession):
"""Streaming session with --enable-mixed-chunk."""
test/registered/unit/mem_cache/test_streaming_session_unit.py
+77 -116
View File
@@ -49,7 +49,13 @@ 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.session = SimpleNamespace(
session_id=session_id,
streaming=True,
finish_req=lambda req: None,
abort_req=lambda: None,
_inflight=False,
)
self.req_pool_idx = req_pool_idx
self.kv_committed_len = committed
self.kv_allocated_len = allocated
@@ -83,7 +89,10 @@ class _FakeReq:
return self.kv_committed_len, self.kv_allocated_len
def test_streaming_release_kv_cache_trims_overallocated_tail(monkeypatch):
def test_streaming_release_kv_cache_defers_tail_free(monkeypatch):
"""Spec tail is NOT trimmed in cache_finished_req; it is deferred to
match_prefix's orphan tail free on the next turn. cache_finished_req
only sets bookkeeping flags and saves the slot as-is."""
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=[])
@@ -101,17 +110,18 @@ def test_streaming_release_kv_cache_trims_overallocated_tail(monkeypatch):
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
# Slot keeps the full allocation — tail free is deferred to match_prefix.
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))
assert slot.kv_allocated_len == 40
assert len(allocator.freed) == 0
def test_match_prefix_abort_does_not_restore_live_session_slot():
def test_preabort_detaches_session_and_preserves_slot():
"""Pre-aborted req (to_finish set before match_prefix) is detached from
the session: session=None, abort_req() called. Slot stays intact."""
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()
@@ -145,132 +155,83 @@ def test_match_prefix_abort_does_not_restore_live_session_slot():
)
)
# Req detached from session.
assert req.session is None
# Slot untouched.
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
def test_first_mid_abort_nukes_ephemeral_slot():
"""First-request mid-processing abort: no slot exists yet, ephemeral
slot is created from req state and nuked via release_session."""
page_size = 1
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()
inner = _FakeInnerCache(req_to_token_pool, allocator, page_size)
tree_cache = SessionAwareCache(inner)
# No slot exists yet (first request).
req = _FakeReq("session-a", req_pool_idx=0, committed=0, allocated=20)
req.finished_reason = FINISH_ABORT("input too long")
tree_cache.cache_finished_req(req)
# Slot must NOT be created.
assert "session-a" not in tree_cache.slots
# Transient pool slot freed.
assert req.req_pool_idx is None
assert req_to_token_pool.free_slots == [0]
assert len(allocator.freed) == 1
assert allocator.freed[0].tolist() == list(range(20))
# Bookkeeping flags set.
assert req.kv_committed_freed is True
assert req.kv_overallocated_freed is True
def test_nth_mid_abort_nukes_session_slot():
"""Nth-request mid-processing abort: slot exists, restore_to_req ran.
ALL KV is wiped (release_session). Slot is deleted. Token IDs stay
in req_nodes for next turn's re-prefill."""
page_size = 1
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)
)
inner = _FakeInnerCache(req_to_token_pool, allocator, page_size)
tree_cache = SessionAwareCache(inner)
# Session already has a slot from a previous turn.
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",
kv_committed_len=50,
kv_allocated_len=50,
last_node=None,
cache_protected_len=0,
)
req = _FakeReq("session-a", req_pool_idx=1, committed=5, allocated=23)
req.finished_reason = FINISH_ABORT("too long")
# Mid-processing abort: req has the SESSION slot's pool_idx (restore_to_req ran).
req = _FakeReq("session-a", req_pool_idx=0, committed=60, allocated=65)
req.finished_reason = FINISH_ABORT("client disconnected")
monkeypatch.setattr(
"sglang.srt.mem_cache.common.get_global_server_args",
lambda: SimpleNamespace(page_size=page_size, speculative_algorithm="eagle"),
)
tree_cache.cache_finished_req(req)
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
# Slot wiped — deleted from slots dict.
assert "session-a" not in tree_cache.slots
# All KV freed: [0, 65) from release_session (slot extended to req's allocated).
assert len(allocator.freed) == 1
assert allocator.freed[0].tolist() == list(range(65))
# Pool slot returned.
assert req_to_token_pool.free_slots == [0]
assert req.req_pool_idx is None
# Bookkeeping flags set.
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))
# Shrink tests removed: streaming sessions are append-only after the
# rollback fix in session_controller (rollback_aborted_req). The shrink
# code path in cache_finished_req no longer exists.