EngineCoreProc:ZMQ 后台进程包装

源码版本v0.25.1

职责

EngineCoreProc 是 EngineCore 的子类(core.py:896-897),把它整个扔进一个独立的后台进程 (background process) 里跑。docstring 写得很直白:ZMQ-wrapper for running EngineCore in background process.。它干三件事:在进程启动时跟前端握手拿到 ZMQ socket 地址、起两个 IO 线程把 socket 和两个 queue.Queue 桥接起来(core.py:915-916)、然后 run_busy_loop 在主线程里反复从 input_queue 拿请求 → step_fn() 跑调度和前向 → 把 outputs 塞进 output_queue(core.py:1259-1267)。

这一层是 v1 多进程架构的关键。前端进程 (SyncMPClient 或 AsyncMPClient) 通过 ZMQ 把 EngineCoreRequest 序列化发过来,EngineCoreProc 的 input 线程收到后 put_nowait 进 input_queue,主循环在调度间隙从队列里取出来交给 _handle_client_request 分发(core.py:1372-1405)。outputs 反方向流:step_fn() 产出的 EngineCoreOutputs 被丢进 output_queue,output 线程从队列拿出来序列化通过 ZMQ 推回前端。ZMQ socket IO 在收发时会释放 GIL,所以这两条 IO 线程可以真正和 GPU 前向重叠(core.py:974-1001),序列化/反序列化也被藏到前向背后。

设计动机

• 进程隔离:EngineCoreProc 跑在独立进程里,前端 Python 崩了不会拖垮已经跑进 GPU 的 worker,worker OOM 也不会把 API 进程一起带走。这就是 v1 把 EngineCore 抽出来的核心理由——v0 的 LLMEngine 在同进程里什么都干,崩一处全崩。
• ZMQ + queue.Queue 两级缓冲:ZMQ socket 不直接被主循环用,中间夹一层 queue.Queue(core.py:915-916)。socket IO 在 IO 线程里,主循环只 queue.get(block=...),免得主循环被 socket recv 阻塞或被 ZMQ 内部锁拖住;同时 socket IO 释放 GIL 时,主循环能立刻拿回 GIL 跑前向。
• 两个 IO 线程:process_input_sockets 和 process_output_sockets(core.py:980-1001)各自独立,输入线程负责反序列化 + put_nowait,输出线程负责 get + 序列化 + send_multipart,两边互不阻塞。
• 启动握手 (handshake):构造函数里 _perform_handshakes(core.py:932-938)拿到一组 EngineZmqAddresses,然后 input 线程在 ROUTER/DEALER 上发一条 EngineCoreReadyResponse 告诉前端"我起来了,这是 max_model_len、num_gpu_blocks、kv_cache_size_tokens"(core.py:1525-1546)。前端在 MPClient.__init__ 里等这条 ready 消息,等不到就当启动失败。
• run_engine_core 进程入口:这函数是 multiprocessing.Process(target=run_engine_core, ...) 的 target(core.py:1154-1224)。它把 EngineCoreProc 实例化、注册 SIGTERM/SIGINT、最后调 run_busy_loop。任何异常都走 except Exception → _send_engine_dead → raise,进程退出时 finally 里调 engine_core.shutdown()(core.py:1226-1242)。

关键文件

