Inferencia por lotes con Ray Data y vLLM

Este ejemplo ejecuta inferencia por lotes de LLM offline con Ray Data y vLLM a través de 4 nodos A10. Un script de arranque inicia un clúster de Ray en los nodos; a continuación, el controlador utiliza la API de LLM de Ray Data (ray.data.llm) para iniciar una réplica de vLLM por nodo y transmitir un conjunto de datos de indicaciones a través de ellas, escribiendo el texto generado en un volumen de Unity Catalog en formato 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 una cabeza de rayo en el nodo 0, une 3 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 4 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

# 4 A10 nodes, one GPU each. Ray Data runs one vLLM replica per node.
compute:
  num_accelerators: 4
  accelerator_type: GPU_1xA10

code_source:
  type: snapshot
  snapshot:
    root_path: .

command: |
  set -e
  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
    trap 'ray stop || true' EXIT
    python batch_inference.py
  else
    echo "NODE_RANK=$NODE_RANK: connecting to Ray head at $MASTER_ADDR:$RAY_HEAD_PORT..."
    joined=""
    for i in $(seq 1 12); do
      if ray start --address="$MASTER_ADDR:$RAY_HEAD_PORT" --num-gpus="$GPUS_PER_NODE" 2>/dev/null; then
        joined=1
        break
      fi
      echo "Attempt $i failed, retrying in 5s..."
      sleep 5
    done
    if [ -z "$joined" ]; then
      echo "Worker failed to join the Ray head after all retries." >&2
      exit 1
    fi
    echo "Worker joined. Waiting for the head to finish..."
    consecutive_failures=0
    for _ in $(seq 1 720); do
      if timeout 5 ray health-check --address "$MASTER_ADDR:$RAY_HEAD_PORT" 2>/dev/null; then
        consecutive_failures=0
      else
        consecutive_failures=$((consecutive_failures + 1))
        if [ "$consecutive_failures" -ge 3 ]; then
          echo "Head is no longer healthy. Stopping local Ray processes..."
          ray stop || true
          exit 0
        fi
        echo "Head health check failed ($consecutive_failures/3). Retrying..."
      fi
      sleep 5
    done
    echo "Timed out waiting for the Ray head to finish." >&2
    ray stop || true
    exit 1
  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: '4' # must match num_accelerators

La instrucción en línea command inicia un Ray Head con la GPU del nodo en el nodo 0 y, a continuación, ejecuta el controlador con python batch_inference.py. Los nodos de trabajo se conectan al nodo principal mediante MASTER_ADDR y NODE_RANK, que la plataforma establece automáticamente. Cada trabajador monitoriza la cabeza y detiene sus procesos Ray locales tras tres fallos consecutivos en los controles de salud.

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. El controlador espera a que todos los nodos se unan antes de leer el recuento de la GPU. AIR proporciona un conjunto fijo de aceleradores, por lo que el controlador establece concurrency como una tupla fija (minimum, maximum) que solicita una réplica por GPU. Como este ejemplo utiliza una carga de trabajo corta y fija, el controlador espera hasta 300 segundos para que todas las réplicas se inicialicen antes de despachar el trabajo. Cada actor procesa hasta dos lotes simultáneamente y tiene como máximo dos tareas de Ray Data enviadas, incluyendo tareas en ejecución y en cola. Esto impide que el primer actor en inicializar reserve la mayor parte de la carga de trabajo. Los 2.000 prompts están divididos en 32 bloques de entrada, con ocho bloques disponibles por réplica. Para cargas de trabajo más largas, ajusta estos ajustes según el tiempo de arranque y los requisitos de rendimiento:

import os
import time

import ray
from ray.data import DataContext
from ray.data.llm import build_processor, vLLMEngineProcessorConfig

ray.init(address="auto")
data_context = DataContext.get_current()
data_context.wait_for_min_actors_s = 300

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}.")

ds = build_prompts().repartition(total_gpus * 8)

config = vLLMEngineProcessorConfig(
    model_source="Qwen/Qwen2.5-7B-Instruct",
    engine_kwargs={"max_model_len": 4096, "tensor_parallel_size": 1},
    concurrency=(total_gpus, total_gpus),
    batch_size=64,
    max_concurrent_batches=2,
    max_tasks_in_flight_per_actor=2,
)

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.

tensor_parallel_size=1 mantiene cada réplica de vLLM en una GPU A10.

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 4 A10 nodes.

The workload `command` starts a Ray head on node 0 and joins 3 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 import DataContext
from ray.data.llm import build_processor, vLLMEngineProcessorConfig

MODEL_SOURCE = "Qwen/Qwen2.5-7B-Instruct"
NUM_PROMPTS = 2000
BATCH_SIZE = 64
BLOCKS_PER_REPLICA = 8
# 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")
    data_context = DataContext.get_current()
    data_context.wait_for_min_actors_s = 300

    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().repartition(total_gpus * BLOCKS_PER_REPLICA)

    # AIR provisions a fixed accelerator pool. Bound prefetching so the first ready
    # actor cannot reserve the small workload before the other actors initialize.
    config = vLLMEngineProcessorConfig(
        model_source=MODEL_SOURCE,
        engine_kwargs={
            "max_model_len": 4096,
            "tensor_parallel_size": 1,
            "enable_chunked_prefill": True,
        },
        concurrency=(total_gpus, total_gpus),
        batch_size=BATCH_SIZE,
        max_concurrent_batches=2,
        max_tasks_in_flight_per_actor=2,
    )

    # 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

• Entrenamiento distribuido con Ray Train
• Referencia de YAML para cargas de trabajo