async_llm_engine.py 28.2 KB
Newer Older
1
2
import asyncio
import time
Antoni Baum's avatar
Antoni Baum committed
3
from functools import partial
4
5
from typing import (Any, AsyncIterator, Callable, Dict, Iterable, List,
                    Optional, Set, Tuple, Type, Union)
6

7
8
from transformers import PreTrainedTokenizer

9
import vllm.envs as envs
10
from vllm.config import DecodingConfig, ModelConfig
11
from vllm.core.scheduler import SchedulerOutputs
Woosuk Kwon's avatar
Woosuk Kwon committed
12
13
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.engine.llm_engine import LLMEngine
14
from vllm.executor.ray_utils import initialize_ray_cluster, ray
Woosuk Kwon's avatar
Woosuk Kwon committed
15
from vllm.logger import init_logger
16
from vllm.lora.request import LoRARequest
Woosuk Kwon's avatar
Woosuk Kwon committed
17
18
from vllm.outputs import RequestOutput
from vllm.sampling_params import SamplingParams
19
from vllm.sequence import MultiModalData, SamplerOutput
yhu422's avatar
yhu422 committed
20
from vllm.usage.usage_lib import UsageContext
21
22

logger = init_logger(__name__)
23
ENGINE_ITERATION_TIMEOUT_S = envs.VLLM_ENGINE_ITERATION_TIMEOUT_S
24

Antoni Baum's avatar
Antoni Baum committed
25

26
27
28
29
class AsyncEngineDeadError(RuntimeError):
    pass


30
31
32
def _raise_exception_on_finish(
        task: asyncio.Task, error_callback: Callable[[Exception],
                                                     None]) -> None:
33
34
    msg = ("Task finished unexpectedly. This should never happen! "
           "Please open an issue on Github.")
35
36

    exception = None
37
    try:
38
39
        task.result()
        # NOTE: This will be thrown if task exits normally (which it should not)
40
        raise AsyncEngineDeadError(msg)
41
42
43
44
45
46
    except Exception as e:
        exception = e
        logger.error("Engine background task failed", exc_info=e)
        error_callback(exception)
        raise AsyncEngineDeadError(
            msg + " See stack trace above for the actual cause.") from e
47
48


Antoni Baum's avatar
Antoni Baum committed
49
50
51
52
53
54
class AsyncStream:
    """A stream of RequestOutputs for a request that can be
    iterated over asynchronously."""

    def __init__(self, request_id: str) -> None:
        self.request_id = request_id
55
        self._queue: asyncio.Queue = asyncio.Queue()
Antoni Baum's avatar
Antoni Baum committed
56
57
        self._finished = False

58
    def put(self, item: Union[RequestOutput, Exception]) -> None:
Antoni Baum's avatar
Antoni Baum committed
59
60
61
62
63
        if self._finished:
            return
        self._queue.put_nowait(item)

    def finish(self) -> None:
64
        self._queue.put_nowait(StopAsyncIteration())
Antoni Baum's avatar
Antoni Baum committed
65
66
67
68
69
70
71
72
73
74
75
        self._finished = True

    @property
    def finished(self) -> bool:
        return self._finished

    def __aiter__(self):
        return self

    async def __anext__(self) -> RequestOutput:
        result = await self._queue.get()
76
        if isinstance(result, Exception):
77
            raise result
Antoni Baum's avatar
Antoni Baum committed
78
79
80
        return result


81
82
83
84
85
86
87
88
class RequestTracker:
    """Synchronous abstraction for tracking requests."""

    def __init__(self) -> None:
        self._request_streams: Dict[str, AsyncStream] = {}
        self._finished_requests: asyncio.Queue[str] = asyncio.Queue()
        self._new_requests: asyncio.Queue[Tuple[AsyncStream,
                                                dict]] = asyncio.Queue()
89
        self.new_requests_event = asyncio.Event()
90
91
92
93

    def __contains__(self, item):
        return item in self._request_streams

94
95
    def __len__(self) -> int:
        return len(self._request_streams)
96
97
98
99
100
101
102
103

    def propagate_exception(self,
                            exc: Exception,
                            request_id: Optional[str] = None) -> None:
        """Propagate an exception to request streams
        (all if request_id is None)."""
        if request_id is not None:
            self._request_streams[request_id].put(exc)
104
            self.abort_request(request_id)
105
        else:
106
            for rid, stream in self._request_streams.items():
107
                stream.put(exc)
108
                self.abort_request(rid)
109
110
111
112
113
114
115
116
117
118
119

    def process_request_output(self,
                               request_output: RequestOutput,
                               *,
                               verbose: bool = False) -> None:
        """Process a request output from the engine."""
        request_id = request_output.request_id

        self._request_streams[request_id].put(request_output)
        if request_output.finished:
            if verbose:
120
                logger.info("Finished request %s.", request_id)
121
122
            self.abort_request(request_id)

123
124
125
126
127
128
129
130
    def process_exception(self,
                          request_id: str,
                          exception: Exception,
                          *,
                          verbose: bool = False) -> None:
        """Propagate an exception from the engine."""
        self._request_streams[request_id].put(exception)
        if verbose:
131
            logger.info("Finished request %s.", request_id)
132
133
        self.abort_request(request_id)

134
135
136
137
138
139
140
141
142
143
144
145
    def add_request(self, request_id: str,
                    **engine_add_request_kwargs) -> AsyncStream:
        """Add a request to be sent to the engine on the next background
        loop iteration."""
        if request_id in self._request_streams:
            raise KeyError(f"Request {request_id} already exists.")

        stream = AsyncStream(request_id)
        self._new_requests.put_nowait((stream, {
            "request_id": request_id,
            **engine_add_request_kwargs
        }))
146
147
148

        self.new_requests_event.set()

149
150
151
152
153
        return stream

    def abort_request(self, request_id: str, *, verbose: bool = False) -> None:
        """Abort a request during next background loop iteration."""
        if verbose:
154
            logger.info("Aborted request %s.", request_id)
155
156
157
158
159
160
161
162
163
164

        self._finished_requests.put_nowait(request_id)

        if request_id not in self._request_streams or self._request_streams[
                request_id].finished:
            # The request has already finished or been aborted.
            return

        self._request_streams[request_id].finish()

165
    def get_new_and_finished_requests(self) -> Tuple[List[Dict], Set[str]]:
166
167
        """Get the new requests and finished requests to be
        sent to the engine."""
168
        new_requests: List[Dict] = []
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
        finished_requests: Set[str] = set()

        while not self._finished_requests.empty():
            request_id = self._finished_requests.get_nowait()
            finished_requests.add(request_id)
            self._request_streams.pop(request_id, None)

        while not self._new_requests.empty():
            stream, new_request = self._new_requests.get_nowait()
            if stream.request_id in finished_requests:
                # The request has already been aborted.
                stream.finish()
                continue
            self._request_streams[stream.request_id] = stream
            new_requests.append(new_request)

        return new_requests, finished_requests
Antoni Baum's avatar
Antoni Baum committed
186

187
    async def wait_for_new_requests(self):
188
189
190
191
192
193
        if not self.has_new_requests():
            await self.new_requests_event.wait()
        self.new_requests_event.clear()

    def has_new_requests(self):
        return not self._new_requests.empty()
194

Antoni Baum's avatar
Antoni Baum committed
195
196
197
198
199
200
201
202
203
204
205
206
207
208

class _AsyncLLMEngine(LLMEngine):
    """Extension of LLMEngine to add async methods."""

    async def step_async(self) -> List[RequestOutput]:
        """Performs one decoding iteration and returns newly generated results.
        The workers are ran asynchronously if possible.

        This function performs one decoding iteration of the engine. It first
        schedules the sequences to be executed in the next iteration and the
        token blocks to be swapped in/out/copy. Then, it executes the model
        and updates the scheduler with the model outputs. Finally, it decodes
        the sequences and returns the newly generated results.
        """
209
        seq_group_metadata_list, scheduler_outputs = self.scheduler.schedule()
Antoni Baum's avatar
Antoni Baum committed
210

211
212
        if not scheduler_outputs.is_empty():
            # Execute the model.
213
            output = await self.model_executor.execute_model_async(
214
215
                seq_group_metadata_list,
                scheduler_outputs.blocks_to_swap_in,
216
                scheduler_outputs.blocks_to_swap_out,
217
218
                scheduler_outputs.blocks_to_copy,
                num_lookahead_slots=scheduler_outputs.num_lookahead_slots)
219
220
        else:
            output = []
Antoni Baum's avatar
Antoni Baum committed
221

