benchmark_serving.py 36.7 KB
Newer Older
1
r"""Benchmark online serving throughput.
2
3

On the server side, run one of the following commands:
4
    vLLM OpenAI API server
Ethan Xu's avatar
Ethan Xu committed
5
6
    vllm serve <your_model> \
        --swap-space 16 \
7
        --disable-log-requests
8
9

    (TGI backend)
Ronen Schaffer's avatar
Ronen Schaffer committed
10
    ./launch_tgi_server.sh <your_model> <max_batch_total_tokens>
11
12
13
14

On the client side, run:
    python benchmarks/benchmark_serving.py \
        --backend <backend> \
15
16
17
18
19
        --model <your_model> \
        --dataset-name sharegpt \
        --dataset-path <path to dataset> \
        --request-rate <request_rate> \ # By default <request_rate> is inf
        --num-prompts <num_prompts> # By default <num_prompts> is 1000
20

21
22
23
    when using tgi backend, add
        --endpoint /generate_stream
    to the end of the command above.
24
25
26
"""
import argparse
import asyncio
27
28
import base64
import io
29
import json
30
import os
31
32
import random
import time
33
import warnings
34
35
from dataclasses import dataclass
from datetime import datetime
36
from typing import Any, AsyncGenerator, Collection, Dict, List, Optional, Tuple
37
38

import numpy as np
39
40
from backend_request_func import (ASYNC_REQUEST_FUNCS, RequestFuncInput,
                                  RequestFuncOutput)
41
42
from datasets import load_dataset
from PIL.Image import Image
43
from tqdm.asyncio import tqdm
44
from transformers import PreTrainedTokenizerBase
45

46
47
48
49
try:
    from vllm.transformers_utils.tokenizer import get_tokenizer
except ImportError:
    from backend_request_func import get_tokenizer
50

51
52
53
54
55
try:
    from vllm.utils import FlexibleArgumentParser
except ImportError:
    from argparse import ArgumentParser as FlexibleArgumentParser

56
57
58
59
60
61
62
63

@dataclass
class BenchmarkMetrics:
    completed: int
    total_input: int
    total_output: int
    request_throughput: float
    output_throughput: float
64
    total_token_throughput: float
65
66
    mean_ttft_ms: float
    median_ttft_ms: float
67
    std_ttft_ms: float
68
    percentiles_ttft_ms: List[Tuple[float, float]]
69
70
    mean_tpot_ms: float
    median_tpot_ms: float
71
    std_tpot_ms: float
72
    percentiles_tpot_ms: List[Tuple[float, float]]
73
74
    mean_itl_ms: float
    median_itl_ms: float
75
    std_itl_ms: float
76
77
78
79
80
81
82
83
    percentiles_itl_ms: List[Tuple[float, float]]
    # E2EL stands for end-to-end latency per request.
    # It is the time taken on the client side from sending
    # a request to receiving a complete response.
    mean_e2el_ms: float
    median_e2el_ms: float
    std_e2el_ms: float
    percentiles_e2el_ms: List[Tuple[float, float]]
84
85


86
def sample_sharegpt_requests(
87
88
89
    dataset_path: str,
    num_requests: int,
    tokenizer: PreTrainedTokenizerBase,
90
    fixed_output_len: Optional[int] = None,
91
) -> List[Tuple[str, int, int, None]]:
92
    # Load the dataset.
93
    with open(dataset_path, encoding='utf-8') as f:
94
95
        dataset = json.load(f)
    # Filter out the conversations with less than 2 turns.
96
    dataset = [data for data in dataset if len(data["conversations"]) >= 2]
97
    # Only keep the first two turns of each conversation.
98
99
    dataset = [(data["conversations"][0]["value"],
                data["conversations"][1]["value"]) for data in dataset]
100

101
102
    # Shuffle the dataset.
    random.shuffle(dataset)
103

104
    # Filter out sequences that are too long or too short
105
    filtered_dataset: List[Tuple[str, int, int]] = []
106
107
108
109
110
111
112
113
114
    for i in range(len(dataset)):
        if len(filtered_dataset) == num_requests:
            break

        # Tokenize the prompts and completions.
        prompt = dataset[i][0]
        prompt_token_ids = tokenizer(prompt).input_ids
        completion = dataset[i][1]
        completion_token_ids = tokenizer(completion).input_ids
115
        prompt_len = len(prompt_token_ids)
116
117
        output_len = len(completion_token_ids
                         ) if fixed_output_len is None else fixed_output_len
118
        if prompt_len < 4 or (fixed_output_len is None and output_len < 4):
119
120
121
122
123
            # Prune too short sequences.
            continue
        if prompt_len > 1024 or prompt_len + output_len > 2048:
            # Prune too long sequences.
            continue
