EngineCoreClient:行程內與多行程兩條客戶端通路

源码版本v0.25.1

職責

vLLM 的前端(LLM / AsyncLLM / serving 入口)不直接摸 EngineCore,而是經一層 client 把請求推進去、把輸出拉出來。這一層就是 EngineCoreClient,它把「同步還是非同步」「同行程還是子行程」這兩條正交維度,統一到同一組抽象方法上。vllm/v1/engine/core_client.py 這個檔案裡塞下了所有實現。

最關鍵的兩個實現是 InprocClient 和 MPClient。前者把 EngineCore 直接構造在本行程裡,調一步算一步,沒有 busy loop、沒有 ZMQ,主要給 V0 風格的 LLMEngine.add_request()/step() 當薄殼。後者把 EngineCore 包成 EngineCoreProc 子行程,通過 ZMQ ROUTER/PULL 雙 socket 收發請求與輸出,前端呼叫執行緒只跟 socket 打交道。MPClient 是基類,具體落地有 SyncMPClient(同步,起後臺執行緒從 output socket 收幀)和 AsyncMPClient(asyncio,起後臺 task 跑同一個 socket),後者另頁講。

EngineCoreClient.make_client 是個靜態工廠,根據 multiprocess_mode 和 asyncio_mode 兩個布爾挑實現。組合 (asyncio=True, multiprocess=False) 直接拋 NotImplementedError——EngineCore 自己不是 asyncio-friendly 的,非同步必須配多行程一起用。

設計動機

• 解耦前端與 EngineCore 的生命週期:同行程調 EngineCore 簡單,但前端要併發跑多個 EngineCore(資料並行)、或者要做彈性 EP(在線加減 rank),就必須把 EngineCore 推到獨立行程裡,前端只持 socket。MPClient.__init__ 裡 weakref.finalize(self, self.resources) 這一句是核心:就算構造中途拋異常,後臺子行程和 ZMQ context 也會被 BackgroundResources.__call__ 兜底回收,不會洩漏。
• 同步介面和非同步介面共享同一套 ZMQ 協議:SyncMPClient 用 queue.Queue 把 output socket 的幀搬到主執行緒,AsyncMPClient 用 asyncio.Queue 搬,但 MPClient 共享了 encoder/decoder、ready 握手、engine monitor、utility future 註冊表。差異被壓在兩個子類裡。
• InprocClient 的 sleep(mode="wait") 直接禁掉(InprocClient.sleep L324-L328),因為本行程模式下沒有可以「等到空閒再 sleep」的併發排程器,只能 abort。這種邊界在子行程模式裡是支持的——子行程有自己的 busy loop 和排程器狀態機。
• utility 呼叫走 future 註冊表:像 collective_rpc、save_sharded_state 這類需要返回值的 RPC,不能像 add_request 那樣 fire-and-forget。call_utility 生成一個 call_id,把 future 塞進 self.utility_results,等 UtilityOutput 從 output socket 回來時 _process_utility_output 用 call_id 取出 future 並 set_result。這套機制在 sync 和 async 兩種 client 上完全一致,只是 future 類型不同(concurrent.futures.Future vs asyncio.Future)。
• engine dead 傳播:子行程意外死亡時,monitor 執行緒把 BackgroundResources.engine_dead 置 True,後續任何 _send_input 都會被 ensure_alive 擋住拋 EngineDeadError。output socket 那邊如果收到 ENGINE_CORE_DEAD 單幀,validate_alive 也會把同一標誌位拉起來,保證發請求側和收輸出側都看得到。

關鍵檔案

• EngineCoreClient + make_client:71-105 — 抽象基類和靜態工廠,根據 multiprocess_mode/asyncio_mode 選 InprocClient/SyncMPClient/AsyncMPClient
• make_async_mp_client:107-132 — 非同步工廠,按 data_parallel_size 與是否外部 LB 選 AsyncMPClient/DPAsyncMPClient/DPLBAsyncMPClient
• InprocClient.__init__ / get_output:276-292 — 直接構造 EngineCore,get_output 就是同步調 step_fn() + post_step()
• InprocClient.add_request / abort / shutdown:297-306 — 直接轉發給 self.engine_core,沒有序列化、沒有 socket
• InprocClient.sleep 拒绝 wait:324-328 — 本行程模式不支持 wait pause,只能 abort
• BackgroundResources:370-458 — dataclass + __call__ 作為 weakref.finalize 回調,統一關 socket / 停 task / shutdown engine manager
• validate_alive:454-457 — 收到單幀 ENGINE_CORE_DEAD 即標記 engine_dead,拋 EngineDeadError
• MPClient.__init__:480-649 — 建 ZMQ context/socket、可選地 launch_core_engines、握手 ready、起 monitor,失敗時回滾
• MPClient.shutdown:651-662 — detach() finalizer 後停 engine manager、跑 BackgroundResources 清理
• start_engine_core_monitor:685-712 — daemon 執行緒 monitor_engine_liveness,引擎死亡時置 engine_dead 並觸發 self.shutdown()
• _apply_ready_response:714-754 — 解碼 EngineCoreReadyResponse,把 max_model_len/num_gpu_blocks/block_size 從子行程回填到前端 vllm_config
• SyncMPClient.__init__:779-847 — 建 queue.Queue,起 process_outputs_socket daemon 執行緒輪詢 output socket + shutdown socket
• SyncMPClient.get_output / _send_input / call_utility:849-881 — 阻塞 outputs_queue.get(),發請求時按需 track MessageTracker 防 tensor 提前回收

