Inferência em lote com Dados de Raios e vLLM

Este exemplo executa inferência em lote offline de LLM com Ray Data e vLLM em 4 nós A10. Um script bootstrap inicia um cluster Ray entre os nós, depois o driver usa a API LLM da Ray Data (ray.data.llm) para lançar uma réplica vLLM por nó e transmitir um conjunto de dados de prompts através deles, escrevendo o texto gerado para um volume do Unity Catalog como Parquet.

Utiliza um modelo público (Qwen2.5-7B-Instruct), pelo que funciona tal como está sem um token da Hugging Face.

A carga de trabalho faz o seguinte:

• Carrega um snapshot do projeto local com code_source.
• Inicia um nó principal do Ray no nó 0, adiciona 3 nós de trabalho e, em seguida, executa o controlador de inferência em lote.
• Utiliza ray.data.llm para executar uma réplica de vLLM por nó e processar prompts em paralelo.
• Escreve os prompts e as saídas geradas num volume do Unity Catalog em formato Parquet.

Pré-requisitos

• A mais recente CLI Databricks instalada e autenticada. Consulte Instalar ou atualizar a CLI do Databricks e Autenticação para a CLI do Databricks.
• Um volume do Unity Catalog onde podes escrever. Define-se o respetivo caminho no ficheiro YAML da carga de trabalho abaixo.

Estrutura do projeto

Crie um diretório com os seguintes ficheiros.

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

Passo 1: Escrever a carga de trabalho em YAML

train.yaml solicita 4 GPU_1xA10 nós. As dependências são declaradas em linha em environment (com a imagem de cliente version), e o command inicia um cluster Ray nos vários nós e depois executa o driver, pelo que a carga de trabalho não precisa de um ficheiro de dependências separado nem de um script de arranque.

O vLLM não está na imagem base, por isso está instalado em linha juntamente com três pinos de que os nós da GPU precisam: hf_transfer (a imagem base permite downloads rápidos do Hugging Face e espera este pacote), um mais fsspec recente (a imagem base envia uma antiga que quebra os downloads) e um fixado opencv-python-headless (o vLLM puxa o OpenCV, cujo volante por defeito faz crashar o auto-teste do OpenSSL FIPS nos nós da GPU).

Define OUTPUT_PATH para um volume do Unity Catalog onde possas escrever. Definir NUM_GPUS para o mesmo 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:
  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

O inline command inicia uma cabeça Ray com a GPU do nó no nó 0, depois executa o driver com python batch_inference.py. Os nós de trabalho ligam-se ao nó principal usando MASTER_ADDR e NODE_RANK, que a plataforma define automaticamente. Cada trabalhador monitoriza a cabeça e interrompe os seus processos Ray locais após três falhas consecutivas nos testes de saúde.

Passo 2: Defina o driver de inferência por lote

batch_inference.py constrói um conjunto de dados Ray de prompts, configura um processador vLLM com ray.data.llm, e escreve os resultados. O controlador aguarda que todos os nós se liguem antes de ler o número de GPUs. O AIR prevê um conjunto fixo de aceleradores, de modo que o driver define concurrency para uma tupla fixa (minimum, maximum) que pede uma réplica por GPU. Como este exemplo utiliza uma carga de trabalho curta e fixa, o driver espera até 300 segundos para que todas as réplicas sejam inicializadas antes de despachar o trabalho. Cada ator processa até dois lotes em simultâneo e tem no máximo duas tarefas Ray Data submetidas, incluindo tarefas em execução e em fila. Isto impede que o primeiro ator a inicializar reserve a maior parte da carga de trabalho. Os 2.000 prompts estão divididos em 32 blocos de entrada, com oito blocos disponíveis por réplica. Para cargas de trabalho mais longas, ajuste estas definições com base no tempo de arranque e nos requisitos de rendimento:

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 transforma cada linha de entrada num pedido de chat e postprocess mantém as colunas para persistir. O Ray Data adiciona uma generated_text coluna com a saída do modelo. O script completo está em script completo do controlador no final desta página.

tensor_parallel_size=1 mantém cada réplica de vLLM numa GPU A10.

Passo 3: Submeter a execução

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

Passo 4: Inspecionar a pista

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

Os registos mostram o prompt do motor vLLM e o throughput de geração à medida que o batch é executado, depois uma Wrote <n> rows linha quando a saída é escrita.

Onde os resultados aparecem

O controlador escreve um conjunto de dados Parquet no volume OUTPUT_PATH, com uma coluna instruction e uma coluna output. Leia-o novamente com Spark ou pandas, por exemplo spark.read.parquet(OUTPUT_PATH).

Script completo do driver

A versão completa batch_inference.py para copiar e colar:

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

Passos seguintes

• Treino distribuído com Ray Train
• Referência YAML de carga de trabalho