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

On the server side, run one of the following commands:
Woosuk Kwon's avatar
Woosuk Kwon committed
4
5
    (vLLM backend)
    python -m vllm.entrypoints.api_server \
6
7
        --model <your_model> --swap-space 16 \
        --disable-log-requests
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

    (TGI backend)
    ./launch_hf_server.sh <your_model>

On the client side, run:
    python benchmarks/benchmark_serving.py \
        --backend <backend> \
        --tokenizer <your_model> --dataset <target_dataset> \
        --request-rate <request_rate>
"""
import argparse
import asyncio
import json
import random
import time
from typing import AsyncGenerator, List, Tuple

import aiohttp
import numpy as np
27
from tqdm.asyncio import tqdm
28
29
from transformers import PreTrainedTokenizerBase
from vllm.transformers_utils.tokenizer import get_tokenizer
30
31
32
33
34
35
36
37
38
39
40
41
42
43

# (prompt len, output len, latency)
REQUEST_LATENCY: List[Tuple[int, int, float]] = []


def sample_requests(
    dataset_path: str,
    num_requests: int,
    tokenizer: PreTrainedTokenizerBase,
) -> List[Tuple[str, int, int]]:
    # Load the dataset.
    with open(dataset_path) as f:
        dataset = json.load(f)
    # Filter out the conversations with less than 2 turns.
44
    dataset = [data for data in dataset if len(data["conversations"]) >= 2]
45
    # Only keep the first two turns of each conversation.
46
47
    dataset = [(data["conversations"][0]["value"],
                data["conversations"][1]["value"]) for data in dataset]
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

    # Tokenize the prompts and completions.
    prompts = [prompt for prompt, _ in dataset]
    prompt_token_ids = tokenizer(prompts).input_ids
    completions = [completion for _, completion in dataset]
    completion_token_ids = tokenizer(completions).input_ids
    tokenized_dataset = []
    for i in range(len(dataset)):
        output_len = len(completion_token_ids[i])
        tokenized_dataset.append((prompts[i], prompt_token_ids[i], output_len))

    # Filter out too long sequences.
    filtered_dataset: List[Tuple[str, int, int]] = []
    for prompt, prompt_token_ids, output_len in tokenized_dataset:
        prompt_len = len(prompt_token_ids)
        if prompt_len < 4 or output_len < 4:
            # Prune too short sequences.
            # This is because TGI causes errors when the input or output length
            # is too short.
            continue
        if prompt_len > 1024 or prompt_len + output_len > 2048:
            # Prune too long sequences.
            continue
        filtered_dataset.append((prompt, prompt_len, output_len))

    # Sample the requests.
    sampled_requests = random.sample(filtered_dataset, num_requests)
    return sampled_requests


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
        # 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)


async def send_request(
    backend: str,
97
    model: str,
98
99
100
101
102
103
    api_url: str,
    prompt: str,
    prompt_len: int,
    output_len: int,
    best_of: int,
    use_beam_search: bool,
104
    pbar: tqdm
105
) -> None:
106
    request_start_time = time.perf_counter()
107
108

    headers = {"User-Agent": "Benchmark Client"}
Woosuk Kwon's avatar
Woosuk Kwon committed
109
    if backend == "vllm":
110
111
112
113
114
115
116
117
118
119
120
        pload = {
            "prompt": prompt,
            "n": 1,
            "best_of": best_of,
            "use_beam_search": use_beam_search,
            "temperature": 0.0 if use_beam_search else 1.0,
            "top_p": 1.0,
            "max_tokens": output_len,
            "ignore_eos": True,
            "stream": False,
        }
121
122
        if model is not None:
            pload["model"] = model
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
    elif backend == "tgi":
        assert not use_beam_search
        params = {
            "best_of": best_of,
            "max_new_tokens": output_len,
            "do_sample": True,
        }
        pload = {
            "inputs": prompt,
            "parameters": params,
        }
    else:
        raise ValueError(f"Unknown backend: {backend}")

    timeout = aiohttp.ClientTimeout(total=3 * 3600)
    async with aiohttp.ClientSession(timeout=timeout) as session:
        while True:
140
141
            async with session.post(api_url, headers=headers,
                                    json=pload) as response:
142
143
144
145
146
147
148
149
150
151
                chunks = []
                async for chunk, _ in response.content.iter_chunks():
                    chunks.append(chunk)
            output = b"".join(chunks).decode("utf-8")
            output = json.loads(output)

            # Re-send the request if it failed.
            if "error" not in output:
                break

152
    request_end_time = time.perf_counter()
153
154
    request_latency = request_end_time - request_start_time
    REQUEST_LATENCY.append((prompt_len, output_len, request_latency))
155
156
    pbar.update(1)

157
158
159
160


async def benchmark(
    backend: str,
161
    model: str,
162
163
164
165
166
167
168
    api_url: str,
    input_requests: List[Tuple[str, int, int]],
    best_of: int,
    use_beam_search: bool,
    request_rate: float,
) -> None:
    tasks: List[asyncio.Task] = []
169
    pbar = tqdm(total=len(input_requests))
170
171
    async for request in get_request(input_requests, request_rate):
        prompt, prompt_len, output_len = request
172
173
        task = asyncio.create_task(
            send_request(backend, model, api_url, prompt, prompt_len,
174
                         output_len, best_of, use_beam_search, pbar))
175
        tasks.append(task)
176
177
    await asyncio.gather(*tasks)
    pbar.close()
178
179
180
181
182
183
184


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

185
    api_url = f"{args.protocol}://{args.host}:{args.port}{args.endpoint}"
186
187
    tokenizer = get_tokenizer(args.tokenizer,
                              trust_remote_code=args.trust_remote_code)
188
189
    input_requests = sample_requests(args.dataset, args.num_prompts, tokenizer)

190
    benchmark_start_time = time.perf_counter()
191
192
193
    asyncio.run(
        benchmark(args.backend, args.model, api_url, input_requests,
                  args.best_of, args.use_beam_search, args.request_rate))
194
    benchmark_end_time = time.perf_counter()
195
196
197
198
199
200
201
202
203
204
205
206
    benchmark_time = benchmark_end_time - benchmark_start_time
    print(f"Total time: {benchmark_time:.2f} s")
    print(f"Throughput: {args.num_prompts / benchmark_time:.2f} requests/s")

    # Compute the latency statistics.
    avg_latency = np.mean([latency for _, _, latency in REQUEST_LATENCY])
    print(f"Average latency: {avg_latency:.2f} s")
    avg_per_token_latency = np.mean([
        latency / (prompt_len + output_len)
        for prompt_len, output_len, latency in REQUEST_LATENCY
    ])
    print(f"Average latency per token: {avg_per_token_latency:.2f} s")
207
208
    avg_per_output_token_latency = np.mean(
        [latency / output_len for _, output_len, latency in REQUEST_LATENCY])
209
210
211
212
213
214
215
    print("Average latency per output token: "
          f"{avg_per_output_token_latency:.2f} s")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Benchmark the online serving throughput.")
216
217
218
    parser.add_argument("--backend",
                        type=str,
                        default="vllm",
Woosuk Kwon's avatar
Woosuk Kwon committed
219
                        choices=["vllm", "tgi"])
220
    parser.add_argument("--protocol", type=str, default="http", choices=["http", "https"])
221
    parser.add_argument("--host", type=str, default="localhost")
222
    parser.add_argument("--port", type=int, default=8000)
223
224
225
226
227
    parser.add_argument("--endpoint", type=str, default="/generate")
    parser.add_argument("--model", type=str, default=None)
    parser.add_argument("--dataset",
                        type=str,
                        required=True,
228
                        help="Path to the dataset.")
229
230
231
    parser.add_argument("--tokenizer",
                        type=str,
                        required=True,
232
                        help="Name or path of the tokenizer.")
233
234
235
    parser.add_argument("--best-of",
                        type=int,
                        default=1,
236
                        help="Generates `best_of` sequences per prompt and "
237
                        "returns the best one.")
238
    parser.add_argument("--use-beam-search", action="store_true")
239
240
241
    parser.add_argument("--num-prompts",
                        type=int,
                        default=1000,
242
                        help="Number of prompts to process.")
243
244
245
    parser.add_argument("--request-rate",
                        type=float,
                        default=float("inf"),
246
                        help="Number of requests per second. If this is inf, "
247
248
249
                        "then all the requests are sent at time 0. "
                        "Otherwise, we use Poisson process to synthesize "
                        "the request arrival times.")
250
    parser.add_argument("--seed", type=int, default=0)
251
252
    parser.add_argument('--trust-remote-code',
                        action='store_true',
253
                        help='trust remote code from huggingface')
254
255
    args = parser.parse_args()
    main(args)