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

Important

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

此範例在 4 個 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 頭部，加入 3 個工作節點，然後執行批次推理驅動程式。
使用 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 請求 4 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

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

code_source:
  type: snapshot
  snapshot:
    root_path: .

command: |
  set -e
  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
    trap 'ray stop || true' EXIT
    python batch_inference.py
  else
    echo "NODE_RANK=$NODE_RANK: connecting to Ray head at $MASTER_ADDR:$RAY_HEAD_PORT..."
    joined=""
    for i in $(seq 1 12); do
      if ray start --address="$MASTER_ADDR:$RAY_HEAD_PORT" --num-gpus="$GPUS_PER_NODE" 2>/dev/null; then
        joined=1
        break
      fi
      echo "Attempt $i failed, retrying in 5s..."
      sleep 5
    done
    if [ -z "$joined" ]; then
      echo "Worker failed to join the Ray head after all retries." >&2
      exit 1
    fi
    echo "Worker joined. Waiting for the head to finish..."
    consecutive_failures=0
    for _ in $(seq 1 720); do
      if timeout 5 ray health-check --address "$MASTER_ADDR:$RAY_HEAD_PORT" 2>/dev/null; then
        consecutive_failures=0
      else
        consecutive_failures=$((consecutive_failures + 1))
        if [ "$consecutive_failures" -ge 3 ]; then
          echo "Head is no longer healthy. Stopping local Ray processes..."
          ray stop || true
          exit 0
        fi
        echo "Head health check failed ($consecutive_failures/3). Retrying..."
      fi
      sleep 5
    done
    echo "Timed out waiting for the Ray head to finish." >&2
    ray stop || true
    exit 1
  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: '4' # must match num_accelerators

內聯的 command 會在節點 0 上使用該節點的 GPU 啟動一個 Ray head 節點，然後用 python batch_inference.py 執行驅動程式。工作節點使用 MASTER_ADDR 和 NODE_RANK 加入 head 節點，而這些值會由平台自動設定。每位工人都會監控頭部，並在連續三次健康檢查失敗後停止本地 Ray 程序。

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

batch_inference.py 建立由提示詞組成的 Ray 資料集，使用 ray.data.llm 設定 vLLM 處理器，並將結果寫出。驅動程式會等所有節點都加入後才讀取 GPU 數量。 AIR 配置一個固定的加速器池，因此驅動程式設定 concurrency 為固定 (minimum, maximum) 元組，每個 GPU 請求一個副本。由於此範例使用短且固定的工作負載，驅動程式需等待最多 300 秒，讓所有複本初始化後才派遣工作。每個 Actor 最多可同時處理兩個批次，且已提交的 Ray Data 任務最多為兩個，包含執行中和佇列中的任務。這可防止第一個完成初始化的執行個體保留大部分的工作負載。這 2,000 個提示被分成 32 個輸入區塊，每個複製品有八個區塊可用。對於較長的工作負載，請根據啟動時間與吞吐量需求調整以下設定：

import os
import time

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

ray.init(address="auto")
data_context = DataContext.get_current()
data_context.wait_for_min_actors_s = 300

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}.")

ds = build_prompts().repartition(total_gpus * 8)

config = vLLMEngineProcessorConfig(
    model_source="Qwen/Qwen2.5-7B-Instruct",
    engine_kwargs={"max_model_len": 4096, "tensor_parallel_size": 1},
    concurrency=(total_gpus, total_gpus),
    batch_size=64,
    max_concurrent_batches=2,
    max_tasks_in_flight_per_actor=2,
)

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=1 將每個 vLLM 副本都放在一顆 A10 GPU 上。

步驟三：提交跑量

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 4 A10 nodes.

The workload `command` starts a Ray head on node 0 and joins 3 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 import DataContext
from ray.data.llm import build_processor, vLLMEngineProcessorConfig

MODEL_SOURCE = "Qwen/Qwen2.5-7B-Instruct"
NUM_PROMPTS = 2000
BATCH_SIZE = 64
BLOCKS_PER_REPLICA = 8
# 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")
    data_context = DataContext.get_current()
    data_context.wait_for_min_actors_s = 300

    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().repartition(total_gpus * BLOCKS_PER_REPLICA)

    # AIR provisions a fixed accelerator pool. Bound prefetching so the first ready
    # actor cannot reserve the small workload before the other actors initialize.
    config = vLLMEngineProcessorConfig(
        model_source=MODEL_SOURCE,
        engine_kwargs={
            "max_model_len": 4096,
            "tensor_parallel_size": 1,
            "enable_chunked_prefill": True,
        },
        concurrency=(total_gpus, total_gpus),
        batch_size=BATCH_SIZE,
        max_concurrent_batches=2,
        max_tasks_in_flight_per_actor=2,
    )

    # 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-09-03