| 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 logging import Logger |
| 10 | from string import Template |
| 11 | from typing import Any |
| 12 | |
| 13 | from fastapi import Request |
| 14 | from fastapi.exceptions import RequestValidationError |
| 15 | from fastapi.responses import JSONResponse, StreamingResponse |
| 16 | from starlette.background import BackgroundTask, BackgroundTasks |
| 17 | |
| 18 | from vllm import envs |
| 19 | from vllm.engine.arg_utils import EngineArgs |
| 20 | from vllm.entrypoints.generate.base.protocol import StreamOptions |
| 21 | from vllm.entrypoints.openai.models.protocol import LoRAModulePath |
| 22 | from vllm.logger import current_formatter_type, init_logger |
| 23 | from vllm.platforms import current_platform |
| 24 | from vllm.utils.argparse_utils import FlexibleArgumentParser |
| 25 | |
| 26 | logger = init_logger(__name__) |
| 27 | |
| 28 | VLLM_SUBCMD_PARSER_EPILOG = ( |
| 29 | "For full list: vllm {subcmd} --help=all\n" |
| 30 | "For a section: vllm {subcmd} --help=ModelConfig (case-insensitive)\n" # noqa: E501 |
| 31 | "For a flag: vllm {subcmd} --help=max-model-len (_ or - accepted)\n" # noqa: E501 |
| 32 | "Documentation: https://docs.vllm.ai\n" |
| 33 | ) |
| 34 | |
| 35 | |
| 36 | async def listen_for_disconnect(request: Request) -> None: |
| 37 | """Returns if a disconnect message is received.""" |
| 38 | while True: |
| 39 | message = await request.receive() |
| 40 | if message["type"] == "http.disconnect": |
| 41 | # If load tracking is enabled *and* the counter exists, decrement |
| 42 | # it. Combines the previous nested checks into a single condition |
| 43 | # to satisfy the linter rule. |
| 44 | if getattr( |
| 45 | request.app.state, "enable_server_load_tracking", False |
| 46 | ) and hasattr(request.app.state, "server_load_metrics"): |
| 47 | request.app.state.server_load_metrics -= 1 |
| 48 | break |
| 49 | |
| 50 | |
| 51 | def with_cancellation(handler_func): |
| 52 | """Decorator that allows a route handler to be cancelled by client |
| 53 | disconnections. |
| 54 | |
| 55 | This does _not_ use request.is_disconnected, which does not work with |
| 56 | middleware. Instead this follows the pattern from |
| 57 | starlette.StreamingResponse, which simultaneously awaits on two tasks- one |
| 58 | to wait for an http disconnect message, and the other to do the work that we |
| 59 | want done. When the first task finishes, the other is cancelled. |
| 60 | |
| 61 | A core assumption of this method is that the body of the request has already |
| 62 | been read. This is a safe assumption to make for fastapi handlers that have |
| 63 | already parsed the body of the request into a pydantic model for us. |
| 64 | This decorator is unsafe to use elsewhere, as it will consume and throw away |
| 65 | all incoming messages for the request while it looks for a disconnect |
| 66 | message. |
| 67 | |
| 68 | In the case where a `StreamingResponse` is returned by the handler, this |
| 69 | wrapper will stop listening for disconnects and instead the response object |
| 70 | will start listening for disconnects. |
| 71 | """ |
| 72 | |
| 73 | # Functools.wraps is required for this wrapper to appear to fastapi as a |
| 74 | # normal route handler, with the correct request type hinting. |
| 75 | @functools.wraps(handler_func) |
| 76 | async def wrapper(*args, **kwargs): |
| 77 | # The request is either the second positional arg or `raw_request` |
| 78 | request = args[1] if len(args) > 1 else kwargs["raw_request"] |
| 79 | |
| 80 | handler_task = asyncio.create_task(handler_func(*args, **kwargs)) |
| 81 | cancellation_task = asyncio.create_task(listen_for_disconnect(request)) |
| 82 | |
| 83 | done, pending = await asyncio.wait( |
| 84 | [handler_task, cancellation_task], return_when=asyncio.FIRST_COMPLETED |
| 85 | ) |
| 86 | for task in pending: |
| 87 | task.cancel() |
| 88 | |
| 89 | if handler_task in done: |
| 90 | return handler_task.result() |
| 91 | return None |
| 92 | |
| 93 | return wrapper |
| 94 | |
| 95 | |
| 96 | def decrement_server_load(request: Request): |
| 97 | request.app.state.server_load_metrics -= 1 |
| 98 | |
| 99 | |
| 100 | def load_aware_call(func): |
| 101 | @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 | if raw_request is None: |
| 106 | raise ValueError( |
| 107 | "raw_request required when server load tracking is enabled" |
| 108 | ) |
| 109 | |
| 110 | if not getattr(raw_request.app.state, "enable_server_load_tracking", False): |
| 111 | return await func(*args, **kwargs) |
| 112 | |
| 113 | # ensure the counter exists |
| 114 | if not hasattr(raw_request.app.state, "server_load_metrics"): |
| 115 | raw_request.app.state.server_load_metrics = 0 |
| 116 | |
| 117 | raw_request.app.state.server_load_metrics += 1 |
| 118 | try: |
| 119 | response = await func(*args, **kwargs) |
| 120 | except Exception: |
| 121 | raw_request.app.state.server_load_metrics -= 1 |
| 122 | raise |
| 123 | |
| 124 | if isinstance(response, (JSONResponse, StreamingResponse)): |
| 125 | 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 | elif isinstance(response.background, BackgroundTask): |
| 130 | # Convert the single BackgroundTask to BackgroundTasks |
| 131 | # and chain the decrement_server_load task to it |
| 132 | tasks = BackgroundTasks() |
| 133 | tasks.add_task( |
| 134 | response.background.func, |
| 135 | *response.background.args, |
| 136 | **response.background.kwargs, |
| 137 | ) |
| 138 | tasks.add_task(decrement_server_load, raw_request) |
| 139 | response.background = tasks |
| 140 | else: |
| 141 | raw_request.app.state.server_load_metrics -= 1 |
| 142 | |
| 143 | return response |
| 144 | |
| 145 | return wrapper |
| 146 | |
| 147 | |
| 148 | def cli_env_setup(): |
| 149 | # The safest multiprocessing method is `spawn`, as the default `fork` method |
| 150 | # is not compatible with some accelerators. The default method will be |
| 151 | # changing in future versions of Python, so we should use it explicitly when |
| 152 | # possible. |
| 153 | # |
| 154 | # We only set it here in the CLI entrypoint, because changing to `spawn` |
| 155 | # could break some existing code using vLLM as a library. `spawn` will cause |
| 156 | # unexpected behavior if the code is not protected by |
| 157 | # `if __name__ == "__main__":`. |
| 158 | # |
| 159 | # References: |
| 160 | # - https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods |
| 161 | # - https://pytorch.org/docs/stable/notes/multiprocessing.html#cuda-in-multiprocessing |
| 162 | # - https://pytorch.org/docs/stable/multiprocessing.html#sharing-cuda-tensors |
| 163 | # - https://docs.habana.ai/en/latest/PyTorch/Getting_Started_with_PyTorch_and_Gaudi/Getting_Started_with_PyTorch.html?highlight=multiprocessing#torch-multiprocessing-for-dataloaders |
| 164 | if "VLLM_WORKER_MULTIPROC_METHOD" not in os.environ: |
| 165 | logger.debug("Setting VLLM_WORKER_MULTIPROC_METHOD to 'spawn'") |
| 166 | os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" |
| 167 | |
| 168 | |
| 169 | def get_max_tokens( |
| 170 | max_model_len: int, |
| 171 | max_tokens: int | None, |
| 172 | input_length: int, |
| 173 | default_sampling_params: dict, |
| 174 | override_max_tokens: int | None = None, |
| 175 | truncate_prompt_tokens: int | None = None, |
| 176 | ) -> int: |
| 177 | if truncate_prompt_tokens is not None: |
| 178 | limit = truncate_prompt_tokens |
| 179 | input_length = min( |
| 180 | input_length, |
| 181 | max_model_len if limit == -1 else limit, |
| 182 | ) |
| 183 | if max_model_len < input_length: |
| 184 | raise ValueError( |
| 185 | f"Input length ({input_length}) exceeds model's maximum " |
| 186 | f"context length ({max_model_len})." |
| 187 | ) |
| 188 | model_max_tokens = max_model_len - input_length |
| 189 | platform_max_tokens = current_platform.get_max_output_tokens(input_length) |
| 190 | fallback_max_tokens = ( |
| 191 | max_tokens |
| 192 | if max_tokens is not None |
| 193 | else default_sampling_params.get("max_tokens") |
| 194 | ) |
| 195 | |
| 196 | return min( |
| 197 | val |
| 198 | for val in ( |
| 199 | model_max_tokens, |
| 200 | fallback_max_tokens, |
| 201 | override_max_tokens, |
| 202 | platform_max_tokens, |
| 203 | ) |
| 204 | if val is not None |
| 205 | ) |
| 206 | |
| 207 | |
| 208 | def get_non_default_args(args: Namespace | EngineArgs) -> dict[str, Any]: |
| 209 | from vllm.entrypoints.launchers.cli_args import make_arg_parser |
| 210 | |
| 211 | non_default_args = {} |
| 212 | |
| 213 | # Handle Namespace |
| 214 | if isinstance(args, Namespace): |
| 215 | parser = make_arg_parser(FlexibleArgumentParser()) |
| 216 | for arg, default in vars(parser.parse_args([])).items(): |
| 217 | if default != getattr(args, arg): |
| 218 | non_default_args[arg] = getattr(args, arg) |
| 219 | |
| 220 | # Handle EngineArgs instance |
| 221 | elif isinstance(args, EngineArgs): |
| 222 | default_args = EngineArgs(model=args.model) # Create default instance |
| 223 | for field in dataclasses.fields(args): |
| 224 | current_val = getattr(args, field.name) |
| 225 | default_val = getattr(default_args, field.name) |
| 226 | if current_val != default_val: |
| 227 | non_default_args[field.name] = current_val |
| 228 | if default_args.model != EngineArgs.model: |
| 229 | non_default_args["model"] = default_args.model |
| 230 | else: |
| 231 | raise TypeError( |
| 232 | "Unsupported argument type. Must be Namespace or EngineArgs instance." |
| 233 | ) |
| 234 | |
| 235 | return non_default_args |
| 236 | |
| 237 | |
| 238 | def _jsonify_arg_value(value: Any) -> Any: |
| 239 | if value is None or isinstance(value, bool | int | float | str): |
| 240 | return value |
| 241 | if dataclasses.is_dataclass(value) and not isinstance(value, type): |
| 242 | return { |
| 243 | key: _jsonify_arg_value(val) |
| 244 | for key, val in dataclasses.asdict(value).items() |
| 245 | } |
| 246 | if isinstance(value, dict): |
| 247 | return {str(key): _jsonify_arg_value(val) for key, val in value.items()} |
| 248 | if isinstance(value, tuple | list): |
| 249 | return [_jsonify_arg_value(item) for item in value] |
| 250 | if (model_dump := getattr(value, "model_dump", None)) is not None: |
| 251 | return _jsonify_arg_value(model_dump(mode="json")) |
| 252 | if (to_dict := getattr(value, "dict", None)) is not None: |
| 253 | return _jsonify_arg_value(to_dict()) |
| 254 | return repr(value) |
| 255 | |
| 256 | |
| 257 | def jsonify_non_default_args( |
| 258 | args: Namespace | EngineArgs, |
| 259 | *, |
| 260 | exclude: set[str] | None = None, |
| 261 | ) -> dict[str, Any]: |
| 262 | non_default_args = get_non_default_args(args) |
| 263 | if exclude is not None: |
| 264 | for key in exclude: |
| 265 | non_default_args.pop(key, None) |
| 266 | |
| 267 | return {key: _jsonify_arg_value(value) for key, value in non_default_args.items()} |
| 268 | |
| 269 | |
| 270 | # Fields whose values must never be logged verbatim. |
| 271 | _SENSITIVE_ARG_FIELDS = frozenset({"api_key", "hf_token", "watermark_config"}) |
| 272 | |
| 273 | |
| 274 | def redact_sensitive_args(args: dict[str, Any]) -> dict[str, Any]: |
| 275 | """Return a copy of `args` with sensitive values redacted for logging.""" |
| 276 | if not any(key in _SENSITIVE_ARG_FIELDS for key in args): |
| 277 | return args |
| 278 | return { |
| 279 | key: ("***" if key in _SENSITIVE_ARG_FIELDS else value) |
| 280 | for key, value in args.items() |
| 281 | } |
| 282 | |
| 283 | |
| 284 | def log_non_default_args(args: Namespace | EngineArgs): |
| 285 | non_default_args = get_non_default_args(args) |
| 286 | logger.info("non-default args: %s", redact_sensitive_args(non_default_args)) |
| 287 | |
| 288 | |
| 289 | def should_include_usage( |
| 290 | stream_options: StreamOptions | None, enable_force_include_usage: bool |
| 291 | ) -> tuple[bool, bool]: |
| 292 | if enable_force_include_usage: |
| 293 | return True, True |
| 294 | if stream_options: |
| 295 | include_usage = bool(stream_options.include_usage) |
| 296 | include_continuous_usage = include_usage and bool( |
| 297 | stream_options.continuous_usage_stats |
| 298 | ) |
| 299 | else: |
| 300 | include_usage, include_continuous_usage = False, False |
| 301 | return include_usage, include_continuous_usage |
| 302 | |
| 303 | |
| 304 | def process_lora_modules( |
| 305 | args_lora_modules: list[LoRAModulePath], default_mm_loras: dict[str, str] | None |
| 306 | ) -> list[LoRAModulePath]: |
| 307 | from vllm.entrypoints.openai.models.serving import LoRAModulePath |
| 308 | |
| 309 | lora_modules = args_lora_modules |
| 310 | if default_mm_loras: |
| 311 | default_mm_lora_paths = [ |
| 312 | LoRAModulePath( |
| 313 | name=modality, |
| 314 | path=lora_path, |
| 315 | ) |
| 316 | for modality, lora_path in default_mm_loras.items() |
| 317 | ] |
| 318 | if args_lora_modules is None: |
| 319 | lora_modules = default_mm_lora_paths |
| 320 | else: |
| 321 | lora_modules += default_mm_lora_paths |
| 322 | return lora_modules |
| 323 | |
| 324 | |
| 325 | def log_version_and_model(lgr: Logger, version: str, model_name: str) -> None: |
| 326 | if envs.VLLM_DISABLE_LOG_LOGO or (formatter := current_formatter_type(lgr)) is None: |
| 327 | message = "vLLM server version %s, serving model %s" |
| 328 | else: |
| 329 | logo_template = Template( |
| 330 | "\n ${w}█ █ █▄ ▄█${r}\n" |
| 331 | " ${o}▄▄${r} ${b}▄█${r} ${w}█ █ █ ▀▄▀ █${r} version ${w}%s${r}\n" |
| 332 | " ${o}█${r}${b}▄█▀${r} ${w}█ █ █ █${r} model ${w}%s${r}\n" |
| 333 | " ${b}▀▀${r} ${w}▀▀▀▀▀ ▀▀▀▀▀ ▀ ▀${r}\n" |
| 334 | ) |
| 335 | colors = { |
| 336 | "w": "\033[1m", # bold, default foreground |
| 337 | "o": "\033[93m", # orange |
| 338 | "b": "\033[94m", # blue |
| 339 | "r": "\033[0m", # reset |
| 340 | } |
| 341 | if formatter != "color": |
| 342 | # monochrome logo (no ansi escape codes) |
| 343 | colors = dict.fromkeys(colors, "") |
| 344 | |
| 345 | message = logo_template.substitute(colors) |
| 346 | |
| 347 | lgr.info(message, version, model_name) |
| 348 | |
| 349 | |
| 350 | async def validate_json_request(raw_request: Request): |
| 351 | content_type = raw_request.headers.get("content-type", "").lower() |
| 352 | media_type = content_type.split(";", maxsplit=1)[0] |
| 353 | if media_type != "application/json": |
| 354 | raise RequestValidationError( |
| 355 | errors=["Unsupported Media Type: Only 'application/json' is allowed"] |
| 356 | ) |