A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Benchmark] Allow arbitrary headers to be passed to benchmarked endpoints (#23937)
Signed-off-by: Clayton Coleman <smarterclayton@gmail.com>
Clayton Coleman committed
Sep 12, 2025 at 16:57 UTC
bc636f21a697a8391ee2767fbf09b9981a0f9604
2 files changed
+35
-1
vllm/benchmarks/lib/endpoint_request_func.py
+7
@@ -68,6 +68,7 @@ class RequestFuncInput:
68
model: str
69
model_name: Optional[str] = None
70
logprobs: Optional[int] = None
71
+ extra_headers: Optional[dict] = None
72
extra_body: Optional[dict] = None
73
multi_modal_content: Optional[Union[dict, list[dict]]] = None
74
ignore_eos: bool = False
@@ -129,6 +130,8 @@ async def async_request_openai_completions(
130
headers = {
131
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"
132
}
133
+ if request_func_input.extra_headers:
134
+ headers |= request_func_input.extra_headers
135
if request_func_input.request_id:
136
headers["x-request-id"] = request_func_input.request_id
137
@@ -258,6 +261,8 @@ async def async_request_openai_chat_completions(
261
"Content-Type": "application/json",
262
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
263
}
264
+ if request_func_input.extra_headers:
265
+ headers |= request_func_input.extra_headers
266
if request_func_input.request_id:
267
headers["x-request-id"] = request_func_input.request_id
268
@@ -364,6 +369,8 @@ async def async_request_openai_audio(
369
headers = {
370
"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
371
}
372
+ if request_func_input.extra_headers:
373
+ headers |= request_func_input.extra_headers
374
if request_func_input.request_id:
375
headers["x-request-id"] = request_func_input.request_id
376
vllm/benchmarks/serve.py
+28
-1
@@ -389,6 +389,7 @@ async def benchmark(
389
goodput_config_dict: dict[str, float],
390
max_concurrency: Optional[int],
391
lora_modules: Optional[Iterable[str]],
392
+ extra_headers: Optional[dict],
393
extra_body: Optional[dict],
394
ramp_up_strategy: Optional[Literal["linear", "exponential"]] = None,
395
ramp_up_start_rps: Optional[int] = None,
@@ -452,6 +453,7 @@ async def benchmark(
453
logprobs=logprobs,
454
multi_modal_content=test_mm_content,
455
ignore_eos=ignore_eos,
456
+ extra_headers=extra_headers,
457
extra_body=extra_body,
458
)
459
@@ -484,6 +486,7 @@ async def benchmark(
486
logprobs=logprobs,
487
multi_modal_content=test_mm_content,
488
ignore_eos=ignore_eos,
489
+ extra_headers=extra_headers,
490
extra_body=extra_body)
491
profile_output = await request_func(
492
request_func_input=profile_input, session=session)
@@ -568,6 +571,7 @@ async def benchmark(
571
logprobs=logprobs,
572
multi_modal_content=mm_content,
573
ignore_eos=ignore_eos,
574
+ extra_headers=extra_headers,
575
extra_body=extra_body,
576
request_id=request_id,)
577
tasks.append(
@@ -815,6 +819,15 @@ def add_cli_args(parser: argparse.ArgumentParser):
819
default="/v1/completions",
820
help="API endpoint.",
821
)
822
+ parser.add_argument(
823
+ "--header",
824
+ metavar="KEY=VALUE",
825
+ nargs="*",
826
+ help="Key-value pairs (e.g, --header x-additional-info=0.3.3) "
827
+ "for headers to be passed with each request. These headers override " \
828
+ "per backend constants and values set via environment variable, and " \
829
+ "will be overriden by other arguments (such as request ids)."
830
+ )
831
parser.add_argument(
832
"--max-concurrency",
833
type=int,
@@ -1104,6 +1117,19 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
1117
api_url = f"http://{args.host}:{args.port}{args.endpoint}"
1118
base_url = f"http://{args.host}:{args.port}"
1119
1120
+ # Headers
1121
+ headers = None
1122
+ if args.header:
1123
+ headers = {}
1124
+ for item in args.header:
1125
+ if "=" in item:
1126
+ kvstring = item.split("=", 1)
1127
+ headers[kvstring[0].strip()] = kvstring[1].strip()
1128
+ else:
1129
+ raise ValueError(
1130
+ "Invalid header format. Please use KEY=VALUE format."
1131
+ )
1132
+
1133
tokenizer = get_tokenizer(tokenizer_id,
1134
tokenizer_mode=tokenizer_mode,
1135
trust_remote_code=args.trust_remote_code)
@@ -1161,6 +1187,7 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
1187
goodput_config_dict=goodput_config_dict,
1188
max_concurrency=args.max_concurrency,
1189
lora_modules=args.lora_modules,
1190
+ extra_headers=headers,
1191
extra_body=sampling_params,
1192
ramp_up_strategy=args.ramp_up_strategy,
1193
ramp_up_start_rps=args.ramp_up_start_rps,
@@ -1184,7 +1211,7 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
1211
if args.metadata:
1212
for item in args.metadata:
1213
if "=" in item:
1187
- kvstring = item.split("=")
1214
+ kvstring = item.split("=", 1)
1215
result_json[kvstring[0].strip()] = kvstring[1].strip()
1216
else:
1217
raise ValueError(