client.py 29.8 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
4
5
6
import asyncio
import copy
import pickle
from contextlib import contextmanager, suppress
7
from typing import (Any, AsyncGenerator, Dict, Iterator, List, Mapping,
8
                    Optional, Union, cast, overload)
9
10

import cloudpickle
11
import psutil
12
13
import zmq
import zmq.asyncio
14
from typing_extensions import deprecated
15
16
17
from zmq import Frame  # type: ignore[attr-defined]
from zmq.asyncio import Socket

18
from vllm import PoolingParams
19
from vllm.config import DecodingConfig, ModelConfig, VllmConfig
20
from vllm.core.scheduler import SchedulerOutputs
21
22
# yapf conflicts with isort for this block
# yapf: disable
23
24
from vllm.engine.async_llm_engine import (
    build_guided_decoding_logits_processor_async)
25
26
27
28
from vllm.engine.multiprocessing import (ENGINE_DEAD_ERROR, IPC_DATA_EXT,
                                         IPC_HEALTH_EXT, IPC_INPUT_EXT,
                                         IPC_OUTPUT_EXT, RPC_REQUEST_T,
                                         VLLM_RPC_SUCCESS_STR, RPCAbortRequest,
29
                                         RPCAdapterLoadedResponse, RPCError,
30
31
                                         RPCIsSleepingRequest,
                                         RPCIsSleepingResponse,
32
                                         RPCLoadAdapterRequest,
33
                                         RPCProcessRequest,
34
                                         RPCResetMultiModalCacheRequest,
35
                                         RPCResetPrefixCacheRequest,
36
37
38
                                         RPCSleepRequest, RPCStartupRequest,
                                         RPCStartupResponse,
                                         RPCUProfileRequest, RPCWakeUpRequest)
39
from vllm.engine.protocol import EngineClient
40
41
# yapf: enable
from vllm.envs import VLLM_RPC_TIMEOUT
42
from vllm.inputs import PromptType
43
from vllm.inputs.preprocess import InputPreprocessor
44
45
from vllm.logger import init_logger
from vllm.lora.request import LoRARequest
46
from vllm.model_executor.layers.sampler import SamplerOutput
47
from vllm.outputs import PoolingRequestOutput, RequestOutput
48
from vllm.prompt_adapter.request import PromptAdapterRequest
49
from vllm.sampling_params import SamplingParams
50
from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
51
from vllm.utils import Device, deprecate_kwargs
52
53
54
55
56
57

logger = init_logger(__name__)


class MQClientClosedError(Exception):
    """Exception class raised when the client is used post-close.
58

59
    The client can be closed, which closes the ZMQ context. This normally
60
61
    happens on server shutdown. In some cases, methods like abort and
    do_log_stats will still be called and then try to open a socket, which
62
63
64
65
66
    causes a ZMQError and creates a huge stack trace.
    So, we throw this error such that we can suppress it.
    """


67
class MQLLMEngineClient(EngineClient):
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
    """A client wrapper for MQLLMEngine that conforms to the
    EngineClient protocol.

    MQLLMEngine and MQLLMEngineClient are intended to run in separate
    processes communicating via zeromq ipc sockets.

    The entrypoint to MQLLMEngineClient is through the generate()
    method. On generate() MQLLMEngine does three things:
        - Creates an asyncio output queue
        - Sends a RPCGenerateRequest to the MQLLMEngine via zmq
        - Pulls RequestOutputs from its queue and yields them

    MQLLMEngine runs two background loops:
        - output_loop: the output loop pulls List[RequestOutput]
            from the MQLLMEngine via zmq (each list is the output
            of one engine_step in the LLMEngine). It then parses
            the list and pushes individual request_outputs into
            the corresponding output_queue such that they can be
            consumed by the .generate() method.
        - health_loop: the health loop queries the health socket
            every N seconds, confirming the engine is healthy
    """

