AsyncMPClient: putting an asyncio event loop around ZMQ

源码版本v0.25.1

Responsibilities

AsyncMPClient is the asynchronous variant of MPClient, used by AsyncLLM and the serving entry points. Like SyncMPClient, it runs EngineCore in a subprocess, sends requests through a ZMQ ROUTER, and receives outputs over PULL. The difference is that the output socket is no longer polled by a daemon thread — it is attached directly to the asyncio event loop, and a long-lived task await output_socket.recv_multipart(copy=False) pulls frames. The frontend awaits, the subprocess computes and replies, and the task decodes the reply and pushes it into an asyncio.Queue. The whole pipeline lives in a single event loop: no thread switches, no cross-thread queue.Queue shuffling.

AsyncMPClient itself is implemented sparingly: it only overrides _send_input (returning an Awaitable instead of blocking), exposes a set of async methods like call_utility_async / add_request_async, and adds the crucial _ensure_output_queue_task. Most of the initialization logic in __init__ reuses MPClient.__init__'s ZMQ handshake and engine monitor; the only difference is that asyncio_mode=True swaps the context to zmq.asyncio.Context and the sockets to zmq.asyncio.Socket. Data-parallel (DP) and elastic-EP extensions are handled by two subclasses, DPAsyncMPClient and DPLBAsyncMPClient, touched on briefly at the end of this page.

Design motivation

• The async loop is single-threaded; a daemon thread cannot be opened: the process_outputs_socket daemon thread from SyncMPClient is unnecessary in async mode — zmq.asyncio.Socket.recv_multipart is already a coroutine, so it can be awaited directly inside an asyncio.create_task. The thread is replaced by a task, and the queue swaps from queue.Queue to asyncio.Queue.
• The loop may not be running at construction time: in AsyncMPClient.__init__, if asyncio.get_running_loop() fails it just passes, deferring the output queue task startup to the first call of _ensure_output_queue_task (lazy task start:974-982). This lets the client be constructed outside an event loop (in tests or during initialization) and lazily attach when actually running. The trade-off is documented: between construction and the first add_request_async, if the subprocess dies and sends back ENGINE_CORE_DEAD, that event is lost because the task is not up yet. The comment explicitly calls this out.
• Avoid the task directly holding a client reference: _ensure_output_queue_task pulls decoder, utility_results, outputs_queue, and output_socket out as closure-local variables; the client itself is held via weakref.ref(self) (_self_ref). If the frontend drops the client, the closure does not keep it alive (no direct ref:989-998). If _self = _self_ref() returns None, the task simply returns and kills itself.
• Utility futures use asyncio.Future: _call_utility_async uses asyncio.get_running_loop().create_future() rather than concurrent.futures.Future, so that await future is driven directly by the event loop. Future fulfillment still goes through the shared _process_utility_output, which is compatible with both future types (_call_utility_async:1104-1116).
• EEP notifications use a special call_id: inside process_outputs_socket, if utility_output.call_id == EEP_NOTIFICATION_CALL_ID, it does not touch the future table — instead it asyncio.create_tasks the eep_process_engine_core_notification callback (EEP notification path:1011-1027). Elastic EP rank add/remove notifications are events, not request-response, so the utility future machinery does not apply.
• output_handler is an optional hook: getattr(self.__class__, "process_engine_outputs", None) (output_handler:994-996). DP subclasses (such as DPLBAsyncMPClient) override this method for load-balancing stats. The base AsyncMPClient does not define it, which is equivalent to not intercepting — but outputs.outputs or outputs.scheduler_stats still goes into the queue (enqueue condition:1042-1043).

Key files

• AsyncMPClient.__init__:950-982 — passes asyncio_mode=True to MPClient.__init__, creates the asyncio.Queue, and tries to lazily start the task if a loop is already running.
• _ensure_output_queue_task:984-1051 — pulls decoder/queue/utility_results/socket into closure locals, calls asyncio.create_task(process_outputs_socket()), idempotent.
• process_outputs_socket:1005-1047 — the await output_socket.recv_multipart(copy=False) main loop, splitting frames into utility / EEP notification / ordinary output.
• CancelledError handling:1042-1047 — when the task is cancelled, pushes EngineDeadError into the queue so the get_output_async side also notices.
• get_output_async:1053-1062 — await outputs_queue.get(); if an Exception comes out, wraps it with _format_exception into an EngineDeadError.
• _send_input / _send_input_message:1064-1099 — returns an Awaitable; with no buffer it returns the send_multipart coroutine directly, with a buffer it uses track=True + add_done_callback to obtain a MessageTracker asynchronously.
• call_utility_async / _call_utility_async:1101-1116 — generates a call_id, registers an asyncio.Future, then await self._send_input_message + await future.
• add_request_async / abort_requests_async:1121-1128 — tags requests with client_index, and after sending calls _ensure_output_queue_task to make sure the task is up.
• collective_rpc_async:1188-1197 — forwards to call_utility_async("collective_rpc", ...), parameters aligned with the sync version.
• DPAsyncMPClient.__init__:1200-1240 — adds current_wave, lb_engines, first_req_send_socket, eep_scaling_cache; lazily starts the stats update task when the loop is already running.
• _ensure_stats_update_task:1242-1249 — starts the DP stats subscription task; not present on the base AsyncMPClient.
• DPLBAsyncMPClient:1380 — internal load-balancing subclass; overrides call_utility_async to dispatch requests across lb_engines to the least-loaded rank.

Data flow

