utils.py 9.75 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
from typing import Any
10
11

from fastapi import Request
12
13
from fastapi.responses import JSONResponse, StreamingResponse
from starlette.background import BackgroundTask, BackgroundTasks
14

15
16
from vllm.engine.arg_utils import EngineArgs
from vllm.entrypoints.openai.cli_args import make_arg_parser
17
18
19
20
21
from vllm.entrypoints.openai.protocol import (
    ChatCompletionRequest,
    CompletionRequest,
    StreamOptions,
)
22
from vllm.logger import init_logger
23
from vllm.platforms import current_platform
24
from vllm.utils import FlexibleArgumentParser
25
26
27

logger = init_logger(__name__)

28
VLLM_SUBCMD_PARSER_EPILOG = (
29
30
31
    "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
32
33
    "Documentation:            https://docs.vllm.ai\n"
)
34

35
36
37
38
39
40

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":
41
42
43
            # 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.
44
45
46
            if getattr(
                request.app.state, "enable_server_load_tracking", False
            ) and hasattr(request.app.state, "server_load_metrics"):
47
                request.app.state.server_load_metrics -= 1
48
49
50
51
52
53
            break


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

55
    This does _not_ use request.is_disconnected, which does not work with
56
    middleware. Instead this follows the pattern from
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
    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))

83
84
85
        done, pending = await asyncio.wait(
            [handler_task, cancellation_task], return_when=asyncio.FIRST_COMPLETED
        )
86
87
88
89
90
91
92
93
        for task in pending:
            task.cancel()

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

    return wrapper
94
95
96
97
98
99
100
101


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


def load_aware_call(func):
    @functools.wraps(func)
102
    async def wrapper(*args, **kwargs):
103
        raw_request = kwargs.get("raw_request", args[1] if len(args) > 1 else None)
104
105
106

        if raw_request is None:
            raise ValueError(
107
108
                "raw_request required when server load tracking is enabled"
            )
109

110
        if not getattr(raw_request.app.state, "enable_server_load_tracking", False):
111
            return await func(*args, **kwargs)
112

113
114
115
116
        # ensure the counter exists
        if not hasattr(raw_request.app.state, "server_load_metrics"):
            raw_request.app.state.server_load_metrics = 0

117
118
        raw_request.app.state.server_load_metrics += 1
        try:
119
            response = await func(*args, **kwargs)
120
121
122
123
124
125
        except Exception:
            raw_request.app.state.server_load_metrics -= 1
            raise

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

        return response

    return wrapper
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166


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"
167
168
169
170


def _validate_truncation_size(
    max_model_len: int,
171
172
173
    truncate_prompt_tokens: int | None,
    tokenization_kwargs: dict[str, Any] | None = None,
) -> int | None:
174
175
176
177
178
179
180
181
    if truncate_prompt_tokens is not None:
        if truncate_prompt_tokens <= -1:
            truncate_prompt_tokens = max_model_len

        if truncate_prompt_tokens > max_model_len:
            raise ValueError(
                f"truncate_prompt_tokens value ({truncate_prompt_tokens}) "
                f"is greater than max_model_len ({max_model_len})."
182
183
                f" Please, select a smaller truncation size."
            )
184
185
186
187
188

        if tokenization_kwargs is not None:
            tokenization_kwargs["truncation"] = True
            tokenization_kwargs["max_length"] = truncate_prompt_tokens

189
190
191
192
    else:
        if tokenization_kwargs is not None:
            tokenization_kwargs["truncation"] = False

193
    return truncate_prompt_tokens
194
195


196
197
def get_max_tokens(
    max_model_len: int,
198
    request: ChatCompletionRequest | CompletionRequest,
199
200
201
202
    input_length: int,
    default_sampling_params: dict,
) -> int:
    max_tokens = getattr(request, "max_completion_tokens", None) or request.max_tokens
203
204
205
    default_max_tokens = max_model_len - input_length
    max_output_tokens = current_platform.get_max_output_tokens(input_length)

206
207
208
209
210
211
212
213
214
215
    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
    )
216
217


218
def log_non_default_args(args: Namespace | EngineArgs):
219
220
    non_default_args = {}

221
222
    # Handle Namespace
    if isinstance(args, Namespace):
223
224
225
226
227
228
229
        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):
230
        default_args = EngineArgs(model=args.model)  # Create default instance
231
232
233
234
235
        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
236
237
        if default_args.model != EngineArgs.model:
            non_default_args["model"] = default_args.model
238
    else:
239
240
241
        raise TypeError(
            "Unsupported argument type. Must be Namespace or EngineArgs instance."
        )
242
243

    logger.info("non-default args: %s", non_default_args)
244
245
246
247
248
249
250
251
252
253
254
255
256


def should_include_usage(
    stream_options: StreamOptions | None, enable_force_include_usage: bool
) -> 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