run_batch.py 30.2 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 sys
6
import tempfile
7
from argparse import Namespace
8
from collections.abc import Awaitable, Callable
9
from http import HTTPStatus
10
from io import BytesIO, StringIO
11
from typing import Any, TypeAlias
12
from urllib.parse import urlparse
13
14

import aiohttp
15
import pybase64 as base64
16
import torch
17
from fastapi import UploadFile
18
from prometheus_client import start_http_server
19
from pydantic import Field, TypeAdapter, field_validator, model_validator
20
from pydantic_core.core_schema import ValidationInfo
21
from starlette.datastructures import State
22
from tqdm import tqdm
23
from urllib3.util import parse_url
24

25
import vllm.envs as envs
26
27
from vllm.config import config
from vllm.engine.arg_utils import AsyncEngineArgs
28
from vllm.engine.protocol import EngineClient
29
from vllm.entrypoints.openai.api_server import init_app_state
30
from vllm.entrypoints.openai.chat_completion.protocol import (
31
    ChatCompletionRequest,
32
    ChatCompletionResponse,
33
)
34
from vllm.entrypoints.openai.cli_args import BaseFrontendArgs
35
from vllm.entrypoints.openai.engine.protocol import (
36
    ErrorInfo,
37
    ErrorResponse,
38
    OpenAIBaseModel,
39
)
40
41
42
43
44
45
46
47
48
49
50
51
from vllm.entrypoints.openai.speech_to_text.protocol import (
    TranscriptionRequest,
    TranscriptionResponse,
    TranscriptionResponseVerbose,
    TranslationRequest,
    TranslationResponse,
    TranslationResponseVerbose,
)
from vllm.entrypoints.pooling.embed.protocol import (
    EmbeddingRequest,
    EmbeddingResponse,
)
52
53
54
55
56
57
from vllm.entrypoints.pooling.score.protocol import (
    RerankRequest,
    RerankResponse,
    ScoreRequest,
    ScoreResponse,
)
58
from vllm.entrypoints.utils import create_error_response
59
from vllm.exceptions import VLLMValidationError
60
from vllm.logger import init_logger
61
from vllm.reasoning import ReasoningParserManager
62
63
from vllm.utils import random_uuid
from vllm.utils.argparse_utils import FlexibleArgumentParser
64
from vllm.version import __version__ as VLLM_VERSION
65

66
67
logger = init_logger(__name__)

68

69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
class BatchTranscriptionRequest(TranscriptionRequest):
    """
    Batch transcription request that uses file_url instead of file.

    This class extends TranscriptionRequest but replaces the file field
    with file_url to support batch processing from audio files written in JSON format.
    """

    file_url: str = Field(
        ...,
        description=(
            "Either a URL of the audio or a data URL with base64 encoded audio data. "
        ),
    )

    # Override file to be optional and unused for batch processing
    file: UploadFile | None = Field(default=None, exclude=True)  # type: ignore[assignment]

    @model_validator(mode="before")
    @classmethod
    def validate_no_file(cls, data: Any):
        """Ensure file field is not provided in batch requests."""
        if isinstance(data, dict) and "file" in data:
92
            raise VLLMValidationError(
93
                "The 'file' field is not supported in batch requests. "
94
95
                "Use 'file_url' instead.",
                parameter="file",
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
            )
        return data


class BatchTranslationRequest(TranslationRequest):
    """
    Batch translation request that uses file_url instead of file.

    This class extends TranslationRequest but replaces the file field
    with file_url to support batch processing from audio files written in JSON format.
    """

    file_url: str = Field(
        ...,
        description=(
            "Either a URL of the audio or a data URL with base64 encoded audio data. "
        ),
    )

    # Override file to be optional and unused for batch processing
    file: UploadFile | None = Field(default=None, exclude=True)  # type: ignore[assignment]

    @model_validator(mode="before")
    @classmethod
    def validate_no_file(cls, data: Any):
        """Ensure file field is not provided in batch requests."""
        if isinstance(data, dict) and "file" in data:
123
            raise VLLMValidationError(
124
                "The 'file' field is not supported in batch requests. "
125
126
                "Use 'file_url' instead.",
                parameter="file",
127
128
129
130
            )
        return data


131
BatchRequestInputBody: TypeAlias = (
132
133
134
135
136
137
    ChatCompletionRequest
    | EmbeddingRequest
    | ScoreRequest
    | RerankRequest
    | BatchTranscriptionRequest
    | BatchTranslationRequest
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
)


