Inferenza batch con Ray Data e vLLM

Questo esempio esegue l'inferenza batch LLM offline con Ray Data e vLLM su 8 GPU H100 in un singolo nodo. Uno script di bootstrap avvia un cluster Ray sul nodo, quindi il driver usa l'API LLM di Ray Data (ray.data.llm) per avviare una replica vLLM per ogni GPU e trasmettere un dataset di prompt attraverso tali repliche, scrivendo il testo generato in un volume di Unity Catalog in formato Parquet.

Usa un modello pubblico (Qwen2.5-7B-Instruct), quindi può essere eseguito così com'è senza bisogno di un token di Hugging Face.

Il carico di lavoro esegue le operazioni seguenti:

• Carica il progetto locale con code_source: snapshot.
• Avvia un'intestazione Ray con tutte e 8 le GPU, quindi esegue il driver di inferenza batch.
• Usa ray.data.llm per eseguire un'istanza vLLM per ogni GPU ed elaborare i prompt in parallelo.
• Scrive i prompt e gli output generati in un volume di Unity Catalog in formato Parquet.

Prerequisiti

• La air CLI è installata e autenticata. Consulta Installazione dell'AI Runtime CLI.
• Volume di Unity Catalog in cui è consentita la scrittura. Ne imposti il percorso nel file YAML del carico di lavoro qui sotto.

Layout del progetto

Creare una directory con i file seguenti.

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

Passaggio 1: Scrivere il carico di lavoro YAML

train.yaml richiede un singolo GPU_8xH100 nodo. Le dipendenze vengono dichiarate inline sotto environment (con l'immagine client version) e command avvia un cluster Ray sul nodo, quindi esegue il driver, perciò il carico di lavoro non richiede un file delle dipendenze o uno script di avvio separato.

VLLM non si trova nell'immagine di base, quindi viene installato inline insieme a tre pin necessari per i nodi GPU: hf_transfer (l'immagine di base consente download rapidi di Hugging Face e prevede questo pacchetto), una versione più recente fsspec (l'immagine di base fornisce un vecchio che interrompe i download) e un vLLM pull in opencv-python-headless OpenCV, il cui volante predefinito arresta il self-test FIPS OpenSSL nei nodi GPU.

Imposta OUTPUT_PATH su un volume di Unity Catalog su cui è possibile scrivere.

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

Il comando inline command avvia un head Ray con tutte le GPU sul nodo, esegue il driver con python batch_inference.py, quindi arresta il cluster. Include anche un ramo worker che si unisce all'head, così lo stesso comando continua a funzionare se si scala il job su più nodi.

Passaggio 2: Definire il driver di inferenza batch

batch_inference.py compila un set di dati Ray di richieste, configura un processore vLLM con ray.data.llme scrive i risultati. concurrency è il numero di repliche vLLM eseguite da Ray Data in parallelo. Impostandolo sul numero di GPU del cluster, viene assegnata una replica per GPU, quindi le richieste vengono elaborate contemporaneamente in ogni GPU e l'esempio viene ridimensionato man mano che si aggiungono nodi:

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 trasforma ogni riga di input in una richiesta di chat e postprocess mantiene le colonne in modo permanente. Ray Data aggiunge una generated_text colonna con l'output del modello. Lo script completo è incluso nello script completo del driver alla fine di questa pagina.

Per i modelli più grandi, impostare tensor_parallel_size in modo da suddividere una replica su più GPU e dividere total_gpus per tale valore, così che le repliche continuino a riempire il cluster, ad esempio concurrency=total_gpus // 2 con tensor_parallel_size=2.

Passaggio 3: Invia il run

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

Passaggio 4: Controllare l'esecuzione

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

I log mostrano il prompt e la velocità effettiva di generazione del motore vLLM durante l'esecuzione del batch, quindi una Wrote <n> rows riga quando viene scritto l'output.

Dove atterrare i risultati

Il driver scrive un dataset Parquet nel volume OUTPUT_PATH, con una colonna instruction e una colonna output. Rileggilo con Spark o pandas, ad esempio spark.read.parquet(OUTPUT_PATH).

Script completo del driver

Il file completo batch_inference.py per la copia-incolla:

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

Passaggi successivi

• Training distribuito con Ray Train
• Informazioni di riferimento su YAML del carico di lavoro