sglang 0.1.14__py3-none-any.whl → 0.1.21__py3-none-any.whl

sglang/srt/managers/detokenizer_manager.py

@@ -1,12 +1,17 @@
+"""DetokenizerManager is a process that detokenizes the token ids."""
+
 import asyncio
+import inspect
 
 import uvloop
 import zmq
 import zmq.asyncio
+
 from sglang.srt.hf_transformers_utils import get_tokenizer
+from sglang.srt.managers.controller.infer_batch import FINISH_MATCHED_STR
 from sglang.srt.managers.io_struct import BatchStrOut, BatchTokenIDOut
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import get_exception_traceback
+from sglang.utils import find_printable_text, get_exception_traceback, graceful_registry
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
@@ -32,48 +37,43 @@ class DetokenizerManager:
 
     async def handle_loop(self):
         while True:
-            recv_obj = await self.recv_from_router.recv_pyobj()
+            recv_obj: BatchTokenIDOut = await self.recv_from_router.recv_pyobj()
+            assert isinstance(recv_obj, BatchTokenIDOut)
 
-            if isinstance(recv_obj, BatchTokenIDOut):
-                output_tokens = recv_obj.output_tokens
+            # TODO(lmzheng): handle skip_special_tokens/spaces_between_special_tokens per request
+            surr_texts = self.tokenizer.batch_decode(
+                recv_obj.surr_output_ids,
+                skip_special_tokens=recv_obj.skip_special_tokens[0],
+                spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
+            )
+            read_texts = self.tokenizer.batch_decode(
+                recv_obj.read_output_ids,
+                skip_special_tokens=recv_obj.skip_special_tokens[0],
+                spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
+            )
 
-                # TODO(lmzheng): handle skip_special_tokens per request
-                output_strs = self.tokenizer.batch_decode(
-                    output_tokens,
-                    skip_special_tokens=recv_obj.skip_special_tokens[0],
-                )
+            # Trim stop str
+            # TODO(lmzheng): handle the case where multiple stop strs are hit
+            output_strs = []
+            for i in range(len(recv_obj.rids)):
+                new_text = read_texts[i][len(surr_texts[i]) :]
+                if recv_obj.finished_reason[i] is None:
+                    new_text = find_printable_text(new_text)
+                output_strs.append(recv_obj.decoded_texts[i] + new_text)
 