124
        filtered_dataset.append((prompt, prompt_len, output_len, None))
125

126
    return filtered_dataset
127
128


129
130
131
132
133
134
135
def sample_sonnet_requests(
    dataset_path: str,
    num_requests: int,
    input_len: int,
    output_len: int,
    prefix_len: int,
    tokenizer: PreTrainedTokenizerBase,
136
) -> List[Tuple[str, str, int, int, None]]:
137
138
139
    assert (
        input_len > prefix_len
    ), "'args.sonnet-input-len' must be greater than 'args.prefix-input-len'."
140
141

    # Load the dataset.
142
    with open(dataset_path, encoding='utf-8') as f:
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
        poem_lines = f.readlines()

    # Tokenize the poem lines.
    poem_token_ids = tokenizer(poem_lines).input_ids
    average_poem_len = sum(
        len(token_ids) for token_ids in poem_token_ids) / len(poem_token_ids)

    # Base prefix for all requests.
    base_prompt = "Pick as many lines as you can from these poem lines:\n"
    base_message = [{
        "role": "user",
        "content": base_prompt,
    }]
    base_prompt_formatted = tokenizer.apply_chat_template(
        base_message, add_generation_prompt=True, tokenize=False)
    base_prompt_offset = len(tokenizer(base_prompt_formatted).input_ids)

160
161
162
    assert (
        input_len > base_prompt_offset
    ), f"Please set 'args.sonnet-input-len' higher than {base_prompt_offset}."
163
164
165
166
167
168
169
    num_input_lines = round(
        (input_len - base_prompt_offset) / average_poem_len)

    # First approximately `prefix_len` number of tokens in the
    # prompt are fixed poem lines.
    assert (
        prefix_len > base_prompt_offset
170
    ), f"Please set 'args.sonnet-prefix-len' higher than {base_prompt_offset}."
171
172
173
174
175
176
177
178

    num_prefix_lines = round(
        (prefix_len - base_prompt_offset) / average_poem_len)
    prefix_lines = poem_lines[:num_prefix_lines]

    # Sample the rest of lines per request.
    sampled_requests: List[Tuple[str, int, int]] = []
    for _ in range(num_requests):
179
180
181
        num_lines_needed = num_input_lines - num_prefix_lines
        sampled_lines = "".join(prefix_lines +
                                random.choices(poem_lines, k=num_lines_needed))
182
183
184
185
186
187
188
189
190
191
192
193

        prompt = f"{base_prompt}{sampled_lines}"
        message = [
            {
                "role": "user",
                "content": prompt,
            },
        ]
        prompt_formatted = tokenizer.apply_chat_template(
            message, add_generation_prompt=True, tokenize=False)
        prompt_len = len(tokenizer(prompt_formatted).input_ids)
        sampled_requests.append(
194
195
196
197
198
199
200
201
202
203
204
            (prompt, prompt_formatted, prompt_len, output_len, None))

    return sampled_requests


def sample_hf_requests(
    dataset_path: str,
    dataset_subset: str,
    dataset_split: str,
    num_requests: int,
    tokenizer: PreTrainedTokenizerBase,
205
    random_seed: int,
206
207
208
209
210
211
212
213
    fixed_output_len: Optional[int] = None,
) -> List[Tuple[str, str, int, Optional[Dict[str, Collection[str]]]]]:
    dataset = load_dataset(dataset_path,
                           name=dataset_subset,
                           split=dataset_split,
                           streaming=True)
    assert "conversations" in dataset.features, (
        "HF Dataset must have 'conversations' column.")
214
215
    filter_func = lambda x: len(x["conversations"]) >= 2
    filtered_dataset = dataset.shuffle(seed=random_seed).filter(filter_func)
216
217
218
219
220
221
222
223
224
225
226
227
228
229
    sampled_requests: List[Tuple[str, int, int, Dict[str,
                                                     Collection[str]]]] = []
    for data in filtered_dataset:
        if len(sampled_requests) == num_requests:
            break

        # Tokenize the prompts and completions.
        prompt = data["conversations"][0]["value"]
        prompt_token_ids = tokenizer(prompt).input_ids
        completion = data["conversations"][1]["value"]
        completion_token_ids = tokenizer(completion).input_ids
        prompt_len = len(prompt_token_ids)
        output_len = len(completion_token_ids
                         ) if fixed_output_len is None else fixed_output_len
230
        if fixed_output_len is None and (prompt_len < 4 or output_len < 4):
231
232
            # Prune too short sequences.
            continue
