utils.py 12.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 dataclasses
6
import functools
7
import os
8
from argparse import Namespace
9
from pathlib import Path
10
from typing import Any
11

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

17
from vllm.config import ModelConfig
18
from vllm.engine.arg_utils import EngineArgs
19
20
21
22
23
24
from vllm.engine.protocol import EngineClient
from vllm.entrypoints.chat_utils import (
    load_chat_template,
    resolve_hf_chat_template,
    resolve_mistral_chat_template,
)
25
from vllm.entrypoints.openai.chat_completion.protocol import (
26
    ChatCompletionRequest,
27
28
29
)
from vllm.entrypoints.openai.cli_args import make_arg_parser
from vllm.entrypoints.openai.engine.protocol import (
30
31
32
    CompletionRequest,
    StreamOptions,
)
33
from vllm.entrypoints.openai.serving_models import LoRAModulePath
34
from vllm.logger import init_logger
35
from vllm.platforms import current_platform
36
from vllm.tokenizers.mistral import MistralTokenizer
37
from vllm.utils.argparse_utils import FlexibleArgumentParser
38
39
40

logger = init_logger(__name__)

41
VLLM_SUBCMD_PARSER_EPILOG = (
42
43
44
    "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
45
46
    "Documentation:            https://docs.vllm.ai\n"
)
47

48
49
50
51
52
53

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":
54
55
56
            # 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.
57
58
59
            if getattr(
                request.app.state, "enable_server_load_tracking", False
            ) and hasattr(request.app.state, "server_load_metrics"):
60
                request.app.state.server_load_metrics -= 1
61
62
63
64
65
66
            break


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

68
    This does _not_ use request.is_disconnected, which does not work with
69
    middleware. Instead this follows the pattern from
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
    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))

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

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

    return wrapper
107
108
109
110
111
112
113
114


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


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

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

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

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

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

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

        return response

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


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"
180
181
182
183


def _validate_truncation_size(
    max_model_len: int,
184
185
186
    truncate_prompt_tokens: int | None,
    tokenization_kwargs: dict[str, Any] | None = None,
) -> int | None:
187
188
189
190
191
192
193
194
    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})."
195
196
                f" Please, select a smaller truncation size."
            )
197
198
199
200
201

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

202
203
204
205
    else:
        if tokenization_kwargs is not None:
            tokenization_kwargs["truncation"] = False

206
    return truncate_prompt_tokens
207
208


209
210
def get_max_tokens(
    max_model_len: int,
211
    request: ChatCompletionRequest | CompletionRequest,
212
213
214
215
    input_length: int,
    default_sampling_params: dict,
) -> int:
    max_tokens = getattr(request, "max_completion_tokens", None) or request.max_tokens
216
217
218
    default_max_tokens = max_model_len - input_length
    max_output_tokens = current_platform.get_max_output_tokens(input_length)

219
220
221
222
223
224
225
226
227
228
    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
    )
229
230


231
def log_non_default_args(args: Namespace | EngineArgs):
232
233
    non_default_args = {}

234
235
    # Handle Namespace
    if isinstance(args, Namespace):
236
237
238
239
240
241
242
        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):
243
        default_args = EngineArgs(model=args.model)  # Create default instance
244
245
246
247
248
        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
249
250
        if default_args.model != EngineArgs.model:
            non_default_args["model"] = default_args.model
251
    else:
252
253
254
        raise TypeError(
            "Unsupported argument type. Must be Namespace or EngineArgs instance."
        )
255
256

    logger.info("non-default args: %s", non_default_args)
257
258
259
260
261
262
263
264
265
266
267
268
269


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
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293


def process_lora_modules(
    args_lora_modules: list[LoRAModulePath], default_mm_loras: dict[str, str] | None
) -> list[LoRAModulePath]:
    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


async def process_chat_template(
    args_chat_template: Path | str | None,
    engine_client: EngineClient,
294
    model_config: ModelConfig,
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
) -> str | None:
    resolved_chat_template = load_chat_template(args_chat_template)
    if resolved_chat_template is not None:
        # Get the tokenizer to check official template
        tokenizer = await engine_client.get_tokenizer()

        if isinstance(tokenizer, MistralTokenizer):
            # The warning is logged in resolve_mistral_chat_template.
            resolved_chat_template = resolve_mistral_chat_template(
                chat_template=resolved_chat_template
            )
        else:
            hf_chat_template = resolve_hf_chat_template(
                tokenizer=tokenizer,
                chat_template=None,
                tools=None,
311
                model_config=model_config,
312
313
314
315
316
317
318
319
            )

            if hf_chat_template != resolved_chat_template:
                logger.warning(
                    "Using supplied chat template: %s\n"
                    "It is different from official chat template '%s'. "
                    "This discrepancy may lead to performance degradation.",
                    resolved_chat_template,
320
                    model_config.model,
321
322
                )
    return resolved_chat_template
323
324
325
326
327


def sanitize_message(message: str) -> str:
    # Avoid leaking memory address from object reprs
    return re.sub(r" at 0x[0-9a-f]+>", ">", message)