Tiny support whitelisted headers in request logging (#16342)

python/sglang/srt/managers/tokenizer_manager.py

@@ -509,7 +509,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
             self._attach_multi_http_worker_info(obj)
 
         # Log the request
-        self.request_logger.log_received_request(obj, self.tokenizer)
+        self.request_logger.log_received_request(obj, self.tokenizer, request)
 
         async with self.is_pause_cond:
             await self.is_pause_cond.wait_for(lambda: not self.is_pause)
@@ -1214,7 +1214,10 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
                         "response_sent_to_client_ts"
                     ] = state.response_sent_to_client_ts
                 self.request_logger.log_finished_request(
-                    obj, out, is_multimodal_gen=self.model_config.is_multimodal_gen
+                    obj,
+                    out,
+                    is_multimodal_gen=self.model_config.is_multimodal_gen,
+                    request=request,
                 )
 
                 if self.request_metrics_exporter_manager.exporter_enabled():

python/sglang/srt/utils/request_logger.py

@@ -21,7 +21,7 @@ import socket
 from datetime import datetime
 from functools import lru_cache
 from logging.handlers import TimedRotatingFileHandler
-from typing import TYPE_CHECKING, Any, List, Optional, Set, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
 
 import torch.distributed as dist
 
@@ -29,10 +29,22 @@ from sglang.srt.environ import envs
 from sglang.srt.utils.common import get_bool_env_var
 
 if TYPE_CHECKING:
+    import fastapi
+
     from sglang.srt.managers.io_struct import EmbeddingReqInput, GenerateReqInput
 
 logger = logging.getLogger(__name__)
 
+WHITELISTED_HEADERS = ["x-smg-routing-key"]
+
+
+def _extract_whitelisted_headers(
+    request: Optional["fastapi.Request"],
+) -> Optional[Dict[str, str]]:
+    if request is None:
+        return None
+    return {h: v for h in WHITELISTED_HEADERS if (v := request.headers.get(h))}
+
 
 class RequestLogger:
     def __init__(
@@ -79,21 +91,28 @@ class RequestLogger:
         self.targets = self._setup_targets()
 
     def log_received_request(
-        self, obj: Union["GenerateReqInput", "EmbeddingReqInput"], tokenizer: Any = None
+        self,
+        obj: Union["GenerateReqInput", "EmbeddingReqInput"],
+        tokenizer: Any = None,
+        request: Optional["fastapi.Request"] = None,
     ) -> None:
         if not self.log_requests:
             return
 
         max_length, skip_names, _ = self.metadata
+        headers = _extract_whitelisted_headers(request)
         if self.log_requests_format == "json":
             log_data = {
                 "rid": obj.rid,
                 "obj": _transform_data_for_logging(obj, max_length, skip_names),
             }
+            if headers:
+                log_data["headers"] = headers
             self._log_json("request.received", log_data)
         else:
+            headers_str = f", headers={headers}" if headers else ""
             self._log(
-                f"Receive: obj={_dataclass_to_string_truncated(obj, max_length, skip_names=skip_names)}"
+                f"Receive: obj={_dataclass_to_string_truncated(obj, max_length, skip_names=skip_names)}{headers_str}"
             )
 
         # FIXME: This is a temporary fix to get the text from the input ids.
@@ -112,6 +131,7 @@ class RequestLogger:
         obj: Union["GenerateReqInput", "EmbeddingReqInput"],
         out: Any,
         is_multimodal_gen: bool = False,
+        request: Optional["fastapi.Request"] = None,
     ) -> None:
         if not self.log_requests:
             return
@@ -121,22 +141,30 @@ class RequestLogger:
             return
 
         max_length, skip_names, out_skip_names = self.metadata
+        headers = _extract_whitelisted_headers(request)
         if self.log_requests_format == "json":
             log_data = {
                 "rid": obj.rid,
                 "obj": _transform_data_for_logging(obj, max_length, skip_names),
             }
+            if headers:
+                log_data["headers"] = headers
             if not is_multimodal_gen:
                 log_data["out"] = _transform_data_for_logging(
                     out, max_length, out_skip_names
                 )
             self._log_json("request.finished", log_data)
         else:
-            if is_multimodal_gen:
-                msg = f"Finish: obj={_dataclass_to_string_truncated(obj, max_length, skip_names=skip_names)}"
-            else:
-                msg = f"Finish: obj={_dataclass_to_string_truncated(obj, max_length, skip_names=skip_names)}, out={_dataclass_to_string_truncated(out, max_length, skip_names=out_skip_names)}"
-            self._log(msg)
+            obj_str = _dataclass_to_string_truncated(
+                obj, max_length, skip_names=skip_names
+            )
+            out_str = (
+                ""
+                if is_multimodal_gen
+                else f", out={_dataclass_to_string_truncated(out, max_length, skip_names=out_skip_names)}"
+            )
+            headers_str = f", headers={headers}" if headers else ""
+            self._log(f"Finish: obj={obj_str}{headers_str}{out_str}")
 
     def _compute_metadata(
         self,