async_llm.py 13.7 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
from vllm.pooling_params import PoolingParams
from vllm.prompt_adapter.request import PromptAdapterRequest
18
from vllm.sampling_params import RequestOutputKind, SamplingParams
19
20
21
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
218
219
220
221
222
223
224
                # Coalesce any additional queued outputs
                while not q.empty():
                    next_out = q.get_nowait()
                    if sampling_params.output_kind == RequestOutputKind.DELTA:
                        out.add(next_out)
                    else:
                        out = next_out

225
                # Note: both OutputProcessor and EngineCore handle their
226
                # own request cleanup based on finished.
227
                finished = out.finished
228
229
230
231
232
233
234
235
                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
236
237
238
239
240
241

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

        try:
            while True:
242
                # 1) Pull EngineCoreOutputs from the EngineCore.
243
244
                outputs = await self.engine_core.get_output_async()

245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
                # 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)
272

273
274
275
                # 4) Logging.
                # TODO(rob): make into a coroutine and launch it in
                # background thread once we add Prometheus.
276
                assert iteration_stats is not None
277
278
                self._log_stats(
                    scheduler_stats=outputs.scheduler_stats,
279
                    iteration_stats=iteration_stats,
280
                )
281

282
283
284
        except Exception as e:
            logger.exception("EngineCore output handler hit an error: %s", e)
            kill_process_tree(os.getpid())
285
286

    async def abort(self, request_id: str) -> None:
287
        """Abort RequestId in OutputProcessor and EngineCore."""
288
289
290

        request_ids = [request_id]
        await self.engine_core.abort_requests_async(request_ids)
291
        self.output_processor.abort_requests(request_ids)
292

293
294
        if self.log_requests:
            logger.info("Aborted request %s.", request_id)
295

296
297
298
299
300
    def _log_stats(
        self,
        scheduler_stats: SchedulerStats,
        iteration_stats: IterationStats,
    ):
301
302
303
304
305
306
        if not self.log_stats:
            return

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

307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
    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.")

324
325
326
    async def get_input_preprocessor(self) -> InputPreprocessor:
        return self.processor.input_preprocessor

327
328
329
330
    async def get_tokenizer(
        self,
        lora_request: Optional[LoRARequest] = None,
    ) -> AnyTokenizer:
331
        return self.tokenizer.get_lora_tokenizer(lora_request)
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346

    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:
347
        await self.engine_core.profile_async(True)
348
349

    async def stop_profile(self) -> None:
350
        await self.engine_core.profile_async(False)
351

352
353
354
    async def reset_prefix_cache(self) -> None:
        await self.engine_core.reset_prefix_cache_async()

355
356
357
358
359
360
361
362
363
364
365
366
367
368
    @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:
369
        return Exception()  # TODO: implement
370
371
372
373

    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")