RayDistributedExecutor:用 Ray 把 worker 拉到多节点多卡

源码版本v0.25.1

职责

单卡场景下 UniProcExecutor 一个进程就把事干完了,但只要 TP>1 或跨节点,就必须有进程间通信。RayDistributedExecutor(ray_executor.py:64-68)就是 Ray 后端的实现:它把每个 worker 包成一个 Ray actor,分布到 placement group 指定的 GPU 上,然后用 Ray Compiled Graph 构一条流水线把前向串起来。uses_ray = True、supports_pp = True(ray_executor.py:67-68),既表明自己用 Ray 编排,又支持流水线并行 (pipeline parallelism)。

它跟 UniProcExecutor 一样继承 Executor,基类的 execute_model/sample_tokens/profile/sleep 这些方法签名都不改,只重写 _init_executor、collective_rpc、execute_model、check_health、shutdown(ray_executor.py:70-97)。区别只在 RPC 怎么走:同进程是直接调方法,Ray 是 worker.execute_method.remote(...)(ray_executor.py:484-489)。

设计动机

• Ray actor 一一对应 worker:每个 GPU 上一个 RayWorkerWrapper actor(ray_executor.py:200-213),Ray 负责调度、容错、资源记账,executor 不自己 spawn 进程。
• placement group 锁定拓扑:actor 用 PlacementGroupSchedulingStrategy 绑到 bundle 上(ray_executor.py:192-196),GPU 亲缘关系在启动前就定好,运行期不会漂。
• 同节点优先排序:actor 创建顺序是 Ray 决定的(可能乱),所以建完所有 actor 后按 IP 重新排序,跟 driver 同节点的排前面(ray_executor.py:234-258),再 collective_rpc("adjust_rank", args=(rerank_mapping,)) 把 rank 拨正。
• Ray Compiled Graph 跑前向:不是每步 RPC 调 execute_model,而是建一条 forward_dag(ray_executor.py:527-620),PP=2/TP=4 时把四个 worker 串成 InputNode -> TP group -> MultiOutputNode,中间张量用 with_tensor_transport(transport="nccl"|"shm") 走 NCCL 或共享内存(ray_executor.py:582-591)。
• 环境变量自动复制:driver 进程里跟 vLLM 相关的环境变量会批量 update_environment_variables 推到每个 worker(ray_executor.py:305-327),WORKER_SPECIFIC_ENV_VARS 除外,避免 worker 间互相污染。
• execute_model 拆两段:开了 sampler 的话,execute_model 不立刻跑,只存 self.scheduler_output,等 sample_tokens 来了再 _execute_dag(ray_executor.py:390-432),让调度器和 sampler 的 grammar bitmask 计算可以跟 forward 重叠。

关键文件

• RayDistributedExecutor._init_executor:64-97 — initialize_ray_cluster、拿 placement group、调 _init_workers_ray,初始化 has_connector、uses_sampler、scheduler_output 占位。
• shutdown:99-113 — forward_dag.teardown() 后逐个 ray.kill(worker),日志提示 SIGTERM 是预期行为。
• _init_workers_ray:143-213 — 按 bundle_indices 逐个 ray.remote(...)(RayWorkerWrapper).remote(rpc_rank=rank),GPU 用 num_gpus=num_gpus,其他平台走 resources={ray_device_key: num_gpus}。
• rerank:217-258 — ray.get 拿每个 actor 的 node IP,按「同 driver 节点优先 → 节点 actor 数少优先 → IP 小优先」排序,然后 collective_rpc("adjust_rank", ...) 改 rpc_rank。
• init_worker + load_model:343-368 — 拼好 all_kwargs 后一次性 collective_rpc("init_worker") + collective_rpc("init_device") + collective_rpc("load_model")。
• pp_tp_workers:370-378 — 按 PP×TP 二维填 self.pp_tp_workers,给后面建 compiled DAG 用。
• execute_model + sample_tokens:390-432 — 两阶段:execute_model 只存 scheduler_output,真正跑在 sample_tokens 里走 _execute_dag。
• _execute_dag:434-468 — 首次调用时 self.forward_dag = self._compiled_ray_dag(...),forward_dag.execute((scheduler_output, grammar_output)) 返回 refs,PP 不开时 refs[0].get() 阻塞拿结果,非阻塞则返回 FutureWrapper。
• collective_rpc:470-495 — 把 method 序列化(cloudpickle.dumps),逐 worker execute_method.remote(...),然后 ray.get(ray_worker_outputs, timeout=timeout);non_block=True 返回 FutureWrapper。
• _compiled_ray_dag:527-620 — 用 InputNode + worker.execute_model_ray.bind(...) 组 DAG,PP 之间用 with_tensor_transport(transport) 传张量,最后 experimental_compile(...)。
• check_health:625-628 — 现版本直接 return,TODO 留着以后做真正的 Ray actor 健康检查。
• RayWorkerWrapper:56-91 — 继承 WorkerWrapperBase,加 adjust_rank、get_node_ip、get_node_and_physical_gpu_ids、execute_method(把异常包成「可能死锁」的日志再抛)。