233
234
        if fixed_output_len is None and \
            (prompt_len > 1024 or prompt_len + output_len > 2048):
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
            # Prune too long sequences.
            continue

        if "image" in data and isinstance(data["image"], Image):
            image: Image = data["image"]
            image = image.convert("RGB")
            image_data = io.BytesIO()
            image.save(image_data, format='JPEG')
            image_base64 = base64.b64encode(
                image_data.getvalue()).decode("utf-8")
            mm_content = {
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/jpeg;base64,{image_base64}"
                },
            }
        else:
            mm_content = None

        sampled_requests.append((prompt, prompt_len, output_len, mm_content))
255
256
257
258

    return sampled_requests


259
def sample_random_requests(
260
261
262
263
264
265
266
267
268
269
    prefix_len: int,
    input_len: int,
    output_len: int,
    num_prompts: int,
    range_ratio: float,
    tokenizer: PreTrainedTokenizerBase,
) -> List[Tuple[str, int, int]]:
    prefix_token_ids = np.random.randint(0,
                                         tokenizer.vocab_size,
                                         size=prefix_len).tolist()
270
271
272
273
274
275
276
277
278
279
280
281
282

    input_lens = np.random.randint(
        int(input_len * range_ratio),
        input_len + 1,
        size=num_prompts,
    )
    output_lens = np.random.randint(
        int(output_len * range_ratio),
        output_len + 1,
        size=num_prompts,
    )
    offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts)
    input_requests = []
283
    for i in range(num_prompts):
284
285
        prompt = tokenizer.decode(prefix_token_ids +
                                  [(offsets[i] + i + j) % tokenizer.vocab_size
286
                                   for j in range(input_lens[i])])
287

288
289
        input_requests.append((prompt, int(prefix_len + input_lens[i]),
                               int(output_lens[i]), None))
290
291
292
293

    return input_requests


294
295
296
297
298
299
300
301
302
303
304
async def get_request(
    input_requests: List[Tuple[str, int, int]],
    request_rate: float,
) -> AsyncGenerator[Tuple[str, int, int], None]:
    input_requests = iter(input_requests)
    for request in input_requests:
        yield request

        if request_rate == float("inf"):
            # If the request rate is infinity, then we don't need to wait.
            continue
305

306
307
308
309
310
311
        # Sample the request interval from the exponential distribution.
        interval = np.random.exponential(1.0 / request_rate)
        # The next request will be sent after the interval.
        await asyncio.sleep(interval)


312
313
314
315
316
def calculate_metrics(
    input_requests: List[Tuple[str, int, int]],
    outputs: List[RequestFuncOutput],
    dur_s: float,
    tokenizer: PreTrainedTokenizerBase,
317
318
    selected_percentile_metrics: List[str],
    selected_percentiles: List[float],
319
) -> Tuple[BenchmarkMetrics, List[int]]:
320
    actual_output_lens: List[int] = []
321
322
    total_input = 0
    completed = 0
323
324
325
    itls: List[float] = []
    tpots: List[float] = []
    ttfts: List[float] = []
326
    e2els: List[float] = []
327
328
    for i in range(len(outputs)):
        if outputs[i].success:
329
330
331
            # We use the tokenizer to count the number of output tokens for all
            # serving backends instead of looking at len(outputs[i].itl) since
            # multiple output tokens may be bundled together
332
            # Note : this may inflate the output token count slightly
333
334
335
            output_len = len(
                tokenizer(outputs[i].generated_text,
                          add_special_tokens=False).input_ids)
336
            actual_output_lens.append(output_len)
337
            total_input += input_requests[i][1]
338
339
340
            if output_len > 1:
                tpots.append(
                    (outputs[i].latency - outputs[i].ttft) / (output_len - 1))
341
            itls += outputs[i].itl
342
            ttfts.append(outputs[i].ttft)
343
            e2els.append(outputs[i].latency)
344
            completed += 1
345
346
        else:
            actual_output_lens.append(0)
347

348
349
350
351
352
    if completed == 0:
        warnings.warn(
            "All requests failed. This is likely due to a misconfiguration "
            "on the benchmark arguments.",
            stacklevel=2)
353
354
355
    metrics = BenchmarkMetrics(
        completed=completed,
        total_input=total_input,
356
        total_output=sum(actual_output_lens),
357
        request_throughput=completed / dur_s,
358
        output_throughput=sum(actual_output_lens) / dur_s,
359
        total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s,
360
361
        mean_ttft_ms=np.mean(ttfts or 0) *
        1000,  # ttfts is empty if streaming is not supported by backend
362
        std_ttft_ms=np.std(ttfts or 0) * 1000,
363
364
365
        median_ttft_ms=np.median(ttfts or 0) * 1000,
        percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000)
                             for p in selected_percentiles],
366
        mean_tpot_ms=np.mean(tpots or 0) * 1000,
367
        std_tpot_ms=np.std(tpots or 0) * 1000,
368
369
370
        median_tpot_ms=np.median(tpots or 0) * 1000,
        percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000)
                             for p in selected_percentiles],
