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:
+            obj_str = _dataclass_to_string_truncated(
-                msg = f"Finish: obj={_dataclass_to_string_truncated(obj, max_length, skip_names=skip_names)}"
+                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)}"
+            out_str = (
-            self._log(msg)
+                ""
+                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,

test/registered/debug_utils/test_request_logger.py

@@ -18,6 +18,8 @@ from sglang.test.test_utils import (
 
 register_cuda_ci(est_time=120, suite="nightly-1-gpu", nightly=True)
 
+TEST_ROUTING_KEY = "test-routing-key-12345"
+
 
 class BaseTestRequestLogger:
     log_requests_format = None
@@ -64,6 +66,7 @@ class BaseTestRequestLogger:
                 "text": "Hello",
                 "sampling_params": {"max_new_tokens": 8, "temperature": 0},
             },
+            headers={"X-SMG-Routing-Key": TEST_ROUTING_KEY},
             timeout=30,
         )
         self.assertEqual(response.status_code, 200)
@@ -85,6 +88,12 @@ class TestRequestLoggerText(BaseTestRequestLogger, CustomTestCase):
     def _verify_logs(self, content: str, source_name: str):
         self.assertIn("Receive:", content, f"'Receive:' not found in {source_name}")
         self.assertIn("Finish:", content, f"'Finish:' not found in {source_name}")
+        self.assertIn(
+            TEST_ROUTING_KEY, content, f"Routing key not found in {source_name}"
+        )
+        self.assertIn(
+            "x-smg-routing-key", content, f"Header name not found in {source_name}"
+        )
 
 
 class TestRequestLoggerJson(BaseTestRequestLogger, CustomTestCase):
@@ -97,14 +106,25 @@ class TestRequestLoggerJson(BaseTestRequestLogger, CustomTestCase):
             if not line.strip() or not line.startswith("{"):
                 continue
             data = json.loads(line)
+
+            rid = data.get("rid", "")
+            if rid.startswith("HEALTH_CHECK"):
+                continue
+
             if data.get("event") == "request.received":
                 self.assertIn("rid", data)
                 self.assertIn("obj", data)
+                self.assertEqual(
+                    data.get("headers", {}).get("x-smg-routing-key"), TEST_ROUTING_KEY
+                )
                 received_found = True
             elif data.get("event") == "request.finished":
                 self.assertIn("rid", data)
                 self.assertIn("obj", data)
                 self.assertIn("out", data)
+                self.assertEqual(
+                    data.get("headers", {}).get("x-smg-routing-key"), TEST_ROUTING_KEY
+                )
                 finished_found = True
 
         self.assertTrue(