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

4
5
import asyncio
from abc import ABC, abstractmethod
6
from typing import AsyncGenerator, Mapping, Optional
7

8
from vllm.beam_search import BeamSearchSequence, create_sort_beams_key_function
9
from vllm.config import DecodingConfig, ModelConfig, VllmConfig
10
from vllm.core.scheduler import SchedulerOutputs
11
from vllm.inputs.data import PromptType, TokensPrompt
12
from vllm.inputs.parse import is_explicit_encoder_decoder_prompt
13
from vllm.inputs.preprocess import InputPreprocessor
14
from vllm.logger import init_logger
15
from vllm.lora.request import LoRARequest
16
from vllm.model_executor.layers.sampler import SamplerOutput
17
from vllm.outputs import CompletionOutput, PoolingRequestOutput, RequestOutput
18
19
from vllm.pooling_params import PoolingParams
from vllm.prompt_adapter.request import PromptAdapterRequest
20
from vllm.sampling_params import BeamSearchParams, SamplingParams
21
from vllm.transformers_utils.tokenizer import AnyTokenizer
22
from vllm.utils import Device, collect_from_async_generator, random_uuid
23

24
logger = init_logger(__name__)
25

26
27

class EngineClient(ABC):
28
    """Protocol class for Clients to Engine"""
29
30

    @property
31
    @abstractmethod
32
33
34
35
    def is_running(self) -> bool:
        ...

    @property
36
    @abstractmethod
37
38
39
40
    def is_stopped(self) -> bool:
        ...

    @property
41
    @abstractmethod
42
43
44
    def errored(self) -> bool:
        ...

45
    @property
46
    @abstractmethod
47
48
    def dead_error(self) -> BaseException:
        ...
49

50
    @abstractmethod
51
    def generate(
52
        self,
53
        prompt: PromptType,
54
55
56
57
        sampling_params: SamplingParams,
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
58
59
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
60
    ) -> AsyncGenerator[RequestOutput, None]:
61
        """Generate outputs for a request."""
62
        ...
63

64
65
    async def beam_search(
        self,
66
        prompt: PromptType,
67
68
        request_id: str,
        params: BeamSearchParams,
69
        lora_request: Optional[LoRARequest] = None,
70
71
72
73
74
75
76
    ) -> AsyncGenerator[RequestOutput, None]:

        beam_width = params.beam_width
        max_tokens = params.max_tokens
        ignore_eos = params.ignore_eos
        temperature = params.temperature
        length_penalty = params.length_penalty
77
        include_stop_str_in_output = params.include_stop_str_in_output
78

79
80
81
        preprocessor = await self.get_input_preprocessor()
        tokenizer_group = preprocessor.get_tokenizer_group()
        tokenizer = await tokenizer_group.get_lora_tokenizer_async()
82

83
84
85
        if is_explicit_encoder_decoder_prompt(prompt):
            raise NotImplementedError
        else:
86
            processed_inputs = preprocessor._prompt_to_llm_inputs(prompt)
87

88
89
90
        if processed_inputs["type"] == "embeds":
            raise NotImplementedError

91
92
93
94
95
        prompt_token_ids = processed_inputs["prompt_token_ids"]
        prompt_text = processed_inputs.get("prompt")
        multi_modal_data = processed_inputs.get("multi_modal_data")
        mm_processor_kwargs = processed_inputs.get("mm_processor_kwargs")

96
        tokenized_length = len(prompt_token_ids)
97
98
99
100

        sort_beams_key = create_sort_beams_key_function(
            tokenizer.eos_token_id, length_penalty)

101
102
103
104
105
        beam_search_params = SamplingParams(
            logprobs=2 * beam_width,
            max_tokens=1,
            temperature=temperature,
        )
106
        all_beams = [
107
108
            BeamSearchSequence(tokens=prompt_token_ids,
                               cum_logprob=0,
109
                               logprobs=[],
110
                               multi_modal_data=multi_modal_data,
111
112
                               mm_processor_kwargs=mm_processor_kwargs,
                               lora_request=lora_request)
113
        ]
114
115
116
        completed = []

        for _ in range(max_tokens):
117
            prompts_batch, lora_req_batch = zip(*[(
118
119
                TokensPrompt(prompt_token_ids=beam.tokens,
                             multi_modal_data=beam.multi_modal_data,
120
121
122
                             mm_processor_kwargs=beam.mm_processor_kwargs),
                beam.lora_request,
            ) for beam in all_beams])
123
124
125
126

            tasks = []

            request_id = f"beam_search-{random_uuid()}"
127
128
            for i, (individual_prompt,
                    lora_req) in enumerate(zip(prompts_batch, lora_req_batch)):
129
130
131
                request_id_item = f"{request_id}-{i}"
                task = asyncio.create_task(
                    collect_from_async_generator(
132
133
134
135
                        self.generate(individual_prompt,
                                      beam_search_params,
                                      request_id_item,
                                      lora_request=lora_req)))
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
                tasks.append(task)

            output = await asyncio.gather(*tasks)

            output = [x[0] for x in output]

            new_beams = []
            for i, current_beam in enumerate(all_beams):
                result = output[i]

                if result.outputs[0].logprobs is not None:
                    logprobs = result.outputs[0].logprobs[0]
                    for token_id, logprob_obj in logprobs.items():
                        if token_id == tokenizer.eos_token_id and \
                            not ignore_eos:
151
152
153
154
155
156
157
158
159
160
161
                            completed.append(
                                BeamSearchSequence(
                                    tokens=current_beam.tokens +
                                    [token_id] if include_stop_str_in_output
                                    else current_beam.tokens,
                                    logprobs=current_beam.logprobs +
                                    [logprobs],
                                    cum_logprob=current_beam.cum_logprob +
                                    logprob_obj.logprob,
                                    finish_reason="stop",
                                    stop_reason=tokenizer.eos_token_id))
162
                        else:
163
164
165
166
167
                            new_beams.append(
                                BeamSearchSequence(
                                    tokens=current_beam.tokens + [token_id],
                                    logprobs=current_beam.logprobs +
                                    [logprobs],
168
                                    lora_request=current_beam.lora_request,
169
170
171
172
173
174
                                    cum_logprob=current_beam.cum_logprob +
                                    logprob_obj.logprob,
                                    multi_modal_data=current_beam.
                                    multi_modal_data,
                                    mm_processor_kwargs=current_beam.
                                    mm_processor_kwargs))
175
176
177
178
179
180
181
182
183

            sorted_beams = sorted(new_beams, key=sort_beams_key, reverse=True)
            all_beams = sorted_beams[:beam_width]

        completed.extend(all_beams)
        sorted_completed = sorted(completed, key=sort_beams_key, reverse=True)
        best_beams = sorted_completed[:beam_width]

        for beam in best_beams:
Robert Shaw's avatar
Robert Shaw committed
184
185
186
187
188
189
            if (beam.tokens[-1] == tokenizer.eos_token_id and not ignore_eos):
                # Skip the eos token in the text.
                tokens = beam.tokens[tokenized_length:-1]
            else:
                tokens = beam.tokens[tokenized_length:]
            beam.text = tokenizer.decode(tokens)
190
191
192

        beam_search_output = RequestOutput(
            request_id=request_id,
193
            prompt=prompt_text,
194
            outputs=[
195
196
197
198
199
200
201
202
203
                CompletionOutput(text=beam.text,
                                 cumulative_logprob=beam.cum_logprob,
                                 token_ids=beam.tokens[tokenized_length:],
                                 index=i,
                                 logprobs=beam.logprobs,
                                 finish_reason=beam.finish_reason if
                                 beam.finish_reason is not None else "length",
                                 stop_reason=beam.stop_reason)
                for (i, beam) in enumerate(best_beams)
204
205
            ],
            finished=True,
206
            prompt_token_ids=prompt_token_ids,
207
208
209
210
211
            prompt_logprobs=None)

        yield beam_search_output

    @abstractmethod
212
    def encode(
213
        self,
214
        prompt: PromptType,
215
216
217
218
        pooling_params: PoolingParams,
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
219
        priority: int = 0,
220
    ) -> AsyncGenerator[PoolingRequestOutput, None]:
221
        """Generate outputs for a request from a pooling model."""
222
        ...
223

224
    @abstractmethod
225
226
227
228
229
230
    async def abort(self, request_id: str) -> None:
        """Abort a request.

        Args:
            request_id: The unique id of the request.
        """
231
        ...
232
233
234
235
236

    @abstractmethod
    async def get_vllm_config(self) -> VllmConfig:
        """Get the vllm configuration of the vLLM engine."""
        ...
237

238
    @abstractmethod
239
240
    async def get_model_config(self) -> ModelConfig:
        """Get the model configuration of the vLLM engine."""
241
        ...
242

243
    @abstractmethod
244
245
    async def get_decoding_config(self) -> DecodingConfig:
        """Get the decoding configuration of the vLLM engine."""
246
247
248
249
250
251
        ...

    @abstractmethod
    async def get_input_preprocessor(self) -> InputPreprocessor:
        """Get the input processor of the vLLM engine."""
        ...
252

253
    @abstractmethod
254
255
256
    async def get_tokenizer(
        self,
        lora_request: Optional[LoRARequest] = None,
257
258
259
    ) -> AnyTokenizer:
        """Get the appropriate tokenizer for the request"""
        ...
260

261
    @abstractmethod
262
    async def is_tracing_enabled(self) -> bool:
263
        ...
264

265
    @abstractmethod
266
267
268
    async def do_log_stats(
        self,
        scheduler_outputs: Optional[SchedulerOutputs] = None,
269
        model_output: Optional[list[SamplerOutput]] = None,
270
    ) -> None:
271
        ...
272

273
    @abstractmethod
274
275
    async def check_health(self) -> None:
        """Raise if unhealthy"""
276
        ...
277

278
    @abstractmethod
279
280
281
282
    async def start_profile(self) -> None:
        """Start profiling the engine"""
        ...

283
    @abstractmethod
284
285
286
    async def stop_profile(self) -> None:
        """Start profiling the engine"""
        ...
287

288
289
290
291
292
    @abstractmethod
    async def reset_mm_cache(self) -> None:
        """Reset the multi-modal cache"""
        ...

293
    @abstractmethod
294
295
    async def reset_prefix_cache(self,
                                 device: Optional[Device] = None) -> None:
296
297
298
        """Reset the prefix cache"""
        ...

299
300
301
302
303
304
    @abstractmethod
    async def sleep(self, level: int = 1) -> None:
        """Sleep the engine"""
        ...

    @abstractmethod
305
    async def wake_up(self, tags: Optional[list[str]] = None) -> None:
306
307
308
        """Wake up the engine"""
        ...

309
310
311
312
313
    @abstractmethod
    async def is_sleeping(self) -> bool:
        """Check whether the engine is sleeping"""
        ...

314
315
316
317
    @abstractmethod
    async def add_lora(self, lora_request: LoRARequest) -> None:
        """Load a new LoRA adapter into the engine for future requests."""
        ...