371
        mean_itl_ms=np.mean(itls or 0) * 1000,
372
        std_itl_ms=np.std(itls or 0) * 1000,
373
374
375
376
377
378
379
380
        median_itl_ms=np.median(itls or 0) * 1000,
        percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000)
                            for p in selected_percentiles],
        mean_e2el_ms=np.median(e2els or 0) * 1000,
        std_e2el_ms=np.std(e2els or 0) * 1000,
        median_e2el_ms=np.mean(e2els or 0) * 1000,
        percentiles_e2el_ms=[(p, np.percentile(e2els or 0, p) * 1000)
                             for p in selected_percentiles],
381
    )
382

383
    return metrics, actual_output_lens
384

385
386
387
388

async def benchmark(
    backend: str,
    api_url: str,
389
    base_url: str,
390
391
    model_id: str,
    tokenizer: PreTrainedTokenizerBase,
392
    input_requests: List[Tuple[str, int, int]],
393
    logprobs: Optional[int],
394
395
    best_of: int,
    request_rate: float,
396
    disable_tqdm: bool,
397
    profile: bool,
398
399
    selected_percentile_metrics: List[str],
    selected_percentiles: List[str],
400
    ignore_eos: bool,
401
    max_concurrency: Optional[int],
402
403
):
    if backend in ASYNC_REQUEST_FUNCS:
404
        request_func = ASYNC_REQUEST_FUNCS[backend]
405
406
407
    else:
        raise ValueError(f"Unknown backend: {backend}")

408
    print("Starting initial single prompt test run...")
409
410
411
412
413
414
    test_prompt, test_prompt_len, test_output_len, test_mm_content = (
        input_requests[0])
    if backend != "openai-chat" and test_mm_content is not None:
        # multi-modal benchmark is only available on OpenAI Chat backend.
        raise ValueError(
            "Multi-modal content is only supported on 'openai-chat' backend.")
415
416
417
418
419
420
    test_input = RequestFuncInput(
        model=model_id,
        prompt=test_prompt,
        api_url=api_url,
        prompt_len=test_prompt_len,
        output_len=test_output_len,
421
        logprobs=logprobs,
422
        best_of=best_of,
423
        multi_modal_content=test_mm_content,
424
        ignore_eos=ignore_eos,
425
426
427
428
429
430
431
432
    )
    test_output = await request_func(request_func_input=test_input)
    if not test_output.success:
        raise ValueError(
            "Initial test run failed - Please make sure benchmark arguments "
            f"are correctly specified. Error: {test_output.error}")
    else:
        print("Initial test run completed. Starting main benchmark run...")
433
434
435

    if profile:
        print("Starting profiler...")
436
437
438
439
440
441
442
443
444
        profile_input = RequestFuncInput(model=model_id,
                                         prompt=test_prompt,
                                         api_url=base_url + "/start_profile",
                                         prompt_len=test_prompt_len,
                                         output_len=test_output_len,
                                         logprobs=logprobs,
                                         best_of=best_of,
                                         multi_modal_content=test_mm_content,
                                         ignore_eos=ignore_eos)
445
446
447
448
        profile_output = await request_func(request_func_input=profile_input)
        if profile_output.success:
            print("Profiler started")

449
    print(f"Traffic request rate: {request_rate}")
450
    print(f"Maximum request concurrency: {max_concurrency}")
451

452
453
    pbar = None if disable_tqdm else tqdm(total=len(input_requests))

454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
    # This can be used once the minimum Python version is 3.10 or higher,
    # and it will simplify the code in limited_request_func.
    #    semaphore = (asyncio.Semaphore(max_concurrency)
    #                 if max_concurrency else contextlib.nullcontext())
    semaphore = (asyncio.Semaphore(max_concurrency)
                 if max_concurrency else None)

    async def limited_request_func(request_func_input, pbar):
        if semaphore is None:
            return await request_func(request_func_input=request_func_input,
                                      pbar=pbar)
        async with semaphore:
            return await request_func(request_func_input=request_func_input,
                                      pbar=pbar)

469
    benchmark_start_time = time.perf_counter()
470
    tasks: List[asyncio.Task] = []
471
    async for request in get_request(input_requests, request_rate):
472
        prompt, prompt_len, output_len, mm_content = request
473
474
475
476
477
478
479
480
481
        request_func_input = RequestFuncInput(model=model_id,
                                              prompt=prompt,
                                              api_url=api_url,
                                              prompt_len=prompt_len,
                                              output_len=output_len,
                                              logprobs=logprobs,
                                              best_of=best_of,
                                              multi_modal_content=mm_content,
                                              ignore_eos=ignore_eos)
