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

8
9
from transformers import PreTrainedTokenizer

10
from vllm.lora.request import LoRARequest
11
from vllm.config import ModelConfig
Woosuk Kwon's avatar
Woosuk Kwon committed
12
13
14
15
16
17
from vllm.engine.arg_utils import AsyncEngineArgs
from vllm.engine.llm_engine import LLMEngine
from vllm.engine.ray_utils import initialize_cluster, ray
from vllm.logger import init_logger
from vllm.outputs import RequestOutput
from vllm.sampling_params import SamplingParams
18
19

logger = init_logger(__name__)
20
21
ENGINE_ITERATION_TIMEOUT_S = int(
    os.environ.get("VLLM_ENGINE_ITERATION_TIMEOUT_S", "60"))
22

Antoni Baum's avatar
Antoni Baum committed
23

24
25
26
27
class AsyncEngineDeadError(RuntimeError):
    pass


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

    exception = None
35
    try:
36
37
        task.result()
        # NOTE: This will be thrown if task exits normally (which it should not)
38
        raise AsyncEngineDeadError(msg)
39
40
41
42
43
44
    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
45
46


Antoni Baum's avatar
Antoni Baum committed
47
48
49
50
51
52
53
54
55
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
        self._queue = asyncio.Queue()
        self._finished = False

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

    def finish(self) -> None:
62
        self._queue.put_nowait(StopAsyncIteration())
Antoni Baum's avatar
Antoni Baum committed
63
64
65
66
67
68
69
70
71
72
73
        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()
74
        if isinstance(result, Exception):
75
            raise result
Antoni Baum's avatar
Antoni Baum committed
76
77
78
        return result


79
80
81
82
83
84
85
86
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()
87
        self.new_requests_event = asyncio.Event()
88
89
90
91

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

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

    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)
102
            self.abort_request(request_id)
103
        else:
104
            for rid, stream in self._request_streams.items():
105
                stream.put(exc)
106
                self.abort_request(rid)
107
108
109
110
111
112
113
114
115
116
117
118
119
120

    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:
                logger.info(f"Finished request {request_id}.")
            self.abort_request(request_id)

121
122
123
124
125
126
127
128
129
130
131
    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:
            logger.info(f"Finished request {request_id}.")
        self.abort_request(request_id)

132
133
134
135
136
137
138
139
140
141
142
143
    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
        }))
144
145
146

        self.new_requests_event.set()

147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
        return stream

    def abort_request(self, request_id: str, *, verbose: bool = False) -> None:
        """Abort a request during next background loop iteration."""
        if verbose:
            logger.info(f"Aborted request {request_id}.")

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

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

185
    async def wait_for_new_requests(self):
186
187
188
189
190
191
        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()
192

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

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.
        """
207
        seq_group_metadata_list, scheduler_outputs = self.scheduler.schedule()
Antoni Baum's avatar
Antoni Baum committed
208

209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
        if not scheduler_outputs.is_empty():
            # Execute the model.
            all_outputs = await self._run_workers_async(
                "execute_model",
                driver_kwargs={
                    "seq_group_metadata_list": seq_group_metadata_list,
                    "blocks_to_swap_in": scheduler_outputs.blocks_to_swap_in,
                    "blocks_to_swap_out": scheduler_outputs.blocks_to_swap_out,
                    "blocks_to_copy": scheduler_outputs.blocks_to_copy,
                })

            # Only the driver worker returns the sampling results.
            output = all_outputs[0]
        else:
            output = []
Antoni Baum's avatar
Antoni Baum committed
224

225
        return self._process_model_outputs(output, scheduler_outputs)
Antoni Baum's avatar
Antoni Baum committed
226

227
228
229
230
231
232
233
234
235
236
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
264
265
266
267
268
269
270
    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,
    ) -> 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)

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

Antoni Baum's avatar
Antoni Baum committed
271
272
273
274
    async def _run_workers_async(
        self,
        method: str,
        *args,
275
276
        driver_args: Optional[List[Any]] = None,
        driver_kwargs: Optional[Dict[str, Any]] = None,
Antoni Baum's avatar
Antoni Baum committed
277
278
279
        **kwargs,
    ) -> Any:
        """Runs the given method on all workers."""
280
        coros = []
Antoni Baum's avatar
Antoni Baum committed
281

282
283
284
285
        if driver_args is None:
            driver_args = args
        if driver_kwargs is None:
            driver_kwargs = kwargs
Antoni Baum's avatar
Antoni Baum committed
286

287
288
289
290
        # Run the driver worker asynchronously.
        driver_executor = getattr(self.driver_worker, method)
        coros.append(asyncio.get_event_loop().run_in_executor(
            None, partial(driver_executor, *driver_args, **driver_kwargs)))
Antoni Baum's avatar
Antoni Baum committed
291

292
293
294
295
296
297
        # Run the ray workers asynchronously.
        for worker in self.workers:
            coros.append(worker.execute_method.remote(method, *args, **kwargs))

        all_outputs = await asyncio.gather(*coros)
        return all_outputs
298

299
300
301
302
    async def check_health_async(self):
        """Raises an error if engine is unhealthy."""
        self._check_if_any_actor_is_dead()

303

304
305
class AsyncLLMEngine:
    """An asynchronous wrapper for LLMEngine.
