diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 6a13e3fec..4106e5cb9 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -22,7 +22,7 @@ import time from collections import deque from dataclasses import dataclass from http import HTTPStatus -from typing import Any, Deque, Dict, List, Optional, Tuple, Union +from typing import Any, Deque, List, Optional, Tuple, Union import psutil import setproctitle @@ -112,7 +112,6 @@ from sglang.srt.managers.io_struct import ( LoadLoRAAdapterReqInput, LoadLoRAAdapterReqOutput, OpenSessionReqInput, - OpenSessionReqOutput, PauseGenerationReqInput, ProfileReq, ReleaseMemoryOccupationReqInput, @@ -167,7 +166,7 @@ from sglang.srt.managers.scheduler_runtime_checker_mixin import ( from sglang.srt.managers.scheduler_update_weights_mixin import ( SchedulerUpdateWeightsMixin, ) -from sglang.srt.managers.session_controller import Session +from sglang.srt.managers.session_controller import SessionController from sglang.srt.managers.utils import GenerationBatchResult, validate_input_length from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.common import release_kv_cache @@ -755,8 +754,7 @@ class Scheduler( self.return_health_check_ct = 0 self.num_retracted_reqs: int = 0 self.num_paused_reqs: int = 0 - self.sessions: Dict[str, Session] = {} - self._last_reap_sessions: float = 0.0 + self.session_controller = SessionController(self.tree_cache) self.forward_sleep_time = None self._engine_paused = False @@ -1131,7 +1129,6 @@ class Scheduler( self.process_batch_result(batch, result) else: # When the server is idle, do self-check and re-init some states. - # Skip if there are any streaming sessions (latency sensitive). self.self_check_during_idle() # Update last_batch @@ -1363,9 +1360,7 @@ class Scheduler( def process_input_requests(self, recv_reqs: List): now = time.monotonic() - if now - self._last_reap_sessions > 1.0: # reap sessions every second - self._last_reap_sessions = now - self.reap_timed_out_sessions() + self.session_controller.maybe_reap(now) for recv_req in recv_reqs: # If it is a health check generation request and there are running requests, ignore it. if is_health_check_generate_req(recv_req) and ( @@ -1485,12 +1480,13 @@ class Scheduler( self, recv_req: TokenizedGenerateReqInput, ): - # Create a new request - if ( - recv_req.session_params is None - or recv_req.session_params.id is None - or recv_req.session_params.id not in self.sessions - ): + # Route: normal request / session request / session-not-found + session_id = ( + recv_req.session_params.id if recv_req.session_params is not None else None + ) + + if session_id is None: + # Normal non-session request if recv_req.input_embeds is not None: # Generate fake input_ids based on the length of input_embeds seq_length = len(recv_req.input_embeds) @@ -1553,21 +1549,14 @@ class Scheduler( self.stream_output([req], req.return_logprob) return - if ( - recv_req.session_params is not None - and recv_req.session_params.id is not None - ): - req.set_finish_with_abort( - f"Invalid request: session id {recv_req.session_params.id} does not exist" - ) - self.init_req_max_new_tokens(req) - self._add_request_to_queue(req) - return - else: - # Create a new request from a previous session - session = self.sessions[recv_req.session_params.id] + elif session_id in self.session_controller: + # Session exists: create request from session + session = self.session_controller.get(session_id) req = session.create_req( - recv_req, self.tokenizer, self.model_config.vocab_size + recv_req, + self.tokenizer, + self.model_config.vocab_size, + eos_token_ids=self.model_config.hf_eos_token_id, ) # TODO: set trace context if self.enable_metrics: @@ -1577,22 +1566,28 @@ class Scheduler( self._add_request_to_queue(req) return + else: + # Session ID provided but session not found + req = Req( + recv_req.rid, + recv_req.input_text, + recv_req.input_ids, + recv_req.sampling_params, + 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" + ) + self.init_req_max_new_tokens(req) + self._add_request_to_queue(req) + return + # Handle multimodal inputs if recv_req.mm_inputs is not None: image_inputs = self._get_multimodal_inputs(recv_req.mm_inputs) - # For session requests, adjust mm_inputs offsets by the prefix length. - # Session.create_req prepends previous context to origin_input_ids, - # so offsets from the new prompt need to be shifted. - if len(recv_req.input_ids) < len(req.origin_input_ids): - assert recv_req.session_params.id in self.sessions - prefix_len = len(req.origin_input_ids) - len(recv_req.input_ids) - for mm_item in image_inputs.mm_items: - if mm_item.offsets: - mm_item.offsets = [ - (start + prefix_len, end + prefix_len) - for start, end in mm_item.offsets - ] + SessionController.adjust_mm_offsets(recv_req, req, image_inputs) # The following steps are already fast, execute locally on each rank. # Expand a single image token into multiple dummy tokens for receiving image embeddings @@ -2938,47 +2933,10 @@ class Scheduler( return ExpertDistributionReqOutput() def open_session(self, recv_req: OpenSessionReqInput): - session_id = recv_req.session_id - if session_id in self.sessions: - logger.warning(f"session id {session_id} already exist, cannot open.") - return OpenSessionReqOutput(session_id, False) - elif session_id is None: - logger.warning("session id is None, cannot open.") - return OpenSessionReqOutput(session_id, False) - else: - self.sessions[session_id] = Session( - recv_req.capacity_of_str_len, - session_id, - streaming=bool(recv_req.streaming), - timeout=recv_req.timeout, - ) - return OpenSessionReqOutput(session_id, True) + return self.session_controller.open(recv_req) def close_session(self, recv_req: CloseSessionReqInput): - session_id = recv_req.session_id - if session_id not in self.sessions: - logger.warning(f"session id {session_id} does not exist, cannot delete.") - else: - self._close_session(session_id) - - def _close_session(self, session_id: str): - session = self.sessions[session_id] - 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 - if isinstance(self.tree_cache, SessionAwareCache): - self.tree_cache.release_session(session_id) - del self.sessions[session_id] - - def reap_timed_out_sessions(self): - 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.") - self._close_session(sid) + self.session_controller.close(recv_req) def maybe_sleep_on_idle(self): if self.idle_sleeper is not None: diff --git a/python/sglang/srt/managers/session_controller.py b/python/sglang/srt/managers/session_controller.py index bbc619191..05cd15cfa 100644 --- a/python/sglang/srt/managers/session_controller.py +++ b/python/sglang/srt/managers/session_controller.py @@ -10,13 +10,26 @@ # limitations under the License. # ============================================================================== +from __future__ import annotations + import logging import time import uuid -from typing import Dict, Optional +from typing import TYPE_CHECKING, Dict, Optional -from sglang.srt.managers.io_struct import TokenizedGenerateReqInput +from sglang.srt.managers.io_struct import ( + CloseSessionReqInput, + OpenSessionReqInput, + OpenSessionReqOutput, + TokenizedGenerateReqInput, +) from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req +from sglang.srt.mem_cache.session_aware_cache import SessionAwareCache + +if TYPE_CHECKING: + from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache + +logger = logging.getLogger(__name__) class SessionReqNode: @@ -85,7 +98,13 @@ class Session: return False return time.monotonic() - self.last_active_time > self.timeout - def create_req(self, req: TokenizedGenerateReqInput, tokenizer, vocab_size: int): + def create_req( + self, + req: TokenizedGenerateReqInput, + tokenizer, + vocab_size: int, + eos_token_ids=None, + ): assert req.session_params is not None self.last_active_time = time.monotonic() session_params = req.session_params @@ -190,6 +209,14 @@ class Session: top_logprobs_num=req.top_logprobs_num, token_ids_logprob=req.token_ids_logprob, vocab_size=vocab_size, + eos_token_ids=eos_token_ids, + require_reasoning=req.require_reasoning, + return_hidden_states=req.return_hidden_states, + return_routed_experts=req.return_routed_experts, + priority=req.priority, + routing_key=req.routing_key, + http_worker_ipc=req.http_worker_ipc, + time_stats=req.time_stats, ) if last_req is not None: new_req.multimodal_inputs = last_req.multimodal_inputs @@ -206,3 +233,77 @@ class Session: self.req_nodes[req.rid] = new_req_node return new_req + + +class SessionController: + def __init__(self, tree_cache: BasePrefixCache): + self.sessions: Dict[str, Session] = {} + self._last_reap_time: float = 0.0 + self.tree_cache = tree_cache + + def __contains__(self, session_id: str) -> bool: + return session_id in self.sessions + + def get(self, session_id: str) -> Optional[Session]: + return self.sessions.get(session_id) + + def open(self, recv_req: OpenSessionReqInput) -> OpenSessionReqOutput: + session_id = recv_req.session_id + if session_id in self.sessions: + logger.warning(f"session id {session_id} already exist, cannot open.") + return OpenSessionReqOutput(session_id, False) + elif session_id is None: + logger.warning("session id is None, cannot open.") + return OpenSessionReqOutput(session_id, False) + else: + self.sessions[session_id] = Session( + recv_req.capacity_of_str_len, + session_id, + streaming=bool(recv_req.streaming), + timeout=recv_req.timeout, + ) + return OpenSessionReqOutput(session_id, True) + + def close(self, recv_req: CloseSessionReqInput): + session_id = recv_req.session_id + if session_id not in self.sessions: + logger.warning(f"session id {session_id} does not exist, cannot delete.") + else: + self._close(session_id) + + def _close(self, session_id: str): + session = self.sessions[session_id] + 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 + if isinstance(self.tree_cache, SessionAwareCache): + self.tree_cache.release_session(session_id) + del self.sessions[session_id] + + 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 + 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.") + self._close(sid) + + @staticmethod + def adjust_mm_offsets(recv_req: TokenizedGenerateReqInput, req: Req, image_inputs): + # For session requests, adjust mm_inputs offsets by the prefix length. + # Session.create_req prepends previous context to origin_input_ids, + # so offsets from the new prompt need to be shifted. + if len(recv_req.input_ids) >= len(req.origin_input_ids): + return + prefix_len = len(req.origin_input_ids) - len(recv_req.input_ids) + for mm_item in image_inputs.mm_items: + if mm_item.offsets: + mm_item.offsets = [ + (start + prefix_len, end + prefix_len) + for start, end in mm_item.offsets + ]