[Session] Update session control interface (#2635)

python/sglang/srt/managers/session_controller.py

@@ -10,41 +10,116 @@
 # limitations under the License.
 # ==============================================================================
 
+import logging
 import uuid
+from typing import Dict, Optional
 
 from sglang.srt.managers.io_struct import TokenizedGenerateReqInput
-from sglang.srt.managers.schedule_batch import FINISH_ABORT, List, Req
+from sglang.srt.managers.schedule_batch import Req
+
+
+class SessionReqNode:
+    def __init__(self, req, parent=None, childs=None):
+        self.req = req
+        self.parent = parent
+        if parent is not None:
+            parent.childs.append(self)
+        self.childs = [] if not childs else childs
+
+    def clear_childs(self, req_dict):
+        for req_node in self.childs:
+            req_node.clear(req_dict)
+        self.childs = []
+
+    def clear(self, req_dict):
+        for req_node in self.childs:
+            req_node.clear(req_dict)
+
+        if self.req.finished_reason == None:
+            self.req.to_abort = True
+        del req_dict[self.req.rid]
+
+    def abort(self):
+        if self.req.finished_reason == None:
+            self.req.to_abort = True
+
+    def __str__(self):
+        return self._str_helper(self.req.rid)
+
+    def _str_helper(self, prefix=""):
+        if len(self.childs) == 0:
+            return prefix + "\n"
+        else:
+            origin_prefix = prefix
+            prefix += " -- " + self.childs[0].req.rid
+            ret = self.childs[0]._str_helper(prefix)
+            for child in self.childs[1:]:
+                prefix = " " * len(origin_prefix) + " \- " + child.req.rid
+                ret += child._str_helper(prefix)
+            return ret
 
 
 class Session:
-    def __init__(self, capacity_of_str_len: int, session_id: str = None):
+    def __init__(self, capacity_of_str_len: int, session_id: Optional[str] = None):
         self.session_id = session_id if session_id is not None else uuid.uuid4().hex
         self.capacity_of_str_len = capacity_of_str_len
-        self.reqs: List[Req] = []
+        self.req_nodes: Dict[str, SessionReqNode] = {}
 
     def create_req(self, req: TokenizedGenerateReqInput, tokenizer):
-        if req.session_rid is not None:
-            while len(self.reqs) > 0:
-                if self.reqs[-1].rid == req.session_rid:
-                    break
-                self.reqs = self.reqs[:-1]
+        assert req.session_params is not None
+        session_params = req.session_params
+
+        last_req_node = None
+        last_req = None
+        abort = False
+        if session_params.replace:
+            if session_params.rid is None:
+                for _, req_node in self.req_nodes.items():
+                    req_node.clear(self.req_nodes)
+            else:
+                if session_params.rid not in self.req_nodes:
+                    abort = True
+                else:
+                    last_req_node = self.req_nodes[session_params.rid]
+                    last_req_node.abort()
+                    last_req = last_req_node.req
+                    last_req_node.clear_childs(self.req_nodes)
         else:
-            self.reqs = []
-        if len(self.reqs) > 0:
+            if session_params.rid is not None:
+                if session_params.rid not in self.req_nodes:
+                    abort = True
+                else:
+                    last_req_node = self.req_nodes[session_params.rid]
+                    last_req = last_req_node.req
+                    if not last_req.finished():
+                        logging.warning(
+                            "The request in a session is appending to a request that hasn't finished."
+                        )
+                        abort = True
+
+        if last_req is not None:
+            # trim bos token if it is an append
+            if req.input_ids[0] == tokenizer.bos_token_id:
+                req.input_ids = req.input_ids[1:]
+
             input_ids = (
-                self.reqs[-1].origin_input_ids
-                + self.reqs[-1].output_ids[
-                    : self.reqs[-1].sampling_params.max_new_tokens
-                ]
-                + req.input_ids
+                last_req.origin_input_ids
+                + last_req.output_ids[: last_req.sampling_params.max_new_tokens]
             )
+            if session_params.offset and session_params.offset != 0:
+                input_ids = input_ids[: session_params.offset] + req.input_ids
+            else:
+                input_ids += req.input_ids
             input_ids_unpadded = (
-                self.reqs[-1].origin_input_ids_unpadded
-                + self.reqs[-1].output_ids[
-                    : self.reqs[-1].sampling_params.max_new_tokens
-                ]
-                + req.input_ids
+                last_req.origin_input_ids_unpadded
+                + last_req.output_ids[: last_req.sampling_params.max_new_tokens]
             )
+            if session_params.offset and session_params.offset != 0:
+                input_ids_unpadded = (
+                    input_ids_unpadded[: session_params.offset] + req.input_ids
+                )
+            else:
+                input_ids_unpadded += req.input_ids
         else:
             input_ids = req.input_ids
             input_ids_unpadded = req.input_ids
@@ -57,13 +132,13 @@ class Session:
             lora_path=req.lora_path,
             session_id=self.session_id,
         )
-        if len(self.reqs) > 0:
-            new_req.image_inputs = self.reqs[-1].image_inputs
+        if last_req is not None:
+            new_req.image_inputs = last_req.image_inputs
         new_req.tokenizer = tokenizer
-        if req.session_rid is not None and len(self.reqs) == 0:
-            new_req.finished_reason = FINISH_ABORT(
-                f"Invalid request: requested session rid {req.session_rid} does not exist in the session history"
-            )
+        if abort:
+            new_req.to_abort = True
         else:
-            self.reqs.append(new_req)
+            new_req_node = SessionReqNode(new_req, last_req_node)
+            self.req_nodes[req.rid] = new_req_node
+
         return new_req