Support gsp send routing id in bench serving (#15588)

python/sglang/bench_serving.py

@@ -88,6 +88,7 @@ class RequestFuncInput:
     image_data: Optional[List[str]]
     extra_request_body: Dict[str, Any]
     timestamp: Optional[float] = None
+    routing_id: Optional[str] = None
 
 
 @dataclass
@@ -233,6 +234,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()
 
         output = RequestFuncOutput.init_new(request_func_input)
@@ -370,6 +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()
 
         output = RequestFuncOutput.init_new(request_func_input)
@@ -574,6 +581,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()
 
         output = RequestFuncOutput.init_new(request_func_input)
@@ -999,6 +1009,7 @@ class DatasetRow:
     vision_prompt_len: Optional[int] = None
     image_data: Optional[List[str]] = None
     timestamp: Optional[float] = None
+    routing_id: Optional[str] = None
 
     def __post_init__(self):
         if self.text_prompt_len is None:
@@ -1650,8 +1661,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)
+
     cache_path = get_gen_prefix_cache_path(args, tokenizer)
-    should_cache = range_ratio == 1
+    should_cache = (range_ratio == 1) and not send_routing_id
 
     # Try to load from cache first
     if cache_path.exists() and should_cache:
@@ -1664,6 +1677,9 @@ def sample_generated_shared_prefix_requests(
         f"({num_groups=}, {prompts_per_group}, {system_prompt_len=}, {question_len=}, {output_len=}, {range_ratio=})"
     )
 
+    run_random_str = uuid.uuid4().hex[:8]
+    run_start_timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
+
     system_prompt_lens = compute_random_lens(
         full_len=system_prompt_len,
         range_ratio=range_ratio,
@@ -1700,6 +1716,11 @@ 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 = (
+            f"{run_random_str}_{run_start_timestamp}_{group_idx}"
+            if send_routing_id
+            else None
+        )
         for prompt_idx in tqdm(
             range(prompts_per_group), desc="Generating questions", leave=False
         ):
@@ -1717,6 +1738,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,
                 )
             )
             total_input_tokens += prompt_len
@@ -2153,6 +2175,7 @@ async def benchmark(
             image_data=request.image_data,
             extra_request_body=extra_request_body,
             timestamp=request.timestamp,
+            routing_id=request.routing_id,
         )
 
         tasks.append(
@@ -2981,6 +3004,11 @@ if __name__ == "__main__":
         action="store_true",
         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",
+        action="store_true",
+        help="Send routing_id in requests. Requests with the same prefix share the same routing_id.",
+    )
     mooncake_group = parser.add_argument_group("mooncake dataset arguments")
     mooncake_group.add_argument(
         "--mooncake-slowdown-factor",