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

python/sglang/srt/managers/tokenizer_manager.py

@@ -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):