482
483
        tasks.append(
            asyncio.create_task(
484
485
                limited_request_func(request_func_input=request_func_input,
                                     pbar=pbar)))
486
    outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks)
487

488
489
490
491
492
493
494
495
    if profile:
        print("Stopping profiler...")
        profile_input = RequestFuncInput(
            model=model_id,
            prompt=test_prompt,
            api_url=base_url + "/stop_profile",
            prompt_len=test_prompt_len,
            output_len=test_output_len,
496
            logprobs=logprobs,
497
498
499
500
501
502
            best_of=best_of,
        )
        profile_output = await request_func(request_func_input=profile_input)
        if profile_output.success:
            print("Profiler stopped")

503
    if pbar is not None:
504
505
506
507
        pbar.close()

    benchmark_duration = time.perf_counter() - benchmark_start_time

508
    metrics, actual_output_lens = calculate_metrics(
509
510
511
512
        input_requests=input_requests,
        outputs=outputs,
        dur_s=benchmark_duration,
        tokenizer=tokenizer,
513
514
        selected_percentile_metrics=selected_percentile_metrics,
        selected_percentiles=selected_percentiles,
515
516
    )

517
518
519
520
521
522
523
524
525
526
527
    print("{s:{c}^{n}}".format(s=' Serving Benchmark Result ', n=50, c='='))
    print("{:<40} {:<10}".format("Successful requests:", metrics.completed))
    print("{:<40} {:<10.2f}".format("Benchmark duration (s):",
                                    benchmark_duration))
    print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input))
    print("{:<40} {:<10}".format("Total generated tokens:",
                                 metrics.total_output))
    print("{:<40} {:<10.2f}".format("Request throughput (req/s):",
                                    metrics.request_throughput))
    print("{:<40} {:<10.2f}".format("Output token throughput (tok/s):",
                                    metrics.output_throughput))
528
529
    print("{:<40} {:<10.2f}".format("Total Token throughput (tok/s):",
                                    metrics.total_token_throughput))
530
531
532
533
534
535

    result = {
        "duration": benchmark_duration,
        "completed": metrics.completed,
        "total_input_tokens": metrics.total_input,
        "total_output_tokens": metrics.total_output,
536
        "request_throughput": metrics.request_throughput,
537
        "output_throughput": metrics.output_throughput,
538
        "total_token_throughput": metrics.total_token_throughput,
539
540
541
542
543
544
        "input_lens": [output.prompt_len for output in outputs],
        "output_lens": actual_output_lens,
        "ttfts": [output.ttft for output in outputs],
        "itls": [output.itl for output in outputs],
        "generated_texts": [output.generated_text for output in outputs],
        "errors": [output.error for output in outputs],
545
    }
546
547
548
549
550
551
552
553
554

    def process_one_metric(
        # E.g., "ttft"
        metric_attribute_name: str,
        # E.g., "TTFT"
        metric_name: str,
        # E.g., "Time to First Token"
        metric_header: str,
    ):
555
        # This function prints and adds statistics of the specified
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
        # metric.
        if metric_attribute_name not in selected_percentile_metrics:
            return
        print("{s:{c}^{n}}".format(s=metric_header, n=50, c='-'))
        print("{:<40} {:<10.2f}".format(
            f"Mean {metric_name} (ms):",
            getattr(metrics, f"mean_{metric_attribute_name}_ms")))
        print("{:<40} {:<10.2f}".format(
            f"Median {metric_name} (ms):",
            getattr(metrics, f"median_{metric_attribute_name}_ms")))
        result[f"mean_{metric_attribute_name}_ms"] = getattr(
            metrics, f"mean_{metric_attribute_name}_ms")
        result[f"median_{metric_attribute_name}_ms"] = getattr(
            metrics, f"median_{metric_attribute_name}_ms")
        result[f"std_{metric_attribute_name}_ms"] = getattr(
            metrics, f"std_{metric_attribute_name}_ms")
        for p, value in getattr(metrics,
                                f"percentiles_{metric_attribute_name}_ms"):
            p_word = str(int(p)) if int(p) == p else str(p)
            print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name} (ms):",
                                            value))
            result[f"p{p_word}_{metric_attribute_name}_ms"] = value

    process_one_metric("ttft", "TTFT", "Time to First Token")
    process_one_metric("tpot", "TPOT",
                       "Time per Output Token (excl. 1st token)")
    process_one_metric("itl", "ITL", "Inter-token Latency")
    process_one_metric("e2el", "E2EL", "End-to-end Latency")

    print("=" * 50)

587
    return result
588
589
590
591
592
593
594


def main(args: argparse.Namespace):
    print(args)
    random.seed(args.seed)
    np.random.seed(args.seed)

595
596
597
598
599
600
    backend = args.backend
    model_id = args.model
    tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model

    if args.base_url is not None:
        api_url = f"{args.base_url}{args.endpoint}"
