run_batch.py 17.4 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import asyncio
5
import tempfile
6
from collections.abc import Awaitable
7
from http import HTTPStatus
8
from io import StringIO
9
from typing import Callable, Optional
10
11

import aiohttp
12
import torch
13
from prometheus_client import start_http_server
14
from tqdm import tqdm
15

16
from vllm.engine.arg_utils import AsyncEngineArgs, optional_type
17
from vllm.engine.async_llm_engine import AsyncLLMEngine
18
from vllm.entrypoints.logger import RequestLogger
19
# yapf: disable
20
21
from vllm.entrypoints.openai.protocol import (BatchRequestInput,
                                              BatchRequestOutput,
22
23
                                              BatchResponseData,
                                              ChatCompletionResponse,
24
                                              EmbeddingResponse, ErrorResponse,
25
                                              RerankResponse, ScoreResponse)
26
# yapf: enable
27
from vllm.entrypoints.openai.serving_chat import OpenAIServingChat
28
from vllm.entrypoints.openai.serving_embedding import OpenAIServingEmbedding
29
30
from vllm.entrypoints.openai.serving_models import (BaseModelPath,
                                                    OpenAIServingModels)
31
from vllm.entrypoints.openai.serving_score import ServingScores
32
from vllm.logger import init_logger
33
from vllm.usage.usage_lib import UsageContext
34
from vllm.utils import FlexibleArgumentParser, random_uuid
35
from vllm.version import __version__ as VLLM_VERSION
36

37
38
logger = init_logger(__name__)

39

40
def make_arg_parser(parser: FlexibleArgumentParser):
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
    parser.add_argument(
        "-i",
        "--input-file",
        required=True,
        type=str,
        help=
        "The path or url to a single input file. Currently supports local file "
        "paths, or the http protocol (http or https). If a URL is specified, "
        "the file should be available via HTTP GET.")
    parser.add_argument(
        "-o",
        "--output-file",
        required=True,
        type=str,
        help="The path or url to a single output file. Currently supports "
        "local file paths, or web (http or https) urls. If a URL is specified,"
        " the file should be available via HTTP PUT.")
58
59
60
61
62
63
64
    parser.add_argument(
        "--output-tmp-dir",
        type=str,
        default=None,
        help="The directory to store the output file before uploading it "
        "to the output URL.",
    )
65
    parser.add_argument("--response-role",
66
                        type=optional_type(str),
67
68
                        default="assistant",
                        help="The role name to return if "
69
                        "`request.add_generation_prompt=True`.")
70
71

    parser = AsyncEngineArgs.add_cli_args(parser)
72
73
74
75
76
77
78
79

    parser.add_argument('--max-log-len',
                        type=int,
                        default=None,
                        help='Max number of prompt characters or prompt '
                        'ID numbers being printed in log.'
                        '\n\nDefault: Unlimited')

80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
    parser.add_argument("--enable-metrics",
                        action="store_true",
                        help="Enable Prometheus metrics")
    parser.add_argument(
        "--url",
        type=str,
        default="0.0.0.0",
        help="URL to the Prometheus metrics server "
        "(only needed if enable-metrics is set).",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=8000,
        help="Port number for the Prometheus metrics server "
        "(only needed if enable-metrics is set).",
    )
97
98
99
100
101
    parser.add_argument(
        "--enable-prompt-tokens-details",
        action='store_true',
        default=False,
        help="If set to True, enable prompt_tokens_details in usage.")
102

103
104
105
106
107
108
109
    return parser


def parse_args():
    parser = FlexibleArgumentParser(
        description="vLLM OpenAI-Compatible batch runner.")
    return make_arg_parser(parser).parse_args()
110
111


112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
# explicitly use pure text format, with a newline at the end
# this makes it impossible to see the animation in the progress bar
# but will avoid messing up with ray or multiprocessing, which wraps
# each line of output with some prefix.
_BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n"  # noqa: E501


