ray-data

Ray Data - 可扩展的 ML 数据处理

用于 ML 和 AI 工作负载的分布式数据处理库。

何时使用 Ray Data

在以下情况下使用 Ray Data：

• 处理用于 ML 训练的大型数据集（>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/")
)

与 ML 框架集成

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	✅	✅	ML 数据（推荐）
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：用于批量推理的 ML 平台

参考资料

• 转换指南 - 映射、筛选、分组操作
• 集成指南 - 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

Agent 操作流程

1. 检查环境。 确认 GPU 内存、模型检查点和许可证、框架版本以及输入格式。
2. 明确输入。 根据数据或用户确认格式、标识符和参数。任何会改变结果的值都应询问，不要猜测。
3. 先运行小规模版本。 先以低分辨率或较短长度对单个样本运行推理。
4. 执行完整任务，遵循上述说明和参考资料。
5. 验证结果。 使用任务指标进行定性检查和评估；检查生成内容的许可证限制。
6. 报告。 说明运行了什么（版本、命令、参数）、检查了什么，以及仍有哪些不确定之处。

如果发生这种情况	请这样做
大尺寸输入导致内存不足	降低分辨率或序列长度，使用半精度，或对输入进行分块处理。
这些说明中的函数、标志或端点在已安装的版本中不存在	查阅该已安装版本自己的文档（help()、--help、官方文档），进行调整并告知用户。切勿编造 API。
必需的输入、标识符或参数含义不明确	询问用户，或在运行前明确说明所作假设。

诚信规则

• 切勿捏造结果、参数、标识符、引用或统计数据。如果某项内容无法运行或验证，请直言说明。
• 商业使用或发布前，请检查模型和数据集的许可证。
• 本文中的版本特定信息可能已过时：请对照已安装版本的官方文档进行确认。
• 对于会产生费用、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，请先征求同意。

相关技能

• ray-train：跨集群的分布式训练编排。
• polars：用于 Python ETL、分析和 pandas 迁移的高性能 DataFrame 库。
• dask：用于处理规模超过内存的 pandas/NumPy 工作流的分布式计算库。