資料流

SyncMPClient 拉輸出靠一個 daemon 執行緒在 ZMQ output_socket 上 poll,解碼後塞進 queue.Queue。前端執行緒調 get_output() 就是 queue.get()——把 ZMQ 的非阻塞 socket 適配成 Python 同步阻塞介面。這段 process_outputs_socket 是整個同步通路的核心:

def process_outputs_socket():
    assert isinstance(out_socket, zmq.Socket)
    shutdown_socket = ctx.socket(zmq.PAIR)
    try:
        shutdown_socket.bind(shutdown_path)
        poller = zmq.Poller()
        poller.register(shutdown_socket, zmq.POLLIN)
        poller.register(out_socket, zmq.POLLIN)
        while True:
            socks = poller.poll()
            if not socks:
                continue
            if len(socks) == 2 or socks[0][0] == shutdown_socket:
                # shutdown signal, exit thread.
                break

            frames = out_socket.recv_multipart(copy=False)
            resources.validate_alive(frames)
            outputs: EngineCoreOutputs = decoder.decode(frames)
            if outputs.utility_output:
                _process_utility_output(outputs.utility_output, utility_results)
            else:
                outputs_queue.put_nowait(outputs)
    except Exception as e:
        outputs_queue.put_nowait(e)
    finally:
        # Close sockets.
        shutdown_socket.close(linger=0)
        out_socket.close(linger=0)

(SyncMPClient.process_outputs_socket:808-836)。注意三個細節:一是 shutdown 走單獨的 zmq.PAIR inproc socket 綁 shutdown_path,BackgroundResources.__call__ 關閉時往這個 socket 發一個空字節通知執行緒退出,避免執行緒死等 out_socket。二是 validate_alive 把 ENGINE_CORE_DEAD 單幀轉成 EngineDeadError。三是 utility 回包(如 collective_rpc 的返回)和普通 outputs 走同一幀,前者直接灌進 utility_results[call_id] 的 future,後者進佇列。

發請求的方向更簡單,SyncMPClient._send_input 直接把 (identity, request_type, *encoder.encode(request)) 多幀發到 ROUTER 上,有 tensor buffer 時用 track=True 拿到 MessageTracker,掛進 pending_messages 防止 Python 端把 backing tensor 提前 GC(_send_input:861-873)。

握手階段在 MPClient.__init__ 裡(等待 ready:615-634):對每個 engine_rank 生成 2 字節的 little-endian identity,在 input_socket 上 poll 等子行程的 EngineCoreReadyResponse,超時由 VLLM_ENGINE_READY_TIMEOUT_S 拉住。_apply_ready_response 把子行程跑出來的 num_gpu_blocks、對齊過的 block_size、自動擬合的 max_model_len 回填到前端的 vllm_config,這是前端真正能拿來排程和限流的真實容量。

邊界與失敗

• asyncio_mode=True, multiprocess_mode=False 直接拒絕 (make_client:91-95)。EngineCore 內部不是 async-friendly 的,這條組合沒有實現。
• 構造失敗必須回收子行程:MPClient.__init__ 用 try/finally + success 標誌,失敗時調 self._finalizer() 觸發 BackgroundResources.__call__,把已經 launch_core_engines 起來的子行程一起停掉(__init__ finally:499-649)。
• InprocClient 不支持 wait pause(InprocClient.sleep:324-328)。本行程模式下沒有誰去等空閒,只能 abort 當前請求後 sleep。
• engine core 子行程死亡會被兩個方向都感知:發送側靠 ensure_alive(ensure_alive:670-672),接收側靠 validate_alive 收到 ENGINE_CORE_DEAD 單幀(validate_alive:454-457)。兩路都會把 BackgroundResources.engine_dead 拉起來。
• utility future 的 InvalidStateError 容忍:_process_utility_output 捕獲 asyncio.InvalidStateError(_process_utility_output:769-776),覆蓋呼叫方 task 被取消導致 future 已 cancelled 的場景。
• output socket 關閉要在 context.term 之前:BackgroundResources.__call__ 顯式 close_sockets 再做後續(sync case:437-452),否則 ZMQ context 終止會卡住。
• SyncMPClient 的 output_socket 由執行緒負責關閉:構造最後把 self.resources.output_socket = None(socket ownership 转交:846-847),避免 BackgroundResources 和 daemon 執行緒同時關。
• ready 握手超時給可讀報錯:Timed out waiting for engine core processes to start. This is often caused by slow weight loading for large models.(ready timeout:622-630)——直接告訴使用者調 VLLM_ENGINE_READY_TIMEOUT_S。

小結

InprocClient 和 MPClient 這兩條路看起來差別巨大,但對外都長得像 EngineCoreClient——add_request 一調、get_output 一拉,前端代碼完全不用關心 EngineCore 是同行程物件還是子行程 + ZMQ。同行程簡單到極致,就是個物件引用;子行程那套把 ZMQ ROUTER/PULL、握手、monitor、future 表、socket ownership 全部封裝好了。非同步變體 AsyncMPClient 在 /client/async-mp 單獨拆;EngineCore 本身見 /engine/engine-core,EngineCoreProc 的 busy loop 和 shutdown 信號見 /engine/engine-core-proc。

對照官方資料:vLLM 文件 · README