-                # Trim stop str
-                # TODO(lmzheng): handle the case where multiple stop strs are hit
-                for i in range(len(output_strs)):
-                    if recv_obj.hit_stop_str[i] is not None:
-                        pos = output_strs[i].find(recv_obj.hit_stop_str[i])
-                        if pos != -1:
-                            output_strs[i] = output_strs[i][:pos]
-
-                    if len(output_tokens[i]) > 0:
-                        first_token = self.tokenizer.convert_ids_to_tokens(
-                            int(output_tokens[i][0])
-                        )
-                        if not isinstance(first_token, str):
-                            first_token = first_token.decode("utf-8", errors="ignore")
-                        if first_token.startswith("▁"):
-                            output_strs[i] = " " + output_strs[i]
-
-                    output_strs[i] = (
-                        recv_obj.output_and_jump_forward_strs[i] + output_strs[i]
-                    )
-
-                self.send_to_tokenizer.send_pyobj(
-                    BatchStrOut(
-                        recv_obj.rids,
-                        output_strs,
-                        recv_obj.meta_info,
-                        recv_obj.finished,
-                    )
+                if isinstance(recv_obj.finished_reason[i], FINISH_MATCHED_STR):
+                    pos = output_strs[i].find(recv_obj.finished_reason[i].matched)
+                    if pos != -1:
+                        output_strs[i] = output_strs[i][:pos]
+
+            self.send_to_tokenizer.send_pyobj(
+                BatchStrOut(
+                    rids=recv_obj.rids,
+                    output_strs=output_strs,
+                    meta_info=recv_obj.meta_info,
+                    finished_reason=recv_obj.finished_reason,
                 )
-            else:
-                raise ValueError(f"Invalid object: {recv_obj}")
+            )
 
 
 def start_detokenizer_process(
@@ -81,9 +81,11 @@ def start_detokenizer_process(
     port_args: PortArgs,
     pipe_writer,
 ):
+    graceful_registry(inspect.currentframe().f_code.co_name)
+
     try:
         manager = DetokenizerManager(server_args, port_args)
-    except Exception as e:
+    except Exception:
         pipe_writer.send(get_exception_traceback())
         raise
     pipe_writer.send("init ok")

sglang/srt/managers/io_struct.py

@@ -1,14 +1,22 @@
+"""
+The definition of objects transfered between different
+processes (TokenizerManager, DetokenizerManager, Controller).
+"""
+
 import uuid
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
 
+from sglang.srt.managers.controller.infer_batch import BaseFinishReason
 from sglang.srt.sampling_params import SamplingParams
 
 
 @dataclass
 class GenerateReqInput:
     # The input prompt
-    text: Union[List[str], str]
+    text: Optional[Union[List[str], str]] = None
+    # The token ids for text; one can either specify text or input_ids
+    input_ids: Optional[Union[List[List[int]], List[int]]] = None
     # The image input
     image_data: Optional[Union[List[str], str]] = None
     # The sampling_params
@@ -19,13 +27,24 @@ class GenerateReqInput:
     return_logprob: Optional[Union[List[bool], bool]] = None
     # The start location of the prompt for return_logprob
     logprob_start_len: Optional[Union[List[int], int]] = None
+    # The number of top logprobs to return
+    top_logprobs_num: Optional[Union[List[int], int]] = None
     # Whether to detokenize tokens in logprobs
     return_text_in_logprobs: bool = False
     # Whether to stream output
     stream: bool = False
 
     def post_init(self):
-        is_single = isinstance(self.text, str)
+        if (self.text is None and self.input_ids is None) or (
+            self.text is not None and self.input_ids is not None
+        ):
+            raise ValueError("Either text or input_ids should be provided.")
+
+        if self.text is not None:
+            is_single = isinstance(self.text, str)
+        else:
+            is_single = isinstance(self.input_ids[0], int)
+        self.is_single = is_single
 
         if is_single:
             if self.sampling_params is None:
@@ -36,8 +55,10 @@ class GenerateReqInput:
                 self.return_logprob = False
             if self.logprob_start_len is None:
                 self.logprob_start_len = 0
+            if self.top_logprobs_num is None:
+                self.top_logprobs_num = 0
         else:
-            num = len(self.text)
+            num = len(self.text) if self.text is not None else len(self.input_ids)
 
             if self.image_data is None:
                 self.image_data = [None] * num
@@ -52,7 +73,8 @@ class GenerateReqInput:
             if self.rid is None:
                 self.rid = [uuid.uuid4().hex for _ in range(num)]
             else:
-                assert isinstance(self.rid, list)
+                if not isinstance(self.rid, list):
+                    raise ValueError("The rid should be a list.")
 
             if self.return_logprob is None:
                 self.return_logprob = [False] * num
@@ -64,6 +86,11 @@ class GenerateReqInput:
             elif not isinstance(self.logprob_start_len, list):
                 self.logprob_start_len = [self.logprob_start_len] * num
 
+            if self.top_logprobs_num is None:
+                self.top_logprobs_num = [0] * num
+            elif not isinstance(self.top_logprobs_num, list):
+                self.top_logprobs_num = [self.top_logprobs_num] * num
+
 
 @dataclass
 class TokenizedGenerateReqInput:
@@ -76,26 +103,28 @@ class TokenizedGenerateReqInput:
     sampling_params: SamplingParams
     return_logprob: bool
     logprob_start_len: int
+    top_logprobs_num: int
     stream: bool
 
 
 @dataclass
 class BatchTokenIDOut:
     rids: List[str]
-    output_tokens: List[List[int]]
-    output_and_jump_forward_strs: List[str]
-    hit_stop_str: List[Optional[str]]
+    decoded_texts: List[str]
+    surr_output_ids: List[List[int]]
+    read_output_ids: List[List[int]]
     skip_special_tokens: List[bool]
+    spaces_between_special_tokens: List[bool]
     meta_info: List[Dict]
-    finished: List[bool]
+    finished_reason: List[BaseFinishReason]
 
 
 @dataclass
 class BatchStrOut:
     rids: List[str]
-    output_str: List[str]
+    output_strs: List[str]
     meta_info: List[Dict]
-    finished: List[bool]
+    finished_reason: List[BaseFinishReason]
 
 
 @dataclass
@@ -103,6 +132,11 @@ class FlushCacheReq:
     pass
 
 
+@dataclass
+class AbortReq:
+    rid: str
+
+
 @dataclass
 class DetokenizeReqInput:
     input_ids: List[int]

sglang/srt/managers/tokenizer_manager.py

@@ -1,15 +1,20 @@
+"""TokenizerManager is a process that tokenizes the text."""
+
 import asyncio
 import concurrent.futures
 import dataclasses
+import logging
 import multiprocessing as mp
 import os
-from typing import List
+from typing import Dict, List
 
 import numpy as np
 import transformers
 import uvloop
 import zmq
 import zmq.asyncio
+from fastapi import BackgroundTasks
+
 from sglang.srt.hf_transformers_utils import (
     get_config,
     get_context_length,
@@ -17,8 +22,9 @@ from sglang.srt.hf_transformers_utils import (
     get_tokenizer,
 )
 from sglang.srt.managers.io_struct import (
+    AbortReq,
     BatchStrOut,
-    DetokenizeReqInput,
+    BatchTokenIDOut,
     FlushCacheReq,
     GenerateReqInput,
     TokenizedGenerateReqInput,
@@ -26,54 +32,19 @@ from sglang.srt.managers.io_struct import (
 from sglang.srt.mm_utils import expand2square, process_anyres_image
 from sglang.srt.sampling_params import SamplingParams
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import get_exception_traceback, is_multimodal_model, load_image
+from sglang.srt.utils import is_multimodal_model, load_image
+from sglang.utils import get_exception_traceback
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
+logger = logging.getLogger(__name__)
+
 
 @dataclasses.dataclass
 class ReqState:
     out_list: List
     finished: bool
     event: asyncio.Event
-    lock: asyncio.Lock
-
-
-global global_processor
-
-
-def init_global_processor(server_args: ServerArgs):
-    global global_processor
-    transformers.logging.set_verbosity_error()
-    global_processor = get_processor(
-        server_args.tokenizer_path,
-        tokenizer_mode=server_args.tokenizer_mode,
-        trust_remote_code=server_args.trust_remote_code,
-    )
-
-
-def get_pixel_values(
-    image_data, image_aspect_ratio=None, image_grid_pinpoints=None, processor=None
-):
-    try:
-        processor = processor or global_processor
-        image = load_image(image_data)
-        image_hash = hash(image_data)
-        if image_aspect_ratio == "pad":
-            image = expand2square(
-                image, tuple(int(x * 255) for x in processor.image_processor.image_mean)
-            )
-            pixel_values = processor.image_processor(image)["pixel_values"][0]
-        elif image_aspect_ratio == "anyres":
-            pixel_values = process_anyres_image(
-                image, processor.image_processor, image_grid_pinpoints
-            )
-        else:
-            pixel_values = processor.image_processor(image)["pixel_values"][0]
-        pixel_values = pixel_values.astype(np.float16)
-        return pixel_values, image_hash, image.size
-    except Exception:
-        print("Exception in TokenizerManager:\n" + get_exception_traceback())
 
 
 class TokenizerManager:
@@ -81,6 +52,7 @@ class TokenizerManager:
         self,
         server_args: ServerArgs,
         port_args: PortArgs,
+        model_overide_args: dict = None,
     ):
         self.server_args = server_args
 
@@ -93,9 +65,10 @@ class TokenizerManager:
 
         self.model_path = server_args.model_path
         self.hf_config = get_config(
-            self.model_path, trust_remote_code=server_args.trust_remote_code
+            self.model_path,
+            trust_remote_code=server_args.trust_remote_code,
+            model_overide_args=model_overide_args,
         )
-
         self.context_len = get_context_length(self.hf_config)
 
         if is_multimodal_model(self.model_path):
@@ -119,7 +92,7 @@ class TokenizerManager:
             )
 
         self.to_create_loop = True
-        self.rid_to_state = {}  # Dict[str -> ReqState]
+        self.rid_to_state: Dict[str, ReqState] = {}
 
     async def get_pixel_values(self, image_data):
         aspect_ratio = getattr(self.hf_config, "image_aspect_ratio", None)
@@ -140,15 +113,26 @@ class TokenizerManager:
                 image_data, aspect_ratio, grid_pinpoints, self.processor
             )
 
-    async def generate_request(self, obj: GenerateReqInput):
+    async def generate_request(self, obj: GenerateReqInput, request=None):
         if self.to_create_loop:
-            await self.create_handle_loop()
-
-        is_single = isinstance(obj.text, str)
+            self.create_handle_loop()
 
+        obj.post_init()
+        is_single = obj.is_single
         if is_single:
             rid = obj.rid
-            input_ids = self.tokenizer.encode(obj.text)
+
+            if obj.input_ids is None:
+                input_ids = self.tokenizer.encode(obj.text)
+            else:
+                input_ids = obj.input_ids
+
+            if len(input_ids) >= self.context_len:
+                raise ValueError(
+                    f"The input ({len(input_ids)} tokens) is longer than the "
+                    f"model's context length ({self.context_len} tokens)."
+                )
+
             sampling_params = SamplingParams(**obj.sampling_params)
             if sampling_params.max_new_tokens != 0:
                 sampling_params.normalize(self.tokenizer)
@@ -174,29 +158,64 @@ class TokenizerManager:
                 sampling_params=sampling_params,
                 return_logprob=obj.return_logprob,
                 logprob_start_len=obj.logprob_start_len,
+                top_logprobs_num=obj.top_logprobs_num,
                 stream=obj.stream,
             )
             self.send_to_router.send_pyobj(tokenized_obj)
 
-            lock = asyncio.Lock()
             event = asyncio.Event()
-            state = ReqState([], False, event, lock)
+            state = ReqState([], False, event)
             self.rid_to_state[rid] = state
 
             while True:
-                await event.wait()
-                yield state.out_list[-1]
+                try:
+                    await asyncio.wait_for(event.wait(), timeout=4)
+                except asyncio.TimeoutError:
+                    if request is not None and await request.is_disconnected():
+                        self.abort_request(rid)
+                        raise ValueError(f"Abort request {rid}")
+                    continue
+
+                out = self.convert_logprob_style(
+                    state.out_list[-1],
+                    obj.return_logprob,
+                    obj.top_logprobs_num,
+                    obj.return_text_in_logprobs,
+                )
+
+                if self.server_args.log_requests and state.finished:
+                    logger.info(f"in={obj.text}, out={out}")
+
                 state.out_list = []
                 if state.finished:
                     del self.rid_to_state[rid]
+
+                    yield out
+
                     break
+
                 event.clear()
+
+                yield out
         else:
-            assert obj.stream is False
-            bs = len(obj.text)
+            if obj.stream:
+                raise ValueError("Do not support stream for batch mode.")
+
+            if obj.input_ids is None:
+                bs = len(obj.text)
+            else:
+                bs = len(obj.input_ids)
+
             for i in range(bs):
                 rid = obj.rid[i]
-                input_ids = self.tokenizer.encode(obj.text[i])
+
+                if obj.input_ids is None:
+                    input_text = obj.text[i]
+                    input_ids = self.tokenizer.encode(obj.text[i])
+                else:
+                    input_text = None
+                    input_ids = obj.input_ids[i]
+
                 sampling_params = SamplingParams(**obj.sampling_params[i])
                 if sampling_params.max_new_tokens != 0:
                     sampling_params.normalize(self.tokenizer)
@@ -209,7 +228,7 @@ class TokenizerManager:
                     )
                 tokenized_obj = TokenizedGenerateReqInput(
                     rid=rid,
-                    input_text=obj.text[i],
+                    input_text=input_text,
                     input_ids=input_ids,
                     pixel_values=pixel_values,
                     image_hash=image_hash,
@@ -217,53 +236,176 @@ class TokenizerManager:
                     sampling_params=sampling_params,
                     return_logprob=obj.return_logprob[i],
                     logprob_start_len=obj.logprob_start_len[i],
+                    top_logprobs_num=obj.top_logprobs_num[i],
                     stream=obj.stream,
                 )
                 self.send_to_router.send_pyobj(tokenized_obj)
 
