sglang 0.1.15__py3-none-any.whl → 0.1.17__py3-none-any.whl

sglang/srt/managers/tokenizer_manager.py

@@ -4,13 +4,14 @@ import dataclasses
 import logging
 import multiprocessing as mp
 import os
-from typing import List
+from typing import List, Dict
 
 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,
@@ -19,16 +20,18 @@ from sglang.srt.hf_transformers_utils import (
     get_tokenizer,
 )
 from sglang.srt.managers.io_struct import (
+    AbortReq,
     BatchStrOut,
-    DetokenizeReqInput,
     FlushCacheReq,
     GenerateReqInput,
     TokenizedGenerateReqInput,
 )
+from sglang.srt.managers.io_struct import BatchTokenIDOut
 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())
 
@@ -42,48 +45,12 @@ class ReqState:
     event: asyncio.Event
 
 
-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:
     def __init__(
         self,
         server_args: ServerArgs,
         port_args: PortArgs,
+        model_overide_args: dict = None,
     ):
         self.server_args = server_args
 
@@ -96,9 +63,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):
@@ -122,7 +90,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)
@@ -143,15 +111,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)
@@ -187,27 +166,54 @@ class TokenizerManager:
             self.rid_to_state[rid] = state
 
             while True:
-                await event.wait()
-                out = self.convert_logprob_style(state.out_list[-1],
-                                                 obj.return_logprob,
-                                                 obj.top_logprobs_num,
-                                                 obj.return_text_in_logprobs)
+                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}")
 
-                yield 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)
@@ -220,7 +226,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,
@@ -241,45 +247,84 @@ class TokenizerManager:
             for i in range(bs):
                 rid = obj.rid[i]
                 state = self.rid_to_state[rid]
-                await state.event.wait()
+
+                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))
+                    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 flush_cache(self):
-        flush_cache_req = FlushCacheReq()
-        self.send_to_router.send_pyobj(flush_cache_req)
+    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):
+        # 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 create_handle_loop(self):
+        background_tasks = BackgroundTasks()
+        background_tasks.add_task(abort_request)
+        return background_tasks
+
+    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()
-            else:
-                raise ValueError(f"Invalid object: {recv_obj}")
-
-    def convert_logprob_style(self, ret, return_logprob, top_logprobs_num, return_text_in_logprobs):
+            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_str[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
@@ -288,11 +333,15 @@ class TokenizerManager:
                 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"]["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
+            ret["meta_info"]["decode_top_logprobs"] = (
+                self.detokenize_top_logprobs_tokens(
+                    ret["meta_info"]["decode_top_logprobs"], return_text_in_logprobs
+                )
             )
         return ret
 
@@ -312,3 +361,49 @@ class TokenizerManager:
             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 != 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:
+                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())

sglang/srt/model_config.py

@@ -10,11 +10,14 @@ class ModelConfig:
         trust_remote_code: bool = True,
         revision: Optional[str] = None,
         context_length: Optional[int] = None,
+        model_overide_args: Optional[dict] = None,
     ) -> None:
         self.path = path
         self.trust_remote_code = trust_remote_code
         self.revision = revision
-        self.hf_config = get_config(self.path, trust_remote_code, revision)
+        self.model_overide_args = model_overide_args
+        self.hf_config = get_config(self.path, trust_remote_code, revision,
+                                    model_overide_args=model_overide_args)
 
         if context_length is not None:
             self.context_len = context_length
@@ -40,4 +43,4 @@ class ModelConfig:
             self.num_key_value_heads = self.num_attention_heads
         self.hidden_size = self.hf_config.hidden_size
         self.num_hidden_layers = self.hf_config.num_hidden_layers
-        self.vocab_size = self.hf_config.vocab_size
+        self.vocab_size = self.hf_config.vocab_size

sglang/srt/models/commandr.py

@@ -18,38 +18,38 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+# Adapted from
+# https://github.com/vllm-project/vllm/blob/c7f2cf2b7f67bce5842fedfdba508440fe257375/vllm/model_executor/models/commandr.py#L1
+
 # This file is based on the LLama model definition file in transformers
 """PyTorch Cohere model."""
-from typing import Optional, Tuple
+from typing import Optional, Tuple, Iterable
 
 import torch
 import torch.utils.checkpoint
 from torch import nn
 from torch.nn.parameter import Parameter
 from transformers import PretrainedConfig
+from vllm.config import CacheConfig
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+)
 from vllm.model_executor.layers.activation import SiluAndMul
 from vllm.model_executor.layers.linear import (
     MergedColumnParallelLinear,
     QKVParallelLinear,
     RowParallelLinear,
 )
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig)
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
-from vllm.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-)
 from vllm.model_executor.utils import set_weight_attrs
-from sglang.srt.weight_utils import (
-    default_weight_loader,
-    hf_model_weights_iterator,
-)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.radix_attention import RadixAttention
-from sglang.srt.managers.router.model_runner import InputMetadata
+from sglang.srt.managers.controller.model_runner import InputMetadata
 
 
 @torch.compile
@@ -305,6 +305,7 @@ class CohereForCausalLM(nn.Module):
         self,
         config: PretrainedConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        cache_config: Optional[CacheConfig] = None,
     ) -> None:
         super().__init__()
         self.config = config
