utils.py 13 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 dataclasses
6
import functools
7
import os
8
9
import sys
import traceback
10
from argparse import Namespace
11
from http import HTTPStatus
12
13
from logging import Logger
from string import Template
14
from typing import TYPE_CHECKING
15

16
import regex as re
17
from fastapi import Request
18
19
from fastapi.responses import JSONResponse, StreamingResponse
from starlette.background import BackgroundTask, BackgroundTasks
20

21
from vllm import envs
22
from vllm.engine.arg_utils import EngineArgs
23
from vllm.exceptions import VLLMValidationError
24
from vllm.logger import current_formatter_type, init_logger
25
from vllm.platforms import current_platform
26
from vllm.utils.argparse_utils import FlexibleArgumentParser
27

28
if TYPE_CHECKING:
29
30
31
32
33
    from vllm.entrypoints.openai.engine.protocol import (
        ErrorInfo,
        ErrorResponse,
        StreamOptions,
    )
34
    from vllm.entrypoints.openai.models.protocol import LoRAModulePath
35
else:
36
37
    ErrorResponse = object
    ErrorInfo = object
38
    LoRAModulePath = object
39
    StreamOptions = object
40

41
42
logger = init_logger(__name__)

43
VLLM_SUBCMD_PARSER_EPILOG = (
44
45
46
    "For full list:            vllm {subcmd} --help=all\n"
    "For a section:            vllm {subcmd} --help=ModelConfig    (case-insensitive)\n"  # noqa: E501
    "For a flag:               vllm {subcmd} --help=max-model-len  (_ or - accepted)\n"  # noqa: E501
47
48
    "Documentation:            https://docs.vllm.ai\n"
)
49

50
51
52
53
54
55

async def listen_for_disconnect(request: Request) -> None:
    """Returns if a disconnect message is received"""
    while True:
        message = await request.receive()
        if message["type"] == "http.disconnect":
56
57
58
            # If load tracking is enabled *and* the counter exists, decrement
            # it. Combines the previous nested checks into a single condition
            # to satisfy the linter rule.
59
60
61
            if getattr(
                request.app.state, "enable_server_load_tracking", False
            ) and hasattr(request.app.state, "server_load_metrics"):
62
                request.app.state.server_load_metrics -= 1
63
64
65
66
67
68
            break


def with_cancellation(handler_func):
    """Decorator that allows a route handler to be cancelled by client
    disconnections.
69

70
    This does _not_ use request.is_disconnected, which does not work with
71
    middleware. Instead this follows the pattern from
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
97
    starlette.StreamingResponse, which simultaneously awaits on two tasks- one
    to wait for an http disconnect message, and the other to do the work that we
    want done. When the first task finishes, the other is cancelled.

    A core assumption of this method is that the body of the request has already
    been read. This is a safe assumption to make for fastapi handlers that have
    already parsed the body of the request into a pydantic model for us.
    This decorator is unsafe to use elsewhere, as it will consume and throw away
    all incoming messages for the request while it looks for a disconnect
    message.

    In the case where a `StreamingResponse` is returned by the handler, this
    wrapper will stop listening for disconnects and instead the response object
    will start listening for disconnects.
    """

    # Functools.wraps is required for this wrapper to appear to fastapi as a
    # normal route handler, with the correct request type hinting.
    @functools.wraps(handler_func)
    async def wrapper(*args, **kwargs):
        # The request is either the second positional arg or `raw_request`
        request = args[1] if len(args) > 1 else kwargs["raw_request"]

        handler_task = asyncio.create_task(handler_func(*args, **kwargs))
        cancellation_task = asyncio.create_task(listen_for_disconnect(request))

98
99
100
        done, pending = await asyncio.wait(
            [handler_task, cancellation_task], return_when=asyncio.FIRST_COMPLETED
        )
101
102
103
104
105
106
107
108
        for task in pending:
            task.cancel()

        if handler_task in done:
            return handler_task.result()
        return None

    return wrapper
109
110
111
112
113
114
115
116


def decrement_server_load(request: Request):
    request.app.state.server_load_metrics -= 1


def load_aware_call(func):
    @functools.wraps(func)
117
    async def wrapper(*args, **kwargs):
118
        raw_request = kwargs.get("raw_request", args[1] if len(args) > 1 else None)
119
120
121

        if raw_request is None:
            raise ValueError(
122
123
                "raw_request required when server load tracking is enabled"
            )
