client.py 26.4 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
23
24
25
26
# yapf conflicts with isort for this block
# yapf: disable
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,
27
                                         RPCAdapterLoadedResponse, RPCError,
28
29
                                         RPCIsSleepingRequest,
                                         RPCIsSleepingResponse,
30
                                         RPCLoadAdapterRequest,
31
                                         RPCProcessRequest,
32
                                         RPCResetMultiModalCacheRequest,
33
                                         RPCResetPrefixCacheRequest,
34
35
36
                                         RPCSleepRequest, RPCStartupRequest,
                                         RPCStartupResponse,
                                         RPCUProfileRequest, RPCWakeUpRequest)
37
from vllm.engine.protocol import EngineClient
38
39
# yapf: enable
from vllm.envs import VLLM_RPC_TIMEOUT
40
from vllm.inputs import PromptType
41
from vllm.inputs.preprocess import InputPreprocessor
42
43
from vllm.logger import init_logger
from vllm.lora.request import LoRARequest
44
from vllm.model_executor.layers.sampler import SamplerOutput
45
from vllm.outputs import PoolingRequestOutput, RequestOutput
46
from vllm.sampling_params import SamplingParams
47
from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
48
from vllm.utils import Device
49
50
51
52
53
54

logger = init_logger(__name__)


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

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


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

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

        # Get the configs.
94
        self.vllm_config = engine_config
95
96
97
        self.model_config = engine_config.model_config
        self.decoding_config = engine_config.decoding_config

98
99
100
101
102
103
104
105
106
107
        if self.vllm_config.model_config.skip_tokenizer_init:
            self.tokenizer = None

        else:
            # Create the tokenizer group.
            self.tokenizer = init_tokenizer_from_configs(
                model_config=self.model_config,
                scheduler_config=engine_config.scheduler_config,
                lora_config=engine_config.lora_config)

108
109
        self.input_preprocessor = InputPreprocessor(self.model_config,
                                                    self.tokenizer)
110
111
112
113
114
115
116
117
118

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

119
120
121
        # 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}")
122
123
124
125
126
127

        # 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] = {}
128
129
130
131
132

        # 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
133
134
135
136

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

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

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

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

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

175
                logger.debug("Heartbeat successful.")
176
177
178
179

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

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

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
231
232
        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
233
234
235
236
237
238
239
                        # 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
240
241

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

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

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

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

275
        # Start output_loop
276
277
278
279
280
281
282
283
        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())
284

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

    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()
304
305
        if self.output_loop is not None:
            self.output_loop.cancel()
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
375
376

    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)

377
378
379
    async def get_input_preprocessor(self) -> InputPreprocessor:
        return self.input_preprocessor

380
    async def get_tokenizer(self, lora_request: Optional[LoRARequest] = None):
381
382
383
384
        if self.tokenizer is None:
            return None
        else:
            return await self.tokenizer.get_lora_tokenizer_async(lora_request)
385

386
387
388
    async def get_vllm_config(self) -> VllmConfig:
        return self.vllm_config

389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
    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)

414
415
416
417
418
419
420
421
    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)
        """
422
423
424
425
426
        pass

    async def check_health(self):
        """
        The check health loop probes the health status of the
427
        Engine's health every N seconds and sets _errored_with
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
        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

445
446
    @property
    def dead_error(self) -> BaseException:
447
        return ENGINE_DEAD_ERROR(self._errored_with)
448

449
    def generate(
450
        self,
451
        prompt: PromptType,
452
453
454
455
        sampling_params: SamplingParams,
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
456
        priority: int = 0,
457
    ) -> AsyncGenerator[RequestOutput, None]:
458
459
460
461
462
463
464
        """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:
465
466
467
            prompt: The prompt to the LLM. See
                [`PromptType`][vllm.inputs.PromptType] for more details about
                the format of each input.
468
469
470
471
            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.
472
473
            priority: Priority of the request (lower means earlier handling).
                Any priority other than 0 will lead to an error if the
474
                scheduling policy is not "priority".
475
        """
476
477
478
        return cast(
            AsyncGenerator[RequestOutput, None],
            self._process_request(prompt, sampling_params, request_id,
479
                                  lora_request, trace_headers, priority))
480
481
482

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

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

    async def _process_request(
        self,
520
        prompt: PromptType,
521
522
523
524
        params: Union[SamplingParams, PoolingParams],
        request_id: str,
        lora_request: Optional[LoRARequest] = None,
        trace_headers: Optional[Mapping[str, str]] = None,
525
        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
        # 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)
546
            if isinstance(params, SamplingParams) and params.logits_processors:
547
                # Defensive shallow copy
548
549
550
                params = copy.copy(params)
                logits_processors = params.logits_processors
                params.logits_processors = None
551
552
553
554
555
                lp_bytes = cloudpickle.dumps(logits_processors)
            else:
                lp_bytes = None

            request_bytes = pickle.dumps(
556
                RPCProcessRequest(
557
                    prompt=prompt,
558
                    params=params,
559
560
561
                    request_id=request_id,
                    lora_request=lora_request,
                    trace_headers=trace_headers,
562
563
                    priority=priority,
                ))
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588

            # 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)
589
590
591
592
593
594
595
596
597
598
599
600

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

602
603
604
605
606
607
608
    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)

609
610
    async def reset_prefix_cache(self,
                                 device: Optional[Device] = None) -> None:
611
612
613
        """Reset the prefix cache"""

        await self._send_one_way_rpc_request(
614
            request=RPCResetPrefixCacheRequest(device),
615
616
            socket=self.input_socket)

617
618
619
620
621
    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)

622
    async def wake_up(self, tags: Optional[list[str]] = None) -> None:
623
624
        """Wake up the engine"""
        return await self._send_one_way_rpc_request(
625
            request=RPCWakeUpRequest(tags), socket=self.input_socket)
626

627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
    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

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