@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f0964e29cb
Fork 0

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] Add benchmark tool for multi turn conversations (#20267)

Daniel Serebrenik committed Aug 8, 2025 at 20:28 UTC f0964e29cb3b2deccdad89f5f8c068d3a629d239
7 files changed +2540
benchmarks/multi_turn/README.md new
+71
@@ -0,0 +1,71 @@
1 +# Benchmark KV Cache Offloading with Multi-Turn Conversations
2 +
3 +The requirements (pip) for `benchmark_serving_multi_turn.py` can be found in `requirements.txt`
4 +
5 +First start serving your model
6 +
7 +```bash
8 +export MODEL_NAME=/models/meta-llama/Meta-Llama-3.1-8B-Instruct/
9 +
10 +vllm serve $MODEL_NAME --disable-log-requests
11 +```
12 +
13 +## Synthetic Multi-Turn Conversations
14 +
15 +Download the following text file (used for generation of synthetic conversations)
16 +
17 +```bash
18 +wget https://www.gutenberg.org/ebooks/1184.txt.utf-8
19 +mv 1184.txt.utf-8 pg1184.txt
20 +```
21 +
22 +The filename `pg1184.txt` is used in `generate_multi_turn.json` (see `"text_files"`).
23 +
24 +But you may use other text files if you prefer (using this specific file is not required).
25 +
26 +Then run the benchmarking script
27 +
28 +```bash
29 +export MODEL_NAME=/models/meta-llama/Meta-Llama-3.1-8B-Instruct/
30 +
31 +python benchmark_serving_multi_turn.py --model $MODEL_NAME --input-file generate_multi_turn.json \
32 +--num-clients 2 --max-active-conversations 6
33 +```
34 +
35 +You can edit the file `generate_multi_turn.json` to change the conversation parameters (number of turns, etc.).
36 +
37 +If successful, you will see the following output
38 +
39 +```bash
40 +----------------------------------------------------------------------------------------------------
41 +Statistics summary:
42 +runtime_sec = 215.810
43 +requests_per_sec = 0.769
44 +----------------------------------------------------------------------------------------------------
45 + count mean std min 25% 50% 75% 90% 99% max
46 +ttft_ms 166.0 78.22 67.63 45.91 59.94 62.26 64.43 69.66 353.18 567.54
47 +tpot_ms 166.0 25.37 0.57 24.40 25.07 25.31 25.50 25.84 27.50 28.05
48 +latency_ms 166.0 2591.07 326.90 1998.53 2341.62 2573.01 2860.10 3003.50 3268.46 3862.94
49 +input_num_turns 166.0 7.43 4.57 1.00 3.00 7.00 11.00 13.00 17.00 17.00
50 +input_num_tokens 166.0 2006.20 893.56 522.00 1247.75 2019.00 2718.00 3233.00 3736.45 3899.00
51 +output_num_tokens 166.0 100.01 11.80 80.00 91.00 99.00 109.75 116.00 120.00 120.00
52 +output_num_chunks 166.0 99.01 11.80 79.00 90.00 98.00 108.75 115.00 119.00 119.00
53 +----------------------------------------------------------------------------------------------------
54 +```
55 +
56 +## ShareGPT Conversations
57 +
58 +To run with the ShareGPT data, download the following ShareGPT dataset:
59 +`https://huggingface.co/datasets/philschmid/sharegpt-raw/blob/main/sharegpt_20230401_clean_lang_split.json`
60 +
61 +Use the `convert_sharegpt_to_openai.py` script to convert the dataset to a format supported by `benchmark_serving_multi_turn.py`
62 +
63 +```bash
64 +python convert_sharegpt_to_openai.py sharegpt_20230401_clean_lang_split.json sharegpt_conv_128.json --seed=99 --max-items=128
65 +```
66 +
67 +The script will convert the ShareGPT dataset to a dataset with the standard user/assistant roles.
68 +
69 +The flag `--max-items=128` is used to sample 128 conversations from the original dataset (change as needed).
70 +
71 +Use the output JSON file `sharegpt_conv_128.json` as the `--input-file` for `benchmark_serving_multi_turn.py`.
benchmarks/multi_turn/bench_dataset.py new
+493
@@ -0,0 +1,493 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from abc import ABC, abstractmethod
4 +from statistics import mean
5 +from typing import Any, NamedTuple, Optional, Union
6 +
7 +import numpy as np # type: ignore
8 +import pandas as pd # type: ignore
9 +from bench_utils import (
10 + TEXT_SEPARATOR,
11 + Color,
12 + logger,
13 +)
14 +from transformers import AutoTokenizer # type: ignore
15 +
16 +# Conversation ID is a string (e.g: "UzTK34D")
17 +ConvId = str
18 +
19 +# A list of dicts (dicts with keys "id" and "messages")
20 +ShareGptConversations = list[dict[str, Any]]
21 +
22 +# A list of dicts (dicts with keys "role" and "content")
23 +MessagesList = list[dict[str, str]]
24 +
25 +# Map conversation ID to conversation messages
26 +ConversationsMap = list[ConvId, MessagesList]
27 +
28 +
29 +class Distribution(ABC):
30 + @abstractmethod
31 + def sample(self, size: int = 1) -> np.ndarray:
32 + pass
33 +
34 +
35 +class UniformDistribution(Distribution):
36 + def __init__(
37 + self,
38 + min_val: Union[int, float],
39 + max_val: Union[int, float],
40 + is_integer: bool = True,
41 + ) -> None:
42 + self.min_val = min_val
43 + self.max_val = max_val
44 + self.is_integer = is_integer
45 +
46 + def sample(self, size: int = 1) -> np.ndarray:
47 + if self.is_integer:
48 + return np.random.randint(
49 + int(self.min_val), int(self.max_val + 1), size=size
50 + )
51 + else:
52 + return np.random.uniform(self.min_val, self.max_val, size=size)
53 +
54 + def __repr__(self) -> str:
55 + return f"UniformDistribution[{self.min_val}, {self.max_val}]"
56 +
57 +
58 +class ConstantDistribution(Distribution):
59 + def __init__(self, value: Union[int, float]) -> None:
60 + self.value = value
61 + self.max_val = value
62 +
63 + def sample(self, size: int = 1) -> np.ndarray:
64 + return np.full(shape=size, fill_value=self.value)
65 +
66 + def __repr__(self) -> str:
67 + return f"Constant[{self.value}]"
68 +
69 +
70 +class ZipfDistribution(Distribution):
71 + def __init__(self, alpha: float, max_val: Optional[int] = None) -> None:
72 + self.alpha = alpha
73 + self.max_val = max_val
74 +
75 + def sample(self, size: int = 1) -> np.ndarray:
76 + samples = np.random.zipf(self.alpha, size=size)
77 + if self.max_val:
78 + samples = np.minimum(samples, self.max_val)
79 + return samples
80 +
81 + def __repr__(self) -> str:
82 + return f"ZipfDistribution[{self.alpha}]"
83 +
84 +
85 +class PoissonDistribution(Distribution):
86 + def __init__(self, alpha: float, max_val: Optional[int] = None) -> None:
87 + self.alpha = alpha
88 + self.max_val = max_val
89 +
90 + def sample(self, size: int = 1) -> np.ndarray:
91 + samples = np.random.poisson(self.alpha, size=size)
92 + if self.max_val:
93 + samples = np.minimum(samples, self.max_val)
94 + return samples
95 +
96 + def __repr__(self) -> str:
97 + return f"PoissonDistribution[{self.alpha}]"
98 +
99 +
100 +class LognormalDistribution(Distribution):
101 + def __init__(
102 + self, mean: float, sigma: float, max_val: Optional[int] = None
103 + ) -> None:
104 + self.mean = mean
105 + self.sigma = sigma
106 + self.max_val = max_val
107 +
108 + def sample(self, size: int = 1) -> np.ndarray:
109 + samples = np.random.lognormal(mean=self.mean, sigma=self.sigma, size=size)
110 + if self.max_val:
111 + samples = np.minimum(samples, self.max_val)
112 +
113 + return np.round(samples).astype(int)
114 +
115 + def __repr__(self) -> str:
116 + return f"LognormalDistribution[{self.mean}, {self.sigma}]"
117 +
118 +
119 +class GenConvArgs(NamedTuple):
120 + num_conversations: int
121 + text_files: list[str]
122 + input_num_turns: Distribution
123 + input_common_prefix_num_tokens: Distribution
124 + input_prefix_num_tokens: Distribution
125 + input_num_tokens: Distribution
126 + output_num_tokens: Distribution
127 + print_stats: bool
128 +
129 +
130 +def verify_field_exists(
131 + conf: dict, field_name: str, section: str, subsection: str
132 +) -> None:
133 + if field_name not in conf:
134 + raise ValueError(
135 + f"Missing field '{field_name}' in {section=} and {subsection=}"
136 + )
137 +
138 +
139 +def get_random_distribution(
140 + conf: dict, section: str, subsection: str, optional: bool = False
141 +) -> Distribution:
142 + # section can be "prompt_input" or "prompt_output" (both required)
143 + conf = conf[section]
144 +
145 + if optional and subsection not in conf:
146 + # Optional subsection, if not found assume the value is always 0
147 + return ConstantDistribution(0)
148 +
149 + # subsection can be "num_turns", "num_tokens" or "prefix_num_tokens"
150 + if subsection not in conf:
151 + raise ValueError(f"Missing subsection {subsection} in section {section}")
152 +
153 + conf = conf[subsection]
154 +
155 + distribution = conf.get("distribution")
156 + if distribution is None:
157 + raise ValueError(
158 + f"Missing field 'distribution' in {section=} and {subsection=}"
159 + )
160 +
161 + if distribution == "constant":
162 + verify_field_exists(conf, "value", section, subsection)
163 + return ConstantDistribution(conf["value"])
164 +
165 + elif distribution == "zipf":
166 + verify_field_exists(conf, "alpha", section, subsection)
167 + max_val = conf.get("max", None)
168 + return ZipfDistribution(conf["alpha"], max_val=max_val)
169 +
170 + elif distribution == "poisson":
171 + verify_field_exists(conf, "alpha", section, subsection)
172 + max_val = conf.get("max", None)
173 + return PoissonDistribution(conf["alpha"], max_val=max_val)
174 +
175 + elif distribution == "lognormal":
176 + verify_field_exists(conf, "mean", section, subsection)
177 + verify_field_exists(conf, "sigma", section, subsection)
178 + max_val = conf.get("max", None)
179 + return LognormalDistribution(conf["mean"], conf["sigma"], max_val=max_val)
180 +
181 + elif distribution == "uniform":
182 + verify_field_exists(conf, "min", section, subsection)
183 + verify_field_exists(conf, "max", section, subsection)
184 +
185 + min_value = conf["min"]
186 + max_value = conf["max"]
187 +
188 + assert min_value > 0
189 + assert min_value <= max_value
190 +
191 + is_integer = isinstance(min_value, int) and isinstance(max_value, int)
192 + return UniformDistribution(min_value, max_value, is_integer)
193 + else:
194 + raise ValueError(f"Unknown distribution: {distribution}")
195 +
196 +
197 +def parse_input_json_file(conf: dict) -> GenConvArgs:
198 + # Validate the input file
199 + assert isinstance(conf, dict)
200 + required_fields = [
201 + "filetype",
202 + "num_conversations",
203 + "text_files",
204 + "prompt_input",
205 + "prompt_output",
206 + ]
207 + for field in required_fields:
208 + assert field in conf, f"Missing field {field} in input {conf}"
209 +
210 + assert conf["filetype"] == "generate_conversations"
211 +
212 + assert conf["num_conversations"] > 0, "num_conversations should be larger than zero"
213 +
214 + text_files = conf["text_files"]
215 +
216 + assert isinstance(text_files, list), "Field 'text_files' should be a list"
217 + assert len(text_files) > 0, (
218 + "Field 'text_files' should be a list with at least one file"
219 + )
220 +
221 + # Parse the parameters for the prompt input/output workload
222 + input_num_turns = get_random_distribution(conf, "prompt_input", "num_turns")
223 + input_num_tokens = get_random_distribution(conf, "prompt_input", "num_tokens")
224 + input_common_prefix_num_tokens = get_random_distribution(
225 + conf, "prompt_input", "common_prefix_num_tokens", optional=True
226 + )
227 + input_prefix_num_tokens = get_random_distribution(
228 + conf, "prompt_input", "prefix_num_tokens"
229 + )
230 + output_num_tokens = get_random_distribution(conf, "prompt_output", "num_tokens")
231 +
232 + print_stats: bool = conf.get("print_stats", False)
233 + assert isinstance(print_stats, bool), (
234 + "Field 'print_stats' should be either 'true' or 'false'"
235 + )
236 +
237 + args = GenConvArgs(
238 + num_conversations=conf["num_conversations"],
239 + text_files=text_files,
240 + input_num_turns=input_num_turns,
241 + input_common_prefix_num_tokens=input_common_prefix_num_tokens,
242 + input_prefix_num_tokens=input_prefix_num_tokens,
243 + input_num_tokens=input_num_tokens,
244 + output_num_tokens=output_num_tokens,
245 + print_stats=print_stats,
246 + )
247 + return args
248 +
249 +
250 +def print_conv_stats(conversations: ConversationsMap, tokenizer: AutoTokenizer) -> None:
251 + # Collect statistics
252 + conv_stats: list[dict[Any, Any]] = []
253 + req_stats: list[int] = []
254 +
255 + print("\nCollecting statistics...")
256 + for messages in conversations.values():
257 + # messages is a list of dicts
258 + user_tokens: list[int] = []
259 + assistant_tokens: list[int] = []
260 + request_tokens: list[int] = []
261 +
262 + req_tokens = 0
263 + for m in messages:
264 + content = m["content"]
265 + num_tokens = len(tokenizer(content).input_ids)
266 +
267 + if m["role"] == "user":
268 + user_tokens.append(num_tokens)
269 + # New user prompt including all chat history
270 + req_tokens += num_tokens
271 + request_tokens.append(req_tokens)
272 +
273 + elif m["role"] == "assistant":
274 + assistant_tokens.append(num_tokens)
275 + # Update assistant answer
276 + # (will be part of chat history for the next user prompt)
277 + req_tokens += num_tokens
278 +
279 + item_stats = {
280 + "conversation_turns": len(messages),
281 + "user_tokens": mean(user_tokens),
282 + "assistant_tokens": mean(assistant_tokens),
283 + }
284 +
285 + conv_stats.append(item_stats)
286 + req_stats.extend(request_tokens)
287 +
288 + # Print statistics
289 + percentiles = [0.25, 0.5, 0.75, 0.9, 0.99]
290 +
291 + print(TEXT_SEPARATOR)
292 + print(f"{Color.YELLOW}Conversations statistics:{Color.RESET}")
293 + print(TEXT_SEPARATOR)
294 + df = pd.DataFrame(conv_stats)
295 + print(df.describe(percentiles=percentiles).transpose())
296 + print(TEXT_SEPARATOR)
297 + print(f"{Color.YELLOW}Request statistics:{Color.RESET}")
298 + print(TEXT_SEPARATOR)
299 + df = pd.DataFrame(req_stats, columns=["request_tokens"])
300 + print(df.describe(percentiles=percentiles).transpose())
301 + print(TEXT_SEPARATOR)
302 +
303 +
304 +def generate_conversations(
305 + args: GenConvArgs, tokenizer: AutoTokenizer
306 +) -> ConversationsMap:
307 + # Text for all user prompts
308 + # (text from the input text files will be appended to this line)
309 + base_prompt_text = "Please rewrite the following text and add more content: "
310 + base_prompt_token_count = len(
311 + tokenizer.encode(base_prompt_text, add_special_tokens=False)
312 + )
313 +
314 + logger.info(f"{Color.PURPLE}Generating conversations...{Color.RESET}")
315 + logger.info(args)
316 +
317 + list_of_tokens = []
318 +
319 + for filename in args.text_files:
320 + # Load text file that will be used to generate prompts
321 + with open(filename) as file:
322 + data = file.read()
323 + tokens_in_file = tokenizer.encode(data, add_special_tokens=False)
324 + list_of_tokens.extend(tokens_in_file)
325 +
326 + conversations: ConversationsMap = {}
327 + conv_id = 0
328 +
329 + # Generate number of turns for every conversation
330 + turn_count: np.ndarray = args.input_num_turns.sample(args.num_conversations)
331 +
332 + # Turn count should be at least 2 (one user prompt and one assistant answer)
333 + turn_count = np.maximum(turn_count, 2)
334 +
335 + # Round up to an even number (every user prompt should have an answer)
336 + turn_count = turn_count + (turn_count % 2)
337 +
338 + # Generate number of prefix tokens for every conversation
339 + conv_prefix_tokens: np.ndarray = args.input_prefix_num_tokens.sample(
340 + args.num_conversations
341 + )
342 +
343 + # Used to reduce shared text between conversations
344 + # (jump/skip over text sections between conversations)
345 + base_offset = 0
346 +
347 + # Common prefix size for all conversations (only 1 sample required)
348 + common_prefix_text = ""
349 + common_prefix_tokens: int = args.input_common_prefix_num_tokens.sample(1)[0]
350 + if common_prefix_tokens > 0:
351 + # Using "." at the end to separate sentences
352 + common_prefix_text = (
353 + tokenizer.decode(list_of_tokens[: common_prefix_tokens - 2]) + "."
354 + )
355 + base_offset += common_prefix_tokens
356 +
357 + for conv_id in range(args.num_conversations):
358 + # Generate a single conversation
359 + messages: MessagesList = []
360 +
361 + nturns = turn_count[conv_id]
362 +
363 + # User prompt token count per turn (with lower limit)
364 + input_token_count: np.ndarray = args.input_num_tokens.sample(nturns)
365 + input_token_count = np.maximum(input_token_count, base_prompt_token_count)
366 +
367 + # Assistant answer token count per turn (with lower limit)
368 + output_token_count: np.ndarray = args.output_num_tokens.sample(nturns)
369 + output_token_count = np.maximum(output_token_count, 1)
370 +
371 + user_turn = True
372 + for turn_id in range(nturns):
373 + if user_turn:
374 + role = "user"
375 + num_tokens = input_token_count[turn_id]
376 +
377 + # Generate the user prompt,
378 + # use a unique prefix (the conv_id) for each conversation
379 + # (to avoid shared prefix between conversations)
380 + content = f"{conv_id} is a nice number... "
381 +
382 + if len(common_prefix_text) > 0 and turn_id == 0:
383 + content = common_prefix_text + content
384 +
385 + # Update the number of tokens left for the content
386 + num_tokens -= len(tokenizer.encode(content, add_special_tokens=False))
387 +
388 + if turn_id == 0:
389 + prefix_num_tokens = conv_prefix_tokens[conv_id]
390 + if prefix_num_tokens > 0:
391 + # Add prefix text (context) to the first turn
392 + start_offset = base_offset
393 + end_offset = start_offset + prefix_num_tokens
394 + assert len(list_of_tokens) > end_offset, (
395 + "Not enough input text to generate "
396 + f"{prefix_num_tokens} tokens for the "
397 + f"prefix text ({start_offset=}, {end_offset=})"
398 + )
399 +
400 + content += f"{conv_id}, " + tokenizer.decode(
401 + list_of_tokens[start_offset:end_offset]
402 + )
403 + base_offset += prefix_num_tokens
404 +
405 + # Add the actual user prompt/question after the prefix text
406 + content += base_prompt_text
407 + num_tokens -= base_prompt_token_count
408 +
409 + if num_tokens > 0:
410 + # Add text from the input file (to reach the desired token count)
411 + start_offset = base_offset + turn_id * input_token_count.max()
412 + end_offset = start_offset + num_tokens
413 + assert len(list_of_tokens) > end_offset, (
414 + f"Not enough input text to generate {num_tokens} tokens "
415 + f"for the prompt ({start_offset=}, {end_offset=})"
416 + )
417 +
418 + # Convert tokens back to text
419 + content += tokenizer.decode(list_of_tokens[start_offset:end_offset])
420 + else:
421 + role = "assistant"
422 + # This content will not be used as input to the LLM server
423 + # (actual answers will be used instead).
424 + # Content is only required to determine the min_tokens/max_tokens
425 + # (inputs to the LLM server).
426 + num_tokens = output_token_count[turn_id]
427 + assert len(list_of_tokens) > num_tokens, (
428 + f"Not enough input text to generate {num_tokens} "
429 + "tokens for assistant content"
430 + )
431 + content = tokenizer.decode(list_of_tokens[:num_tokens])
432 +
433 + # Append the user/assistant message to the list of messages
434 + messages.append({"role": role, "content": content})
435 + user_turn = not user_turn
436 +
437 + # Add the new conversation
438 + conversations[f"CONV_ID_{conv_id}"] = messages
439 +
440 + # Increase base offset for the next conversation
441 + base_offset += nturns
442 +
443 + if args.print_stats:
444 + print_conv_stats(conversations, tokenizer)
445 +
446 + return conversations
447 +
448 +
449 +def conversations_list_to_dict(input_list: ShareGptConversations) -> ConversationsMap:
450 + conversations: ConversationsMap = {}
451 +
452 + for item in input_list:
453 + conv_id: str = item["id"]
454 + assert isinstance(conv_id, str)
455 +
456 + assert conv_id not in conversations, (
457 + f"Conversation ID {conv_id} found more than once in the input"
458 + )
459 +
460 + messages: MessagesList = item["messages"]
461 + assert isinstance(messages, list), (
462 + f"Conversation messages should be a list (ID: {conv_id})"
463 + )
464 + assert len(messages) > 0, f"Conversation with no messages (ID: {conv_id})"
465 +
466 + conversations[conv_id] = messages
467 +
468 + logger.info(f"Using {len(conversations)} unique conversations (IDs)")
469 + assert len(conversations) == len(input_list)
470 +
471 + # Print statistics about the selected conversations
472 + stats: list[dict[str, Any]] = []
473 + for conv_data in conversations.values():
474 + stats.append({"num_turns": len(conv_data)})
475 +
476 + print(TEXT_SEPARATOR)
477 + print(f"{Color.YELLOW}Conversations statistics:{Color.RESET}")
478 + print(TEXT_SEPARATOR)
479 + percentiles = [0.25, 0.5, 0.75, 0.9, 0.99, 0.999, 0.9999]
480 + conv_stats = pd.DataFrame(stats).describe(percentiles=percentiles)
481 + print(conv_stats.transpose())
482 + print(TEXT_SEPARATOR)
483 +
484 + return conversations
485 +
486 +
487 +def conversations_dict_to_list(input_dict: ConversationsMap) -> ShareGptConversations:
488 + output: ShareGptConversations = []
489 + for conv_id, conv_data in input_dict.items():
490 + new_item = {"id": conv_id, "messages": conv_data}
491 + output.append(new_item)
492 +
493 + return output
benchmarks/multi_turn/bench_utils.py new
+25
@@ -0,0 +1,25 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import logging
4 +from enum import Enum
5 +
6 +
7 +class Color(str, Enum):
8 + RED = "\033[91m"
9 + GREEN = "\033[92m"
10 + BLUE = "\033[94m"
11 + PURPLE = "\033[95m"
12 + CYAN = "\033[96m"
13 + YELLOW = "\033[93m"
14 + RESET = "\033[0m"
15 +
16 +
17 +TEXT_SEPARATOR = "-" * 100
18 +
19 +# Configure the logger
20 +logging.basicConfig(
21 + level=logging.INFO,
22 + format="%(asctime)s [%(levelname)s] - %(message)s",
23 + datefmt="%d-%m-%Y %H:%M:%S",
24 +)
25 +logger = logging.getLogger(__name__)
benchmarks/multi_turn/benchmark_serving_multi_turn.py new
+1557
@@ -0,0 +1,1557 @@
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, Optional, Union
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 +NUM_TOKENS_FROM_DATASET = 0
36 +TERM_SIGNAL = None
37 +
38 +
39 +class ConversationSampling(str, Enum):
40 + ROUND_ROBIN = "round_robin"
41 + RANDOM = "random"
42 +
43 + def __str__(self):
44 + return self.value
45 +
46 +
47 +class ClientArgs(NamedTuple):
48 + seed: int
49 + max_num_requests: Optional[int]
50 + skip_first_turn: bool
51 + max_turns: Optional[int]
52 + max_active_conversations: int
53 + verbose: bool
54 + print_content: bool
55 + verify_output: bool
56 + conversation_sampling: ConversationSampling
57 + request_rate: float
58 +
59 +
60 +class RequestArgs(NamedTuple):
61 + chat_url: str
62 + model: str
63 + stream: bool
64 + limit_min_tokens: int # Use negative value for no limit
65 + limit_max_tokens: int # Use negative value for no limit
66 +
67 +
68 +class BenchmarkArgs(NamedTuple):
69 + url: str
70 + num_clients: int
71 + early_stop: bool
72 +
73 +
74 +class ServerResponse(NamedTuple):
75 + valid: bool
76 + ttft_ms: float # time to first chunk
77 + tpot_ms: float # time per output chunk (one or more tokens)
78 + latency_ms: float
79 + start_time_ms: float
80 + first_chunk: str # first chunk of the content
81 + content: str # includes the first_chunk
82 + num_chunks: int
83 +
84 + def __str__(self) -> str:
85 + return f"ttft_ms {self.ttft_ms:.2f}, tpot_ms {self.tpot_ms:.2f}, latency_ms {self.latency_ms:.2f}" # noqa: E501
86 +
87 +
88 +class RequestStats(NamedTuple):
89 + ttft_ms: float
90 + tpot_ms: float
91 + latency_ms: float
92 + start_time_ms: float
93 + input_num_turns: int
94 + input_num_tokens: int
95 + output_num_tokens: int
96 + output_num_chunks: int
97 + output_num_first_chunk_tokens: int
98 + approx_cached_percent: float
99 + conversation_id: str
100 + client_id: int
101 +
102 + def __str__(self) -> str:
103 + return (
104 + 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
105 + 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
106 + f"approx_cached_percent {self.approx_cached_percent:.2f}%"
107 + )
108 +
109 +
110 +class MetricStats:
111 + def __init__(self) -> None:
112 + self.min: Optional[float] = None
113 + self.max: Optional[float] = None
114 + self.avg: Optional[float] = None
115 + self.sum = 0.0
116 + self.count = 0
117 +
118 + def update(self, value: float) -> None:
119 + if self.min is None:
120 + self.min = value
121 + else:
122 + self.min = min(self.min, value)
123 +
124 + if self.max is None:
125 + self.max = value
126 + else:
127 + self.max = max(self.max, value)
128 +
129 + self.sum += value
130 + self.count += 1
131 + self.avg = self.sum / self.count
132 +
133 + def __repr__(self) -> str:
134 + if self.count == 0:
135 + return "no data"
136 + return f"avg: {self.avg:>10.3f}, min: {self.min:>10.3f}, max: {self.max:>10.3f}"
137 +
138 +
139 +class MovingAverage:
140 + def __init__(self, window_size: int) -> None:
141 + self.window_size = window_size
142 + self.window = np.zeros(window_size)
143 + self.index = 0
144 + self.sum = 0.0
145 + self.count = 0
146 + self.avg: Optional[float] = None
147 +
148 + def update(self, new_value: float) -> None:
149 + if self.count < self.window_size:
150 + # Filling up the window
151 + self.sum += new_value
152 + self.window[self.count] = new_value
153 + self.count += 1
154 + else:
155 + # Window is full, start replacing old values
156 + old_value = self.window[self.index]
157 + self.sum = self.sum - old_value + new_value
158 + self.window[self.index] = new_value
159 + self.index = (self.index + 1) % self.window_size
160 +
161 + self.avg = self.sum / self.count
162 +
163 + def __repr__(self) -> str:
164 + if self.count == 0:
165 + return "no data"
166 + return f"avg: {self.avg:>10.3f} ({self.count} samples)"
167 +
168 +
169 +class DebugStats:
170 + def __init__(self, logger: logging.Logger, window_size: int) -> None:
171 + self.logger = logger
172 + self.metrics: dict[str, Union[MovingAverage, MetricStats]] = {
173 + "moving_avg_ttft_ms": MovingAverage(window_size),
174 + "moving_avg_tpot_ms": MovingAverage(window_size),
175 + "ttft_ms": MetricStats(),
176 + "tpot_ms": MetricStats(),
177 + "latency_ms": MetricStats(),
178 + "input_num_turns": MetricStats(),
179 + "input_num_tokens": MetricStats(),
180 + "output_num_tokens": MetricStats(),
181 + }
182 +
183 + def update(self, data: RequestStats) -> None:
184 + self.metrics["ttft_ms"].update(data.ttft_ms)
185 + self.metrics["moving_avg_ttft_ms"].update(data.ttft_ms)
186 + self.metrics["tpot_ms"].update(data.tpot_ms)
187 + self.metrics["moving_avg_tpot_ms"].update(data.tpot_ms)
188 + self.metrics["latency_ms"].update(data.latency_ms)
189 + self.metrics["input_num_turns"].update(data.input_num_turns)
190 + self.metrics["input_num_tokens"].update(data.input_num_tokens)
191 + self.metrics["output_num_tokens"].update(data.output_num_tokens)
192 +
193 + def print(self) -> None:
194 + self.logger.info("-" * 50)
195 + for k, v in self.metrics.items():
196 + kv_info = f"[{k:25}] {v}"
197 + self.logger.info(kv_info)
198 + self.logger.info("-" * 50)
199 +
200 +
201 +# Must support Python 3.8, we can't use str.removeprefix(prefix)
202 +# introduced in Python 3.9
203 +def remove_prefix(text: str, prefix: str) -> str:
204 + if text.startswith(prefix):
205 + return text[len(prefix) :]
206 + return text
207 +
208 +
209 +def nanosec_to_millisec(value: float) -> float:
210 + return value / 1000000.0
211 +
212 +
213 +def nanosec_to_sec(value: float) -> float:
214 + return value / 1000000000.0
215 +
216 +
217 +async def send_request(
218 + session: aiohttp.ClientSession,
219 + messages: list[dict[str, str]],
220 + chat_url: str,
221 + model: str,
222 + stream: bool = True,
223 + min_tokens: Optional[int] = None,
224 + max_tokens: Optional[int] = None,
225 +) -> ServerResponse:
226 + payload = {
227 + "model": model,
228 + "messages": messages,
229 + "seed": 0,
230 + "temperature": 0.0,
231 + }
232 +
233 + if stream:
234 + payload["stream"] = True
235 + payload["stream_options"] = {"include_usage": False}
236 +
237 + if min_tokens is not None:
238 + payload["min_tokens"] = min_tokens
239 +
240 + if max_tokens is not None:
241 + payload["max_tokens"] = max_tokens
242 +
243 + headers = {"Content-Type": "application/json"}
244 +
245 + # Calculate the timeout for the request
246 + timeout_sec = 120
247 + if max_tokens is not None:
248 + # Assume TPOT of 200ms and use max_tokens to determine timeout
249 + timeout_sec = max(timeout_sec, int(max_tokens * 0.2))
250 + timeout = aiohttp.ClientTimeout(total=timeout_sec)
251 +
252 + valid_response = True
253 + ttft: Optional[float] = None
254 + chunk_delay: list[int] = []
255 + latency: Optional[float] = None
256 + first_chunk = ""
257 + generated_text = ""
258 +
259 + start_time: int = time.perf_counter_ns()
260 + most_recent_timestamp: int = start_time
261 +
262 + async with session.post(
263 + url=chat_url, json=payload, headers=headers, timeout=timeout
264 + ) as response:
265 + http_status = HTTPStatus(response.status)
266 + if http_status == HTTPStatus.OK:
267 + async for chunk_bytes in response.content:
268 + chunk_bytes = chunk_bytes.strip()
269 + if not chunk_bytes:
270 + continue
271 +
272 + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ")
273 + if chunk == "[DONE]":
274 + # End of stream
275 + latency = time.perf_counter_ns() - start_time
276 + elif stream is False:
277 + data = json.loads(chunk)
278 + message = data["choices"][0]["message"]
279 + assert message["role"] == "assistant"
280 + generated_text += message["content"]
281 + else:
282 + timestamp: int = time.perf_counter_ns()
283 + data = json.loads(chunk)
284 +
285 + # Delta is the new content/text/data
286 + delta = data["choices"][0]["delta"]
287 + if delta.get("content", None):
288 + if ttft is None:
289 + # First token
290 + first_token_time = time.perf_counter_ns()
291 + ttft = first_token_time - start_time
292 + first_chunk = delta["content"]
293 + else:
294 + # Decoding phase
295 + chunk_delay.append(timestamp - most_recent_timestamp)
296 +
297 + generated_text += delta["content"]
298 +
299 + most_recent_timestamp = timestamp
300 + else:
301 + valid_response = False
302 + content = await response.text()
303 + logger.warning(
304 + f"{Color.YELLOW}Received HTTP status {http_status.value} "
305 + f"({http_status.phrase}): {content}{Color.RESET}"
306 + )
307 +
308 + if latency is None:
309 + latency = -1.0
310 + if valid_response:
311 + # Streaming is disabled, latency was not set
312 + latency = time.perf_counter_ns() - start_time
313 +
314 + if ttft is None:
315 + # The response was a single chunk
316 + ttft = latency
317 +
318 + # Each chunk may include more than one token
319 + tpot: float = mean(chunk_delay) if len(chunk_delay) > 0 else 0.0
320 + num_chunks: int = len(chunk_delay)
321 +
322 + sr = ServerResponse(
323 + valid=valid_response,
324 + ttft_ms=nanosec_to_millisec(ttft) if ttft > 0.0 else -1.0,
325 + tpot_ms=nanosec_to_millisec(tpot),
326 + latency_ms=nanosec_to_millisec(latency),
327 + start_time_ms=nanosec_to_millisec(start_time),
328 + first_chunk=first_chunk,
329 + content=generated_text,
330 + num_chunks=num_chunks,
331 + )
332 + return sr
333 +
334 +
335 +def get_short_string(input: str) -> str:
336 + n = 20
337 + if len(input) < 400:
338 + return input
339 +
340 + return f"{input[:n]}...{input[-n:]}"
341 +
342 +
343 +def get_token_count(tokenizer: AutoTokenizer, text: str) -> int:
344 + return len(tokenizer(text, add_special_tokens=False).input_ids)
345 +
346 +
347 +def get_messages_token_count(
348 + tokenizer: AutoTokenizer, messages: list[dict[str, str]]
349 +) -> int:
350 + token_count = 0
351 + for m in messages:
352 + token_count += get_token_count(tokenizer, m["content"])
353 +
354 + return token_count
355 +
356 +
357 +async def send_turn(
358 + session: aiohttp.ClientSession,
359 + client_id: int,
360 + conv_id: str,
361 + conversation_messages: MessagesList,
362 + messages_to_use: int,
363 + tokenizer: AutoTokenizer,
364 + req_args: RequestArgs,
365 + verbose: bool,
366 + verify_output: bool,
367 +) -> Optional[RequestStats]:
368 + assert messages_to_use > 0
369 + assert messages_to_use <= len(conversation_messages)
370 +
371 + messages = conversation_messages[:messages_to_use]
372 +
373 + # Index of the next message (the role should be "user")
374 + index = messages_to_use - 1
375 +
376 + # Verify that the message has only two keys, "role" and "content"
377 + assert len(messages[index].keys()) == 2
378 + assert "role" in messages[index] and "content" in messages[index]
379 + assert messages[index]["role"] == "user", (
380 + f"Failed on conversation ID {conv_id}, message role should be user"
381 + )
382 +
383 + if verbose:
384 + print(
385 + f"{Color.CYAN}Messages (conversation ID {conv_id},"
386 + f" {len(messages)} turns):{Color.RESET}",
387 + messages,
388 + )
389 +
390 + # None means that there is no upper/lower limit for the output token count
391 + min_tokens = None if req_args.limit_min_tokens < 0 else req_args.limit_min_tokens
392 + max_tokens = None if req_args.limit_max_tokens < 0 else req_args.limit_max_tokens
393 +
394 + if len(conversation_messages) > messages_to_use:
395 + # The conversation contains an assistant answer for the next user prompt
396 + if (
397 + min_tokens == NUM_TOKENS_FROM_DATASET
398 + or max_tokens == NUM_TOKENS_FROM_DATASET
399 + ):
400 + # Compute number of tokens in the answer (from the input conversation)
401 + assistant_answer = conversation_messages[messages_to_use]
402 + answer_num_tokens = get_token_count(tokenizer, assistant_answer["content"])
403 + assert assistant_answer["role"] == "assistant"
404 +
405 + if min_tokens == NUM_TOKENS_FROM_DATASET:
406 + min_tokens = max(1, answer_num_tokens)
407 +
408 + if max_tokens == NUM_TOKENS_FROM_DATASET:
409 + max_tokens = max(1, answer_num_tokens)
410 +
411 + # Send the current conversation to LLM and get a response
412 + response: ServerResponse = await send_request(
413 + session,
414 + messages,
415 + req_args.chat_url,
416 + req_args.model,
417 + req_args.stream,
418 + min_tokens,
419 + max_tokens,
420 + )
421 +
422 + if response.valid is False:
423 + # Request failed
424 + return None
425 +
426 + # Compute number of tokens in input / output
427 + input_num_tokens = get_messages_token_count(tokenizer, messages)
428 +
429 + # Num tokens in the user's last question
430 + question_num_tokens = get_token_count(tokenizer, messages[index]["content"])
431 +
432 + # Num tokens in the history/context of the question
433 + assert input_num_tokens >= question_num_tokens
434 + history_num_tokens = input_num_tokens - question_num_tokens
435 +
436 + # Num tokens in the LLM's answer (first chunk and full answer)
437 + first_chunk_tokens = get_token_count(tokenizer, response.first_chunk)
438 +
439 + output_content = response.content
440 + output_num_tokens = get_token_count(tokenizer, output_content)
441 +
442 + # Prefix caching approximated cached percent
443 + approx_cached_percent = (
444 + 100.0 * (history_num_tokens / input_num_tokens) if input_num_tokens > 0 else 0.0
445 + )
446 +
447 + # Compute the correct TTFT and TPOT (based on tokens and not chunks).
448 + # Required because multiple output tokens may be bundled in a single chunk.
449 + if output_num_tokens > 1 and output_num_tokens > first_chunk_tokens:
450 + # More than one token and more than one chunk in the output
451 + decode_ms = response.latency_ms - response.ttft_ms
452 + decode_num_tokens = output_num_tokens - first_chunk_tokens
453 + tpot_ms = decode_ms / decode_num_tokens
454 + else:
455 + # In this case: output_num_tokens == first_chunk_tokens
456 + # Output was a single chunk (output_num_tokens > 1)
457 + # or even a single token (output_num_tokens == 1)
458 + tpot_ms = 0.0
459 +
460 + if first_chunk_tokens > 1:
461 + # First chunk had multiple tokens, adjust TTFT for a single token
462 + delta_ms = (first_chunk_tokens - 1) * tpot_ms
463 + ttft_ms = max(0.1, response.ttft_ms - delta_ms)
464 + else:
465 + # First chunk had only one token
466 + ttft_ms = response.ttft_ms
467 +
468 + rs = RequestStats(
469 + ttft_ms=ttft_ms,
470 + tpot_ms=tpot_ms,
471 + latency_ms=response.latency_ms,
472 + start_time_ms=response.start_time_ms,
473 + input_num_turns=len(messages),
474 + input_num_tokens=input_num_tokens,
475 + output_num_tokens=output_num_tokens,
476 + output_num_chunks=response.num_chunks,
477 + output_num_first_chunk_tokens=first_chunk_tokens,
478 + approx_cached_percent=approx_cached_percent,
479 + conversation_id=conv_id,
480 + client_id=client_id,
481 + )
482 +
483 + if verbose:
484 + print(
485 + f"\n{Color.YELLOW}Response ({output_num_tokens} tokens):{Color.RESET}",
486 + output_content,
487 + )
488 + print(f"{Color.YELLOW}Response metrics: {rs}{Color.RESET}")
489 + print("-" * 70)
490 +
491 + # Save the LLM's answer (will be used as part of the context for the next user turn)
492 + answer_index = messages_to_use
493 + if len(conversation_messages) > answer_index:
494 + assert conversation_messages[answer_index]["role"] == "assistant", (
495 + f"Failed on conversation ID {conv_id}, message role should be assistant"
496 + )
497 +
498 + orig_content = conversation_messages[answer_index]["content"]
499 + if verify_output:
500 + # Compare the new answer to the answer from the input file
501 + debug_info = (
502 + f"LLM/dataset answers do not match ({conv_id}):"
503 + f"\n'{get_short_string(output_content)}' (len: {len(output_content)}),"
504 + f"\n'{get_short_string(orig_content)}' (len: {len(orig_content)})"
505 + )
506 + if orig_content != output_content:
507 + raise ValueError(debug_info)
508 +
509 + # Update the answer
510 + conversation_messages[answer_index]["content"] = output_content
511 + else:
512 + # A user prompt that has no answer, add the answer as a new message
513 + new_answer = {"role": "assistant", "content": output_content}
514 + conversation_messages.append(new_answer)
515 +
516 + return rs
517 +
518 +
519 +async def poisson_sleep(request_rate: float, verbose: bool = False) -> None:
520 + # Generate a random time interval from the Poisson distribution
521 + assert request_rate > 0
522 +
523 + interval = np.random.exponential(1.0 / request_rate)
524 + if verbose:
525 + logger.info(f"Sleeping for {interval:.3f} seconds...")
526 + await asyncio.sleep(interval)
527 +
528 +
529 +async def client_main(
530 + args: ClientArgs,
531 + req_args: RequestArgs,
532 + client_id: int,
533 + tokenizer: AutoTokenizer,
534 + stop_event: mp.Event, # type: ignore
535 + task_queue: mp.Queue,
536 + result_queue: mp.Queue,
537 + conv_queue: mp.Queue,
538 +) -> None:
539 + logger.info(
540 + 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
541 + )
542 +
543 + random.seed(args.seed)
544 + np.random.seed(args.seed)
545 +
546 + # Active conversations
547 + active_convs: ConversationsMap = {}
548 + conv_id_queue: deque = deque(maxlen=args.max_active_conversations)
549 +
550 + # Keep track of how many messages have been used for each conversation
551 + turns_count: Counter = Counter()
552 + num_successes = 0
553 + num_failures = 0
554 +
555 + # Track the timestamp (time.perf_counter())
556 + # of the last turn per conversation (only for debug)
557 + time_of_last_turn: dict[ConvId, float] = {}
558 +
559 + # Flag that indicates that there are no new tasks (conversations) for the client
560 + task_queue_empty = False
561 +
562 + async with aiohttp.ClientSession() as session:
563 + # Print progress
564 +
565 + while task_queue_empty is False:
566 + result = None
567 +
568 + if (
569 + args.max_num_requests
570 + and num_successes + num_failures == args.max_num_requests
571 + ):
572 + logger.info(
573 + f"{Color.YELLOW}Client {client_id} reached "
574 + f"request limit{Color.RESET}"
575 + )
576 + break
577 +
578 + if stop_event.is_set(): # type: ignore
579 + logger.info(
580 + f"{Color.YELLOW}Client {client_id} received "
581 + f"a termination signal{Color.RESET}"
582 + )
583 + break
584 +
585 + while (
586 + len(active_convs) < args.max_active_conversations
587 + and task_queue_empty is False
588 + ):
589 + # Get a new conversation from the task queue
590 + conv_id, messages = task_queue.get()
591 +
592 + if conv_id is TERM_SIGNAL:
593 + task_queue_empty = True
594 + break
595 +
596 + if args.skip_first_turn:
597 + # Skip the first turn (both user and assistant),
598 + # relevant if warmup was enabled.
599 + # Default turns_count[conv_id] will be zero if conv_id
600 + # was never inserted/updated in turns_count.
601 + turns_count[conv_id] += 2
602 +
603 + if turns_count[conv_id] < len(messages):
604 + # Add new conversation
605 + active_convs[conv_id] = messages
606 + conv_id_queue.append(conv_id)
607 +
608 + if args.verbose:
609 + logger.info(
610 + f"{Color.GREEN}Client {client_id} will use conversation ID {conv_id} (active conversations {len(active_convs)}){Color.RESET}" # noqa: E501
611 + )
612 +
613 + elif args.verbose:
614 + # No more messages (conversation finished during the warmup)
615 + logger.info(
616 + f"{Color.YELLOW}Client {client_id} will not use conversation ID {conv_id} (all {len(messages)} messages already sent){Color.RESET}" # noqa: E501
617 + )
618 +
619 + if len(active_convs) == 0 or task_queue_empty:
620 + logger.info(
621 + f"{Color.YELLOW}Client {client_id} has no more work{Color.RESET}"
622 + )
623 + break
624 +
625 + # Pick an active conversation for the next request
626 + if args.conversation_sampling == ConversationSampling.ROUND_ROBIN:
627 + conv_id = conv_id_queue.pop()
628 + else:
629 + # ConversationSampling.RANDOM
630 + active_ids = list(active_convs.keys())
631 + conv_id = random.choice(active_ids)
632 +
633 + messages = active_convs[conv_id]
634 + assert isinstance(messages, list) and len(messages) > 0
635 +
636 + # Update the amount of messages to use
637 + turns_count[conv_id] += 1
638 + current_turn = turns_count[conv_id]
639 +
640 + assert current_turn < len(messages), (
641 + f"Turn number {current_turn} is invalid for conversation ID {conv_id}"
642 + f" that has only {len(messages)} messages"
643 + )
644 +
645 + if args.verbose:
646 + curr_time_sec: float = time.perf_counter()
647 + time_since_last_turn: Union[str, float] = "N/A"
648 + if conv_id in time_of_last_turn:
649 + time_since_last_turn = round(
650 + curr_time_sec - time_of_last_turn[conv_id], 3
651 + )
652 + logger.info(
653 + f"Client {client_id} using conversation ID {conv_id} (turn: {current_turn}, time since last turn [sec]: {time_since_last_turn})" # noqa: E501
654 + )
655 + time_of_last_turn[conv_id] = curr_time_sec
656 +
657 + success = True
658 + try:
659 + result = await send_turn(
660 + session,
661 + client_id,
662 + conv_id,
663 + messages,
664 + current_turn,
665 + tokenizer,
666 + req_args,
667 + args.print_content,
668 + args.verify_output,
669 + )
670 + if result is not None:
671 + result_queue.put(result)
672 + else:
673 + # None means that the request failed,
674 + # and should not be added to the statistics.
675 + success = False
676 + num_failures += 1
677 +
678 + logger.warning(
679 + f"{Color.YELLOW}Client {client_id} - Request rejected during conversation ID {conv_id} (turn: {current_turn}){Color.RESET}" # noqa: E501
680 + )
681 +
682 + # Remove the conversation (should not be used again)
683 + active_convs.pop(conv_id)
684 +
685 + except asyncio.exceptions.TimeoutError:
686 + num_failures += 1
687 + logger.exception(
688 + f"{Color.RED}Client {client_id} - Timeout during conversation ID {conv_id} (turn: {current_turn}){Color.RESET}" # noqa: E501
689 + )
690 + break # Exit gracefully instead of raising an error
691 +
692 + except Exception:
693 + num_failures += 1
694 + logger.exception(
695 + f"{Color.RED}Client {client_id} - Exception during conversation ID {conv_id} (turn: {current_turn}){Color.RESET}" # noqa: E501
696 + )
697 + break # Exit gracefully instead of raising an error
698 +
699 + if success:
700 + num_successes += 1
701 +
702 + # Update the turns counter to include the LLM response
703 + # The LLM response will be used as context for the next user turn
704 + turns_count[conv_id] += 1
705 +
706 + max_turns = len(messages)
707 + if args.max_turns is not None:
708 + # Limit the number of turns in the conversation
709 + max_turns = min(args.max_turns, max_turns)
710 +
711 + if turns_count[conv_id] >= max_turns:
712 + # Conversation has no more turns (no longer active)
713 + # save the updated conversation (with the LLM server's answer)
714 + conv_queue.put((conv_id, active_convs.pop(conv_id)))
715 + if args.verbose:
716 + logger.info(
717 + f"{Color.GREEN}Client {client_id} finished "
718 + f"conversation ID {conv_id}{Color.RESET}"
719 + )
720 + else:
721 + # Conversation is not finished, insert it at the back of the queue
722 + conv_id_queue.appendleft(conv_id)
723 +
724 + # Sleep between requests (if lambda is positive)
725 + if args.request_rate > 0:
726 + await poisson_sleep(args.request_rate, args.verbose)
727 +
728 + # Send indication that the client is done
729 + conv_queue.put((TERM_SIGNAL, TERM_SIGNAL))
730 +
731 + logger.info(
732 + f"{Color.CYAN}Client {client_id} is done "
733 + f"({num_successes=}, {num_failures=}){Color.RESET}"
734 + )
735 +
736 +
737 +def worker_function(
738 + client_id: int,
739 + tokenizer: AutoTokenizer,
740 + client_args: ClientArgs,
741 + req_args: RequestArgs,
742 + stop_event: mp.Event, # type: ignore
743 + task_queue: mp.Queue,
744 + result_queue: mp.Queue,
745 + conv_queue: mp.Queue,
746 +) -> None:
747 + asyncio.run(
748 + client_main(
749 + client_args,
750 + req_args,
751 + client_id,
752 + tokenizer,
753 + stop_event,
754 + task_queue,
755 + result_queue,
756 + conv_queue,
757 + )
758 + )
759 +
760 +
761 +def get_client_config(
762 + args: argparse.Namespace, input_conv: ConversationsMap
763 +) -> tuple[ClientArgs, RequestArgs]:
764 + if args.num_clients < 1:
765 + raise ValueError("Number of clients must be a positive number")
766 +
767 + if len(input_conv) < args.num_clients:
768 + raise ValueError(
769 + "Number of conversations must be equal or larger than the number of clients"
770 + )
771 +
772 + max_req_per_client: Optional[int] = None
773 + if args.max_num_requests is not None:
774 + # Max number of requests per client
775 + req_per_client = args.max_num_requests // args.num_clients
776 + if req_per_client < 1:
777 + raise ValueError("Number of requests should be at least one per client")
778 + max_req_per_client = req_per_client
779 +
780 + max_active_conversations = args.max_active_conversations
781 + if max_active_conversations is None:
782 + # Each client will have only one active conversation at a time
783 + max_active_conversations = args.num_clients
784 +
785 + if max_active_conversations > len(input_conv):
786 + raise ValueError(
787 + f"Max active conversations {max_active_conversations} "
788 + "must be equal or less than the total number of conversations"
789 + )
790 +
791 + # Max number of active conversations per client
792 + max_active_conv_per_client = max_active_conversations // args.num_clients
793 + if max_active_conv_per_client < 1:
794 + raise ValueError(
795 + f"Max active conversations {max_active_conversations} "
796 + "must be equal or greater than the number of clients"
797 + )
798 +
799 + # Skip the first user turn (as part of the warmup)
800 + skip_first_turn = args.warmup_step
801 +
802 + # Common arguments for all clients
803 + client_args = ClientArgs(
804 + seed=args.seed,
805 + max_num_requests=max_req_per_client,
806 + skip_first_turn=skip_first_turn,
807 + max_turns=args.max_turns,
808 + max_active_conversations=max_active_conv_per_client,
809 + verbose=args.verbose,
810 + print_content=args.print_content,
811 + verify_output=args.verify_output,
812 + conversation_sampling=args.conversation_sampling,
813 + request_rate=args.request_rate,
814 + )
815 +
816 + if args.limit_min_tokens > 0 or args.limit_max_tokens > 0:
817 + if args.limit_min_tokens < 1 or args.limit_max_tokens < 1:
818 + raise ValueError(
819 + "Invalid min/max tokens limits (both limits should be provided)"
820 + )
821 + if args.limit_min_tokens > args.limit_max_tokens:
822 + raise ValueError(
823 + "Invalid min/max tokens limits (min should not be larger than max)"
824 + )
825 +
826 + # Arguments for API requests
827 + chat_url = f"{args.url}/v1/chat/completions"
828 + req_args = RequestArgs(
829 + chat_url=chat_url,
830 + model=args.model,
831 + stream=not args.no_stream,
832 + limit_min_tokens=args.limit_min_tokens,
833 + limit_max_tokens=args.limit_max_tokens,
834 + )
835 +
836 + return client_args, req_args
837 +
838 +
839 +async def main_mp(
840 + client_args: ClientArgs,
841 + req_args: RequestArgs,
842 + bench_args: BenchmarkArgs,
843 + tokenizer: AutoTokenizer,
844 + input_conv: ConversationsMap,
845 +) -> tuple[ConversationsMap, list[RequestStats]]:
846 + # An event that will trigger graceful termination of all the clients
847 + stop_event = mp.Event()
848 +
849 + # Queue for input conversations (from the input file/dataset)
850 + task_queue: mp.Queue = mp.Queue()
851 +
852 + # Queue for client measurements (TTFT, TPOT, etc. for each request)
853 + result_queue: mp.Queue = mp.Queue()
854 +
855 + # Queue for output conversations (with the LLM answers, sent by the server)
856 + conv_queue: mp.Queue = mp.Queue()
857 + output_conv: ConversationsMap = {}
858 + client_metrics: list[RequestStats] = []
859 +
860 + # Start all clients
861 + start_time = time.perf_counter_ns()
862 + logger.info(f"{Color.GREEN}Starting {bench_args.num_clients} clients{Color.RESET}")
863 +
864 + clients = []
865 + for client_id in range(bench_args.num_clients):
866 + client = mp.Process(
867 + name=f"client_{client_id}",
868 + target=worker_function,
869 + args=(
870 + client_id,
871 + tokenizer,
872 + client_args,
873 + req_args,
874 + stop_event,
875 + task_queue,
876 + result_queue,
877 + conv_queue,
878 + ),
879 + )
880 + clients.append(client)
881 + client.start()
882 +
883 + # Submit all the input conversations as tasks for the clients
884 + for conv_id, messages in input_conv.items():
885 + task_queue.put((conv_id, messages))
886 +
887 + # Add termination signals for clients
888 + for _ in range(bench_args.num_clients):
889 + task_queue.put((TERM_SIGNAL, TERM_SIGNAL))
890 +
891 + # Collect the updated conversations from all clients
892 + num_clients_finished = 0
893 + total_convs = len(input_conv)
894 +
895 + debug_stats = DebugStats(logger, min(15 * bench_args.num_clients, 500))
896 +
897 + while num_clients_finished < bench_args.num_clients:
898 + # Collect updated conversation
899 + conv_id, messages = conv_queue.get()
900 +
901 + # Collect results (measurements)
902 + while not result_queue.empty():
903 + new_data = result_queue.get()
904 + client_metrics.append(new_data)
905 + debug_stats.update(new_data)
906 +
907 + if conv_id is TERM_SIGNAL:
908 + num_clients_finished += 1
909 + logger.info(
910 + f"{Color.CYAN}{num_clients_finished} out of "
911 + f"{bench_args.num_clients} clients finished{Color.RESET}"
912 + )
913 +
914 + if bench_args.early_stop and not stop_event.is_set():
915 + # Once one client finished, stop all other clients.
916 + # there is no reason to continue the benchmark with fewer clients.
917 + logger.info(
918 + f"{Color.YELLOW}Sending termination signal to clients{Color.RESET}"
919 + )
920 + stop_event.set()
921 + else:
922 + output_conv[conv_id] = messages
923 +
924 + finished_convs = len(output_conv)
925 + percent = finished_convs / total_convs
926 +
927 + # Tuned to control the print rate (can be changed if required)
928 + print_cycle = max(3, int(bench_args.num_clients / 4))
929 +
930 + if finished_convs % print_cycle == 0:
931 + runtime_sec = nanosec_to_sec(time.perf_counter_ns() - start_time)
932 + logger.info(
933 + f"{Color.CYAN}Finished {finished_convs} out of {total_convs} conversations ({percent:.0%}), " # noqa: E501
934 + 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
935 + )
936 +
937 + rps: Union[str, float] = round(len(client_metrics) / runtime_sec, 3)
938 + if len(client_metrics) < (5 * bench_args.num_clients):
939 + # Do not estimate the RPS if the number of samples is very low
940 + # (threshold can be tuned if needed)
941 + rps = "N/A"
942 +
943 + runtime_left_sec: Union[str, float] = round(
944 + (runtime_sec / finished_convs) * (total_convs - finished_convs), 3
945 + )
946 + if percent < 0.05:
947 + # If less than 5% of the conversations were not finished,
948 + # the estimation will probably be very inaccurate
949 + # (threshold can be tuned if needed).
950 + runtime_left_sec = "N/A"
951 +
952 + logger.info(
953 + f"{Color.CYAN}Estimated req/sec {rps}, estimated runtime left {runtime_left_sec} sec{Color.RESET}" # noqa: E501
954 + )
955 + debug_stats.print()
956 +
957 + logger.info(
958 + f"{Color.CYAN}All {bench_args.num_clients} clients finished{Color.RESET}"
959 + )
960 +
961 + # At this point all the clients finished,
962 + # collect results (TTFT, TPOT, etc.) from all the clients.
963 + # This needs to happens before calling join on the clients
964 + # (result_queue should be emptied).
965 + while not result_queue.empty():
966 + client_metrics.append(result_queue.get())
967 +
968 + logger.info(f"Collected {len(client_metrics)} samples from all the clients")
969 +
970 + # Wait for all clients to finish
971 + for client in clients:
972 + logger.info(
973 + f"{Color.CYAN}Waiting for client {client.name} "
974 + f"(is alive: {client.is_alive()}){Color.RESET}"
975 + )
976 +
977 + client.join(timeout=120)
978 +
979 + if client.is_alive():
980 + logger.warning(
981 + f"{Color.YELLOW}Client {client.name} will be terminated{Color.RESET}"
982 + )
983 + client.terminate()
984 +
985 + exitcode = client.exitcode
986 + if exitcode != 0:
987 + logger.error(
988 + f"{Color.RED}Client {client.name} exited "
989 + f"with exit code {exitcode}{Color.RESET}"
990 + )
991 +
992 + logger.info(
993 + f"All {bench_args.num_clients} clients exited (successfully "
994 + f"finished {len(output_conv)} out of {total_convs} conversations)"
995 + )
996 +
997 + # Queues should be closed, required to avoid hang at interpreter shutdown
998 + unfinished_tasks = 0
999 + while not task_queue.empty():
1000 + task_queue.get()
1001 + unfinished_tasks += 1
1002 +
1003 + if unfinished_tasks > 0:
1004 + # Can happen if not all tasks (conversations) have finished.
1005 + # May happen if --max-num-requests was used,
1006 + # or if an error occurred in one of the clients.
1007 + logger.debug(f"Discarding {unfinished_tasks} unfinished tasks")
1008 +
1009 + task_queue.close()
1010 + task_queue.join_thread()
1011 +
1012 + result_queue.close()
1013 + result_queue.join_thread()
1014 +
1015 + conv_queue.close()
1016 + conv_queue.join_thread()
1017 +
1018 + return output_conv, client_metrics
1019 +
1020 +
1021 +def get_filename_with_timestamp(label: str, extension: str) -> str:
1022 + time_now = datetime.now()
1023 + timestamp = time_now.strftime("%d-%m-%Y_%H-%M-%S")
1024 + filename = f"{label}__{timestamp}.{extension}"
1025 + return filename
1026 +
1027 +
1028 +def process_statistics(
1029 + client_metrics: list[RequestStats],
1030 + warmup_percentages: list[float],
1031 + test_params: dict,
1032 + verbose: bool,
1033 + gen_conv_args: Optional[GenConvArgs] = None,
1034 + excel_output: bool = False,
1035 +) -> None:
1036 + if len(client_metrics) == 0:
1037 + logger.info("No samples to process")
1038 + return
1039 +
1040 + logger.info(f"Processing {len(client_metrics)} samples...")
1041 +
1042 + raw_data = pd.DataFrame(client_metrics)
1043 +
1044 + if verbose:
1045 + # Calculate the time between user turns in each conversation (in a new column)
1046 + raw_data = raw_data.sort_values(by=["conversation_id", "start_time_ms"])
1047 + raw_data["time_between_user_turns_sec"] = raw_data.groupby("conversation_id")[
1048 + "start_time_ms"
1049 + ].diff()
1050 +
1051 + # Convert milliseconds to seconds
1052 + raw_data["time_between_user_turns_sec"] = (
1053 + raw_data["time_between_user_turns_sec"] / 1000.0
1054 + )
1055 +
1056 + # Final raw data should be sorted by time
1057 + raw_data = raw_data.sort_values(by=["start_time_ms"])
1058 + raw_data["end_time_ms"] = raw_data["start_time_ms"] + raw_data["latency_ms"]
1059 +
1060 + percentiles = [0.25, 0.5, 0.75, 0.9]
1061 +
1062 + # Add more percentiles if there are enough samples
1063 + if len(raw_data) >= 100:
1064 + percentiles.append(0.99)
1065 +
1066 + if len(raw_data) >= 1000:
1067 + percentiles.append(0.999)
1068 +
1069 + if len(raw_data) >= 10000:
1070 + percentiles.append(0.9999)
1071 +
1072 + # Set precision for numbers in the output text (the dataframes)
1073 + pd.set_option("display.precision", 2)
1074 +
1075 + # Exclude parameters from RequestStats
1076 + exclude = [
1077 + "start_time_ms",
1078 + "end_time_ms",
1079 + "output_num_first_chunk_tokens",
1080 + "approx_cached_percent",
1081 + "conversation_id",
1082 + "client_id",
1083 + ]
1084 +
1085 + print(TEXT_SEPARATOR)
1086 + print(f"{Color.YELLOW}Parameters:{Color.RESET}")
1087 + for k, v in test_params.items():
1088 + print(f"{k}={v}")
1089 +
1090 + # conversations generation parameters
1091 + if gen_conv_args is not None:
1092 + gen_params = {
1093 + "text_files": ", ".join(gen_conv_args.text_files),
1094 + "input_num_turns": str(gen_conv_args.input_num_turns),
1095 + "input_common_prefix_num_tokens": str(
1096 + gen_conv_args.input_common_prefix_num_tokens
1097 + ),
1098 + "input_prefix_num_tokens": str(gen_conv_args.input_prefix_num_tokens),
1099 + "input_num_tokens": str(gen_conv_args.input_num_tokens),
1100 + "output_num_tokens": str(gen_conv_args.output_num_tokens),
1101 + }
1102 +
1103 + print(f"{Color.YELLOW}Conversations Generation Parameters:{Color.RESET}")
1104 + for k, v in gen_params.items():
1105 + print(f"{k}={v}")
1106 +
1107 + print(TEXT_SEPARATOR)
1108 +
1109 + params_list = []
1110 + df_list = []
1111 + for percent in warmup_percentages:
1112 + # Select samples from the end (tail) of the dataframe
1113 + warmup_count = int(percent * len(raw_data))
1114 + tail_count = len(raw_data) - warmup_count
1115 + if tail_count == 0:
1116 + # No reason to process if the count of samples is zero
1117 + break
1118 +
1119 + df = raw_data.tail(tail_count)
1120 +
1121 + # Runtime is the diff between the end of the last request
1122 + # and the start of the first request
1123 + runtime_sec = df["end_time_ms"].iloc[-1] - df["start_time_ms"].iloc[0]
1124 +
1125 + # Convert milliseconds to seconds
1126 + runtime_sec = runtime_sec / 1000.0
1127 + requests_per_sec = float(len(df)) / runtime_sec
1128 +
1129 + params = {"runtime_sec": runtime_sec, "requests_per_sec": requests_per_sec}
1130 +
1131 + # Generate a summary of relevant metrics (and drop irrelevant data)
1132 + df = df.drop(columns=exclude).describe(percentiles=percentiles).transpose()
1133 +
1134 + # List for Excel file
1135 + params_list.append(params)
1136 + df_list.append(df)
1137 +
1138 + # Print the statistics summary
1139 + if percent > 0 or len(warmup_percentages) > 1:
1140 + print(
1141 + f"{Color.YELLOW}Statistics summary "
1142 + f"(assuming {percent:.0%} warmup samples):{Color.RESET}"
1143 + )
1144 + else:
1145 + print(f"{Color.YELLOW}Statistics summary:{Color.RESET}")
1146 +
1147 + for k, v in params.items():
1148 + if isinstance(v, float):
1149 + print(f"{k} = {v:.3f}")
1150 + else:
1151 + print(f"{k} = {v}")
1152 + print(TEXT_SEPARATOR)
1153 + print(df)
1154 + print(TEXT_SEPARATOR)
1155 +
1156 + if excel_output:
1157 + prefix = f"statistics_{test_params['num_clients']}_clients"
1158 + filename = get_filename_with_timestamp(prefix, "xlsx")
1159 +
1160 + with pd.ExcelWriter(filename, engine="xlsxwriter") as writer:
1161 + startrow = 0
1162 + test_params_df = pd.DataFrame([test_params])
1163 + test_params_df.to_excel(
1164 + writer, sheet_name="Summary", index=False, startrow=startrow
1165 + )
1166 + startrow += len(test_params_df) + 3
1167 +
1168 + if gen_conv_args is not None:
1169 + gen_params_df = pd.DataFrame([gen_params])
1170 + gen_params_df.to_excel(
1171 + writer, sheet_name="Summary", index=False, startrow=(startrow - 1)
1172 + )
1173 + startrow += len(gen_params_df) + 3
1174 +
1175 + for params, df_stats in zip(params_list, df_list):
1176 + df_params = pd.DataFrame([params])
1177 + df_params.to_excel(
1178 + writer, sheet_name="Summary", index=False, startrow=startrow
1179 + )
1180 + startrow += len(df_params) + 2
1181 + df_stats.to_excel(
1182 + writer, sheet_name="Summary", index=True, startrow=startrow
1183 + )
1184 + startrow += len(df_stats) + 3
1185 +
1186 + raw_data.to_excel(writer, sheet_name="Raw data", index=False, startrow=0)
1187 +
1188 + logger.info(
1189 + f"{Color.GREEN}Client metrics exported to file: {filename}{Color.RESET}"
1190 + )
1191 +
1192 +
1193 +async def get_server_info(url: str) -> None:
1194 + logger.info(f"{Color.BLUE}Collecting information from server: {url}{Color.RESET}")
1195 + async with aiohttp.ClientSession() as session:
1196 + # Get server version (not mandatory, "version" endpoint may not exist)
1197 + url_version = f"{url}/version"
1198 + async with session.get(url_version) as response:
1199 + if HTTPStatus(response.status) == HTTPStatus.OK:
1200 + text = await response.text()
1201 + logger.info(f"{Color.BLUE}Server version: {text}{Color.RESET}")
1202 +
1203 + # Get available models
1204 + url_models = f"{url}/v1/models"
1205 + async with session.get(url_models) as response:
1206 + if HTTPStatus(response.status) == HTTPStatus.OK:
1207 + text = await response.text()
1208 + logger.info(f"{Color.BLUE}Models:{Color.RESET}")
1209 + models_data = json.loads(text)
1210 + models_list = models_data["data"]
1211 + for model in models_list:
1212 + model_id = model["id"]
1213 + max_model_len = model.get("max_model_len", "N/A")
1214 + logger.info(
1215 + f"{Color.BLUE}\t{model_id=}, {max_model_len=}{Color.RESET}"
1216 + )
1217 + else:
1218 + logger.info(f"{Color.RED}Failed to get models{Color.RESET}")
1219 +
1220 +
1221 +async def main() -> None:
1222 + parser = argparse.ArgumentParser(
1223 + prog="Benchmark serving with multi-turn conversations",
1224 + description="Benchmark online inference using REST API",
1225 + )
1226 + parser.add_argument("--version", action="version", version="%(prog)s 1.0")
1227 +
1228 + parser.add_argument(
1229 + "-i",
1230 + "--input-file",
1231 + type=str,
1232 + required=True,
1233 + help="Input JSON file with ShareGPT conversations or "
1234 + "configuration file for generation of synthetic conversations",
1235 + )
1236 + parser.add_argument(
1237 + "-o",
1238 + "--output-file",
1239 + type=str,
1240 + default=None,
1241 + help="Output JSON file containing conversations with updated assistant answers",
1242 + )
1243 +
1244 + parser.add_argument(
1245 + "--seed",
1246 + type=int,
1247 + default=0,
1248 + help="Seed for random number generators (default: 0)",
1249 + )
1250 + parser.add_argument(
1251 + "-m", "--model", type=str, required=True, help="Path of the LLM model"
1252 + )
1253 + parser.add_argument(
1254 + "-u",
1255 + "--url",
1256 + type=str,
1257 + default="http://localhost:8000",
1258 + help="Base URL for the LLM API server",
1259 + )
1260 +
1261 + parser.add_argument(
1262 + "-p",
1263 + "--num-clients",
1264 + type=int,
1265 + default=1,
1266 + help="Number of clients that will send requests in parallel",
1267 + )
1268 + parser.add_argument(
1269 + "-k",
1270 + "--max-active-conversations",
1271 + type=int,
1272 + default=None,
1273 + help="Max number of active conversations at a time (for all clients)",
1274 + )
1275 + parser.add_argument(
1276 + "-n",
1277 + "--max-num-requests",
1278 + type=int,
1279 + default=None,
1280 + help="Max number of requests to send (total for all clients)",
1281 + )
1282 +
1283 + parser.add_argument(
1284 + "--warmup-step",
1285 + default=False,
1286 + action="store_true",
1287 + help="Run a warmup step (using only the first turn of every conversation), "
1288 + "measurements will not be included in the final benchmark results",
1289 + )
1290 +
1291 + parser.add_argument(
1292 + "--max-turns",
1293 + type=int,
1294 + default=None,
1295 + help="Maximum number of turns/messages per conversation, "
1296 + "includes both user and assistant messages "
1297 + "(a positive number, e.g: 2, 4, 6, etc.), disabled by default",
1298 + )
1299 + parser.add_argument(
1300 + "--no-early-stop",
1301 + default=False,
1302 + action="store_true",
1303 + help="By default, the benchmark will stop if at least one client exits."
1304 + " Use this flag to disable this behavior",
1305 + )
1306 +
1307 + parser.add_argument(
1308 + "--limit-max-tokens",
1309 + type=int,
1310 + default=NUM_TOKENS_FROM_DATASET,
1311 + help="Set max_tokens for the output token count of each request "
1312 + "(must also set --limit-min-tokens). "
1313 + "Overrides output token count from the input dataset. "
1314 + "Use a negative value to disable this limit.",
1315 + )
1316 + parser.add_argument(
1317 + "--limit-min-tokens",
1318 + type=int,
1319 + default=NUM_TOKENS_FROM_DATASET,
1320 + help="Set min_tokens for the output token count of each request "
1321 + "(must also set --limit-max-tokens). "
1322 + "Overrides output token count from the input dataset. "
1323 + "Use a negative value to disable this limit.",
1324 + )
1325 +
1326 + parser.add_argument(
1327 + "--request-rate",
1328 + type=float,
1329 + default=0,
1330 + help="Expected request rate (Poisson process) per client in requests/sec."
1331 + "Set to 0 for no delay between requests.",
1332 + )
1333 + parser.add_argument(
1334 + "--conversation-sampling",
1335 + type=ConversationSampling,
1336 + choices=list(ConversationSampling),
1337 + default=ConversationSampling.ROUND_ROBIN,
1338 + help=(
1339 + "Strategy for selecting which conversation to use for the next request. "
1340 + "Options: 'round_robin' (cycle through conversations), "
1341 + "'random' (pick randomly)."
1342 + ),
1343 + )
1344 + parser.add_argument(
1345 + "--verify-output",
1346 + default=False,
1347 + action="store_true",
1348 + help="Verify the LLM output (compare to the answers in the input JSON file)",
1349 + )
1350 +
1351 + parser.add_argument(
1352 + "--no-stream",
1353 + default=False,
1354 + action="store_true",
1355 + help="Disable stream/streaming mode (set 'stream' to False in the API request)",
1356 + )
1357 +
1358 + parser.add_argument(
1359 + "-e",
1360 + "--excel-output",
1361 + default=False,
1362 + action="store_true",
1363 + help="Export summary to Excel file (optional)",
1364 + )
1365 + parser.add_argument(
1366 + "-v",
1367 + "--verbose",
1368 + default=False,
1369 + action="store_true",
1370 + help="Enable verbose output",
1371 + )
1372 + parser.add_argument(
1373 + "--print-content",
1374 + default=False,
1375 + action="store_true",
1376 + help="Print the user prompts and the server's answers",
1377 + )
1378 +
1379 + parser.add_argument(
1380 + "--warmup-percentages",
1381 + type=str,
1382 + default="0%",
1383 + help="Ignore the first X samples as warmup (X is a percentage)."
1384 + " A comma separated list of percentages can be used "
1385 + "(for example: --warmup-percentages=0%%,50%%)",
1386 + )
1387 +
1388 + args = parser.parse_args()
1389 +
1390 + logger.info(args)
1391 +
1392 + logger.info(f"{Color.GREEN}Input parameters:{Color.RESET}")
1393 + logger.info(f"url={args.url}")
1394 + logger.info(f"model={args.model}")
1395 + logger.info(f"num_clients={args.num_clients}")
1396 +
1397 + if args.verify_output:
1398 + logger.info(f"{Color.PURPLE}Verify is enabled{Color.RESET}")
1399 +
1400 + # Calculate the amount of samples to filter (as warmup samples/measurements).
1401 + try:
1402 + warmup_percentages: list[float] = [0.0]
1403 + if not args.warmup_step:
1404 + # Warmup percentage can be used only if the warmup step was used
1405 + warmup_strings: list[str] = args.warmup_percentages.split(",")
1406 + warmup_strings = [x.replace("%", "") for x in warmup_strings]
1407 + warmup_percentages = [float(x) / 100 for x in warmup_strings]
1408 +
1409 + # Check for valid range (0 to 1)
1410 + for p in warmup_percentages:
1411 + assert p >= 0.0 and p < 1.0
1412 +
1413 + # Sort from high to low warmup percentage
1414 + warmup_percentages.sort()
1415 +
1416 + logger.info(
1417 + f"Warmup percentages (percentage of samples): {warmup_percentages}"
1418 + )
1419 +
1420 + except Exception:
1421 + raise ValueError(
1422 + f"Invalid --warmup-percentage={args.warmup_percentage}"
1423 + ) from None
1424 +
1425 + random.seed(args.seed)
1426 + np.random.seed(args.seed)
1427 +
1428 + if not os.path.exists(args.model):
1429 + raise OSError(f"Path does not exist: {args.model}")
1430 + logger.info("Loading tokenizer")
1431 + tokenizer = AutoTokenizer.from_pretrained(args.model)
1432 +
1433 + await get_server_info(args.url)
1434 +
1435 + # Load the input file (either conversations of configuration file)
1436 + logger.info(f"Reading input file: {args.input_file}")
1437 + with open(args.input_file) as f:
1438 + input_data = json.load(f)
1439 +
1440 + gen_conv_args = None
1441 + if isinstance(input_data, list):
1442 + # The conversations are stored as a list of dicts
1443 + logger.info(f"Found {len(input_data)} items in the input file")
1444 +
1445 + # Convert the list to a ConversationsMap
1446 + conversations = conversations_list_to_dict(input_data)
1447 +
1448 + elif isinstance(input_data, dict):
1449 + # The input file is a configuration file
1450 + # (type is determined by the field 'filetype')
1451 + if "filetype" not in input_data:
1452 + raise Exception(
1453 + f"Input file {args.input_file} is invalid (missing 'filetype')"
1454 + )
1455 +
1456 + logger.info(f"Using input file with filetype: {input_data['filetype']}")
1457 +
1458 + gen_conv_args = parse_input_json_file(input_data)
1459 +
1460 + # Disable warning from "huggingface/tokenizers"
1461 + # (when using python multiprocessing and tokenizers)
1462 + os.environ["TOKENIZERS_PARALLELISM"] = "true"
1463 +
1464 + # Generate synthetic conversations
1465 + conversations = generate_conversations(gen_conv_args, tokenizer)
1466 +
1467 + else:
1468 + raise Exception(f"Input file {args.input_file} is invalid")
1469 +
1470 + if args.max_turns is not None:
1471 + if args.max_turns < 1:
1472 + raise ValueError("Max turns must be a positive number")
1473 + logger.info(
1474 + f"{Color.PURPLE}Max turns per conversation "
1475 + f"is limited to {args.max_turns}{Color.RESET}"
1476 + )
1477 +
1478 + # Create benchmark configurations
1479 + client_args, req_args = get_client_config(args, conversations)
1480 +
1481 + bench_args = BenchmarkArgs(
1482 + url=args.url, num_clients=args.num_clients, early_stop=not args.no_early_stop
1483 + )
1484 +
1485 + # Warm-up step
1486 + if args.warmup_step:
1487 + # Only send a single user prompt from every conversation.
1488 + # max_active_conversations must be 1,
1489 + # otherwise the clients may exit after sending a single request
1490 + # (because the task queue is empty).
1491 + warmup_client_args = client_args._replace(
1492 + skip_first_turn=False, max_turns=1, max_active_conversations=1
1493 + )
1494 +
1495 + # Early stop should be disabled,
1496 + # all clients should finish their work before exiting
1497 + warmup_bench_args = bench_args._replace(early_stop=False)
1498 +
1499 + logger.info(f"{Color.PURPLE}Warmup start{Color.RESET}")
1500 + conversations, _ = await main_mp(
1501 + warmup_client_args, req_args, warmup_bench_args, tokenizer, conversations
1502 + )
1503 + logger.info(f"{Color.PURPLE}Warmup done{Color.RESET}")
1504 +
1505 + # Run the benchmark
1506 + start_time = time.perf_counter_ns()
1507 + client_convs, client_metrics = await main_mp(
1508 + client_args, req_args, bench_args, tokenizer, conversations
1509 + )
1510 + total_runtime_ms = nanosec_to_millisec(time.perf_counter_ns() - start_time)
1511 +
1512 + # Calculate requests per second
1513 + total_runtime_sec = total_runtime_ms / 1000.0
1514 + rps = len(client_metrics) / total_runtime_sec
1515 + logger.info(
1516 + f"{Color.GREEN}All clients finished, total runtime: {total_runtime_sec:.3f} sec"
1517 + f" ({total_runtime_ms:.3f} ms), requests per second: {rps:.3f}{Color.RESET}"
1518 + )
1519 +
1520 + # Benchmark parameters
1521 + params = {
1522 + "model": args.model,
1523 + "num_clients": args.num_clients,
1524 + "num_conversations": len(conversations),
1525 + "active_conversations": args.max_active_conversations,
1526 + "seed": args.seed,
1527 + }
1528 +
1529 + if args.limit_min_tokens > 0:
1530 + params["min_tokens"] = args.limit_min_tokens
1531 +
1532 + if args.limit_max_tokens > 0:
1533 + params["max_tokens"] = args.limit_max_tokens
1534 +
1535 + # Process and print statistics (and save excel file with the statistics)
1536 + process_statistics(
1537 + client_metrics,
1538 + test_params=params,
1539 + warmup_percentages=warmup_percentages,
1540 + verbose=args.verbose,
1541 + gen_conv_args=gen_conv_args,
1542 + excel_output=args.excel_output,
1543 + )
1544 +
1545 + if args.output_file is not None:
1546 + # Write a JSON file with the updated conversations
1547 + # The "assistant" content will contain the answers from the tested LLM
1548 + output_data: ShareGptConversations = conversations_dict_to_list(client_convs)
1549 + logger.info(
1550 + f"{Color.GREEN}Writing conversations file: {args.output_file}{Color.RESET}"
1551 + )
1552 + with open(args.output_file, "w") as f:
1553 + json.dump(output_data, f, indent=4)
1554 +
1555 +
1556 +if __name__ == "__main__":
1557 + asyncio.run(main())
benchmarks/multi_turn/convert_sharegpt_to_openai.py new
+354
@@ -0,0 +1,354 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Download dataset from:
5 +https://huggingface.co/datasets/philschmid/sharegpt-raw/blob/main/sharegpt_20230401_clean_lang_split.json
6 +
7 +Convert to OpenAI API:
8 +export INPUT_FILE=sharegpt_20230401_clean_lang_split.json
9 +python convert_sharegpt_to_openai.py $INPUT_FILE sharegpt_conv_128.json --max-items=128
10 +"""
11 +
12 +import argparse
13 +import json
14 +import random
15 +from statistics import mean
16 +from typing import Any, Optional
17 +
18 +import pandas as pd # type: ignore
19 +import tqdm # type: ignore
20 +from transformers import AutoTokenizer # type: ignore
21 +
22 +
23 +def has_non_english_chars(text: str) -> bool:
24 + return not text.isascii()
25 +
26 +
27 +def content_is_valid(
28 + content: str, min_content_len: Optional[int], max_content_len: Optional[int]
29 +) -> bool:
30 + if min_content_len and len(content) < min_content_len:
31 + return False
32 +
33 + if max_content_len and len(content) > max_content_len:
34 + return False
35 +
36 + return has_non_english_chars(content)
37 +
38 +
39 +def print_stats(
40 + conversations: "list[dict[Any, Any]]", tokenizer: Optional[AutoTokenizer] = None
41 +) -> None:
42 + # Collect statistics
43 + stats = []
44 +
45 + print("\nCollecting statistics...")
46 + for item in tqdm.tqdm(conversations):
47 + # item has "id" and "messages"
48 + messages = item["messages"]
49 +
50 + user_turns = 0
51 + assistant_turns = 0
52 + user_words = 0
53 + assistant_words = 0
54 + conv_chars = 0
55 +
56 + user_tokens: list[int] = []
57 + assistant_tokens: list[int] = []
58 +
59 + for m in messages:
60 + content = m["content"]
61 + conv_chars += len(content)
62 + content_num_words = content.count(" ") + 1
63 +
64 + num_tokens = 0
65 + if tokenizer:
66 + num_tokens = len(tokenizer(m["content"]).input_ids)
67 +
68 + if m["role"] == "user":
69 + user_turns += 1
70 + user_words += content_num_words
71 + if tokenizer:
72 + user_tokens.append(num_tokens)
73 +
74 + elif m["role"] == "assistant":
75 + assistant_turns += 1
76 + assistant_words += content_num_words
77 + if tokenizer:
78 + assistant_tokens.append(num_tokens)
79 +
80 + # assert user_turns == assistant_turns, \
81 + # f"Invalid conversation ID {item['id']}"
82 +
83 + conv_words = user_words + assistant_words
84 + item_stats = {
85 + "user_turns": user_turns,
86 + "assistant_turns": assistant_turns,
87 + "user_words": user_words,
88 + "assistant_words": assistant_words,
89 + "conv_turns": len(messages),
90 + "conv_words": conv_words,
91 + "conv_characters": conv_chars,
92 + }
93 +
94 + if len(user_tokens) > 0:
95 + item_stats["user_tokens"] = int(mean(user_tokens))
96 +
97 + if len(assistant_tokens) > 0:
98 + item_stats["assistant_tokens"] = int(mean(assistant_tokens))
99 +
100 + stats.append(item_stats)
101 +
102 + print("\nStatistics:")
103 + percentiles = [0.25, 0.5, 0.75, 0.9, 0.99, 0.999, 0.9999]
104 + df = pd.DataFrame(stats)
105 + print(df.describe(percentiles=percentiles).transpose())
106 +
107 +
108 +def convert_sharegpt_to_openai(
109 + seed: int,
110 + input_file: str,
111 + output_file: str,
112 + max_items: Optional[int],
113 + min_content_len: Optional[int] = None,
114 + max_content_len: Optional[int] = None,
115 + min_turns: Optional[int] = None,
116 + max_turns: Optional[int] = None,
117 + model: Optional[str] = None,
118 +) -> None:
119 + if min_turns and max_turns:
120 + assert min_turns <= max_turns
121 +
122 + if min_content_len and max_content_len:
123 + # Verify that min is not larger than max if both were given
124 + assert min_content_len <= max_content_len
125 +
126 + print(
127 + f"Input parameters:\n{seed=}, {max_items=}, {min_content_len=},"
128 + f" {max_content_len=}, {min_turns=}, {max_turns=}\n"
129 + )
130 +
131 + random.seed(seed)
132 +
133 + tokenizer = None
134 + if model is not None:
135 + print(f"Loading tokenizer from: {model}")
136 + tokenizer = AutoTokenizer.from_pretrained(model)
137 +
138 + # Read the ShareGPT JSON file
139 + print(f"Reading file: {input_file}")
140 + with open(input_file, encoding="utf-8") as f:
141 + # Should be a list of dicts
142 + # Each dict should have "id" (string) and "conversations" (list of dicts)
143 + sharegpt_data = json.load(f)
144 +
145 + assert isinstance(sharegpt_data, list), "Input file should contain a list of dicts"
146 +
147 + print(f"Total items in input file: {len(sharegpt_data):,}")
148 +
149 + print(f"Shuffling dataset with seed {seed}")
150 + random.shuffle(sharegpt_data)
151 +
152 + # Map conversation ID to the all the messages
153 + conversation_parts: dict[str, list[Any]] = {}
154 +
155 + for item in tqdm.tqdm(sharegpt_data):
156 + assert "id" in item, "Missing key 'id'"
157 + assert "conversations" in item, "Missing key 'conversations'"
158 +
159 + # Conversation ID (e.g: "hiWPlMD") and part/session (0, 1, 2, etc.)
160 + conv_id, _ = item["id"].split("_")
161 + new_turns = item["conversations"]
162 +
163 + if conv_id not in conversation_parts:
164 + # Start new conversation
165 + conversation_parts[conv_id] = []
166 + elif len(conversation_parts[conv_id]) > 0 and len(new_turns) > 0:
167 + prev_turns = conversation_parts[conv_id][-1]
168 + if prev_turns[-1]["from"] == new_turns[0]["from"]:
169 + new_turns = new_turns[1:]
170 +
171 + if len(new_turns) > 0:
172 + # We assume that parts are in order in the ShareGPT dataset
173 + conversation_parts[conv_id].append(new_turns)
174 +
175 + dataset: list[dict[str, Any]] = []
176 + for conv_id, conv_parts in conversation_parts.items():
177 + new_item = {"id": conv_id}
178 +
179 + conversations: list[dict[str, str]] = []
180 +
181 + # Merge all parts
182 + for conv_part in conv_parts:
183 + conversations.extend(conv_part)
184 +
185 + if len(conversations) > 0:
186 + new_item["conversations"] = conversations
187 + dataset.append(new_item)
188 +
189 + print(f"Total unique conversations (IDs) in input file: {len(dataset):,}")
190 +
191 + # Final output data
192 + final_openai_dataset: list[dict] = []
193 +
194 + # Filter conversations from the ShareGPT dataset and convert to OpenAI format
195 + for item in tqdm.tqdm(dataset):
196 + messages: list[dict] = []
197 +
198 + assert "id" in item, "Missing key 'id'"
199 + assert "conversations" in item, "Missing key 'conversations'"
200 +
201 + conv_id = item["id"]
202 + conversations = item["conversations"]
203 +
204 + if min_turns is not None and len(conversations) < min_turns:
205 + # Skip short conversations
206 + continue
207 +
208 + # Convert each message in the conversation, up to max_turns if specified
209 + for i, turn in enumerate(conversations):
210 + assert "from" in turn and "value" in turn, (
211 + f"Invalid conversation ID {conv_id} - missing 'from' or 'value'"
212 + )
213 +
214 + role = None
215 + turn_from = turn["from"]
216 +
217 + if turn_from in {"human", "user"}:
218 + role = "user"
219 + elif turn_from in {"gpt", "bing", "chatgpt", "bard"}:
220 + role = "assistant"
221 + elif turn_from == "system":
222 + role = "system"
223 +
224 + assert role is not None, (
225 + f"Invalid conversation ID {conv_id} - 'from'='{turn_from}' is invalid"
226 + )
227 +
228 + if i == 0 and role != "user":
229 + # If the first message is from assistant (gpt), skip it.
230 + # this happens when the conversation is a follow-up
231 + # to a previous conversation (from the same user).
232 + continue
233 +
234 + if max_turns is not None and i >= max_turns:
235 + break
236 +
237 + # Convert message to OpenAI format (with "role" and "content")
238 + content = turn["value"]
239 + messages.append({"role": role, "content": content})
240 +
241 + # Add the converted conversation to the OpenAI format
242 + if len(messages) > 0:
243 + valid_messages = True
244 +
245 + # First turn should always be from the user
246 + user_turn = True
247 +
248 + for m in messages:
249 + # Make sure that turns alternate between user and assistant
250 + if (user_turn and m["role"] != "user") or (
251 + not user_turn and m["role"] != "assistant"
252 + ):
253 + valid_messages = False
254 + break
255 +
256 + user_turn = not user_turn
257 +
258 + content = m["content"]
259 + valid_messages = content_is_valid(
260 + content, min_content_len, max_content_len
261 + )
262 + if not valid_messages:
263 + break
264 +
265 + if valid_messages is True:
266 + final_openai_dataset.append({"id": conv_id, "messages": messages})
267 +
268 + assert len(final_openai_dataset) > 0, "Final number of conversations is zero"
269 +
270 + print_stats(final_openai_dataset)
271 +
272 + print_stats_again = False
273 + if max_items is not None and len(final_openai_dataset) > max_items:
274 + print(f"\n\nSampling {max_items} items from the dataset...")
275 + print_stats_again = True
276 + final_openai_dataset = random.sample(final_openai_dataset, max_items)
277 +
278 + if print_stats_again:
279 + # Print stats after the dataset changed
280 + print_stats(final_openai_dataset, tokenizer)
281 +
282 + # Write the converted data to a new JSON file
283 + final_size = len(final_openai_dataset)
284 + print(f"\nTotal conversations converted (after filtering): {final_size:,}")
285 + print(f"\nWriting file: {output_file}")
286 + with open(output_file, "w", encoding="utf-8") as f:
287 + json.dump(final_openai_dataset, f, ensure_ascii=False, indent=2)
288 +
289 +
290 +def main() -> None:
291 + parser = argparse.ArgumentParser(
292 + description="Convert ShareGPT dataset to OpenAI API format"
293 + )
294 + parser.add_argument("input_file", help="Path to the input ShareGPT JSON file")
295 + parser.add_argument(
296 + "output_file", help="Path to the output OpenAI format JSON file"
297 + )
298 + parser.add_argument(
299 + "--seed", type=int, default=0, help="Seed for random number generators"
300 + )
301 + parser.add_argument(
302 + "--max-items",
303 + type=int,
304 + default=None,
305 + help="Maximum number of items in the output file",
306 + )
307 + parser.add_argument(
308 + "--min-turns",
309 + type=int,
310 + default=None,
311 + help="Minimum number of turns per conversation",
312 + )
313 + parser.add_argument(
314 + "--max-turns",
315 + type=int,
316 + default=None,
317 + help="Maximum number of turns per conversation",
318 + )
319 + parser.add_argument(
320 + "--min-content-len",
321 + type=int,
322 + default=None,
323 + help="Min number of characters in the messages' content",
324 + )
325 + parser.add_argument(
326 + "--max-content-len",
327 + type=int,
328 + default=None,
329 + help="Max number of characters in the messages' content",
330 + )
331 + parser.add_argument(
332 + "--model",
333 + type=str,
334 + default=None,
335 + help="LLM model, only the tokenizer will be used",
336 + )
337 +
338 + args = parser.parse_args()
339 +
340 + convert_sharegpt_to_openai(
341 + args.seed,
342 + args.input_file,
343 + args.output_file,
344 + args.max_items,
345 + args.min_content_len,
346 + args.max_content_len,
347 + args.min_turns,
348 + args.max_turns,
349 + args.model,
350 + )
351 +
352 +
353 +if __name__ == "__main__":
354 + main()
benchmarks/multi_turn/generate_multi_turn.json new
+35
@@ -0,0 +1,35 @@
1 +{
2 + "filetype": "generate_conversations",
3 + "num_conversations": 24,
4 + "text_files": ["pg1184.txt"],
5 + "print_stats": false,
6 + "prompt_input": {
7 + "num_turns": {
8 + "distribution": "uniform",
9 + "min": 12,
10 + "max": 18
11 + },
12 + "common_prefix_num_tokens": {
13 + "distribution": "constant",
14 + "value": 500
15 + },
16 + "prefix_num_tokens": {
17 + "distribution": "lognormal",
18 + "mean": 6,
19 + "sigma": 4,
20 + "max": 1500
21 + },
22 + "num_tokens": {
23 + "distribution": "uniform",
24 + "min": 120,
25 + "max": 160
26 + }
27 + },
28 + "prompt_output": {
29 + "num_tokens": {
30 + "distribution": "uniform",
31 + "min": 80,
32 + "max": 120
33 + }
34 + }
35 +}
\ No newline at end of file
benchmarks/multi_turn/requirements.txt new
+5
@@ -0,0 +1,5 @@
1 +numpy>=1.24
2 +pandas>=2.0.0
3 +aiohttp>=3.10
4 +transformers>=4.46
5 +xlsxwriter>=3.2.1
\ No newline at end of file