91
    def __init__(self, ipc_path: str, engine_config: VllmConfig,
92
                 engine_pid: int):
93
94
95
96
        self.context = zmq.asyncio.Context()
        self._errored_with: Optional[BaseException] = None

        # Get the configs.
97
        self.vllm_config = engine_config
98
99
100
101
102
103
104
        self.model_config = engine_config.model_config
        self.decoding_config = engine_config.decoding_config

        # Create the tokenizer group.
        self.tokenizer = init_tokenizer_from_configs(
            model_config=self.model_config,
            scheduler_config=engine_config.scheduler_config,
105
            lora_config=engine_config.lora_config)
106
107
        self.input_preprocessor = InputPreprocessor(self.model_config,
                                                    self.tokenizer)
108
109
110
111
112
113
114
115
116

        # Send RPCGenerateRequest to the MQLLMEngine.
        self.input_socket: Socket = self.context.socket(zmq.constants.PUSH)
        self.input_socket.connect(f"{ipc_path}{IPC_INPUT_EXT}")

        # Receive streams of RequestOutput from the MQLLMEngine.
        self.output_socket: Socket = self.context.socket(zmq.constants.PULL)
        self.output_socket.connect(f"{ipc_path}{IPC_OUTPUT_EXT}")

117
118
119
        # IPC path for acking heartbeats.
        self.heartbeat_socket: Socket = self.context.socket(zmq.constants.PULL)
        self.heartbeat_socket.connect(f"{ipc_path}{IPC_HEALTH_EXT}")
120
121
122
123
124
125

        # IPC path for the data socket.
        self.data_ipc_path = f"{ipc_path}{IPC_DATA_EXT}"

        # Stream for each individual request.
        self.output_queues: Dict[str, asyncio.Queue] = {}
126
127
128
129
130

        # Loop to handle output of the LLMEngine periodically.
        # Started after the MQLLMEngine is ready so that we can
        # build the Client in an executor to enable clean shutdown.
        self.output_loop: Optional[asyncio.Task] = None
131
132
133
134

        # Loop to check health of the LLMEngine periodically.
        # Started after the MQLLMEngine is ready.
        self.health_loop: Optional[asyncio.Task] = None
135
        self._engine_process = psutil.Process(engine_pid)
136
137

    @staticmethod
138
    def is_unsupported_config(vllm_config: VllmConfig):
139
        # Pipeline parallel not yet supported
140
        return vllm_config.parallel_config.pipeline_parallel_size > 1
141
142
143
144
145
146
147
148
149
150

    @contextmanager
    def get_data_socket(self) -> Iterator[Socket]:
        socket = self.context.socket(zmq.constants.DEALER)
        try:
            socket.connect(self.data_ipc_path)
            yield socket
        finally:
            socket.close(linger=0)

151
    async def run_heartbeat_loop(self, timeout: int):
152
153
        """Background loop that continually checks to ensure the engine process
        is still alive.
154
155
156
        """
        try:
            while True:
157
158
159
160
                # Check if the engine process is running:
                if not self._engine_process.is_running() or (
                        self._engine_process.status() == psutil.STATUS_ZOMBIE):
                    # NB: is_running() returns True for zombies
161
                    self._set_errored(
162
163
164
                        RuntimeError(
                            f"Engine process (pid {self._engine_process.pid}) "
                            "died."))
165
166
                    break

167
                if await self.heartbeat_socket.poll(timeout=timeout):
168
                    # Heartbeat received- check the message
169
                    await self._check_success(
170
171
                        error_message="Heartbeat failed.",
                        socket=self.heartbeat_socket)
172

173
                logger.debug("Heartbeat successful.")
174
175
176
177

        except asyncio.CancelledError:
            logger.debug("Shutting down MQLLMEngineClient check health loop.")

178
179
180
181
182
        except psutil.NoSuchProcess:
            self._set_errored(
                RuntimeError(
                    f"Engine process (pid {self._engine_process.pid}) died."))