-                lock = asyncio.Lock()
                 event = asyncio.Event()
-                state = ReqState([], False, event, lock)
+                state = ReqState([], False, event)
                 self.rid_to_state[rid] = state
 
             output_list = []
             for i in range(bs):
                 rid = obj.rid[i]
                 state = self.rid_to_state[rid]
-                await state.event.wait()
-                output_list.append(state.out_list[-1])
+
+                while True:
+                    try:
+                        await asyncio.wait_for(state.event.wait(), timeout=4)
+                        break
+                    except asyncio.TimeoutError:
+                        if request is not None and await request.is_disconnected():
+                            for rid in obj.rid:
+                                self.abort_request(rid)
+                            raise ValueError(f"Abort request {rid}")
+                        continue
+
+                output_list.append(
+                    self.convert_logprob_style(
+                        state.out_list[-1],
+                        obj.return_logprob[i],
+                        obj.top_logprobs_num[i],
+                        obj.return_text_in_logprobs,
+                    )
+                )
                 assert state.finished
                 del self.rid_to_state[rid]
 
             yield output_list
 
-    async def detokenize(self, obj: DetokenizeReqInput):
-        token_texts = self.tokenizer.convert_ids_to_tokens(obj.input_ids)
-        return [t.decode() if isinstance(t, bytes) else t for t in token_texts]
+    def flush_cache(self):
+        req = FlushCacheReq()
+        self.send_to_router.send_pyobj(req)
+
+    def abort_request(self, rid):
+        if rid not in self.rid_to_state:
+            return
+        del self.rid_to_state[rid]
+        req = AbortReq(rid)
+        self.send_to_router.send_pyobj(req)
+
+    def create_abort_task(self, obj: GenerateReqInput):
+        # Abort the request if the client is disconnected.
+        async def abort_request():
+            await asyncio.sleep(3)
+            if obj.is_single:
+                self.abort_request(obj.rid)
+            else:
+                for rid in obj.rids:
+                    self.abort_request(rid)
 