306

307
    This class is used to wrap the LLMEngine class to make it asynchronous. It
308
    uses asyncio to create a background loop that keeps processing incoming
309
    requests. The LLMEngine is kicked by the generate method when there
310
    are requests in the waiting queue. The generate method yields the outputs
311
    from the LLMEngine to the caller.
312

313
    NOTE: For the comprehensive list of arguments, see `LLMEngine`.
314
315
316
317
318

    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
319
        engine_use_ray: Whether to make LLMEngine a Ray actor. If so, the
320
321
            async frontend will be executed in a separate process as the
            model workers.
322
        log_requests: Whether to log the requests.
zspo's avatar
zspo committed
323
324
        max_log_len: Maximum number of prompt characters or prompt ID numbers
            being printed in log.
325
326
        start_engine_loop: If True, the background task to run the engine
            will be automatically started in the generate call.
327
328
        *args: Arguments for LLMEngine.
        *kwargs: Arguments for LLMEngine.
329
    """
330

Antoni Baum's avatar
Antoni Baum committed
331
332
    _engine_class: Type[_AsyncLLMEngine] = _AsyncLLMEngine

333
334
335
336
337
    def __init__(self,
                 worker_use_ray: bool,
                 engine_use_ray: bool,
                 *args,
                 log_requests: bool = True,
338
                 max_log_len: Optional[int] = None,
339
                 start_engine_loop: bool = True,
340
                 **kwargs) -> None:
341
        self.worker_use_ray = worker_use_ray
Zhuohan Li's avatar
Zhuohan Li committed
342
        self.engine_use_ray = engine_use_ray
343
        self.log_requests = log_requests
344
        self.max_log_len = max_log_len
Antoni Baum's avatar
Antoni Baum committed
345
346
347
        self.engine = self._init_engine(*args, **kwargs)

        self.background_loop = None
348
349
350
351
        # We need to keep a reference to unshielded
        # task as well to prevent it from being garbage
        # collected
        self._background_loop_unshielded = None
352
        self.start_engine_loop = start_engine_loop
353
354
        self._request_tracker: Optional[RequestTracker] = None
        self._errored_with: Optional[BaseException] = None
Antoni Baum's avatar
Antoni Baum committed
355

356
357
    @property
    def is_running(self) -> bool:
358
        return (self.background_loop is not None
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
                and not self._background_loop_unshielded.done())

    @property
    def is_stopped(self) -> bool:
        return self.errored or (self.background_loop is not None
                                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)
376

377
378
379
380
381
    async def get_tokenizer(self) -> "PreTrainedTokenizer":
        if self.engine_use_ray:
            return await self.engine.get_tokenizer.remote()
        else:
            return self.engine.get_tokenizer()
382

383
    def start_background_loop(self) -> None:
Antoni Baum's avatar
Antoni Baum committed
384
        """Start the background loop."""
385
386
387
        if self.errored:
            raise AsyncEngineDeadError(
                "Background loop has errored already.") from self._errored_with
388
        if self.is_running:
Antoni Baum's avatar
Antoni Baum committed
389
            raise RuntimeError("Background loop is already running.")
390
391
        # Initialize the RequestTracker here so it uses the right event loop.
        self._request_tracker = RequestTracker()
392
393
394
395

        self._background_loop_unshielded = asyncio.get_event_loop(
        ).create_task(self.run_engine_loop())
        self._background_loop_unshielded.add_done_callback(
396
            partial(_raise_exception_on_finish,
397
                    error_callback=self._error_callback))
398
        self.background_loop = asyncio.shield(self._background_loop_unshielded)
Antoni Baum's avatar
Antoni Baum committed
399
400
401

    def _init_engine(self, *args,
                     **kwargs) -> Union[_AsyncLLMEngine, "ray.ObjectRef"]:
Zhuohan Li's avatar
Zhuohan Li committed
402
        if not self.engine_use_ray:
Antoni Baum's avatar
Antoni Baum committed
403
            engine_class = self._engine_class
404
        elif self.worker_use_ray:
Antoni Baum's avatar
Antoni Baum committed
405
            engine_class = ray.remote(num_cpus=0)(self._engine_class).remote
406
        else:
Woosuk Kwon's avatar
Woosuk Kwon committed
407
408
409
410
411
412
413
414
415
416
            # FIXME(woosuk): This is a bit hacky. Be careful when changing the
            # order of the arguments.
            cache_config = args[1]
            parallel_config = args[2]
            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
417
418
        return engine_class(*args, **kwargs)

419
420
421
422
    async def engine_step(self) -> bool:
        """Kick the engine to process the waiting requests.

        Returns True if there are in-progress requests."""
423
424

        new_requests, finished_requests = (
425
            self._request_tracker.get_new_and_finished_requests())
426
427
428
429

        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
430
431
432
433
434
435
436
437
438
439
440
441
            try:
                if self.engine_use_ray:
                    await self.engine.add_request.remote(**new_request)
                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,
                )
442
443
444
445

        if finished_requests:
            await self._engine_abort(finished_requests)

Zhuohan Li's avatar
Zhuohan Li committed
446
447
        if self.engine_use_ray:
            request_outputs = await self.engine.step.remote()
448
        else:
Antoni Baum's avatar
Antoni Baum committed
449
            request_outputs = await self.engine.step_async()
450

Antoni Baum's avatar
Antoni Baum committed
451
        # Put the outputs into the corresponding streams.
452
        for request_output in request_outputs:
453
            self._request_tracker.process_request_output(
454
                request_output, verbose=self.log_requests)
Antoni Baum's avatar
Antoni Baum committed
455

456
457
        return len(request_outputs) > 0

Antoni Baum's avatar
Antoni Baum committed
458
459
460
461
462
463
464
    async def _engine_abort(self, request_ids: Iterable[str]):
        if self.engine_use_ray:
            await self.engine.abort_request.remote(request_ids)
        else:
            self.engine.abort_request(request_ids)

    async def run_engine_loop(self):
465
        has_requests_in_progress = False
Antoni Baum's avatar
Antoni Baum committed
466
        while True:
467
            if not has_requests_in_progress:
468
                logger.debug("Waiting for new requests...")
469
                await self._request_tracker.wait_for_new_requests()
470
471
472
473
474
475
476
477
478
479
480
481
                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
482
483
484
485
486
487
488
489
490
            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,
491
        lora_request: Optional[LoRARequest] = None,
Antoni Baum's avatar
Antoni Baum committed
492
493
    ) -> AsyncStream:
        if self.log_requests:
494
495
496
497
498
499
500
501
            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]
Antoni Baum's avatar
Antoni Baum committed
502
            logger.info(f"Received request {request_id}: "
503
                        f"prompt: {shortened_prompt!r}, "
zspo's avatar
zspo committed
504
505
                        f"sampling_params: {sampling_params}, "
                        f"prompt_token_ids: {shortened_token_ids}, "
506
                        f"lora_request: {lora_request}.")
Antoni Baum's avatar
Antoni Baum committed
507

508
        if not self.is_running:
509
510
511
512
513
514
515
516
            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
517

518
519
        if arrival_time is None:
            arrival_time = time.time()
520
521
522
523
524
525
526
527
528
529
530
531
532

        if self.engine_use_ray:
            prompt_token_ids = await self.engine.encode_request_async.remote(
                request_id=request_id,
                prompt=prompt,
                prompt_token_ids=prompt_token_ids,
                lora_request=lora_request)
        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)
533

534
        stream = self._request_tracker.add_request(
535
536
537
538
            request_id,
            prompt=prompt,
            sampling_params=sampling_params,
            prompt_token_ids=prompt_token_ids,
539
            arrival_time=arrival_time,
540
            lora_request=lora_request)
Antoni Baum's avatar
Antoni Baum committed
541
542

        return stream
543

544
    async def generate(
545
546
547
548
        self,
        prompt: Optional[str],
        sampling_params: SamplingParams,
        request_id: str,
549
        prompt_token_ids: Optional[List[int]] = None,
550
        lora_request: Optional[LoRARequest] = None,
551
    ) -> AsyncIterator[RequestOutput]:
552
553
554
        """Generate outputs for a request.

        Generate outputs for a request. This method is a coroutine. It adds the