数据流

启动期 _init_executor 走 _init_workers_ray,这一段最长:选 bundle、按 placement group 创建 RayWorkerWrapper actor,然后 ray.get 拿每个 actor 的 node IP 给排序、rerank:

# vllm/v1/executor/ray_executor.py L191-L213
for rank, bundle_id in enumerate(bundle_indices):
    scheduling_strategy = PlacementGroupSchedulingStrategy(
        placement_group=placement_group,
        placement_group_capture_child_tasks=True,
        placement_group_bundle_index=bundle_id,
    )

    if current_platform.ray_device_key == "GPU":
        worker = ray.remote(
            num_cpus=0,
            num_gpus=num_gpus,
            scheduling_strategy=scheduling_strategy,
            **ray_remote_kwargs,
        )(RayWorkerWrapper).remote(rpc_rank=rank)
    else:
        worker = ray.remote(
            num_cpus=0,
            num_gpus=0,
            resources={current_platform.ray_device_key: num_gpus},
            scheduling_strategy=scheduling_strategy,
            **ray_remote_kwargs,
        )(RayWorkerWrapper).remote(rpc_rank=rank)

    worker_metadata.append(RayWorkerMetaData(worker=worker, created_rank=rank))

运行期 EngineCore.step 调到 execute_model 时,如果 uses_sampler=True 且本批有 token,executor 只暂存 scheduler_output 并返回(ray_executor.py:401-407);sample_tokens 被调到时才真正走 _execute_dag,首次会建 compiled DAG:

# vllm/v1/executor/ray_executor.py L434-L452
def _execute_dag(self, scheduler_output, grammar_output, non_block=False):
    if self.forward_dag is None:
        self.forward_dag = self._compiled_ray_dag(enable_asyncio=False)

    refs = self.forward_dag.execute((scheduler_output, grammar_output))

    if not self.has_connector:
        if not non_block:
            output = refs[0].get()
            detach_zero_copy_from_model_runner_output(output)
            return output
        return FutureWrapper(refs[0])
    ...

控制平面的 RPC(initialize_from_config、compile_or_warm_up_model、determine_available_memory 等)走的都是 collective_rpc,即逐 worker execute_method.remote(...)(ray_executor.py:484-489),再 ray.get 收回来。这跟数据平面(forward 走 compiled DAG)是两套独立通路。

边界与失败

• 节点 IP 必须唯一:_init_workers_ray 显式检查 n_nodes == n_ips,不满足就抛 RuntimeError,并提示查 VLLM_HOST_IP(ray_executor.py:295-303)。
• 单节点用回环地址:同节点多 actor 时 driver_ip = "127.0.0.1",避免 get_ip() 取到错误的网卡导致 worker 间通信失败(ray_executor.py:329-341)。
• 状态机顺序错就报错:execute_model 看到 self.scheduler_output is not None 直接抛「sample_tokens() must be called after execute_model() returns None」(ray_executor.py:395-399)。
• Ray Compiled Graph 版本和依赖:_check_ray_cgraph_installation 要求 ray>=2.43.0、ray[cgraph]、可选 cupy(ray_executor.py:497-525),否则启动期就 fail-fast。
• TPU/XPU 走共享内存:_init_executor 检测到 TPU/XPU 时强制 VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE=shm(ray_executor.py:73-75),避开 NVIDIA NCCL。
• check_health 当前是空:RayDistributedExecutor.check_health 直接 return(ray_executor.py:625-628),actor 真挂了要靠 ray.get 抛异常上抛,目前没有主动心跳。
• shutdown 强 kill:forward_dag.teardown() 之后 ray.kill(worker)(ray_executor.py:111-112),worker 内未落盘状态会丢,文档里特别说明 SIGTERM 日志是预期的。

小结

RayDistributedExecutor 是 v1 在多节点多卡场景下的执行器实现。它把每个 worker 包成 Ray actor,靠 placement group 锁定拓扑,用 Ray Compiled Graph 建一条前向 DAG 把 TP/PP 串起来;控制平面 RPC 还是走 actor 的 execute_method。它跟 UniProcExecutor 暴露给 EngineCore 的接口完全一致,只是底层换了一套 RPC 机制。执行器怎么抽象出来见 /executor/executor,worker 内部怎么跑前向见 /worker/worker 和 /worker/gpu-model-runner。

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