async_llm.py 13.4 KB
Newer Older
1
import asyncio
2
import os
3
from typing import AsyncGenerator, List, Mapping, Optional, Type, Union
4

5
6
import numpy as np

7
8
9
from vllm.config import ModelConfig, VllmConfig
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.engine.protocol import EngineClient
10
from vllm.envs import VLLM_V1_OUTPUT_PROC_CHUNK_SIZE
11
from vllm.inputs import INPUT_REGISTRY, InputRegistry, PromptType
12
from vllm.inputs.preprocess import InputPreprocessor
13
14
from vllm.logger import init_logger
from vllm.lora.request import LoRARequest
15
from vllm.outputs import RequestOutput
16
17
18
19
20
21
from vllm.pooling_params import PoolingParams
from vllm.prompt_adapter.request import PromptAdapterRequest
from vllm.sampling_params import SamplingParams
from vllm.transformers_utils.tokenizer import AnyTokenizer
from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
from vllm.usage.usage_lib import UsageContext
22
from vllm.utils import cdiv, kill_process_tree
23
from vllm.v1.engine.core_client import EngineCoreClient
24
from vllm.v1.engine.output_processor import OutputProcessor
25
from vllm.v1.engine.processor import Processor
26
from vllm.v1.executor.abstract import Executor
27
from vllm.v1.metrics.loggers import LoggingStatLogger, StatLoggerBase
28
from vllm.v1.metrics.stats import IterationStats, SchedulerStats
29
30
31
32
33
34
35
36
37

logger = init_logger(__name__)


class AsyncLLM(EngineClient):

    def __init__(
        self,
        vllm_config: VllmConfig,
38
        executor_class: Type[Executor],
39
40
41
42
43
44
45
        log_stats: bool,
        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
        input_registry: InputRegistry = INPUT_REGISTRY,
        use_cached_outputs: bool = False,
        log_requests: bool = True,
        start_engine_loop: bool = True,
    ) -> None:
46

47
48
49
50
        assert start_engine_loop

        self.log_requests = log_requests
        self.log_stats = log_stats
51
52
53
54
        self.stat_loggers: List[StatLoggerBase] = [
            LoggingStatLogger(),
            # TODO(rob): PrometheusStatLogger(),
        ]
55
56
57
58
59
60
61
        self.model_config = vllm_config.model_config

        # Tokenizer (+ ensure liveness if running in another process).
        self.tokenizer = init_tokenizer_from_configs(
            model_config=vllm_config.model_config,
            scheduler_config=vllm_config.scheduler_config,
            parallel_config=vllm_config.parallel_config,
62
            lora_config=vllm_config.lora_config)
63
64
65
        self.tokenizer.ping()

        # Processor (converts Inputs --> EngineCoreRequests).
66
67
68
69
70
71
72
        self.processor = Processor(
            model_config=vllm_config.model_config,
            cache_config=vllm_config.cache_config,
            lora_config=vllm_config.lora_config,
            tokenizer=self.tokenizer,
            input_registry=input_registry,
        )
73

74
75
76
        # OutputProcessor (converts EngineCoreOutputs --> RequestOutput).
        self.output_processor = OutputProcessor(self.tokenizer,
                                                log_stats=self.log_stats)
77
78
79
80
81

        # EngineCore (starts the engine in background process).
        self.engine_core = EngineCoreClient.make_client(
            multiprocess_mode=True,
            asyncio_mode=True,
82
83
            vllm_config=vllm_config,
            executor_class=executor_class,
84
85
        )

86
        self.output_handler: Optional[asyncio.Task] = None
87
88
89
90
91
92
93
94

    @classmethod
    def from_engine_args(
        cls,
        engine_args: AsyncEngineArgs,
        engine_config: Optional[VllmConfig] = None,
        start_engine_loop: bool = True,
        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
95
    ) -> "AsyncLLM":
96
97
98
99
        """Create an AsyncLLM from the EngineArgs."""

        # Create the engine configs.
        if engine_config is None:
100
            vllm_config = engine_args.create_engine_config(usage_context)
101
102
103
        else:
            vllm_config = engine_config

104
        executor_class = Executor.get_class(vllm_config)
