Ray Data ve vLLM ile toplu çıkarım

Important

Bu özellik Genel Önizleme aşamasındadır.

Bu örnek, 4 A10 düğümünde Ray Data ve vLLM ile çevrimdışı LLM toplu çıkarımı çalıştırır. Bir önyükleme betiği, düğümler genelinde bir Ray kümesi başlatır; ardından driver, Ray Data'nın LLM API'sini (ray.data.llm) kullanarak düğüm başına bir vLLM kopyası başlatır ve istemlerden oluşan bir veri kümesini bunlar üzerinden akıtarak oluşturulan metni Parquet biçiminde bir Unity Catalog birimine yazar.

Herkese açık bir model (Qwen2.5-7B-Instruct) kullanır; bu nedenle Hugging Face erişim belirtecine gerek duymadan olduğu gibi çalışır.

İş yükü aşağıdakileri yapar:

Yerel projeyi code_source: snapshot ile karşıya yükler.
0. düğümde bir Ray baş düğümü başlatır, 3 işçi düğümü ekler ve ardından toplu çıkarım sürücüsünü çalıştırır.
ray.data.llm kullanarak düğüm başına bir vLLM replikası çalıştırır ve istemleri paralel olarak işler.
İstemleri ve üretilen çıktıları Unity Catalog birimine Parquet olarak yazar.

Prerequisites

air CLI yüklendi ve kimliği doğrulandı. Bkz. AI Runtime CLI'yı yükleme.
Yazma erişiminizin olduğu bir Unity Catalog depo birimi. Yolunu aşağıdaki iş yükü YAML dosyasında belirtirsiniz.

Proje düzeni

Aşağıdaki dosyaları içeren bir dizin oluşturun.

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

1. Adım: İş yükü YAML'sini yazma

train.yaml 4 GPU_1xA10 düğüm talep ediyor. Bağımlılıklar, version altında (istemci görüntüsü environment ile) satır içinde tanımlanır ve command, düğümler arasında bir Ray kümesi başlatıp ardından sürücüyü çalıştırır; böylece iş yükü ayrı bir bağımlılık dosyasına veya başlatıcı betiğe ihtiyaç duymaz.

vLLM temel imajda bulunmadığından, GPU düğümlerinin ihtiyaç duyduğu üç sürüm sabitlemesiyle birlikte satır içinde kuruluyor: hf_transfer (temel imaj, Hugging Face indirmelerini hızlandırır ve bu paketi bekler), daha yeni bir fsspec (temel imaj, indirmeleri bozan eski bir sürümle gelir) ve sabitlenmiş bir opencv-python-headless (vLLM, OpenCV'yi bağımlılık olarak getirir; bunun varsayılan wheel paketi GPU düğümlerinde OpenSSL FIPS öz testinin çökmesine neden olur).

OUTPUT_PATH öğesini, yazma izninizin olduğu bir Unity Catalog birimi olarak ayarlayın. NUM_GPUS öğesini, num_accelerators ile aynı değere ayarlayın.

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

Satır içi command, 0 numaralı düğümde düğümün GPU'sunu kullanarak bir Ray head düğümü başlatır, ardından sürücüyü python batch_inference.py ile çalıştırır. İşçi düğümleri, platformun otomatik olarak ayarladığı NODE_RANK ve MASTER_ADDR kullanarak ana düğüme katılır. Her çalışan kafayı izler ve üç üst üste sağlık kontrolü başarısız olduktan sonra yerel Ray süreçlerini durdurur.

2. Adım: Toplu çıkarım sürücüsünü tanımlama

batch_inference.py istemlerden oluşan bir Ray Veri Kümesi oluşturur, ile ray.data.llmbir vLLM işlemcisi yapılandırıp sonuçları yazar. Sürücü, tüm düğümlerin katılmasını bekler, sonra GPU sayısını okuyur. AIR sabit bir hızlandırıcı havuzu sunar, bu nedenle sürücü, GPU başına bir replika isteyen sabit bir (minimum, maximum) demeti concurrency ayarlar. Bu örnek kısa ve sabit bir iş yükü kullandığı için, sürücü tüm replikaların başlatılması için 300 saniyeye kadar bekler. Her aktör aynı anda en fazla iki toplu işlemi işler ve çalışan ve kuyrukta bekleyen görevler dahil olmak üzere en fazla iki gönderilmiş Ray Data görevine sahip olur. Bu, ilk başlatılan aktörün iş yükünün büyük kısmını kendine ayırmasını engeller. 2.000 istem 32 giriş bloğuna ayrılmış, her replikada sekiz blok mevcuttur. Daha uzun iş yükleri için, bu ayarları başlatma süresi ve veri kapasitesi gereksinimlerine göre ayarlayın:

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 her giriş satırını bir sohbet isteğine dönüştürür ve postprocess sütunların kalıcı olmasını sağlar. Ray Data, modelin çıkışını içeren bir generated_text sütun ekler. Betiğin tamamı, bu sayfanın sonundaki Tam sürücü betiğindedir .

tensor_parallel_size=1 Her vLLM replikasını bir A10 GPU'da tutuyor.

3. Adım: Çalıştırmayı gönderme

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

4. Adım: Çalıştırmayı inceleyin

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

Günlükler, toplu işlem çalışırken vLLM motorunun istem ve üretim verimini, ardından çıktı yazıldığında bir Wrote <n> rows satırını gösterir.

Sonuçların geldiği yer

Sürücü, OUTPUT_PATH birimine, instruction sütunu ve output sütunu içeren bir Parquet veri kümesi yazar. Örneğin, Spark ya da pandas kullanarak bunu yeniden okuyun spark.read.parquet(OUTPUT_PATH).

Tam sürücü betiği

Kopyala-yapıştır için eksiksiz 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()

Sonraki Adımlar

Geri Bildirim

Bu sayfayı yararlı buldunuz mu?

Last updated on 2026-09-03