class BatchProgressTracker:

    def __init__(self):
        self._total = 0
        self._pbar: Optional[tqdm] = None

    def submitted(self):
        self._total += 1

    def completed(self):
        if self._pbar:
            self._pbar.update()

    def pbar(self) -> tqdm:
        enable_tqdm = not torch.distributed.is_initialized(
        ) or torch.distributed.get_rank() == 0
        self._pbar = tqdm(total=self._total,
                          unit="req",
                          desc="Running batch",
                          mininterval=5,
                          disable=not enable_tqdm,
                          bar_format=_BAR_FORMAT)
        return self._pbar


144
145
146
147
148
149
async def read_file(path_or_url: str) -> str:
    if path_or_url.startswith("http://") or path_or_url.startswith("https://"):
        async with aiohttp.ClientSession() as session, \
                   session.get(path_or_url) as resp:
            return await resp.text()
    else:
150
        with open(path_or_url, encoding="utf-8") as f:
151
152
153
            return f.read()


154
async def write_local_file(output_path: str,
155
                           batch_outputs: list[BatchRequestOutput]) -> None:
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
    """
    Write the responses to a local file.
    output_path: The path to write the responses to.
    batch_outputs: The list of batch outputs to write.
    """
    # We should make this async, but as long as run_batch runs as a
    # standalone program, blocking the event loop won't effect performance.
    with open(output_path, "w", encoding="utf-8") as f:
        for o in batch_outputs:
            print(o.model_dump_json(), file=f)


async def upload_data(output_url: str, data_or_file: str,
                      from_file: bool) -> None:
    """
    Upload a local file to a URL.
    output_url: The URL to upload the file to.
    data_or_file: Either the data to upload or the path to the file to upload.
    from_file: If True, data_or_file is the path to the file to upload.
    """
    # Timeout is a common issue when uploading large files.
    # We retry max_retries times before giving up.
    max_retries = 5
    # Number of seconds to wait before retrying.
    delay = 5

    for attempt in range(1, max_retries + 1):
        try:
            # We increase the timeout to 1000 seconds to allow
            # for large files (default is 300).
            async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(
                    total=1000)) as session:
                if from_file:
                    with open(data_or_file, "rb") as file:
                        async with session.put(output_url,
                                               data=file) as response:
                            if response.status != 200:
                                raise Exception(f"Failed to upload file.\n"
                                                f"Status: {response.status}\n"
                                                f"Response: {response.text()}")
                else:
                    async with session.put(output_url,
                                           data=data_or_file) as response:
                        if response.status != 200:
                            raise Exception(f"Failed to upload data.\n"
                                            f"Status: {response.status}\n"
                                            f"Response: {response.text()}")

        except Exception as e:
            if attempt < max_retries:
                logger.error(
207
208
209
210
                    "Failed to upload data (attempt %d). Error message: %s.\nRetrying in %d seconds...",  # noqa: E501
                    attempt,
                    e,
                    delay,
211
212
213
                )
                await asyncio.sleep(delay)
            else:
214
215
216
                raise Exception(
                    f"Failed to upload data (attempt {attempt}). Error message: {str(e)}."  # noqa: E501
                ) from e
217
218


219
async def write_file(path_or_url: str, batch_outputs: list[BatchRequestOutput],
220
221
222
223
224
225
226
227
                     output_tmp_dir: str) -> None:
    """
    Write batch_outputs to a file or upload to a URL.
    path_or_url: The path or URL to write batch_outputs to.
    batch_outputs: The list of batch outputs to write.
    output_tmp_dir: The directory to store the output file before uploading it
    to the output URL.
    """
228
    if path_or_url.startswith("http://") or path_or_url.startswith("https://"):
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
        if output_tmp_dir is None:
            logger.info("Writing outputs to memory buffer")
            output_buffer = StringIO()
            for o in batch_outputs:
                print(o.model_dump_json(), file=output_buffer)
            output_buffer.seek(0)
            logger.info("Uploading outputs to %s", path_or_url)
            await upload_data(
                path_or_url,
                output_buffer.read().strip().encode("utf-8"),
                from_file=False,
            )
        else:
            # Write responses to a temporary file and then upload it to the URL.
            with tempfile.NamedTemporaryFile(
                    mode="w",
                    encoding="utf-8",
                    dir=output_tmp_dir,
                    prefix="tmp_batch_output_",
                    suffix=".jsonl",
            ) as f:
                logger.info("Writing outputs to temporary local file %s",
                            f.name)
                await write_local_file(f.name, batch_outputs)
                logger.info("Uploading outputs to %s", path_or_url)
                await upload_data(path_or_url, f.name, from_file=True)