183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
        except Exception as e:
            self._set_errored(e)

    async def run_output_handler_loop(self):
        """Get RequestOutputs from Engine and stream to Request Queues"""

        try:
            while True:
                # Poll, checking for ENGINE_DEAD
                while await self.output_socket.poll(timeout=VLLM_RPC_TIMEOUT
                                                    ) == 0:
                    logger.debug("Waiting for output from MQLLMEngine.")

                    # If errored, alert all running requests.
                    if self.errored:
                        for queue_j in tuple(self.output_queues.values()):
                            queue_j.put_nowait(
                                ENGINE_DEAD_ERROR(self._errored_with))
                        return

                message: Frame = await self.output_socket.recv(copy=False)
                request_outputs = pickle.loads(message.buffer)

                is_error = isinstance(request_outputs,
                                      (BaseException, RPCError))
                if is_error:
                    if isinstance(request_outputs, RPCError):
                        rpc_error: RPCError = request_outputs
                        request_id = rpc_error.request_id
                        exception = rpc_error.exception
                        is_engine_errored = rpc_error.is_engine_errored
                    else:
                        # MPLLMEngine should always return an RPCError to
                        # the output_socket when an issue arises.
                        # If we are here, we are in a bad state and
                        # should shut down the server.
                        error: BaseException = request_outputs
                        logger.error(
                            "Received Exception %s rather than RPCError from "
                            "MPLLMEngine. This should never happen.", error)
                        request_id = None
                        exception = error
                        is_engine_errored = True

                    # Set to error state only on engine critical error
                    # (and record only the first one)
                    if is_engine_errored and not self._errored_with:
                        self._errored_with = exception
231
232
233
234
235
236
237
                        # If engine is errored, no matter the type of exception
                        # it will no longer be able to receive new requests,
                        # therefore we have to inform that the current
                        # processed requests failed as well. Send back a dead
                        # engine error give this feedback and also give a
                        # 'hint' to the server to shutdown next.
                        exception = self.dead_error
238
239

                    if request_id is None:
240
241
242
243
244
                        # If request_id is None, then the engine raised an
                        # exception for a batch, and we may not know the
                        # request that caused it, neither if it was actually
                        # caused by any of them (e.g. CUDA OOM). Therefore we
                        # broadcast the same exception for all requests.
245
246
247
248
249
250
                        for queue_i in tuple(self.output_queues.values()):
                            queue_i.put_nowait(exception)
                    else:
                        queue = self.output_queues.get(request_id)
                        if queue is not None:
                            queue.put_nowait(exception)
251
                # Put each output into the appropriate queue.
252
253
254
                elif isinstance(
                        request_outputs,
                    (RPCAdapterLoadedResponse, RPCIsSleepingResponse)):
255
                    self._add_output(request_outputs)
256
257
                else:
                    for request_output in request_outputs:
258
                        self._add_output(request_output)
259
260
261
262

        except asyncio.CancelledError:
            logger.debug("Shutting down MQLLMEngineClient output handler.")

263
    def _add_output(self, request_output: Union[RequestOutput,
264
265
                                                RPCAdapterLoadedResponse,
                                                RPCIsSleepingResponse]):
266
267
268
269
        queue = self.output_queues.get(request_output.request_id)
        if queue is not None:
            queue.put_nowait(request_output)

270
271
272
    async def setup(self):
        """Setup the client before it starts sending server requests."""

273
        # Start output_loop
274
275
276
277
278
279
280
281
        if self.output_loop is None:
            # only generate once to avoid multiple concurrent output_loops
            # this will lead to race conditions and wrong orders of tokens
            # returned by the engine
            # setup will be called multiple times during the startup of
            # the engine
            self.output_loop = asyncio.create_task(
                self.run_output_handler_loop())
282

