sglang 0.4.1__py3-none-any.whl → 0.4.1.post2__py3-none-any.whl

sglang/srt/managers/scheduler.py

@@ -22,7 +22,7 @@ import warnings
 from collections import deque
 from concurrent import futures
 from types import SimpleNamespace
-from typing import Callable, Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import psutil
 import setproctitle
@@ -52,6 +52,8 @@ from sglang.srt.managers.io_struct import (
     UpdateWeightFromDiskReqOutput,
     UpdateWeightsFromDistributedReqInput,
     UpdateWeightsFromDistributedReqOutput,
+    UpdateWeightsFromTensorReqInput,
+    UpdateWeightsFromTensorReqOutput,
 )
 from sglang.srt.managers.schedule_batch import (
     FINISH_ABORT,
@@ -88,7 +90,7 @@ from sglang.utils import get_exception_traceback
 
 logger = logging.getLogger(__name__)
 
-# Test retract decode
+# Test retract decode for debugging purposes
 test_retract = get_bool_env_var("SGLANG_TEST_RETRACT")
 
 
@@ -127,12 +129,12 @@ class Scheduler:
             )
 
             if server_args.skip_tokenizer_init:
-                # Directly send to the tokenizer/api
+                # Directly send to the TokenizerManager
                 self.send_to_detokenizer = get_zmq_socket(
                     context, zmq.PUSH, port_args.tokenizer_ipc_name
                 )
             else:
-                # Send to the detokenizer
+                # Send to the DetokenizerManager
                 self.send_to_detokenizer = get_zmq_socket(
                     context, zmq.PUSH, port_args.detokenizer_ipc_name
                 )
@@ -383,7 +385,8 @@ class Scheduler:
             self.process_input_requests(recv_reqs)
 
             batch = self.get_next_batch_to_run()
-            if self.server_args.enable_dp_attention:
+
+            if self.server_args.enable_dp_attention:  # TODO: simplify this
                 batch = self.prepare_dp_attn_batch(batch)
 
             self.cur_batch = batch
@@ -392,7 +395,7 @@ class Scheduler:
                 result = self.run_batch(batch)
                 self.process_batch_result(batch, result)
             else:
-                # Self-check and re-init some states when the server is idle
+                # When the server is idle, so self-check and re-init some states
                 self.check_memory()
                 self.new_token_ratio = self.init_new_token_ratio
 
@@ -409,12 +412,13 @@ class Scheduler:
 
             batch = self.get_next_batch_to_run()
             self.cur_batch = batch
+
             if batch:
                 result = self.run_batch(batch)
                 result_queue.append((batch.copy(), result))
 
                 if self.last_batch is None:
-                    # A dummy first batch to start the pipeline for overlap scheduler.
+                    # Create a dummy first batch to start the pipeline for overlap scheduler.
                     # It is now used for triggering the sampling_info_done event.
                     tmp_batch = ScheduleBatch(
                         reqs=None,
@@ -424,19 +428,21 @@ class Scheduler:
                     self.process_batch_result(tmp_batch, None)
 
             if self.last_batch:
+                # Process the results of the last batch
                 tmp_batch, tmp_result = result_queue.popleft()
                 tmp_batch.next_batch_sampling_info = (
                     self.tp_worker.cur_sampling_info if batch else None
                 )
                 self.process_batch_result(tmp_batch, tmp_result)
             elif batch is None:
-                # Self-check and re-init some states when the server is idle
+                # When the server is idle, so self-check and re-init some states
                 self.check_memory()
                 self.new_token_ratio = self.init_new_token_ratio
 
             self.last_batch = batch
 
-    def recv_requests(self):
+    def recv_requests(self) -> List[Req]:
+        """Receive results at tp_rank = 0 and broadcast it to all other TP ranks."""
         if self.tp_rank == 0 or self.server_args.enable_dp_attention:
             recv_reqs = []
 
@@ -468,9 +474,6 @@ class Scheduler:
                 self.send_to_tokenizer.send_pyobj(
                     UpdateWeightFromDiskReqOutput(success, message)
                 )
-            elif isinstance(recv_req, GetWeightsByNameReqInput):
-                parameter = self.get_weights_by_name(recv_req)
-                self.send_to_tokenizer.send_pyobj(GetWeightsByNameReqOutput(parameter))
             elif isinstance(recv_req, InitWeightsUpdateGroupReqInput):
                 success, message = self.init_weights_update_group(recv_req)
                 self.send_to_tokenizer.send_pyobj(
@@ -481,6 +484,11 @@ class Scheduler:
                 self.send_to_tokenizer.send_pyobj(
                     UpdateWeightsFromDistributedReqOutput(success, message)
                 )
+            elif isinstance(recv_req, UpdateWeightsFromTensorReqInput):
+                success, message = self.update_weights_from_tensor(recv_req)
+                self.send_to_tokenizer.send_pyobj(
+                    UpdateWeightsFromTensorReqOutput(success, message)
+                )
             elif isinstance(recv_req, GetWeightsByNameReqInput):
                 parameter = self.get_weights_by_name(recv_req)
                 self.send_to_tokenizer.send_pyobj(GetWeightsByNameReqOutput(parameter))
@@ -490,8 +498,10 @@ class Scheduler:
                 else:
                     self.stop_profile()
             elif isinstance(recv_req, OpenSessionReqInput):
-                session_id = self.open_session(recv_req)
-                self.send_to_tokenizer.send_pyobj(OpenSessionReqOutput(session_id))
+                session_id, success = self.open_session(recv_req)
+                self.send_to_tokenizer.send_pyobj(
+                    OpenSessionReqOutput(session_id=session_id, success=success)
+                )
             elif isinstance(recv_req, CloseSessionReqInput):
                 self.close_session(recv_req)
             else:
@@ -502,7 +512,11 @@ class Scheduler:
         recv_req: TokenizedGenerateReqInput,
     ):
         # Create a new request
-        if recv_req.session_id is None or recv_req.session_id not in self.sessions:
+        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
+        ):
 
             if recv_req.input_embeds is not None:
                 # Generate fake input_ids based on the length of input_embeds
@@ -520,18 +534,22 @@ class Scheduler:
                 stream=recv_req.stream,
                 lora_path=recv_req.lora_path,
                 input_embeds=recv_req.input_embeds,
+                eos_token_ids=self.model_config.hf_eos_token_id,
             )
             req.tokenizer = self.tokenizer
 
-            if recv_req.session_id is not None:
+            if (
+                recv_req.session_params is not None
+                and recv_req.session_params.id is not None
+            ):
                 req.finished_reason = FINISH_ABORT(
-                    f"Invalid request: session id {recv_req.session_id} does not exist"
+                    f"Invalid request: session id {recv_req.session_params.id} does not exist"
                 )
                 self.waiting_queue.append(req)
                 return
         else:
-            # Create a new request from a previsou session
-            session = self.sessions[recv_req.session_id]
+            # Create a new request from a previous session
+            session = self.sessions[recv_req.session_params.id]
             req = session.create_req(recv_req, self.tokenizer)
             if isinstance(req.finished_reason, FINISH_ABORT):
                 self.waiting_queue.append(req)
@@ -565,7 +583,7 @@ class Scheduler:
 
         if req.logprob_start_len == -1:
             # By default, only return the logprobs for output tokens
-            req.logprob_start_len = len(recv_req.input_ids) - 1
+            req.logprob_start_len = len(req.origin_input_ids) - 1
 
         # Truncate prompts that are too long
         if len(req.origin_input_ids) > self.max_req_input_len:
@@ -589,12 +607,15 @@ class Scheduler:
         if (
             req.sampling_params.json_schema is not None
             or req.sampling_params.regex is not None
+            or req.sampling_params.ebnf is not None
         ):
             assert self.grammar_backend is not None
             if req.sampling_params.json_schema is not None:
                 key = ("json", req.sampling_params.json_schema)
             elif req.sampling_params.regex is not None:
                 key = ("regex", req.sampling_params.regex)
