Tiny change bench-serving to use routing key header (#15827)

python/sglang/bench_serving.py

@@ -50,6 +50,7 @@ from transformers import (
 )
 
 ASSISTANT_SUFFIX = "Assistant:"
+_ROUTING_KEY_HEADER = "X-SMG-Routing-Key"
 
 TERM_PLOTLIB_AVAILABLE = (importlib.util.find_spec("termplotlib") is not None) and (
     shutil.which("gnuplot") is not None
@@ -88,7 +89,7 @@ class RequestFuncInput:
     image_data: Optional[List[str]]
     extra_request_body: Dict[str, Any]
     timestamp: Optional[float] = None
-    routing_id: Optional[str] = None
+    routing_key: Optional[str] = None
 
 
 @dataclass
@@ -234,10 +235,9 @@ async def async_request_openai_completions(
         if request_func_input.image_data:
             payload.update({"image_data": request_func_input.image_data})
 
-        if request_func_input.routing_id:
-            payload["routing_id"] = request_func_input.routing_id
-
         headers = get_auth_headers()
+        if request_func_input.routing_key:
+            headers[_ROUTING_KEY_HEADER] = request_func_input.routing_key
 
         output = RequestFuncOutput.init_new(request_func_input)
 
@@ -374,10 +374,9 @@ async def async_request_openai_chat_completions(
             payload["model"] = request_func_input.lora_name
             payload["lora_path"] = request_func_input.lora_name
 
-        if request_func_input.routing_id:
-            payload["routing_id"] = request_func_input.routing_id
-
         headers = get_auth_headers()
+        if request_func_input.routing_key:
+            headers[_ROUTING_KEY_HEADER] = request_func_input.routing_key
 
         output = RequestFuncOutput.init_new(request_func_input)
 
@@ -581,10 +580,9 @@ async def async_request_sglang_generate(
         if request_func_input.image_data:
             payload["image_data"] = request_func_input.image_data
 
-        if request_func_input.routing_id:
-            payload["routing_id"] = request_func_input.routing_id
-
         headers = get_auth_headers()
+        if request_func_input.routing_key:
+            headers[_ROUTING_KEY_HEADER] = request_func_input.routing_key
 
         output = RequestFuncOutput.init_new(request_func_input)
 
@@ -1009,7 +1007,7 @@ class DatasetRow:
     vision_prompt_len: Optional[int] = None
     image_data: Optional[List[str]] = None
     timestamp: Optional[float] = None
-    routing_id: Optional[str] = None
+    routing_key: Optional[str] = None
 
     def __post_init__(self):
         if self.text_prompt_len is None:
@@ -1661,10 +1659,10 @@ def sample_generated_shared_prefix_requests(
     args: argparse.Namespace,
 ) -> List[DatasetRow]:
     """Generate benchmark requests with shared system prompts using random tokens and caching."""
-    send_routing_id = getattr(args, "gsp_send_routing_id", False)
+    send_routing_key = getattr(args, "gsp_send_routing_key", False)
 
     cache_path = get_gen_prefix_cache_path(args, tokenizer)
-    should_cache = (range_ratio == 1) and not send_routing_id
+    should_cache = (range_ratio == 1) and not send_routing_key
 
     # Try to load from cache first
     if cache_path.exists() and should_cache:
@@ -1716,9 +1714,9 @@ def sample_generated_shared_prefix_requests(
 
     for group_idx in tqdm(range(num_groups), desc="Generating system prompt"):
         system_prompt = system_prompts[group_idx]
-        routing_id = (
+        routing_key = (
             f"{run_random_str}_{run_start_timestamp}_{group_idx}"
-            if send_routing_id
+            if send_routing_key
             else None
         )
         for prompt_idx in tqdm(
@@ -1738,7 +1736,7 @@ def sample_generated_shared_prefix_requests(
                     prompt=full_prompt,
                     prompt_len=prompt_len,
                     output_len=output_lens[flat_index].item(),
-                    routing_id=routing_id,
+                    routing_key=routing_key,
                 )
             )
             total_input_tokens += prompt_len
@@ -2175,7 +2173,7 @@ async def benchmark(
             image_data=request.image_data,
             extra_request_body=extra_request_body,
             timestamp=request.timestamp,
-            routing_id=request.routing_id,
+            routing_key=request.routing_key,
         )
 
         tasks.append(
@@ -3005,9 +3003,9 @@ if __name__ == "__main__":
         help="Speedup preparing by removing statistics computation, which will make some output statistics inaccurate but suitable for pressure tests.",
     )
     group.add_argument(
-        "--gsp-send-routing-id",
+        "--gsp-send-routing-key",
         action="store_true",
-        help="Send routing_id in requests. Requests with the same prefix share the same routing_id.",
+        help="Send routing key in requests via X-SMG-Routing-Key header. Requests with the same prefix share the same routing key.",
     )
     mooncake_group = parser.add_argument_group("mooncake dataset arguments")
     mooncake_group.add_argument(