283
284
285
286
287
288
289
        with self.get_data_socket() as socket:
            # Wait until server is ready.
            response = await self._wait_for_server_rpc(socket)

            self.tracing_flag = response.tracing_enabled

            # Start health_loop.
290
291
292
            if self.health_loop is None:
                self.health_loop = asyncio.create_task(
                    self.run_heartbeat_loop(timeout=VLLM_RPC_TIMEOUT))
293
294
295
296
297
298
299
300
301

    def close(self):
        """Destroy the ZeroMQ Context."""
        # Close all sockets and terminate the context.
        self.context.destroy(linger=0)

        # Cancel background tasks.
        if self.health_loop is not None:
            self.health_loop.cancel()
302
303
        if self.output_loop is not None:
            self.output_loop.cancel()
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374

    def _set_errored(self, e: BaseException):
        logger.exception(repr(e))
        if self._errored_with is None:
            self._errored_with = e

    @staticmethod
    async def _send_get_data_rpc_request(request: RPCStartupRequest,
                                         expected_type: Any,
                                         error_message: str,
                                         socket: Socket) -> Any:
        """Send an RPC request that is expecting data back."""

        # Ping RPCServer with a request.
        await socket.send_multipart((pickle.dumps(request), ), copy=False)

        # Make sure the server responds in time.
        if await socket.poll(timeout=VLLM_RPC_TIMEOUT) == 0:
            raise TimeoutError("RPCServer didn't reply within "
                               f"{VLLM_RPC_TIMEOUT} ms")

        # Await the data from the Server.
        frame = await socket.recv(copy=False)
        data = pickle.loads(frame.buffer)

        if isinstance(data, BaseException):
            raise data
        elif not isinstance(data, expected_type):
            raise ValueError(error_message)

        return data

    @staticmethod
    async def _send_one_way_rpc_request(request: RPC_REQUEST_T,
                                        socket: Socket):
        """Send one-way RPC request to trigger an action."""

        if socket.closed:
            raise MQClientClosedError()

        await socket.send_multipart((pickle.dumps(request), ))

    async def _await_ack(self, error_message: str, socket: Socket):
        """Await acknowledgement that a request succeeded."""

        if socket.closed:
            raise MQClientClosedError()

        if await socket.poll(timeout=VLLM_RPC_TIMEOUT) == 0:
            raise TimeoutError("MQLLMEngine didn't reply within "
                               f"{VLLM_RPC_TIMEOUT}ms")

        await self._check_success(error_message, socket)

    @staticmethod
    async def _check_success(error_message: str, socket: Socket):
        """Confirm that socket has a VLLM_RPC_SUCCESS_STR message"""

        if socket.closed:
            raise MQClientClosedError()

        frame = await socket.recv(copy=False)
        response = pickle.loads(frame.buffer)

        # Raise error if unsuccessful
        if isinstance(response, BaseException):
            raise response
        elif (not isinstance(response, str)
              or response != VLLM_RPC_SUCCESS_STR):
            raise ValueError(error_message)

375
376
377
    async def get_input_preprocessor(self) -> InputPreprocessor:
        return self.input_preprocessor

378
    async def get_tokenizer(self, lora_request: Optional[LoRARequest] = None):
379
380
        return await self.tokenizer.get_lora_tokenizer_async(lora_request)

381
382
383
    async def get_vllm_config(self) -> VllmConfig:
        return self.vllm_config

384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
    async def get_decoding_config(self) -> DecodingConfig:
        return self.decoding_config

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

    async def is_tracing_enabled(self) -> bool:
        return self.tracing_flag

    async def _wait_for_server_rpc(self, socket: Socket) -> RPCStartupResponse:
        """Wait for the RPCServer to start up."""

        return await self._send_get_data_rpc_request(
            request=RPCStartupRequest.IS_SERVER_READY,
            expected_type=RPCStartupResponse,
            error_message="Unable to start RPC Server",
            socket=socket)

    async def abort(self, request_id: str):
        """Send an ABORT_REQUEST signal to the RPC Server"""

        with suppress(MQClientClosedError):
            await self._send_one_way_rpc_request(
                request=RPCAbortRequest(request_id), socket=self.input_socket)

