@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 201 lines 5.37 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import time
4 from collections.abc import AsyncGenerator, Callable, Generator, Sequence
5 from dataclasses import dataclass, field
6 from typing import Any, Generic, TypeAlias, TypedDict, TypeVar
7
8 from fastapi import Request
9 from pydantic import ConfigDict
10
11 from vllm import PoolingParams, PoolingRequestOutput, PromptType
12 from vllm.entrypoints.chat_utils import ChatCompletionMessageParam
13 from vllm.inputs import DataPrompt, EngineInput
14 from vllm.lora.request import LoRARequest
15 from vllm.renderers import ChatParams, TokenizeParams
16 from vllm.renderers.inputs import DictPrompt
17
18 from ...tasks import PoolingTask
19 from .classify.protocol import (
20 ClassificationChatRequest,
21 ClassificationCompletionRequest,
22 ClassificationResponse,
23 )
24 from .embed.protocol import (
25 CohereEmbedRequest,
26 EmbeddingBytesResponse,
27 EmbeddingChatInputRequest,
28 EmbeddingChatRequest,
29 EmbeddingCompletionRequest,
30 EmbeddingRequest,
31 EmbeddingResponse,
32 )
33 from .pooling.protocol import (
34 IOProcessorRequest,
35 PoolingBytesResponse,
36 PoolingChatRequest,
37 PoolingCompletionRequest,
38 PoolingResponse,
39 )
40 from .scoring.protocol import ScoringRequest, ScoringResponse
41 from .scoring.typing import ScoreData, ScoringData
42
43 PoolingCompletionLikeRequest: TypeAlias = (
44 EmbeddingCompletionRequest
45 | ClassificationCompletionRequest
46 | PoolingCompletionRequest
47 )
48
49 PoolingChatLikeRequest: TypeAlias = (
50 EmbeddingChatRequest
51 | EmbeddingChatInputRequest
52 | ClassificationChatRequest
53 | PoolingChatRequest
54 )
55
56 AnyPoolingRequest: TypeAlias = (
57 EmbeddingRequest
58 | PoolingCompletionLikeRequest
59 | PoolingChatLikeRequest
60 | IOProcessorRequest
61 | ScoringRequest
62 | CohereEmbedRequest
63 )
64
65 AnyPoolingResponse: TypeAlias = (
66 ClassificationResponse
67 | EmbeddingResponse
68 | EmbeddingBytesResponse
69 | PoolingResponse
70 | PoolingBytesResponse
71 | ScoringResponse
72 )
73
74 PoolingRequestT = TypeVar("PoolingRequestT", bound=AnyPoolingRequest)
75
76
77 @dataclass(kw_only=True)
78 class ChunkedEmbeddingMetadata:
79 prompt_index: int
80 chunk_index: int
81
82
83 @dataclass(kw_only=True)
84 class PoolingServeContext(Generic[PoolingRequestT]):
85 model_config = ConfigDict(arbitrary_types_allowed=True)
86
87 request: PoolingRequestT
88 raw_request: Request | None = None
89 model_name: str
90 request_id: str
91 pooling_params: PoolingParams
92 lora_request: LoRARequest | None
93 priorities: int | Sequence[int] | None
94 prompt_extras: dict[str, Any] | None
95
96 created_time: int = field(default_factory=lambda: int(time.time()))
97 engine_inputs: Sequence["PoolingEngineInput"] | None = None
98 prompt_request_ids: list[str] | None = None
99
100 result_generator: AsyncGenerator[tuple[int, PoolingRequestOutput], None] | None = (
101 None
102 )
103 final_res_batch: list[PoolingRequestOutput] = field(default_factory=list)
104
105 ## for Long Text Embedding with Chunked Processing
106 original_engine_inputs: Sequence["PoolingEngineInput"] | None = None
107 chunked_embedding_metadata: list[ChunkedEmbeddingMetadata] | None = None
108
109 ## for bi-encoder & late-interaction
110 n_queries: int | None = None
111
112 ## for IOProcessorResponse
113 response: Any | None = None
114
115 ## for flash-late-interaction
116 query_final_res_batch: list[PoolingRequestOutput] | None = None
117 late_interaction_query_keys: list[str] | None = None
118 late_interaction_doc_keys: list[str] | None = None
119
120
121 @dataclass
122 class OfflineInputsContext:
123 pooling_task: PoolingTask
124 tokenization_kwargs: dict[str, Any] | None
125 lora_request: Sequence[LoRARequest | None] | None
126 priorities: int | Sequence[int] | None
127
128
129 @dataclass
130 class OfflineEncodeInputsContext(OfflineInputsContext):
131 prompts: PromptType | Sequence[PromptType]
132 pooling_params: PoolingParams | Sequence[PoolingParams] | None
133
134
135 @dataclass
136 class OfflineScoringInputsContext(OfflineInputsContext):
137 scoring_data: ScoringData
138 chat_template: str | None
139 pooling_params: PoolingParams
140
141
142 @dataclass
143 class OfflinePluginInputsContext(OfflineInputsContext):
144 prompts: DataPrompt
145 pooling_params: PoolingParams | Sequence[PoolingParams] | None
146
147
148 AnyOfflineInputsContext: TypeAlias = (
149 OfflineEncodeInputsContext
150 | OfflineScoringInputsContext
151 | OfflinePluginInputsContext
152 )
153
154
155 @dataclass
156 class OfflineOutputsContext:
157 outputs: list[PoolingRequestOutput]
158
159 ## for bi-encoder & late-interaction
160 n_queries: int | None = None
161
162
163 class RenderParams(TypedDict):
164 tok_params: TokenizeParams
165 prompt_extras: dict[str, Any] | None
166 skip_mm_cache: bool
167
168 params: PoolingParams
169 lora_requests: LoRARequest | None
170 priorities: int
171
172
173 class EncodeCMPLRenderParams(RenderParams):
174 prompts: DictPrompt
175
176
177 class EncodeChatRenderParams(RenderParams):
178 conversations: list["ChatCompletionMessageParam"]
179 chat_params: ChatParams
180
181
182 class ScoringRenderParams(RenderParams):
183 data_1: ScoreData
184 data_2: ScoreData
185 chat_template: str | None
186 max_tokens_per_query: int
187 max_tokens_per_doc: int
188
189
190 AnyRenderParam: TypeAlias = (
191 EncodeCMPLRenderParams | EncodeChatRenderParams | ScoringRenderParams
192 )
193 RequestGenerator: TypeAlias = Generator[AnyRenderParam]
194 RequestFactory: TypeAlias = Callable[[], RequestGenerator]
195
196
197 class PoolingEngineInput(TypedDict):
198 prompts: EngineInput
199 params: PoolingParams
200 lora_requests: LoRARequest | None
201 priorities: int