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