409
410
411
412
413
414
415
416
    async def do_log_stats(
        self,
        scheduler_outputs: Optional[SchedulerOutputs] = None,
        model_output: Optional[List[SamplerOutput]] = None,
    ) -> None:
        """
        Ignore do_log_stats (handled on MQLLMEngine polling)
        """
417
418
419
420
421
        pass

    async def check_health(self):
        """
        The check health loop probes the health status of the
422
        Engine's health every N seconds and sets _errored_with
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
        if the engine is unhealthy.
        """
        if self._errored_with is not None:
            raise self._errored_with

    @property
    def is_running(self) -> bool:
        return not self.errored

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

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

440
441
    @property
    def dead_error(self) -> BaseException:
442
        return ENGINE_DEAD_ERROR(self._errored_with)
443

444
    @overload
445
    def generate(
446
        self,
447
        prompt: PromptType,
448
449
450
451
        sampling_params: SamplingParams,
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
452
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
453
        priority: int = 0,
454
455
456
457
    ) -> AsyncGenerator[RequestOutput, None]:
        ...

    @overload
458
    @deprecated("'inputs' will be renamed to 'prompt")
459
460
    def generate(
        self,
461
462
        *,
        inputs: PromptType,
463
464
465
466
467
        sampling_params: SamplingParams,
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
468
        priority: int = 0,
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
    ) -> AsyncGenerator[RequestOutput, None]:
        ...

    @deprecate_kwargs(
        "inputs",
        additional_message="Please use the 'prompt' parameter instead.",
    )
    def generate(
        self,
        prompt: Optional[PromptType] = None,
        sampling_params: Optional[SamplingParams] = None,
        request_id: Optional[str] = None,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
484
        priority: int = 0,
485
486
        *,
        inputs: Optional[PromptType] = None  # DEPRECATED
487
    ) -> AsyncGenerator[RequestOutput, None]:
488
489
490
491
492
493
494
        """Generate outputs for a request.

        Generate outputs for a request. This method is a coroutine. It adds the
        request into the waiting queue of the LLMEngine and streams the outputs
        from the LLMEngine to the caller.

        Args:
495
            prompt: The prompt to the LLM. See {class}`~vllm.inputs.PromptType`
496
497
498
499
500
501
502
                for more details about the format of each input.
            sampling_params: The sampling parameters of the request.
            request_id: The unique id of the request.
            lora_request: LoRA request to use for generation, if any.
            trace_headers: OpenTelemetry trace headers.
            prompt_adapter_request: Prompt Adapter request to use
                                            for generation, if any.
503
504
            priority: Priority of the request (lower means earlier handling).
                Any priority other than 0 will lead to an error if the
505
                scheduling policy is not "priority".
506
        """
507
508
509
510
511
512
        if inputs is not None:
            prompt = inputs
        assert (prompt is not None and sampling_params is not None
                and request_id is not None)

        return self._process_request(prompt, sampling_params, request_id,
513
                                     lora_request, trace_headers,
514
                                     prompt_adapter_request, priority)
515

516
    @overload
517
518
    def encode(
        self,
519
        prompt: PromptType,
520
521
522
523
        pooling_params: PoolingParams,
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
524
        priority: int = 0,
525
    ) -> AsyncGenerator[PoolingRequestOutput, None]:
526
527
528
        ...

    @overload
529
    @deprecated("'inputs' will be renamed to 'prompt")
530
531
    def encode(
        self,
532
533
        *,
        inputs: PromptType,
534
535
536
537
        pooling_params: PoolingParams,
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
538
        priority: int = 0,
539
    ) -> AsyncGenerator[PoolingRequestOutput, None]:
