client.py 29.6 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
34
                                         RPCProcessRequest,
                                         RPCResetPrefixCacheRequest,
35
36
37
                                         RPCSleepRequest, RPCStartupRequest,
                                         RPCStartupResponse,
                                         RPCUProfileRequest, RPCWakeUpRequest)
38
from vllm.engine.protocol import EngineClient
39
40
# yapf: enable
from vllm.envs import VLLM_RPC_TIMEOUT
41
from vllm.inputs import PromptType
42
from vllm.inputs.preprocess import InputPreprocessor
43
44
from vllm.logger import init_logger
from vllm.lora.request import LoRARequest
45
from vllm.model_executor.layers.sampler import SamplerOutput
46
from vllm.outputs import PoolingRequestOutput, RequestOutput
47
from vllm.prompt_adapter.request import PromptAdapterRequest
48
from vllm.sampling_params import SamplingParams
49
from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
50
from vllm.utils import Device, deprecate_kwargs
51
52
53
54
55
56

logger = init_logger(__name__)


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

58
    The client can be closed, which closes the ZMQ context. This normally
59
60
    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
61
62
63
64
65
    causes a ZMQError and creates a huge stack trace.
    So, we throw this error such that we can suppress it.
    """


66
class MQLLMEngineClient(EngineClient):
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
    """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
    """

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

        # Get the configs.
96
        self.vllm_config = engine_config
97
98
99
100
101
102
103
        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,
104
            lora_config=engine_config.lora_config)
105
106
        self.input_preprocessor = InputPreprocessor(self.model_config,
                                                    self.tokenizer)
107
108
109
110
111
112
113
114
115

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

116
117
118
        # 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}")
119
120
121
122
123
124

        # 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] = {}
125
126
127
128
129

        # 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
130
131
132
133

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

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

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

150
    async def run_heartbeat_loop(self, timeout: int):
151
152
        """Background loop that continually checks to ensure the engine process
        is still alive.
153
154
155
        """
        try:
            while True:
156
157
158
159
                # 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
160
                    self._set_errored(
161
162
163
                        RuntimeError(
                            f"Engine process (pid {self._engine_process.pid}) "
                            "died."))
164
165
                    break

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

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

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

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

182
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
        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
230
231
232
233
234
235
236
                        # 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
237
238

                    if request_id is None:
239
240
241
242
243
                        # 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.
244
245
246
247
248
249
                        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)
250
                # Put each output into the appropriate queue.
251
252
253
                elif isinstance(
                        request_outputs,
                    (RPCAdapterLoadedResponse, RPCIsSleepingResponse)):
254
                    self._add_output(request_outputs)
255
256
                else:
                    for request_output in request_outputs:
257
                        self._add_output(request_output)
258
259
260
261

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

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

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

272
        # Start output_loop
273
274
275
276
277
278
279
280
        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())
281

282
283
284
285
286
287
288
        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.
289
290
291
            if self.health_loop is None:
                self.health_loop = asyncio.create_task(
                    self.run_heartbeat_loop(timeout=VLLM_RPC_TIMEOUT))
292
293
294
295
296
297
298
299
300

    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()
301
302
        if self.output_loop is not None:
            self.output_loop.cancel()
303
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

    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)

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

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

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

383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
    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)

408
409
410
411
412
413
414
415
    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)
        """
416
417
418
419
420
        pass

    async def check_health(self):
        """
        The check health loop probes the health status of the
421
        Engine's health every N seconds and sets _errored_with
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
        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

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

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

    @overload
457
    @deprecated("'inputs' will be renamed to 'prompt")
458
459
    def generate(
        self,
460
461
        *,
        inputs: PromptType,
462
463
464
465
466
        sampling_params: SamplingParams,
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
467
        priority: int = 0,
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
    ) -> 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,
483
        priority: int = 0,
484
485
        *,
        inputs: Optional[PromptType] = None  # DEPRECATED
486
    ) -> AsyncGenerator[RequestOutput, None]:
487
488
489
490
491
492
493
        """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:
494
            prompt: The prompt to the LLM. See :class:`~vllm.inputs.PromptType`
495
496
497
498
499
500
501
                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.
502
503
            priority: Priority of the request (lower means earlier handling).
                Any priority other than 0 will lead to an error if the
504
                scheduling policy is not "priority".
505
        """
506
507
508
509
510
511
        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,
512
                                     lora_request, trace_headers,
513
                                     prompt_adapter_request, priority)
514

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

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

    @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,
552
        priority: int = 0,
553
554
        *,
        inputs: Optional[PromptType] = None  # DEPRECATED
555
    ) -> AsyncGenerator[PoolingRequestOutput, None]:
556
        """Generate outputs for a request from a pooling model.
557
558
559
560
561
562

        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:
563
            prompt: The prompt to the LLM. See :class:`~vllm.inputs.PromptType`
564
565
566
567
568
569
570
                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:
571
            The output `PoolingRequestOutput` objects from the LLMEngine
572
573
            for the request.
        """
574
575
576
577
578
        if inputs is not None:
            prompt = inputs
        assert (prompt is not None and pooling_params is not None
                and request_id is not None)

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

    async def _process_request(
        self,
590
        prompt: PromptType,
591
592
593
594
        params: Union[SamplingParams, PoolingParams],
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
595
596
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
597
    ) -> Union[AsyncGenerator[RequestOutput, None], AsyncGenerator[
598
            PoolingRequestOutput, None]]:
599
600
601
602
603
604
        """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)

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

609
610
611
612
613
614
615
616
617
        # 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),
618
619
620
                    default_guided_backend=(self.decoding_config.guided_decoding_backend
                        if self.decoding_config
                        else DecodingConfig.guided_decoding_backend),
621
622
                    model_config=self.model_config,
                    reasoning_backend=self.decoding_config.reasoning_backend,
623
624
                )

625
626
627
628
629
630
631
632
        # 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)
633
            if isinstance(params, SamplingParams) and params.logits_processors:
634
                # Defensive shallow copy
635
636
637
                params = copy.copy(params)
                logits_processors = params.logits_processors
                params.logits_processors = None
638
639
640
641
642
                lp_bytes = cloudpickle.dumps(logits_processors)
            else:
                lp_bytes = None

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

            # 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)
677
678
679
680
681
682
683
684
685
686
687
688

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

690
691
    async def reset_prefix_cache(self,
                                 device: Optional[Device] = None) -> None:
692
693
694
        """Reset the prefix cache"""

        await self._send_one_way_rpc_request(
695
            request=RPCResetPrefixCacheRequest(device),
696
697
            socket=self.input_socket)

698
699
700
701
702
    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)

703
    async def wake_up(self, tags: Optional[list[str]] = None) -> None:
704
705
        """Wake up the engine"""
        return await self._send_one_way_rpc_request(
706
            request=RPCWakeUpRequest(tags), socket=self.input_socket)
707

708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
    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

726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
    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