124

125
        if not getattr(raw_request.app.state, "enable_server_load_tracking", False):
126
            return await func(*args, **kwargs)
127

128
129
130
131
        # ensure the counter exists
        if not hasattr(raw_request.app.state, "server_load_metrics"):
            raw_request.app.state.server_load_metrics = 0

132
133
        raw_request.app.state.server_load_metrics += 1
        try:
134
            response = await func(*args, **kwargs)
135
136
137
138
139
140
        except Exception:
            raw_request.app.state.server_load_metrics -= 1
            raise

        if isinstance(response, (JSONResponse, StreamingResponse)):
            if response.background is None:
141
                response.background = BackgroundTask(decrement_server_load, raw_request)
142
            elif isinstance(response.background, BackgroundTasks):
143
                response.background.add_task(decrement_server_load, raw_request)
144
145
146
147
            elif isinstance(response.background, BackgroundTask):
                # Convert the single BackgroundTask to BackgroundTasks
                # and chain the decrement_server_load task to it
                tasks = BackgroundTasks()
148
149
150
151
152
                tasks.add_task(
                    response.background.func,
                    *response.background.args,
                    **response.background.kwargs,
                )
153
154
155
156
157
158
159
160
                tasks.add_task(decrement_server_load, raw_request)
                response.background = tasks
        else:
            raw_request.app.state.server_load_metrics -= 1

        return response

    return wrapper
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181


def cli_env_setup():
    # The safest multiprocessing method is `spawn`, as the default `fork` method
    # is not compatible with some accelerators. The default method will be
    # changing in future versions of Python, so we should use it explicitly when
    # possible.
    #
    # We only set it here in the CLI entrypoint, because changing to `spawn`
    # could break some existing code using vLLM as a library. `spawn` will cause
    # unexpected behavior if the code is not protected by
    # `if __name__ == "__main__":`.
    #
    # References:
    # - https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods
    # - https://pytorch.org/docs/stable/notes/multiprocessing.html#cuda-in-multiprocessing
    # - https://pytorch.org/docs/stable/multiprocessing.html#sharing-cuda-tensors
    # - https://docs.habana.ai/en/latest/PyTorch/Getting_Started_with_PyTorch_and_Gaudi/Getting_Started_with_PyTorch.html?highlight=multiprocessing#torch-multiprocessing-for-dataloaders
    if "VLLM_WORKER_MULTIPROC_METHOD" not in os.environ:
        logger.debug("Setting VLLM_WORKER_MULTIPROC_METHOD to 'spawn'")
        os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
182
183


184
185
def get_max_tokens(
    max_model_len: int,
186
    max_tokens: int | None,
187
    input_length: int,
188
    default_sampling_params: dict,
189
    override_max_tokens: int | None = None,
190
) -> int:
191
192
193
194
195
196
197
    model_max_tokens = max_model_len - input_length
    platform_max_tokens = current_platform.get_max_output_tokens(input_length)
    fallback_max_tokens = (
        max_tokens
        if max_tokens is not None
        else default_sampling_params.get("max_tokens")
    )
198

199
200
201
    return min(
        val
        for val in (
202
203
204
205
            model_max_tokens,
            fallback_max_tokens,
            override_max_tokens,
            platform_max_tokens,
206
207
208
        )
        if val is not None
    )
209
210


211
def log_non_default_args(args: Namespace | EngineArgs):
212
213
    from vllm.entrypoints.openai.cli_args import make_arg_parser

214
215
    non_default_args = {}

216
217
    # Handle Namespace
    if isinstance(args, Namespace):
218
219
220
221
222
223
224
        parser = make_arg_parser(FlexibleArgumentParser())
        for arg, default in vars(parser.parse_args([])).items():
            if default != getattr(args, arg):
                non_default_args[arg] = getattr(args, arg)

    # Handle EngineArgs instance
    elif isinstance(args, EngineArgs):
225
        default_args = EngineArgs(model=args.model)  # Create default instance
226
227
228
229
230
        for field in dataclasses.fields(args):
            current_val = getattr(args, field.name)
            default_val = getattr(default_args, field.name)
            if current_val != default_val:
                non_default_args[field.name] = current_val
231
232
        if default_args.model != EngineArgs.model:
            non_default_args["model"] = default_args.model
233
    else:
234
235
236
        raise TypeError(
            "Unsupported argument type. Must be Namespace or EngineArgs instance."
        )
237
238

    logger.info("non-default args: %s", non_default_args)
239
240
241


def should_include_usage(
242
    stream_options: "StreamOptions | None", enable_force_include_usage: bool
243
244
245
246
247
248
249
250
251
) -> tuple[bool, bool]:
    if stream_options:
        include_usage = stream_options.include_usage or enable_force_include_usage
        include_continuous_usage = include_usage and bool(
            stream_options.continuous_usage_stats
        )
    else:
        include_usage, include_continuous_usage = enable_force_include_usage, False
    return include_usage, include_continuous_usage
252
253
254
255
256


def process_lora_modules(
    args_lora_modules: list[LoRAModulePath], default_mm_loras: dict[str, str] | None
) -> list[LoRAModulePath]:
257
    from vllm.entrypoints.openai.models.serving import LoRAModulePath
258

259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
    lora_modules = args_lora_modules
    if default_mm_loras:
        default_mm_lora_paths = [
            LoRAModulePath(
                name=modality,
                path=lora_path,
            )
            for modality, lora_path in default_mm_loras.items()
        ]
        if args_lora_modules is None:
            lora_modules = default_mm_lora_paths
        else:
            lora_modules += default_mm_lora_paths
    return lora_modules


275
276
277
def sanitize_message(message: str) -> str:
    # Avoid leaking memory address from object reprs
    return re.sub(r" at 0x[0-9a-f]+>", ">", message)
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302


def log_version_and_model(lgr: Logger, version: str, model_name: str) -> None:
    if envs.VLLM_DISABLE_LOG_LOGO or (formatter := current_formatter_type(lgr)) is None:
        message = "vLLM server version %s, serving model %s"
    else:
        logo_template = Template(
            "\n       ${w}█     █     █▄   ▄█${r}\n"
            " ${o}▄▄${r} ${b}▄█${r} ${w}█     █     █ ▀▄▀ █${r}  version ${w}%s${r}\n"
            "  ${o}█${r}${b}▄█▀${r} ${w}█     █     █     █${r}  model   ${w}%s${r}\n"
            "   ${b}▀▀${r}  ${w}▀▀▀▀▀ ▀▀▀▀▀ ▀     ▀${r}\n"
        )
        colors = {
            "w": "\033[97;1m",  # white
            "o": "\033[93m",  # orange
            "b": "\033[94m",  # blue
            "r": "\033[0m",  # reset
        }
        if formatter != "color":
            # monochrome logo (no ansi escape codes)
            colors = dict.fromkeys(colors, "")

        message = logo_template.substitute(colors)

    lgr.info(message, version, model_name)
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358


def create_error_response(
    message: str | Exception,
    err_type: str = "BadRequestError",
    status_code: HTTPStatus = HTTPStatus.BAD_REQUEST,
    param: str | None = None,
    log_error_stack: bool = False,
) -> "ErrorResponse":
    exc: Exception | None = None

    from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse

    if isinstance(message, Exception):
        exc = message

        if isinstance(exc, VLLMValidationError):
            err_type = "BadRequestError"
            status_code = HTTPStatus.BAD_REQUEST
            param = exc.parameter
        elif isinstance(exc, (ValueError, TypeError, RuntimeError, OverflowError)):
            # Common validation errors from user input
            err_type = "BadRequestError"
            status_code = HTTPStatus.BAD_REQUEST
            param = None
        elif isinstance(exc, NotImplementedError):
            err_type = "NotImplementedError"
            status_code = HTTPStatus.NOT_IMPLEMENTED
            param = None
        elif exc.__class__.__name__ == "TemplateError":
            # jinja2.TemplateError (avoid importing jinja2)
            err_type = "BadRequestError"
            status_code = HTTPStatus.BAD_REQUEST
            param = None
        else:
            err_type = "InternalServerError"
            status_code = HTTPStatus.INTERNAL_SERVER_ERROR
            param = None

        message = str(exc)

    if log_error_stack:
        exc_type, _, _ = sys.exc_info()
        if exc_type is not None:
            traceback.print_exc()
        else:
            traceback.print_stack()

    return ErrorResponse(
        error=ErrorInfo(
            message=sanitize_message(message),
            type=err_type,
            code=status_code.value,
            param=param,
        )
    )