utils.py 10.8 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
from argparse import Namespace
9
10
from logging import Logger
from string import Template
11
from typing import TYPE_CHECKING
12

13
import regex as re
14
from fastapi import Request
15
16
from fastapi.responses import JSONResponse, StreamingResponse
from starlette.background import BackgroundTask, BackgroundTasks
17

18
from vllm import envs
19
from vllm.engine.arg_utils import EngineArgs
20
from vllm.logger import current_formatter_type, init_logger
21
from vllm.platforms import current_platform
22
from vllm.utils.argparse_utils import FlexibleArgumentParser
23

24
if TYPE_CHECKING:
25
    from vllm.entrypoints.openai.engine.protocol import StreamOptions
26
    from vllm.entrypoints.openai.models.protocol import LoRAModulePath
27
28
29
30
31
else:
    StreamOptions = object
    LoRAModulePath = object


32
33
logger = init_logger(__name__)

34
VLLM_SUBCMD_PARSER_EPILOG = (
35
36
37
    "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
38
39
    "Documentation:            https://docs.vllm.ai\n"
)
40

41
42
43
44
45
46

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":
47
48
49
            # 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.
50
51
52
            if getattr(
                request.app.state, "enable_server_load_tracking", False
            ) and hasattr(request.app.state, "server_load_metrics"):
53
                request.app.state.server_load_metrics -= 1
54
55
56
57
58
59
            break


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

61
    This does _not_ use request.is_disconnected, which does not work with
62
    middleware. Instead this follows the pattern from
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
    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))

89
90
91
        done, pending = await asyncio.wait(
            [handler_task, cancellation_task], return_when=asyncio.FIRST_COMPLETED
        )
92
93
94
95
96
97
98
99
        for task in pending:
            task.cancel()

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

    return wrapper
100
101
102
103
104
105
106
107


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


def load_aware_call(func):
    @functools.wraps(func)
108
    async def wrapper(*args, **kwargs):
109
        raw_request = kwargs.get("raw_request", args[1] if len(args) > 1 else None)
110
111
112

        if raw_request is None:
            raise ValueError(
113
114
                "raw_request required when server load tracking is enabled"
            )
115

116
        if not getattr(raw_request.app.state, "enable_server_load_tracking", False):
117
            return await func(*args, **kwargs)
118

119
120
121
122
        # ensure the counter exists
        if not hasattr(raw_request.app.state, "server_load_metrics"):
            raw_request.app.state.server_load_metrics = 0

123
124
        raw_request.app.state.server_load_metrics += 1
        try:
125
            response = await func(*args, **kwargs)
126
127
128
129
130
131
        except Exception:
            raw_request.app.state.server_load_metrics -= 1
            raise

        if isinstance(response, (JSONResponse, StreamingResponse)):
            if response.background is None:
132
                response.background = BackgroundTask(decrement_server_load, raw_request)
133
            elif isinstance(response.background, BackgroundTasks):
134
                response.background.add_task(decrement_server_load, raw_request)
135
136
137
138
            elif isinstance(response.background, BackgroundTask):
                # Convert the single BackgroundTask to BackgroundTasks
                # and chain the decrement_server_load task to it
                tasks = BackgroundTasks()
139
140
141
142
143
                tasks.add_task(
                    response.background.func,
                    *response.background.args,
                    **response.background.kwargs,
                )
144
145
146
147
148
149
150
151
                tasks.add_task(decrement_server_load, raw_request)
                response.background = tasks
        else:
            raw_request.app.state.server_load_metrics -= 1

        return response

    return wrapper
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172


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"
173
174


175
176
def get_max_tokens(
    max_model_len: int,
177
    max_tokens: int | None,
178
    input_length: int,
179
180
    default_sampling_params: dict,
) -> int:
181
182
183
    default_max_tokens = max_model_len - input_length
    max_output_tokens = current_platform.get_max_output_tokens(input_length)

184
185
186
187
188
189
190
191
192
193
    return min(
        val
        for val in (
            default_max_tokens,
            max_tokens,
            max_output_tokens,
            default_sampling_params.get("max_tokens"),
        )
        if val is not None
    )
194
195


196
def log_non_default_args(args: Namespace | EngineArgs):
197
198
    from vllm.entrypoints.openai.cli_args import make_arg_parser

199
200
    non_default_args = {}

201
202
    # Handle Namespace
    if isinstance(args, Namespace):
203
204
205
206
207
208
209
        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):
210
        default_args = EngineArgs(model=args.model)  # Create default instance
211
212
213
214
215
        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
216
217
        if default_args.model != EngineArgs.model:
            non_default_args["model"] = default_args.model
218
    else:
219
220
221
        raise TypeError(
            "Unsupported argument type. Must be Namespace or EngineArgs instance."
        )
222
223

    logger.info("non-default args: %s", non_default_args)
224
225
226


def should_include_usage(
227
    stream_options: "StreamOptions | None", enable_force_include_usage: bool
228
229
230
231
232
233
234
235
236
) -> 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
237
238
239
240
241


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

244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
    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


260
261
262
def sanitize_message(message: str) -> str:
    # Avoid leaking memory address from object reprs
    return re.sub(r" at 0x[0-9a-f]+>", ">", message)
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287


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)