555
556
        request into the waiting queue of the LLMEngine and streams the outputs
        from the LLMEngine to the caller.
557
558
559
560
561
562
563
564

        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.
565
            lora_request: LoRA request to use for generation, if any.
566
567

        Yields:
568
            The output `RequestOutput` objects from the LLMEngine for the
569
            request.
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612

        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
            >>> ...
613
        """
614
        # Preprocess the request.
615
616
        # This should not be used for logging, as it is monotonic time.
        arrival_time = time.monotonic()
617

Antoni Baum's avatar
Antoni Baum committed
618
        try:
619
620
621
622
623
624
625
626
            stream = await self.add_request(
                request_id,
                prompt,
                sampling_params,
                prompt_token_ids=prompt_token_ids,
                arrival_time=arrival_time,
                lora_request=lora_request,
            )
627

Antoni Baum's avatar
Antoni Baum committed
628
629
            async for request_output in stream:
                yield request_output
630
631
632
        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
633
634
            self._abort(request_id)
            raise e
635

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

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

Antoni Baum's avatar
Antoni Baum committed
642
643
644
        Args:
            request_id: The unique id of the request.
        """
645
646
647
648
649
650
651
        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
652
        return self._abort(request_id)
653

Antoni Baum's avatar
Antoni Baum committed
654
    def _abort(self, request_id: str) -> None:
