| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import TYPE_CHECKING |
| 5 | |
| 6 | from fastapi import FastAPI |
| 7 | |
| 8 | from vllm.config import ModelConfig, VllmConfig |
| 9 | from vllm.entrypoints.chat_utils import ChatTemplateConfig |
| 10 | from vllm.logger import init_logger |
| 11 | from vllm.plugins.io_processors import has_io_processor |
| 12 | from vllm.renderers import BaseRenderer |
| 13 | from vllm.tasks import POOLING_TASKS, SCORE_TYPE_MAP, SupportedTask |
| 14 | |
| 15 | from .base.io_processor import PoolingIOProcessor |
| 16 | from .utils import enable_scoring_api |
| 17 | |
| 18 | if TYPE_CHECKING: |
| 19 | from argparse import Namespace |
| 20 | |
| 21 | from starlette.datastructures import State |
| 22 | |
| 23 | from vllm.engine.protocol import EngineClient |
| 24 | from vllm.entrypoints.serve.sagemaker.api_router import ( |
| 25 | EndpointFn, |
| 26 | GetHandlerFn, |
| 27 | RequestType, |
| 28 | ) |
| 29 | from vllm.entrypoints.serve.utils.request_logger import RequestLogger |
| 30 | |
| 31 | else: |
| 32 | RequestLogger = object |
| 33 | |
| 34 | |
| 35 | logger = init_logger(__name__) |
| 36 | |
| 37 | |
| 38 | def init_pooling_io_processors( |
| 39 | supported_tasks: tuple[SupportedTask, ...], |
| 40 | vllm_config: VllmConfig, |
| 41 | renderer: BaseRenderer, |
| 42 | chat_template_config: ChatTemplateConfig, |
| 43 | ) -> dict[str, PoolingIOProcessor]: |
| 44 | model_config = vllm_config.model_config |
| 45 | processors: dict[str, type[PoolingIOProcessor]] = {} |
| 46 | pooling_task = model_config.get_pooling_task(supported_tasks) |
| 47 | |
| 48 | if pooling_task == "classify": |
| 49 | from .classify.io_processor import ClassifyIOProcessor |
| 50 | |
| 51 | processors["classify"] = ClassifyIOProcessor |
| 52 | |
| 53 | if pooling_task == "token_classify": |
| 54 | from .classify.io_processor import TokenClassifyIOProcessor |
| 55 | |
| 56 | processors["token_classify"] = TokenClassifyIOProcessor |
| 57 | |
| 58 | if pooling_task == "embed": |
| 59 | from .embed.io_processor import EmbedIOProcessor |
| 60 | |
| 61 | processors["embed"] = EmbedIOProcessor |
| 62 | |
| 63 | if pooling_task == "token_embed": |
| 64 | from .embed.io_processor import TokenEmbedIOProcessor |
| 65 | |
| 66 | processors["token_embed"] = TokenEmbedIOProcessor |
| 67 | |
| 68 | if pooling_task == "embed&token_classify": |
| 69 | from .pooling.io_processor import UnsupportedCombinedTaskIOProcessor |
| 70 | |
| 71 | processors[pooling_task] = UnsupportedCombinedTaskIOProcessor |
| 72 | |
| 73 | has_plugin = has_io_processor( |
| 74 | vllm_config, |
| 75 | model_config.io_processor_plugin, |
| 76 | ) |
| 77 | if has_plugin: |
| 78 | from .pooling.io_processor import PluginWithIOProcessorPlugins |
| 79 | |
| 80 | processors["plugin"] = PluginWithIOProcessorPlugins |
| 81 | elif pooling_task == "plugin": |
| 82 | from .pooling.io_processor import PluginWithoutIOProcessorPlugins |
| 83 | |
| 84 | processors["plugin"] = PluginWithoutIOProcessorPlugins |
| 85 | |
| 86 | if enable_scoring_api(supported_tasks, model_config): |
| 87 | from .scoring.io_processor import ScoringIOProcessors |
| 88 | |
| 89 | score_type: str | None = SCORE_TYPE_MAP.get(pooling_task, None) # type: ignore[arg-type] |
| 90 | if score_type is not None and score_type in ScoringIOProcessors: |
| 91 | processors[score_type] = ScoringIOProcessors[score_type] |
| 92 | |
| 93 | if model_config.architecture == "JinaForRanking": |
| 94 | from .embed.io_processor import JinaRankingTokenEmbedIOProcessor |
| 95 | from .scoring.io_processor import ScoringIOProcessors |
| 96 | |
| 97 | processors["token_embed"] = JinaRankingTokenEmbedIOProcessor |
| 98 | processors["late-interaction"] = ScoringIOProcessors["jina-reranking-scoring"] |
| 99 | |
| 100 | return { |
| 101 | task: processor_cls( |
| 102 | vllm_config=vllm_config, |
| 103 | renderer=renderer, |
| 104 | chat_template_config=chat_template_config, |
| 105 | ) |
| 106 | for task, processor_cls in processors.items() |
| 107 | } |
| 108 | |
| 109 | |
| 110 | def register_pooling_api_routers( |
| 111 | app: FastAPI, |
| 112 | supported_tasks: tuple["SupportedTask", ...], |
| 113 | model_config: ModelConfig | None = None, |
| 114 | ): |
| 115 | if model_config is None: |
| 116 | return |
| 117 | |
| 118 | pooling_task = model_config.get_pooling_task(supported_tasks) |
| 119 | |
| 120 | if pooling_task is not None: |
| 121 | from .pooling.api_router import router as pooling_router |
| 122 | |
| 123 | app.include_router(pooling_router) |
| 124 | |
| 125 | if "classify" in supported_tasks: |
| 126 | from .classify.api_router import ( |
| 127 | router as classify_router, |
| 128 | ) |
| 129 | |
| 130 | app.include_router(classify_router) |
| 131 | |
| 132 | if "embed" in supported_tasks: |
| 133 | from .embed.api_router import router as embed_router |
| 134 | |
| 135 | app.include_router(embed_router) |
| 136 | |
| 137 | if enable_scoring_api(supported_tasks, model_config): |
| 138 | from .scoring.api_router import router as score_router |
| 139 | |
| 140 | app.include_router(score_router) |
| 141 | |
| 142 | |
| 143 | def init_pooling_state( |
| 144 | engine_client: "EngineClient", |
| 145 | state: "State", |
| 146 | args: "Namespace", |
| 147 | request_logger: RequestLogger | None, |
| 148 | supported_tasks: tuple["SupportedTask", ...], |
| 149 | ): |
| 150 | model_config = engine_client.model_config |
| 151 | if model_config is None: |
| 152 | return |
| 153 | |
| 154 | from vllm.entrypoints.chat_utils import load_chat_template |
| 155 | from vllm.tasks import POOLING_TASKS |
| 156 | |
| 157 | from .classify.serving import ServingClassification |
| 158 | from .embed.serving import ServingEmbedding |
| 159 | from .pooling.serving import ServingPooling |
| 160 | from .scoring.serving import ServingScores |
| 161 | |
| 162 | resolved_chat_template = load_chat_template(args.chat_template) |
| 163 | pooling_task = model_config.get_pooling_task(supported_tasks) |
| 164 | |
| 165 | chat_template_config = ChatTemplateConfig( |
| 166 | chat_template=resolved_chat_template, |
| 167 | chat_template_content_format=args.chat_template_content_format, |
| 168 | trust_request_chat_template=args.trust_request_chat_template, |
| 169 | trust_request_mm_kwargs=args.trust_request_mm_kwargs, |
| 170 | ) |
| 171 | |
| 172 | state.serving_pooling = ( |
| 173 | ( |
| 174 | ServingPooling( |
| 175 | engine_client, |
| 176 | state.openai_serving_models, |
| 177 | supported_tasks=supported_tasks, |
| 178 | request_logger=request_logger, |
| 179 | chat_template_config=chat_template_config, |
| 180 | ) |
| 181 | ) |
| 182 | if any(t in supported_tasks for t in POOLING_TASKS) |
| 183 | else None |
| 184 | ) |
| 185 | state.serving_embedding = ( |
| 186 | ServingEmbedding( |
| 187 | engine_client, |
| 188 | state.openai_serving_models, |
| 189 | request_logger=request_logger, |
| 190 | chat_template_config=chat_template_config, |
| 191 | ) |
| 192 | if pooling_task == "embed" |
| 193 | else None |
| 194 | ) |
| 195 | state.serving_classification = ( |
| 196 | ServingClassification( |
| 197 | engine_client, |
| 198 | state.openai_serving_models, |
| 199 | request_logger=request_logger, |
| 200 | chat_template_config=chat_template_config, |
| 201 | ) |
| 202 | if pooling_task == "classify" |
| 203 | else None |
| 204 | ) |
| 205 | state.serving_scores = ( |
| 206 | ServingScores( |
| 207 | engine_client, |
| 208 | state.openai_serving_models, |
| 209 | supported_tasks=supported_tasks, |
| 210 | request_logger=request_logger, |
| 211 | chat_template_config=chat_template_config, |
| 212 | enable_flash_late_interaction=getattr( |
| 213 | args, "enable_flash_late_interaction", True |
| 214 | ), |
| 215 | ) |
| 216 | if enable_scoring_api(supported_tasks, model_config) |
| 217 | else None |
| 218 | ) |
| 219 | |
| 220 | |
| 221 | def get_pooling_invocation_types( |
| 222 | supported_tasks: tuple["SupportedTask", ...], |
| 223 | model_config: ModelConfig | None = None, |
| 224 | ): |
| 225 | # NOTE: Items defined earlier take higher priority |
| 226 | invocation_types: list[tuple[RequestType, tuple[GetHandlerFn, EndpointFn]]] = [] |
| 227 | |
| 228 | if model_config is None: |
| 229 | return invocation_types |
| 230 | |
| 231 | pooling_task = model_config.get_pooling_task(supported_tasks) |
| 232 | |
| 233 | if pooling_task == "embed": |
| 234 | from .embed.api_router import create_embedding, embedding |
| 235 | from .embed.protocol import EmbeddingRequest |
| 236 | |
| 237 | invocation_types += [ |
| 238 | (EmbeddingRequest, (embedding, create_embedding)), |
| 239 | ] |
| 240 | |
| 241 | if pooling_task == "classify": |
| 242 | from .classify.api_router import classify, create_classify |
| 243 | from .classify.protocol import ClassificationRequest |
| 244 | |
| 245 | invocation_types += [ |
| 246 | (ClassificationRequest, (classify, create_classify)), |
| 247 | ] |
| 248 | |
| 249 | if enable_scoring_api(supported_tasks, model_config): |
| 250 | from .scoring.api_router import do_rerank, rerank |
| 251 | from .scoring.protocol import RerankRequest |
| 252 | |
| 253 | invocation_types += [ |
| 254 | (RerankRequest, (rerank, do_rerank)), |
| 255 | ] |
| 256 | |
| 257 | from .scoring.api_router import create_score, score |
| 258 | from .scoring.protocol import ScoreRequest |
| 259 | |
| 260 | invocation_types += [ |
| 261 | (ScoreRequest, (score, create_score)), |
| 262 | ] |
| 263 | |
| 264 | if any(task in POOLING_TASKS for task in supported_tasks): |
| 265 | from .pooling.api_router import create_pooling, pooling |
| 266 | from .pooling.protocol import PoolingRequest |
| 267 | |
| 268 | invocation_types += [ |
| 269 | (PoolingRequest, (pooling, create_pooling)), |
| 270 | ] |
| 271 | |
| 272 | return invocation_types |