client.py 27.5 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

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

import cloudpickle
12
import psutil
13
14
15
16
17
import zmq
import zmq.asyncio
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
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
    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
    ) -> AsyncGenerator[RequestOutput, None]:
454
455
456
457
458
459
460
        """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:
461
462
463
            prompt: The prompt to the LLM. See
                [`PromptType`][vllm.inputs.PromptType] for more details about
                the format of each input.
464
465
466
467
468
469
            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.
470
471
            priority: Priority of the request (lower means earlier handling).
                Any priority other than 0 will lead to an error if the
472
                scheduling policy is not "priority".
473
        """
474
475
476
477
478
        return cast(
            AsyncGenerator[RequestOutput, None],
            self._process_request(prompt, sampling_params, request_id,
                                  lora_request, trace_headers,
                                  prompt_adapter_request, priority))
479
480
481

    def encode(
        self,
482
        prompt: PromptType,
483
484
485
486
        pooling_params: PoolingParams,
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
487
        priority: int = 0,
488
    ) -> AsyncGenerator[PoolingRequestOutput, None]:
489
        """Generate outputs for a request from a pooling model.
490
491
492
493
494
495

        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:
496
497
498
            prompt: The prompt to the LLM. See
                [`PromptType`][vllm.inputs.PromptType] for more details about
                the format of each input.
499
500
501
502
503
504
            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:
505
            The output `PoolingRequestOutput` objects from the LLMEngine
506
507
            for the request.
        """
508
        return cast(
509
            AsyncGenerator[PoolingRequestOutput, None],
510
511
512
513
514
515
            self._process_request(prompt,
                                  pooling_params,
                                  request_id,
                                  lora_request,
                                  trace_headers,
                                  priority=priority))
516
517
518

    async def _process_request(
        self,
519
        prompt: PromptType,
520
521
522
523
        params: Union[SamplingParams, PoolingParams],
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
524
525
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        priority: int = 0,
526
    ) -> Union[AsyncGenerator[RequestOutput, None], AsyncGenerator[
527
            PoolingRequestOutput, None]]:
528
529
530
531
532
533
        """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)

534
535
536
537
        # Ensure the request id is unique among running requests
        if request_id in self.output_queues:
            raise ValueError(f"Request {request_id} already exists")

538
539
540
541
542
543
544
545
546
        # 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),
547
                    default_guided_backend=(self.decoding_config.backend
548
                        if self.decoding_config
549
                        else DecodingConfig.backend),
550
551
                    model_config=self.model_config,
                    reasoning_backend=self.decoding_config.reasoning_backend,
552
553
                )

554
555
556
557
558
559
560
561
        # 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)
562
            if isinstance(params, SamplingParams) and params.logits_processors:
563
                # Defensive shallow copy
564
565
566
                params = copy.copy(params)
                logits_processors = params.logits_processors
                params.logits_processors = None
567
568
569
570
571
                lp_bytes = cloudpickle.dumps(logits_processors)
            else:
                lp_bytes = None

            request_bytes = pickle.dumps(
572
                RPCProcessRequest(
573
                    prompt=prompt,
574
                    params=params,
575
576
577
                    request_id=request_id,
                    lora_request=lora_request,
                    trace_headers=trace_headers,
578
579
580
                    prompt_adapter_request=prompt_adapter_request,
                    priority=priority,
                ))
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

            # 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)
606
607
608
609
610
611
612
613
614
615
616
617

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

619
620
621
622
623
624
625
    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)

626
627
    async def reset_prefix_cache(self,
                                 device: Optional[Device] = None) -> None:
628
629
630
        """Reset the prefix cache"""

        await self._send_one_way_rpc_request(
631
            request=RPCResetPrefixCacheRequest(device),
632
633
            socket=self.input_socket)

634
635
636
637
638
    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)

639
    async def wake_up(self, tags: Optional[list[str]] = None) -> None:
640
641
        """Wake up the engine"""
        return await self._send_one_way_rpc_request(
642
            request=RPCWakeUpRequest(tags), socket=self.input_socket)
643

644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
    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

662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
    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