class BatchRequestInput(OpenAIBaseModel):
    """
    The per-line object of the batch input file.

    NOTE: Currently only the `/v1/chat/completions` endpoint is supported.
    """

    # A developer-provided per-request id that will be used to match outputs to
    # inputs. Must be unique for each request in a batch.
    custom_id: str

    # The HTTP method to be used for the request. Currently only POST is
    # supported.
    method: str

    # The OpenAI API relative URL to be used for the request. Currently
    # /v1/chat/completions is supported.
    url: str

    # The parameters of the request.
    body: BatchRequestInputBody

    @field_validator("body", mode="plain")
    @classmethod
    def check_type_for_url(cls, value: Any, info: ValidationInfo):
        # Use url to disambiguate models
        url: str = info.data["url"]
        if url == "/v1/chat/completions":
            return ChatCompletionRequest.model_validate(value)
        if url == "/v1/embeddings":
            return TypeAdapter(EmbeddingRequest).validate_python(value)
        if url.endswith("/score"):
173
            return TypeAdapter(ScoreRequest).validate_python(value)
174
175
        if url.endswith("/rerank"):
            return RerankRequest.model_validate(value)
176
177
178
179
        if url == "/v1/audio/transcriptions":
            return BatchTranscriptionRequest.model_validate(value)
        if url == "/v1/audio/translations":
            return BatchTranslationRequest.model_validate(value)
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
        return TypeAdapter(BatchRequestInputBody).validate_python(value)


class BatchResponseData(OpenAIBaseModel):
    # HTTP status code of the response.
    status_code: int = 200

    # An unique identifier for the API request.
    request_id: str

    # The body of the response.
    body: (
        ChatCompletionResponse
        | EmbeddingResponse
        | ScoreResponse
        | RerankResponse
196
197
198
199
        | TranscriptionResponse
        | TranscriptionResponseVerbose
        | TranslationResponse
        | TranslationResponseVerbose
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
        | None
    ) = None


class BatchRequestOutput(OpenAIBaseModel):
    """
    The per-line object of the batch output and error files
    """

    id: str

    # A developer-provided per-request id that will be used to match outputs to
    # inputs.
    custom_id: str

    response: BatchResponseData | None

    # For requests that failed with a non-HTTP error, this will contain more
    # information on the cause of the failure.
    error: Any | None


222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
@config
class BatchFrontendArgs(BaseFrontendArgs):
    """Arguments for the batch runner frontend."""

    input_file: str | None = None
    """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."""
    output_file: str | None = None
    """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."""
    output_tmp_dir: str | None = None
    """The directory to store the output file before uploading it
    to the output URL."""
    enable_metrics: bool = False
    """Enable Prometheus metrics"""
    host: str | None = None
    """Host name for the Prometheus metrics server
    (only needed if enable-metrics is set)."""
    port: int = 8000
    """Port number for the Prometheus metrics server
    (only needed if enable-metrics is set)."""
    url: str = "0.0.0.0"
    """[DEPRECATED] Host name for the Prometheus metrics server
    (only needed if enable-metrics is set). Use --host instead."""
248

249
250
251
252
253
254
    @classmethod
    def _customize_cli_kwargs(
        cls,
        frontend_kwargs: dict[str, Any],
    ) -> dict[str, Any]:
        frontend_kwargs = super()._customize_cli_kwargs(frontend_kwargs)
255

256
257
258
259
260
261
262
263
264
        frontend_kwargs["input_file"]["flags"] = ["-i"]
        frontend_kwargs["input_file"]["required"] = True
        frontend_kwargs["output_file"]["flags"] = ["-o"]
        frontend_kwargs["output_file"]["required"] = True

        frontend_kwargs["enable_metrics"]["action"] = "store_true"

        frontend_kwargs["url"]["deprecated"] = True
        return frontend_kwargs
265

266

267
268
269
def make_arg_parser(parser: FlexibleArgumentParser):
    parser = BatchFrontendArgs.add_cli_args(parser)
    parser = AsyncEngineArgs.add_cli_args(parser)
270
271
272
273
    return parser


def parse_args():
274
    parser = FlexibleArgumentParser(description="vLLM OpenAI-Compatible batch runner.")
