Fix Session for multimodal and expose it through Engine (#18152)

python/sglang/srt/entrypoints/engine.py

@@ -42,6 +42,7 @@ from sglang.srt.managers.data_parallel_controller import (
 )
 from sglang.srt.managers.detokenizer_manager import run_detokenizer_process
 from sglang.srt.managers.io_struct import (
+    CloseSessionReqInput,
     DestroyWeightsUpdateGroupReqInput,
     EmbeddingReqInput,
     GenerateReqInput,
@@ -50,6 +51,7 @@ from sglang.srt.managers.io_struct import (
     LoadLoRAAdapterFromTensorsReqInput,
     LoadLoRAAdapterReqInput,
     MultimodalDataInputFormat,
+    OpenSessionReqInput,
     ReleaseMemoryOccupationReqInput,
     ResumeMemoryOccupationReqInput,
     RpcReqInput,
@@ -233,6 +235,7 @@ class Engine(EngineBase):
         data_parallel_rank: Optional[int] = None,
         external_trace_header: Optional[Dict] = None,
         rid: Optional[Union[List[str], str]] = None,
+        session_params: Optional[Dict] = None,
     ) -> Union[Dict, Iterator[Dict]]:
         """
         The arguments of this function is the same as `sglang/srt/managers/io_struct.py::GenerateReqInput`.
@@ -270,6 +273,7 @@ class Engine(EngineBase):
             data_parallel_rank=data_parallel_rank,
             external_trace_header=external_trace_header,
             rid=rid,
+            session_params=session_params,
         )
         generator = self.tokenizer_manager.generate_request(obj, None)
 
@@ -320,6 +324,7 @@ class Engine(EngineBase):
         data_parallel_rank: Optional[int] = None,
         external_trace_header: Optional[Dict] = None,
         rid: Optional[Union[List[str], str]] = None,
+        session_params: Optional[Dict] = None,
     ) -> Union[Dict, AsyncIterator[Dict]]:
         """
         The arguments of this function is the same as `sglang/srt/managers/io_struct.py::GenerateReqInput`.
@@ -358,6 +363,7 @@ class Engine(EngineBase):
             data_parallel_rank=data_parallel_rank,
             external_trace_header=external_trace_header,
             rid=rid,
+            session_params=session_params,
         )
         generator = self.tokenizer_manager.generate_request(obj, None)
 
@@ -448,6 +454,37 @@ class Engine(EngineBase):
     def flush_cache(self):
         return self.loop.run_until_complete(self.tokenizer_manager.flush_cache())
 
+    def open_session(
+        self,
+        capacity_of_str_len: int,
+        session_id: Optional[str] = None,
+    ) -> str:
+        """Open a session for multi-turn conversation with shared context.
+
+        Args:
+            capacity_of_str_len: Maximum string length capacity for the session.
+            session_id: Optional session ID. If not provided, a UUID will be generated.
+
+        Returns:
+            The session ID (either the provided one or a newly generated UUID).
+        """
+        obj = OpenSessionReqInput(
+            capacity_of_str_len=capacity_of_str_len,
+            session_id=session_id,
+        )
+        return self.loop.run_until_complete(
+            self.tokenizer_manager.open_session(obj, None)
+        )
+
+    def close_session(self, session_id: str) -> None:
+        """Close a session and release its resources.
+
+        Args:
+            session_id: The session ID to close.
+        """
+        obj = CloseSessionReqInput(session_id=session_id)
+        self.loop.run_until_complete(self.tokenizer_manager.close_session(obj, None))
+
     def start_profile(self, **kwargs):
         self.loop.run_until_complete(self.tokenizer_manager.start_profile(**kwargs))
 

python/sglang/srt/managers/schedule_batch.py

@@ -1107,7 +1107,6 @@ class Req:
         self.last_node = None
         self.swa_uuid_for_lock = None
         self.extend_input_len = 0
-        self.customized_info = None
         self.is_retracted = True
         self.retracted_stain = True
         self.input_token_logprobs = None

python/sglang/srt/managers/scheduler.py

@@ -1415,6 +1415,10 @@ class Scheduler(
         for req in batch.reqs:
             if not req.finished() or not (mm_inputs := req.multimodal_inputs):
                 continue
+            # For session requests, keep mm_inputs for the next request
+            if req.session_id:
+                continue
+            # For non-session requests, clear features and mm_inputs
             for item in mm_inputs.mm_items:
                 item.feature = None
             req.multimodal_inputs = None
@@ -1508,6 +1512,19 @@ class Scheduler(
         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
+                        ]
+
             # The following steps are already fast, execute locally on each rank.
             # Expand a single image token into multiple dummy tokens for receiving image embeddings
             req.origin_input_ids = self.pad_input_ids_func(