105
106
107
108
109
110
111
112
113
114
115
116
117
118

        # Create the AsyncLLM.
        return cls(
            vllm_config=vllm_config,
            executor_class=executor_class,
            log_requests=not engine_args.disable_log_requests,
            log_stats=not engine_args.disable_log_stats,
            start_engine_loop=start_engine_loop,
            usage_context=usage_context,
        )

    def shutdown(self):
        """Shutdown, cleaning up the background proc and IPC."""

119
120
        if engine_core := getattr(self, "engine_core", None):
            engine_core.shutdown()
121
122
123
124
125
126
127
128
129
130
131
132
133
134

        if handler := getattr(self, "output_handler", None):
            handler.cancel()

    async def add_request(
        self,
        request_id: str,
        prompt: PromptType,
        params: Union[SamplingParams, PoolingParams],
        arrival_time: Optional[float] = None,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
135
    ) -> asyncio.Queue[RequestOutput]:
136
137
        """Add new request to the AsyncLLM."""

138
        # 1) Create a new output queue for the request.
139
        if self.output_processor.is_request_active(request_id):
140
            raise ValueError(f"Request id {request_id} already running.")
141
        queue: asyncio.Queue[RequestOutput] = asyncio.Queue()
142

143
144
145
146
147
148
        # 2) Convert Input --> Request.
        request = self.processor.process_inputs(request_id, prompt, params,
                                                arrival_time, lora_request,
                                                trace_headers,
                                                prompt_adapter_request,
                                                priority)
149

150
151
        # 3) Add the request to OutputProcessor (this process).
        self.output_processor.add_request(request, queue)
152
153

        # 4) Add the EngineCoreRequest to EngineCore (separate process).
154
        await self.engine_core.add_request_async(request)
155

156
157
158
        if self.log_requests:
            logger.info("Added request %s.", request_id)

159
        return queue
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178

    # TODO: we should support multiple prompts in one call, as you
    # can do with LLM.generate. So that for multi-prompt completion
    # requests we don't need to send multiple messages to core proc,
    # and so we don't need multiple streams which then get
    # re-multiplexed in the API server anyhow.
    async def generate(
        self,
        prompt: PromptType,
        sampling_params: SamplingParams,
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
    ) -> AsyncGenerator[RequestOutput, None]:
        """
        Main function called by the API server to kick off a request
            * 1) Making an AsyncStream corresponding to the Request.
179
            * 2) Processing the Input.
180
181
182
183
184
185
186
187
188
189
190
            * 3) Adding the Request to the Detokenizer.
            * 4) Adding the Request to the EngineCore (separate process).

        A separate output_handler loop runs in a background AsyncIO task, 
        pulling outputs from EngineCore and putting them into the 
        per-request AsyncStream.

        The caller of generate() iterates the returned AsyncGenerator,
        returning the RequestOutput back to the caller.
        """

191
192
193
194
195
196
197
198
199
        try:
            # We start the output_handler on the first call to generate() so
            # we can call __init__ before the event loop, which enables us
            # to handle startup failure gracefully in the OpenAI server.
            if self.output_handler is None:
                self.output_handler = asyncio.create_task(
                    self._run_output_handler())

            q = await self.add_request(
200
201
202
203
204
205
206
                request_id,
                prompt,
                sampling_params,
                lora_request=lora_request,
                trace_headers=trace_headers,
                prompt_adapter_request=prompt_adapter_request,
                priority=priority,
207
            )
208

209
210
            # The output_handler task pushes items into the queue.
            # This task pulls from the queue and yields to caller.
211
212
            finished = False
            while not finished:
213
214
                # Note: drain queue without await if possible (avoids
                # task switching under load which helps performance).
215
                out = q.get_nowait() if not q.empty() else await q.get()
216

217
                # Note: both OutputProcessor and EngineCore handle their
218
                # own request cleanup based on finished.
219
                finished = out.finished
220
221
222
223
224
225
226
227
                yield out

        # If the request is disconnected by the client, the
        # generate() task will be canceled. So, we abort the
        # request if we end up here.
        except asyncio.CancelledError:
            await self.abort(request_id)
            raise
228
229
230
231
232
233

    async def _run_output_handler(self):
        """Background loop: pulls from EngineCore and pushes to AsyncStreams."""

        try:
            while True:
234
                # 1) Pull EngineCoreOutputs from the EngineCore.