275
276
277
278
279
280
281
282
283
284
285
286
287
288
    args = make_arg_parser(parser).parse_args()

    # Backward compatibility: If --url is set, use it for host
    url_explicit = any(arg == "--url" or arg.startswith("--url=") for arg in sys.argv)
    host_explicit = any(
        arg == "--host" or arg.startswith("--host=") for arg in sys.argv
    )
    if url_explicit and hasattr(args, "url") and not host_explicit:
        args.host = args.url
        logger.warning_once(
            "Using --url for metrics is deprecated. Please use --host instead."
        )

    return args
289
290


291
292
293
294
295
296
297
298
299
300
# 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
301
        self._pbar: tqdm | None = None
302
303
304
305
306
307
308
309
310

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

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

    def pbar(self) -> tqdm:
311
312
313
314
315
316
317
318
319
320
321
        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,
        )
322
323
324
        return self._pbar


325
326
async def read_file(path_or_url: str) -> str:
    if path_or_url.startswith("http://") or path_or_url.startswith("https://"):
327
        async with aiohttp.ClientSession() as session, session.get(path_or_url) as resp:
328
            resp.raise_for_status()
329
330
            return await resp.text()
    else:
331
        with open(path_or_url, encoding="utf-8") as f:
332
333
334
            return f.read()


335
336
337
async def write_local_file(
    output_path: str, batch_outputs: list[BatchRequestOutput]
) -> None:
338
339
340
341
342
343
    """
    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
344
    # standalone program, blocking the event loop won't affect performance.
345
346
347
348
349
    with open(output_path, "w", encoding="utf-8") as f:
        for o in batch_outputs:
            print(o.model_dump_json(), file=f)


350
async def upload_data(output_url: str, data_or_file: str, from_file: bool) -> None:
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
    """
    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).
367
368
369
            async with aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=1000)
            ) as session:
370
371
                if from_file:
                    with open(data_or_file, "rb") as file:
372
                        async with session.put(output_url, data=file) as response:
373
                            if response.status != 200:
374
375
376
377
378
                                raise Exception(
                                    f"Failed to upload file.\n"
                                    f"Status: {response.status}\n"
                                    f"Response: {response.text()}"
                                )
379
                else:
380
                    async with session.put(output_url, data=data_or_file) as response:
381
                        if response.status != 200:
382
383
384
385
386
                            raise Exception(
                                f"Failed to upload data.\n"
                                f"Status: {response.status}\n"
                                f"Response: {response.text()}"
                            )
387
388
389
390

        except Exception as e:
            if attempt < max_retries:
                logger.error(
391
392
393
394
                    "Failed to upload data (attempt %d). Error message: %s.\nRetrying in %d seconds...",  # noqa: E501
                    attempt,
                    e,
                    delay,
395
396
397
                )
                await asyncio.sleep(delay)
            else:
398
399
400
                raise Exception(
                    f"Failed to upload data (attempt {attempt}). Error message: {str(e)}."  # noqa: E501
                ) from e
401
402


403
404
405
async def write_file(
    path_or_url: str, batch_outputs: list[BatchRequestOutput], output_tmp_dir: str
) -> None:
406
407
408
409
410
411
412
    """
    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.
    """
413
    if path_or_url.startswith("http://") or path_or_url.startswith("https://"):
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
        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(
429
430
431
432
433
                mode="w",
                encoding="utf-8",
                dir=output_tmp_dir,
                prefix="tmp_batch_output_",
                suffix=".jsonl",
434
            ) as f:
435
                logger.info("Writing outputs to temporary local file %s", f.name)
436
437
438
                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)
439
    else:
440
441
        logger.info("Writing outputs to local file %s", path_or_url)
        await write_local_file(path_or_url, batch_outputs)
442
443


444
445
446
447
async def download_bytes_from_url(
    url: str,
    allowed_media_domains: list[str] | None = None,
) -> bytes:
448
449
450
451
452
    """
    Download data from a URL or decode from a data URL.

    Args:
        url: Either an HTTP/HTTPS URL or a data URL (data:...;base64,...)
453
454
455
        allowed_media_domains: If set, only HTTP/HTTPS URLs whose hostname
            is in this list are permitted. data: URLs are not subject to
            this restriction.
456
457
458
459
460
461

    Returns:
        Data as bytes
    """
    parsed = urlparse(url)

462
    # Handle data URLs (base64 encoded) - not subject to domain restrictions
