| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import asyncio |
| 4 | from abc import ABC, abstractmethod |
| 5 | from collections.abc import AsyncGenerator, Mapping |
| 6 | from concurrent.futures import Executor |
| 7 | from http import HTTPStatus |
| 8 | from typing import ClassVar |
| 9 | |
| 10 | import torch |
| 11 | from fastapi import Request |
| 12 | from fastapi.responses import Response |
| 13 | from starlette.datastructures import Headers |
| 14 | |
| 15 | from vllm import PoolingRequestOutput, envs |
| 16 | from vllm.config import VllmConfig |
| 17 | from vllm.engine.protocol import EngineClient |
| 18 | from vllm.entrypoints.chat_utils import ChatTemplateConfig |
| 19 | from vllm.entrypoints.generate.base.protocol import validate_request_mm_kwargs |
| 20 | from vllm.entrypoints.openai.models.serving import OpenAIServingModels |
| 21 | from vllm.entrypoints.serve.engine.serving import BaseServing |
| 22 | from vllm.entrypoints.serve.engine.typing import AnyRequest |
| 23 | from vllm.entrypoints.serve.utils.request_logger import RequestLogger |
| 24 | from vllm.exceptions import GenerationError, RetryableRequestError |
| 25 | from vllm.lora.request import LoRARequest |
| 26 | from vllm.renderers.base import BaseRenderer |
| 27 | from vllm.tracing import ( |
| 28 | contains_trace_headers, |
| 29 | extract_trace_headers, |
| 30 | log_tracing_disabled_warning, |
| 31 | ) |
| 32 | from vllm.utils.async_utils import make_async, merge_async_iterators |
| 33 | |
| 34 | from ...serve.engine.protocol import ErrorResponse |
| 35 | from ..typing import AnyPoolingRequest, PoolingServeContext |
| 36 | from .io_processor import PoolingIOProcessor |
| 37 | |
| 38 | |
| 39 | class PoolingBaseServing(ABC, BaseServing): |
| 40 | request_id_prefix: ClassVar[str] |
| 41 | |
| 42 | def __init__( |
| 43 | self, |
| 44 | engine_client: EngineClient, |
| 45 | models: OpenAIServingModels, |
| 46 | *, |
| 47 | request_logger: RequestLogger | None, |
| 48 | chat_template_config: ChatTemplateConfig, |
| 49 | return_tokens_as_token_ids: bool = False, |
| 50 | log_error_stack: bool = False, |
| 51 | ): |
| 52 | super().__init__( |
| 53 | models=models, |
| 54 | model_config=models.model_config, |
| 55 | request_logger=request_logger, |
| 56 | ) |
| 57 | |
| 58 | self.engine_client = engine_client |
| 59 | self.renderer = engine_client.renderer |
| 60 | self.vllm_config = engine_client.vllm_config |
| 61 | self.max_model_len = self.model_config.max_model_len |
| 62 | self.return_tokens_as_token_ids = return_tokens_as_token_ids |
| 63 | self.log_error_stack = log_error_stack |
| 64 | self.chat_template_config = chat_template_config |
| 65 | self.trust_request_mm_kwargs = chat_template_config.trust_request_mm_kwargs |
| 66 | |
| 67 | # Shared thread pool executor for preprocessing and postprocessing. |
| 68 | self._executor: Executor = self.renderer._executor |
| 69 | self._postprocessing_async = make_async( |
| 70 | self._postprocessing, executor=self._executor |
| 71 | ) |
| 72 | |
| 73 | async def __call__( |
| 74 | self, |
| 75 | request: AnyPoolingRequest, |
| 76 | raw_request: Request | None = None, |
| 77 | ) -> Response: |
| 78 | io_processor = self.get_io_processor(request) |
| 79 | ctx = await self._init_ctx(io_processor, request, raw_request) |
| 80 | await self._preprocessing(io_processor, ctx) |
| 81 | await self._prepare_generators(ctx) |
| 82 | await self._collect_batch(ctx) |
| 83 | return await self._postprocessing_async(io_processor, ctx) |
| 84 | |
| 85 | @abstractmethod |
| 86 | def get_io_processor(self, request: AnyPoolingRequest) -> PoolingIOProcessor: |
| 87 | raise NotImplementedError |
| 88 | |
| 89 | async def _preprocessing( |
| 90 | self, io_processor: PoolingIOProcessor, ctx: PoolingServeContext |
| 91 | ): |
| 92 | requests = io_processor.get_request_factory_online(ctx) |
| 93 | |
| 94 | if len(requests) == 0: |
| 95 | raise ValueError("You must pass at least one prompt") |
| 96 | |
| 97 | ctx.engine_inputs = await asyncio.gather( |
| 98 | *[io_processor.render_async(request) for request in requests] |
| 99 | ) |
| 100 | |
| 101 | @torch.inference_mode() |
| 102 | def _postprocessing( |
| 103 | self, io_processor: PoolingIOProcessor, ctx: PoolingServeContext |
| 104 | ): |
| 105 | io_processor.post_process_online(ctx) |
| 106 | return self._build_response(ctx) |
| 107 | |
| 108 | async def _init_ctx( |
| 109 | self, |
| 110 | io_processor: PoolingIOProcessor, |
| 111 | request: AnyPoolingRequest, |
| 112 | raw_request: Request | None = None, |
| 113 | ): |
| 114 | validate_request_mm_kwargs( |
| 115 | mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None), |
| 116 | media_io_kwargs=getattr(request, "media_io_kwargs", None), |
| 117 | trust_request_mm_kwargs=self.trust_request_mm_kwargs, |
| 118 | ) |
| 119 | base_request_id = self._base_request_id( |
| 120 | raw_request, getattr(request, "request_id", None) |
| 121 | ) |
| 122 | request_id = f"{self.request_id_prefix}-{base_request_id}" |
| 123 | await self._check_model(request) |
| 124 | |
| 125 | pooling_params = io_processor.create_pooling_params(request) |
| 126 | lora_request = self._maybe_get_adapters(request) |
| 127 | model_name = self.models.model_name(lora_request) |
| 128 | priorities = getattr(request, "priority", 0) |
| 129 | prompt_extras = { |
| 130 | k: v |
| 131 | for k in ("mm_processor_kwargs", "cache_salt", "chat_template_kwargs") |
| 132 | if (v := getattr(request, k, None)) is not None |
| 133 | } |
| 134 | |
| 135 | ctx = PoolingServeContext( |
| 136 | request=request, |
| 137 | raw_request=raw_request, |
| 138 | model_name=model_name, |
| 139 | pooling_params=pooling_params, |
| 140 | request_id=request_id, |
| 141 | lora_request=lora_request, |
| 142 | priorities=priorities, |
| 143 | prompt_extras=prompt_extras, |
| 144 | ) |
| 145 | |
| 146 | self._validate_request(ctx) |
| 147 | return ctx |
| 148 | |
| 149 | async def _prepare_generators( |
| 150 | self, |
| 151 | ctx: PoolingServeContext, |
| 152 | ): |
| 153 | if ctx.engine_inputs is None: |
| 154 | raise ValueError("Engine inputs not available") |
| 155 | |
| 156 | generators: list[AsyncGenerator[PoolingRequestOutput, None]] = [] |
| 157 | |
| 158 | trace_headers = ( |
| 159 | None |
| 160 | if ctx.raw_request is None |
| 161 | else await self._get_trace_headers(ctx.raw_request.headers) |
| 162 | ) |
| 163 | |
| 164 | assert ctx.pooling_params is not None |
| 165 | pooling_params = ctx.pooling_params |
| 166 | pooling_params.verify(self.model_config) |
| 167 | |
| 168 | for i, engine_input in enumerate(ctx.engine_inputs): |
| 169 | prompt_request_id = ( |
| 170 | f"{ctx.request_id}-{i}" |
| 171 | if ctx.prompt_request_ids is None |
| 172 | else ctx.prompt_request_ids[i] |
| 173 | ) |
| 174 | |
| 175 | self._log_inputs( |
| 176 | prompt_request_id, |
| 177 | engine_input["prompts"], |
| 178 | params=engine_input["params"], |
| 179 | lora_request=ctx.lora_request, |
| 180 | ) |
| 181 | |
| 182 | generator = self.engine_client.encode( |
| 183 | request_id=prompt_request_id, |
| 184 | prompt=engine_input["prompts"], |
| 185 | pooling_params=engine_input["params"], |
| 186 | lora_request=engine_input["lora_requests"], |
| 187 | priority=engine_input["priorities"], |
| 188 | trace_headers=trace_headers, |
| 189 | ) |
| 190 | |
| 191 | generators.append(generator) |
| 192 | |
| 193 | ctx.result_generator = merge_async_iterators(*generators) |
| 194 | |
| 195 | async def _collect_batch( |
| 196 | self, |
| 197 | ctx: PoolingServeContext, |
| 198 | ): |
| 199 | if ctx.engine_inputs is None: |
| 200 | raise ValueError("Engine prompts not available") |
| 201 | |
| 202 | if ctx.result_generator is None: |
| 203 | raise ValueError("Result generator not available") |
| 204 | |
| 205 | num_inputs = len(ctx.engine_inputs) |
| 206 | final_res_batch: list[PoolingRequestOutput | None] |
| 207 | final_res_batch = [None] * num_inputs |
| 208 | |
| 209 | async for i, res in ctx.result_generator: |
| 210 | if res.error is not None: |
| 211 | if res.error.retryable: |
| 212 | raise RetryableRequestError(res.error.message) |
| 213 | raise GenerationError(res.error.message) |
| 214 | final_res_batch[i] = res |
| 215 | |
| 216 | if None in final_res_batch: |
| 217 | raise ValueError("Failed to generate results for all prompts") |
| 218 | |
| 219 | ctx.final_res_batch = [res for res in final_res_batch if res is not None] |
| 220 | |
| 221 | @abstractmethod |
| 222 | def _build_response( |
| 223 | self, |
| 224 | ctx: PoolingServeContext, |
| 225 | ) -> Response: |
| 226 | raise NotImplementedError |
| 227 | |
| 228 | async def _check_model( |
| 229 | self, |
| 230 | request: AnyRequest | AnyPoolingRequest, |
| 231 | ) -> ErrorResponse | None: |
| 232 | if self._is_model_supported(request.model): |
| 233 | return None |
| 234 | if request.model in self.models.lora_requests: |
| 235 | return None |
| 236 | if ( |
| 237 | envs.VLLM_ALLOW_RUNTIME_LORA_UPDATING |
| 238 | and request.model |
| 239 | and (load_result := await self.models.resolve_lora(request.model)) |
| 240 | ): |
| 241 | if isinstance(load_result, LoRARequest): |
| 242 | return None |
| 243 | if ( |
| 244 | isinstance(load_result, ErrorResponse) |
| 245 | and load_result.error.code == HTTPStatus.BAD_REQUEST.value |
| 246 | ): |
| 247 | raise ValueError(load_result.error.message) |
| 248 | return None |
| 249 | |
| 250 | def _validate_request(self, ctx: PoolingServeContext) -> None: |
| 251 | truncate_prompt_tokens = getattr(ctx.request, "truncate_prompt_tokens", None) |
| 252 | |
| 253 | if ( |
| 254 | truncate_prompt_tokens is not None |
| 255 | and truncate_prompt_tokens > self.max_model_len |
| 256 | ): |
| 257 | raise ValueError( |
| 258 | "truncate_prompt_tokens value is " |
| 259 | "greater than max_model_len." |
| 260 | " Please request a smaller truncation size." |
| 261 | ) |
| 262 | |
| 263 | return None |
| 264 | |
| 265 | async def _get_trace_headers( |
| 266 | self, |
| 267 | headers: Headers, |
| 268 | ) -> Mapping[str, str] | None: |
| 269 | is_tracing_enabled = await self.engine_client.is_tracing_enabled() |
| 270 | |
| 271 | if is_tracing_enabled: |
| 272 | return extract_trace_headers(headers) |
| 273 | |
| 274 | if contains_trace_headers(headers): |
| 275 | log_tracing_disabled_warning() |
| 276 | |
| 277 | return None |
| 278 | |
| 279 | |
| 280 | class PoolingServing(PoolingBaseServing, ABC): |
| 281 | def __init__(self, *args, **kwargs): |
| 282 | super().__init__(*args, **kwargs) |
| 283 | |
| 284 | self.io_processor = self.init_io_processor( |
| 285 | vllm_config=self.vllm_config, |
| 286 | renderer=self.renderer, |
| 287 | chat_template_config=self.chat_template_config, |
| 288 | ) |
| 289 | |
| 290 | @abstractmethod |
| 291 | def init_io_processor( |
| 292 | self, |
| 293 | vllm_config: VllmConfig, |
| 294 | renderer: BaseRenderer, |
| 295 | chat_template_config: ChatTemplateConfig, |
| 296 | ) -> PoolingIOProcessor: |
| 297 | raise NotImplementedError |
| 298 | |
| 299 | def get_io_processor(self, request: AnyPoolingRequest) -> PoolingIOProcessor: |
| 300 | return self.io_processor |