| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import argparse |
| 4 | import asyncio |
| 5 | import json |
| 6 | import logging |
| 7 | import multiprocessing as mp |
| 8 | import os |
| 9 | import random |
| 10 | import time |
| 11 | from collections import Counter, deque |
| 12 | from datetime import datetime |
| 13 | from enum import Enum |
| 14 | from http import HTTPStatus |
| 15 | from statistics import mean |
| 16 | from typing import NamedTuple |
| 17 | |
| 18 | import aiohttp # type: ignore |
| 19 | import numpy as np # type: ignore |
| 20 | import pandas as pd # type: ignore |
| 21 | from bench_dataset import ( |
| 22 | ConversationsMap, |
| 23 | ConvId, |
| 24 | GenConvArgs, |
| 25 | MessagesList, |
| 26 | ShareGptConversations, |
| 27 | conversations_dict_to_list, |
| 28 | conversations_list_to_dict, |
| 29 | generate_conversations, |
| 30 | parse_input_json_file, |
| 31 | ) |
| 32 | from bench_utils import TEXT_SEPARATOR, Color, logger |
| 33 | from transformers import AutoTokenizer # type: ignore |
| 34 | |
| 35 | from vllm.benchmarks.lib.utils import redact_sensitive_namespace |
| 36 | |
| 37 | NUM_TOKENS_FROM_DATASET = 0 |
| 38 | TERM_SIGNAL = None |
| 39 | _SENSITIVE_ARG_FIELDS = ("api_key", "header") |
| 40 | |
| 41 | |
| 42 | class ConversationSampling(str, Enum): |
| 43 | ROUND_ROBIN = "round_robin" |
| 44 | RANDOM = "random" |
| 45 | |
| 46 | def __str__(self): |
| 47 | return self.value |
| 48 | |
| 49 | |
| 50 | class ClientArgs(NamedTuple): |
| 51 | seed: int |
| 52 | max_num_requests: int | None |
| 53 | skip_first_turn: bool |
| 54 | max_turns: int | None |
| 55 | max_active_conversations: int |
| 56 | verbose: bool |
| 57 | print_content: bool |
| 58 | verify_output: bool |
| 59 | conversation_sampling: ConversationSampling |
| 60 | request_rate: float |
| 61 | max_retries: int |
| 62 | |
| 63 | |
| 64 | class RequestArgs(NamedTuple): |
| 65 | chat_url: str |
| 66 | model: str |
| 67 | stream: bool |
| 68 | limit_min_tokens: int # Use negative value for no limit |
| 69 | limit_max_tokens: int # Use negative value for no limit |
| 70 | timeout_sec: int |
| 71 | send_conversation_id: bool |
| 72 | headers: dict[str, str] |
| 73 | |
| 74 | |
| 75 | def parse_custom_header(header: str) -> tuple[str, str]: |
| 76 | separators = (":", "=") |
| 77 | for separator in separators: |
| 78 | if separator in header: |
| 79 | key, value = header.split(separator, 1) |
| 80 | key = key.strip() |
| 81 | value = value.strip() |
| 82 | if key: |
| 83 | return key, value |
| 84 | break |
| 85 | raise argparse.ArgumentTypeError( |
| 86 | "Headers must be provided as 'Header-Name: value' or 'Header-Name=value'" |
| 87 | ) |
| 88 | |
| 89 | |
| 90 | def build_request_headers( |
| 91 | api_key: str | None, custom_headers: list[tuple[str, str]] | None |
| 92 | ) -> dict[str, str]: |
| 93 | headers = dict(custom_headers or []) |
| 94 | if api_key: |
| 95 | headers["Authorization"] = f"Bearer {api_key}" |
| 96 | return headers |
| 97 | |
| 98 | |
| 99 | class BenchmarkArgs(NamedTuple): |
| 100 | url: str |
| 101 | num_clients: int |
| 102 | early_stop: bool |
| 103 | |
| 104 | |
| 105 | class ServerResponse(NamedTuple): |
| 106 | valid: bool |
| 107 | ttft_ms: float # time to first chunk |
| 108 | tpot_ms: float # time per output chunk (one or more tokens) |
| 109 | latency_ms: float |
| 110 | start_time_ms: float |
| 111 | first_chunk: str # first chunk of the content |
| 112 | content: str # includes the first_chunk |
| 113 | num_chunks: int |
| 114 | |
| 115 | def __str__(self) -> str: |
| 116 | return f"ttft_ms {self.ttft_ms:.2f}, tpot_ms {self.tpot_ms:.2f}, latency_ms {self.latency_ms:.2f}" # noqa: E501 |
| 117 | |
| 118 | |
| 119 | class RequestStats(NamedTuple): |
| 120 | ttft_ms: float |
| 121 | tpot_ms: float |
| 122 | latency_ms: float |
| 123 | start_time_ms: float |
| 124 | input_num_turns: int |
| 125 | input_num_tokens: int |
| 126 | output_num_tokens: int |
| 127 | output_num_chunks: int |
| 128 | output_num_first_chunk_tokens: int |
| 129 | approx_cached_percent: float |
| 130 | conversation_id: str |
| 131 | client_id: int |
| 132 | |
| 133 | def __str__(self) -> str: |
| 134 | return ( |
| 135 | f"ttft_ms {self.ttft_ms:.2f}, tpot_ms {self.tpot_ms:.2f}, latency_ms {self.latency_ms:.2f}, input_num_tokens {self.input_num_tokens}, " # noqa: E501 |
| 136 | f"output_num_tokens {self.output_num_tokens} ({self.output_num_chunks} chunks, {self.output_num_first_chunk_tokens} tokens in first chunk), " # noqa: E501 |
| 137 | f"approx_cached_percent {self.approx_cached_percent:.2f}%" |
| 138 | ) |
| 139 | |
| 140 | |
| 141 | class MetricStats: |
| 142 | def __init__(self) -> None: |
| 143 | self.min: float | None = None |
| 144 | self.max: float | None = None |
| 145 | self.avg: float | None = None |
| 146 | self.sum = 0.0 |
| 147 | self.count = 0 |
| 148 | |
| 149 | def update(self, value: float) -> None: |
| 150 | if self.min is None: |
| 151 | self.min = value |
| 152 | else: |
| 153 | self.min = min(self.min, value) |
| 154 | |
| 155 | if self.max is None: |
| 156 | self.max = value |
| 157 | else: |
| 158 | self.max = max(self.max, value) |
| 159 | |
| 160 | self.sum += value |
| 161 | self.count += 1 |
| 162 | self.avg = self.sum / self.count |
| 163 | |
| 164 | def __repr__(self) -> str: |
| 165 | if self.count == 0: |
| 166 | return "no data" |
| 167 | return f"avg: {self.avg:>10.3f}, min: {self.min:>10.3f}, max: {self.max:>10.3f}" |
| 168 | |
| 169 | |
| 170 | class MovingAverage: |
| 171 | def __init__(self, window_size: int) -> None: |
| 172 | self.window_size = window_size |
| 173 | self.window = np.zeros(window_size) |
| 174 | self.index = 0 |
| 175 | self.sum = 0.0 |
| 176 | self.count = 0 |
| 177 | self.avg: float | None = None |
| 178 | |
| 179 | def update(self, new_value: float) -> None: |
| 180 | if self.count < self.window_size: |
| 181 | # Filling up the window |
| 182 | self.sum += new_value |
| 183 | self.window[self.count] = new_value |
| 184 | self.count += 1 |
| 185 | else: |
| 186 | # Window is full, start replacing old values |
| 187 | old_value = self.window[self.index] |
| 188 | self.sum = self.sum - old_value + new_value |
| 189 | self.window[self.index] = new_value |
| 190 | self.index = (self.index + 1) % self.window_size |
| 191 | |
| 192 | self.avg = self.sum / self.count |
| 193 | |
| 194 | def __repr__(self) -> str: |
| 195 | if self.count == 0: |
| 196 | return "no data" |
| 197 | return f"avg: {self.avg:>10.3f} ({self.count} samples)" |
| 198 | |
| 199 | |
| 200 | class DebugStats: |
| 201 | def __init__(self, logger: logging.Logger, window_size: int) -> None: |
| 202 | self.logger = logger |
| 203 | self.metrics: dict[str, MovingAverage | MetricStats] = { |
| 204 | "moving_avg_ttft_ms": MovingAverage(window_size), |
| 205 | "moving_avg_tpot_ms": MovingAverage(window_size), |
| 206 | "ttft_ms": MetricStats(), |
| 207 | "tpot_ms": MetricStats(), |
| 208 | "latency_ms": MetricStats(), |
| 209 | "input_num_turns": MetricStats(), |
| 210 | "input_num_tokens": MetricStats(), |
| 211 | "output_num_tokens": MetricStats(), |
| 212 | } |
| 213 | |
| 214 | def update(self, data: RequestStats) -> None: |
| 215 | self.metrics["ttft_ms"].update(data.ttft_ms) |
| 216 | self.metrics["moving_avg_ttft_ms"].update(data.ttft_ms) |
| 217 | self.metrics["tpot_ms"].update(data.tpot_ms) |
| 218 | self.metrics["moving_avg_tpot_ms"].update(data.tpot_ms) |
| 219 | self.metrics["latency_ms"].update(data.latency_ms) |
| 220 | self.metrics["input_num_turns"].update(data.input_num_turns) |
| 221 | self.metrics["input_num_tokens"].update(data.input_num_tokens) |
| 222 | self.metrics["output_num_tokens"].update(data.output_num_tokens) |
| 223 | |
| 224 | def print(self) -> None: |
| 225 | self.logger.info("-" * 50) |
| 226 | for k, v in self.metrics.items(): |
| 227 | kv_info = f"[{k:25}] {v}" |
| 228 | self.logger.info(kv_info) |
| 229 | self.logger.info("-" * 50) |
| 230 | |
| 231 | |
| 232 | def nanosec_to_millisec(value: float) -> float: |
| 233 | return value / 1000000.0 |
| 234 | |
| 235 | |
| 236 | def nanosec_to_sec(value: float) -> float: |
| 237 | return value / 1000000000.0 |
| 238 | |
| 239 | |
| 240 | async def send_request( |
| 241 | session: aiohttp.ClientSession, |
| 242 | messages: list[dict[str, str]], |
| 243 | chat_url: str, |
| 244 | model: str, |
| 245 | stream: bool = True, |
| 246 | min_tokens: int | None = None, |
| 247 | max_tokens: int | None = None, |
| 248 | timeout_sec: int = 120, |
| 249 | conversation_id: str | None = None, |
| 250 | headers: dict[str, str] | None = None, |
| 251 | ) -> ServerResponse: |
| 252 | payload = { |
| 253 | "model": model, |
| 254 | "messages": messages, |
| 255 | } |
| 256 | |
| 257 | if conversation_id is not None: |
| 258 | payload["conversation_id"] = conversation_id |
| 259 | |
| 260 | if stream: |
| 261 | payload["stream"] = True |
| 262 | payload["stream_options"] = {"include_usage": False} |
| 263 | |
| 264 | # if min_tokens is not None: |
| 265 | # payload["min_tokens"] = min_tokens |
| 266 | |
| 267 | if max_tokens is not None: |
| 268 | payload["max_tokens"] = max_tokens |
| 269 | |
| 270 | request_headers = {"Content-Type": "application/json"} |
| 271 | if conversation_id is not None: |
| 272 | request_headers["X-Session-ID"] = str(conversation_id) |
| 273 | if headers is not None: |
| 274 | request_headers.update(headers) |
| 275 | |
| 276 | # Calculate the timeout for the request |
| 277 | if max_tokens is not None: |
| 278 | # Assume TPOT of 200ms and use max_tokens to determine timeout |
| 279 | token_based_timeout = int(max_tokens * 0.2) |
| 280 | if token_based_timeout > timeout_sec: |
| 281 | timeout_sec = token_based_timeout |
| 282 | logger.info( |
| 283 | "Using timeout of %ds based on max_tokens %d", |
| 284 | timeout_sec, |
| 285 | max_tokens, |
| 286 | ) |
| 287 | timeout = aiohttp.ClientTimeout(total=timeout_sec) |
| 288 | |
| 289 | valid_response = True |
| 290 | ttft: float | None = None |
| 291 | chunk_delay: list[int] = [] |
| 292 | latency: float | None = None |
| 293 | first_chunk = "" |
| 294 | generated_text = "" |
| 295 | |
| 296 | start_time: int = time.perf_counter_ns() |
| 297 | most_recent_timestamp: int = start_time |
| 298 | |
| 299 | async with session.post( |
| 300 | url=chat_url, json=payload, headers=request_headers, timeout=timeout |
| 301 | ) as response: |
| 302 | http_status = HTTPStatus(response.status) |
| 303 | if http_status == HTTPStatus.OK: |
| 304 | async for chunk_bytes in response.content: |
| 305 | chunk_bytes = chunk_bytes.strip() |
| 306 | if not chunk_bytes: |
| 307 | continue |
| 308 | |
| 309 | chunk = chunk_bytes.decode("utf-8").removeprefix("data: ") |
| 310 | if chunk == "[DONE]": |
| 311 | # End of stream |
| 312 | latency = time.perf_counter_ns() - start_time |
| 313 | elif stream is False: |
| 314 | data = json.loads(chunk) |
| 315 | message = data["choices"][0]["message"] |
| 316 | assert message["role"] == "assistant" |
| 317 | generated_text += message["content"] |
| 318 | else: |
| 319 | timestamp: int = time.perf_counter_ns() |
| 320 | data = json.loads(chunk) |
| 321 | |
| 322 | # Delta is the new content/text/data |
| 323 | delta = data["choices"][0]["delta"] |
| 324 | if delta.get("content", None): |
| 325 | if ttft is None: |
| 326 | # First token |
| 327 | first_token_time = time.perf_counter_ns() |
| 328 | ttft = first_token_time - start_time |
| 329 | first_chunk = delta["content"] |
| 330 | else: |
| 331 | # Decoding phase |
| 332 | chunk_delay.append(timestamp - most_recent_timestamp) |
| 333 | |
| 334 | generated_text += delta["content"] |
| 335 | |
| 336 | most_recent_timestamp = timestamp |
| 337 | else: |
| 338 | valid_response = False |
| 339 | content = await response.text() |
| 340 | logger.warning( |
| 341 | f"{Color.YELLOW}Received HTTP status {http_status.value} " |
| 342 | f"({http_status.phrase}): {content}{Color.RESET}" |
| 343 | ) |
| 344 | |
| 345 | if latency is None: |
| 346 | latency = -1.0 |
| 347 | if valid_response: |
| 348 | # Streaming is disabled, latency was not set |
| 349 | latency = time.perf_counter_ns() - start_time |
| 350 | |
| 351 | if ttft is None: |
| 352 | if stream: |
| 353 | valid_response = False |
| 354 | # The response was a single chunk |
| 355 | ttft = latency |
| 356 | |
| 357 | # Each chunk may include more than one token |
| 358 | tpot: float = mean(chunk_delay) if len(chunk_delay) > 0 else 0.0 |
| 359 | num_chunks: int = len(chunk_delay) |
| 360 | |
| 361 | sr = ServerResponse( |
| 362 | valid=valid_response, |
| 363 | ttft_ms=nanosec_to_millisec(ttft) if ttft > 0.0 else -1.0, |
| 364 | tpot_ms=nanosec_to_millisec(tpot), |
| 365 | latency_ms=nanosec_to_millisec(latency), |
| 366 | start_time_ms=nanosec_to_millisec(start_time), |
| 367 | first_chunk=first_chunk, |
| 368 | content=generated_text, |
| 369 | num_chunks=num_chunks, |
| 370 | ) |
| 371 | return sr |
| 372 | |
| 373 | |
| 374 | def get_short_string(input: str) -> str: |
| 375 | n = 20 |
| 376 | if len(input) < 400: |
| 377 | return input |
| 378 | |
| 379 | return f"{input[:n]}...{input[-n:]}" |
| 380 | |
| 381 | |
| 382 | def get_token_count(tokenizer: AutoTokenizer, text: str) -> int: |
| 383 | return len(tokenizer(text, add_special_tokens=False).input_ids) |
| 384 | |
| 385 | |
| 386 | def get_messages_token_count( |
| 387 | tokenizer: AutoTokenizer, messages: list[dict[str, str]] |
| 388 | ) -> int: |
| 389 | token_count = 0 |
| 390 | for m in messages: |
| 391 | token_count += get_token_count(tokenizer, m["content"]) |
| 392 | |
| 393 | return token_count |
| 394 | |
| 395 | |
| 396 | async def send_turn( |
| 397 | session: aiohttp.ClientSession, |
| 398 | client_id: int, |
| 399 | conv_id: str, |
| 400 | conversation_messages: MessagesList, |
| 401 | messages_to_use: int, |
| 402 | tokenizer: AutoTokenizer, |
| 403 | req_args: RequestArgs, |
| 404 | verbose: bool, |
| 405 | verify_output: bool, |
| 406 | ) -> RequestStats | None: |
| 407 | assert messages_to_use > 0 |
| 408 | assert messages_to_use <= len(conversation_messages) |
| 409 | |
| 410 | messages = conversation_messages[:messages_to_use] |
| 411 | |
| 412 | # Index of the next message (the role should be "user") |
| 413 | index = messages_to_use - 1 |
| 414 | |
| 415 | # Verify that the message has only two keys, "role" and "content" |
| 416 | assert len(messages[index].keys()) == 2 |
| 417 | assert "role" in messages[index] and "content" in messages[index] |
| 418 | assert messages[index]["role"] == "user", ( |
| 419 | f"Failed on conversation ID {conv_id}, message role should be user" |
| 420 | ) |
| 421 | |
| 422 | if verbose: |
| 423 | print( |
| 424 | f"{Color.CYAN}Messages (conversation ID {conv_id}," |
| 425 | f" {len(messages)} turns):{Color.RESET}", |
| 426 | messages, |
| 427 | ) |
| 428 | |
| 429 | # None means that there is no upper/lower limit for the output token count |
| 430 | min_tokens = None if req_args.limit_min_tokens < 0 else req_args.limit_min_tokens |
| 431 | max_tokens = None if req_args.limit_max_tokens < 0 else req_args.limit_max_tokens |
| 432 | |
| 433 | if len(conversation_messages) > messages_to_use: |
| 434 | # The conversation contains an assistant answer for the next user prompt |
| 435 | if ( |
| 436 | min_tokens == NUM_TOKENS_FROM_DATASET |
| 437 | or max_tokens == NUM_TOKENS_FROM_DATASET |
| 438 | ): |
| 439 | # Compute number of tokens in the answer (from the input conversation) |
| 440 | assistant_answer = conversation_messages[messages_to_use] |
| 441 | answer_num_tokens = get_token_count(tokenizer, assistant_answer["content"]) |
| 442 | assert assistant_answer["role"] == "assistant" |
| 443 | |
| 444 | if min_tokens == NUM_TOKENS_FROM_DATASET: |
| 445 | min_tokens = max(1, answer_num_tokens) |
| 446 | |
| 447 | if max_tokens == NUM_TOKENS_FROM_DATASET: |
| 448 | max_tokens = max(1, answer_num_tokens) |
| 449 | |
| 450 | # Send the current conversation to LLM and get a response |
| 451 | response: ServerResponse = await send_request( |
| 452 | session, |
| 453 | messages, |
| 454 | req_args.chat_url, |
| 455 | req_args.model, |
| 456 | req_args.stream, |
| 457 | min_tokens, |
| 458 | max_tokens, |
| 459 | req_args.timeout_sec, |
| 460 | conversation_id=conv_id if req_args.send_conversation_id else None, |
| 461 | headers=req_args.headers, |
| 462 | ) |
| 463 | |
| 464 | if response.valid is False: |
| 465 | # Request failed |
| 466 | return None |
| 467 | |
| 468 | # Compute number of tokens in input / output |
| 469 | input_num_tokens = get_messages_token_count(tokenizer, messages) |
| 470 | |
| 471 | # Num tokens in the user's last question |
| 472 | question_num_tokens = get_token_count(tokenizer, messages[index]["content"]) |
| 473 | |
| 474 | # Num tokens in the history/context of the question |
| 475 | assert input_num_tokens >= question_num_tokens |
| 476 | history_num_tokens = input_num_tokens - question_num_tokens |
| 477 | |
| 478 | # Num tokens in the LLM's answer (first chunk and full answer) |
| 479 | first_chunk_tokens = get_token_count(tokenizer, response.first_chunk) |
| 480 | |
| 481 | output_content = response.content |
| 482 | output_num_tokens = get_token_count(tokenizer, output_content) |
| 483 | |
| 484 | # Prefix caching approximated cached percent |
| 485 | approx_cached_percent = ( |
| 486 | 100.0 * (history_num_tokens / input_num_tokens) if input_num_tokens > 0 else 0.0 |
| 487 | ) |
| 488 | |
| 489 | # Compute the correct TTFT and TPOT (based on tokens and not chunks). |
| 490 | # Required because multiple output tokens may be bundled in a single chunk. |
| 491 | if output_num_tokens > 1 and output_num_tokens > first_chunk_tokens: |
| 492 | # More than one token and more than one chunk in the output |
| 493 | decode_ms = response.latency_ms - response.ttft_ms |
| 494 | decode_num_tokens = output_num_tokens - first_chunk_tokens |
| 495 | tpot_ms = decode_ms / decode_num_tokens |
| 496 | else: |
| 497 | # In this case: output_num_tokens == first_chunk_tokens |
| 498 | # Output was a single chunk (output_num_tokens > 1) |
| 499 | # or even a single token (output_num_tokens == 1) |
| 500 | tpot_ms = 0.0 |
| 501 | |
| 502 | if first_chunk_tokens > 1: |
| 503 | # First chunk had multiple tokens, adjust TTFT for a single token |
| 504 | delta_ms = (first_chunk_tokens - 1) * tpot_ms |
| 505 | ttft_ms = max(0.1, response.ttft_ms - delta_ms) |
| 506 | else: |
| 507 | # First chunk had only one token |
| 508 | ttft_ms = response.ttft_ms |
| 509 | |
| 510 | rs = RequestStats( |
| 511 | ttft_ms=ttft_ms, |
| 512 | tpot_ms=tpot_ms, |
| 513 | latency_ms=response.latency_ms, |
| 514 | start_time_ms=response.start_time_ms, |
| 515 | input_num_turns=len(messages), |
| 516 | input_num_tokens=input_num_tokens, |
| 517 | output_num_tokens=output_num_tokens, |
| 518 | output_num_chunks=response.num_chunks, |
| 519 | output_num_first_chunk_tokens=first_chunk_tokens, |
| 520 | approx_cached_percent=approx_cached_percent, |
| 521 | conversation_id=conv_id, |
| 522 | client_id=client_id, |
| 523 | ) |
| 524 | |
| 525 | if verbose: |
| 526 | print( |
| 527 | f"\n{Color.YELLOW}Response ({output_num_tokens} tokens):{Color.RESET}", |
| 528 | output_content, |
| 529 | ) |
| 530 | print(f"{Color.YELLOW}Response metrics: {rs}{Color.RESET}") |
| 531 | print("-" * 70) |
| 532 | |
| 533 | # Save the LLM's answer (will be used as part of the context for the next user turn) |
| 534 | answer_index = messages_to_use |
| 535 | if len(conversation_messages) > answer_index: |
| 536 | assert conversation_messages[answer_index]["role"] == "assistant", ( |
| 537 | f"Failed on conversation ID {conv_id}, message role should be assistant" |
| 538 | ) |
| 539 | |
| 540 | orig_content = conversation_messages[answer_index]["content"] |
| 541 | if verify_output: |
| 542 | # Compare the new answer to the answer from the input file |
| 543 | debug_info = ( |
| 544 | f"LLM/dataset answers do not match ({conv_id}):" |
| 545 | f"\n'{get_short_string(output_content)}' (len: {len(output_content)})," |
| 546 | f"\n'{get_short_string(orig_content)}' (len: {len(orig_content)})" |
| 547 | ) |
| 548 | if orig_content != output_content: |
| 549 | raise ValueError(debug_info) |
| 550 | |
| 551 | # Update the answer |
| 552 | conversation_messages[answer_index]["content"] = output_content |
| 553 | else: |
| 554 | # A user prompt that has no answer, add the answer as a new message |
| 555 | new_answer = {"role": "assistant", "content": output_content} |
| 556 | conversation_messages.append(new_answer) |
| 557 | |
| 558 | return rs |
| 559 | |
| 560 | |
| 561 | async def poisson_sleep(request_rate: float, verbose: bool = False) -> None: |
| 562 | # Generate a random time interval from the Poisson distribution |
| 563 | assert request_rate > 0 |
| 564 | |
| 565 | interval = np.random.exponential(1.0 / request_rate) |
| 566 | if verbose: |
| 567 | logger.info(f"Sleeping for {interval:.3f} seconds...") |
| 568 | await asyncio.sleep(interval) |
| 569 | |
| 570 | |
| 571 | async def exponential_backoff_sleep( |
| 572 | attempt_cnt: int, |
| 573 | base_rate: float = 1.0, |
| 574 | backoff_factor: float = 2.0, |
| 575 | jitter_fraction: float = 0.10, |
| 576 | verbose: bool = False, |
| 577 | ) -> None: |
| 578 | # Sleep with exponential backoff and jitter after a failed request. |
| 579 | backoff_delay = base_rate * (backoff_factor**attempt_cnt) |
| 580 | jittered_delay = backoff_delay * ( |
| 581 | 1 + np.random.uniform(-jitter_fraction, jitter_fraction) |
| 582 | ) |
| 583 | |
| 584 | if verbose: |
| 585 | logger.info(f"Backoff for {jittered_delay:.3f} seconds...") |
| 586 | |
| 587 | await asyncio.sleep(jittered_delay) |
| 588 | |
| 589 | |
| 590 | async def client_main( |
| 591 | args: ClientArgs, |
| 592 | req_args: RequestArgs, |
| 593 | client_id: int, |
| 594 | tokenizer: AutoTokenizer, |
| 595 | stop_event: mp.Event, # type: ignore |
| 596 | task_queue: mp.Queue, |
| 597 | result_queue: mp.Queue, |
| 598 | conv_queue: mp.Queue, |
| 599 | ) -> None: |
| 600 | logger.info( |
| 601 | f"{Color.CYAN}Started client {client_id}: max_num_requests={args.max_num_requests}, max_active_conversations={args.max_active_conversations}{Color.RESET}" # noqa: E501 |
| 602 | ) |
| 603 | |
| 604 | # Set unique seed per client (each client runs in its own process) |
| 605 | # Add 1 to ensure no client uses the same seed as the main process |
| 606 | client_seed = args.seed + client_id + 1 |
| 607 | random.seed(client_seed) |
| 608 | np.random.seed(client_seed) |
| 609 | |
| 610 | # Active conversations |
| 611 | active_convs: ConversationsMap = {} |
| 612 | conv_id_queue: deque = deque(maxlen=args.max_active_conversations) |
| 613 | |
| 614 | # Keep track of how many messages have been used for each conversation |
| 615 | turns_count: Counter = Counter() |
| 616 | num_successes = 0 |
| 617 | num_failures = 0 |
| 618 | |
| 619 | # Track the timestamp (time.perf_counter()) |
| 620 | # of the last turn per conversation (only for debug) |
| 621 | time_of_last_turn: dict[ConvId, float] = {} |
| 622 | |
| 623 | # Flag that indicates that there are no new tasks (conversations) for the client |
| 624 | task_queue_empty = False |
| 625 | |
| 626 | async with aiohttp.ClientSession() as session: |
| 627 | # Print progress |
| 628 | |
| 629 | while task_queue_empty is False: |
| 630 | result = None |
| 631 | |
| 632 | if ( |
| 633 | args.max_num_requests |
| 634 | and num_successes + num_failures == args.max_num_requests |
| 635 | ): |
| 636 | logger.info( |
| 637 | f"{Color.YELLOW}Client {client_id} reached " |
| 638 | f"request limit{Color.RESET}" |
| 639 | ) |
| 640 | break |
| 641 | |
| 642 | if stop_event.is_set(): # type: ignore |
| 643 | logger.info( |
| 644 | f"{Color.YELLOW}Client {client_id} received " |
| 645 | f"a termination signal{Color.RESET}" |
| 646 | ) |
| 647 | break |
| 648 | |
| 649 | while ( |
| 650 | len(active_convs) < args.max_active_conversations |
| 651 | and task_queue_empty is False |
| 652 | ): |
| 653 | # Get a new conversation from the task queue |
| 654 | conv_id, messages = task_queue.get() |
| 655 | |
| 656 | if conv_id is TERM_SIGNAL: |
| 657 | task_queue_empty = True |
| 658 | break |
| 659 | |
| 660 | if args.skip_first_turn: |
| 661 | # Skip the first turn (both user and assistant), |
| 662 | # relevant if warmup was enabled. |
| 663 | # Default turns_count[conv_id] will be zero if conv_id |
| 664 | # was never inserted/updated in turns_count. |
| 665 | turns_count[conv_id] += 2 |
| 666 | |
| 667 | if turns_count[conv_id] < len(messages): |
| 668 | # Add new conversation |
| 669 | active_convs[conv_id] = messages |
| 670 | conv_id_queue.append(conv_id) |
| 671 | |
| 672 | if args.verbose: |
| 673 | logger.info( |
| 674 | f"{Color.GREEN}Client {client_id} will use conversation ID {conv_id} (active conversations {len(active_convs)}){Color.RESET}" # noqa: E501 |
| 675 | ) |
| 676 | |
| 677 | elif args.verbose: |
| 678 | # No more messages (conversation finished during the warmup) |
| 679 | logger.info( |
| 680 | f"{Color.YELLOW}Client {client_id} will not use conversation ID {conv_id} (all {len(messages)} messages already sent){Color.RESET}" # noqa: E501 |
| 681 | ) |
| 682 | |
| 683 | if len(active_convs) == 0 or task_queue_empty: |
| 684 | logger.info( |
| 685 | f"{Color.YELLOW}Client {client_id} has no more work{Color.RESET}" |
| 686 | ) |
| 687 | break |
| 688 | |
| 689 | # Pick an active conversation for the next request |
| 690 | if args.conversation_sampling == ConversationSampling.ROUND_ROBIN: |
| 691 | conv_id = conv_id_queue.pop() |
| 692 | else: |
| 693 | # ConversationSampling.RANDOM |
| 694 | active_ids = list(active_convs.keys()) |
| 695 | conv_id = random.choice(active_ids) |
| 696 | |
| 697 | messages = active_convs[conv_id] |
| 698 | assert isinstance(messages, list) and len(messages) > 0 |
| 699 | |
| 700 | # Update the amount of messages to use |
| 701 | turns_count[conv_id] += 1 |
| 702 | current_turn = turns_count[conv_id] |
| 703 | |
| 704 | assert current_turn < len(messages), ( |
| 705 | f"Turn number {current_turn} is invalid for conversation ID {conv_id}" |
| 706 | f" that has only {len(messages)} messages" |
| 707 | ) |
| 708 | |
| 709 | if args.verbose: |
| 710 | curr_time_sec: float = time.perf_counter() |
| 711 | time_since_last_turn: str | float = "N/A" |
| 712 | if conv_id in time_of_last_turn: |
| 713 | time_since_last_turn = round( |
| 714 | curr_time_sec - time_of_last_turn[conv_id], 3 |
| 715 | ) |
| 716 | logger.info( |
| 717 | f"Client {client_id} using conversation ID {conv_id} (turn: {current_turn}, time since last turn [sec]: {time_since_last_turn})" # noqa: E501 |
| 718 | ) |
| 719 | time_of_last_turn[conv_id] = curr_time_sec |
| 720 | |
| 721 | success = False |
| 722 | for attempt_cnt in range(args.max_retries + 1): |
| 723 | try: |
| 724 | exception = False |
| 725 | result = await send_turn( |
| 726 | session, |
| 727 | client_id, |
| 728 | conv_id, |
| 729 | messages, |
| 730 | current_turn, |
| 731 | tokenizer, |
| 732 | req_args, |
| 733 | args.print_content, |
| 734 | args.verify_output, |
| 735 | ) |
| 736 | if result is not None: |
| 737 | result_queue.put(result) |
| 738 | success = True |
| 739 | break |
| 740 | else: |
| 741 | logger.warning( |
| 742 | f"{Color.YELLOW}Client {client_id} - Request rejected during conversation ID {conv_id} (turn: {current_turn}){Color.RESET}" # noqa: E501 |
| 743 | ) |
| 744 | except asyncio.exceptions.TimeoutError: |
| 745 | exception = True |
| 746 | logger.error( |
| 747 | "%sClient %d - Timeout during conversation ID %s (turn: %d). " |
| 748 | "Base timeout is %ss (set with --request-timeout-sec), but the " |
| 749 | "effective timeout may be longer based on max_tokens. If this " |
| 750 | "is unexpected, consider increasing the timeout or checking " |
| 751 | "model performance.%s", |
| 752 | Color.RED, |
| 753 | client_id, |
| 754 | conv_id, |
| 755 | current_turn, |
| 756 | req_args.timeout_sec, |
| 757 | Color.RESET, |
| 758 | ) |
| 759 | except Exception: |
| 760 | exception = True |
| 761 | logger.exception( |
| 762 | f"{Color.RED}Client {client_id} - Exception during conversation ID {conv_id} (turn: {current_turn}){Color.RESET}" # noqa: E501 |
| 763 | ) |
| 764 | |
| 765 | # Sleep before retry if not last attempt |
| 766 | if not success and attempt_cnt < args.max_retries: |
| 767 | await exponential_backoff_sleep(attempt_cnt, verbose=args.verbose) |
| 768 | |
| 769 | if not success: |
| 770 | num_failures += 1 |
| 771 | # Remove the conversation (should not be used again) |
| 772 | active_convs.pop(conv_id) |
| 773 | if exception: |
| 774 | break # Exit gracefully instead of raising an error |
| 775 | |
| 776 | else: |
| 777 | num_successes += 1 |
| 778 | |
| 779 | # Update the turns counter to include the LLM response |
| 780 | # The LLM response will be used as context for the next user turn |
| 781 | turns_count[conv_id] += 1 |
| 782 | |
| 783 | max_turns = len(messages) |
| 784 | if args.max_turns is not None: |
| 785 | # Limit the number of turns in the conversation |
| 786 | max_turns = min(args.max_turns, max_turns) |
| 787 | |
| 788 | if turns_count[conv_id] >= max_turns: |
| 789 | # Conversation has no more turns (no longer active) |
| 790 | # save the updated conversation (with the LLM server's answer) |
| 791 | conv_queue.put((conv_id, active_convs.pop(conv_id))) |
| 792 | if args.verbose: |
| 793 | logger.info( |
| 794 | f"{Color.GREEN}Client {client_id} finished " |
| 795 | f"conversation ID {conv_id}{Color.RESET}" |
| 796 | ) |
| 797 | else: |
| 798 | # Conversation is not finished, insert it at the back of the queue |
| 799 | conv_id_queue.appendleft(conv_id) |
| 800 | |
| 801 | # Sleep between requests (if lambda is positive) |
| 802 | if args.request_rate > 0: |
| 803 | await poisson_sleep(args.request_rate, args.verbose) |
| 804 | |
| 805 | # Send indication that the client is done |
| 806 | conv_queue.put((TERM_SIGNAL, TERM_SIGNAL)) |
| 807 | |
| 808 | logger.info( |
| 809 | f"{Color.CYAN}Client {client_id} is done " |
| 810 | f"({num_successes=}, {num_failures=}){Color.RESET}" |
| 811 | ) |
| 812 | |
| 813 | |
| 814 | def worker_function( |
| 815 | client_id: int, |
| 816 | tokenizer: AutoTokenizer, |
| 817 | client_args: ClientArgs, |
| 818 | req_args: RequestArgs, |
| 819 | stop_event: mp.Event, # type: ignore |
| 820 | task_queue: mp.Queue, |
| 821 | result_queue: mp.Queue, |
| 822 | conv_queue: mp.Queue, |
| 823 | ) -> None: |
| 824 | asyncio.run( |
| 825 | client_main( |
| 826 | client_args, |
| 827 | req_args, |
| 828 | client_id, |
| 829 | tokenizer, |
| 830 | stop_event, |
| 831 | task_queue, |
| 832 | result_queue, |
| 833 | conv_queue, |
| 834 | ) |
| 835 | ) |
| 836 | |
| 837 | |
| 838 | def get_client_config( |
| 839 | args: argparse.Namespace, input_conv: ConversationsMap |
| 840 | ) -> tuple[ClientArgs, RequestArgs]: |
| 841 | if args.num_clients < 1: |
| 842 | raise ValueError("Number of clients must be a positive number") |
| 843 | |
| 844 | if len(input_conv) < args.num_clients: |
| 845 | raise ValueError( |
| 846 | "Number of conversations must be equal or larger than the number of clients" |
| 847 | ) |
| 848 | |
| 849 | max_req_per_client: int | None = None |
| 850 | if args.max_num_requests is not None: |
| 851 | # Max number of requests per client |
| 852 | req_per_client = args.max_num_requests // args.num_clients |
| 853 | if req_per_client < 1: |
| 854 | raise ValueError("Number of requests should be at least one per client") |
| 855 | max_req_per_client = req_per_client |
| 856 | |
| 857 | max_active_conversations = args.max_active_conversations |
| 858 | if max_active_conversations is None: |
| 859 | # Each client will have only one active conversation at a time |
| 860 | max_active_conversations = args.num_clients |
| 861 | |
| 862 | if max_active_conversations > len(input_conv): |
| 863 | raise ValueError( |
| 864 | f"Max active conversations {max_active_conversations} " |
| 865 | "must be equal or less than the total number of conversations" |
| 866 | ) |
| 867 | |
| 868 | # Max number of active conversations per client |
| 869 | max_active_conv_per_client = max_active_conversations // args.num_clients |
| 870 | if max_active_conv_per_client < 1: |
| 871 | raise ValueError( |
| 872 | f"Max active conversations {max_active_conversations} " |
| 873 | "must be equal or greater than the number of clients" |
| 874 | ) |
| 875 | |
| 876 | # Skip the first user turn (as part of the warmup) |
| 877 | skip_first_turn = args.warmup_step |
| 878 | |
| 879 | # Common arguments for all clients |
| 880 | client_args = ClientArgs( |
| 881 | seed=args.seed, |
| 882 | max_num_requests=max_req_per_client, |
| 883 | skip_first_turn=skip_first_turn, |
| 884 | max_turns=args.max_turns, |
| 885 | max_active_conversations=max_active_conv_per_client, |
| 886 | verbose=args.verbose, |
| 887 | print_content=args.print_content, |
| 888 | verify_output=args.verify_output, |
| 889 | conversation_sampling=args.conversation_sampling, |
| 890 | request_rate=args.request_rate, |
| 891 | max_retries=args.max_retries, |
| 892 | ) |
| 893 | |
| 894 | if args.limit_min_tokens > 0 or args.limit_max_tokens > 0: |
| 895 | if args.limit_min_tokens < 1 or args.limit_max_tokens < 1: |
| 896 | raise ValueError( |
| 897 | "Invalid min/max tokens limits (both limits should be provided)" |
| 898 | ) |
| 899 | if args.limit_min_tokens > args.limit_max_tokens: |
| 900 | raise ValueError( |
| 901 | "Invalid min/max tokens limits (min should not be larger than max)" |
| 902 | ) |
| 903 | |
| 904 | if args.request_timeout_sec <= 0: |
| 905 | raise ValueError("Request timeout must be a positive number") |
| 906 | |
| 907 | # Arguments for API requests |
| 908 | chat_url = f"{args.url}/v1/chat/completions" |
| 909 | model_name = args.served_model_name if args.served_model_name else args.model |
| 910 | headers = build_request_headers(args.api_key, args.header) |
| 911 | |
| 912 | req_args = RequestArgs( |
| 913 | chat_url=chat_url, |
| 914 | model=model_name, |
| 915 | stream=not args.no_stream, |
| 916 | limit_min_tokens=args.limit_min_tokens, |
| 917 | limit_max_tokens=args.limit_max_tokens, |
| 918 | timeout_sec=args.request_timeout_sec, |
| 919 | send_conversation_id=args.send_conversation_id, |
| 920 | headers=headers, |
| 921 | ) |
| 922 | |
| 923 | return client_args, req_args |
| 924 | |
| 925 | |
| 926 | async def main_mp( |
| 927 | client_args: ClientArgs, |
| 928 | req_args: RequestArgs, |
| 929 | bench_args: BenchmarkArgs, |
| 930 | tokenizer: AutoTokenizer, |
| 931 | input_conv: ConversationsMap, |
| 932 | ) -> tuple[ConversationsMap, list[RequestStats]]: |
| 933 | # An event that will trigger graceful termination of all the clients |
| 934 | stop_event = mp.Event() |
| 935 | |
| 936 | # Queue for input conversations (from the input file/dataset) |
| 937 | task_queue: mp.Queue = mp.Queue() |
| 938 | |
| 939 | # Queue for client measurements (TTFT, TPOT, etc. for each request) |
| 940 | result_queue: mp.Queue = mp.Queue() |
| 941 | |
| 942 | # Queue for output conversations (with the LLM answers, sent by the server) |
| 943 | conv_queue: mp.Queue = mp.Queue() |
| 944 | output_conv: ConversationsMap = {} |
| 945 | client_metrics: list[RequestStats] = [] |
| 946 | |
| 947 | # Start all clients |
| 948 | start_time = time.perf_counter_ns() |
| 949 | logger.info(f"{Color.GREEN}Starting {bench_args.num_clients} clients{Color.RESET}") |
| 950 | |
| 951 | clients = [] |
| 952 | for client_id in range(bench_args.num_clients): |
| 953 | client = mp.Process( |
| 954 | name=f"client_{client_id}", |
| 955 | target=worker_function, |
| 956 | args=( |
| 957 | client_id, |
| 958 | tokenizer, |
| 959 | client_args, |
| 960 | req_args, |
| 961 | stop_event, |
| 962 | task_queue, |
| 963 | result_queue, |
| 964 | conv_queue, |
| 965 | ), |
| 966 | ) |
| 967 | clients.append(client) |
| 968 | client.start() |
| 969 | |
| 970 | # Submit all the input conversations as tasks for the clients |
| 971 | for conv_id, messages in input_conv.items(): |
| 972 | task_queue.put((conv_id, messages)) |
| 973 | |
| 974 | # Add termination signals for clients |
| 975 | for _ in range(bench_args.num_clients): |
| 976 | task_queue.put((TERM_SIGNAL, TERM_SIGNAL)) |
| 977 | |
| 978 | # Collect the updated conversations from all clients |
| 979 | num_clients_finished = 0 |
| 980 | total_convs = len(input_conv) |
| 981 | |
| 982 | debug_stats = DebugStats(logger, min(15 * bench_args.num_clients, 500)) |
| 983 | |
| 984 | while num_clients_finished < bench_args.num_clients: |
| 985 | # Collect updated conversation |
| 986 | conv_id, messages = conv_queue.get() |
| 987 | |
| 988 | # Collect results (measurements) |
| 989 | while not result_queue.empty(): |
| 990 | new_data = result_queue.get() |
| 991 | client_metrics.append(new_data) |
| 992 | debug_stats.update(new_data) |
| 993 | |
| 994 | if conv_id is TERM_SIGNAL: |
| 995 | num_clients_finished += 1 |
| 996 | logger.info( |
| 997 | f"{Color.CYAN}{num_clients_finished} out of " |
| 998 | f"{bench_args.num_clients} clients finished{Color.RESET}" |
| 999 | ) |
| 1000 | |
| 1001 | if bench_args.early_stop and not stop_event.is_set(): |
| 1002 | # Once one client finished, stop all other clients. |
| 1003 | # there is no reason to continue the benchmark with fewer clients. |
| 1004 | logger.info( |
| 1005 | f"{Color.YELLOW}Sending termination signal to clients{Color.RESET}" |
| 1006 | ) |
| 1007 | stop_event.set() |
| 1008 | else: |
| 1009 | output_conv[conv_id] = messages |
| 1010 | |
| 1011 | finished_convs = len(output_conv) |
| 1012 | percent = finished_convs / total_convs |
| 1013 | |
| 1014 | # Tuned to control the print rate (can be changed if required) |
| 1015 | print_cycle = max(3, int(bench_args.num_clients / 4)) |
| 1016 | |
| 1017 | if finished_convs % print_cycle == 0: |
| 1018 | runtime_sec = nanosec_to_sec(time.perf_counter_ns() - start_time) |
| 1019 | logger.info( |
| 1020 | f"{Color.CYAN}Finished {finished_convs} out of {total_convs} conversations ({percent:.0%}), " # noqa: E501 |
| 1021 | f"{num_clients_finished} out of {bench_args.num_clients} clients finished, collected {len(client_metrics)} measurements, runtime {runtime_sec:.3f} sec{Color.RESET}" # noqa: E501 |
| 1022 | ) |
| 1023 | |
| 1024 | rps: str | float = round(len(client_metrics) / runtime_sec, 3) |
| 1025 | if len(client_metrics) < (5 * bench_args.num_clients): |
| 1026 | # Do not estimate the RPS if the number of samples is very low |
| 1027 | # (threshold can be tuned if needed) |
| 1028 | rps = "N/A" |
| 1029 | |
| 1030 | runtime_left_sec: str | float = round( |
| 1031 | (runtime_sec / finished_convs) * (total_convs - finished_convs), 3 |
| 1032 | ) |
| 1033 | if percent < 0.05: |
| 1034 | # If less than 5% of the conversations were not finished, |
| 1035 | # the estimation will probably be very inaccurate |
| 1036 | # (threshold can be tuned if needed). |
| 1037 | runtime_left_sec = "N/A" |
| 1038 | |
| 1039 | logger.info( |
| 1040 | f"{Color.CYAN}Estimated req/sec {rps}, estimated runtime left {runtime_left_sec} sec{Color.RESET}" # noqa: E501 |
| 1041 | ) |
| 1042 | debug_stats.print() |
| 1043 | |
| 1044 | logger.info( |
| 1045 | f"{Color.CYAN}All {bench_args.num_clients} clients finished{Color.RESET}" |
| 1046 | ) |
| 1047 | |
| 1048 | # At this point all the clients finished, |
| 1049 | # collect results (TTFT, TPOT, etc.) from all the clients. |
| 1050 | # This needs to happen before calling join on the clients |
| 1051 | # (result_queue should be emptied). |
| 1052 | while not result_queue.empty(): |
| 1053 | client_metrics.append(result_queue.get()) |
| 1054 | |
| 1055 | logger.info(f"Collected {len(client_metrics)} samples from all the clients") |
| 1056 | |
| 1057 | # Wait for all clients to finish |
| 1058 | for client in clients: |
| 1059 | logger.info( |
| 1060 | f"{Color.CYAN}Waiting for client {client.name} " |
| 1061 | f"(is alive: {client.is_alive()}){Color.RESET}" |
| 1062 | ) |
| 1063 | |
| 1064 | client.join(timeout=req_args.timeout_sec + 1) |
| 1065 | |
| 1066 | if client.is_alive(): |
| 1067 | logger.warning( |
| 1068 | f"{Color.YELLOW}Client {client.name} will be terminated{Color.RESET}" |
| 1069 | ) |
| 1070 | client.terminate() |
| 1071 | |
| 1072 | exitcode = client.exitcode |
| 1073 | if exitcode != 0: |
| 1074 | logger.error( |
| 1075 | f"{Color.RED}Client {client.name} exited " |
| 1076 | f"with exit code {exitcode}{Color.RESET}" |
| 1077 | ) |
| 1078 | |
| 1079 | logger.info( |
| 1080 | f"All {bench_args.num_clients} clients exited (successfully " |
| 1081 | f"finished {len(output_conv)} out of {total_convs} conversations)" |
| 1082 | ) |
| 1083 | |
| 1084 | # Queues should be closed, required to avoid hang at interpreter shutdown |
| 1085 | unfinished_tasks = 0 |
| 1086 | while not task_queue.empty(): |
| 1087 | task_queue.get() |
| 1088 | unfinished_tasks += 1 |
| 1089 | |
| 1090 | if unfinished_tasks > 0: |
| 1091 | # Can happen if not all tasks (conversations) have finished. |
| 1092 | # May happen if --max-num-requests was used, |
| 1093 | # or if an error occurred in one of the clients. |
| 1094 | logger.debug(f"Discarding {unfinished_tasks} unfinished tasks") |
| 1095 | |
| 1096 | task_queue.close() |
| 1097 | task_queue.join_thread() |
| 1098 | |
| 1099 | result_queue.close() |
| 1100 | result_queue.join_thread() |
| 1101 | |
| 1102 | conv_queue.close() |
| 1103 | conv_queue.join_thread() |
| 1104 | |
| 1105 | return output_conv, client_metrics |
| 1106 | |
| 1107 | |
| 1108 | def get_filename_with_timestamp(label: str, extension: str) -> str: |
| 1109 | time_now = datetime.now() |
| 1110 | timestamp = time_now.strftime("%d-%m-%Y_%H-%M-%S") |
| 1111 | filename = f"{label}__{timestamp}.{extension}" |
| 1112 | return filename |
| 1113 | |
| 1114 | |
| 1115 | def process_statistics( |
| 1116 | client_metrics: list[RequestStats], |
| 1117 | warmup_percentages: list[float], |
| 1118 | test_params: dict, |
| 1119 | verbose: bool, |
| 1120 | gen_conv_args: GenConvArgs | None = None, |
| 1121 | excel_output: bool = False, |
| 1122 | warmup_runtime_sec: float | None = None, |
| 1123 | ) -> None: |
| 1124 | if len(client_metrics) == 0: |
| 1125 | logger.info("No samples to process") |
| 1126 | return |
| 1127 | |
| 1128 | logger.info(f"Processing {len(client_metrics)} samples...") |
| 1129 | |
| 1130 | raw_data = pd.DataFrame(client_metrics) |
| 1131 | |
| 1132 | if verbose: |
| 1133 | # Calculate the time between user turns in each conversation (in a new column) |
| 1134 | raw_data = raw_data.sort_values(by=["conversation_id", "start_time_ms"]) |
| 1135 | raw_data["time_between_user_turns_sec"] = raw_data.groupby("conversation_id")[ |
| 1136 | "start_time_ms" |
| 1137 | ].diff() |
| 1138 | |
| 1139 | # Convert milliseconds to seconds |
| 1140 | raw_data["time_between_user_turns_sec"] = ( |
| 1141 | raw_data["time_between_user_turns_sec"] / 1000.0 |
| 1142 | ) |
| 1143 | |
| 1144 | # Final raw data should be sorted by time |
| 1145 | raw_data = raw_data.sort_values(by=["start_time_ms"]) |
| 1146 | raw_data["end_time_ms"] = raw_data["start_time_ms"] + raw_data["latency_ms"] |
| 1147 | |
| 1148 | percentiles = [0.25, 0.5, 0.75, 0.9] |
| 1149 | |
| 1150 | # Add more percentiles if there are enough samples |
| 1151 | if len(raw_data) >= 100: |
| 1152 | percentiles.append(0.99) |
| 1153 | |
| 1154 | if len(raw_data) >= 1000: |
| 1155 | percentiles.append(0.999) |
| 1156 | |
| 1157 | if len(raw_data) >= 10000: |
| 1158 | percentiles.append(0.9999) |
| 1159 | |
| 1160 | # Set precision for numbers in the output text (the dataframes) |
| 1161 | pd.set_option("display.precision", 2) |
| 1162 | |
| 1163 | # Exclude parameters from RequestStats |
| 1164 | exclude = [ |
| 1165 | "start_time_ms", |
| 1166 | "end_time_ms", |
| 1167 | "output_num_first_chunk_tokens", |
| 1168 | "approx_cached_percent", |
| 1169 | "conversation_id", |
| 1170 | "client_id", |
| 1171 | ] |
| 1172 | |
| 1173 | print(TEXT_SEPARATOR) |
| 1174 | print(f"{Color.YELLOW}Parameters:{Color.RESET}") |
| 1175 | for k, v in test_params.items(): |
| 1176 | print(f"{k}={v}") |
| 1177 | |
| 1178 | # conversations generation parameters |
| 1179 | if gen_conv_args is not None: |
| 1180 | gen_params = { |
| 1181 | "text_files": ", ".join(gen_conv_args.text_files), |
| 1182 | "input_num_turns": str(gen_conv_args.input_num_turns), |
| 1183 | "input_common_prefix_num_tokens": str( |
| 1184 | gen_conv_args.input_common_prefix_num_tokens |
| 1185 | ), |
| 1186 | "input_prefix_num_tokens": str(gen_conv_args.input_prefix_num_tokens), |
| 1187 | "input_num_tokens": str(gen_conv_args.input_num_tokens), |
| 1188 | "output_num_tokens": str(gen_conv_args.output_num_tokens), |
| 1189 | } |
| 1190 | |
| 1191 | print(f"{Color.YELLOW}Conversations Generation Parameters:{Color.RESET}") |
| 1192 | for k, v in gen_params.items(): |
| 1193 | print(f"{k}={v}") |
| 1194 | |
| 1195 | print(TEXT_SEPARATOR) |
| 1196 | |
| 1197 | params_list = [] |
| 1198 | df_list = [] |
| 1199 | for percent in warmup_percentages: |
| 1200 | # Select samples from the end (tail) of the dataframe |
| 1201 | warmup_count = int(percent * len(raw_data)) |
| 1202 | tail_count = len(raw_data) - warmup_count |
| 1203 | if tail_count == 0: |
| 1204 | # No reason to process if the count of samples is zero |
| 1205 | break |
| 1206 | |
| 1207 | df = raw_data.tail(tail_count) |
| 1208 | |
| 1209 | # Runtime is the diff between the end of the last request |
| 1210 | # and the start of the first request |
| 1211 | runtime_sec = df["end_time_ms"].iloc[-1] - df["start_time_ms"].iloc[0] |
| 1212 | |
| 1213 | # Convert milliseconds to seconds |
| 1214 | runtime_sec = runtime_sec / 1000.0 |
| 1215 | requests_per_sec = float(len(df)) / runtime_sec |
| 1216 | params = { |
| 1217 | "runtime_sec": runtime_sec, |
| 1218 | "requests_per_sec": requests_per_sec, |
| 1219 | } |
| 1220 | if warmup_runtime_sec is not None: |
| 1221 | params["warmup_runtime_sec"] = warmup_runtime_sec |
| 1222 | params["total_runtime_incl_warmup_sec"] = runtime_sec + warmup_runtime_sec |
| 1223 | |
| 1224 | # Generate a summary of relevant metrics (and drop irrelevant data) |
| 1225 | df = df.drop(columns=exclude).describe(percentiles=percentiles).transpose() |
| 1226 | |
| 1227 | # List for Excel file |
| 1228 | params_list.append(params) |
| 1229 | df_list.append(df) |
| 1230 | |
| 1231 | # Print the statistics summary |
| 1232 | if percent > 0 or len(warmup_percentages) > 1: |
| 1233 | print( |
| 1234 | f"{Color.YELLOW}Statistics summary " |
| 1235 | f"(assuming {percent:.0%} warmup samples):{Color.RESET}" |
| 1236 | ) |
| 1237 | else: |
| 1238 | print(f"{Color.YELLOW}Statistics summary:{Color.RESET}") |
| 1239 | |
| 1240 | for k, v in params.items(): |
| 1241 | if isinstance(v, float): |
| 1242 | print(f"{k} = {v:.3f}") |
| 1243 | else: |
| 1244 | print(f"{k} = {v}") |
| 1245 | print(TEXT_SEPARATOR) |
| 1246 | print(df) |
| 1247 | print(TEXT_SEPARATOR) |
| 1248 | |
| 1249 | if excel_output: |
| 1250 | prefix = f"statistics_{test_params['num_clients']}_clients" |
| 1251 | filename = get_filename_with_timestamp(prefix, "xlsx") |
| 1252 | |
| 1253 | with pd.ExcelWriter(filename, engine="xlsxwriter") as writer: |
| 1254 | startrow = 0 |
| 1255 | test_params_df = pd.DataFrame([test_params]) |
| 1256 | test_params_df.to_excel( |
| 1257 | writer, sheet_name="Summary", index=False, startrow=startrow |
| 1258 | ) |
| 1259 | startrow += len(test_params_df) + 3 |
| 1260 | |
| 1261 | if gen_conv_args is not None: |
| 1262 | gen_params_df = pd.DataFrame([gen_params]) |
| 1263 | gen_params_df.to_excel( |
| 1264 | writer, sheet_name="Summary", index=False, startrow=(startrow - 1) |
| 1265 | ) |
| 1266 | startrow += len(gen_params_df) + 3 |
| 1267 | |
| 1268 | for params, df_stats in zip(params_list, df_list): |
| 1269 | df_params = pd.DataFrame([params]) |
| 1270 | df_params.to_excel( |
| 1271 | writer, sheet_name="Summary", index=False, startrow=startrow |
| 1272 | ) |
| 1273 | startrow += len(df_params) + 2 |
| 1274 | df_stats.to_excel( |
| 1275 | writer, sheet_name="Summary", index=True, startrow=startrow |
| 1276 | ) |
| 1277 | startrow += len(df_stats) + 3 |
| 1278 | |
| 1279 | raw_data.to_excel(writer, sheet_name="Raw data", index=False, startrow=0) |
| 1280 | |
| 1281 | logger.info( |
| 1282 | f"{Color.GREEN}Client metrics exported to file: {filename}{Color.RESET}" |
| 1283 | ) |
| 1284 | |
| 1285 | |
| 1286 | async def get_server_info(url: str, headers: dict[str, str] | None = None) -> None: |
| 1287 | logger.info(f"{Color.BLUE}Collecting information from server: {url}{Color.RESET}") |
| 1288 | async with aiohttp.ClientSession() as session: |
| 1289 | # Get server version (not mandatory, "version" endpoint may not exist) |
| 1290 | url_version = f"{url}/version" |
| 1291 | async with session.get(url_version, headers=headers) as response: |
| 1292 | if HTTPStatus(response.status) == HTTPStatus.OK: |
| 1293 | text = await response.text() |
| 1294 | logger.info(f"{Color.BLUE}Server version: {text}{Color.RESET}") |
| 1295 | |
| 1296 | # Get available models |
| 1297 | url_models = f"{url}/v1/models" |
| 1298 | async with session.get(url_models, headers=headers) as response: |
| 1299 | if HTTPStatus(response.status) == HTTPStatus.OK: |
| 1300 | text = await response.text() |
| 1301 | logger.info(f"{Color.BLUE}Models:{Color.RESET}") |
| 1302 | models_data = json.loads(text) |
| 1303 | models_list = models_data["data"] |
| 1304 | for model in models_list: |
| 1305 | model_id = model["id"] |
| 1306 | max_model_len = model.get("max_model_len", "N/A") |
| 1307 | logger.info( |
| 1308 | f"{Color.BLUE}\t{model_id=}, {max_model_len=}{Color.RESET}" |
| 1309 | ) |
| 1310 | else: |
| 1311 | logger.info(f"{Color.RED}Failed to get models{Color.RESET}") |
| 1312 | |
| 1313 | |
| 1314 | async def main() -> None: |
| 1315 | parser = argparse.ArgumentParser( |
| 1316 | prog="Benchmark serving with multi-turn conversations", |
| 1317 | description="Benchmark online inference using REST API", |
| 1318 | ) |
| 1319 | parser.add_argument("--version", action="version", version="%(prog)s 1.0") |
| 1320 | |
| 1321 | parser.add_argument( |
| 1322 | "-i", |
| 1323 | "--input-file", |
| 1324 | type=str, |
| 1325 | required=True, |
| 1326 | help="Input JSON file with ShareGPT conversations or " |
| 1327 | "configuration file for generation of synthetic conversations", |
| 1328 | ) |
| 1329 | parser.add_argument( |
| 1330 | "-o", |
| 1331 | "--output-file", |
| 1332 | type=str, |
| 1333 | default=None, |
| 1334 | help="Output JSON file containing conversations with updated assistant answers", |
| 1335 | ) |
| 1336 | |
| 1337 | parser.add_argument( |
| 1338 | "--seed", |
| 1339 | type=int, |
| 1340 | default=0, |
| 1341 | help="Seed for random number generators (default: 0)", |
| 1342 | ) |
| 1343 | |
| 1344 | parser.add_argument( |
| 1345 | "-m", "--model", type=str, required=True, help="Path of the LLM model" |
| 1346 | ) |
| 1347 | parser.add_argument( |
| 1348 | "--served-model-name", |
| 1349 | type=str, |
| 1350 | default=None, |
| 1351 | help="The model name used in the API. " |
| 1352 | "If not specified, the model name will be the " |
| 1353 | "same as the `--model` argument. ", |
| 1354 | ) |
| 1355 | |
| 1356 | parser.add_argument( |
| 1357 | "-u", |
| 1358 | "--url", |
| 1359 | type=str, |
| 1360 | default="http://localhost:8000", |
| 1361 | help="Base URL for the LLM API server", |
| 1362 | ) |
| 1363 | |
| 1364 | parser.add_argument( |
| 1365 | "--api-key", |
| 1366 | type=str, |
| 1367 | default=None, |
| 1368 | help="API key to send as an Authorization bearer token", |
| 1369 | ) |
| 1370 | parser.add_argument( |
| 1371 | "--header", |
| 1372 | action="append", |
| 1373 | type=parse_custom_header, |
| 1374 | default=None, |
| 1375 | metavar="KEY=VALUE", |
| 1376 | help="Custom request header. Can be specified multiple times. " |
| 1377 | "Accepts 'Header-Name: value' or 'Header-Name=value'.", |
| 1378 | ) |
| 1379 | |
| 1380 | parser.add_argument( |
| 1381 | "-p", |
| 1382 | "--num-clients", |
| 1383 | type=int, |
| 1384 | default=1, |
| 1385 | help="Number of clients that will send requests in parallel", |
| 1386 | ) |
| 1387 | parser.add_argument( |
| 1388 | "-k", |
| 1389 | "--max-active-conversations", |
| 1390 | type=int, |
| 1391 | default=None, |
| 1392 | help="Max number of active conversations at a time (for all clients)", |
| 1393 | ) |
| 1394 | parser.add_argument( |
| 1395 | "-n", |
| 1396 | "--max-num-requests", |
| 1397 | type=int, |
| 1398 | default=None, |
| 1399 | help="Max number of requests to send (total for all clients)", |
| 1400 | ) |
| 1401 | |
| 1402 | parser.add_argument( |
| 1403 | "--warmup-step", |
| 1404 | default=False, |
| 1405 | action="store_true", |
| 1406 | help="Run a warmup step (using only the first turn of every conversation), " |
| 1407 | "measurements will not be included in the final benchmark results", |
| 1408 | ) |
| 1409 | |
| 1410 | parser.add_argument( |
| 1411 | "--max-turns", |
| 1412 | type=int, |
| 1413 | default=None, |
| 1414 | help="Maximum number of turns/messages per conversation, " |
| 1415 | "includes both user and assistant messages " |
| 1416 | "(a positive number, e.g: 2, 4, 6, etc.), disabled by default", |
| 1417 | ) |
| 1418 | parser.add_argument( |
| 1419 | "--no-early-stop", |
| 1420 | default=False, |
| 1421 | action="store_true", |
| 1422 | help="By default, the benchmark will stop if at least one client exits." |
| 1423 | " Use this flag to disable this behavior", |
| 1424 | ) |
| 1425 | |
| 1426 | parser.add_argument( |
| 1427 | "--limit-max-tokens", |
| 1428 | type=int, |
| 1429 | default=NUM_TOKENS_FROM_DATASET, |
| 1430 | help="Set max_tokens for the output token count of each request " |
| 1431 | "(must also set --limit-min-tokens). " |
| 1432 | "Overrides output token count from the input dataset. " |
| 1433 | "Use a negative value to disable this limit.", |
| 1434 | ) |
| 1435 | parser.add_argument( |
| 1436 | "--limit-min-tokens", |
| 1437 | type=int, |
| 1438 | default=NUM_TOKENS_FROM_DATASET, |
| 1439 | help="Set min_tokens for the output token count of each request " |
| 1440 | "(must also set --limit-max-tokens). " |
| 1441 | "Overrides output token count from the input dataset. " |
| 1442 | "Use a negative value to disable this limit.", |
| 1443 | ) |
| 1444 | |
| 1445 | parser.add_argument( |
| 1446 | "--request-rate", |
| 1447 | type=float, |
| 1448 | default=0, |
| 1449 | help="Expected request rate (Poisson process) per client in requests/sec." |
| 1450 | "Set to 0 for no delay between requests.", |
| 1451 | ) |
| 1452 | parser.add_argument( |
| 1453 | "--max-retries", |
| 1454 | type=int, |
| 1455 | default=int(os.environ.get("MULTITURN_BENCH_MAX_RETRIES", "0")), |
| 1456 | help="Maximum number of retry attempts for timed-out requests. " |
| 1457 | "Default is 0 (no retries). " |
| 1458 | "Set to higher values to retry failed requests and maintain " |
| 1459 | "fair workload distribution. " |
| 1460 | "Can also be set via MULTITURN_BENCH_MAX_RETRIES environment variable.", |
| 1461 | ) |
| 1462 | parser.add_argument( |
| 1463 | "--conversation-sampling", |
| 1464 | type=ConversationSampling, |
| 1465 | choices=list(ConversationSampling), |
| 1466 | default=ConversationSampling.ROUND_ROBIN, |
| 1467 | help=( |
| 1468 | "Strategy for selecting which conversation to use for the next request. " |
| 1469 | "Options: 'round_robin' (cycle through conversations), " |
| 1470 | "'random' (pick randomly)." |
| 1471 | ), |
| 1472 | ) |
| 1473 | parser.add_argument( |
| 1474 | "--verify-output", |
| 1475 | default=False, |
| 1476 | action="store_true", |
| 1477 | help="Verify the LLM output (compare to the answers in the input JSON file)", |
| 1478 | ) |
| 1479 | parser.add_argument( |
| 1480 | "--request-timeout-sec", |
| 1481 | type=int, |
| 1482 | default=120, |
| 1483 | help="Timeout in seconds for each API request (default: 120). " |
| 1484 | "Automatically increased if max tokens imply longer decoding.", |
| 1485 | ) |
| 1486 | |
| 1487 | parser.add_argument( |
| 1488 | "--no-stream", |
| 1489 | default=False, |
| 1490 | action="store_true", |
| 1491 | help="Disable stream/streaming mode (set 'stream' to False in the API request)", |
| 1492 | ) |
| 1493 | |
| 1494 | parser.add_argument( |
| 1495 | "--send-conversation-id", |
| 1496 | default=False, |
| 1497 | action="store_true", |
| 1498 | help=( |
| 1499 | "Inject a `conversation_id` field into each Chat Completions " |
| 1500 | "payload. This is a non-standard OpenAI extension consumed by " |
| 1501 | "vLLM's disaggregated multi-turn proxy " |
| 1502 | "(examples/disaggregated/disaggregated_serving/" |
| 1503 | "disagg_proxy_multiturn.py) to key cross-turn KV cache reuse. " |
| 1504 | "Leave disabled (default) when targeting strict " |
| 1505 | "OpenAI-compatible endpoints; enable when benchmarking the " |
| 1506 | "disaggregated proxy." |
| 1507 | ), |
| 1508 | ) |
| 1509 | |
| 1510 | parser.add_argument( |
| 1511 | "-e", |
| 1512 | "--excel-output", |
| 1513 | default=False, |
| 1514 | action="store_true", |
| 1515 | help="Export summary to Excel file (optional)", |
| 1516 | ) |
| 1517 | parser.add_argument( |
| 1518 | "--stats-json-output", |
| 1519 | type=str, |
| 1520 | default=None, |
| 1521 | help="Export per-request stats (ttft_ms, tpot_ms, etc.) to a JSON file", |
| 1522 | ) |
| 1523 | parser.add_argument( |
| 1524 | "-v", |
| 1525 | "--verbose", |
| 1526 | default=False, |
| 1527 | action="store_true", |
| 1528 | help="Enable verbose output", |
| 1529 | ) |
| 1530 | parser.add_argument( |
| 1531 | "--print-content", |
| 1532 | default=False, |
| 1533 | action="store_true", |
| 1534 | help="Print the user prompts and the server's answers", |
| 1535 | ) |
| 1536 | |
| 1537 | parser.add_argument( |
| 1538 | "--warmup-percentages", |
| 1539 | type=str, |
| 1540 | default="0%", |
| 1541 | help="Ignore the first X samples as warmup (X is a percentage)." |
| 1542 | " A comma separated list of percentages can be used " |
| 1543 | "(for example: --warmup-percentages=0%%,50%%)", |
| 1544 | ) |
| 1545 | |
| 1546 | parser.add_argument( |
| 1547 | "--trust-remote-code", |
| 1548 | action="store_true", |
| 1549 | help="Trust remote code when loading the tokenizer.", |
| 1550 | ) |
| 1551 | |
| 1552 | args = parser.parse_args() |
| 1553 | |
| 1554 | logger.info(redact_sensitive_namespace(args, _SENSITIVE_ARG_FIELDS)) |
| 1555 | |
| 1556 | logger.info(f"{Color.GREEN}Input parameters:{Color.RESET}") |
| 1557 | logger.info(f"url={args.url}") |
| 1558 | logger.info(f"model={args.model}") |
| 1559 | logger.info(f"num_clients={args.num_clients}") |
| 1560 | |
| 1561 | if args.verify_output: |
| 1562 | logger.info(f"{Color.PURPLE}Verify is enabled{Color.RESET}") |
| 1563 | |
| 1564 | # Calculate the amount of samples to filter (as warmup samples/measurements). |
| 1565 | try: |
| 1566 | warmup_percentages: list[float] = [0.0] |
| 1567 | if not args.warmup_step: |
| 1568 | # Warmup percentage can be used only if the warmup step was used |
| 1569 | warmup_strings: list[str] = args.warmup_percentages.split(",") |
| 1570 | warmup_strings = [x.replace("%", "") for x in warmup_strings] |
| 1571 | warmup_percentages = [float(x) / 100 for x in warmup_strings] |
| 1572 | |
| 1573 | # Check for valid range (0 to 1) |
| 1574 | for p in warmup_percentages: |
| 1575 | assert p >= 0.0 and p < 1.0 |
| 1576 | |
| 1577 | # Sort from high to low warmup percentage |
| 1578 | warmup_percentages.sort() |
| 1579 | |
| 1580 | logger.info( |
| 1581 | f"Warmup percentages (percentage of samples): {warmup_percentages}" |
| 1582 | ) |
| 1583 | |
| 1584 | except Exception: |
| 1585 | raise ValueError( |
| 1586 | f"Invalid --warmup-percentage={args.warmup_percentage}" |
| 1587 | ) from None |
| 1588 | |
| 1589 | # Set global seeds for main process |
| 1590 | random.seed(args.seed) |
| 1591 | np.random.seed(args.seed) |
| 1592 | |
| 1593 | logger.info("Loading tokenizer") |
| 1594 | tokenizer = AutoTokenizer.from_pretrained( |
| 1595 | args.model, trust_remote_code=args.trust_remote_code |
| 1596 | ) |
| 1597 | |
| 1598 | headers = build_request_headers(args.api_key, args.header) |
| 1599 | await get_server_info(args.url, headers=headers) |
| 1600 | |
| 1601 | # Load the input file (either conversations of configuration file) |
| 1602 | logger.info(f"Reading input file: {args.input_file}") |
| 1603 | with open(args.input_file) as f: |
| 1604 | input_data = json.load(f) |
| 1605 | |
| 1606 | gen_conv_args = None |
| 1607 | if isinstance(input_data, list): |
| 1608 | # The conversations are stored as a list of dicts |
| 1609 | logger.info(f"Found {len(input_data)} items in the input file") |
| 1610 | |
| 1611 | # Convert the list to a ConversationsMap |
| 1612 | conversations = conversations_list_to_dict(input_data) |
| 1613 | |
| 1614 | elif isinstance(input_data, dict): |
| 1615 | # The input file is a configuration file |
| 1616 | # (type is determined by the field 'filetype') |
| 1617 | if "filetype" not in input_data: |
| 1618 | raise Exception( |
| 1619 | f"Input file {args.input_file} is invalid (missing 'filetype')" |
| 1620 | ) |
| 1621 | |
| 1622 | logger.info(f"Using input file with filetype: {input_data['filetype']}") |
| 1623 | |
| 1624 | gen_conv_args = parse_input_json_file(input_data) |
| 1625 | |
| 1626 | # Disable warning from "huggingface/tokenizers" |
| 1627 | # (when using python multiprocessing and tokenizers) |
| 1628 | os.environ["TOKENIZERS_PARALLELISM"] = "true" |
| 1629 | |
| 1630 | # Generate synthetic conversations |
| 1631 | conversations = generate_conversations(gen_conv_args, tokenizer) |
| 1632 | |
| 1633 | else: |
| 1634 | raise Exception(f"Input file {args.input_file} is invalid") |
| 1635 | |
| 1636 | if args.max_turns is not None: |
| 1637 | if args.max_turns < 1: |
| 1638 | raise ValueError("Max turns must be a positive number") |
| 1639 | logger.info( |
| 1640 | f"{Color.PURPLE}Max turns per conversation " |
| 1641 | f"is limited to {args.max_turns}{Color.RESET}" |
| 1642 | ) |
| 1643 | |
| 1644 | # Create benchmark configurations |
| 1645 | client_args, req_args = get_client_config(args, conversations) |
| 1646 | |
| 1647 | bench_args = BenchmarkArgs( |
| 1648 | url=args.url, num_clients=args.num_clients, early_stop=not args.no_early_stop |
| 1649 | ) |
| 1650 | |
| 1651 | warmup_runtime_sec: float | None = None |
| 1652 | |
| 1653 | # Warm-up step |
| 1654 | if args.warmup_step: |
| 1655 | # Only send a single user prompt from every conversation. |
| 1656 | # max_active_conversations must be 1, |
| 1657 | # otherwise the clients may exit after sending a single request |
| 1658 | # (because the task queue is empty). |
| 1659 | warmup_client_args = client_args._replace( |
| 1660 | skip_first_turn=False, max_turns=1, max_active_conversations=1 |
| 1661 | ) |
| 1662 | |
| 1663 | # Early stop should be disabled, |
| 1664 | # all clients should finish their work before exiting |
| 1665 | warmup_bench_args = bench_args._replace(early_stop=False) |
| 1666 | |
| 1667 | logger.info("%sWarmup start%s", Color.PURPLE, Color.RESET) |
| 1668 | warmup_start_ns = time.perf_counter_ns() |
| 1669 | conversations, _ = await main_mp( |
| 1670 | warmup_client_args, req_args, warmup_bench_args, tokenizer, conversations |
| 1671 | ) |
| 1672 | warmup_runtime_sec = nanosec_to_sec(time.perf_counter_ns() - warmup_start_ns) |
| 1673 | logger.info( |
| 1674 | "%sWarmup runtime: %.3f sec (%.3f ms)%s", |
| 1675 | Color.PURPLE, |
| 1676 | warmup_runtime_sec, |
| 1677 | warmup_runtime_sec * 1000, |
| 1678 | Color.RESET, |
| 1679 | ) |
| 1680 | logger.info("%sWarmup done%s", Color.PURPLE, Color.RESET) |
| 1681 | |
| 1682 | # Run the benchmark |
| 1683 | benchmark_start_ns = time.perf_counter_ns() |
| 1684 | client_convs, client_metrics = await main_mp( |
| 1685 | client_args, req_args, bench_args, tokenizer, conversations |
| 1686 | ) |
| 1687 | benchmark_runtime_sec = nanosec_to_sec(time.perf_counter_ns() - benchmark_start_ns) |
| 1688 | |
| 1689 | # Calculate requests per second |
| 1690 | requests_per_sec = len(client_metrics) / benchmark_runtime_sec |
| 1691 | benchmark_runtime_ms = benchmark_runtime_sec * 1000.0 |
| 1692 | logger.info( |
| 1693 | "%sAll clients finished, benchmark runtime: %.3f sec (%.3f ms), " |
| 1694 | "requests per second: %.3f%s", |
| 1695 | Color.GREEN, |
| 1696 | benchmark_runtime_sec, |
| 1697 | benchmark_runtime_ms, |
| 1698 | requests_per_sec, |
| 1699 | Color.RESET, |
| 1700 | ) |
| 1701 | if warmup_runtime_sec is not None: |
| 1702 | total_runtime_sec = benchmark_runtime_sec + warmup_runtime_sec |
| 1703 | logger.info( |
| 1704 | "%sWarmup runtime: %.3f sec (%.3f ms)%s", |
| 1705 | Color.GREEN, |
| 1706 | warmup_runtime_sec, |
| 1707 | warmup_runtime_sec * 1000, |
| 1708 | Color.RESET, |
| 1709 | ) |
| 1710 | logger.info( |
| 1711 | "%sTotal runtime (including warmup): %.3f sec (%.3f ms)%s", |
| 1712 | Color.GREEN, |
| 1713 | total_runtime_sec, |
| 1714 | total_runtime_sec * 1000, |
| 1715 | Color.RESET, |
| 1716 | ) |
| 1717 | |
| 1718 | # Benchmark parameters |
| 1719 | params = { |
| 1720 | "model": args.model, |
| 1721 | "num_clients": args.num_clients, |
| 1722 | "num_conversations": len(conversations), |
| 1723 | "active_conversations": args.max_active_conversations, |
| 1724 | "seed": args.seed, |
| 1725 | } |
| 1726 | |
| 1727 | if args.limit_min_tokens > 0: |
| 1728 | params["min_tokens"] = args.limit_min_tokens |
| 1729 | |
| 1730 | if args.limit_max_tokens > 0: |
| 1731 | params["max_tokens"] = args.limit_max_tokens |
| 1732 | |
| 1733 | # Process and print statistics (and save excel file with the statistics) |
| 1734 | process_statistics( |
| 1735 | client_metrics, |
| 1736 | test_params=params, |
| 1737 | warmup_percentages=warmup_percentages, |
| 1738 | verbose=args.verbose, |
| 1739 | gen_conv_args=gen_conv_args, |
| 1740 | excel_output=args.excel_output, |
| 1741 | warmup_runtime_sec=warmup_runtime_sec, |
| 1742 | ) |
| 1743 | |
| 1744 | if args.stats_json_output is not None: |
| 1745 | # Export per-request metrics as a JSON array for downstream analysis. |
| 1746 | stats_data = [s._asdict() for s in client_metrics] |
| 1747 | logger.info( |
| 1748 | f"{Color.GREEN}Writing per-request stats JSON: " |
| 1749 | f"{args.stats_json_output}{Color.RESET}" |
| 1750 | ) |
| 1751 | os.makedirs( |
| 1752 | os.path.dirname(os.path.abspath(args.stats_json_output)), exist_ok=True |
| 1753 | ) |
| 1754 | with open(args.stats_json_output, "w") as f: |
| 1755 | json.dump(stats_data, f, indent=2) |
| 1756 | |
| 1757 | if args.output_file is not None: |
| 1758 | # Write a JSON file with the updated conversations |
| 1759 | # The "assistant" content will contain the answers from the tested LLM |
| 1760 | output_data: ShareGptConversations = conversations_dict_to_list(client_convs) |
| 1761 | logger.info( |
| 1762 | f"{Color.GREEN}Writing conversations file: {args.output_file}{Color.RESET}" |
| 1763 | ) |
| 1764 | with open(args.output_file, "w") as f: |
| 1765 | json.dump(output_data, f, indent=4) |
| 1766 | |
| 1767 | |
| 1768 | if __name__ == "__main__": |
| 1769 | asyncio.run(main()) |