backend_request_func.py 15.4 KB
Newer Older
1
2
import json
import os
3
import sys
4
import time
5
6
import traceback
from dataclasses import dataclass, field
7
from typing import List, Optional, Union
8
9

import aiohttp
10
import huggingface_hub.constants
11
from tqdm.asyncio import tqdm
12
13
from transformers import (AutoTokenizer, PreTrainedTokenizer,
                          PreTrainedTokenizerFast)
14
15
16
17
18
19
20
21
22
23
24
25
26

AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60)


@dataclass
class RequestFuncInput:
    prompt: str
    api_url: str
    prompt_len: int
    output_len: int
    model: str
    best_of: int = 1
    use_beam_search: bool = False
27
    logprobs: Optional[int] = None
28
29
30
31
32
33


@dataclass
class RequestFuncOutput:
    generated_text: str = ""
    success: bool = False
34
35
    latency: float = 0.0
    ttft: float = 0.0  # Time to first token
36
37
    itl: List[float] = field(
        default_factory=list)  # List of inter-token latencies
38
    prompt_len: int = 0
39
    error: str = ""
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64


async def async_request_tgi(
    request_func_input: RequestFuncInput,
    pbar: Optional[tqdm] = None,
) -> RequestFuncOutput:
    api_url = request_func_input.api_url
    assert api_url.endswith("generate_stream")

    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
        assert not request_func_input.use_beam_search
        params = {
            "best_of": request_func_input.best_of,
            "max_new_tokens": request_func_input.output_len,
            "do_sample": True,
            "temperature": 0.01,  # TGI does not accept 0.0 temperature.
            "top_p": 0.99,  # TGI does not accept 1.0 top_p.
        }
        payload = {
            "inputs": request_func_input.prompt,
            "parameters": params,
        }
        output = RequestFuncOutput()
        output.prompt_len = request_func_input.prompt_len

65
        ttft = 0.0
66
        st = time.perf_counter()
67
        most_recent_timestamp = st
68
69
70
        try:
            async with session.post(url=api_url, json=payload) as response:
                if response.status == 200:
71
72
73
                    async for chunk_bytes in response.content:
                        chunk_bytes = chunk_bytes.strip()
                        if not chunk_bytes:
74
                            continue
75
                        chunk_bytes = chunk_bytes.decode("utf-8")
76

77
78
79
80
81
                        #NOTE: Sometimes TGI returns a ping response without
                        # any data, we should skip it.
                        if chunk_bytes.startswith(":"):
                            continue
                        chunk = remove_prefix(chunk_bytes, "data:")
82

83
84
85
                        data = json.loads(chunk)
                        timestamp = time.perf_counter()
                        # First token
86
                        if ttft == 0.0:
87
88
89
                            ttft = time.perf_counter() - st
                            output.ttft = ttft

90
91
92
93
                        # Decoding phase
                        else:
                            output.itl.append(timestamp -
                                              most_recent_timestamp)
94

95
96
97
98
99
                        most_recent_timestamp = timestamp

                    output.latency = most_recent_timestamp - st
                    output.success = True
                    output.generated_text = data["generated_text"]
100
101
102
                else:
                    output.error = response.reason or ""
                    output.success = False
103
        except Exception:
104
            output.success = False
105
106
            exc_info = sys.exc_info()
            output.error = "".join(traceback.format_exception(*exc_info))
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133

        if pbar:
            pbar.update(1)
        return output


async def async_request_trt_llm(
    request_func_input: RequestFuncInput,
    pbar: Optional[tqdm] = None,
) -> RequestFuncOutput:
    api_url = request_func_input.api_url
    assert api_url.endswith("generate_stream")

    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
        assert not request_func_input.use_beam_search
        assert request_func_input.best_of == 1
        payload = {
            "accumulate_tokens": True,
            "text_input": request_func_input.prompt,
            "temperature": 0.0,
            "top_p": 1.0,
            "max_tokens": request_func_input.output_len,
            "stream": True,
        }
        output = RequestFuncOutput()
        output.prompt_len = request_func_input.prompt_len

134
        ttft = 0.0
135
        st = time.perf_counter()
136
        most_recent_timestamp = st
137
        try:
138
139
            async with session.post(url=api_url, json=payload) as response:
                if response.status == 200:
140
141
142
                    async for chunk_bytes in response.content:
                        chunk_bytes = chunk_bytes.strip()
                        if not chunk_bytes:
143
144
                            continue

145
146
                        chunk = remove_prefix(chunk_bytes.decode("utf-8"),
                                              "data:")
147
148

                        data = json.loads(chunk)
