serving.py 23.2 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import json
4
5
6
from collections.abc import AsyncGenerator, Callable, Mapping
from functools import partial
from typing import Any, Final, Literal, TypeAlias, cast
7

8
import torch
9
from fastapi import Request
10
from typing_extensions import assert_never
11

12
from vllm.engine.protocol import EngineClient
13
from vllm.entrypoints.chat_utils import ChatTemplateContentFormatOption
14
from vllm.entrypoints.logger import RequestLogger
15
16
from vllm.entrypoints.openai.engine.protocol import ErrorResponse, UsageInfo
from vllm.entrypoints.openai.engine.serving import OpenAIServing, ServeContext
17
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
18
19
20
21
22
23
24
25
from vllm.entrypoints.pooling.embed.protocol import (
    EmbeddingBytesResponse,
    EmbeddingChatRequest,
    EmbeddingCompletionRequest,
    EmbeddingRequest,
    EmbeddingResponse,
    EmbeddingResponseData,
)
26
27
28
29
30
from vllm.entrypoints.pooling.utils import (
    encode_pooling_bytes,
    encode_pooling_output_base64,
    encode_pooling_output_float,
)
31
from vllm.inputs.data import ProcessorInputs, TokensPrompt, token_inputs
32
from vllm.logger import init_logger
33
from vllm.outputs import PoolingOutput, PoolingRequestOutput
34
from vllm.pooling_params import PoolingParams
35
36
from vllm.utils.async_utils import merge_async_iterators
from vllm.utils.collection_utils import chunk_list
37
from vllm.utils.serial_utils import EmbedDType, Endianness
38
39
40
41

logger = init_logger(__name__)


42
EmbeddingServeContext: TypeAlias = ServeContext[EmbeddingRequest]
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66


class OpenAIServingEmbedding(OpenAIServing):
    request_id_prefix = "embd"

    def __init__(
        self,
        engine_client: EngineClient,
        models: OpenAIServingModels,
        *,
        request_logger: RequestLogger | None,
        chat_template: str | None,
        chat_template_content_format: ChatTemplateContentFormatOption,
        trust_request_chat_template: bool = False,
    ) -> None:
        super().__init__(
            engine_client=engine_client,
            models=models,
            request_logger=request_logger,
        )

        self.chat_template = chat_template
        self.chat_template_content_format: Final = chat_template_content_format
        self.trust_request_chat_template = trust_request_chat_template
67
68

        pooler_config = self.model_config.pooler_config
69
70
        assert pooler_config is not None
        self.pooler_config = pooler_config
71

72
    async def _preprocess(
73
        self,
74
        ctx: EmbeddingServeContext,
75
    ) -> ErrorResponse | None:
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
        ctx.lora_request = self._maybe_get_adapters(ctx.request)

        if isinstance(ctx.request, EmbeddingChatRequest):
            error_check_ret = self._validate_chat_template(
                request_chat_template=ctx.request.chat_template,
                chat_template_kwargs=ctx.request.chat_template_kwargs,
                trust_request_chat_template=self.trust_request_chat_template,
            )
            if error_check_ret is not None:
                return error_check_ret

            _, ctx.engine_prompts = await self._preprocess_chat(
                ctx.request,
                ctx.request.messages,
                default_template=self.chat_template,
                default_template_content_format=self.chat_template_content_format,
                default_template_kwargs=None,
            )
        elif isinstance(ctx.request, EmbeddingCompletionRequest):
            ctx.engine_prompts = await self._preprocess_completion(
                ctx.request,
                prompt_input=ctx.request.input,
                prompt_embeds=None,
            )
        else:
            return self.create_error_response("Invalid classification request type")
102

103
        return None
104