601
        base_url = f"{args.base_url}"
602
603
    else:
        api_url = f"http://{args.host}:{args.port}{args.endpoint}"
604
        base_url = f"http://{args.host}:{args.port}"
605
606

    tokenizer = get_tokenizer(tokenizer_id,
607
                              trust_remote_code=args.trust_remote_code)
608
609
610
611
612
613
614
615
616
617
618

    if args.dataset is not None:
        warnings.warn(
            "The '--dataset' argument will be deprecated in the next "
            "release. Please use '--dataset-name' and "
            "'--dataset-path' in the future runs.",
            stacklevel=2)
        input_requests = sample_sharegpt_requests(
            dataset_path=args.dataset,
            num_requests=args.num_prompts,
            tokenizer=tokenizer,
619
            fixed_output_len=args.sharegpt_output_len,
620
621
622
623
624
625
626
        )

    elif args.dataset_name == "sharegpt":
        input_requests = sample_sharegpt_requests(
            dataset_path=args.dataset_path,
            num_requests=args.num_prompts,
            tokenizer=tokenizer,
627
            fixed_output_len=args.sharegpt_output_len,
628
629
630
631
632
633
634
635
        )

    elif args.dataset_name == "sonnet":
        # Do not format the prompt, pass to message directly
        if args.backend == "openai-chat":
            input_requests = sample_sonnet_requests(
                dataset_path=args.dataset_path,
                num_requests=args.num_prompts,
636
637
638
                input_len=args.sonnet_input_len,
                output_len=args.sonnet_output_len,
                prefix_len=args.sonnet_prefix_len,
639
640
                tokenizer=tokenizer,
            )
641
            input_requests = [(prompt, prompt_len, output_len, None)
642
                              for prompt, prompt_formatted, prompt_len,
643
                              output_len, _ in input_requests]
644
645
646
647
648
649
650
        else:
            assert (
                tokenizer.chat_template or tokenizer.default_chat_template
            ), "Tokenizer/model must have chat template for sonnet dataset."
            input_requests = sample_sonnet_requests(
                dataset_path=args.dataset_path,
                num_requests=args.num_prompts,
651
652
653
                input_len=args.sonnet_input_len,
                output_len=args.sonnet_output_len,
                prefix_len=args.sonnet_prefix_len,
654
655
                tokenizer=tokenizer,
            )
656
            input_requests = [(prompt_formatted, prompt_len, output_len, None)
657
                              for prompt, prompt_formatted, prompt_len,
658
                              output_len, _ in input_requests]
659

660
661
662
663
664
665
666
    elif args.dataset_name == "hf":
        input_requests = sample_hf_requests(
            dataset_path=args.dataset_path,
            dataset_subset=args.hf_subset,
            dataset_split=args.hf_split,
            num_requests=args.num_prompts,
            tokenizer=tokenizer,
667
            random_seed=args.seed,
668
669
670
            fixed_output_len=args.hf_output_len,
        )

671
672
    elif args.dataset_name == "random":
        input_requests = sample_random_requests(
673
            prefix_len=args.random_prefix_len,
674
675
            input_len=args.random_input_len,
            output_len=args.random_output_len,
676
            num_prompts=args.num_prompts,
677
            range_ratio=args.random_range_ratio,
678
679
680
            tokenizer=tokenizer,
        )

681
682
    else:
        raise ValueError(f"Unknown dataset: {args.dataset_name}")
683

684
685
686
687
    benchmark_result = asyncio.run(
        benchmark(
            backend=backend,
            api_url=api_url,
688
            base_url=base_url,
689
690
691
            model_id=model_id,
            tokenizer=tokenizer,
            input_requests=input_requests,
692
            logprobs=args.logprobs,
693
694
695
            best_of=args.best_of,
            request_rate=args.request_rate,
            disable_tqdm=args.disable_tqdm,
696
            profile=args.profile,
697
698
699
700
            selected_percentile_metrics=args.percentile_metrics.split(","),
            selected_percentiles=[
                float(p) for p in args.metric_percentiles.split(",")
            ],
701
            ignore_eos=args.ignore_eos,
702
            max_concurrency=args.max_concurrency,
703
704
705
706
        ))

    # Save config and results to json
    if args.save_result:
707
        result_json: Dict[str, Any] = {}
708
709
710
711
712
713
714
715
716
717

        # Setup
        current_dt = datetime.now().strftime("%Y%m%d-%H%M%S")
        result_json["date"] = current_dt
        result_json["backend"] = backend
        result_json["model_id"] = model_id
        result_json["tokenizer_id"] = tokenizer_id
        result_json["best_of"] = args.best_of
        result_json["num_prompts"] = args.num_prompts

