Batchinferentie met Ray Data en vLLM

Important

Deze functie bevindt zich in openbare preview-versie.

Dit voorbeeld voert offline LLM-batchinferentie uit met Ray Data en vLLM over 8 A10-knooppunten. Een bootstrap-script start een Ray-cluster over de nodes, waarna de driver de LLM API (ray.data.llm) van Ray Data gebruikt om per node één vLLM-replica te starten en een dataset prompts erdoorheen te streamen, waarbij de gegenereerde tekst als Parquet naar een Unity Catalog-volume wordt geschreven.

Het gebruikt een openbaar model (Qwen2.5-7B-Instruct), dus het werkt direct zonder een Hugging Face-token.

De werkbelasting doet het volgende:

  • Uploadt het lokale project met code_source: snapshot.
  • Start een Ray-headnode op node 0, voegt vervolgens 7 workernodes toe en voert daarna de driver voor batchinferentie uit.
  • Gebruikt ray.data.llm om één vLLM-replica per knoop te draaien en prompts parallel te verwerken.
  • Schrijft de prompts en gegenereerde uitvoer naar een Unity Catalog-volume als Parquet.

Prerequisites

  • De air CLI is geïnstalleerd en geverifieerd. Zie De AI Runtime CLI installeren.
  • Een Unity Catalog-volume waarnaar u kunt schrijven. Je stelt het pad in de workload-YAML hieronder in.

Projectindeling

Maak een map met de volgende bestanden.

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

Stap 1: de YAML van de workload schrijven

train.yaml vraagt 8 GPU_1xA10 knooppunten aan. Afhankelijkheden worden inline onder environment opgegeven (met de clientimage version), en command start vervolgens een Ray-cluster over de knooppunten heen en voert daarna de driver uit, zodat de werklast geen apart afhankelijkhedenbestand of startscript nodig heeft.

vLLM bevindt zich niet in de basisinstallatiekopie, dus deze wordt inline geïnstalleerd, samen met drie pinnen die de GPU-knooppunten nodig hebben: hf_transfer (de basisinstallatiekopie maakt snelle Hugging Face-downloads mogelijk en verwacht dit pakket), een nieuwere fsspec (de basisinstallatiekopie verzendt een oude die downloads onderbreekt) en een vastgemaakte opencv-python-headless (vLLM haalt OpenCV op, waarvan het standaardwiel de Zelftest van OpenSSL FIPS op de GPU-knooppunten vastloopt).

Stel OUTPUT_PATH in op een Unity Catalog-volume waarnaar u kunt schrijven. Stel NUM_GPUS in op dezelfde waarde als 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

Het inline-commando command start een Ray-headnode met de GPU van het knooppunt op node 0, voert vervolgens de driver uit met python batch_inference.py en stopt daarna het cluster. Workerknooppunten maken verbinding met het hoofdknooppunt met behulp van MASTER_ADDR en NODE_RANK, die door het platform automatisch worden ingesteld.

Stap 2: Het batchinferentiestuurprogramma definiëren

batch_inference.py bouwt een Ray Dataset met prompts, configureert een vLLM-processor met ray.data.llmen schrijft de resultaten. concurrency is het aantal vLLM-replica's dat Ray Data parallel uitvoert. De driver wacht tot alle nodes zijn toegevoegd voordat hij de GPU-telling leest, dus elke node wordt gebruikt:

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 zet elke invoerrij om in een chataanvraag en postprocess houdt de kolommen behouden. Ray Data voegt een generated_text kolom toe met de uitvoer van het model. Het volledige script bevindt zich aan het einde van deze pagina in het volledige stuurprogrammascript .

Voor grotere modellen stelt u tensor_parallel_size zo in dat één replica over meerdere GPU's wordt verdeeld, en deelt u total_gpus door die waarde zodat de replica's het cluster nog steeds vullen, bijvoorbeeld concurrency=total_gpus // 2 met tensor_parallel_size=2.

Stap 3: De uitvoering verzenden

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

Stap 4: De uitvoering controleren

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

In de logboeken wordt de prompt- en generatiedoorvoer van de vLLM-engine weergegeven terwijl de batch wordt uitgevoerd en vervolgens een Wrote <n> rows regel wanneer de uitvoer wordt geschreven.

Waar resultaten terechtkomen

De driver schrijft één Parquet-dataset naar het volume OUTPUT_PATH, met een kolom instruction en een kolom output. Lees het terug met Spark of pandas, bijvoorbeeld spark.read.parquet(OUTPUT_PATH).

Volledig stuurprogrammascript

Het volledige batch_inference.py voor kopiëren en plakken:

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

Volgende stappen 

  • Last updated on