463
464
465
466
467
468
469
470
471
472
473
474
475
    if parsed.scheme == "data":
        # Format: data:...;base64,<base64_data>
        if "," in url:
            header, data = url.split(",", 1)
            if "base64" in header:
                return base64.b64decode(data)
            else:
                raise ValueError(f"Unsupported data URL encoding: {header}")
        else:
            raise ValueError(f"Invalid data URL format: {url}")

    # Handle HTTP/HTTPS URLs
    elif parsed.scheme in ("http", "https"):
476
477
478
479
480
481
482
483
484
485
486
487
        if allowed_media_domains is not None:
            url_spec = parse_url(url)
            if url_spec.hostname not in allowed_media_domains:
                raise ValueError(
                    f"The URL must be from one of the allowed domains: "
                    f"{allowed_media_domains}. Input URL domain: "
                    f"{url_spec.hostname}"
                )
            # Use the normalized URL to prevent parsing discrepancies
            # between urllib3 and aiohttp (e.g. backslash-@ attacks).
            url = url_spec.url

488
489
        async with (
            aiohttp.ClientSession() as session,
490
491
492
493
            session.get(
                url,
                allow_redirects=envs.VLLM_MEDIA_URL_ALLOW_REDIRECTS,
            ) as resp,
494
495
496
497
498
499
500
501
502
503
504
505
506
507
        ):
            if resp.status != 200:
                raise Exception(
                    f"Failed to download data from URL: {url}. Status: {resp.status}"
                )
            return await resp.read()

    else:
        raise ValueError(
            f"Unsupported URL scheme: {parsed.scheme}. "
            "Supported schemes: http, https, data"
        )


508
509
510
def make_error_request_output(
    request: BatchRequestInput, error_msg: str
) -> BatchRequestOutput:
511
512
513
514
515
516
517
518
519
520
521
522
523
    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(
524
525
    request: BatchRequestInput, error_msg: str
) -> BatchRequestOutput:
526
527
528
    return make_error_request_output(request, error_msg)


529
530
531
532
533
async def run_request(
    serving_engine_func: Callable,
    request: BatchRequestInput,
    tracker: BatchProgressTracker,
) -> BatchRequestOutput:
534
535
536
537
    try:
        response = await serving_engine_func(request.body)
    except Exception as e:
        response = create_error_response(e)
538

539
    if isinstance(
540
        response,
541
542
543
544
545
546
547
548
549
550
        (
            ChatCompletionResponse,
            EmbeddingResponse,
            ScoreResponse,
            RerankResponse,
            TranscriptionResponse,
            TranscriptionResponseVerbose,
            TranslationResponse,
            TranslationResponseVerbose,
        ),
551
    ):
552
553
554
        batch_output = BatchRequestOutput(
            id=f"vllm-{random_uuid()}",
            custom_id=request.custom_id,
555
            response=BatchResponseData(
556
557
                body=response, request_id=f"vllm-batch-{random_uuid()}"
            ),
558
559
            error=None,
        )
560
    elif isinstance(response, ErrorResponse):
561
562
563
        batch_output = BatchRequestOutput(
            id=f"vllm-{random_uuid()}",
            custom_id=request.custom_id,
564
            response=BatchResponseData(
565
                status_code=response.error.code,
566
567
                request_id=f"vllm-batch-{random_uuid()}",
            ),
568
            error=response,
569
        )
570
    else:
571
        batch_output = make_error_request_output(
572
573
            request, error_msg="Request must not be sent in stream mode"
        )
574

575
    tracker.completed()
576
577
578
    return batch_output


579
580
581
WrapperFn: TypeAlias = Callable[[Callable], Callable]


582
583
584
585
586
def handle_endpoint_request(
    request: BatchRequestInput,
    tracker: BatchProgressTracker,
    url_matcher: Callable[[str], bool],
    handler_getter: Callable[[], Callable | None],
587
    wrapper_fn: WrapperFn | None = None,
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
) -> Awaitable[BatchRequestOutput] | None:
    """
    Generic handler for endpoint requests.

    Args:
        request: The batch request input
        tracker: Progress tracker for the batch
        url_matcher: Function that takes a URL and returns True if it matches
        handler_getter: Function that returns the handler function or None
        wrapper_fn: Optional function to wrap the handler (e.g., for transcriptions)

    Returns:
        Awaitable[BatchRequestOutput] if the request was handled,
        None if URL didn't match
    """
    if not url_matcher(request.url):
        return None
605

606
607
608
609
    handler_fn = handler_getter()
    if handler_fn is None:
        error_msg = f"Model does not support endpoint: {request.url}"
        return make_async_error_request_output(request, error_msg=error_msg)
610