• EngineCoreProc class:896-910 — class EngineCoreProc(EngineCore): 定义 + 类常量 ENGINE_CORE_DEAD = b"ENGINE_CORE_DEAD"。
• EngineCoreProc.__init__:903-1009 — 起 input_queue/output_queue、tensor IPC receiver、_perform_handshakes、super().__init__、起两个 IO 线程。
• queues + executor_fail_callback:915-919 — input_queue = queue.Queue[...]、output_queue = queue.Queue[...],executor 失败回调往 input_queue 塞 EXECUTOR_FAILED。
• IO threads startup:974-1001 — process_input_sockets 和 process_output_sockets 两个 daemon thread 启动,ready_event 在 input 线程 ready 后才 set。
• run_engine_core:1154-1224 — 进程入口,选 EngineCoreProc 或 DPEngineCoreProc,注册信号 handler,调 run_busy_loop。
• signal handlers:1204-1222 — wakeup_engine 往 input_queue 塞 WAKEUP,SIGTERM/SIGINT 把 shutdown_state 改成 REQUESTED。
• run_busy_loop:1259-1267 — while self._handle_shutdown(): _process_input_queue(); _process_engine_step() 跑到 shutdown 后 raise SystemExit。
• _process_input_queue:1269-1298 — 没活时阻塞 input_queue.get(block=True),有活时非阻塞 drain。
• _process_engine_step:1300-1317 — 调 step_fn()、output_queue.put_nowait(output)、post_step,没跑前向但 scheduler 还有活时 time.sleep(0.001) 让 GIL。
• _handle_client_request:1372-1405 — EngineCoreRequestType 分发:WAKEUP/ADD/ABORT/UTILITY/EXECUTOR_FAILED。
• _send_engine_dead:1470-1482 — 往 output_queue 塞 ENGINE_CORE_DEAD,等 output 线程把消息发出去再退出。
• process_input_sockets:1484-1587 — input 线程主体,用 zmq.Poller 监听 DEALER + 可选 coord XSUB,反序列化后 put_nowait 进 input_queue。
• process_output_sockets:1589-1654 — output 线程主体,从 output_queue.get() 拿 outputs,encoder.encode_into + send_multipart(copy=False, track=True) 零拷贝发回前端。

数据流

请求从 ZMQ input socket 进来,先被 input 线程反序列化塞进 input_queue,主循环在 _process_input_queue 里取出来交给 _handle_client_request 分发:

# vllm/v1/engine/core.py L1556-L1587
while True:
    for input_socket, _ in poller.poll():
        # (RequestType, RequestData)
        type_frame, *data_frames = input_socket.recv_multipart(copy=False)
        # NOTE(yongji): ignore READY message sent by DP coordinator
        # that is used to notify newly started engines
        if type_frame.buffer == b"READY":
            assert input_socket == coord_socket
            continue
        request_type = EngineCoreRequestType(bytes(type_frame.buffer))

        # Deserialize the request data.
        request: Any
        if request_type == EngineCoreRequestType.ADD:
            req: EngineCoreRequest = add_request_decoder.decode(data_frames)
            try:
                request = self.preprocess_add_request(req)
            except Exception:
                self._handle_request_preproc_error(req)
                continue
        else:
            request = generic_decoder.decode(data_frames)

            if request_type == EngineCoreRequestType.ABORT:
                # Aborts are added to *both* queues, allows us to eagerly
                # process aborts while also ensuring ordering in the input
                # queue to avoid leaking requests. This is ok because
                # aborting in the scheduler is idempotent.
                self.aborts_queue.put_nowait(request)

        # Push to input queue for core busy loop.
        self.input_queue.put_nowait((request_type, request))

注意 ABORT 被同时塞进 aborts_queue 和 input_queue,因为 abort 在调度器里是幂等的 (core.py:1276-1279),让主循环在 step 间隙也能从 aborts_queue 里 drain 掉,避免 leak 已经 abort 的请求。

主循环拿到请求后,_handle_client_request 按 EngineCoreRequestType 分发(core.py:1372-1405):

# vllm/v1/engine/core.py L1377-L1401
if request_type == EngineCoreRequestType.WAKEUP:
    return
elif request_type == EngineCoreRequestType.ADD:
    req, request_wave = request
    if self._reject_add_in_shutdown(req):
        return
    self.add_request(req, request_wave)
elif request_type == EngineCoreRequestType.ABORT:
    self.abort_requests(request)
elif request_type == EngineCoreRequestType.UTILITY:
    client_idx, call_id, method_name, args = request
    if self._reject_utility_in_shutdown(client_idx, call_id, method_name):
        return
    output = UtilityOutput(call_id)
    # Lazily look-up utility method so that failure will be handled/returned.
    get_result = lambda: (
        (method := getattr(self, method_name))
        and method(*self._convert_msgspec_args(method, args))
    )
    enqueue_output = lambda out: self.output_queue.put_nowait(
        (client_idx, EngineCoreOutputs(utility_output=out))
    )
    self._invoke_utility_method(method_name, get_result, output, enqueue_output)
