A high-throughput and memory-efficient inference and serving engine for LLMs
feat(multi-turn-bench): add api_key and custom headers for multi turn benchmark (#44516)
Signed-off-by: Jimmy <jinmingyi1998@sina.cn> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: simon-mo <simon.mo@hey.com>
Jimmy committed
Jun 10, 2026 at 05:00 UTC
1c2ffc6f8891809fa819d71ba360d08129d476ac
1 file changed
+59
-11
benchmarks/multi_turn/benchmark_serving_multi_turn.py
+59
-11
@@ -65,6 +65,31 @@ class RequestArgs(NamedTuple):
65
limit_min_tokens: int # Use negative value for no limit
66
limit_max_tokens: int # Use negative value for no limit
67
timeout_sec: int
68
+ headers: dict[str, str]
69
+
70
+
71
+def parse_custom_header(header: str) -> tuple[str, str]:
72
+ separators = (":", "=")
73
+ for separator in separators:
74
+ if separator in header:
75
+ key, value = header.split(separator, 1)
76
+ key = key.strip()
77
+ value = value.strip()
78
+ if key:
79
+ return key, value
80
+ break
81
+ raise argparse.ArgumentTypeError(
82
+ "Headers must be provided as 'Header-Name: value' or 'Header-Name=value'"
83
+ )
84
+
85
+
86
+def build_request_headers(
87
+ api_key: str | None, custom_headers: list[tuple[str, str]] | None
88
+) -> dict[str, str]:
89
+ headers = dict(custom_headers or [])
90
+ if api_key:
91
+ headers["Authorization"] = f"Bearer {api_key}"
92
+ return headers
93
94
95
class BenchmarkArgs(NamedTuple):
@@ -218,12 +243,11 @@ async def send_request(
243
max_tokens: int | None = None,
244
timeout_sec: int = 120,
245
conversation_id: str | None = None,
246
+ headers: dict[str, str] | None = None,
247
) -> ServerResponse:
248
payload = {
249
"model": model,
250
"messages": messages,
225
- "seed": 0,
226
- "temperature": 0.0,
251
}
252
253
if conversation_id is not None:
@@ -233,15 +257,17 @@ async def send_request(
257
payload["stream"] = True
258
payload["stream_options"] = {"include_usage": False}
259
236
- if min_tokens is not None:
237
- payload["min_tokens"] = min_tokens
260
+ # if min_tokens is not None:
261
+ # payload["min_tokens"] = min_tokens
262
263
if max_tokens is not None:
264
payload["max_tokens"] = max_tokens
265
242
- headers = {"Content-Type": "application/json"}
266
+ request_headers = {"Content-Type": "application/json"}
267
if conversation_id is not None:
244
- headers["X-Session-ID"] = str(conversation_id)
268
+ request_headers["X-Session-ID"] = str(conversation_id)
269
+ if headers is not None:
270
+ request_headers.update(headers)
271
272
# Calculate the timeout for the request
273
if max_tokens is not None:
@@ -267,7 +293,7 @@ async def send_request(
293
most_recent_timestamp: int = start_time
294
295
async with session.post(
270
- url=chat_url, json=payload, headers=headers, timeout=timeout
296
+ url=chat_url, json=payload, headers=request_headers, timeout=timeout
297
) as response:
298
http_status = HTTPStatus(response.status)
299
if http_status == HTTPStatus.OK:
@@ -319,6 +345,8 @@ async def send_request(
345
latency = time.perf_counter_ns() - start_time
346
347
if ttft is None:
348
+ if stream:
349
+ valid_response = False
350
# The response was a single chunk
351
ttft = latency
352
@@ -426,6 +454,7 @@ async def send_turn(
454
max_tokens,
455
req_args.timeout_sec,
456
conversation_id=conv_id,
457
+ headers=req_args.headers,
458
)
459
460
if response.valid is False:
@@ -874,6 +903,7 @@ def get_client_config(
903
# Arguments for API requests
904
chat_url = f"{args.url}/v1/chat/completions"
905
model_name = args.served_model_name if args.served_model_name else args.model
906
+ headers = build_request_headers(args.api_key, args.header)
907
908
req_args = RequestArgs(
909
chat_url=chat_url,
@@ -882,6 +912,7 @@ def get_client_config(
912
limit_min_tokens=args.limit_min_tokens,
913
limit_max_tokens=args.limit_max_tokens,
914
timeout_sec=args.request_timeout_sec,
915
+ headers=headers,
916
)
917
918
return client_args, req_args
@@ -1247,19 +1278,19 @@ def process_statistics(
1278
)
1279
1280
1250
-async def get_server_info(url: str) -> None:
1281
+async def get_server_info(url: str, headers: dict[str, str] | None = None) -> None:
1282
logger.info(f"{Color.BLUE}Collecting information from server: {url}{Color.RESET}")
1283
async with aiohttp.ClientSession() as session:
1284
# Get server version (not mandatory, "version" endpoint may not exist)
1285
url_version = f"{url}/version"
1255
- async with session.get(url_version) as response:
1286
+ async with session.get(url_version, headers=headers) as response:
1287
if HTTPStatus(response.status) == HTTPStatus.OK:
1288
text = await response.text()
1289
logger.info(f"{Color.BLUE}Server version: {text}{Color.RESET}")
1290
1291
# Get available models
1292
url_models = f"{url}/v1/models"
1262
- async with session.get(url_models) as response:
1293
+ async with session.get(url_models, headers=headers) as response:
1294
if HTTPStatus(response.status) == HTTPStatus.OK:
1295
text = await response.text()
1296
logger.info(f"{Color.BLUE}Models:{Color.RESET}")
@@ -1325,6 +1356,22 @@ async def main() -> None:
1356
help="Base URL for the LLM API server",
1357
)
1358
1359
+ parser.add_argument(
1360
+ "--api-key",
1361
+ type=str,
1362
+ default=None,
1363
+ help="API key to send as an Authorization bearer token",
1364
+ )
1365
+ parser.add_argument(
1366
+ "--header",
1367
+ action="append",
1368
+ type=parse_custom_header,
1369
+ default=None,
1370
+ metavar="KEY=VALUE",
1371
+ help="Custom request header. Can be specified multiple times. "
1372
+ "Accepts 'Header-Name: value' or 'Header-Name=value'.",
1373
+ )
1374
+
1375
parser.add_argument(
1376
"-p",
1377
"--num-clients",
@@ -1527,7 +1574,8 @@ async def main() -> None:
1574
args.model, trust_remote_code=args.trust_remote_code
1575
)
1576
1530
- await get_server_info(args.url)
1577
+ headers = build_request_headers(args.api_key, args.header)
1578
+ await get_server_info(args.url, headers=headers)
1579
1580
# Load the input file (either conversations of configuration file)
1581
logger.info(f"Reading input file: {args.input_file}")