serving.py 9.98 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4


5
6
7
8
9
10
11
12
13
import asyncio
import time
from collections.abc import AsyncGenerator
from collections.abc import Sequence as GenericSequence

from fastapi import Request

from vllm.engine.protocol import EngineClient
from vllm.entrypoints.logger import RequestLogger
14
from vllm.entrypoints.openai.chat_completion.protocol import (
15
16
17
    ChatCompletionLogProb,
    ChatCompletionLogProbs,
    ChatCompletionLogProbsContent,
18
19
)
from vllm.entrypoints.openai.engine.protocol import (
20
21
22
23
24
    ErrorResponse,
    PromptTokenUsageInfo,
    RequestResponseMetadata,
    UsageInfo,
)
25
from vllm.entrypoints.openai.engine.serving import OpenAIServing, clamp_prompt_logprobs
26
from vllm.entrypoints.openai.models.serving import OpenAIServingModels
27
28
29
30
31
from vllm.entrypoints.serve.disagg.protocol import (
    GenerateRequest,
    GenerateResponse,
    GenerateResponseChoice,
)
32
from vllm.entrypoints.serve.render.serving import OpenAIServingRender
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
from vllm.logger import init_logger
from vllm.logprobs import Logprob
from vllm.outputs import RequestOutput
from vllm.sampling_params import SamplingParams
from vllm.utils.collection_utils import as_list

logger = init_logger(__name__)


class ServingTokens(OpenAIServing):
    """Provides Tokens IN <> Tokens OUT functionality to vLLM API."""

    def __init__(
        self,
        engine_client: EngineClient,
        models: OpenAIServingModels,
49
        openai_serving_render: OpenAIServingRender,
50
51
52
53
54
55
56
        *,
        request_logger: RequestLogger | None,
        force_no_detokenize: bool = False,
        return_tokens_as_token_ids: bool = False,
        enable_prompt_tokens_details: bool = False,
        enable_log_outputs: bool = False,
    ):
57
58
59
60
61
62
        super().__init__(
            engine_client=engine_client,
            models=models,
            request_logger=request_logger,
            return_tokens_as_token_ids=return_tokens_as_token_ids,
        )
63
        self.openai_serving_render = openai_serving_render
64
65
66
67
        self.enable_prompt_tokens_details = enable_prompt_tokens_details
        self.enable_log_outputs = enable_log_outputs
        self.force_no_detokenize = force_no_detokenize
        if force_no_detokenize:
68
69
70
71
            logger.info(
                "Tokens-only mode is enabled, skipping detokenization "
                "step for incoming requests."
            )
72
73
74
75

    async def serve_tokens(
        self,
        request: GenerateRequest,
76
        raw_request: Request | None = None,
77
78
79
80
81
82
83
84
85
86
87
88
89
    ) -> GenerateResponse | ErrorResponse:
        error_check_ret = await self._check_model(request)
        if error_check_ret is not None:
            logger.error("Error with model %s", error_check_ret)
            return error_check_ret

        # If the engine is dead, raise the engine's DEAD_ERROR.
        # This is required for the streaming case, where we return a
        # success status before we actually start generating text :).
        if self.engine_client.errored:
            raise self.engine_client.dead_error

        lora_request = None
90
        lora_request = self._maybe_get_adapters(request, supports_default_mm_loras=True)
91
92
93

        model_name = self.models.model_name(lora_request)

94
95
96
        request_id = (
            f"generate-tokens-{self._base_request_id(raw_request, request.request_id)}"
        )
97
98
99
100
101

        request_metadata = RequestResponseMetadata(request_id=request_id)
        if raw_request:
            raw_request.state.request_metadata = request_metadata

102
        engine_prompts = await self.openai_serving_render.preprocess_completion(
103
104
105
106
107
108
            request,
            prompt_input=request.token_ids,
            prompt_embeds=None,
        )
        assert len(engine_prompts) == 1
        engine_prompt = engine_prompts[0]
109
110
111

        # Schedule the request and get the result generator.
        result_generator: AsyncGenerator[RequestOutput, None] | None = None
112
113
114
115
116
117
118
119
120
121
        sampling_params = request.sampling_params
        if self.force_no_detokenize:
            sampling_params.detokenize = False

        self._log_inputs(
            request_id,
            engine_prompt,
            params=sampling_params,
            lora_request=lora_request,
        )
122

123
124
125
126
127
        trace_headers = (
            None
            if raw_request is None
            else await self._get_trace_headers(raw_request.headers)
        )
128

129
130
131
132
133
134
135
136
        result_generator = self.engine_client.generate(
            engine_prompt,
            sampling_params,
            request_id,
            lora_request=lora_request,
            trace_headers=trace_headers,
            priority=request.priority,
        )
137
138
139

        # TODO(NickLucche): Implement streaming response