105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
    def request_output_to_embed_json_response(
        self,
        final_res_batch: list[PoolingRequestOutput],
        request_id: str,
        created_time: int,
        model_name: str,
        encoding_format: Literal["float", "base64"],
        embed_dtype: EmbedDType,
        endianness: Endianness,
    ) -> EmbeddingResponse:
        encode_fn = cast(
            Callable[[PoolingRequestOutput], list[float] | str],
            (
                encode_pooling_output_float
                if encoding_format == "float"
                else partial(
                    encode_pooling_output_base64,
                    embed_dtype=embed_dtype,
                    endianness=endianness,
                )
            ),
        )

        items: list[EmbeddingResponseData] = []
        num_prompt_tokens = 0

        for idx, final_res in enumerate(final_res_batch):
            item = EmbeddingResponseData(
                index=idx,
                embedding=encode_fn(final_res),
            )
            prompt_token_ids = final_res.prompt_token_ids

            items.append(item)
            num_prompt_tokens += len(prompt_token_ids)

        usage = UsageInfo(
            prompt_tokens=num_prompt_tokens,
            total_tokens=num_prompt_tokens,
        )

        return EmbeddingResponse(
            id=request_id,
            created=created_time,
            model=model_name,
            data=items,
            usage=usage,
        )

    def request_output_to_embed_bytes_response(
        self,
        final_res_batch: list[PoolingRequestOutput],
        request_id: str,
        created_time: int,
        model_name: str,
        encoding_format: Literal["bytes", "bytes_only"],
        embed_dtype: EmbedDType,
        endianness: Endianness,
    ) -> EmbeddingBytesResponse:
        content, items, usage = encode_pooling_bytes(
            pooling_outputs=final_res_batch,
            embed_dtype=embed_dtype,
            endianness=endianness,
        )

        headers = (
            None
            if encoding_format == "bytes_only"
            else {
                "metadata": json.dumps(
                    {
                        "id": request_id,
                        "created": created_time,
                        "model": model_name,
                        "data": items,
                        "usage": usage,
                    }
                )
            }
        )

        return EmbeddingBytesResponse(content=content, headers=headers)

188
    def _build_response(
189
        self,
190
191
192
193
194
        ctx: EmbeddingServeContext,
    ) -> EmbeddingResponse | EmbeddingBytesResponse | ErrorResponse:
        encoding_format = ctx.request.encoding_format
        embed_dtype = ctx.request.embed_dtype
        endianness = ctx.request.endianness
195

196
197
198
199
200
201
202
203
204
        if encoding_format == "float" or encoding_format == "base64":
            return self.request_output_to_embed_json_response(
                ctx.final_res_batch,
                ctx.request_id,
                ctx.created_time,
                ctx.model_name,
                encoding_format,
                embed_dtype,
                endianness,
205
            )
206

207
208
209
210
211
212
213
214
215
        if encoding_format == "bytes" or encoding_format == "bytes_only":
            return self.request_output_to_embed_bytes_response(
                ctx.final_res_batch,
                ctx.request_id,
                ctx.created_time,
                ctx.model_name,
                encoding_format,
                embed_dtype,
                endianness,
216
217
            )

218
        assert_never(encoding_format)
219

220
221
222
223
224
225
    def _get_max_position_embeddings(self) -> int:
        """Get the model's effective maximum sequence length for chunking."""
        return self.model_config.max_model_len

    def _should_use_chunked_processing(self, request) -> bool:
        """Check if chunked processing should be used for this request."""
226
227
        return (
            isinstance(request, (EmbeddingCompletionRequest, EmbeddingChatRequest))
228
            and self.pooler_config.enable_chunked_processing
229
        )
230
231
232
233

    async def _process_chunked_request(
        self,
        ctx: EmbeddingServeContext,
234
        token_ids: list[int],
235
236
        pooling_params: PoolingParams,
        trace_headers: Mapping[str, str] | None,
237
238
239
240
241
242
243
244
245
        prompt_idx: int,
    ) -> list[AsyncGenerator[PoolingRequestOutput, None]]:
        """Process a single prompt using chunked processing."""
        generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []

        # Split into chunks using max_position_embeddings
        max_pos_embeddings = self._get_max_position_embeddings()
        # Process all chunks for MEAN aggregation
        for chunk_idx, chunk_tokens in enumerate(
246
247
            chunk_list(token_ids, max_pos_embeddings)
        ):
248
            # Create a request ID for this chunk
249
            chunk_request_id = f"{ctx.request_id}-prompt-{prompt_idx}-chunk-{chunk_idx}"
250
251

            # Create engine prompt for this chunk
252
            chunk_engine_prompt = token_inputs(chunk_tokens)
253
254

            # Log the chunk
255
256
            self._log_inputs(
                chunk_request_id,
257
                chunk_engine_prompt,
258
259
260
                params=pooling_params,
                lora_request=ctx.lora_request,
            )