235
236
                outputs = await self.engine_core.get_output_async()

237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
                # Split outputs into chunks of at most
                # VLLM_V1_OUTPUT_PROC_CHUNK_SIZE, so that we don't block the
                # event loop for too long.
                num_outputs = len(outputs.outputs)
                if num_outputs <= VLLM_V1_OUTPUT_PROC_CHUNK_SIZE:
                    slices = (outputs.outputs, )
                else:
                    slices = np.array_split(
                        outputs.outputs,
                        cdiv(num_outputs, VLLM_V1_OUTPUT_PROC_CHUNK_SIZE))

                iteration_stats = None
                for i, outputs_slice in enumerate(slices):
                    # 2) Process EngineCoreOutputs.
                    processed_outputs = self.output_processor.process_outputs(
                        outputs_slice, iteration_stats)
                    # NOTE: RequestOutputs are pushed to their queues.
                    assert not processed_outputs.request_outputs
                    iteration_stats = processed_outputs.iteration_stats

                    # Allow other asyncio tasks to run between chunks
                    if i + 1 < len(slices):
                        await asyncio.sleep(0)

                    # 3) Abort any reqs that finished due to stop strings.
                    await self.engine_core.abort_requests_async(
                        processed_outputs.reqs_to_abort)
264

265
266
267
                # 4) Logging.
                # TODO(rob): make into a coroutine and launch it in
                # background thread once we add Prometheus.
268
                assert iteration_stats is not None
269
270
                self._log_stats(
                    scheduler_stats=outputs.scheduler_stats,
271
                    iteration_stats=iteration_stats,
272
                )
273

274
275
276
        except Exception as e:
            logger.exception("EngineCore output handler hit an error: %s", e)
            kill_process_tree(os.getpid())
277
278

    async def abort(self, request_id: str) -> None:
279
        """Abort RequestId in OutputProcessor and EngineCore."""
280
281
282

        request_ids = [request_id]
        await self.engine_core.abort_requests_async(request_ids)
283
        self.output_processor.abort_requests(request_ids)
284

285
286
        if self.log_requests:
            logger.info("Aborted request %s.", request_id)
287

288
289
290
291
292
    def _log_stats(
        self,
        scheduler_stats: SchedulerStats,
        iteration_stats: IterationStats,
    ):
293
294
295
296
297
298
        if not self.log_stats:
            return

        for logger in self.stat_loggers:
            logger.log(scheduler_stats=scheduler_stats)

299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
    def encode(
        self,
        prompt: PromptType,
        pooling_params: PoolingParams,
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
        priority: int = 0,
    ):
        raise ValueError("Not Supported on V1 yet.")

    async def get_model_config(self) -> ModelConfig:
        return self.model_config

    async def get_decoding_config(self):
        raise ValueError("Not Supported on V1 yet.")

316
317
318
    async def get_input_preprocessor(self) -> InputPreprocessor:
        return self.processor.input_preprocessor

319
320
321
322
    async def get_tokenizer(
        self,
        lora_request: Optional[LoRARequest] = None,
    ) -> AnyTokenizer:
323
        return self.tokenizer.get_lora_tokenizer(lora_request)
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338

    async def is_tracing_enabled(self) -> bool:
        return False

    async def do_log_stats(
        self,
        scheduler_outputs=None,
        model_output=None,
    ) -> None:
        logger.debug("Called do_log_stats.")

    async def check_health(self) -> None:
        logger.debug("Called check_health.")

    async def start_profile(self) -> None:
339
        await self.engine_core.profile_async(True)
340
341

    async def stop_profile(self) -> None:
342
        await self.engine_core.profile_async(False)
343

344
345
346
    async def reset_prefix_cache(self) -> None:
        await self.engine_core.reset_prefix_cache_async()

347
348
349
350
351
352
353
354
355
356
357
358
359
360
    @property
    def is_running(self) -> bool:
        return True

    @property
    def is_stopped(self) -> bool:
        return False

    @property
    def errored(self) -> bool:
        return False

    @property
    def dead_error(self) -> BaseException:
361
        return Exception()  # TODO: implement
362
363
364
365

    async def add_lora(self, lora_request: LoRARequest) -> None:
        """Load a new LoRA adapter into the engine for future requests."""
        raise NotImplementedError("LoRA not yet supported in V1")