140
141
142
143
        assert result_generator is not None
        return await self.serve_tokens_full_generator(
            request, result_generator, request_id, model_name, request_metadata
        )
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

    async def serve_tokens_full_generator(
        self,
        request: GenerateRequest,
        result_generator: AsyncGenerator[RequestOutput, None],
        request_id: str,
        model_name: str,
        request_metadata: RequestResponseMetadata,
    ) -> ErrorResponse | GenerateResponse:
        created_time = int(time.time())
        final_res: RequestOutput | None = None
        sampling_params: SamplingParams = request.sampling_params

        try:
            async for res in result_generator:
                final_res = res
        except asyncio.CancelledError:
            return self.create_error_response("Client disconnected")

        assert final_res is not None

        choices: list[GenerateResponseChoice] = []
        num_generated_tokens = 0
        for output in final_res.outputs:
            token_ids = output.token_ids
            out_logprobs = output.logprobs

            # This is top_logprobs in completions API
172
            if sampling_params.logprobs is not None:
173
174
175
176
177
178
179
180
181
182
183
184
                assert out_logprobs is not None, "Did not output logprobs"
                logprobs = self._create_tokens_logprobs(
                    token_ids=token_ids,
                    top_logprobs=out_logprobs,
                    num_output_top_logprobs=sampling_params.logprobs,
                )
            else:
                logprobs = None

            choice_data = GenerateResponseChoice(
                index=output.index,
                logprobs=logprobs,
185
186
187
                finish_reason=output.finish_reason if output.finish_reason else "stop",
                token_ids=as_list(output.token_ids),
            )
188
189
190
191
192
193
194
195
196

            choices.append(choice_data)
            num_generated_tokens += len(output.token_ids)

        assert final_res.prompt_token_ids is not None
        num_prompt_tokens = len(final_res.prompt_token_ids)
        if final_res.encoder_prompt_token_ids is not None:
            num_prompt_tokens += len(final_res.encoder_prompt_token_ids)

197
198
199
200
201
        usage = UsageInfo(
            prompt_tokens=num_prompt_tokens,
            completion_tokens=num_generated_tokens,
            total_tokens=num_prompt_tokens + num_generated_tokens,
        )
202
203
204
        if self.enable_prompt_tokens_details and final_res.num_cached_tokens:
            # This info is not available at the /coordinator level
            usage.prompt_tokens_details = PromptTokenUsageInfo(
205
206
                cached_tokens=final_res.num_cached_tokens
            )
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225

        request_metadata.final_usage_info = usage

        response = GenerateResponse(
            id=request_id,
            created=created_time,
            model=model_name,
            choices=choices,
            usage=usage,
            prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs),
            kv_transfer_params=final_res.kv_transfer_params,
        )

        # Log complete response if output logging is enabled
        if self.enable_log_outputs and self.request_logger:
            for choice in choices:
                # Get the corresponding output token IDs
                output_token_ids = None
                if choice.index < len(final_res.outputs):
226
                    output_token_ids = final_res.outputs[choice.index].token_ids
227
228
229
230
231

                if output_token_ids:
                    # Log token_ids only.
                    self.request_logger.log_outputs(
                        request_id=request_id,
232
                        outputs="",
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
                        output_token_ids=output_token_ids,
                        finish_reason=choice.finish_reason,
                        is_streaming=False,
                        delta=False,
                    )

        return response

    def _create_tokens_logprobs(
        self,
        token_ids: GenericSequence[int],
        top_logprobs: GenericSequence[dict[int, Logprob] | None],
        num_output_top_logprobs: int | None = None,
    ) -> ChatCompletionLogProbs:
        """Create OpenAI-style logprobs."""
        logprobs_content: list[ChatCompletionLogProbsContent] = []

        for i, token_id in enumerate(token_ids):
            token = f"token_id:{token_id}"
            step_top_logprobs = top_logprobs[i]
253
            if step_top_logprobs is None or step_top_logprobs.get(token_id) is None:
254
                logprobs_content.append(
255
256
257
258
                    ChatCompletionLogProbsContent(
                        token=token,
                    )
                )
259
260
261
262
263
264
265
266
267
268
269
            else:
                step_token = step_top_logprobs[token_id]

                logprobs_content.append(
                    ChatCompletionLogProbsContent(
                        token=token,
                        logprob=max(step_token.logprob, -9999.0),
                        top_logprobs=[
                            ChatCompletionLogProb(
                                token=token,
                                logprob=max(p[1].logprob, -9999.0),
270
271
                            )
                            for i, p in enumerate(step_top_logprobs.items())
272
273
                            if num_output_top_logprobs is not None
                            and i < max(num_output_top_logprobs, 1)
274
275
276
                        ],
                    )
                )
277
278

        return ChatCompletionLogProbs(content=logprobs_content)