718
719
720
721
722
723
724
725
726
727
728
        # Metadata
        if args.metadata:
            for item in args.metadata:
                if "=" in item:
                    kvstring = item.split("=")
                    result_json[kvstring[0].strip()] = kvstring[1].strip()
                else:
                    raise ValueError(
                        "Invalid metadata format. Please use KEY=VALUE format."
                    )

729
730
731
        # Traffic
        result_json["request_rate"] = (
            args.request_rate if args.request_rate < float("inf") else "inf")
732
        result_json["max_concurrency"] = args.max_concurrency
733
734
735
736
737
738

        # Merge with benchmark result
        result_json = {**result_json, **benchmark_result}

        # Save to file
        base_model_id = model_id.split("/")[-1]
739
740
741
        max_concurrency_str = (f"-concurrency{args.max_concurrency}"
                               if args.max_concurrency is not None else "")
        file_name = f"{backend}-{args.request_rate}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json"  #noqa
742
743
        if args.result_filename:
            file_name = args.result_filename
744
745
        if args.result_dir:
            file_name = os.path.join(args.result_dir, file_name)
746
        with open(file_name, "w", encoding='utf-8') as outfile:
747
            json.dump(result_json, outfile)
748
749
750


if __name__ == "__main__":
751
    parser = FlexibleArgumentParser(
752
        description="Benchmark the online serving throughput.")
753
754
755
756
757
758
759
760
761
762
763
764
    parser.add_argument(
        "--backend",
        type=str,
        default="vllm",
        choices=list(ASYNC_REQUEST_FUNCS.keys()),
    )
    parser.add_argument(
        "--base-url",
        type=str,
        default=None,
        help="Server or API base url if not using http host and port.",
    )
765
    parser.add_argument("--host", type=str, default="localhost")
766
    parser.add_argument("--port", type=int, default=8000)
767
768
769
    parser.add_argument(
        "--endpoint",
        type=str,
770
        default="/v1/completions",
771
772
        help="API endpoint.",
    )
773
774
775
776
777
778
779
780
781
782
783
    parser.add_argument(
        "--dataset",
        type=str,
        default=None,
        help="Path to the ShareGPT dataset, will be deprecated in the "
        "next release.",
    )
    parser.add_argument(
        "--dataset-name",
        type=str,
        default="sharegpt",
784
        choices=["sharegpt", "sonnet", "random", "hf"],
785
786
787
        help="Name of the dataset to benchmark on.",
    )
    parser.add_argument("--dataset-path",
788
                        type=str,
789
                        default=None,
790
791
                        help="Path to the sharegpt/sonnet dataset. "
                        "Or the huggingface dataset ID if using HF dataset.")
792
793
794
795
796
797
798
799
800
801
802
803
804
    parser.add_argument(
        "--max-concurrency",
        type=int,
        default=None,
        help="Maximum number of concurrent requests. This can be used "
        "to help simulate an environment where a higher level component "
        "is enforcing a maximum number of concurrent requests. While the "
        "--request-rate argument controls the rate at which requests are "
        "initiated, this argument will control how many are actually allowed "
        "to execute at a time. This means that when used in combination, the "
        "actual request rate may be lower than specified with --request-rate, "
        "if the server is not processing requests fast enough to keep up.")

805
806
807
808
809
810
811
812
813
814
    parser.add_argument(
        "--model",
        type=str,
        required=True,
        help="Name of the model.",
    )
    parser.add_argument(
        "--tokenizer",
        type=str,
        help=
815
        "Name or path of the tokenizer, if not using the default tokenizer.",  # noqa: E501
816
817
818
819
820
821
822
823
    )
    parser.add_argument(
        "--best-of",
        type=int,
        default=1,
        help="Generates `best_of` sequences per prompt and "
        "returns the best one.",
    )
824
    parser.add_argument("--use-beam-search", action="store_true")
825
826
827
828
829
830
    parser.add_argument(
        "--num-prompts",
        type=int,
        default=1000,
        help="Number of prompts to process.",
    )
831
832
833
834
835
836
837
838
839
840
    parser.add_argument(
        "--logprobs",
        type=int,
        default=None,
        help=("Number of logprobs-per-token to compute & return as part of "
              "the request. If unspecified, then either (1) if beam search "
              "is disabled, no logprobs are computed & a single dummy "
              "logprob is returned for each token; or (2) if beam search "
              "is enabled 1 logprob per token is computed"),
    )
841
842
843
844
845
846
847
848
849
    parser.add_argument(
        "--request-rate",
        type=float,
        default=float("inf"),
        help="Number of requests per second. If this is inf, "
        "then all the requests are sent at time 0. "
        "Otherwise, we use Poisson process to synthesize "
        "the request arrival times.",
    )