261
262
263
264
265
266
267
268

            # Create generator for this chunk and wrap it to return indices
            original_generator = self.engine_client.encode(
                chunk_engine_prompt,
                pooling_params,
                chunk_request_id,
                lora_request=ctx.lora_request,
                trace_headers=trace_headers,
269
                priority=ctx.request.priority,
270
271
272
273
274
275
276
277
            )

            generators.append(original_generator)

        return generators

    def _validate_input(
        self,
278
        request: object,
279
280
        input_ids: list[int],
        input_text: str,
281
    ) -> TokensPrompt:
282
283
284
285
        """Override to support chunked processing for embedding requests."""
        token_num = len(input_ids)

        # Note: EmbeddingRequest doesn't have max_tokens
286
        if isinstance(request, (EmbeddingCompletionRequest, EmbeddingChatRequest)):
287
288
289
290
291
292
293
            # Check if chunked processing is enabled for pooling models
            enable_chunked = self._should_use_chunked_processing(request)

            # Use max_position_embeddings for chunked processing decisions
            max_pos_embeddings = self._get_max_position_embeddings()

            # Determine the effective max length for validation
294
            if self.pooler_config.max_embed_len:
295
296
                # Use max_embed_len for validation instead of max_model_len
                length_type = "maximum embedding input length"
297
                max_length_value = self.pooler_config.max_embed_len
298
299
300
            else:
                # Fall back to max_model_len validation (original behavior)
                length_type = "maximum context length"
301
                max_length_value = self.model_config.max_model_len
302
303
304
305

            validation_error_msg = (
                "This model's {length_type} is {max_length_value} tokens. "
                "However, you requested {token_num} tokens in the input for "
306
307
                "embedding generation. Please reduce the length of the input."
            )
308
309
310
311
312

            chunked_processing_error_msg = (
                "This model's {length_type} is {max_length_value} tokens. "
                "However, you requested {token_num} tokens in the input for "
                "embedding generation. Please reduce the length of the input "
313
314
                "or enable chunked processing."
            )
315
316
317
318
319
320
321

            # Check if input exceeds max length
            if token_num > max_length_value:
                raise ValueError(
                    validation_error_msg.format(
                        length_type=length_type,
                        max_length_value=max_length_value,
322
323
324
                        token_num=token_num,
                    )
                )
325
326
327
328
329
330
331
332

            # Check for chunked processing
            # when exceeding max_position_embeddings
            if token_num > max_pos_embeddings:
                if enable_chunked:
                    # Allow long inputs when chunked processing is enabled
                    logger.info(
                        "Input length %s exceeds max_position_embeddings "
333
334
335
336
                        "%s, will use chunked processing",
                        token_num,
                        max_pos_embeddings,
                    )
337
338
339
340
341
                else:
                    raise ValueError(
                        chunked_processing_error_msg.format(
                            length_type="maximum position embeddings length",
                            max_length_value=max_pos_embeddings,
342
343
344
                            token_num=token_num,
                        )
                    )
345

346
            return TokensPrompt(prompt=input_text, prompt_token_ids=input_ids)
347
348
349
350
351
352
353

        # For other request types, use the parent's implementation
        return super()._validate_input(request, input_ids, input_text)

    async def _create_single_prompt_generator(
        self,
        ctx: EmbeddingServeContext,
354
        engine_prompt: ProcessorInputs,
355
        pooling_params: PoolingParams,
356
        trace_headers: Mapping[str, str] | None,
357
        prompt_index: int,
358
    ) -> AsyncGenerator[PoolingRequestOutput, None]:
359
360
361
        """Create a generator for a single prompt using standard processing."""
        request_id_item = f"{ctx.request_id}-{prompt_index}"

362
363
364
365
366
367
        self._log_inputs(
            request_id_item,
            engine_prompt,
            params=pooling_params,
            lora_request=ctx.lora_request,
        )
368
369
370
371
372
373
374
375

        # Return the original generator without wrapping
        return self.engine_client.encode(
            engine_prompt,
            pooling_params,
            request_id_item,
            lora_request=ctx.lora_request,
            trace_headers=trace_headers,
376
            priority=ctx.request.priority,
377
378
379
380
        )

    async def _prepare_generators(
        self,
381
        ctx: EmbeddingServeContext,
382
    ) -> ErrorResponse | None:
383
384
385
386
387
388
389
390
391
        """Override to support chunked processing."""
        # Check if we should use chunked processing
        use_chunked = self._should_use_chunked_processing(ctx.request)

        # If no chunked processing needed, delegate to parent class
        if not use_chunked:
            return await super()._prepare_generators(ctx)

        # Custom logic for chunked processing
392
        generators: list[AsyncGenerator[PoolingRequestOutput, None]] = []
393

394
395
396
397
398
        trace_headers = (
            None
            if ctx.raw_request is None
            else await self._get_trace_headers(ctx.raw_request.headers)
        )
399

400
401
402
        pooling_params = self._create_pooling_params(ctx)
        if isinstance(pooling_params, ErrorResponse):
            return pooling_params
403

404
405
        if ctx.engine_prompts is None:
            return self.create_error_response("Engine prompts not available")
406

407
        max_pos_embeddings = self._get_max_position_embeddings()
408

409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
        for i, engine_prompt in enumerate(ctx.engine_prompts):
            # Check if this specific prompt needs chunked processing
            if "prompt_token_ids" in engine_prompt:
                prompt_token_ids = engine_prompt["prompt_token_ids"]  # type: ignore[typeddict-item]

                if len(prompt_token_ids) > max_pos_embeddings:
                    # Use chunked processing for this prompt
                    chunk_generators = await self._process_chunked_request(
                        ctx,
                        prompt_token_ids,
                        pooling_params,
                        trace_headers,
                        i,
                    )
                    generators.extend(chunk_generators)
                    continue
425

426
427
428
429
430
            # Normal processing for short prompts or non-token prompts
            generator = await self._create_single_prompt_generator(
                ctx, engine_prompt, pooling_params, trace_headers, i
            )
            generators.append(generator)
431

432
        ctx.result_generator = merge_async_iterators(*generators)
433

434
        return None
435
436
437

    async def _collect_batch(
        self,
438
        ctx: EmbeddingServeContext,
439
    ) -> ErrorResponse | None:
440
441
        """Collect and aggregate batch results
        with support for chunked processing.
442
443

        For chunked requests, performs online aggregation to
444
445
446
        minimize memory usage.
        For regular requests, collects results normally.
        """
447
448
        if ctx.engine_prompts is None:
            return self.create_error_response("Engine prompts not available")
449

450
451
        # Check if we used chunked processing
        use_chunked = self._should_use_chunked_processing(ctx.request)
452

453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
        if not use_chunked:
            return await super()._collect_batch(ctx=ctx)

        if ctx.result_generator is None:
            return self.create_error_response("Result generator not available")

        # Online aggregation for chunked requests to
        # minimize memory usage
        # Track aggregation state for each prompt
        prompt_aggregators: dict[int, dict[str, Any]] = {}
        short_prompts_results: dict[int, PoolingRequestOutput] = {}

        async for result_idx, result in ctx.result_generator:
            if "-chunk-" in result.request_id:
                # Extract prompt_idx from chunked request_id
                parts = result.request_id.split("-")
                try:
                    prompt_idx = int(parts[parts.index("prompt") + 1])
                except (ValueError, IndexError):
                    # Fallback: extract from result_idx if parsing fails
                    prompt_idx = result_idx

                # Initialize aggregator for this prompt if needed
                if prompt_idx not in prompt_aggregators:
                    prompt_aggregators[prompt_idx] = {
                        "weighted_sum": None,
                        "total_weight": 0,
                        "chunk_count": 0,
                        "request_id": result.request_id.split("-chunk-")[0],
                    }
483

484
485
486
487
488
489
490
491
492
493
494
                aggregator = prompt_aggregators[prompt_idx]

                # MEAN pooling with online weighted averaging
                # Ensure result is PoolingRequestOutput
                # for embedding processing
                if not isinstance(result, PoolingRequestOutput):
                    return self.create_error_response(
                        f"Expected PoolingRequestOutput for "
                        f"chunked embedding, got "
                        f"{type(result).__name__}"
                    )
495

496
497
498
499
500
501
502
503
504
505
506
507
508
                # Handle both PoolingOutput and
                # EmbeddingOutput types
                if hasattr(result.outputs, "data"):
                    # PoolingOutput case
                    embedding_data = result.outputs.data
                elif hasattr(result.outputs, "embedding"):
                    # EmbeddingOutput case -
                    # convert embedding list to tensor
                    embedding_data = result.outputs.embedding
                else:
                    return self.create_error_response(
                        f"Unsupported output type: {type(result.outputs).__name__}"
                    )
