Inférence par lots avec Ray Data et vLLM

Important

Cette fonctionnalité est en version bêta. Les administrateurs d’espace de travail peuvent contrôler l’accès à cette fonctionnalité à partir de la page Aperçus . Consultez Gérer les préversions d’Azure Databricks.

Cet exemple exécute l’inférence de lot LLM hors connexion avec Ray Data et vLLM sur 8 GPU H100 sur un seul nœud. Un script d’amorçage lance un cluster Ray sur le nœud, puis le programme pilote utilise l’API LLM de Ray Data (ray.data.llm) pour lancer un réplica vLLM par GPU et diffuser un jeu de données de requêtes à travers celles-ci, en écrivant le texte généré dans un volume Unity Catalog au format Parquet.

Il utilise un modèle public (Qwen2.5-7B-Instruct), de sorte qu’il s’exécute as-is sans jeton Hugging Face.

La charge de travail effectue les opérations suivantes :

  • Charge le projet local avec code_source: snapshot.
  • Démarre une tête Ray avec 8 GPU, puis exécute le pilote d’inférence par lots.
  • Utilise ray.data.llm pour exécuter une instance vLLM par GPU et traiter les invites en parallèle.
  • Écrit les prompts et les résultats générés dans un volume Unity Catalog au format Parquet.

Prerequisites

  • L’interface air CLI a été installée et authentifiée. Consultez Installer l’interface CLI d’AI Runtime.
  • Un volume de catalogue Unity dans lequel vous pouvez écrire. Vous définissez son chemin dans la charge de travail YAML ci-dessous.

Disposition du projet

Créez un répertoire avec les fichiers suivants.

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

Étape 1 : Écrire la charge de travail YAML

train.yaml demande un nœud unique GPU_8xH100 . Les dépendances sont déclarées directement sous environment (avec l’image cliente version), et command démarre un cluster Ray sur le nœud, puis exécute le driver, de sorte que la charge de travail n’a pas besoin d’un fichier de dépendances ni d’un script de lancement distinct.

vLLM n’est pas inclus dans l’image de base ; il est donc installé en ligne, avec trois dépendances épinglées nécessaires aux nœuds GPU : hf_transfer (l’image de base permet des téléchargements rapides depuis Hugging Face et attend ce paquet), une version plus récente de fsspec (l’image de base inclut une version ancienne qui empêche les téléchargements) et une version épinglée de opencv-python-headless (vLLM dépend d’OpenCV, dont le wheel par défaut provoque l’échec du test FIPS OpenSSL sur les nœuds GPU).

Définissez OUTPUT_PATH sur un volume Unity Catalog sur lequel vous disposez d’un accès en écriture.

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

L’inline command démarre une tête Ray avec toutes les GPU sur le nœud, exécute le pilote avec python batch_inference.py, puis arrête le cluster. Inclut également une branche worker qui rejoint le nœud principal, de sorte que la même commande continue de fonctionner si vous exécutez la tâche sur plusieurs nœuds.

Étape 2 : Définir le pilote d’inférence par lots

batch_inference.py crée un jeu de données Ray de prompts, configure un processeur vLLM avec ray.data.llm et écrit les résultats. concurrency est le nombre de réplicas vLLM que Ray Data exécute en parallèle. Le fait de le définir sur le nombre de GPU du cluster donne une réplique par GPU, de sorte que les requêtes sont réparties simultanément sur tous les GPU et que l’exemple passe à l’échelle à mesure que vous ajoutez des nœuds :

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 transforme chaque ligne d’entrée en requête de chat et postprocess conserve les colonnes. Ray Data ajoute une generated_text colonne avec la sortie du modèle. Le script complet se trouve dans le script de pilote complet à la fin de cette page.

Pour les modèles plus volumineux, définissez tensor_parallel_size de façon à répartir une réplique sur plusieurs GPU, et divisez total_gpus par cette valeur afin que les répliques occupent toujours tout le cluster, par exemple concurrency=total_gpus // 2 avec tensor_parallel_size=2.

Étape 3 : Envoyer l’exécution

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

Étape 4 : Inspecter l’exécution

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

Les journaux affichent la requête et le débit de génération du moteur vLLM pendant l’exécution du lot, puis une ligne Wrote <n> rows quand la sortie est enregistrée.

Où s’affichent les résultats

Le pilote écrit un jeu de données Parquet sur le volume OUTPUT_PATH, avec une colonne instruction et une colonne output. Vous pouvez les relire avec Spark ou pandas, par exemple spark.read.parquet(OUTPUT_PATH).

Script de pilote complet

L’intégralité batch_inference.py à copier-coller :

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

Étapes suivantes

  • Last updated on