ray-data

Ray Data - 可扩展的机器学习数据处理

用于机器学习和 AI 工作负载的分布式数据处理库。

何时使用 Ray Data

以下情况适合使用 Ray Data：

• 为机器学习训练处理大型数据集（>100GB）
• 需要在集群中进行分布式数据预处理
• 构建批量推理流水线
• 加载多模态数据（图像、音频、视频）
• 将数据处理从笔记本电脑扩展到集群

主要功能：

• 流式执行：处理超出内存容量的数据
• GPU 支持：使用 GPU 加速转换
• 框架集成：PyTorch、TensorFlow、HuggingFace
• 多模态：图像、Parquet、CSV、JSON、音频、视频

以下情况建议使用其他方案：

• Pandas：在单台机器上处理小型数据（<1GB）
• Dask：表格数据、类似 SQL 的操作
• Spark：企业级 ETL、SQL 查询

快速开始

安装

pip install -U 'ray[data]'

加载和转换数据

import ray

# Read Parquet files
ds = ray.data.read_parquet("s3://bucket/data/*.parquet")

# Transform data (lazy execution)
ds = ds.map_batches(lambda batch: {"processed": batch["text"].str.lower()})

# Consume data
for batch in ds.iter_batches(batch_size=100):
    print(batch)

与 Ray Train 集成

import ray
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer

# Create dataset
train_ds = ray.data.read_parquet("s3://bucket/train/*.parquet")

def train_func(config):
    # Access dataset in training
    train_ds = ray.train.get_dataset_shard("train")

    for epoch in range(10):
        for batch in train_ds.iter_batches(batch_size=32):
            # Train on batch
            pass

# Train with Ray
trainer = TorchTrainer(
    train_func,
    datasets={"train": train_ds},
    scaling_config=ScalingConfig(num_workers=4, use_gpu=True)
)
trainer.fit()

读取数据

从云存储读取

import ray

# Parquet (recommended for ML)
ds = ray.data.read_parquet("s3://bucket/data/*.parquet")

# CSV
ds = ray.data.read_csv("s3://bucket/data/*.csv")

# JSON
ds = ray.data.read_json("gs://bucket/data/*.json")

# Images
ds = ray.data.read_images("s3://bucket/images/")

从 Python 对象读取

# From list
ds = ray.data.from_items([{"id": i, "value": i * 2} for i in range(1000)])

# From range
ds = ray.data.range(1000000)  # Synthetic data

# From pandas
import pandas as pd
df = pd.DataFrame({"col1": [1, 2, 3], "col2": [4, 5, 6]})
ds = ray.data.from_pandas(df)

转换

映射批次（向量化）

# Batch transformation (fast)
def process_batch(batch):
    batch["doubled"] = batch["value"] * 2
    return batch

ds = ds.map_batches(process_batch, batch_size=1000)

行转换

# Row-by-row (slower)
def process_row(row):
    row["squared"] = row["value"] ** 2
    return row

ds = ds.map(process_row)

筛选

# Filter rows
ds = ds.filter(lambda row: row["value"] > 100)

分组和聚合

# Group by column
ds = ds.groupby("category").count()

# Custom aggregation
ds = ds.groupby("category").map_groups(lambda group: {"sum": group["value"].sum()})

GPU 加速转换

# Use GPU for preprocessing
def preprocess_images_gpu(batch):
    import torch
    images = torch.tensor(batch["image"]).cuda()
    # GPU preprocessing
    processed = images * 255
    return {"processed": processed.cpu().numpy()}

ds = ds.map_batches(
    preprocess_images_gpu,
    batch_size=64,
    num_gpus=1  # Request GPU
)

写入数据

# Write to Parquet
ds.write_parquet("s3://bucket/output/")

# Write to CSV
ds.write_csv("output/")

# Write to JSON
ds.write_json("output/")

性能优化

重新分区

# Control parallelism
ds = ds.repartition(100)  # 100 blocks for 100-core cluster

调整批次大小

# Larger batches = faster vectorized ops
ds.map_batches(process_fn, batch_size=10000)  # vs batch_size=100

流式执行

# Process data larger than memory
ds = ray.data.read_parquet("s3://huge-dataset/")
for batch in ds.iter_batches(batch_size=1000):
    process(batch)  # Streamed, not loaded to memory

常见模式

批量推理

import ray

# Load model
def load_model():
    # Load once per worker
    return MyModel()

# Inference function
class BatchInference:
    def __init__(self):
        self.model = load_model()

    def __call__(self, batch):
        predictions = self.model(batch["input"])
        return {"prediction": predictions}

# Run distributed inference
ds = ray.data.read_parquet("s3://data/")
predictions = ds.map_batches(BatchInference, batch_size=32, num_gpus=1)
predictions.write_parquet("s3://output/")

数据预处理流水线

# Multi-step pipeline
ds = (
    ray.data.read_parquet("s3://raw/")
    .map_batches(clean_data)
    .map_batches(tokenize)
    .map_batches(augment)
    .write_parquet("s3://processed/")
)

与机器学习框架集成

PyTorch

# Convert to PyTorch
torch_ds = ds.to_torch(label_column="label", batch_size=32)

for batch in torch_ds:
    # batch is dict with tensors
    inputs, labels = batch["features"], batch["label"]

TensorFlow

# Convert to TensorFlow
tf_ds = ds.to_tf(feature_columns=["image"], label_column="label", batch_size=32)

for features, labels in tf_ds:
    # Train model
    pass

支持的数据格式

格式	读取	写入	使用场景
Parquet	✅	✅	机器学习数据（推荐）
CSV	✅	✅	表格数据
JSON	✅	✅	半结构化数据
图像	✅	❌	计算机视觉
NumPy	✅	✅	数组
Pandas	✅	❌	DataFrame

性能基准

扩展性（处理 100GB 数据）：

• 1 个节点（16 个核心）：约 30 分钟
• 4 个节点（64 个核心）：约 8 分钟
• 16 个节点（256 个核心）：约 2 分钟

GPU 加速（图像预处理）：

• 仅 CPU：每秒 1,000 张图像
• 1 个 GPU：每秒 5,000 张图像
• 4 个 GPU：每秒 18,000 张图像

使用场景

生产部署：

• Pinterest：用于模型训练的最后一公里数据处理
• ByteDance：通过多模态 LLM 扩展离线推理
• Spotify：用于批量推理的机器学习平台

参考资料

• 转换指南 - 映射、筛选、分组操作
• 集成指南 - Ray Train、PyTorch、TensorFlow

资源

• 文档：https://docs.ray.io/en/latest/data/data.html
• GitHub：https://github.com/ray-project/ray ⭐ 36,000+
• 版本：Ray 2.40.0+
• 示例：https://docs.ray.io/en/latest/data/examples/overview.html