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

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>

python/sglang/srt/managers/scheduler.py

@@ -1848,8 +1848,11 @@ class Scheduler(
                     self.stream_output([req], req.return_logprob)
                     return
 
-        elif session_id in self.session_controller:
-            # Session exists: create request from session
+        elif (
+            session_id in self.session_controller
+            and not self.session_controller.get(session_id).close_on_finish
+        ):
+            # Session exists and is not closing: create request from session
             session = self.session_controller.get(session_id)
             req = session.create_req(
                 recv_req,
@@ -1866,7 +1869,13 @@ class Scheduler(
                 return
 
         else:
-            # Session ID provided but session not found
+            # Session not found, or session is closing
+            if session_id in self.session_controller:
+                error_msg = (
+                    f"Invalid request: close was requested for session {session_id}"
+                )
+            else:
+                error_msg = f"Invalid request: session id {session_id} does not exist"
             req = Req(
                 recv_req.rid,
                 recv_req.input_text,
@@ -1875,9 +1884,7 @@ class Scheduler(
                 vocab_size=self.model_config.vocab_size,
             )
             req.tokenizer = self.tokenizer
-            req.set_finish_with_abort(
-                f"Invalid request: session id {session_id} does not exist"
-            )
+            req.set_finish_with_abort(error_msg)
             self.init_req_max_new_tokens(req)
             self._add_request_to_queue(req)
             return
@@ -3461,7 +3468,10 @@ class Scheduler(
         return ExpertDistributionReqOutput()
 
     def open_session(self, recv_req: OpenSessionReqInput):
-        return self.session_controller.open(recv_req)
+        output = self.session_controller.open(recv_req)
+        if self.pp_rank == 0 and self.tp_rank == 0 and self.attn_cp_rank == 0:
+            return output
+        return None
 
     def close_session(self, recv_req: CloseSessionReqInput):
         self.session_controller.close(recv_req)

python/sglang/srt/managers/scheduler_runtime_checker_mixin.py

@@ -117,6 +117,13 @@ class PoolStats:
 
 
 class SchedulerRuntimeCheckerMixin:
+    def _alive_streaming_session_count(self: Scheduler) -> int:
+        return sum(
+            1
+            for session in self.session_controller.sessions.values()
+            if session.streaming
+        )
+
     def _session_held_tokens(self: Scheduler) -> int:
         if isinstance(self.tree_cache, SessionAwareCache):
             return self.tree_cache.session_held_tokens()
@@ -451,6 +458,8 @@ class SchedulerRuntimeCheckerMixin:
             return
 
         self.get_pool_stats().update_scheduler_stats(self.stats)
+        self.stats.num_streaming_sessions = self._alive_streaming_session_count()
+        self.stats.streaming_session_held_tokens = self._session_held_tokens()
 
         priority_enabled = self.enable_priority_scheduling
         self.stats.num_running_reqs = QueueCount.from_reqs(

python/sglang/srt/managers/session_controller.py

@@ -25,6 +25,7 @@ from sglang.srt.managers.io_struct import (
 )
 from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req
 from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache
+from sglang.srt.utils.common import log_info_on_rank0
 
 if TYPE_CHECKING:
     from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
@@ -92,6 +93,7 @@ class Session:
         self.timeout = timeout
         self.last_active_time: float = time.monotonic()
         self.req_nodes: Dict[str, SessionReqNode] = {}
+        self.close_on_finish: bool = False
 
     def is_timed_out(self) -> bool:
         if self.timeout is None:
@@ -275,6 +277,9 @@ class SessionController:
                 streaming=bool(recv_req.streaming),
                 timeout=recv_req.timeout,
             )
+            log_info_on_rank0(
+                logger, f"Session opened: {session_id} (active={len(self.sessions)})"
+            )
             return OpenSessionReqOutput(session_id, True)
 
     def close(self, recv_req: CloseSessionReqInput):
@@ -286,11 +291,31 @@ class SessionController:
 
     def _close(self, session_id: str):
         session = self.sessions[session_id]
+        req = None
+        has_unfinished_request = False
         if session.streaming and session.req_nodes:
             assert len(session.req_nodes) == 1
             req = next(iter(session.req_nodes.values())).req
             if not req.finished():
-                req.session = None
+                has_unfinished_request = True
+
+        if has_unfinished_request:
+            # An in-flight request is still decoding on this session's KV
+            # memory. Freeing now would corrupt the scheduler. Mark the
+            # session for deferred cleanup: the request keeps its session
+            # reference so cache_finished_req takes the streaming path,
+            # and we schedule release_session for after it completes.
+            session.close_on_finish = True
+            logger.info(
+                "Deferring session close for %s (unfinished request)",
+                session_id,
+            )
+            return
+
+        # No active request -- safe to release immediately.
+        if session.streaming and session.req_nodes:
+            req = next(iter(session.req_nodes.values())).req
+            req.session = None
 
         # Release multimodal features held by session requests.
         # Session reqs skip the normal mm cleanup path (scheduler and
@@ -304,20 +329,46 @@ class SessionController:
             node.req.multimodal_inputs = None
 
         if isinstance(self.tree_cache, SessionAwareCache):
-            self.tree_cache.release_session(session_id)
+            self.tree_cache.release_session(
+                session_id, req if session.streaming else None
+            )
         del self.sessions[session_id]
+        log_info_on_rank0(
+            logger, f"Session closed: {session_id} (active={len(self.sessions)})"
+        )
 
     def maybe_reap(self, now: float, interval: float = 1.0):
         # reap sessions every second
         if now - self._last_reap_time > interval:
             self._last_reap_time = now
+
+            # Finish deferred closes for sessions whose requests completed.
+            pending = [
+                sid
+                for sid, session in self.sessions.items()
+                if session.close_on_finish and self._all_requests_finished(session)
+            ]
+            for sid in pending:
+                log_info_on_rank0(
+                    logger, f"Deferred close ready for session {sid}, releasing."
+                )
+                # Reset close_on_finish so _close proceeds with the release.
+                self.sessions[sid].close_on_finish = False
+                self._close(sid)
+
             timed_out = [
                 sid for sid, session in self.sessions.items() if session.is_timed_out()
             ]
             for sid in timed_out:
-                logger.info(f"Session {sid} timed out, closing.")
+                log_info_on_rank0(logger, f"Session {sid} timed out, closing.")
                 self._close(sid)
 
+    @staticmethod
+    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())
+
     @staticmethod
     def adjust_mm_offsets(recv_req: TokenizedGenerateReqInput, req: Req, image_inputs):
         # For session requests, adjust mm_inputs offsets by the prefix length.

python/sglang/srt/managers/tokenizer_communicator_mixin.py

@@ -1090,12 +1090,14 @@ class TokenizerCommunicatorMixin:
         elif obj.session_id in self.session_futures:
             return None
 
+        future = asyncio.Future()
+        self.session_futures[obj.session_id] = future
         self.send_to_scheduler.send_pyobj(obj)
 
-        self.session_futures[obj.session_id] = asyncio.Future()
-        session_id = await self.session_futures[obj.session_id]
-        del self.session_futures[obj.session_id]
-        return session_id
+        try:
+            return await future
+        finally:
+            self.session_futures.pop(obj.session_id, None)
 
     async def close_session(
         self: TokenizerManager,

python/sglang/srt/managers/tokenizer_manager.py

@@ -2365,9 +2365,15 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin):
         self.send_to_scheduler.send_pyobj(ranks)
 
     def _handle_open_session_req_output(self, recv_obj):
-        self.session_futures[recv_obj.session_id].set_result(
-            recv_obj.session_id if recv_obj.success else None
-        )
+        future = self.session_futures.get(recv_obj.session_id)
+        if future is None:
+            logger.warning(
+                "Open session response arrived after waiter cleanup: %s",
+                recv_obj.session_id,
+            )
+            return
+        if not future.done():
+            future.set_result(recv_obj.session_id if recv_obj.success else None)
 
     def _handle_update_weights_from_disk_req_output(self, recv_obj):
         if self.server_args.dp_size == 1:

python/sglang/srt/mem_cache/common.py

@@ -9,6 +9,7 @@ import triton.language as tl
 
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams
 from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
+from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache
 from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
 from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils import support_triton
@@ -476,13 +477,52 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
             req.mamba_pool_idx = None
         return
 
+    # Streaming sessions transfer req_pool ownership into SessionSlot objects.
+    # Trim any speculative tail before that transfer, otherwise later turns
+    # restore only the committed prefix and can strand unreachable KV pages.
+    #
+    # Aborted streaming-session requests (e.g. input too long) skip the
+    # streaming path entirely.  match_prefix did not restore the slot's KV
+    # state, so the request has a fresh pool slot that should be freed by
+    # cache_finished_req below (which also sets req_pool_idx = None).
+    from sglang.srt.managers.schedule_batch import FINISH_ABORT
+
+    is_streaming_session = (
+        isinstance(tree_cache, SessionAwareCache)
+        and getattr(req, "session", None) is not None
+        and req.session.streaming
+    )
+    is_aborted_streaming = is_streaming_session and isinstance(
+        getattr(req, "finished_reason", None), FINISH_ABORT
+    )
+    if is_streaming_session and not is_aborted_streaming:
+        start_p, end_p = req.pop_overallocated_kv_cache()
+        page_size = get_global_server_args().page_size
+        if page_size > 1:
+            start_p = ceil_align(start_p, page_size)
+        if start_p < end_p:
+            indices_to_free = tree_cache.req_to_token_pool.req_to_token[
+                req.req_pool_idx
+            ][start_p:end_p]
+            tree_cache.token_to_kv_pool_allocator.free(indices_to_free)
+        req.kv_allocated_len = req.kv_committed_len
+
     tree_cache.cache_finished_req(req, is_insert=is_insert)
 
-    # FIXME: SessionAwareCache.cache_finished_req sets req_pool_idx = None to
-    # transfer KV ownership to the SessionSlot, so we skip the remaining
-    # cleanup (overalloc free + pool slot free). This means over-allocated
-    # tokens from speculative decoding are NOT freed between turns.
+    # SessionAwareCache.cache_finished_req sets req_pool_idx = None to transfer
+    # KV ownership to the SessionSlot, so the remaining cleanup is skipped.
+    # Streaming-session specific overalloc trimming must therefore happen
+    # before cache_finished_req above.
     if req.req_pool_idx is None:
+        if is_streaming_session:
+            # The request no longer owns any KV once SessionAwareCache either
+            # transfers it into the session slot or frees it on abort. Mark
+            # both bookkeeping flags so busy-time memory checks do not keep
+            # counting this finished request as uncached KV.
+            if not req.kv_committed_freed:
+                req.pop_committed_kv_cache()
+            if not req.kv_overallocated_freed:
+                req.pop_overallocated_kv_cache()
         return
 
     start_p, end_p = req.pop_overallocated_kv_cache()

python/sglang/srt/mem_cache/session_aware_cache.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import logging
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Dict, Optional
 
@@ -22,6 +23,9 @@ if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
 
 
+logger = logging.getLogger(__name__)
+
+
 class _VirtualNode:
     """Sentinel node for streaming session requests.
 
@@ -187,6 +191,16 @@ 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.
+        if req.to_finish is not None:
+            return self.inner.match_prefix(params)
+
         slot.restore_to_req(req)
 
         # logprob_start_len is already forced to -1 for streaming sessions
@@ -208,13 +222,56 @@ class SessionAwareCache(BasePrefixCache):
         if not _is_streaming(req):
             return self.inner.cache_finished_req(req, is_insert=is_insert, **kwargs)
 
+        from sglang.srt.managers.schedule_batch import FINISH_ABORT
+
         session_id = req.session.session_id
         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
+            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)
 
     def cache_unfinished_req(self, req: Req, **kwargs):
@@ -251,23 +308,99 @@ class SessionAwareCache(BasePrefixCache):
 
     # -- Session lifecycle --
 
-    def release_session(self, session_id: str):
+    def _resolve_release_state(
+        self, slot: SessionSlot, req: Optional[Req]
+    ) -> tuple[int, Any]:
+        """Resolve the currently tree-owned prefix for a session slot.
+
+        A long-lived session can outlive radix-tree splits caused by unrelated
+        traffic. In that case, the saved `last_node` may no longer represent the
+        full protected prefix even though the slot's req_to_token row still
+        contains tree-owned indices at the front. Re-match the current request
+        text, then intersect the returned tree indices with the slot's row so
+        release uses the prefix that is still actually backed by the tree.
+        """
+        protected_len = slot.cache_protected_len
+        lock_node = slot.last_node
+
+        # TODO: re-match logic disabled — match_prefix has side effects
+        # (splits) that disturb tree accounting. Directly using
+        # slot.last_node + cache_protected_len is safe after split analysis.
+        return protected_len, lock_node
+
+        if (
+            req is None
+            or not slot.is_holding_kv
+            or slot.req_pool_idx is None
+            or protected_len <= 0
+        ):
+            return protected_len, lock_node
+
+        from sglang.srt.mem_cache.radix_cache import RadixKey
+
+        token_ids = (req.origin_input_ids + req.output_ids)[: slot.kv_committed_len]
+        if not token_ids:
+            return 0, None
+
+        match = self.inner.match_prefix(
+            MatchPrefixParams(
+                key=RadixKey(token_ids=token_ids, extra_key=req.extra_key),
+                req=None,
+            )
+        )
+        if len(match.device_indices) == 0:
+            return 0, None
+
+        max_protected_len = min(len(match.device_indices), protected_len)
+        row_indices = self.req_to_token_pool.req_to_token[
+            slot.req_pool_idx, :max_protected_len
+        ].to(dtype=torch.int64)
+        match_indices = match.device_indices[:max_protected_len]
+        mismatches = (match_indices != row_indices).nonzero(as_tuple=False)
+        if mismatches.numel() == 0 and max_protected_len == len(match.device_indices):
+            common_len = max_protected_len
+            return common_len, match.last_device_node
+
+        common_len = (
+            int(mismatches[0].item()) if mismatches.numel() > 0 else max_protected_len
+        )
+        if self.page_size > 1:
+            common_len = (common_len // self.page_size) * self.page_size
+        if common_len <= 0:
+            return 0, None
+
+        rematch = self.inner.match_prefix(
+            MatchPrefixParams(
+                key=RadixKey(token_ids=token_ids[:common_len], extra_key=req.extra_key),
+                req=None,
+            )
+        )
+        return len(rematch.device_indices), rematch.last_device_node
+
+    def release_session(self, session_id: str, req: Optional[Req] = None):
         """Release all KV resources held by a streaming session."""
         slot = self.slots.pop(session_id, None)
         if slot is None:
             return
+        protected_len, lock_node = self._resolve_release_state(slot, req)
+        tokens_freed = (
+            max(0, slot.kv_allocated_len - protected_len) if slot.is_holding_kv else 0
+        )
+        logger.info(
+            "Session KV released: %s (%d tokens freed)", session_id, tokens_freed
+        )
 
-        if slot.last_node is not None:
+        if lock_node is not None:
             if slot.swa_uuid_for_lock is not None:
                 self.inner.dec_lock_ref(
-                    slot.last_node,
+                    lock_node,
                     DecLockRefParams(swa_uuid_for_lock=slot.swa_uuid_for_lock),
                 )
             else:
-                self.inner.dec_lock_ref(slot.last_node)
+                self.inner.dec_lock_ref(lock_node)
 
         if slot.is_holding_kv:
-            start = slot.cache_protected_len
+            start = protected_len
             end = slot.kv_allocated_len
             if start < end:
                 kv_indices = self.req_to_token_pool.req_to_token[

python/sglang/srt/observability/metrics_collector.py

@@ -132,6 +132,10 @@ class SchedulerStats:
     hicache_host_used_tokens: int = 0
     hicache_host_total_tokens: int = 0
 
+    # Streaming session metrics
+    num_streaming_sessions: int = 0
+    streaming_session_held_tokens: int = 0
+
     # Routing key metrics
     num_unique_running_routing_keys: int = 0
     routing_key_running_req_counts: List[int] = field(default_factory=list)
@@ -176,6 +180,7 @@ class SchedulerMetricsCollector:
         labels: Dict[str, str],
         enable_lora: bool = False,
         enable_hierarchical_cache: bool = False,
+        enable_streaming_session: bool = False,
         server_args: Optional["ServerArgs"] = None,
     ) -> None:
         # We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR`
@@ -184,6 +189,7 @@ class SchedulerMetricsCollector:
         self.labels = labels
         self.enable_lora = enable_lora
         self.enable_hierarchical_cache = enable_hierarchical_cache
+        self.enable_streaming_session = enable_streaming_session
         self.last_log_time = time.perf_counter()
         self._known_priorities: Set[int] = set()
 
@@ -654,6 +660,21 @@ class SchedulerMetricsCollector:
                 multiprocess_mode="mostrecent",
             )
 
+        # Streaming session metrics (only created when streaming sessions are enabled)
+        if self.enable_streaming_session:
+            self.num_streaming_sessions = Gauge(
+                name="sglang:num_streaming_sessions",
+                documentation="The number of active streaming sessions.",
+                labelnames=labels.keys(),
+                multiprocess_mode="mostrecent",
+            )
+            self.streaming_session_held_tokens = Gauge(
+                name="sglang:streaming_session_held_tokens",
+                documentation="The number of KV tokens currently held by streaming session slots.",
+                labelnames=labels.keys(),
+                multiprocess_mode="mostrecent",
+            )
+
         self.num_unique_running_routing_keys = Gauge(
             name="sglang:num_unique_running_routing_keys",
             documentation="Number of unique routing keys in running batch.",
@@ -1049,6 +1070,13 @@ class SchedulerMetricsCollector:
                 self.hicache_host_total_tokens, stats.hicache_host_total_tokens
             )
 
+        # Streaming session metrics (only logged if streaming sessions are enabled)
+        if self.enable_streaming_session:
+            self._log_gauge(self.num_streaming_sessions, stats.num_streaming_sessions)
+            self._log_gauge(
+                self.streaming_session_held_tokens, stats.streaming_session_held_tokens
+            )
+
         self._log_gauge(
             self.num_unique_running_routing_keys, stats.num_unique_running_routing_keys
         )

python/sglang/srt/observability/scheduler_metrics_mixin.py

@@ -148,6 +148,7 @@ class SchedulerMetricsMixin:
                 labels=labels,
                 enable_lora=self.enable_lora,
                 enable_hierarchical_cache=self.enable_hierarchical_cache,
+                enable_streaming_session=self.server_args.enable_streaming_session,
                 server_args=self.server_args,
             )
             self.enable_mfu_metrics = bool(self.server_args.enable_mfu_metrics)
@@ -602,6 +603,8 @@ class SchedulerMetricsMixin:
             self.stats.cache_hit_rate = cache_hit_rate
 
             self.stats.max_total_num_tokens = self.max_total_num_tokens
+            self.stats.num_streaming_sessions = self._alive_streaming_session_count()
+            self.stats.streaming_session_held_tokens = self._session_held_tokens()
 
             # Speculative decoding
             self.stats.spec_accept_rate = spec_accept_rate