222
        request_outputs = self._process_model_outputs(
223
            output, scheduler_outputs.scheduled_seq_groups,
224
            scheduler_outputs.ignored_seq_groups, seq_group_metadata_list)
Antoni Baum's avatar
Antoni Baum committed
225

226
        # Log stats.
227
        self.do_log_stats(scheduler_outputs, output)
228
229
230

        return request_outputs

231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
    async def encode_request_async(
        self,
        request_id: str,  # pylint: disable=unused-argument
        prompt: Optional[str],
        prompt_token_ids: Optional[List[int]] = None,
        lora_request: Optional[LoRARequest] = None,
    ):
        if prompt_token_ids is None:
            assert prompt is not None
            prompt_token_ids = await self.tokenizer.encode_async(
                request_id=request_id,
                prompt=prompt,
                lora_request=lora_request)
        return prompt_token_ids

    async def add_request_async(
        self,
        request_id: str,
        prompt: Optional[str],
        sampling_params: SamplingParams,
        prompt_token_ids: Optional[List[int]] = None,
        arrival_time: Optional[float] = None,
        lora_request: Optional[LoRARequest] = None,
254
        multi_modal_data: Optional[MultiModalData] = None,
255
256
257
258
259
260
261
262
263
264
265
266
    ) -> None:
        if lora_request is not None and not self.lora_config:
            raise ValueError(f"Got lora_request {lora_request} but LoRA is "
                             "not enabled!")
        if arrival_time is None:
            arrival_time = time.time()
        prompt_token_ids = await self.encode_request_async(
            request_id=request_id,
            prompt=prompt,
            prompt_token_ids=prompt_token_ids,
            lora_request=lora_request)

267
268
269
270
271
272
273
        return self.add_request(request_id,
                                prompt=prompt,
                                prompt_token_ids=prompt_token_ids,
                                sampling_params=sampling_params,
                                arrival_time=arrival_time,
                                lora_request=lora_request,
                                multi_modal_data=multi_modal_data)
274

275
276
    async def check_health_async(self) -> None:
        self.model_executor.check_health()
277

278

279
280
class AsyncLLMEngine:
    """An asynchronous wrapper for LLMEngine.
281

282
    This class is used to wrap the LLMEngine class to make it asynchronous. It
283
    uses asyncio to create a background loop that keeps processing incoming
284
    requests. The LLMEngine is kicked by the generate method when there
285
    are requests in the waiting queue. The generate method yields the outputs
286
    from the LLMEngine to the caller.
287

288
    NOTE: For the comprehensive list of arguments, see `LLMEngine`.
289
290
291
292
293

    Args:
        worker_use_ray: Whether to use Ray for model workers. Required for
            distributed execution. Should be the same as
            `parallel_config.worker_use_ray`.
Zhuohan Li's avatar
Zhuohan Li committed
294
        engine_use_ray: Whether to make LLMEngine a Ray actor. If so, the
295
296
            async frontend will be executed in a separate process as the
            model workers.
297
        log_requests: Whether to log the requests.
zspo's avatar
zspo committed
298
299
        max_log_len: Maximum number of prompt characters or prompt ID numbers
            being printed in log.
300
301
        start_engine_loop: If True, the background task to run the engine
            will be automatically started in the generate call.
302
303
        *args: Arguments for LLMEngine.
        *kwargs: Arguments for LLMEngine.
304
    """
305

Antoni Baum's avatar
Antoni Baum committed
306
307
    _engine_class: Type[_AsyncLLMEngine] = _AsyncLLMEngine

308
309
310
311
312
    def __init__(self,
                 worker_use_ray: bool,
                 engine_use_ray: bool,
                 *args,
                 log_requests: bool = True,
313
                 max_log_len: Optional[int] = None,
314
                 start_engine_loop: bool = True,
315
                 **kwargs) -> None:
316
        self.worker_use_ray = worker_use_ray
Zhuohan Li's avatar
Zhuohan Li committed
317
        self.engine_use_ray = engine_use_ray
318
        self.log_requests = log_requests
319
        self.max_log_len = max_log_len
Antoni Baum's avatar
Antoni Baum committed
320
321
        self.engine = self._init_engine(*args, **kwargs)

322
        self.background_loop: Optional[asyncio.Future] = None
323
324
325
        # We need to keep a reference to unshielded
        # task as well to prevent it from being garbage
        # collected
