chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Minor] Improve the function organization in TokenizerManager & improve loggers (#1208)

Browse Source
This commit is contained in:
Lianmin Zheng
2024-08-25 14:46:34 -07:00
committed by GitHub
parent 30b4f771b0
commit 902278008a
12 changed files with 137 additions and 134 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/controller_multi.py
+2 -5
View File
@@ -35,7 +35,7 @@ from sglang.srt.managers.io_struct import (
TokenizedGenerateReqInput,
)
from sglang.srt.server_args import PortArgs, ServerArgs
from sglang.srt.utils import kill_parent_process
from sglang.srt.utils import configure_logger, kill_parent_process
from sglang.utils import get_exception_traceback
logger = logging.getLogger(__name__)
@@ -193,10 +193,7 @@ def start_controller_process(
):
"""Start a controller process."""
logging.basicConfig(
level=getattr(logging, server_args.log_level.upper()),
format="%(message)s",
)
configure_logger(server_args)
try:
controller = ControllerMulti(server_args, port_args, model_overide_args)
python/sglang/srt/managers/controller_single.py
+7 -7
View File
@@ -27,7 +27,7 @@ from sglang.srt.managers.tp_worker import (
launch_tp_servers,
)
from sglang.srt.server_args import PortArgs, ServerArgs
from sglang.srt.utils import kill_parent_process
from sglang.srt.utils import configure_logger, kill_parent_process
from sglang.utils import get_exception_traceback
logger = logging.getLogger(__name__)
@@ -52,7 +52,7 @@ class ControllerSingle:
self.dp_worker_id = dp_worker_id
self.mp_queue = mp_queue
# Init communication
# Init inter-process communication
context = zmq.Context(2)
if not self.is_dp_worker:
@@ -133,11 +133,11 @@ def start_controller_process(
queue: multiprocessing.connection.Connection = None,
):
"""Start a controller process."""
logging.basicConfig(
level=getattr(logging, server_args.log_level.upper()),
format="%(message)s",
)
if is_data_parallel_worker:
logger_prefix = f" DP{dp_worker_id} TP0"
else:
logger_prefix = " TP0"
configure_logger(server_args, prefix=logger_prefix)
if not is_data_parallel_worker:
tp_size_local = server_args.tp_size // server_args.nnodes
python/sglang/srt/managers/detokenizer_manager.py
+12 -8
View File
@@ -56,6 +56,7 @@ class DetokenizerManager:
server_args: ServerArgs,
port_args: PortArgs,
):
# Init inter-process communication
context = zmq.asyncio.Context(2)
self.recv_from_router = context.socket(zmq.PULL)
self.recv_from_router.bind(f"tcp://127.0.0.1:{port_args.detokenizer_port}")
@@ -75,10 +76,13 @@ class DetokenizerManager:
self.decode_status = {}
async def handle_loop(self):
"""The event loop that handles requests"""
while True:
recv_obj: BatchTokenIDOut = await self.recv_from_router.recv_pyobj()
recv_obj = await self.recv_from_router.recv_pyobj()
if isinstance(recv_obj, BatchEmbeddingOut):
# If it is embedding model, no detokenization is needed.
self.send_to_tokenizer.send_pyobj(
BatchEmbeddingOut(
rids=recv_obj.rids,
@@ -88,19 +92,18 @@ class DetokenizerManager:
)
)
continue
if isinstance(recv_obj, UpdateWeightReqOutput):
elif isinstance(recv_obj, UpdateWeightReqOutput):
# If it is a weight update request, no detokenization is needed.
self.send_to_tokenizer.send_pyobj(recv_obj)
continue
elif self.tokenizer is None:
# If the tokenizer is skipped, no detokenization is needed
self.send_to_tokenizer.send_pyobj(recv_obj)
continue
assert isinstance(recv_obj, BatchTokenIDOut)
bs = len(recv_obj.rids)
if self.tokenizer is None:
# Send BatchTokenIDOut if no tokenizer init'ed.
self.send_to_tokenizer.send_pyobj(recv_obj)
continue
# Initialize decode status
read_ids, surr_ids = [], []
for i in range(bs):
@@ -134,6 +137,7 @@ class DetokenizerManager:
spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
)
# Incremental decoding
output_strs = []
for i in range(bs):
s = self.decode_status[recv_obj.rids[i]]
python/sglang/srt/managers/tokenizer_manager.py
+81 -75
View File
@@ -21,7 +21,7 @@ import dataclasses
import logging
import multiprocessing as mp
import os
from typing import Dict, List, Tuple, Union
from typing import Dict, List, Optional, Tuple, Union
import numpy as np
import transformers
@@ -80,6 +80,7 @@ class TokenizerManager:
):
self.server_args = server_args
# Init inter-process communication
context = zmq.asyncio.Context(2)
self.recv_from_detokenizer = context.socket(zmq.PULL)
self.recv_from_detokenizer.bind(f"tcp://127.0.0.1:{port_args.tokenizer_port}")
@@ -87,6 +88,7 @@ class TokenizerManager:
self.send_to_router = context.socket(zmq.PUSH)
self.send_to_router.connect(f"tcp://127.0.0.1:{port_args.controller_port}")
# Read model args
self.model_path = server_args.model_path
self.served_model_name = server_args.served_model_name
self.hf_config = get_config(
@@ -104,6 +106,7 @@ class TokenizerManager:
else:
self.context_len = get_context_length(self.hf_config)
# Create tokenizer
if server_args.skip_tokenizer_init:
self.tokenizer = self.processor = None
else:
@@ -127,6 +130,7 @@ class TokenizerManager:
trust_remote_code=server_args.trust_remote_code,
)
# Store states
self.to_create_loop = True
self.rid_to_state: Dict[str, ReqState] = {}
@@ -134,63 +138,6 @@ class TokenizerManager:
self.model_update_lock = asyncio.Lock()
self.model_update_result = None
async def get_pixel_values(self, image_data, aspect_ratio=None):
aspect_ratio = (
getattr(self.hf_config, "image_aspect_ratio", None)
if aspect_ratio is None
else aspect_ratio
)
grid_pinpoints = (
self.hf_config.image_grid_pinpoints
if hasattr(self.hf_config, "image_grid_pinpoints")
and "anyres" in aspect_ratio
else None
)
if isinstance(image_data, list) and len(image_data) > 0:
pixel_values, image_hash, image_size = [], [], []
if len(image_data) > 1:
aspect_ratio = "pad" # LLaVA OneVision Handling: more than one image --> interleaved image mode or video mode. We do not use anyres
for img_data in image_data:
pixel_v, image_h, image_s = await self._process_single_image(
img_data, aspect_ratio, grid_pinpoints
)
pixel_values.append(pixel_v)
image_hash.append(image_h)
image_size.append(image_s)
pixel_values = np.stack(pixel_values, axis=0)
else:
pixel_values, image_hash, image_size = await self._process_single_image(
image_data[0], aspect_ratio, grid_pinpoints
)
image_hash = [image_hash]
image_size = [image_size]
elif isinstance(image_data, str):
pixel_values, image_hash, image_size = await self._process_single_image(
image_data, aspect_ratio, grid_pinpoints
)
image_hash = [image_hash]
image_size = [image_size]
else:
pixel_values, image_hash, image_size = None, None, None
return pixel_values, image_hash, image_size
async def _process_single_image(self, image_data, aspect_ratio, grid_pinpoints):
if self.executor is not None:
loop = asyncio.get_event_loop()
return await loop.run_in_executor(
self.executor,
get_pixel_values,
image_data,
aspect_ratio,
grid_pinpoints,
)
else:
return get_pixel_values(
image_data, aspect_ratio, grid_pinpoints, self.processor
)
async def generate_request(
self, obj: Union[GenerateReqInput, EmbeddingReqInput], request=None
):
@@ -198,7 +145,7 @@ class TokenizerManager:
self.create_handle_loop()
while self.model_update_lock.locked():
await asyncio.sleep(0)
await asyncio.sleep(0.001)
obj.post_init()
is_single = obj.is_single
@@ -214,8 +161,8 @@ class TokenizerManager:
self,
obj: Union[GenerateReqInput, EmbeddingReqInput],
request,
index=None,
is_cache_for_prefill=False,
index: Optional[int] = None,
is_cache_for_prefill: Optional[bool] = False,
):
if not is_cache_for_prefill: # The normal case with a single prompt
not_use_index = index is None
@@ -235,7 +182,7 @@ class TokenizerManager:
)
if self.is_generation:
pixel_values, image_hash, image_size = await self.get_pixel_values(
pixel_values, image_hash, image_size = await self._get_pixel_values(
obj.image_data
)
return_logprob = (
@@ -345,7 +292,7 @@ class TokenizerManager:
parallel_sample_num = obj.parallel_sample_num
if parallel_sample_num != 1:
# Send prefill requests to cache the common input
# Send prefill requests to cache the common prefix
parallel_sample_num += 1
input_id_result = [] if obj.input_ids is None else None
for i in range(batch_size):
@@ -436,7 +383,6 @@ class TokenizerManager:
)
# Then process the responses based on streaming option
is_stream = hasattr(obj, "stream") and obj.stream
tasks = [asyncio.create_task(gen.__anext__()) for gen in generators]
@@ -482,9 +428,9 @@ class TokenizerManager:
async def _get_pixel_values(self, image_data):
if isinstance(image_data, list) and len(image_data) > 0:
return await self.get_pixel_values(image_data[0])
return await self._get_pixel_values_internal(image_data[0])
elif isinstance(image_data, str):
return await self.get_pixel_values(image_data)
return await self._get_pixel_values_internal(image_data)
else:
return None, None, None
@@ -563,6 +509,13 @@ class TokenizerManager:
req = FlushCacheReq()
self.send_to_router.send_pyobj(req)
def abort_request(self, rid: str):
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)
async def update_weights(self, obj: UpdateWeightReqInput, request):
if self.to_create_loop:
self.create_handle_loop()
@@ -587,13 +540,6 @@ class TokenizerManager:
else:
return False, "Another update is in progress. Please try again later."
def abort_request(self, rid: str):
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():
@@ -617,6 +563,8 @@ class TokenizerManager:
loop.create_task(self.handle_loop())
async def handle_loop(self):
"""The event loop that handles requests"""
while True:
recv_obj: Union[
BatchStrOut, BatchEmbeddingOut, BatchTokenIDOut, UpdateWeightReqOutput
@@ -713,11 +661,69 @@ class TokenizerManager:
)
return top_logprobs
async def _get_pixel_values_internal(self, image_data, aspect_ratio=None):
aspect_ratio = (
getattr(self.hf_config, "image_aspect_ratio", None)
if aspect_ratio is None
else aspect_ratio
)
grid_pinpoints = (
self.hf_config.image_grid_pinpoints
if hasattr(self.hf_config, "image_grid_pinpoints")
and "anyres" in aspect_ratio
else None
)
if isinstance(image_data, list) and len(image_data) > 0:
pixel_values, image_hash, image_size = [], [], []
if len(image_data) > 1:
aspect_ratio = "pad" # LLaVA OneVision Handling: more than one image --> interleaved image mode or video mode. We do not use anyres
for img_data in image_data:
pixel_v, image_h, image_s = await self._process_single_image(
img_data, aspect_ratio, grid_pinpoints
)
pixel_values.append(pixel_v)
image_hash.append(image_h)
image_size.append(image_s)
pixel_values = np.stack(pixel_values, axis=0)
else:
pixel_values, image_hash, image_size = await self._process_single_image(
image_data[0], aspect_ratio, grid_pinpoints
)
image_hash = [image_hash]
image_size = [image_size]
elif isinstance(image_data, str):
pixel_values, image_hash, image_size = await self._process_single_image(
image_data, aspect_ratio, grid_pinpoints
)
image_hash = [image_hash]
image_size = [image_size]
else:
pixel_values, image_hash, image_size = None, None, None
return pixel_values, image_hash, image_size
async def _process_single_image(self, image_data, aspect_ratio, grid_pinpoints):
if self.executor is not None:
loop = asyncio.get_event_loop()
return await loop.run_in_executor(
self.executor,
_process_single_image_task,
image_data,
aspect_ratio,
grid_pinpoints,
)
else:
return _process_single_image_task(
image_data, aspect_ratio, grid_pinpoints, self.processor
)
global global_processor
def init_global_processor(server_args: ServerArgs):
"""Init the global processor for multi modal models."""
global global_processor
transformers.logging.set_verbosity_error()
global_processor = get_processor(
@@ -727,7 +733,7 @@ def init_global_processor(server_args: ServerArgs):
)
def get_pixel_values(
def _process_single_image_task(
image_data, image_aspect_ratio=None, image_grid_pinpoints=None, processor=None
):
try:
@@ -759,4 +765,4 @@ def get_pixel_values(
pixel_values = pixel_values.astype(np.float16)
return pixel_values, image_hash, image.size
except Exception:
print("Exception in TokenizerManager:\n" + get_exception_traceback())
logger.error("Exception in TokenizerManager:\n" + get_exception_traceback())
python/sglang/srt/managers/tp_worker.py
+8 -6
View File
@@ -56,6 +56,7 @@ from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import (
configure_logger,
is_multimodal_model,
set_random_seed,
suppress_other_loggers,
@@ -145,7 +146,6 @@ class ModelTpServer:
# Print info
logger.info(
f"[gpu={self.gpu_id}] "
f"max_total_num_tokens={self.max_total_num_tokens}, "
f"max_prefill_tokens={self.max_prefill_tokens}, "
f"max_running_requests={self.max_running_requests}, "
@@ -284,7 +284,7 @@ class ModelTpServer:
self.num_generated_tokens = 0
self.last_stats_tic = time.time()
logger.info(
f"[gpu={self.gpu_id}] Decode batch. "
f"Decode batch. "
f"#running-req: {len(self.running_batch.reqs)}, "
f"#token: {num_used}, "
f"token usage: {num_used / self.max_total_num_tokens:.2f}, "
@@ -443,7 +443,7 @@ class ModelTpServer:
if num_mixed_running > 0:
logger.info(
f"[gpu={self.gpu_id}] Prefill batch"
f"Prefill batch"
f"(mixed #running-req: {num_mixed_running}). "
f"#new-seq: {len(can_run_list)}, "
f"#new-token: {adder.log_input_tokens}, "
@@ -453,7 +453,7 @@ class ModelTpServer:
)
else:
logger.info(
f"[gpu={self.gpu_id}] Prefill batch. "
f"Prefill batch. "
f"#new-seq: {len(can_run_list)}, "
f"#new-token: {adder.log_input_tokens}, "
f"#cached-token: {adder.log_hit_tokens}, "
@@ -631,7 +631,7 @@ class ModelTpServer:
self.new_token_ratio = new_token_ratio
logger.info(
"decode out of memory happened, "
"Decode out of memory happened. "
f"#retracted_reqs: {len(retracted_reqs)}, "
f"#new_token_ratio: {old_ratio:.4f} -> {self.new_token_ratio:.4f}"
)
@@ -848,7 +848,9 @@ def run_tp_server(
nccl_port: int,
model_overide_args: dict,
):
"""Run a tensor parallel server."""
"""Run a tensor parallel model server."""
configure_logger(server_args, prefix=f" TP{tp_rank}")
try:
model_server = ModelTpServer(
gpu_id,