The entire chain runs in a single asyncio loop. The frontend await client.add_request_async(req) → _send_input does send_multipart of the frame to the ZMQ ROUTER (returning the send_multipart awaitable directly when there is no buffer, or wrapping with track=True + a done callback when there is a buffer). The subprocess side's EngineCoreProc.run_busy_loop computes and sends EngineCoreOutputs back to the PULL socket. The long-lived task process_outputs_socket picks up the frame at await output_socket.recv_multipart(copy=False), decodes it, and dispatches by type:

async def process_outputs_socket():
    try:
        while True:
            frames = await output_socket.recv_multipart(copy=False)
            resources.validate_alive(frames)
            outputs: EngineCoreOutputs = decoder.decode(frames)
            if outputs.utility_output:
                if (
                    outputs.utility_output.call_id == EEP_NOTIFICATION_CALL_ID
                    and notification_callback_handler is not None
                ):
                    assert _self_ref is not None
                    _self = _self_ref()
                    if not _self:
                        return
                    if outputs.utility_output.result is None:
                        continue
                    notification_data = outputs.utility_output.result.result
                    assert isinstance(notification_data, Sequence)
                    assert len(notification_data) == 2
                    asyncio.create_task(
                        notification_callback_handler(_self, notification_data)
                    )
                else:
                    _process_utility_output(
                        outputs.utility_output, utility_results
                    )
                continue

            if output_handler is not None:
                assert _self_ref is not None
                _self = _self_ref()
                if not _self:
                    # Client has been garbage collected, abort.
                    return
                await output_handler(_self, outputs)

            if outputs.outputs or outputs.scheduler_stats:
                outputs_queue.put_nowait(outputs)
    except Exception as e:
        outputs_queue.put_nowait(e)
    except asyncio.CancelledError:
        outputs_queue.put_nowait(EngineDeadError())

(process_outputs_socket:1005-1047). Note the three-way split: utility replies go through the future table, EEP notifications spawn a new task, and ordinary outputs go through the output_handler hook (if present) before being enqueued. output_handler is designed for DP subclasses to plug in stats — the base class does not define it, so frames go straight to outputs_queue, and the frontend picks them up via await get_output_async().

The future path for utility calls deserves a separate look, because it unifies the sync and async future types:

async def _call_utility_async(
    self, method: str, *args, engine: EngineIdentity
) -> Any:
    call_id = uuid.uuid1().int >> 64
    future = asyncio.get_running_loop().create_future()
    self.utility_results[call_id] = future
    message = (
        EngineCoreRequestType.UTILITY.value,
        *self.encoder.encode((self.client_index, call_id, method, args)),
    )
    await self._send_input_message(message, engine, args)
    self._ensure_output_queue_task()
    return await future

(_call_utility_async:1104-1116). call_id uses uuid.uuid1().int >> 64 to take the upper 64 bits — large and sparse enough. The future is registered before the request is sent, so a fast reply cannot beat the registration. _ensure_output_queue_task() is called once more to make sure the task is up — otherwise await future would never be fulfilled. Fulfillment happens in process_outputs_socket via the _process_utility_output(outputs.utility_output, utility_results) branch.

Boundaries and failure modes

• Loop not running at construction loses early death signals: in __init__, if asyncio.get_running_loop() fails the task is started lazily (lazy:974-982). This means that if the subprocess dies before the first _ensure_output_queue_task, the ENGINE_CORE_DEAD frame is never consumed. The comment explicitly warns about this trade-off.
• Cancelled task still pushes EngineDeadError: except asyncio.CancelledError pushes EngineDeadError into the queue (CancelledError:1046-1047), so the get_output_async side raises immediately instead of hanging silently.
• Task self-terminates when the client is GC'd: when _self = _self_ref() returns None the task returns (GC self-kill:1037-1039), so it does not keep holding released resources.
• With a tensor buffer, returns a future instead of a coroutine: in _send_input_message, if len(msg) > 3, send_multipart(track=True) returns an asyncio.Future[zmq.MessageTracker], and add_done_callback(add_pending) asynchronously attaches the tracker to pending_messages (track path:1091-1099). The caller gets an Awaitable that only completes the send after being awaited.
• EEP notification call_id is separate: EEP_NOTIFICATION_CALL_ID does not go through the utility_results table (EEP branch:1011-1027) — it spawns a new task to call eep_process_engine_core_notification, because notifications are one-way events.
• _ensure_output_queue_task is idempotent: if resources.output_queue_task is not None: return (idempotent:986-987). add_request_async and call_utility_async call it every time to make sure the task is up; if it is already up, it is a no-op.
• DP subclass extension points are explicit: DPAsyncMPClient adds current_wave / lb_engines / eep_scaling_cache and _ensure_stats_update_task; DPLBAsyncMPClient overrides call_utility_async to do internal LB (DP init:1200-1240). The base class does not carry any of this DP logic.
• The output_handler hook is a classmethod-style lookup: getattr(self.__class__, "process_engine_outputs", None) (output_handler lookup:994-996). A subclass that defines this method gets called; the base class skips it. Note that it can only be invoked after _self_ref successfully resolves the client — it cannot close over self directly.

Summary

The core of AsyncMPClient is replacing SyncMPClient's daemon thread + queue.Queue with an asyncio task + asyncio.Queue; the ZMQ context also switches to zmq.asyncio.Context. Everything else — the ready handshake, engine monitor, utility future table, MessageTracker to prevent early tensor GC, and ENGINE_CORE_DEAD propagation — is inherited from the MPClient base. DP and elastic-EP extensions live in DPAsyncMPClient / DPLBAsyncMPClient, plugged in via the output_handler hook and by overriding call_utility_async. For the sync variant see /client/inproc-mp; for EngineCore and EngineCoreProc themselves see /engine/engine-core and /engine/engine-core-proc.

See official docs: vLLM docs · README