diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py index 71d35a413..5cfefda66 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/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(): diff --git a/python/sglang/srt/utils/request_logger.py b/python/sglang/srt/utils/request_logger.py index c861e2421..0adb4b4e0 100644 --- a/python/sglang/srt/utils/request_logger.py +++ b/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, diff --git a/test/registered/debug_utils/test_request_logger.py b/test/registered/debug_utils/test_request_logger.py index 286ce0ba0..a9b04637d 100644 --- a/test/registered/debug_utils/test_request_logger.py +++ b/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(