使用 Ray Data 與 vLLM 進行批次推論

Important

這項功能目前處於公開預覽版。

此範例在 8 個 A10 節點上運行離線批次推論，使用 Ray Data 與 vLLM 。啟動腳本會在各節點間啟動 Ray 叢集，接著由驅動程式使用 Ray Data 的 LLM API（ray.data.llm），在每個節點上啟動一個 vLLM 複本，並讓提示資料集以串流方式通過這些複本，然後將產生的文字以 Parquet 格式寫入 Unity Catalog 磁碟區。

它使用公開模型（Qwen2.5-7B-Instruct），因此無需 Hugging Face 權杖即可直接執行。

工作量會做出以下效果：

使用 code_source: snapshot 上傳本地專案。
在節點 0 啟動一個 Ray 頭部，加入 7 個工作節點，然後執行批次推理驅動程式。
使用 ray.data.llm 在每個節點上執行一個 vLLM 複本，並平行處理提示。
以 Parquet 格式將提示詞和產生的輸出寫入 Unity Catalog 磁碟區。

先決條件

air CLI 已安裝並完成驗證。請參閱安裝 AI 執行階段 CLI。
可寫入的 Unity Catalog 磁碟區。你可以在下面的工作負載 YAML 裡設定它的路徑。

專案版面配置

建立一個包含以下檔案的目錄。

ray_batch_inference/
├── train.yaml            # air workload config (inline dependencies + Ray bootstrap)
└── batch_inference.py    # Ray Data + vLLM batch inference driver

步驟 1：撰寫 YAML 工作負載

train.yaml 請求 8 GPU_1xA10 個節點。相依性會直接在 version 下宣告（搭配用戶端映像 environment），而 command 會先在各節點上啟動 Ray 叢集，再執行 driver，因此工作負載不需要額外的相依性檔案或啟動指令碼。

vLLM 不在基礎映像裡，所以它是內建安裝的，並且和 GPU 節點需要的三個腳位一起安裝： hf_transfer （基礎映像檔支援快速的 Hugging Face 下載，並期待這個套件）、一個較 fsspec 新的（基礎映像檔出貨一個舊的，會破壞下載），以及一個釘 opencv-python-headless 選（vLLM 會拉入 OpenCV，OpenCV 的預設輪輪會讓 GPU 節點的 OpenSSL FIPS 自我測試當機）。

將 OUTPUT_PATH 設為您可寫入的 Unity Catalog 磁碟區。設定 NUM_GPUS 為與相同的值 num_accelerators。

experiment_name: air-ray-batch-inference

environment:
  version: '5'
  dependencies:
    - ray[data]==2.56.1
    - vllm
    - datasets>=3.0
    - huggingface_hub>=0.34
    # The base image sets HF_HUB_ENABLE_HF_TRANSFER=1; install the package it expects
    # so model and dataset downloads don't error out.
    - hf_transfer
    # The base image ships fsspec 2023.5.0, which is too old for modern
    # huggingface_hub and breaks dataset/model downloads. Pin a newer fsspec.
    - fsspec>=2024.6.1
    # vLLM pulls in opencv; its default wheel crashes the OpenSSL FIPS self-test
    # on the GPU nodes. This pinned headless build avoids the crash.
    - opencv-python-headless==4.12.0.88

# 8 A10 nodes, one GPU each. Ray Data runs one vLLM replica per node.
compute:
  num_accelerators: 8
  accelerator_type: GPU_1xA10

code_source:
  type: snapshot
  snapshot:
    root_path: .

command: |
  cd $CODE_SOURCE_PATH
  RAY_HEAD_PORT=6379
  GPUS_PER_NODE=${LOCAL_WORLD_SIZE:-1}
  if [ "${NODE_RANK:-0}" = "0" ]; then
    echo "NODE_RANK=0: starting Ray head with $GPUS_PER_NODE GPU(s)..."
    ray start --head --port=$RAY_HEAD_PORT --num-gpus="$GPUS_PER_NODE" --dashboard-host=0.0.0.0
    python batch_inference.py
    ray stop
  else
    echo "NODE_RANK=$NODE_RANK: connecting to Ray head at $MASTER_ADDR:$RAY_HEAD_PORT..."
    for i in $(seq 1 12); do
      if ray start --address="$MASTER_ADDR:$RAY_HEAD_PORT" --num-gpus="$GPUS_PER_NODE" --block 2>/dev/null; then
        break
      fi
      echo "Attempt $i failed, retrying in 5s..."
      sleep 5
    done
  fi

max_retries: 0
timeout_minutes: 60
env_variables:
  NCCL_SOCKET_IFNAME: eth0
  # Unity Catalog volume where results land as Parquet. Replace with your volume.
  OUTPUT_PATH: /Volumes/main/default/air_examples/ray_batch_inference
  NUM_GPUS: '8' # must match num_accelerators

內嵌的 command 會在節點 0 上使用該節點的 GPU 啟動一個 Ray head 節點，接著使用 python batch_inference.py 執行驅動程式，然後停止叢集。工作節點使用 MASTER_ADDR 和 NODE_RANK 加入 head 節點，而這些值會由平台自動設定。

步驟 2：定義批次推論驅動程式

batch_inference.py 建立由提示詞組成的 Ray 資料集，使用 ray.data.llm 設定 vLLM 處理器，並將結果寫出。 concurrency 是 Ray Data 平行執行的 vLLM 複本數量。驅動程式會等待所有節點加入後再讀取 GPU 計數，因此每個節點都會被使用：

import os
import time

import ray
from ray.data.llm import build_processor, vLLMEngineProcessorConfig

ray.init(address="auto")
num_gpus = int(os.environ["NUM_GPUS"])
for _ in range(60):
    if int(ray.cluster_resources().get("GPU", 0)) >= num_gpus:
        break
    time.sleep(5)
