From 4c7136bb364d4fdf9ce0f0d31a6d7a92a3457822 Mon Sep 17 00:00:00 2001 From: Douglas Yang Date: Fri, 23 Jan 2026 16:04:28 -0800 Subject: [PATCH] feature: adding openai compatible API request to bench_serving (#17219) --- python/sglang/bench_serving.py | 138 +++++++++++++++++++++++++++++++-- 1 file changed, 130 insertions(+), 8 deletions(-) diff --git a/python/sglang/bench_serving.py b/python/sglang/bench_serving.py index a178f6f62..a00dec41c 100644 --- a/python/sglang/bench_serving.py +++ b/python/sglang/bench_serving.py @@ -228,17 +228,26 @@ async def async_request_openai_completions( prompt = request_func_input.prompt async with _create_bench_client_session() as session: + # Build payload with defaults that can be overridden by extra_request_body payload = { "model": request_func_input.model, "prompt": prompt, - "temperature": 0.0, "best_of": 1, "max_tokens": request_func_input.output_len, "stream": not args.disable_stream, - "ignore_eos": not args.disable_ignore_eos, - **request_func_input.extra_request_body, } + # Add temperature default only if not specified in extra_request_body + if "temperature" not in request_func_input.extra_request_body: + payload["temperature"] = 0.0 + + # Add ignore_eos default only if not specified in extra_request_body + if "ignore_eos" not in request_func_input.extra_request_body: + payload["ignore_eos"] = not args.disable_ignore_eos + + # Merge in extra parameters - these will override defaults if present + payload.update(request_func_input.extra_request_body) + # hack to accommodate different LoRA conventions between SGLang and vLLM. if request_func_input.lora_name: payload["model"] = request_func_input.lora_name @@ -373,16 +382,27 @@ async def async_request_openai_chat_completions( messages = [{"role": "user", "content": request_func_input.prompt}] async with _create_bench_client_session() as session: + # Build payload with defaults that can be overridden by extra_request_body payload = { "model": request_func_input.model, "messages": messages, - "temperature": 0.0, "max_completion_tokens": request_func_input.output_len, "stream": not args.disable_stream, - "ignore_eos": not args.disable_ignore_eos, - **request_func_input.extra_request_body, } + # Add temperature default only if not specified in extra_request_body + if "temperature" not in request_func_input.extra_request_body: + payload["temperature"] = 0.0 + + # Add ignore_eos default only if not specified in extra_request_body + # Default to False for more realistic behavior (respect EOS tokens) + if "ignore_eos" not in request_func_input.extra_request_body: + payload["ignore_eos"] = not args.disable_ignore_eos + + # Merge in extra parameters (tools, temperature, top_p, etc.) + # These will override defaults if present + payload.update(request_func_input.extra_request_body) + # hack to accommodate different LoRA conventions between SGLang and vLLM. if request_func_input.lora_name: payload["model"] = request_func_input.lora_name @@ -895,6 +915,13 @@ def get_dataset(args, tokenizer, model_id=None): prompt_suffix=args.prompt_suffix, apply_chat_template=args.apply_chat_template, ) + elif args.dataset_name == "openai": + input_requests = sample_openai_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + fixed_output_len=args.sharegpt_output_len, + ) else: raise ValueError(f"Unknown dataset: {args.dataset_name}") return input_requests @@ -1034,12 +1061,15 @@ class DatasetRow: image_data: Optional[List[str]] = None timestamp: Optional[float] = None routing_key: Optional[str] = None + extra_request_body: Optional[Dict[str, Any]] = None # Per-request API parameters def __post_init__(self): if self.text_prompt_len is None: self.text_prompt_len = self.prompt_len if self.vision_prompt_len is None: self.vision_prompt_len = 0 + if self.extra_request_body is None: + self.extra_request_body = {} async def get_mooncake_request_over_time( @@ -1300,6 +1330,85 @@ def sample_sharegpt_requests( return filtered_dataset +def sample_openai_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + fixed_output_len: Optional[int] = None, +) -> List[DatasetRow]: + """ + Load OpenAI-compatible chat completion requests from a JSONL file. + + Each line should be a JSON object with: + - "messages": list of {"role": str, "content": str} + - "max_tokens": int (used as output_len if fixed_output_len not set) + - "tools": optional list of tool definitions + - "temperature": optional temperature value + - "top_p": optional top_p value + - Other OpenAI API parameters are also extracted and passed through + """ + dataset = [] + with open(dataset_path, "r") as f: + for line in f: + if num_requests > 0 and len(dataset) >= num_requests: + break + if line.strip(): + try: + dataset.append(json.loads(line)) + except json.JSONDecodeError: + # Skip invalid JSON lines + continue + + # Fields that should NOT be passed through extra_request_body + # These are either handled separately or are metadata + # max_tokens is excluded because it's handled via output_len -> max_completion_tokens + # max_completion_tokens is also excluded to avoid conflicts + EXCLUDED_FIELDS = {"messages", "max_tokens", "max_completion_tokens", "model"} + + filtered_dataset: List[DatasetRow] = [] + for data in dataset: + messages = data.get("messages", []) + if not messages: + continue + + # Use max_tokens from the request, or fall back to fixed_output_len + output_len = fixed_output_len or data.get("max_tokens", 256) + + # Extract extra request body parameters (tools, temperature, top_p, etc.) + extra_body = {k: v for k, v in data.items() if k not in EXCLUDED_FIELDS} + + # Calculate prompt length by applying chat template + # This includes the messages but not the tools + prompt_len = len( + tokenizer.apply_chat_template( + messages, tokenize=True, add_generation_prompt=True + ) + ) + + # If tools are present, we need to add their token count + # Tools are sent as part of the request and count toward input tokens + if "tools" in extra_body: + # Encode tools as JSON string to estimate token count + tools_str = json.dumps(extra_body["tools"]) + tools_tokens = len(tokenizer.encode(tools_str)) + prompt_len += tools_tokens + + # Pass messages list directly - bench_serving handles List[Dict] prompts + filtered_dataset.append( + DatasetRow( + prompt=messages, + prompt_len=prompt_len, + output_len=output_len, + extra_request_body=extra_body, # Store per-request parameters + ) + ) + + print(f"Loaded {len(filtered_dataset)} OpenAI-format requests") + print(f"#Input tokens: {np.sum([x.prompt_len for x in filtered_dataset])}") + print(f"#Output tokens: {np.sum([x.output_len for x in filtered_dataset])}") + return filtered_dataset + + def sample_custom_requests( dataset_path: str, num_requests: int, @@ -2183,7 +2292,15 @@ async def benchmark( else: raise ValueError(f"Unknown backend: {backend}") - is_multi_turn = isinstance(input_requests[0].prompt, list) + # Check for multi-turn: prompt is a list of strings (not OpenAI messages dicts) + # Multi-turn format: ["turn1", "turn2", ...] - list of strings + # OpenAI format: [{"role": "user", "content": "..."}, ...] - list of dicts + first_prompt = input_requests[0].prompt + is_multi_turn = ( + isinstance(first_prompt, list) + and len(first_prompt) > 0 + and isinstance(first_prompt[0], str) + ) if is_multi_turn: request_func = wrap_multi_turn_request_func(request_func, backend=backend) @@ -2340,6 +2457,10 @@ async def benchmark( else: lora_name = None + # Merge global extra_request_body with per-request extras + # Per-request parameters take precedence over global ones + merged_extra_body = {**extra_request_body, **request.extra_request_body} + request_func_input = RequestFuncInput( model=model_id, prompt=request.prompt, @@ -2348,7 +2469,7 @@ async def benchmark( output_len=request.output_len, lora_name=lora_name, image_data=request.image_data, - extra_request_body=extra_request_body, + extra_request_body=merged_extra_body, timestamp=request.timestamp, routing_key=request.routing_key, ) @@ -2877,6 +2998,7 @@ if __name__ == "__main__": choices=[ "sharegpt", "custom", + "openai", "random", "random-ids", "generated-shared-prefix",