655
656
657
658
659
660
661
662
        """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.
        """
663
664
        self._request_tracker.abort_request(request_id,
                                            verbose=self.log_requests)
665

666
667
668
669
670
671
672
    async def get_model_config(self) -> ModelConfig:
        """Get the model configuration of the vLLM engine."""
        if self.engine_use_ray:
            return await self.engine.get_model_config.remote()
        else:
            return self.engine.get_model_config()

Zhuohan Li's avatar
Zhuohan Li committed
673
    @classmethod
674
    def from_engine_args(cls,
675
                         engine_args: AsyncEngineArgs,
676
                         start_engine_loop: bool = True) -> "AsyncLLMEngine":
Zhuohan Li's avatar
Zhuohan Li committed
677
678
679
680
        """Creates an async LLM engine from the engine arguments."""
        # Create the engine configs.
        engine_configs = engine_args.create_engine_configs()
        parallel_config = engine_configs[2]
Zhuohan Li's avatar
Zhuohan Li committed
681
        # Initialize the cluster.
682
683
        placement_group = initialize_cluster(parallel_config,
                                             engine_args.engine_use_ray)
Zhuohan Li's avatar
Zhuohan Li committed
684
        # Create the async LLM engine.
685
        engine = cls(parallel_config.worker_use_ray,
Zhuohan Li's avatar
Zhuohan Li committed
686
687
                     engine_args.engine_use_ray,
                     *engine_configs,
688
                     placement_group,
689
                     log_requests=not engine_args.disable_log_requests,
690
                     log_stats=not engine_args.disable_log_stats,
691
                     max_log_len=engine_args.max_log_len,
692
                     start_engine_loop=start_engine_loop)
Zhuohan Li's avatar
Zhuohan Li committed
693
        return engine
694
695
696
697
698
699

    async def do_log_stats(self) -> None:
        if self.engine_use_ray:
            await self.engine.do_log_stats.remote()
        else:
            self.engine.do_log_stats()
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715

    async def check_health(self):
        """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:
                await self.engine.check_health.remote()
            except ray.exceptions.RayActorError as e:
                raise RuntimeError("Engine is dead.") from e
        else:
            await self.engine.check_health_async()
        logger.debug(f"Health check took {time.perf_counter()-t}s")