-    async def flush_cache(self):
-        flush_cache_req = FlushCacheReq()
-        self.send_to_router.send_pyobj(flush_cache_req)
+        background_tasks = BackgroundTasks()
+        background_tasks.add_task(abort_request)
+        return background_tasks
 
-    async def create_handle_loop(self):
+    def create_handle_loop(self):
         self.to_create_loop = False
         loop = asyncio.get_event_loop()
         loop.create_task(self.handle_loop())
 
     async def handle_loop(self):
         while True:
-            recv_obj = await self.recv_from_detokenizer.recv_pyobj()
-
-            if isinstance(recv_obj, BatchStrOut):
-                for i, rid in enumerate(recv_obj.rids):
-                    recv_obj.meta_info[i]["id"] = rid
-                    out_dict = {
-                        "text": recv_obj.output_str[i],
-                        "meta_info": recv_obj.meta_info[i],
-                    }
-                    state = self.rid_to_state[rid]
-                    state.out_list.append(out_dict)
-                    state.finished = recv_obj.finished[i]
-                    state.event.set()
+            recv_obj: BatchTokenIDOut = await self.recv_from_detokenizer.recv_pyobj()
+            assert isinstance(recv_obj, BatchStrOut)
+
+            for i, rid in enumerate(recv_obj.rids):
+                state = self.rid_to_state.get(rid, None)
+                if state is None:
+                    continue
+
+                recv_obj.meta_info[i]["id"] = rid
+                out_dict = {
+                    "text": recv_obj.output_strs[i],
+                    "meta_info": recv_obj.meta_info[i],
+                }
+                state.out_list.append(out_dict)
+                state.finished = recv_obj.finished_reason[i] is not None
+                state.event.set()
+
+    def convert_logprob_style(
+        self, ret, return_logprob, top_logprobs_num, return_text_in_logprobs
+    ):
+        if return_logprob:
+            ret["meta_info"]["prefill_token_logprobs"] = self.detokenize_logprob_tokens(
+                ret["meta_info"]["prefill_token_logprobs"], return_text_in_logprobs
+            )
+            ret["meta_info"]["decode_token_logprobs"] = self.detokenize_logprob_tokens(
+                ret["meta_info"]["decode_token_logprobs"], return_text_in_logprobs
+            )
+
+            if top_logprobs_num > 0:
+                ret["meta_info"][
+                    "prefill_top_logprobs"
+                ] = self.detokenize_top_logprobs_tokens(
+                    ret["meta_info"]["prefill_top_logprobs"], return_text_in_logprobs
+                )
+                ret["meta_info"][
+                    "decode_top_logprobs"
+                ] = self.detokenize_top_logprobs_tokens(
+                    ret["meta_info"]["decode_top_logprobs"], return_text_in_logprobs
+                )
+        return ret
+
+    def detokenize_logprob_tokens(self, token_logprobs, decode_to_text):
+        if not decode_to_text:
+            return [(logprob, token_id, None) for logprob, token_id in token_logprobs]
+
+        token_ids = [tid for _, tid in token_logprobs]
+        token_texts = self.tokenizer.batch_decode(token_ids)
+        return [
+            (logprob, token_id, token_text)
+            for (logprob, token_id), token_text, in zip(token_logprobs, token_texts)
+        ]
+
+    def detokenize_top_logprobs_tokens(self, top_logprobs, decode_to_text):
+        for i, t in enumerate(top_logprobs):
+            if t:
+                top_logprobs[i] = self.detokenize_logprob_tokens(t, decode_to_text)
+        return top_logprobs
+
+
+global global_processor
+
+
+def init_global_processor(server_args: ServerArgs):
+    global global_processor
+    transformers.logging.set_verbosity_error()
+    global_processor = get_processor(
+        server_args.tokenizer_path,
+        tokenizer_mode=server_args.tokenizer_mode,
+        trust_remote_code=server_args.trust_remote_code,
+    )
+
+
+def get_pixel_values(
+    image_data, image_aspect_ratio=None, image_grid_pinpoints=None, processor=None
+):
+    try:
+        processor = processor or global_processor
+        image, image_size = load_image(image_data)
+        if image_size is not None:
+            image_hash = hash(image_data)
+            pixel_values = processor.image_processor(image)["pixel_values"]
+            for _ in range(len(pixel_values)):
+                pixel_values[_] = pixel_values[_].astype(np.float16)
+            pixel_values = np.stack(pixel_values, axis=0)
+            return pixel_values, image_hash, image_size
+        else:
+            image_hash = hash(image_data)
+            if image_aspect_ratio == "pad":
+                image = expand2square(
+                    image,
+                    tuple(int(x * 255) for x in processor.image_processor.image_mean),
+                )
+                pixel_values = processor.image_processor(image)["pixel_values"][0]
+            elif image_aspect_ratio == "anyres":
+                pixel_values = process_anyres_image(
+                    image, processor.image_processor, image_grid_pinpoints
+                )
             else:
-                raise ValueError(f"Invalid object: {recv_obj}")
+                pixel_values = processor.image_processor(image)["pixel_values"][0]
+            pixel_values = pixel_values.astype(np.float16)
+            return pixel_values, image_hash, image.size
+    except Exception:
+        print("Exception in TokenizerManager:\n" + get_exception_traceback())