+            elif req.sampling_params.ebnf is not None:
+                key = ("ebnf", req.sampling_params.ebnf)
 
             req.grammar = self.grammar_backend.get_cached_value(key)
             if not req.grammar:
@@ -629,16 +650,13 @@ class Scheduler:
         self.waiting_queue.append(req)
 
     def log_prefill_stats(self, adder, can_run_list, running_bs, has_being_chunked):
-        if isinstance(self.tree_cache, RadixCache):
-            self.tree_cache_metrics["total"] += (
-                adder.log_input_tokens + adder.log_hit_tokens
-            ) / 10**9
-            self.tree_cache_metrics["hit"] += (adder.log_hit_tokens) / 10**9
-            tree_cache_hit_rate = (
-                self.tree_cache_metrics["hit"] / self.tree_cache_metrics["total"]
-            )
-        else:
-            tree_cache_hit_rate = 0.0
+        self.tree_cache_metrics["total"] += (
+            adder.log_input_tokens + adder.log_hit_tokens
+        ) / 10**9
+        self.tree_cache_metrics["hit"] += (adder.log_hit_tokens) / 10**9
+        tree_cache_hit_rate = (
+            self.tree_cache_metrics["hit"] / self.tree_cache_metrics["total"]
+        )
 
         num_used = self.max_total_num_tokens - (
             self.token_to_kv_pool.available_size() + self.tree_cache.evictable_size()
@@ -807,6 +825,8 @@ class Scheduler:
                 if res == AddReqResult.NO_TOKEN:
                     self.batch_is_full = True
                 break
+            if self.server_args.prefill_only_one_req:
+                break
 
         # Update waiting queue
         can_run_list = adder.can_run_list
@@ -1460,6 +1480,17 @@ class Scheduler:
             logger.error(message)
         return success, message
 
+    def update_weights_from_tensor(self, recv_req: UpdateWeightsFromTensorReqInput):
+        """Update the online model parameter from tensors."""
+        success, message = self.tp_worker.update_weights_from_tensor(recv_req)
+        # TODO extract common code b/t update_weights_from_distributed and update_weights_from_tensor later
+        if success:
+            flash_cache_success = self.flush_cache()
+            assert flash_cache_success, "Cache flush failed after updating weights"
+        else:
+            logger.error(message)
+        return success, message
+
     def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput):
         parameter = self.tp_worker.get_weights_by_name(recv_req)
         return parameter
@@ -1478,16 +1509,20 @@ class Scheduler:
         )
         logger.info("Profiler is done")
 
-    def open_session(self, recv_req: OpenSessionReqInput) -> str:
+    def open_session(self, recv_req: OpenSessionReqInput) -> Tuple[Optional[str], bool]:
         # handle error
         session_id = recv_req.session_id
         if session_id in self.sessions:
             logger.warning(f"session id {session_id} already exist, cannot open.")
+            return session_id, False
+        elif session_id is None:
+            logger.warning(f"session id is None, cannot open.")
+            return session_id, False
         else:
             self.sessions[session_id] = Session(
                 recv_req.capacity_of_str_len, session_id
             )
-        return session_id
+            return session_id, True
 
     def close_session(self, recv_req: CloseSessionReqInput):
         # handle error
@@ -1512,18 +1547,20 @@ def run_scheduler_process(
     if dp_rank is None and "SGLANG_DP_RANK" in os.environ:
         dp_rank = int(os.environ["SGLANG_DP_RANK"])
 
+    # Configue the logger
     if dp_rank is None:
         configure_logger(server_args, prefix=f" TP{tp_rank}")
     else:
         configure_logger(server_args, prefix=f" DP{dp_rank} TP{tp_rank}")
+    suppress_other_loggers()
 
-    # set cpu affinity to this gpu process
+    # Set cpu affinity to this gpu process
     if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"):
         set_gpu_proc_affinity(server_args.tp_size, server_args.nnodes, gpu_id)
 
-    suppress_other_loggers()
     parent_process = psutil.Process().parent()
 
+    # Create a scheduler and run the event loop
     try:
         scheduler = Scheduler(server_args, port_args, gpu_id, tp_rank, dp_rank)
         pipe_writer.send(

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