feature: adding openai compatible API request to bench_serving (#17219)

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",