total_gpus = int(ray.cluster_resources().get("GPU", 0))
if total_gpus < num_gpus:
    raise SystemExit(f"Expected {num_gpus} GPU(s) but Ray only sees {total_gpus}.")

config = vLLMEngineProcessorConfig(
    model_source="Qwen/Qwen2.5-7B-Instruct",
    engine_kwargs={"max_model_len": 4096, "tensor_parallel_size": 1},
    concurrency=total_gpus,   # one vLLM replica per GPU in the cluster
    batch_size=64,
)

processor = build_processor(
    config,
    preprocess=lambda row: dict(
        messages=[{"role": "user", "content": row["instruction"]}],
        sampling_params=dict(max_tokens=256, temperature=0.7),
    ),
    postprocess=lambda row: dict(instruction=row["instruction"], output=row["generated_text"]),
)

out = processor(ds)       # ds is a Ray Dataset with an "instruction" column
out.write_parquet(OUTPUT_PATH)

preprocess 將每個輸入列都轉為聊天請求，並 postprocess 保留欄位的持續存在。 Ray Data 會新增一個 generated_text 欄位，其中包含模型的輸出。完整腳本以完整驅動腳本格式在本頁末尾。

對於較大的模型，將 tensor_parallel_size 設為把一個副本分片到多個 GPU，並將 total_gpus 除以該值，使副本仍能填滿叢集，例如 concurrency=total_gpus // 2 搭配 tensor_parallel_size=2。

步驟三：提交跑量

air run -f train.yaml --dry-run
air run -f train.yaml --watch

步驟 4：檢查執行結果

air get run <run-id>
air logs <run-id>

日誌顯示批次執行時 vLLM 引擎的提示音與產生吞吐量，然後在寫入輸出時顯示一 Wrote <n> rows 行。

結果落在哪裡

驅動程式會將一個 Parquet 資料集寫入 OUTPUT_PATH 磁碟區，其中包含 instruction 資料行和 output 資料行。例如，可用 Spark 或 pandas 將它讀取回來 spark.read.parquet(OUTPUT_PATH)。

完整驅動腳本

完整的 batch_inference.py，供複製貼上：

#!/usr/bin/env python3
"""Offline batch inference with Ray Data + vLLM across 8 A10 nodes.

The workload `command` starts a Ray head on node 0 and joins 7 worker nodes, each
contributing 1 GPU. Ray Data's LLM API (`ray.data.llm`) launches one vLLM replica
per GPU and streams a dataset of prompts through them, then writes the generated text
to a Unity Catalog volume as Parquet.

Uses a public model (no Hugging Face token required) so the example runs as-is.
"""

import os
import time

import ray
from datasets import load_dataset
from ray.data.llm import build_processor, vLLMEngineProcessorConfig

MODEL_SOURCE = "Qwen/Qwen2.5-7B-Instruct"
NUM_PROMPTS = 1000
# Unity Catalog volume path where results land as Parquet. Set this in train.yaml.
OUTPUT_PATH = os.environ.get("OUTPUT_PATH", "/Volumes/main/default/air_examples/ray_batch_inference")


def build_prompts():
    """Build a Ray Dataset of prompts from a public instruction dataset."""
    raw = load_dataset("tatsu-lab/alpaca", split=f"train[:{NUM_PROMPTS}]")
    items = []
    for row in raw:
        instruction = row["instruction"]
        if row.get("input"):
            instruction = f"{instruction}\n\n{row['input']}"
        items.append({"instruction": instruction})
    return ray.data.from_items(items)


def main():
    ray.init(address="auto")
    num_gpus = int(os.environ["NUM_GPUS"])
    for _ in range(60):
        if int(ray.cluster_resources().get("GPU", 0)) >= num_gpus:
            break
        time.sleep(5)
    total_gpus = int(ray.cluster_resources().get("GPU", 0))
    if total_gpus < num_gpus:
        raise SystemExit(
            f"Expected {num_gpus} GPU(s) but Ray only sees {total_gpus}; "
            "check GPU discovery / node join on all nodes."
        )
    print(f"Ray cluster ready: {total_gpus} GPU(s)", flush=True)

    ds = build_prompts()

    # vLLM engine config. concurrency = number of replicas Ray Data runs in parallel;
    # one per GPU in the cluster here. engine_kwargs are passed through to the vLLM engine.
    config = vLLMEngineProcessorConfig(
        model_source=MODEL_SOURCE,
        engine_kwargs={
            "max_model_len": 4096,
            "tensor_parallel_size": 1,
            "enable_chunked_prefill": True,
        },
        concurrency=total_gpus,
        batch_size=64,
    )

    # preprocess maps each input row to a chat request; postprocess keeps the columns
    # we want to persist. ray.data.llm adds a `generated_text` column.
    processor = build_processor(
        config,
        preprocess=lambda row: dict(
            messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": row["instruction"]},
            ],
            sampling_params=dict(max_tokens=256, temperature=0.7),
        ),
        postprocess=lambda row: dict(
            instruction=row["instruction"],
            output=row["generated_text"],
        ),
    )

    # materialize once so the write and the sample print don't re-run inference.
    out = processor(ds).materialize()
    out.write_parquet(OUTPUT_PATH)
    print(f"Wrote {out.count()} rows to {OUTPUT_PATH}", flush=True)

    for row in out.take(2):
        print("INSTRUCTION:", row["instruction"][:120], flush=True)
        print("OUTPUT:", row["output"][:200], flush=True)

    ray.shutdown()


if __name__ == "__main__":
    main()

下一步

意見反應

此頁面對您有幫助嗎？

Last updated on 2026-08-21