serving.py 9.78 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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
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,
        *,
        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,
    ):
55
56
57
58
59
60
        super().__init__(
            engine_client=engine_client,
            models=models,
            request_logger=request_logger,
            return_tokens_as_token_ids=return_tokens_as_token_ids,
        )
61
62
63
64
        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:
65
66
67
68
            logger.info(
                "Tokens-only mode is enabled, skipping detokenization "
                "step for incoming requests."
            )
69
70
71
72

    async def serve_tokens(
        self,
        request: GenerateRequest,
73
        raw_request: Request | None = None,
74
75
76
77
78
79
80
81
82
83
84
85
86
    ) -> 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
87
        lora_request = self._maybe_get_adapters(request, supports_default_mm_loras=True)
88
89
90

        model_name = self.models.model_name(lora_request)

91
92
93
        request_id = (
            f"generate-tokens-{self._base_request_id(raw_request, request.request_id)}"
        )
94
95
96
97
98

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

99
100
101
102
103
104
105
        engine_prompts = await self._preprocess_completion(
            request,
            prompt_input=request.token_ids,
            prompt_embeds=None,
        )
        assert len(engine_prompts) == 1
        engine_prompt = engine_prompts[0]
106
107
108

        # Schedule the request and get the result generator.
        result_generator: AsyncGenerator[RequestOutput, None] | None = None
109
110
111
112
113
114
115
116
117
118
        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,
        )
119

120
121
122
123
124
        trace_headers = (
            None
            if raw_request is None
            else await self._get_trace_headers(raw_request.headers)
        )
125

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

        # TODO(NickLucche): Implement streaming response

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

    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
169
            if sampling_params.logprobs is not None:
170
171
172
173
174
175
176
177
178
179
180
181
                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,
182
183
184
                finish_reason=output.finish_reason if output.finish_reason else "stop",
                token_ids=as_list(output.token_ids),
            )
185
186
187
188
189
190
191
192
193

            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)

194
195
196
197
198
        usage = UsageInfo(
            prompt_tokens=num_prompt_tokens,
            completion_tokens=num_generated_tokens,
            total_tokens=num_prompt_tokens + num_generated_tokens,
        )
199
200
201
        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(
202
203
                cached_tokens=final_res.num_cached_tokens
            )
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222

        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):
223
                    output_token_ids = final_res.outputs[choice.index].token_ids
224
225
226
227
228

                if output_token_ids:
                    # Log token_ids only.
                    self.request_logger.log_outputs(
                        request_id=request_id,
229
                        outputs="",
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
                        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]
250
            if step_top_logprobs is None or step_top_logprobs.get(token_id) is None:
251
                logprobs_content.append(
252
253
254
255
                    ChatCompletionLogProbsContent(
                        token=token,
                    )
                )
256
257
258
259
260
261
262
263
264
265
266
            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),
267
268
                            )
                            for i, p in enumerate(step_top_logprobs.items())
269
270
                            if num_output_top_logprobs is not None
                            and i < max(num_output_top_logprobs, 1)
271
272
273
                        ],
                    )
                )
274
275

        return ChatCompletionLogProbs(content=logprobs_content)