509

510
511
                if not isinstance(embedding_data, torch.Tensor):
                    embedding_data = torch.tensor(embedding_data, dtype=torch.float32)
512

513
514
515
516
517
518
519
520
521
522
523
                if result.prompt_token_ids is None:
                    return self.create_error_response(
                        "prompt_token_ids cannot be None for chunked processing"
                    )
                weight = len(result.prompt_token_ids)

                weighted_embedding = embedding_data.to(dtype=torch.float32) * weight

                if aggregator["weighted_sum"] is None:
                    # First chunk
                    aggregator["weighted_sum"] = weighted_embedding
524
                else:
525
526
                    # Accumulate
                    aggregator["weighted_sum"] += weighted_embedding
527

528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
                aggregator["total_weight"] += weight
                aggregator["chunk_count"] += 1
            else:
                # Non-chunked result - extract prompt_idx from request_id
                parts = result.request_id.split("-")
                try:
                    # Last part should be prompt index
                    prompt_idx = int(parts[-1])
                except (ValueError, IndexError):
                    prompt_idx = result_idx  # Fallback to result_idx

                short_prompts_results[prompt_idx] = result

        # Finalize aggregated results
        final_res_batch: list[PoolingRequestOutput] = []
        num_prompts = len(ctx.engine_prompts)

        for prompt_idx in range(num_prompts):
            if prompt_idx in prompt_aggregators:
                # Finalize MEAN aggregation for this chunked prompt
                aggregator = prompt_aggregators[prompt_idx]

                weighted_sum = aggregator["weighted_sum"]
                total_weight = aggregator["total_weight"]

                if (
                    weighted_sum is not None
                    and isinstance(weighted_sum, torch.Tensor)
                    and isinstance(total_weight, (int, float))
                    and total_weight > 0
                ):
                    # Compute final mean embedding
                    final_embedding = weighted_sum / total_weight

                    # Create a PoolingRequestOutput
                    # for the aggregated result
                    pooling_output_data = PoolingOutput(data=final_embedding)

                    # Get original prompt token IDs for this prompt
                    original_prompt = ctx.engine_prompts[prompt_idx]
                    if "prompt_token_ids" not in original_prompt:
569
                        return self.create_error_response(
570
                            f"Chunked prompt {prompt_idx} does not contain token IDs"
571
                        )
572
573
574
575
576
577
578
579
580
581
582
583

                    original_token_ids = original_prompt["prompt_token_ids"]  # type: ignore[typeddict-item]

                    pooling_request_output = PoolingRequestOutput(
                        request_id=aggregator["request_id"],
                        prompt_token_ids=original_token_ids,
                        outputs=pooling_output_data,
                        num_cached_tokens=0,
                        finished=True,
                    )

                    final_res_batch.append(pooling_request_output)
584
585
                else:
                    return self.create_error_response(
586
                        f"Failed to aggregate chunks for prompt {prompt_idx}"
587
                    )
588
589
590
591
592
593
            elif prompt_idx in short_prompts_results:
                final_res_batch.append(short_prompts_results[prompt_idx])
            else:
                return self.create_error_response(
                    f"Result not found for prompt {prompt_idx}"
                )
594

595
        ctx.final_res_batch = final_res_batch
596

597
        return None
598

599
600
601
    async def create_embedding(
        self,
        request: EmbeddingRequest,
602
603
        raw_request: Request | None = None,
    ) -> EmbeddingResponse | ErrorResponse:
604
605
606
607
608
609
        """
        Embedding API similar to OpenAI's API.

        See https://platform.openai.com/docs/api-reference/embeddings/create
        for the API specification. This API mimics the OpenAI Embedding API.
        """
610
        model_name = self.models.model_name()
611
612
        request_id = (
            f"{self.request_id_prefix}-"
613
614
            f"{self._base_request_id(raw_request, request.request_id)}"
        )
615
616
617
618
619
620
621
622

        ctx = EmbeddingServeContext(
            request=request,
            raw_request=raw_request,
            model_name=model_name,
            request_id=request_id,
        )

623
        return await self.handle(ctx)  # type: ignore[return-value]