326
        self._background_loop_unshielded: Optional[asyncio.Task[Any]] = None
327
        self.start_engine_loop = start_engine_loop
328
        self._errored_with: Optional[BaseException] = None
Antoni Baum's avatar
Antoni Baum committed
329

330
331
332
        # Lazy initialized fields
        self._request_tracker: RequestTracker

333
    @classmethod
yhu422's avatar
yhu422 committed
334
335
336
337
338
339
    def from_engine_args(
        cls,
        engine_args: AsyncEngineArgs,
        start_engine_loop: bool = True,
        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
    ) -> "AsyncLLMEngine":
340
341
        """Creates an async LLM engine from the engine arguments."""
        # Create the engine configs.
342
        engine_config = engine_args.create_engine_config()
343

344
        if engine_config.device_config.device_type == "neuron":
345
346
            from vllm.executor.neuron_executor import NeuronExecutorAsync
            executor_class = NeuronExecutorAsync
347
348
349
350
351
        elif engine_config.device_config.device_type == "cpu":
            assert not engine_config.parallel_config.worker_use_ray, (
                "Ray is not supported with the CPU backend.")
            from vllm.executor.cpu_executor import CPUExecutorAsync
            executor_class = CPUExecutorAsync
352
        elif engine_config.parallel_config.worker_use_ray:
353
            initialize_ray_cluster(engine_config.parallel_config)
354
355
356
            from vllm.executor.ray_gpu_executor import RayGPUExecutorAsync
            executor_class = RayGPUExecutorAsync
        else:
357
            assert engine_config.parallel_config.world_size == 1, (
358
359
360
361
                "Ray is required if parallel_config.world_size > 1.")
            from vllm.executor.gpu_executor import GPUExecutorAsync
            executor_class = GPUExecutorAsync
        # Create the async LLM engine.
yhu422's avatar
yhu422 committed
362
        engine = cls(
363
            engine_config.parallel_config.worker_use_ray,
yhu422's avatar
yhu422 committed
364
            engine_args.engine_use_ray,
365
366
            **engine_config.to_dict(),
            executor_class=executor_class,
yhu422's avatar
yhu422 committed
367
368
369
370
371
372
            log_requests=not engine_args.disable_log_requests,
            log_stats=not engine_args.disable_log_stats,
            max_log_len=engine_args.max_log_len,
            start_engine_loop=start_engine_loop,
            usage_context=usage_context,
        )
373
374
        return engine

375
376
    @property
    def is_running(self) -> bool:
377
        return (self.background_loop is not None
378
                and self._background_loop_unshielded is not None
379
380
381
382
                and not self._background_loop_unshielded.done())

    @property
    def is_stopped(self) -> bool:
383
384
        return self.errored or (self.background_loop is not None and
                                self._background_loop_unshielded is not None
385
386
387
388
389
390
391
392
393
394
395
396
                                and self._background_loop_unshielded.done())

    @property
    def errored(self) -> bool:
        return self._errored_with is not None

    def set_errored(self, exc: Exception) -> None:
        self._errored_with = exc

    def _error_callback(self, exc: Exception) -> None:
        self.set_errored(exc)
        self._request_tracker.propagate_exception(exc)
397

398
399
    async def get_tokenizer(self) -> "PreTrainedTokenizer":
        if self.engine_use_ray:
400
            return await self.engine.get_tokenizer.remote()  # type: ignore
401
402
        else:
            return self.engine.get_tokenizer()
403

404
    def start_background_loop(self) -> None:
Antoni Baum's avatar
Antoni Baum committed
405
        """Start the background loop."""
406
407
408
        if self.errored:
            raise AsyncEngineDeadError(
                "Background loop has errored already.") from self._errored_with
409
        if self.is_running:
Antoni Baum's avatar
Antoni Baum committed
410
            raise RuntimeError("Background loop is already running.")
411
412
        # Initialize the RequestTracker here so it uses the right event loop.
        self._request_tracker = RequestTracker()
413
414
415
416

        self._background_loop_unshielded = asyncio.get_event_loop(
        ).create_task(self.run_engine_loop())
        self._background_loop_unshielded.add_done_callback(
417
            partial(_raise_exception_on_finish,
418
                    error_callback=self._error_callback))
419
        self.background_loop = asyncio.shield(self._background_loop_unshielded)