611
612
613
614
615
616
617
618
    # Apply wrapper if provided (e.g., for transcriptions/translations)
    if wrapper_fn is not None:
        handler_fn = wrapper_fn(handler_fn)

    tracker.submitted()
    return run_request(handler_fn, request, tracker)


619
620
621
622
def make_transcription_wrapper(
    is_translation: bool,
    allowed_media_domains: list[str] | None = None,
) -> WrapperFn:
623
624
625
626
627
628
629
630
    """
    Factory function to create a wrapper for transcription/translation handlers.
    The wrapper converts BatchTranscriptionRequest or BatchTranslationRequest
    to TranscriptionRequest or TranslationRequest and calls the appropriate handler.

    Args:
        is_translation: If True, process as translation; otherwise process
            as transcription
631
632
        allowed_media_domains: If set, only URLs from these domains are
            permitted for HTTP/HTTPS fetches.
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649

    Returns:
        A function that takes a handler and returns a wrapped handler
    """

    def wrapper(handler_fn: Callable):
        async def transcription_wrapper(
            batch_request_body: (BatchTranscriptionRequest | BatchTranslationRequest),
        ) -> (
            TranscriptionResponse
            | TranscriptionResponseVerbose
            | TranslationResponse
            | TranslationResponseVerbose
            | ErrorResponse
        ):
            try:
                # Download data from URL
650
651
652
653
                audio_data = await download_bytes_from_url(
                    batch_request_body.file_url,
                    allowed_media_domains=allowed_media_domains,
                )
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692

                # Create a mock file from the downloaded audio data
                mock_file = UploadFile(
                    file=BytesIO(audio_data),
                    filename="audio.bin",
                )

                # Convert batch request to regular request
                # by copying all fields except file_url and setting file to mock_file
                request_dict = batch_request_body.model_dump(exclude={"file_url"})
                request_dict["file"] = mock_file

                if is_translation:
                    # Create TranslationRequest from BatchTranslationRequest
                    translation_request = TranslationRequest.model_validate(
                        request_dict
                    )
                    return await handler_fn(audio_data, translation_request)
                else:
                    # Create TranscriptionRequest from BatchTranscriptionRequest
                    transcription_request = TranscriptionRequest.model_validate(
                        request_dict
                    )
                    return await handler_fn(audio_data, transcription_request)
            except Exception as e:
                operation = "translation" if is_translation else "transcription"
                return ErrorResponse(
                    error=ErrorInfo(
                        message=f"Failed to process {operation}: {str(e)}",
                        type="BadRequestError",
                        code=HTTPStatus.BAD_REQUEST.value,
                    )
                )

        return transcription_wrapper

    return wrapper


693
async def build_endpoint_registry(
694
695
    engine_client: EngineClient,
    args: Namespace,
696
697
698
) -> dict[str, dict[str, Any]]:
    """
    Build the endpoint registry with all serving objects and handler configurations.
699

700
701
702
    Args:
        engine_client: The engine client
        args: Command line arguments
703

704
705
706
    Returns:
        Dictionary mapping endpoint keys to their configurations
    """
707
708
    supported_tasks = await engine_client.get_supported_tasks()
    logger.info("Supported tasks: %s", supported_tasks)
709

710
711
    # Create a state object to hold serving objects
    state = State()
712

713
714
715
716
    # Initialize all serving objects using init_app_state
    # This provides full functionality including chat template processing,
    # LoRA support, tool servers, etc.
    await init_app_state(engine_client, state, args, supported_tasks)
717

718
719
720
721
    # Get serving objects from state (defaulting to None if not set)
    openai_serving_chat = getattr(state, "openai_serving_chat", None)
    openai_serving_transcription = getattr(state, "openai_serving_transcription", None)
    openai_serving_translation = getattr(state, "openai_serving_translation", None)
722
723
    serving_embedding = getattr(state, "serving_embedding", None)
    serving_scores = getattr(state, "serving_scores", None)
724

725
726
    allowed_media_domains = getattr(args, "allowed_media_domains", None)