149
                        output.generated_text += data["text_output"]
150
151
                        timestamp = time.perf_counter()
                        # First token
152
                        if ttft == 0.0:
153
154
155
                            ttft = time.perf_counter() - st
                            output.ttft = ttft

156
157
158
159
160
161
162
163
                        # Decoding phase
                        else:
                            output.itl.append(timestamp -
                                              most_recent_timestamp)

                        most_recent_timestamp = timestamp

                    output.latency = most_recent_timestamp - st
164
165
166
                    output.success = True

                else:
167
                    output.error = response.reason or ""
168
                    output.success = False
169
        except Exception:
170
            output.success = False
171
172
            exc_info = sys.exc_info()
            output.error = "".join(traceback.format_exception(*exc_info))
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187

        if pbar:
            pbar.update(1)
        return output


async def async_request_deepspeed_mii(
    request_func_input: RequestFuncInput,
    pbar: Optional[tqdm] = None,
) -> RequestFuncOutput:
    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
        assert request_func_input.best_of == 1
        assert not request_func_input.use_beam_search

        payload = {
188
189
190
            "prompt": request_func_input.prompt,
            "max_tokens": request_func_input.output_len,
            "temperature": 0.01,  # deepspeed-mii does not accept 0.0 temp.
191
192
193
194
195
            "top_p": 1.0,
        }
        output = RequestFuncOutput()
        output.prompt_len = request_func_input.prompt_len

196
        # NOTE: DeepSpeed-MII doesn't support streaming as of Jan 28 2024,
197
        # will use 0 as placeholder.
198
        # See https://github.com/microsoft/DeepSpeed-MII/pull/311
199
200
201
202
203
        output.ttft = 0

        st = time.perf_counter()
        try:
            async with session.post(url=request_func_input.api_url,
204
205
206
                                    json=payload) as response:
                if response.status == 200:
                    parsed_resp = await response.json()
207
                    output.latency = time.perf_counter() - st
208
                    output.generated_text = parsed_resp["text"][0]
209
210
                    output.success = True
                else:
211
                    output.error = response.reason or ""
212
                    output.success = False
213
        except Exception:
214
            output.success = False
215
216
            exc_info = sys.exc_info()
            output.error = "".join(traceback.format_exception(*exc_info))
217
218
219
220
221
222
223
224
225
226
227

        if pbar:
            pbar.update(1)
        return output


async def async_request_openai_completions(
    request_func_input: RequestFuncInput,
    pbar: Optional[tqdm] = None,
) -> RequestFuncOutput:
    api_url = request_func_input.api_url
228
    assert api_url.endswith(
229
230
        ("completions", "profile")
    ), "OpenAI Completions API URL must end with 'completions' or 'profile'."
231
232
233
234
235
236
237
238
239

    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
        assert not request_func_input.use_beam_search
        payload = {
            "model": request_func_input.model,
            "prompt": request_func_input.prompt,
            "temperature": 0.0,
            "best_of": request_func_input.best_of,
            "max_tokens": request_func_input.output_len,
240
            "logprobs": request_func_input.logprobs,
241
242
243
244
245
246
247
248
249
250
            "stream": True,
        }
        headers = {
            "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"
        }

        output = RequestFuncOutput()
        output.prompt_len = request_func_input.prompt_len

        generated_text = ""
251
        ttft = 0.0
252
        st = time.perf_counter()
253
        most_recent_timestamp = st
254
255
256
257
        try:
            async with session.post(url=api_url, json=payload,
                                    headers=headers) as response:
                if response.status == 200:
258
259
260
                    async for chunk_bytes in response.content:
                        chunk_bytes = chunk_bytes.strip()
                        if not chunk_bytes:
261
262
                            continue

263
264
                        chunk = remove_prefix(chunk_bytes.decode("utf-8"),
                                              "data: ")
265
266
267
                        if chunk == "[DONE]":
                            latency = time.perf_counter() - st
                        else:
268
269
                            data = json.loads(chunk)

270
271
272
                            # NOTE: Some completion API might have a last
                            # usage summary response without a token so we
                            # want to check a token was generated
273
274
275
                            if data["choices"][0]["text"]:
                                timestamp = time.perf_counter()
                                # First token
276
                                if ttft == 0.0:
277
278
279
280
                                    ttft = time.perf_counter() - st
                                    output.ttft = ttft

                                # Decoding phase
281
282
283
                                else:
                                    output.itl.append(timestamp -
                                                      most_recent_timestamp)
284
285
286

                                most_recent_timestamp = timestamp
                                generated_text += data["choices"][0]["text"]
