diff --git a/python/sglang/bench_serving.py b/python/sglang/bench_serving.py index ecf036f99..20430f4e1 100644 --- a/python/sglang/bench_serving.py +++ b/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",