540
541
542
543
544
545
546
547
548
549
550
551
552
        ...

    @deprecate_kwargs(
        "inputs",
        additional_message="Please use the 'prompt' parameter instead.",
    )
    def encode(
        self,
        prompt: Optional[PromptType] = None,
        pooling_params: Optional[PoolingParams] = None,
        request_id: Optional[str] = None,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
553
        priority: int = 0,
554
555
        *,
        inputs: Optional[PromptType] = None  # DEPRECATED
556
    ) -> AsyncGenerator[PoolingRequestOutput, None]:
557
        """Generate outputs for a request from a pooling model.
558
559
560
561
562
563

        Generate outputs for a request. This method is a coroutine. It adds the
        request into the waiting queue of the LLMEngine and streams the outputs
        from the LLMEngine to the caller.

        Args:
564
            prompt: The prompt to the LLM. See {class}`~vllm.inputs.PromptType`
565
566
567
568
569
570
571
                for more details about the format of each input.
            pooling_params: The pooling parameters of the request.
            request_id: The unique id of the request.
            lora_request: LoRA request to use for generation, if any.
            trace_headers: OpenTelemetry trace headers.

        Yields:
572
            The output `PoolingRequestOutput` objects from the LLMEngine
573
574
            for the request.
        """
575
576
577
578
579
        if inputs is not None:
            prompt = inputs
        assert (prompt is not None and pooling_params is not None
                and request_id is not None)

580
        return cast(
581
            AsyncGenerator[PoolingRequestOutput, None],
582
583
584
585
586
587
            self._process_request(prompt,
                                  pooling_params,
                                  request_id,
                                  lora_request,
                                  trace_headers,
                                  priority=priority))
588
589
590

    async def _process_request(
        self,
591
        prompt: PromptType,
592
593
594
595
        params: Union[SamplingParams, PoolingParams],
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
596
597
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
598
    ) -> Union[AsyncGenerator[RequestOutput, None], AsyncGenerator[
599
            PoolingRequestOutput, None]]:
600
601
602
603
604
605
        """Send an RPCGenerateRequest to the RPCServer and stream responses."""

        # If already dead, error out.
        if self._errored_with is not None:
            raise ENGINE_DEAD_ERROR(self._errored_with)

606
607
608
609
        # Ensure the request id is unique among running requests
        if request_id in self.output_queues:
            raise ValueError(f"Request {request_id} already exists")

610
611
612
613
614
615
616
617
618
        # Constructing guided decoding logits processors is expensive, so we do
        # it here to avoid contending with cpu resources and the GIL on the
        # backend process.
        if isinstance(params, SamplingParams) and \
            params.guided_decoding is not None:
            params = await \
                build_guided_decoding_logits_processor_async(
                    sampling_params=params,
                    tokenizer=await self.get_tokenizer(lora_request),
619
                    default_guided_backend=(self.decoding_config.backend
620
                        if self.decoding_config
621
                        else DecodingConfig.backend),
622
623
                    model_config=self.model_config,
                    reasoning_backend=self.decoding_config.reasoning_backend,
624
625
                )

626
627
628
629
630
631
632
633
        # 1) Create output queue for this requests.
        queue: asyncio.Queue[Union[RequestOutput,
                                   BaseException]] = asyncio.Queue()
        self.output_queues[request_id] = queue

        try:
            # 2) Detach logits processors so that they can be pickled
            # separately (may require cloudpickle which is slower)
634
            if isinstance(params, SamplingParams) and params.logits_processors:
635
                # Defensive shallow copy
636
637
638
                params = copy.copy(params)
                logits_processors = params.logits_processors
                params.logits_processors = None
