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
python/sglang/srt/managers/session_controller.py
+29 -8
View File
@@ -94,6 +94,7 @@ class Session:
self.last_active_time: float = time.monotonic()
self.req_nodes: Dict[str, SessionReqNode] = {}
self.close_on_finish: bool = False
self._inflight: bool = False
def is_timed_out(self) -> bool:
if self.timeout is None:
@@ -117,7 +118,10 @@ class Session:
abort_message = ""
if self.streaming:
# Streaming sessions: only simple appends allowed; reject otherwise.
if session_params.replace:
if self._inflight:
abort = True
abort_message = "Streaming session already has an active request."
elif session_params.replace:
abort = True
abort_message = "Streaming sessions do not support replace."
elif session_params.drop_previous_output:
@@ -130,7 +134,9 @@ class Session:
abort_message = "Streaming sessions do not support offset."
elif self.req_nodes:
assert len(self.req_nodes) == 1
_, last_req_node = self.req_nodes.popitem()
# Peek (don't pop) the single req_node. req_nodes is updated
# only in finish_req after the request completes successfully.
[last_req_node] = self.req_nodes.values()
last_req = last_req_node.req
elif session_params.replace:
if session_params.rid is None:
@@ -240,15 +246,27 @@ class Session:
if abort:
new_req.set_finish_with_abort(abort_message)
elif self.streaming:
if last_req is not None:
last_req.session = None
self.req_nodes[req.rid] = SessionReqNode(new_req)
# req_nodes is NOT updated here — finish_req() handles it.
self._inflight = True
else:
new_req_node = SessionReqNode(new_req, last_req_node)
self.req_nodes[req.rid] = new_req_node
return new_req
def finish_req(self, req):
"""Update req_nodes after a streaming request finishes successfully."""
self._inflight = False
if self.req_nodes:
[prev_node] = self.req_nodes.values()
prev_node.req.session = None
self.req_nodes.clear()
self.req_nodes[req.rid] = SessionReqNode(req)
def abort_req(self):
"""Clear inflight flag on abort (req_nodes stays unchanged)."""
self._inflight = False
class SessionController:
def __init__(self, tree_cache: BasePrefixCache):
@@ -293,9 +311,12 @@ class SessionController:
session = self.sessions[session_id]
req = None
has_unfinished_request = False
if session.streaming and session.req_nodes:
if session.streaming and session._inflight:
has_unfinished_request = True
elif session.streaming and session.req_nodes:
assert len(session.req_nodes) == 1
req = next(iter(session.req_nodes.values())).req
[last_node] = session.req_nodes.values()
req = last_node.req
if not req.finished():
has_unfinished_request = True
@@ -362,7 +383,7 @@ class SessionController:
self._close(sid)
@staticmethod
def _all_requests_finished(session: "Session") -> bool:
def _all_requests_finished(session: Session) -> bool:
if not session.req_nodes:
return True
return all(node.req.finished() for node in session.req_nodes.values())
python/sglang/srt/mem_cache/session_aware_cache.py
+42 -45
View File
@@ -183,14 +183,12 @@ class SessionAwareCache(BasePrefixCache):
if slot is None or slot.req_pool_idx is None:
return self.inner.match_prefix(params)
# If the request is destined for abort (e.g. input too long),
# do NOT restore the slot's KV state. set_finish_with_abort
# truncates origin_input_ids to [0], so alloc_for_extend would
# overwrite the slot's req_to_token row with a 1-token prefix,
# destroying the session's accumulated KV mapping. By skipping
# restore, the request gets a fresh pool slot from alloc_for_extend
# and the session slot remains untouched.
# Pre-aborted req (scheduler-level abort, e.g. input too long):
# detach from session so cache_finished_req treats it as a normal
# req. The slot stays intact for the next request.
if req.to_finish is not None:
req.session.abort_req()
req.session = None
return self.inner.match_prefix(params)
slot.restore_to_req(req)
@@ -220,52 +218,51 @@ class SessionAwareCache(BasePrefixCache):
slot = self.slots.get(session_id)
is_first = slot is None
# When an aborted streaming-session request was scheduled (e.g.
# input too long), match_prefix skipped restore_to_req so the
# request got a fresh pool slot from alloc_for_extend. Don't
# overwrite the session slot -- free the transient KV and pool slot.
if not is_first and isinstance(req.finished_reason, FINISH_ABORT):
if req.req_pool_idx is not None:
# Free all KV pages allocated for this aborted request.
end = req.kv_allocated_len
if end > 0:
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :end
]
self.token_to_kv_pool_allocator.free(kv_indices)
self.req_to_token_pool.free_slots.append(req.req_pool_idx)
req.req_pool_idx = None
# Mid-processing abort only. Pre-aborted reqs have session=None
# (set in match_prefix) and never reach here.
# Nuke all KV via release_session, delete slot. Token IDs stay
# in req_nodes (finish_req was never called -> last successful
# req). Next request re-prefills from scratch.
if isinstance(req.finished_reason, FINISH_ABORT):
if slot is None:
# First-request mid-processing abort: create ephemeral
# slot from req state so release_session handles cleanup.
# Include last_node/cache_protected_len from the req so
# release_session calls dec_lock_ref on the tree lock.
slot = SessionSlot(
req_pool_idx=req.req_pool_idx,
kv_allocated_len=req.kv_allocated_len,
last_node=req.last_node,
cache_protected_len=req.cache_protected_len,
swa_uuid_for_lock=req.swa_uuid_for_lock,
)
self.slots[session_id] = slot
slot.kv_allocated_len = max(slot.kv_allocated_len, req.kv_allocated_len)
self.release_session(session_id)
req.req_pool_idx = None
req.session.abort_req()
self._mark_kv_freed(req)
return
if is_first:
slot = SessionSlot()
self.slots[session_id] = slot
# If the session's KV is shrinking (e.g. client sent a shorter
# prompt after an abort), free the orphaned tail pages before
# save_from_req overwrites the slot's committed length.
# Never free tree-protected tokens — those are managed by the tree.
if (
not is_first
and slot.is_holding_kv
and req.kv_committed_len < slot.kv_committed_len
):
old_end = slot.kv_allocated_len
new_end = req.kv_committed_len
if self.page_size > 1:
new_end = ceil_align(new_end, self.page_size)
new_end = max(new_end, slot.cache_protected_len)
if new_end < old_end:
kv_indices = self.req_to_token_pool.req_to_token[
slot.req_pool_idx, new_end:old_end
]
self.token_to_kv_pool_allocator.free(kv_indices)
slot.cache_protected_len = min(
slot.cache_protected_len, req.kv_committed_len
)
slot.save_from_req(req, is_first=is_first)
# Update req_nodes to this successfully finished request.
req.session.finish_req(req)
self._mark_kv_freed(req)
@staticmethod
def _mark_kv_freed(req: Req):
"""Set bookkeeping flags so busy check skips this finished req."""
if not req.kv_committed_freed:
req.pop_committed_kv_cache()
if not req.kv_overallocated_freed:
req.pop_overallocated_kv_cache()
def cache_unfinished_req(self, req: Req, **kwargs):
if _is_streaming(req):
# in chunked_prefill for streaming, we skip the stash path which triggers radix.