Ray Data：面向 AI 工作负载的可扩展数据处理#

Ray Data 是一个基于 Ray 构建的、面向 AI 工作负载的可扩展数据处理库。Ray Data 为常见的操作提供了灵活且高性能的 API，例如批量推理、数据预处理以及用于 ML 训练的数据加载。与其他分布式数据系统不同，Ray Data 具有一个流式执行引擎，可高效处理大型数据集并保持 CPU 和 GPU 工作负载的高利用率。

快速入门#

首先，安装 Ray Data。要了解有关安装 Ray 及其库的更多信息，请参阅安装 Ray

$ pip install -U 'ray[data]'

以下是一个使用 Ray Data 进行简单批量文本分类任务的示例

import ray
import pandas as pd

class ClassificationModel:
    def __init__(self):
        from transformers import pipeline
        self.pipe = pipeline("text-classification")

    def __call__(self, batch: pd.DataFrame):
        results = self.pipe(list(batch["text"]))
        result_df = pd.DataFrame(results)
        return pd.concat([batch, result_df], axis=1)

ds = ray.data.read_text("s3://anonymous@ray-example-data/sms_spam_collection_subset.txt")
ds = ds.map_batches(
    ClassificationModel,
    compute=ray.data.ActorPoolStrategy(size=2),
    batch_size=64,
    batch_format="pandas"
    # num_gpus=1  # this will set 1 GPU per worker
)
ds.show(limit=1)

{'text': 'ham\tGo until jurong point, crazy.. Available only in bugis n great world la e buffet... Cine there got amore wat...', 'label': 'NEGATIVE', 'score': 0.9935141801834106}

为什么选择 Ray Data？#

现代 AI 工作负载主要围绕深度学习模型的使用，这些模型计算密集且通常需要 GPU 等专用硬件。与 CPU 不同，GPU 通常内存较少，调度语义不同，并且运行成本更高。为支持传统数据处理管道而构建的系统通常无法很好地利用这些资源。

Ray Data 将 AI 工作负载作为一等公民进行支持，并提供了几项关键优势：

深度学习更快更便宜：Ray Data 在 CPU 预处理和 GPU 推理/训练任务之间流式传输数据，最大限度地提高资源利用率，并通过保持 GPU 处于活动状态来降低成本。
框架友好：Ray Data 为常见的 AI 框架（vLLM、PyTorch、HuggingFace、TensorFlow）和常见云提供商（AWS、GCP、Azure）提供了高性能的一等集成。
支持多模态数据：Ray Data 利用 Apache Arrow 和 Pandas，并支持 ML 工作负载中使用的许多数据格式，如 Parquet、Lance、图像、JSON、CSV、音频、视频等。
默认可扩展：基于 Ray 构建，可在具有不同 CPU 和 GPU 机器的异构集群之间自动扩展。代码可直接在单台机器或数百个节点上运行，处理数百 TB 的数据。