Пакетный вывод с помощью Ray Data и vLLM

Important

Эта функция доступна в общедоступной предварительной версии.

В этом примере выполняется пакетный вывод LLM в офлайн-режиме с помощью Ray Data и vLLM на 4 узлах A10. Bootstrap-скрипт запускает кластер Ray на всех узлах, после чего драйвер использует API LLM в Ray Data (ray.data.llm), чтобы запустить по одной реплике vLLM на каждом узле и пропустить через них набор данных с запросами, записывая сгенерированный текст в том Unity Catalog в формате Parquet.

Она использует публичную модель (Qwen2.5-7B-Instruct), поэтому работает как есть без токена Hugging Face.

Рабочая нагрузка выполняет следующие действия:

  • Загружает локальный проект с помощью code_source: snapshot.
  • Запускает лучевую головку на узле 0, соединяет 3 рабочих узла, затем запускает драйвер пакетного вывода.
  • Используется ray.data.llm для запуска одной реплики vLLM на узел и параллельной обработки запросов.
  • Записывает запросы и созданные выходные данные в том каталога Unity в виде Parquet.

Необходимые условия

Структура проекта

Создайте каталог со следующими файлами.

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

Шаг 1. Напишите YAML-файл рабочей нагрузки

train.yaml запрашивает 4 GPU_1xA10 узла. Зависимости объявляются непосредственно под environment (с клиентским образом version), а command запускает кластер Ray на всех узлах, а затем драйвер, поэтому рабочей нагрузке не требуется отдельный файл зависимостей или скрипт запуска.

vLLM отсутствует в базовом образе, поэтому он устанавливается прямо здесь вместе с тремя жёстко закреплёнными зависимостями, которые нужны узлам GPU: hf_transfer (базовый образ включает быстрые загрузки из Hugging Face и ожидает наличие этого пакета), более новая версия fsspec (в базовом образе поставляется старая версия, которая ломает загрузку), и закреплённая версия opencv-python-headless (vLLM подтягивает OpenCV, а его wheel-пакет по умолчанию вызывает сбой самопроверки OpenSSL FIPS на узлах GPU).

Задайте OUTPUT_PATH для тома каталога Unity, на который можно написать. Установите NUM_GPUS то же значение, что и 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

Эта однострочная команда command запускает головной узел Ray на узле 0, используя его GPU, а затем запускает драйвер с помощью python batch_inference.py. Рабочие узлы присоединяются к головному узлу, используя MASTER_ADDR и NODE_RANK, которые автоматически задаются платформой. Каждый работник контролирует головку и останавливает локальные лучевые процессы после трёх подряд неудачных проверок здоровья.

Шаг 2. Определить драйвер пакетного инференса

batch_inference.py создает набор данных Ray из промптов, настраивает процессор vLLM с помощью ray.data.llm и записывает результаты. Драйвер ждёт, пока все узлы соединятся, прежде чем считывать счёт GPU. AIR предоставляет фиксированный пул ускорителей, поэтому драйвер задаёт concurrency как фиксированный кортеж (minimum, maximum), который запрашивает по одной реплике на каждый GPU. Поскольку в этом примере используется короткая, фиксированная нагрузка, драйвер ждёт до 300 секунд, пока все реплики инициализируются перед отправкой работ. Каждый актор одновременно обрабатывает до двух пакетов и имеет не более двух сданных задач Ray Data, включая выполненные и поставленные в очередь задачи. Это мешает первому актору, инициализировавшемусь, зарезервировать большую часть рабочей нагрузки. 2000 подсказок разделены на 32 входных блока, при этом на каждую реплику доступно восемь блоков. Для более длительных рабочих нагрузок настройте эти параметры в зависимости от времени запуска и требований к пропускной способности:

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 преобразует каждую входную строку в запрос чата и postprocess сохраняет столбцы. Ray Data добавляет generated_text столбец с выходными данными модели. Полный скрипт находится в скрипте полного драйвера в конце этой страницы.

tensor_parallel_size=1 размещает каждый экземпляр vLLM на одном GPU A10.

Шаг 3: Отправьте запуск

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

Шаг 4. Проверка выполнения

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

В журналах отображаются показатели пропускной способности промпта и генерации для ядра vLLM во время выполнения пакета, а затем — строка Wrote <n> rows, когда записываются выходные данные.

Куда попадают результаты

Драйвер записывает один набор данных Parquet в OUTPUT_PATH том с столбцом instruction и столбцом output . Считайте это обратно с помощью Spark или pandas, например spark.read.parquet(OUTPUT_PATH).

Полный скрипт драйвера

Полный текст batch_inference.py для копирования и вставки:

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

Дальнейшие действия

  • Last updated on