diff --git a/benchmarks/multi_turn/benchmark_serving_multi_turn.py b/benchmarks/multi_turn/benchmark_serving_multi_turn.py index 750adf797e..5f0c194af6 100644 --- a/benchmarks/multi_turn/benchmark_serving_multi_turn.py +++ b/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}")