elif request_type == EngineCoreRequestType.EXECUTOR_FAILED:
    raise RuntimeError("Executor failed.")

UTILITY 是个通用 RPC 通道:前端把方法名 + msgspec args 发过来,EngineCoreProc 用 getattr 反射调用,结果走 output_queue 回送。_invoke_utility_method 还能处理返回 Future 的情况(core.py:1440-1446),异步工具完成后才回送 output。

边界与失败

• 进程崩溃通知:run_engine_core 在 except Exception 里调 engine_core._send_engine_dead()(core.py:1229-1235),往 output_queue 塞一个 ENGINE_CORE_DEAD 字节串(core.py:1470-1474),output 线程收到这个特殊值后向所有 output socket 广播(core.py:1625-1628),前端 MPClient 收到后抛错而不是无限等下去。
• Executor 失败回调:__init__ 里把 executor_fail_callback 闭包注册到 model_executor(core.py:917-919),worker 内部报错时回调往 input_queue 塞一条 EXECUTOR_FAILED,主循环 _handle_client_request 拿到就 raise RuntimeError("Executor failed.")(core.py:1400-1401),整个进程会被 run_engine_core 的 except 捕到走 _send_engine_dead。
• Shutdown 模式:_handle_shutdown 看 shutdown_state(core.py:1324-1360),REQUESTED 状态下 shutdown_timeout==0 走 abort 立刻杀所有在飞请求,非 0 走 drain 等所有请求跑完。shutdown 期间新进来的 ADD 和 UTILITY 会被显式拒绝(core.py:1407-1432),_reject_add_in_shutdown 给客户端回一条 abort output,_reject_utility_in_shutdown 回一条带 failure_message="Server shutting down" 的 UtilityOutput。
• input 线程死了要报错:__init__ 里 ready_event.wait(timeout=10) 循环等 input 线程 ready(core.py:1003-1009),10 秒还没 ready 且线程死了就 raise RuntimeError("Input socket thread died during startup");如果只是没 ready 还活着,就继续等 DP Coordinator 的 READY 消息。
• 零拷贝发送 + MessageTracker:output 线程用 send_multipart(buffers, copy=False, track=True)(core.py:1646-1654),返回 MessageTracker,没发送完的 tracker + outputs 引用 + buffer 被挂到 pending deque 上,避免 GC 把还没发出去的 backing buffer 回收。reuse_buffers 限制最多 len(sockets) + 1 个复用 buffer,防止内存无限增长。
• WAKEUP 是空操作:信号 handler 通过 wakeup_engine 往 input_queue 塞 WAKEUP,只是把阻塞在 input_queue.get(block=True) 的主循环唤醒(core.py:1377-1378),本身不做事,真正的 shutdown 处理在 _handle_shutdown 里看 shutdown_state 完成。
• finally 里恢复信号:run_engine_core 在 finally 里 signal.signal(signal.SIGTERM, signal.SIG_DFL)(core.py:1236-1242),把信号 handler 恢复成默认,然后调 engine_core.shutdown() 释放 model_executor / scheduler / 分布式状态。shutdown 也会 gc.unfreeze()(core.py:652-656),把启动时 gc.freeze() 冻住的 startup heap 重新交给 GC,否则同进程模式(unit test)会泄漏显存。

小结

EngineCoreProc 是 v1 多进程架构的承重墙:它继承 EngineCore,在独立进程里跑 ZMQ + 两个 queue.Queue + 两个 IO 线程,让 socket IO 和 GPU 前向真正重叠;run_engine_core 是进程入口,负责实例化、信号注册、异常兜底和 shutdown。所有前端 (SyncMPClient / AsyncMPClient) 和后端 (EngineCore) 的通信都经过这一层。继续往里看可以去 /engine/engine-core 看引擎核心本身怎么装配,去 /engine/llm-engine 看前端薄壳,客户端三种实现见 /client/inproc-mp,调度器内部见 /scheduler/scheduler。

对照官方资料:vLLM 文档 · README。