850
    parser.add_argument("--seed", type=int, default=0)
851
852
853
854
855
856
857
858
    parser.add_argument(
        "--trust-remote-code",
        action="store_true",
        help="Trust remote code from huggingface",
    )
    parser.add_argument(
        "--disable-tqdm",
        action="store_true",
859
        help="Specify to disable tqdm progress bar.",
860
861
    )
    parser.add_argument(
862
863
864
865
866
867
        "--profile",
        action="store_true",
        help="Use Torch Profiler. The endpoint must be launched with "
        "VLLM_TORCH_PROFILER_DIR to enable profiler.",
    )
    parser.add_argument(
868
869
870
871
        "--save-result",
        action="store_true",
        help="Specify to save benchmark results to a json file",
    )
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
    parser.add_argument(
        "--metadata",
        metavar="KEY=VALUE",
        nargs="*",
        help="Key-value pairs (e.g, --metadata version=0.3.3 tp=1) "
        "for metadata of this run to be saved in the result JSON file "
        "for record keeping purposes.",
    )
    parser.add_argument(
        "--result-dir",
        type=str,
        default=None,
        help="Specify directory to save benchmark json results."
        "If not specified, results are saved in the current directory.",
    )
887
888
889
890
891
892
893
894
895
    parser.add_argument(
        "--result-filename",
        type=str,
        default=None,
        help="Specify the filename to save benchmark json results."
        "If not specified, results will be saved in "
        "{backend}-{args.request_rate}qps-{base_model_id}-{current_dt}.json"
        " format.",
    )
896
897
898
899
900
    parser.add_argument(
        "--ignore-eos",
        action="store_true",
        help="Set ignore_eos flag when sending the benchmark request."
        "Warning: ignore_eos is not supported in deepspeed_mii and tgi.")
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
    parser.add_argument(
        "--percentile-metrics",
        type=str,
        default="ttft,tpot,itl",
        help="Comma-seperated list of selected metrics to report percentils. "
        "This argument specifies the metrics to report percentiles. "
        "Allowed metric names are \"ttft\", \"tpot\", \"itl\", \"e2el\". "
        "Default value is \"ttft,tpot,itl\".")
    parser.add_argument(
        "--metric-percentiles",
        type=str,
        default="99",
        help="Comma-seperated list of percentiles for selected metrics. "
        "To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". "
        "Default value is \"99\". "
        "Use \"--percentile-metrics\" to select metrics.",
    )
918

919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
    # group for dataset specific arguments
    sonnet_group = parser.add_argument_group("sonnet dataset options")
    sonnet_group.add_argument(
        "--sonnet-input-len",
        type=int,
        default=550,
        help=
        "Number of input tokens per request, used only for sonnet dataset.",
    )
    sonnet_group.add_argument(
        "--sonnet-output-len",
        type=int,
        default=150,
        help=
        "Number of output tokens per request, used only for sonnet dataset.",
    )
    sonnet_group.add_argument(
        "--sonnet-prefix-len",
        type=int,
        default=200,
        help=
        "Number of prefix tokens per request, used only for sonnet dataset.",
    )

    sharegpt_group = parser.add_argument_group("sharegpt dataset options")
    sharegpt_group.add_argument(
        "--sharegpt-output-len",
        type=int,
        default=None,
        help="Output length for each request. Overrides the output length "
        "from the ShareGPT dataset.")

    random_group = parser.add_argument_group("random dataset options")
    random_group.add_argument(
        "--random-input-len",
        type=int,
        default=1024,
        help=
        "Number of input tokens per request, used only for random sampling.",
    )
    random_group.add_argument(
        "--random-output-len",
        type=int,
        default=128,
        help=
        "Number of output tokens per request, used only for random sampling.",
    )
    random_group.add_argument(
        "--random-range-ratio",
        type=float,
        default=1.0,
        help="Range of sampled ratio of input/output length, "
        "used only for random sampling.",
    )
    random_group.add_argument(
        "--random-prefix-len",
        type=int,
        default=0,
        help="Number of fixed prefix tokens before random "
        " context. The length range of context in a random "
        " request is [random-prefix-len, "
        " random-prefix-len + random-prefix-len * random-range-ratio).")

    hf_group = parser.add_argument_group("hf dataset options")
    hf_group.add_argument("--hf-subset",
                          type=str,
                          default=None,
                          help="Subset of the HF dataset.")
    hf_group.add_argument("--hf-split",
                          type=str,
                          default=None,
                          help="Split of the HF dataset.")
    hf_group.add_argument(
        "--hf-output-len",
        type=int,
        default=None,
        help="Output length for each request. Overrides the output lengths "
        "from the sampled HF dataset.",
    )

999
    args = parser.parse_args()
1000
    main(args)