Antoni Baum's avatar
Antoni Baum committed
420
421
422

    def _init_engine(self, *args,
                     **kwargs) -> Union[_AsyncLLMEngine, "ray.ObjectRef"]:
Zhuohan Li's avatar
Zhuohan Li committed
423
        if not self.engine_use_ray:
Antoni Baum's avatar
Antoni Baum committed
424
            engine_class = self._engine_class
425
        elif self.worker_use_ray:
Antoni Baum's avatar
Antoni Baum committed
426
            engine_class = ray.remote(num_cpus=0)(self._engine_class).remote
427
        else:
Woosuk Kwon's avatar
Woosuk Kwon committed
428
429
            # FIXME(woosuk): This is a bit hacky. Be careful when changing the
            # order of the arguments.
430
431
            cache_config = kwargs["cache_config"]
            parallel_config = kwargs["parallel_config"]
Woosuk Kwon's avatar
Woosuk Kwon committed
432
433
434
435
436
437
            if parallel_config.tensor_parallel_size == 1:
                num_gpus = cache_config.gpu_memory_utilization
            else:
                num_gpus = 1
            engine_class = ray.remote(num_gpus=num_gpus)(
                self._engine_class).remote
Antoni Baum's avatar
Antoni Baum committed
438
439
        return engine_class(*args, **kwargs)

440
441
442
443
    async def engine_step(self) -> bool:
        """Kick the engine to process the waiting requests.

        Returns True if there are in-progress requests."""
444
445

        new_requests, finished_requests = (
446
            self._request_tracker.get_new_and_finished_requests())
447
448
449
450

        for new_request in new_requests:
            # Add the request into the vLLM engine's waiting queue.
            # TODO: Maybe add add_request_batch to reduce Ray overhead
451
452
            try:
                if self.engine_use_ray:
453
454
                    await self.engine.add_request.remote(  # type: ignore
                        **new_request)
455
456
457
458
459
460
461
462
463
                else:
                    await self.engine.add_request_async(**new_request)
            except ValueError as e:
                # TODO: use a vLLM specific error for failed validation
                self._request_tracker.process_exception(
                    new_request["request_id"],
                    e,
                    verbose=self.log_requests,
                )
464
465
466
467

        if finished_requests:
            await self._engine_abort(finished_requests)

Zhuohan Li's avatar
Zhuohan Li committed
468
        if self.engine_use_ray:
469
            request_outputs = await self.engine.step.remote()  # type: ignore
470
        else:
Antoni Baum's avatar
Antoni Baum committed
471
            request_outputs = await self.engine.step_async()
472

Antoni Baum's avatar
Antoni Baum committed
473
        # Put the outputs into the corresponding streams.
474
        for request_output in request_outputs:
475
            self._request_tracker.process_request_output(
476
                request_output, verbose=self.log_requests)
Antoni Baum's avatar
Antoni Baum committed
477

478
479
        return len(request_outputs) > 0

Antoni Baum's avatar
Antoni Baum committed
480
481
    async def _engine_abort(self, request_ids: Iterable[str]):
        if self.engine_use_ray:
482
            await self.engine.abort_request.remote(request_ids)  # type: ignore
Antoni Baum's avatar
Antoni Baum committed
483
484
485
486
        else:
            self.engine.abort_request(request_ids)

    async def run_engine_loop(self):
487
        has_requests_in_progress = False
Antoni Baum's avatar
Antoni Baum committed
488
        while True:
489
            if not has_requests_in_progress:
490
                logger.debug("Waiting for new requests...")
491
                await self._request_tracker.wait_for_new_requests()
492
493
494
495
496
497
498
499
500
501
502
503
                logger.debug("Got new requests!")

            # Abort if iteration takes too long due to unrecoverable errors
            # (eg. NCCL timeouts).
            try:
                has_requests_in_progress = await asyncio.wait_for(
                    self.engine_step(), ENGINE_ITERATION_TIMEOUT_S)
            except asyncio.TimeoutError as exc:
                logger.error(
                    "Engine iteration timed out. This should never happen!")
                self.set_errored(exc)
                raise
Antoni Baum's avatar
Antoni Baum committed
504
505
506
507
508
509
510
511
512
            await asyncio.sleep(0)

    async def add_request(
        self,
        request_id: str,
        prompt: Optional[str],
        sampling_params: SamplingParams,
        prompt_token_ids: Optional[List[int]] = None,
        arrival_time: Optional[float] = None,
513
        lora_request: Optional[LoRARequest] = None,
514
        multi_modal_data: Optional[MultiModalData] = None,
Antoni Baum's avatar
Antoni Baum committed
515
516
    ) -> AsyncStream:
        if self.log_requests:
517
518
519
520
521
522
523
524
            shortened_prompt = prompt
            shortened_token_ids = prompt_token_ids
            if self.max_log_len is not None:
                if shortened_prompt is not None:
                    shortened_prompt = shortened_prompt[:self.max_log_len]
                if shortened_token_ids is not None:
                    shortened_token_ids = shortened_token_ids[:self.
                                                              max_log_len]
525
526
527
528
529
            logger.info(
                "Received request %s: prompt: %r, "
                "sampling_params: %s, prompt_token_ids: %s, "
                "lora_request: %s.", request_id, shortened_prompt,
                sampling_params, shortened_token_ids, lora_request)
Antoni Baum's avatar
Antoni Baum committed
530

531
        if not self.is_running:
532
533
534
535
536
537
538
539
            if self.start_engine_loop:
                self.start_background_loop()
            else:
                raise AsyncEngineDeadError(
                    "Background loop is not running. If it was running, "
                    "inspect the output to find the stacktrace of the "
                    "error that caused the background loop to stop "
                    "(AsyncEngineDeadError).")
Antoni Baum's avatar
Antoni Baum committed
540

541
542
        if arrival_time is None:
            arrival_time = time.time()
543
544

        if self.engine_use_ray:
545
546
547
548
549
550
            prompt_token_ids = await (
                self.engine.encode_request_async.remote(  # type: ignore
                    request_id=request_id,
                    prompt=prompt,
                    prompt_token_ids=prompt_token_ids,
                    lora_request=lora_request))
551
552
553
554
555
556
        else:
            prompt_token_ids = await self.engine.encode_request_async(
                request_id=request_id,
                prompt=prompt,
                prompt_token_ids=prompt_token_ids,
                lora_request=lora_request)
557

558
        stream = self._request_tracker.add_request(
559
560
561
562
            request_id,
            prompt=prompt,
            sampling_params=sampling_params,
            prompt_token_ids=prompt_token_ids,
563
            arrival_time=arrival_time,
564
565
566
            lora_request=lora_request,
            multi_modal_data=multi_modal_data,
        )
Antoni Baum's avatar
Antoni Baum committed
567
568

        return stream
569

570
    async def generate(
571
572
573
574
        self,
        prompt: Optional[str],
        sampling_params: SamplingParams,
        request_id: str,
575
        prompt_token_ids: Optional[List[int]] = None,
576
        lora_request: Optional[LoRARequest] = None,
577
        multi_modal_data: Optional[MultiModalData] = None
578
    ) -> AsyncIterator[RequestOutput]:
579
580
581
        """Generate outputs for a request.

        Generate outputs for a request. This method is a coroutine. It adds the
582
583
        request into the waiting queue of the LLMEngine and streams the outputs
        from the LLMEngine to the caller.
584
585
586
587
588
589
590
591

        Args:
            prompt: The prompt string. Can be None if prompt_token_ids is
                provided.
            sampling_params: The sampling parameters of the request.
            request_id: The unique id of the request.
            prompt_token_ids: The token IDs of the prompt. If None, we
                use the tokenizer to convert the prompts to token IDs.
592
            lora_request: LoRA request to use for generation, if any.
593
            multi_modal_data: Multi modal data per request.
594
595

        Yields:
596
            The output `RequestOutput` objects from the LLMEngine for the
597
            request.
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640

        Details:
            - If the engine is not running, start the background loop,
              which iteratively invokes
              :meth:`~vllm.engine.async_llm_engine.AsyncLLMEngine.engine_step`
              to process the waiting requests.
            - Add the request to the engine's `RequestTracker`.
              On the next background loop, this request will be sent to
              the underlying engine.
              Also, a corresponding `AsyncStream` will be created.
            - Wait for the request outputs from `AsyncStream` and yield them.

        Example:
            >>> # Please refer to entrypoints/api_server.py for
            >>> # the complete example.
            >>>
            >>> # initialize the engine and the example input
            >>> engine = AsyncLLMEngine.from_engine_args(engine_args)
            >>> example_input = {
            >>>     "prompt": "What is LLM?",
            >>>     "stream": False, # assume the non-streaming case
            >>>     "temperature": 0.0,
            >>>     "request_id": 0,
            >>> }
            >>>
            >>> # start the generation
            >>> results_generator = engine.generate(
            >>>    example_input["prompt"],
            >>>    SamplingParams(temperature=example_input["temperature"]),
            >>>    example_input["request_id"])
            >>>
            >>> # get the results
            >>> final_output = None
            >>> async for request_output in results_generator:
            >>>     if await request.is_disconnected():
            >>>         # Abort the request if the client disconnects.
            >>>         await engine.abort(request_id)
            >>>         # Return or raise an error
            >>>         ...
            >>>     final_output = request_output
            >>>
            >>> # Process and return the final output
            >>> ...
641
        """
