Inferencia por lotes con Ray Data y vLLM

Importante

Esta característica está en versión preliminar pública.

Este ejemplo ejecuta inferencia por lotes de LLM offline con Ray Data y vLLM en 8 nodos A10. Un script de arranque inicia un clúster Ray a través de los nodos, luego el controlador utiliza la API LLM de Ray Data (ray.data.llm) para lanzar una réplica vLLM por nodo y transmitir un conjunto de datos de prompts a través de ellos, escribiendo el texto generado en un volumen de Unity Catalog como Parquet.

Usa un modelo público (Qwen2.5-7B-Instruct), por lo que se ejecuta as-is sin un token de Hugging Face.

La carga de trabajo hace lo siguiente:

  • Carga el proyecto local con code_source: snapshot.
  • Inicia un cabezal de rayo en el nodo 0, une 7 nodos trabajadores y luego ejecuta el controlador de inferencia por lotes.
  • Se usa ray.data.llm para ejecutar una réplica de vLLM por nodo y procesar los prompts en paralelo.
  • Escribe las instrucciones y las salidas generadas en un volumen de Unity Catalog en formato Parquet.

Prerrequisitos

  • La CLI air está instalada y autenticada. Consulte Instalación de la CLI de AI Runtime.
  • Un volumen de Unity Catalog en el que se puede escribir. Puede establecer la ruta en el YAML de la carga de trabajo que aparece a continuación.

Diseño del proyecto

Cree un directorio con los siguientes archivos.

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

Paso 1: Escribir la carga de trabajo YAML

train.yaml Solicita 8 GPU_1xA10 nodos. Las dependencias se declaran en línea bajo environment (con la imagen versiondel cliente ), y luego command inicia un clúster Ray entre los nodos y luego ejecuta el controlador, por lo que la carga de trabajo no necesita un archivo de dependencia separado ni un script de lanzador.

vLLM no está en la imagen base, por lo que se instala directamente junto con tres versiones fijadas que necesitan los nodos con GPU: hf_transfer (la imagen base habilita descargas rápidas de Hugging Face y requiere este paquete), una versión más reciente de fsspec (la imagen base incluye una versión antigua que interrumpe las descargas) y una versión fijada de opencv-python-headless (vLLM instala OpenCV como dependencia, cuyo paquete wheel predeterminado hace fallar la autocomprobación FIPS de OpenSSL en los nodos con GPU).

Establezca OUTPUT_PATH en un volumen de catálogo de Unity en el que puede escribir. Especifique NUM_GPUS el mismo valor que 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

El comando en línea command inicia un nodo principal de Ray usando la GPU del nodo 0, ejecuta el programa controlador con python batch_inference.py y, a continuación, detiene el clúster. Los nodos de trabajo se conectan al nodo principal mediante MASTER_ADDR y NODE_RANK, que la plataforma establece automáticamente.

Paso 2: Definir el controlador de inferencia por lotes

batch_inference.py compila un conjunto de datos ray de mensajes, configura un procesador vLLM con ray.data.llmy escribe los resultados. concurrency es el número de réplicas de vLLM que Ray Data ejecuta en paralelo. El controlador espera a que todos los nodos se unan antes de leer el recuento de GPU, así que se usa cada nodo:

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 convierte cada fila de entrada en una solicitud de chat y postprocess mantiene las columnas para conservarlas. Ray Data agrega una generated_text columna con la salida del modelo. El script completo está en script de controlador completo al final de esta página.

Para los modelos más grandes, establezca tensor_parallel_size en particionar una réplica en varias GPU y divida total_gpus por ese valor para que las réplicas sigan rellenando el clúster, por ejemplo concurrency=total_gpus // 2 con tensor_parallel_size=2.

Paso 3: Enviar la ejecución

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

Paso 4: Inspección de la ejecución

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

Los registros muestran el rendimiento del prompt y de la generación del motor vLLM durante la ejecución del lote, y luego una línea Wrote <n> rows cuando se escribe la salida.

Dónde llegan los resultados

El controlador escribe un conjunto de datos Parquet en el volumen OUTPUT_PATH, con una columna instruction y una columna output. Léalo de nuevo con Spark o pandas, por ejemplo spark.read.parquet(OUTPUT_PATH).

Script de controlador completo

El batch_inference.py completo para copiar y pegar:

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

Pasos siguientes

  • Last updated on