chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Douglas Yang
2026-01-23 16:04:28 -08:00
committed by GitHub
parent ad05782160
commit 4c7136bb36
1 changed files with 130 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

138
python/sglang/bench_serving.py
View File

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