AsyncMPClient: asyncio-Event-Loop um ZMQ

源码版本v0.25.1

Verantwortung

AsyncMPClient ist die asynchrone Variante von MPClient und wird von AsyncLLM und den Serving-Einstiegen verwendet. Wie SyncMPClient verlagert es EngineCore in einen Unterprozess und sendet Anfragen über ZMQ ROUTER und empfängt Outputs über PULL. Der Unterschied: Das Output-Socket wird nicht mehr von einem Daemon-Thread polliert, sondern hängt direkt an der asyncio-Event-Loop; ein dauerhafter Task await output_socket.recv_multipart(copy=False) zieht die Frames. So läuft das Frontend mit await, der Unterprozess rechnet und schickt zurück, der Task dekodiert und steckt die Daten in eine asyncio.Queue — die gesamte Kette läuft in einer Event-Loop, ohne Thread-Wechsel, ohne threadübergreifendes Umkopieren aus queue.Queue.

AsyncMPClient selbst ist sparsam implementiert: Es überschreibt nur _send_input (gibt ein Awaitable statt blockierend zurück), eine Gruppe async-Methoden wie call_utility_async/add_request_async sowie den zentralen _ensure_output_queue_task. Die meisten Initialisierungen in __init__ verwenden die ZMQ-Handshake- und Engine-Monitor-Logik von MPClient.__init__; nur asyncio_mode=True sorgt dafür, dass der Context zu zmq.asyncio.Context und der Socket zu zmq.asyncio.Socket wird. Datenparallel (DP) und elastisches EP werden von den beiden Unterklassen DPAsyncMPClient und DPLBAsyncMPClient übernommen; diese Seite behandelt sie nur am Rande.

Entwurfsmotivation

Async-Loop ist single-Threaded, keine Daemon-Threads: Der process_outputs_socket-Daemon-Thread von SyncMPClient ist im asynchronen Modus unnötig — zmq.asyncio.Socket.recv_multipart ist selbst eine Coroutine und kann direkt awaited werden; in asyncio.create_task gesteckt, ist die Sache erledigt. Der Thread wird durch einen Task ersetzt; die Queue wird von queue.Queue zu asyncio.Queue.
Bei Konstruktion ist die Loop vielleicht noch nicht oben: In AsyncMPClient.__init__ fängt asyncio.get_running_loop() einen Fehler ab und passt; der Start des Output-Queue-Tasks wird auf den ersten Aufruf von _ensure_output_queue_task verschoben (lazy task start:974-982). So kann der Client außerhalb der Event-Loop konstruiert werden (z. B. in Test oder Init); beim echten Start wird lazy angehängt. Das hat einen Preis — zwischen Konstruktion und dem ersten add_request_async geht möglicherweise ein ENGINE_CORE_DEAD verloren, weil der Task noch nicht läuft; der Kommentar weist explizit auf diesen Trade-off hin.
Task darf den Client nicht direkt halten: In _ensure_output_queue_task werden decoder, utility_results, outputs_queue, output_socket als Closure-lokale Variablen herausgezogen; der Client selbst wird über weakref.ref(self) (_self_ref) referenziert. So wird der Client, wenn das Frontend ihn fallen lässt, nicht durch eine Closure-Referenz am Leben erhalten (no direct ref:989-998). _self = _self_ref() liefert None, so returnt der Task und beendet sich selbst.
Utility-Future als asyncio.Future: _call_utility_async verwendet asyncio.get_running_loop().create_future() und nicht concurrent.futures.Future, damit await future direkt von der Event-Loop angetrieben wird. Das Füllen des Futures erfolgt weiterhin über die geteilte Funktion _process_utility_output, die mit beiden Future-Typen kompatibel ist (_call_utility_async:1104-1116).
EEP-Benachrichtigungen mit spezieller call_id: Wenn in process_outputs_socket utility_output.call_id == EEP_NOTIFICATION_CALL_ID ist, wird das Future-Registry umgangen und stattdessen per asyncio.create_task der Callback eep_process_engine_core_notification getriggert (EEP notification path:1011-1027). Rank-Änderungen beim elastischen EP sind ein "Ereignis" und keine "Anfrage-Antwort" und können daher nicht über das Utility-Future abgebildet werden.
output_handler ist ein optionaler Hook: getattr(self.__class__, "process_engine_outputs", None) (output_handler:994-996). DP-Unterklassen (z. B. DPLBAsyncMPClient) überschreiben diese Methode, um Lastverteilungs-Statistiken zu erfassen; die Basisklasse AsyncMPClient hat diese Methode nicht und fängt nichts ab — aber outputs.outputs or outputs.scheduler_stats geht dennoch in die Queue (enqueue condition:1042-1043).