255
    else:
256
257
        logger.info("Writing outputs to local file %s", path_or_url)
        await write_local_file(path_or_url, batch_outputs)
258
259


260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
def make_error_request_output(request: BatchRequestInput,
                              error_msg: str) -> BatchRequestOutput:
    batch_output = BatchRequestOutput(
        id=f"vllm-{random_uuid()}",
        custom_id=request.custom_id,
        response=BatchResponseData(
            status_code=HTTPStatus.BAD_REQUEST,
            request_id=f"vllm-batch-{random_uuid()}",
        ),
        error=error_msg,
    )
    return batch_output


async def make_async_error_request_output(
        request: BatchRequestInput, error_msg: str) -> BatchRequestOutput:
    return make_error_request_output(request, error_msg)


279
async def run_request(serving_engine_func: Callable,
280
281
                      request: BatchRequestInput,
                      tracker: BatchProgressTracker) -> BatchRequestOutput:
282
    response = await serving_engine_func(request.body)
283

284
285
286
287
288
    if isinstance(
            response,
        (ChatCompletionResponse, EmbeddingResponse, ScoreResponse,
         RerankResponse),
    ):
289
290
291
        batch_output = BatchRequestOutput(
            id=f"vllm-{random_uuid()}",
            custom_id=request.custom_id,
292
            response=BatchResponseData(
293
                body=response, request_id=f"vllm-batch-{random_uuid()}"),
294
295
            error=None,
        )
296
    elif isinstance(response, ErrorResponse):
297
298
299
        batch_output = BatchRequestOutput(
            id=f"vllm-{random_uuid()}",
            custom_id=request.custom_id,
300
            response=BatchResponseData(
301
                status_code=response.code,
302
                request_id=f"vllm-batch-{random_uuid()}"),
303
            error=response,
304
        )
305
    else:
306
307
        batch_output = make_error_request_output(
            request, error_msg="Request must not be sent in stream mode")
308

309
    tracker.completed()
310
311
312
313
314
315
316
317
318
319
320
    return batch_output


async def main(args):
    if args.served_model_name is not None:
        served_model_names = args.served_model_name
    else:
        served_model_names = [args.model]

    engine_args = AsyncEngineArgs.from_cli_args(args)
    engine = AsyncLLMEngine.from_engine_args(
321
        engine_args, usage_context=UsageContext.OPENAI_BATCH_RUNNER)
322
323

    model_config = await engine.get_model_config()
324
325
326
327
    base_model_paths = [
        BaseModelPath(name=name, model_path=args.model)
        for name in served_model_names
    ]
328

329
330
331
332
333
    if args.disable_log_requests:
        request_logger = None
    else:
        request_logger = RequestLogger(max_log_len=args.max_log_len)

334
    # Create the openai serving objects.
335
    openai_serving_models = OpenAIServingModels(
336
        engine_client=engine,
337
338
339
340
341
        model_config=model_config,
        base_model_paths=base_model_paths,
        lora_modules=None,
        prompt_adapters=None,
    )
342
343
344
    openai_serving_chat = OpenAIServingChat(
        engine,
        model_config,
345
        openai_serving_models,
346
        args.response_role,
347
348
        request_logger=request_logger,
        chat_template=None,
349
        chat_template_content_format="auto",
350
        enable_prompt_tokens_details=args.enable_prompt_tokens_details,
351
    ) if model_config.runner_type == "generate" else None
352
353
354
    openai_serving_embedding = OpenAIServingEmbedding(
        engine,
        model_config,
355
        openai_serving_models,
356
        request_logger=request_logger,
357
        chat_template=None,
358
        chat_template_content_format="auto",
359
    ) if model_config.task == "embed" else None