642
        # Preprocess the request.
643
        arrival_time = time.time()
644

Antoni Baum's avatar
Antoni Baum committed
645
        try:
646
647
648
649
650
651
652
            stream = await self.add_request(
                request_id,
                prompt,
                sampling_params,
                prompt_token_ids=prompt_token_ids,
                arrival_time=arrival_time,
                lora_request=lora_request,
653
                multi_modal_data=multi_modal_data,
654
            )
655

Antoni Baum's avatar
Antoni Baum committed
656
657
            async for request_output in stream:
                yield request_output
658
659
660
        except (Exception, asyncio.CancelledError) as e:
            # If there is an exception or coroutine is cancelled, abort the
            # request.
Antoni Baum's avatar
Antoni Baum committed
661
662
            self._abort(request_id)
            raise e
663

Antoni Baum's avatar
Antoni Baum committed
664
665
    async def abort(self, request_id: str) -> None:
        """Abort a request.
666

Antoni Baum's avatar
Antoni Baum committed
667
668
        Abort a submitted request. If the request is finished or not found,
        this method will be a no-op.
669

Antoni Baum's avatar
Antoni Baum committed
670
671
672
        Args:
            request_id: The unique id of the request.
        """
673
674
675
676
677
678
679
        if not self.is_running:
            raise AsyncEngineDeadError(
                "Background loop is not running. If it was running, "
                "inspect the output to find the stacktrace of the "
                "error that caused the background loop to stop "
                "(AsyncEngineDeadError).")

Antoni Baum's avatar
Antoni Baum committed
680
        return self._abort(request_id)
681

Antoni Baum's avatar
Antoni Baum committed
682
    def _abort(self, request_id: str) -> None:
683
684
685
686
687
688
689
690
        """Abort a request.

        Abort a submitted request. If the request is finished or not found,
        this method will be a no-op.

        Args:
            request_id: The unique id of the request.
        """
691
692
        self._request_tracker.abort_request(request_id,
                                            verbose=self.log_requests)
693

694
695
696
    async def get_model_config(self) -> ModelConfig:
        """Get the model configuration of the vLLM engine."""
        if self.engine_use_ray:
697
            return await self.engine.get_model_config.remote()  # type: ignore
698
699
700
        else:
            return self.engine.get_model_config()

701
702
703
704
705
706
707
708
    async def get_decoding_config(self) -> DecodingConfig:
        """Get the decoding configuration of the vLLM engine."""
        if self.engine_use_ray:
            return await self.engine.get_decoding_config.remote(  # type: ignore
            )
        else:
            return self.engine.get_decoding_config()

709
710
711
712
    async def do_log_stats(
            self,
            scheduler_outputs: Optional[SchedulerOutputs] = None,
            model_output: Optional[List[SamplerOutput]] = None) -> None:
713
        if self.engine_use_ray:
714
715
            await self.engine.do_log_stats.remote(  # type: ignore
                scheduler_outputs, model_output)
716
717
        else:
            self.engine.do_log_stats()
718

719
    async def check_health(self) -> None:
720
721
722
723
724
725
726
727
        """Raises an error if engine is unhealthy."""
        t = time.perf_counter()
        logger.debug("Starting health check...")
        if self.is_stopped:
            raise AsyncEngineDeadError("Background loop is stopped.")

        if self.engine_use_ray:
            try:
728
                await self.engine.check_health.remote()  # type: ignore
729
730
731
732
            except ray.exceptions.RayActorError as e:
                raise RuntimeError("Engine is dead.") from e
        else:
            await self.engine.check_health_async()
733
        logger.debug("Health check took %fs", time.perf_counter() - t)