Schlüsseldateien

AsyncMPClient.__init__:950-982 — reicht asyncio_mode=True an MPClient.__init__ weiter; baut asyncio.Queue und versucht lazy, in einer bestehenden Loop einen Task zu starten.
_ensure_output_queue_task:984-1051 — zieht decoder/queue/utility_results/socket als Closure-locals; asyncio.create_task(process_outputs_socket()); idempotent.
process_outputs_socket:1005-1047 — Hauptschleife await output_socket.recv_multipart(copy=False); unterscheidet drei Frame-Typen: Utility / EEP-Benachrichtigung / normale Outputs.
CancelledError handling:1042-1047 — wenn der Task abgebrochen wird, wird EngineDeadError in die Queue gelegt, sodass die Seite get_output_async es ebenfalls merkt.
get_output_async:1053-1062 — await outputs_queue.get(); enthält der Wert eine Exception, wird sie per _format_exception in EngineDeadError verpackt.
_send_input / _send_input_message:1064-1099 — gibt Awaitable zurück; ohne Buffer direkt die send_multipart-Coroutine; mit Buffer track=True + add_done_callback, um asynchron den MessageTracker zu holen.
call_utility_async / _call_utility_async:1101-1116 — erzeugt call_id, registriert asyncio.Future, await self._send_input_message + await future.
add_request_async / abort_requests_async:1121-1128 — stempelt client_index auf die Anfrage und ruft nach dem Senden _ensure_output_queue_task auf, um den Task zu sichern.
collective_rpc_async:1188-1197 — reicht an call_utility_async("collective_rpc", ...) weiter; Parameter stimmen mit der sync-Version überein.
DPAsyncMPClient.__init__:1200-1240 — zusätzlich current_wave, lb_engines, first_req_send_socket, eep_scaling_cache; bei laufender Loop lazy ein Stats-Update-Task.
_ensure_stats_update_task:1242-1249 — startet den DP-Stats-Subscribe-Task; in der Basisklasse AsyncMPClient nicht vorhanden.
DPLBAsyncMPClient:1380 — interne Load-Balancer-Unterklasse; überschreibt call_utility_async, um Anfragen gemäß lb_engines an den am wenigsten ausgelasteten Rank zu senden.

Datenfluss

Die gesamte Kette läuft in einer asyncio-Loop. Frontend await client.add_request_async(req) → _send_input sendet die Frames per send_multipart an ZMQ ROUTER (ohne Buffer direkt als Awaitable von send_multipart; mit Buffer mit track=True + Done-Callback verpackt). In EngineCoreProc.run_busy_loop berechnet der Unterprozess und sendet EngineCoreOutputs an den PULL-Socket zurück. Der dauerhafte Task process_outputs_socket holt die Frames an await output_socket.recv_multipart(copy=False), dekodiert und leitet sie nach Typ weiter:

python

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). Drei Wege: Utility-Antwort geht in die Future-Registry; EEP-Benachrichtigung startet einen neuen Task; normale Outputs laufen über den output_handler-Hook (falls vorhanden) und dann in die Queue. output_handler ist für DP-Unterklassen zum Einklinken von Statistiken gedacht — die Basisklasse hat diese Methode nicht; Frames gehen direkt in die outputs_queue, und das Frontend await get_output_async() erhält diese.

Der Future-Pfad für Utility-Aufrufe muss separat betrachtet werden, weil er sync und async Future-Typen vereinheitlicht:

python

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 nimmt uuid.uuid1().int >> 64 und die oberen 64 Bit — groß und dünn genug verteilt; das Future wird registriert, bevor die Anfrage gesendet wird, damit das Antwort nicht vor der Registrierung ankommt; _ensure_output_queue_task() wird noch einmal aufgerufen, um den Task zu sichern — sonst würde await future nie gefüllt. Das Füllen erfolgt in process_outputs_socket im Zweig _process_utility_output(outputs.utility_output, utility_results).

Grenzen und Fehler

Bei Konstruktion ohne Loop gehen frühe Todsignale verloren: In __init__ schlägt asyncio.get_running_loop() fehl und der Task wird lazy gestartet (lazy:974-982); stirbt der Unterprozess vor dem ersten _ensure_output_queue_task, wird der ENGINE_CORE_DEAD-Frame nicht konsumiert. Der Kommentar weist explizit auf diesen Trade-off hin.
Bei Task-Abbruch wird EngineDeadError in die Queue gelegt: except asyncio.CancelledError steckt EngineDeadError in die Queue (CancelledError:1046-1047), sodass auf der get_output_async-Seite beim await direkt ein Fehler geworfen wird und nicht stumm hängen bleibt.
Bei GC des Clients beendet sich der Task selbst: _self = _self_ref() liefert None, dann return (GC suicide:1037-1039), damit der Task keine bereits freigegebenen Ressourcen weiterhält.
Mit Tensor-Buffer wird ein Future statt einer Coroutine zurückgegeben: In _send_input_message gilt: wenn len(msg) > 3, gibt send_multipart(track=True) ein asyncio.Future[zmq.MessageTracker] zurück; add_done_callback(add_pending) hängt den Tracker asynchron an pending_messages (track path:1091-1099). Der Aufrufer erhält ein Awaitable; erst nach dem await ist das Senden wirklich abgeschlossen.
EEP-Benachrichtigung hat eigene call_id: EEP_NOTIFICATION_CALL_ID geht nicht über utility_results (EEP branch:1011-1027), sondern spawnt einen neuen Task, der eep_process_engine_core_notification aufruft, weil Benachrichtigungen einseitige Ereignisse sind.
_ensure_output_queue_task ist idempotent: if resources.output_queue_task is not None: return (idempotent:986-987). add_request_async und call_utility_async rufen es jedes Mal auf, um den Task zu sichern; läuft er bereits, ist es ein No-Op.
DP-Unterklassen haben klare Erweiterungspunkte: DPAsyncMPClient fügt current_wave / lb_engines / eep_scaling_cache und _ensure_stats_update_task hinzu; DPLBAsyncMPClient überschreibt call_utility_async für das interne LB (DP init:1200-1240). Die Basisklasse trägt diese DP-Logik nicht.
output_handler-Hook ist ein classmethod-artiger Lookup: getattr(self.__class__, "process_engine_outputs", None) (output_handler lookup:994-996). Definiert die Unterklasse diese Methode, wird sie aufgerufen; definiert die Basisklasse sie nicht, wird übersprungen. Achtung: Sie kann nur aufgerufen werden, wenn _self_ref den Client liefert; self darf nicht direkt geschlossen werden.

Zusammenfassung

Der Kern von AsyncMPClient ist, den Daemon-Thread + queue.Queue von SyncMPClient gegen asyncio-Task + asyncio.Queue zu tauschen; auch der ZMQ-Context wechselt zu zmq.asyncio.Context. Alles andere — Ready-Handshake, Engine-Monitor, Utility-Future-Registry, MessageTracker gegen vorzeitige GC des Tensors, ENGINE_CORE_DEAD-Ausbreitung — wird von der MPClient-Basisklasse geerbt. DP- und elastisches EP-Erweiterungen leben in DPAsyncMPClient / DPLBAsyncMPClient und binden über den output_handler-Hook und das Überschreiben von call_utility_async ein. Die synchrone Variante steht in /client/inproc-mp; EngineCore und EngineCoreProc selbst in /engine/engine-core und /engine/engine-core-proc.

Siehe offizielle Dokumentation: vLLM 文档 · README.

AsyncMPClient: asyncio-Event-Loop um ZMQ ​

Verantwortung ​

Entwurfsmotivation ​

Schlüsseldateien ​

Datenfluss ​

Grenzen und Fehler ​

Zusammenfassung ​