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>

benchmarks/multi_turn/benchmark_serving_multi_turn.py

@@ -65,6 +65,31 @@ class RequestArgs(NamedTuple):
     limit_min_tokens: int  # Use negative value for no limit
     limit_max_tokens: int  # Use negative value for no limit
     timeout_sec: int
+    headers: dict[str, str]
+
+
+def parse_custom_header(header: str) -> tuple[str, str]:
+    separators = (":", "=")
+    for separator in separators:
+        if separator in header:
+            key, value = header.split(separator, 1)
+            key = key.strip()
+            value = value.strip()
+            if key:
+                return key, value
+            break
+    raise argparse.ArgumentTypeError(
+        "Headers must be provided as 'Header-Name: value' or 'Header-Name=value'"
+    )
+
+
+def build_request_headers(
+    api_key: str | None, custom_headers: list[tuple[str, str]] | None
+) -> dict[str, str]:
+    headers = dict(custom_headers or [])
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+    return headers
 
 
 class BenchmarkArgs(NamedTuple):
@@ -218,12 +243,11 @@ async def send_request(
     max_tokens: int | None = None,
     timeout_sec: int = 120,
     conversation_id: str | None = None,
+    headers: dict[str, str] | None = None,
 ) -> ServerResponse:
     payload = {
         "model": model,
         "messages": messages,
-        "seed": 0,
-        "temperature": 0.0,
     }
 
     if conversation_id is not None:
@@ -233,15 +257,17 @@ async def send_request(
         payload["stream"] = True
         payload["stream_options"] = {"include_usage": False}
 
-    if min_tokens is not None:
-        payload["min_tokens"] = min_tokens
+    # if min_tokens is not None:
+    #     payload["min_tokens"] = min_tokens
 
     if max_tokens is not None:
         payload["max_tokens"] = max_tokens
 
-    headers = {"Content-Type": "application/json"}
+    request_headers = {"Content-Type": "application/json"}
     if conversation_id is not None:
-        headers["X-Session-ID"] = str(conversation_id)
+        request_headers["X-Session-ID"] = str(conversation_id)
+    if headers is not None:
+        request_headers.update(headers)
 
     # Calculate the timeout for the request
     if max_tokens is not None:
@@ -267,7 +293,7 @@ async def send_request(
     most_recent_timestamp: int = start_time
 
     async with session.post(
-        url=chat_url, json=payload, headers=headers, timeout=timeout
+        url=chat_url, json=payload, headers=request_headers, timeout=timeout
     ) as response:
         http_status = HTTPStatus(response.status)
         if http_status == HTTPStatus.OK:
@@ -319,6 +345,8 @@ async def send_request(
             latency = time.perf_counter_ns() - start_time
 
     if ttft is None:
+        if stream:
+            valid_response = False
         # The response was a single chunk
         ttft = latency
 
@@ -426,6 +454,7 @@ async def send_turn(
         max_tokens,
         req_args.timeout_sec,
         conversation_id=conv_id,
+        headers=req_args.headers,
     )
 
     if response.valid is False:
@@ -874,6 +903,7 @@ def get_client_config(
     # Arguments for API requests
     chat_url = f"{args.url}/v1/chat/completions"
     model_name = args.served_model_name if args.served_model_name else args.model
+    headers = build_request_headers(args.api_key, args.header)
 
     req_args = RequestArgs(
         chat_url=chat_url,
@@ -882,6 +912,7 @@ def get_client_config(
         limit_min_tokens=args.limit_min_tokens,
         limit_max_tokens=args.limit_max_tokens,
         timeout_sec=args.request_timeout_sec,
+        headers=headers,
     )
 
     return client_args, req_args
@@ -1247,19 +1278,19 @@ def process_statistics(
         )
 
 
-async def get_server_info(url: str) -> None:
+async def get_server_info(url: str, headers: dict[str, str] | None = None) -> None:
     logger.info(f"{Color.BLUE}Collecting information from server: {url}{Color.RESET}")
     async with aiohttp.ClientSession() as session:
         # Get server version (not mandatory, "version" endpoint may not exist)
         url_version = f"{url}/version"
-        async with session.get(url_version) as response:
+        async with session.get(url_version, headers=headers) as response:
             if HTTPStatus(response.status) == HTTPStatus.OK:
                 text = await response.text()
                 logger.info(f"{Color.BLUE}Server version: {text}{Color.RESET}")
 
         # Get available models
         url_models = f"{url}/v1/models"
-        async with session.get(url_models) as response:
+        async with session.get(url_models, headers=headers) as response:
             if HTTPStatus(response.status) == HTTPStatus.OK:
                 text = await response.text()
                 logger.info(f"{Color.BLUE}Models:{Color.RESET}")
@@ -1325,6 +1356,22 @@ async def main() -> None:
         help="Base URL for the LLM API server",
     )
 
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default=None,
+        help="API key to send as an Authorization bearer token",
+    )
+    parser.add_argument(
+        "--header",
+        action="append",
+        type=parse_custom_header,
+        default=None,
+        metavar="KEY=VALUE",
+        help="Custom request header. Can be specified multiple times. "
+        "Accepts 'Header-Name: value' or 'Header-Name=value'.",
+    )
+
     parser.add_argument(
         "-p",
         "--num-clients",
@@ -1527,7 +1574,8 @@ async def main() -> None:
         args.model, trust_remote_code=args.trust_remote_code
     )
 
-    await get_server_info(args.url)
+    headers = build_request_headers(args.api_key, args.header)
+    await get_server_info(args.url, headers=headers)
 
     # Load the input file (either conversations of configuration file)
     logger.info(f"Reading input file: {args.input_file}")