287
288
289
290

                    output.generated_text = generated_text
                    output.success = True
                    output.latency = latency
291
292
293
                else:
                    output.error = response.reason or ""
                    output.success = False
294
        except Exception:
295
            output.success = False
296
297
            exc_info = sys.exc_info()
            output.error = "".join(traceback.format_exception(*exc_info))
298
299
300
301
302
303

    if pbar:
        pbar.update(1)
    return output


304
305
306
307
308
309
async def async_request_openai_chat_completions(
    request_func_input: RequestFuncInput,
    pbar: Optional[tqdm] = None,
) -> RequestFuncOutput:
    api_url = request_func_input.api_url
    assert api_url.endswith(
310
311
        "chat/completions"
    ), "OpenAI Chat Completions API URL must end with 'chat/completions'."
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328

    async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session:
        assert not request_func_input.use_beam_search
        payload = {
            "model": request_func_input.model,
            "messages": [
                {
                    "role": "user",
                    "content": request_func_input.prompt,
                },
            ],
            "temperature": 0.0,
            "max_tokens": request_func_input.output_len,
            "stream": True,
        }
        headers = {
            "Content-Type": "application/json",
329
            "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}",
330
331
332
333
334
335
        }

        output = RequestFuncOutput()
        output.prompt_len = request_func_input.prompt_len

        generated_text = ""
336
        ttft = 0.0
337
        st = time.perf_counter()
338
        most_recent_timestamp = st
339
340
341
342
        try:
            async with session.post(url=api_url, json=payload,
                                    headers=headers) as response:
                if response.status == 200:
343
344
345
                    async for chunk_bytes in response.content:
                        chunk_bytes = chunk_bytes.strip()
                        if not chunk_bytes:
346
347
                            continue

348
349
                        chunk = remove_prefix(chunk_bytes.decode("utf-8"),
                                              "data: ")
350
351
352
                        if chunk == "[DONE]":
                            latency = time.perf_counter() - st
                        else:
353
354
355
                            timestamp = time.perf_counter()
                            data = json.loads(chunk)

356
357
                            delta = data["choices"][0]["delta"]
                            if delta.get("content", None):
358
                                # First token
359
                                if ttft == 0.0:
360
361
362
363
364
365
366
367
                                    ttft = time.perf_counter() - st
                                    output.ttft = ttft

                                # Decoding phase
                                else:
                                    output.itl.append(timestamp -
                                                      most_recent_timestamp)

368
                                generated_text += delta["content"]
369

370
371
                            most_recent_timestamp = timestamp

372
373
374
375
                    output.generated_text = generated_text
                    output.success = True
                    output.latency = latency
                else:
376
                    output.error = response.reason or ""
377
                    output.success = False
378
        except Exception:
379
            output.success = False
380
381
            exc_info = sys.exc_info()
            output.error = "".join(traceback.format_exception(*exc_info))
382
383
384
385
386
387

    if pbar:
        pbar.update(1)
    return output


388
389
# Since vllm must support Python 3.8, we can't use str.removeprefix(prefix)
# introduced in Python 3.9
390
391
392
393
394
395
def remove_prefix(text: str, prefix: str) -> str:
    if text.startswith(prefix):
        return text[len(prefix):]
    return text


396
def get_model(pretrained_model_name_or_path: str) -> str:
397
398
    if os.getenv('VLLM_USE_MODELSCOPE', 'False').lower() == 'true':
        from modelscope import snapshot_download
399
400
401
402
403
404
405
406

        model_path = snapshot_download(
            model_id=pretrained_model_name_or_path,
            local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
            ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"])

        return model_path
    return pretrained_model_name_or_path
407
408
409
410
411
412
413
414
415
416
417
418
419


def get_tokenizer(
    pretrained_model_name_or_path: str, trust_remote_code: bool
) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:
    if pretrained_model_name_or_path is not None and not os.path.exists(
            pretrained_model_name_or_path):
        pretrained_model_name_or_path = get_model(
            pretrained_model_name_or_path)
    return AutoTokenizer.from_pretrained(pretrained_model_name_or_path,
                                         trust_remote_code=trust_remote_code)


420
421
ASYNC_REQUEST_FUNCS = {
    "tgi": async_request_tgi,
422
423
    "vllm": async_request_openai_completions,
    "lmdeploy": async_request_openai_completions,
424
425
    "deepspeed-mii": async_request_deepspeed_mii,
    "openai": async_request_openai_completions,
426
    "openai-chat": async_request_openai_chat_completions,
427
    "tensorrt-llm": async_request_trt_llm,
428
    "scalellm": async_request_openai_completions,
429
}