727
728
729
730
731
    # Registry of endpoint configurations
    endpoint_registry: dict[str, dict[str, Any]] = {
        "completions": {
            "url_matcher": lambda url: url == "/v1/chat/completions",
            "handler_getter": lambda: (
732
733
734
                openai_serving_chat.create_chat_completion
                if openai_serving_chat is not None
                else None
735
736
737
738
739
740
            ),
            "wrapper_fn": None,
        },
        "embeddings": {
            "url_matcher": lambda url: url == "/v1/embeddings",
            "handler_getter": lambda: (
741
                serving_embedding if serving_embedding is not None else None
742
743
744
745
746
747
            ),
            "wrapper_fn": None,
        },
        "score": {
            "url_matcher": lambda url: url.endswith("/score"),
            "handler_getter": lambda: (
748
                serving_scores.create_score if serving_scores is not None else None
749
750
751
752
753
754
            ),
            "wrapper_fn": None,
        },
        "rerank": {
            "url_matcher": lambda url: url.endswith("/rerank"),
            "handler_getter": lambda: (
755
                serving_scores.do_rerank if serving_scores is not None else None
756
757
758
759
760
761
762
763
764
765
            ),
            "wrapper_fn": None,
        },
        "transcriptions": {
            "url_matcher": lambda url: url == "/v1/audio/transcriptions",
            "handler_getter": lambda: (
                openai_serving_transcription.create_transcription
                if openai_serving_transcription is not None
                else None
            ),
766
767
768
769
            "wrapper_fn": make_transcription_wrapper(
                is_translation=False,
                allowed_media_domains=allowed_media_domains,
            ),
770
771
772
773
774
775
776
777
        },
        "translations": {
            "url_matcher": lambda url: url == "/v1/audio/translations",
            "handler_getter": lambda: (
                openai_serving_translation.create_translation
                if openai_serving_translation is not None
                else None
            ),
778
779
780
781
            "wrapper_fn": make_transcription_wrapper(
                is_translation=True,
                allowed_media_domains=allowed_media_domains,
            ),
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
        },
    }

    return endpoint_registry


def validate_run_batch_args(args):
    valid_reasoning_parsers = ReasoningParserManager.list_registered()
    if (
        reasoning_parser := args.structured_outputs_config.reasoning_parser
    ) and reasoning_parser not in valid_reasoning_parsers:
        raise KeyError(
            f"invalid reasoning parser: {reasoning_parser} "
            f"(chose from {{ {','.join(valid_reasoning_parsers)} }})"
        )


async def run_batch(
    engine_client: EngineClient,
    args: Namespace,
) -> None:
803
    endpoint_registry = await build_endpoint_registry(
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
        engine_client=engine_client,
        args=args,
    )

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

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

        request = BatchRequestInput.model_validate_json(request_json)

        # Use the last segment of the URL as the endpoint key.
        # More advanced URL matching is done in url_matcher of endpoint_registry.
        endpoint_key = request.url.split("/")[-1]

        result = None
        if endpoint_key in endpoint_registry:
            endpoint_config = endpoint_registry[endpoint_key]
            result = handle_endpoint_request(
                request,
                tracker,
                url_matcher=endpoint_config["url_matcher"],
                handler_getter=endpoint_config["handler_getter"],
                wrapper_fn=endpoint_config["wrapper_fn"],
834
            )
835

836
837
        if result is not None:
            response_futures.append(result)
838
        else:
839
840
841
            response_futures.append(
                make_async_error_request_output(
                    request,
842
843
                    error_msg=f"URL {request.url} was used. "
                    "Supported endpoints: /v1/chat/completions, /v1/embeddings,"
844
845
846
                    " /v1/audio/transcriptions, /v1/audio/translations, /score, "
                    " /rerank. See vllm/entrypoints/openai/api_server.py "
                    "for supported score/rerank versions.",
847
848
                )
            )
849

850
851
    with tracker.pbar():
        responses = await asyncio.gather(*response_futures)
852

853
    await write_file(args.output_file, responses, args.output_tmp_dir)
854
855


856
async def main(args: Namespace):
857
858
859
    from vllm.entrypoints.openai.api_server import build_async_engine_client
    from vllm.usage.usage_lib import UsageContext

860
861
    validate_run_batch_args(args)

862
    async with build_async_engine_client(
863
864
        args,
        usage_context=UsageContext.OPENAI_BATCH_RUNNER,
865
    ) as engine_client:
866
        await run_batch(engine_client, args)
867
868


869
870
871
if __name__ == "__main__":
    args = parse_args()

872
    logger.info("vLLM batch processing API version %s", VLLM_VERSION)
873
874
    logger.info("args: %s", args)

875
876
877
878
    # 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")
879
        start_http_server(port=args.port, addr=args.host)
880
881
882
    else:
        logger.info("Prometheus metrics disabled")

883
    asyncio.run(main(args))