Ray Data ve vLLM ile toplu çıkarım

Important

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

Bu örnek, tek bir düğümde 8 adet H100 GPU üzerinde Ray Data ve vLLM ile çevrimdışı LLM toplu çıkarımı çalıştırır. Bootstrap betiği düğüm üzerinde bir Ray kümesi başlatır; ardından sürücü, GPU başına bir vLLM çoğaltısı başlatmak ve istemlerden oluşan bir veri kümesini bunlar üzerinden akış halinde geçirip oluşturulan metni Parquet biçiminde bir Unity Catalog birimine yazmak için Ray Data'nın LLM API'sini (ray.data.llm) kullanır.

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.
8 GPU'yu da içeren bir Ray head başlatır, ardından toplu çıkarım sürücüsünü çalıştırır.
GPU başına bir vLLM çoğaltması çalıştırmak ve istemleri paralel olarak işlemek için kullanır ray.data.llm .
İ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 tek GPU_8xH100 bir düğüm isteğinde bulunur. Bağımlılıklar, environment altında satır içi olarak bildirilir (istemci imajı version ile) ve command, düğüm üzerinde bir Ray kümesi başlatıp ardından sürücüyü çalıştırır; bu nedenle iş yükü için ayrı bir bağımlılık dosyası veya başlatma betiği gerekmez.

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.

experiment_name: air-ray-batch-inference

environment:
  version: '4'
  dependencies:
    - ray[data]>=2.44
    - 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 H100 on a single node. Ray Data runs one vLLM replica per GPU.
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100

code_source:
  type: snapshot
  snapshot:
    root_path: .

command: |
  cd $CODE_SOURCE_PATH
  RAY_HEAD_PORT=6379
  GPUS_PER_NODE=${LOCAL_WORLD_SIZE:-8}
  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

Satır içi command, düğümdeki tüm GPU'larla bir Ray head düğümü başlatır, sürücüyü python batch_inference.py ile çalıştırır ve ardından kümeyi durdurur. Ayrıca, ana düğüme katılan bir çalışan dalı da içerir; böylece işi birden fazla düğüme ölçeklendirseniz bile aynı komut çalışmaya devam eder.

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. concurrency , Ray Data'nın paralel çalıştığı vLLM çoğaltmalarının sayısıdır. Bu değeri kümenin GPU sayısı olarak ayarlamak, GPU başına bir replika sağlar; böylece istemler aynı anda tüm GPU'lara dağıtılarak işlenir ve düğüm ekledikçe örnek de ölçeklenir:

from ray.data.llm import build_llm_processor, vLLMEngineProcessorConfig

# Read the GPU count from the live Ray cluster so concurrency scales with the cluster.
total_gpus = int(ray.cluster_resources().get("GPU", 0))

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_llm_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 .

Daha büyük modeller için, tensor_parallel_size değerini bir replikayı birkaç GPU'ya bölecek şekilde ayarlayın ve replikaların kümeyi doldurmaya devam etmesi için total_gpus değerini bu değere bölün; örneğin, concurrency=total_gpus // 2 ile tensor_parallel_size=2.

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 on a single 8x H100 node.

The workload `command` starts a Ray head with 8 GPUs and runs this script. 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 ray
from datasets import load_dataset
from ray.data.llm import build_llm_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")
    # Derive replicas from the live cluster so the example scales when nodes are added.
    total_gpus = int(ray.cluster_resources().get("GPU", 0))
    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_llm_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-07-16