639
640
641
642
643
                lp_bytes = cloudpickle.dumps(logits_processors)
            else:
                lp_bytes = None

            request_bytes = pickle.dumps(
644
                RPCProcessRequest(
645
                    prompt=prompt,
646
                    params=params,
647
648
649
                    request_id=request_id,
                    lora_request=lora_request,
                    trace_headers=trace_headers,
650
651
652
                    prompt_adapter_request=prompt_adapter_request,
                    priority=priority,
                ))
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677

            # 3) Send the RPCGenerateRequest to the MQLLMEngine.
            parts = (request_bytes,
                     lp_bytes) if lp_bytes else (request_bytes, )
            await self.input_socket.send_multipart(parts, copy=False)

            # 4) Stream the RequestOutputs from the output queue. Note
            # that the output_loop pushes RequestOutput objects to this
            # queue after pulling them from the zmq socket.
            finished = False
            try:
                while not finished:
                    request_output = await queue.get()

                    if isinstance(request_output, BaseException):
                        raise request_output

                    finished = request_output.finished
                    yield request_output
            finally:
                # Request was canceled by the client.
                if not finished and not self.errored:
                    await self.abort(request_id)
        finally:
            self.output_queues.pop(request_id)
678
679
680
681
682
683
684
685
686
687
688
689

    async def start_profile(self) -> None:
        """Start profiling the engine"""

        await self._send_one_way_rpc_request(
            request=RPCUProfileRequest.START_PROFILE, socket=self.input_socket)

    async def stop_profile(self) -> None:
        """Stop profiling the engine"""

        await self._send_one_way_rpc_request(
            request=RPCUProfileRequest.STOP_PROFILE, socket=self.input_socket)
690

691
692
693
694
695
696
697
    async def reset_mm_cache(self) -> None:
        """Reset the multi-modal cache"""

        await self._send_one_way_rpc_request(
            request=RPCResetMultiModalCacheRequest.RESET,
            socket=self.input_socket)

698
699
    async def reset_prefix_cache(self,
                                 device: Optional[Device] = None) -> None:
700
701
702
        """Reset the prefix cache"""

        await self._send_one_way_rpc_request(
703
            request=RPCResetPrefixCacheRequest(device),
704
705
            socket=self.input_socket)

706
707
708
709
710
    async def sleep(self, level: int = 1) -> None:
        """Sleep the engine for a given level"""
        return await self._send_one_way_rpc_request(
            request=RPCSleepRequest(level), socket=self.input_socket)

711
    async def wake_up(self, tags: Optional[list[str]] = None) -> None:
712
713
        """Wake up the engine"""
        return await self._send_one_way_rpc_request(
714
            request=RPCWakeUpRequest(tags), socket=self.input_socket)
715

716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
    async def is_sleeping(self) -> bool:
        """Check whether the engine is sleeping"""
        request = RPCIsSleepingRequest()

        queue: asyncio.Queue[Union[BaseException,
                                   RPCIsSleepingResponse]] = asyncio.Queue()
        self.output_queues[request.request_id] = queue

        request_bytes = pickle.dumps(request)
        await self.input_socket.send_multipart((request_bytes, ), copy=False)

        request_output = await queue.get()
        self.output_queues.pop(request.request_id)

        if isinstance(request_output, BaseException):
            raise request_output
        return request_output.is_sleeping

734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
    async def add_lora(self, lora_request: LoRARequest) -> None:
        """Load a new LoRA adapter into the engine for future requests."""
        # Uses the same I/O as generate requests
        request = RPCLoadAdapterRequest(lora_request)

        # Create output queue for this requests.
        queue: asyncio.Queue[Union[None, BaseException]] = asyncio.Queue()
        self.output_queues[request.request_id] = queue

        # Send the request
        request_bytes = pickle.dumps(request)
        await self.input_socket.send_multipart((request_bytes, ), copy=False)

        # Wait for the response
        request_output = await queue.get()
        self.output_queues.pop(request.request_id)

        # Raise on error, otherwise happily return None
        if isinstance(request_output, BaseException):
            raise request_output