Batchinferentie met Ray Data en vLLM

In dit voorbeeld wordt offline LLM-batchdeductie uitgevoerd met Ray Data en vLLM over 8 H100 GPU's op één knooppunt. Een bootstrapscript start een Ray-cluster op het knooppunt en vervolgens gebruikt het stuurprogramma de LLM-API (ray.data.llm) van Ray Data om één vLLM-replica per GPU te starten en een gegevensset met prompts erdoor te streamen, waardoor de gegenereerde tekst naar een Unity Catalog-volume wordt geschreven als Parquet.

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-hoofd met alle 8 GPU's en voert vervolgens het batchdeductiestuurprogramma uit.
• Wordt gebruikt ray.data.llm om één vLLM-replica per GPU uit te voeren en procesprompts parallel uit te voeren.
• 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 één GPU_8xH100 knooppunt aan. Afhankelijkheden worden inline opgegeven onder environment (met de clientimage version), en command start een Ray-cluster op het knooppunt en voert daarna de driver uit, zodat de workload geen afzonderlijk afhankelijkhedenbestand of launcherscript 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.

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

De inline command start een Ray-hoofd met alle GPU's op het knooppunt, voert het stuurprogramma uit met python batch_inference.pyen stopt vervolgens het cluster. Het bevat ook een worker-branch die zich bij de head voegt, zodat hetzelfde commando blijft werken als u de taak opschaalt naar meerdere knooppunten.

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. Als u dit instelt op het AANTAL GPU's van het cluster, krijgt u één replica per GPU, zodat de prompts in elke GPU tegelijk worden verwerkt en het voorbeeld wordt geschaald wanneer u knooppunten toevoegt:

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 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 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()

Volgende stappen 

• Gedistribueerde training met Ray Train
• YAML-verwijzing voor workload