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