Add a new api configure_logging to allow dumping the requests (#2875)

This commit is contained in:
Lianmin Zheng
2025-01-13 14:24:00 -08:00
committed by GitHub
parent 923f518337
commit 46d4431889
13 changed files with 164 additions and 71 deletions

View File

@@ -0,0 +1,43 @@
"""
Copyright 2023-2025 SGLang Team
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
"""
Configure the logging settings of a server.
Usage:
python3 -m sglang.srt.managers.configure_logging --url http://localhost:30000
"""
import argparse
import requests
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--url", type=str, default="http://localhost:30000")
parser.add_argument(
"--dump-requests-folder", type=str, default="/tmp/sglang_request_dump"
)
parser.add_argument("--dump-requests-threshold", type=int, default=1000)
args = parser.parse_args()
response = requests.post(
args.url + "/configure_logging",
json={
"dump_requests_folder": args.dump_requests_folder,
"dump_requests_threshold": args.dump_requests_threshold,
},
)
assert response.status_code == 200

View File

@@ -488,6 +488,13 @@ class ProfileReq(Enum):
STOP_PROFILE = 2
@dataclass
class ConfigureLoggingReq:
log_requests: Optional[bool] = None
dump_requests_folder: Optional[str] = None
dump_requests_threshold: Optional[int] = None
@dataclass
class OpenSessionReqInput:
capacity_of_str_len: int

View File

@@ -82,6 +82,7 @@ from sglang.srt.metrics.collector import SchedulerMetricsCollector, SchedulerSta
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.server_args import PortArgs, ServerArgs
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.torch_memory_saver_adapter import TorchMemorySaverAdapter
from sglang.srt.utils import (
broadcast_pyobj,
configure_logger,
@@ -92,7 +93,6 @@ from sglang.srt.utils import (
set_random_seed,
suppress_other_loggers,
)
from sglang.torch_memory_saver_adapter import TorchMemorySaverAdapter
from sglang.utils import get_exception_traceback
logger = logging.getLogger(__name__)

View File

@@ -18,10 +18,12 @@ import copy
import dataclasses
import logging
import os
import pickle
import signal
import sys
import time
import uuid
from datetime import datetime
from typing import Any, Awaitable, Dict, Generic, List, Optional, Tuple, TypeVar, Union
import fastapi
@@ -43,6 +45,7 @@ from sglang.srt.managers.io_struct import (
BatchStrOut,
BatchTokenIDOut,
CloseSessionReqInput,
ConfigureLoggingReq,
EmbeddingReqInput,
FlushCacheReq,
GenerateReqInput,
@@ -109,6 +112,7 @@ class TokenizerManager:
# Parse args
self.server_args = server_args
self.enable_metrics = server_args.enable_metrics
self.log_requests = server_args.log_requests
# Init inter-process communication
context = zmq.asyncio.Context(2)
@@ -167,6 +171,9 @@ class TokenizerManager:
# Store states
self.to_create_loop = True
self.rid_to_state: Dict[str, ReqState] = {}
self.dump_requests_folder = "" # By default do not dump
self.dump_requests_threshold = 1000
self.dump_request_list: List[Tuple] = []
# The event to notify the weight sync is finished.
self.model_update_lock = RWLock()
@@ -225,7 +232,7 @@ class TokenizerManager:
obj.normalize_batch_and_arguments()
if self.server_args.log_requests:
if self.log_requests:
logger.info(f"Receive: obj={dataclass_to_string_truncated(obj)}")
async with self.model_update_lock.reader_lock:
@@ -346,7 +353,7 @@ class TokenizerManager:
state.out_list = []
if state.finished:
if self.server_args.log_requests:
if self.log_requests:
msg = f"Finish: obj={dataclass_to_string_truncated(obj)}, out={dataclass_to_string_truncated(out)}"
logger.info(msg)
del self.rid_to_state[obj.rid]
@@ -597,6 +604,15 @@ class TokenizerManager:
assert not self.to_create_loop, "close session should not be the first request"
await self.send_to_scheduler.send_pyobj(obj)
def configure_logging(self, obj: ConfigureLoggingReq):
if obj.log_requests is not None:
self.log_requests = obj.log_requests
if obj.dump_requests_folder is not None:
self.dump_requests_folder = obj.dump_requests_folder
if obj.dump_requests_threshold is not None:
self.dump_requests_threshold = obj.dump_requests_threshold
logging.info(f"Config logging: {obj=}")
def create_abort_task(self, obj: GenerateReqInput):
# Abort the request if the client is disconnected.
async def abort_request():
@@ -708,6 +724,8 @@ class TokenizerManager:
if self.enable_metrics:
self.collect_metrics(state, recv_obj, i)
if self.dump_requests_folder and state.finished:
self.dump_requests(state, out_dict)
elif isinstance(recv_obj, OpenSessionReqOutput):
self.session_futures[recv_obj.session_id].set_result(
recv_obj.session_id if recv_obj.success else None
@@ -850,6 +868,25 @@ class TokenizerManager:
(time.time() - state.created_time) / completion_tokens
)
def dump_requests(self, state: ReqState, out_dict: dict):
self.dump_request_list.append(
(state.obj, out_dict, state.created_time, time.time())
)
if len(self.dump_request_list) >= self.dump_requests_threshold:
to_dump = self.dump_request_list
self.dump_request_list = []
def background_task():
os.makedirs(self.dump_requests_folder, exist_ok=True)
current_time = datetime.now()
filename = current_time.strftime("%Y-%m-%d_%H-%M-%S") + ".pkl"
with open(os.path.join(self.dump_requests_folder, filename), "wb") as f:
pickle.dump(to_dump, f)
# Schedule the task to run in the background without awaiting it
asyncio.create_task(asyncio.to_thread(background_task))
class SignalHandler:
def __init__(self, tokenizer_manager):