@@ -328,13 +329,7 @@ class CohereForCausalLM(nn.Module):
             input_ids, hidden_states, self.model.embed_tokens.weight, input_metadata
         )
 
-    def load_weights(
-        self,
-        model_name_or_path: str,
-        cache_dir: Optional[str] = None,
-        load_format: str = "auto",
-        revision: Optional[str] = None,
-    ):
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         stacked_params_mapping = [
             # (param_name, shard_name, shard_id)
             ("qkv_proj", "q_proj", "q"),
@@ -345,9 +340,7 @@ class CohereForCausalLM(nn.Module):
         ]
         params_dict = dict(self.named_parameters())
         loaded_params = set()
-        for name, loaded_weight in hf_model_weights_iterator(
-            model_name_or_path, cache_dir, load_format, revision
-        ):
+        for name, loaded_weight in weights:
             for param_name, shard_name, shard_id in stacked_params_mapping:
                 if shard_name not in name:
                     continue

sglang/srt/models/dbrx.py

@@ -1,41 +1,36 @@
 # Adapted from:
-# https://github.com/vllm-project/vllm/blob/14ccd94c89d0ffd9da283545d93ab1dfea5da340/vllm/model_executor/models/dbrx.py
+# https://github.com/vllm-project/vllm/blob/c7f2cf2b7f67bce5842fedfdba508440fe257375/vllm/model_executor/models/dbrx.py#L1
 # coding=utf-8
-from typing import Optional
+from typing import Iterable, Optional, Tuple
 
 import torch
 import torch.nn as nn
+from vllm.config import CacheConfig
+from vllm.distributed import (
+    get_tensor_model_parallel_rank,
+    get_tensor_model_parallel_world_size,
+    tensor_model_parallel_all_reduce,
+)
 from vllm.model_executor.layers.fused_moe import fused_moe
 from vllm.model_executor.layers.linear import (
     QKVParallelLinear,
     ReplicatedLinear,
     RowParallelLinear,
 )
-from vllm.model_executor.layers.quantization.base_config import (
-    QuantizationConfig)
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 from vllm.model_executor.layers.rotary_embedding import get_rope
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     DEFAULT_VOCAB_PADDING_SIZE,
     ParallelLMHead,
     VocabParallelEmbedding,
 )
-from vllm.distributed import (
-    tensor_model_parallel_all_reduce,
-)
-from vllm.distributed import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-)
 from vllm.model_executor.utils import set_weight_attrs
-from sglang.srt.weight_utils import (
-    default_weight_loader,
-    hf_model_weights_iterator,
-)
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.transformers_utils.configs.dbrx import DbrxConfig
 
 from sglang.srt.layers.logits_processor import LogitsProcessor
 from sglang.srt.layers.radix_attention import RadixAttention
-from sglang.srt.managers.router.model_runner import InputMetadata
-from sglang.srt.models.dbrx_config import DbrxConfig
+from sglang.srt.managers.controller.model_runner import InputMetadata
 
 
 class DbrxRouter(nn.Module):
@@ -291,7 +286,9 @@ class DbrxBlock(nn.Module):
         quant_config: Optional[QuantizationConfig] = None,
     ):
         super().__init__()
-        self.norm_attn_norm = DbrxFusedNormAttention(config, layer_id, quant_config=quant_config)
+        self.norm_attn_norm = DbrxFusedNormAttention(
+            config, layer_id, quant_config=quant_config
+        )
         self.ffn = DbrxExperts(config, quant_config=quant_config)
 
     def forward(
@@ -322,7 +319,10 @@ class DbrxModel(nn.Module):
             config.d_model,
         )
         self.blocks = nn.ModuleList(
-            [DbrxBlock(config, i, quant_config=quant_config) for i in range(config.n_layers)]
+            [
+                DbrxBlock(config, i, quant_config=quant_config)
+                for i in range(config.n_layers)
+            ]
         )
         self.norm_f = nn.LayerNorm(config.d_model, eps=1e-5)
         for module in self.modules():
@@ -353,6 +353,7 @@ class DbrxForCausalLM(nn.Module):
         self,
         config: DbrxConfig,
         quant_config: Optional[QuantizationConfig] = None,
+        cache_config: Optional[CacheConfig] = None,
     ):
         super().__init__()
         self.config = config
@@ -378,13 +379,7 @@ class DbrxForCausalLM(nn.Module):
             input_ids, hidden_states, self.lm_head.weight, input_metadata
         )
 
-    def load_weights(
-        self,
-        model_name_or_path: str,
-        cache_dir: Optional[str] = None,
-        load_format: str = "auto",
-        revision: Optional[str] = None,
-    ):
+    def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
         expert_params_mapping = [
             (
                 "ws" if weight_name in ["w1", "v1"] else "w2s",
@@ -393,9 +388,7 @@ class DbrxForCausalLM(nn.Module):
             for weight_name in ["w1", "v1", "w2"]
         ]
         params_dict = dict(self.named_parameters(remove_duplicate=False))
-        for name, loaded_weight in hf_model_weights_iterator(
-            model_name_or_path, cache_dir, load_format, revision
-        ):
+        for name, loaded_weight in weights:
             for param_name, weight_name in expert_params_mapping:
                 if weight_name not in name:
                     continue