Inferensi batch dengan Ray Data dan vLLM

Important

Fitur ini ada di Pratinjau Umum.

Contoh ini menjalankan inferensi batch LLM offline dengan Ray Data dan vLLM di 8 node A10. Skrip bootstrap memulai klaster Ray di seluruh node, lalu driver menggunakan API LLM Ray Data (ray.data.llm) untuk meluncurkan satu replika vLLM per node dan mengalirkan dataset prompt melaluinya, lalu menulis teks yang dihasilkan ke volume Unity Catalog dalam format Parquet.

Ini menggunakan model publik (Qwen2.5-7B-Instruct), sehingga dapat berjalan langsung tanpa token Hugging Face.

Beban kerja ini melakukan hal berikut:

• Mengunggah proyek lokal dengan code_source: snapshot.
• Memulai kepala Ray pada node 0, bergabung dengan 7 node pekerja, lalu menjalankan driver inferensi batch.
• Digunakan ray.data.llm untuk menjalankan satu replika vLLM per node dan memproses prompt secara paralel.
• Menulis prompt dan output yang dihasilkan ke volume Unity Catalog dalam format Parquet.

Prasyarat

• air CLI diinstal dan diautentikasi. Lihat Menginstal CLI Runtime AI.
• Volume Unity Catalog yang dapat Anda tulis. Anda mengatur jalurnya dalam YAML beban kerja di bawah ini.

Tata letak proyek

Buat direktori dengan file berikut.

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

Langkah 1: Tulis YAML beban kerja

train.yaml meminta 8 GPU_1xA10 node. Dependensi dideklarasikan secara inline di bawah environment (dengan citra klien version), dan command memulai cluster Ray di seluruh node lalu menjalankan driver, sehingga beban kerja tidak memerlukan berkas dependensi atau skrip peluncur terpisah.

vLLM tidak ada dalam image dasar, jadi diinstal secara inline bersama tiga versi paket yang dikunci yang dibutuhkan node GPU: hf_transfer (image dasar mengaktifkan unduhan Hugging Face yang cepat dan mengharapkan paket ini), fsspec yang lebih baru (image dasar menyertakan versi lama yang menyebabkan unduhan gagal), dan opencv-python-headless yang dikunci versinya (vLLM turut menarik dependensi OpenCV, yang wheel default-nya menyebabkan uji mandiri OpenSSL FIPS gagal pada node GPU).

Atur OUTPUT_PATH ke volume Unity Catalog yang dapat ditulisi. Atur NUM_GPUS ke nilai yang sama dengan 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

Perintah inline command memulai head Ray dengan GPU milik node tersebut pada node 0, menjalankan driver dengan python batch_inference.py, lalu menghentikan klaster. Node pekerja bergabung dengan node kepala menggunakan MASTER_ADDR dan NODE_RANK, yang diatur secara otomatis oleh platform.

Langkah 2: Tentukan driver inferensi batch

batch_inference.py membangun Himpunan Data Ray dari perintah, mengonfigurasi prosesor vLLM dengan ray.data.llm, dan menulis hasilnya. concurrency adalah jumlah replika vLLM yang dijalankan Ray Data secara paralel. Driver menunggu semua node bergabung sebelum membaca jumlah GPU, sehingga setiap node digunakan:

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 mengubah setiap baris input menjadi permintaan obrolan, dan postprocess mempertahankan kolom untuk bertahan. Ray Data menambahkan kolom generated_text dengan hasil model. Skrip lengkap ada di skrip Driver lengkap di akhir halaman ini.

Untuk model yang lebih besar, setel tensor_parallel_size agar satu replika dipecah ke beberapa GPU, lalu bagi total_gpus dengan nilai itu agar replika tetap memenuhi kluster, misalnya concurrency=total_gpus // 2 dengan tensor_parallel_size=2.

Langkah 3: Kirim eksekusi

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

Langkah 4: Periksa proses eksekusi

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

Log menunjukkan perintah mesin vLLM dan throughput pembuatan saat batch berjalan, lalu Wrote <n> rows garis saat output ditulis.

Tempat hasil ditampilkan

Driver tersebut menulis satu himpunan data Parquet ke volume OUTPUT_PATH, dengan kolom instruction dan kolom output. Baca kembali dengan Spark atau panda, misalnya spark.read.parquet(OUTPUT_PATH).

Skrip driver lengkap

batch_inference.py lengkap untuk salin-tempel:

#!/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()

Langkah berikutnya

• Pelatihan terdistribusi dengan Ray Train
• Referensi YAML beban kerja