Inferencia por lotes con Ray Data y vLLM

En este ejemplo se ejecuta la inferencia por lotes LLM sin conexión con Ray Data y vLLM en 8 GPU H100 en un solo nodo. Un script de arranque inicia un clúster de Ray en el nodo y, a continuación, el proceso controlador usa la API de LLM de Ray Data (ray.data.llm) para lanzar una réplica de vLLM por GPU y transmitir a través de estas un conjunto de datos de indicaciones, 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 un nodo principal de Ray con las 8 GPU y, a continuación, ejecuta el controlador de inferencia por lotes.
• Usa ray.data.llm para ejecutar una réplica de vLLM por GPU y procesar solicitudes 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 un único GPU_8xH100 nodo. Las dependencias se declaran en línea en environment (con la imagen versionde cliente) y command inicia un clúster de Ray en el nodo y, a continuación, ejecuta el controlador, por lo que la carga de trabajo no necesita un archivo de dependencia o script de iniciador independiente.

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.

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

La instrucción inline command inicia un nodo principal de Ray con todas las GPU del nodo, ejecuta el proceso controlador con python batch_inference.py y, a continuación, detiene el clúster. También incluye una rama de nodo de trabajo que se conecta al nodo principal, por lo que el mismo comando sigue funcionando si se escala el trabajo a varios nodos.

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. Si se establece en el número de GPU del clúster, se obtiene una réplica por cada GPU, por lo que los prompts se procesan de forma simultánea en todas las GPU y el ejemplo escala a medida que añade nodos:

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

Pasos siguientes

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