360
    openai_serving_scores = (ServingScores(
361
362
363
364
365
        engine,
        model_config,
        openai_serving_models,
        request_logger=request_logger,
    ) if model_config.task == "score" else None)
366

367
368
369
    tracker = BatchProgressTracker()
    logger.info("Reading batch from %s...", args.input_file)

370
    # Submit all requests in the file to the engine "concurrently".
371
    response_futures: list[Awaitable[BatchRequestOutput]] = []
372
    for request_json in (await read_file(args.input_file)).strip().split("\n"):
373
374
375
376
377
        # Skip empty lines.
        request_json = request_json.strip()
        if not request_json:
            continue

378
        request = BatchRequestInput.model_validate_json(request_json)
379
380
381

        # Determine the type of request and run it.
        if request.url == "/v1/chat/completions":
382
383
            chat_handler_fn = openai_serving_chat.create_chat_completion if \
                openai_serving_chat is not None else None
384
            if chat_handler_fn is None:
385
386
387
388
389
390
391
392
                response_futures.append(
                    make_async_error_request_output(
                        request,
                        error_msg=
                        "The model does not support Chat Completions API",
                    ))
                continue

393
394
            response_futures.append(
                run_request(chat_handler_fn, request, tracker))
395
            tracker.submitted()
396
        elif request.url == "/v1/embeddings":
397
398
            embed_handler_fn = openai_serving_embedding.create_embedding if \
                openai_serving_embedding is not None else None
399
            if embed_handler_fn is None:
400
401
402
403
404
405
406
                response_futures.append(
                    make_async_error_request_output(
                        request,
                        error_msg="The model does not support Embeddings API",
                    ))
                continue

407
408
            response_futures.append(
                run_request(embed_handler_fn, request, tracker))
409
            tracker.submitted()
410
        elif request.url.endswith("/score"):
411
412
            score_handler_fn = openai_serving_scores.create_score if \
                openai_serving_scores is not None else None
413
            if score_handler_fn is None:
414
415
416
417
418
419
420
                response_futures.append(
                    make_async_error_request_output(
                        request,
                        error_msg="The model does not support Scores API",
                    ))
                continue

421
422
            response_futures.append(
                run_request(score_handler_fn, request, tracker))
423
            tracker.submitted()
424
425
426
427
428
429
430
431
432
433
434
435
436
437
        elif request.url.endswith("/rerank"):
            rerank_handler_fn = openai_serving_scores.do_rerank if \
                openai_serving_scores is not None else None
            if rerank_handler_fn is None:
                response_futures.append(
                    make_async_error_request_output(
                        request,
                        error_msg="The model does not support Rerank API",
                    ))
                continue

            response_futures.append(
                run_request(rerank_handler_fn, request, tracker))
            tracker.submitted()
438
        else:
439
440
441
            response_futures.append(
                make_async_error_request_output(
                    request,
442
443
444
445
446
                    error_msg=f"URL {request.url} was used. "
                    "Supported endpoints: /v1/chat/completions, /v1/embeddings,"
                    " /score, /rerank ."
                    "See vllm/entrypoints/openai/api_server.py for supported "
                    "score/rerank versions.",
447
                ))
448

449
450
    with tracker.pbar():
        responses = await asyncio.gather(*response_futures)
451

452
    await write_file(args.output_file, responses, args.output_tmp_dir)
453
454
455
456
457


if __name__ == "__main__":
    args = parse_args()

458
    logger.info("vLLM batch processing API version %s", VLLM_VERSION)
459
460
    logger.info("args: %s", args)

461
462
463
464
465
466
467
468
    # Start the Prometheus metrics server. LLMEngine uses the Prometheus client
    # to publish metrics at the /metrics endpoint.
    if args.enable_metrics:
        logger.info("Prometheus metrics enabled")
        start_http_server(port=args.port, addr=args.url)
    else:
        logger.info("Prometheus metrics disabled")

469
    asyncio.run(main(args))