Data-Juicer 分布式数据处理#

概览#

Data-Juicer 支持基于 Ray 和阿里巴巴 PAI 的大规模分布式数据处理。

经过专门的设计后，几乎所有在单机模式下实现的 Data-Juicer 算子都可以无缝地运行在 Ray 的分布式模式下。对于大规模场景，我们继续进行了针对计算引擎的特定优化，例如用于平衡文件和进程数目的数据子集分割策略，针对 Ray 和 Apache Arrow的 JSON 文件流式 I/O 补丁等。

作为参考，我们在 25 到 100 个阿里云节点上进行了实验，使用 Ray 模式下的 Data-Juicer 处理不同的数据集。在 6,400 个 CPU 核上处理包含 700 亿条样本的数据集只需要花费 2 小时，在 3,200 个 CPU 核上处理包含 70 亿条样本的数据集只需要花费 0.45 小时。此外，在 Ray 模式下，对 TB 大小级别的数据集，Data-Juicer 的 MinHash-LSH 去重算子在 1,280 个 CPU 核的 8 节点集群上进行去重只需 3 小时。

更多细节请参考我们的论文：Data-Juicer 2.0: Cloud-Scale Adaptive Data Processing for Foundation Models 。

实现与优化#

Data-Juicer 的 Ray 处理模式#

  • 对于 Data-Juicer 的大部分算子实现，其核心处理函数是引擎无关的。RayDataset 和 RayExecutor 封装了与Ray引擎的具体互操作，它们分别是基类 DJDataset 和 BaseExecutor 的子类，并且都支持 Ray Tasks 和 Actors 。

  • 其中，去重算子是例外。它们在单机模式下很难规模化。因此我们提供了针对它们的 Ray 优化版本算子，并以特殊前缀开头：ray_xx_deduplicator 。

数据子集分割#

当在上万个节点中处理仅有若干个文件的数据集时， Ray 会根据可用资源分割数据集文件，并将它们分发到所有节点上，这可能带来极大的网络通信开销并减少 CPU 利用率。更多细节可以参考文档 Ray's autodetect_parallelism 和 tuning output blocks for Ray 。

这种默认执行计划可能非常低效，尤其是在节点数量较多的情况下。为了优化此类情况的性能，我们考虑到 Ray 和 Arrow 的特性，提前将原始数据集自动拆分为较小的文件。当用户遇到此类性能问题时，他们可以利用此功能或根据偏好自己拆分数据集。在我们的自动拆分策略中，单个文件大小设置为 128MB，且结果应确保 拆分后的子文件数量 至少是 集群中可用CPU核心总数 的两倍。对应工具可在tools/data_resplit.py获取。

JSON 文件的流式读取#

为了解决 Ray Dataset 类底层框架 Arrow 对流式读取 JSON 数据的原生支持的缺失，我们开发了一个流式载入的接口并贡献到了一个针对 Apache Arrow 的内部 补丁（ 相关 PR ） 。这个补丁可以缓解内存不够的问题。

流式读取 JSON 文件是基础模型数据处理中的常见要求，因为许多数据集都以 JSONL 格式存储，并且尺寸巨大。 但是，Ray Datasets 中当前的实现不支持流式读取 JSON 文件，根因来源于其底层 Arrow 库（截至 Ray 版本 2.40 和 Arrow 版本 18.1.0）。

为了解决不支持流式 JSON 数据的原生读取问题，我们开发了一个流式加载接口，并为 Apache Arrow 贡献了一个第三方 补丁（PR 到 repo）。这将有助于缓解内存不足问题。使用此补丁后， Data-Juicer 的Ray模式将默认使用流式加载接口加载 JSON 文件。此外，如果输入变为 CSV 和 Parquet 文件，Ray模式下流式读取已经会自动开启。

去重#

在 Ray 模式下，我们提供了一个优化过的基于 MinHash-LSH 的去重算子。我们使用 Ray Actors 实现了一个多进程的并查集和一个负载均衡的分布式算法 BTS 来完成等价类合并操作。这个算子在 1,280 个CPU核上对 TB 大小级别的数据集去重只需要 3 个小时。我们的消融实验还表明相比于这个去重算子的初始实现版本，这些专门的优化项可以带来 2-3 倍的提速。

性能结果#

不同数据规模的数据处理#

我们在十亿样本规模的数据集上进行了实验。我们先准备了一个 56 万条样本的多模态数据集，并用不同的倍数（1-125,000倍）将其扩展来创建不同大小的数据集。下图的实验结果展示出了 Data-Juicer 的高扩展性。

Overview

大规模数据集分布式去重#

我们在 200GB、1TB、5TB 的数据集上测试了我们的基于 MinHash 的 Ray 去重算子，测试机器的 CPU 核数从 640 核到 1280 核。如下表所示，当数据集大小增长 5 倍，处理时间增长 4.02 到 5.62 倍。当 CPU 核数翻倍，处理时间较原来减少了 58.9% 到 67.1%。

CPU 核数

200GB 耗时

1TB 耗时

5TB 耗时

4 * 160

11.13 分钟

50.83 分钟

285.43 分钟

8 * 160

7.47 分钟

30.08 分钟

168.10 分钟

分布式数据分析#

除了分布式数据处理外，Data-Juicer 还通过 RayAnalyzer 支持分布式数据分析，它是本地 Analyzer 的分布式版本。

工作原理#

RayAnalyzer 使用与本地 Analyzer 相同的 dj-analyze 命令行入口。当配置文件中 executor_type 设置为 ray 时，dj-analyze 会自动调度到 RayAnalyzer。工作流程如下：

  1. 加载数据 — 通过 Ray 的分布式数据加载能力（支持 JSON、CSV、Parquet）。

  2. 计算统计信息 — 通过 Ray map_batches 分布式执行每个 Filter 的 compute_stats 函数，与本地模式使用相同的统计计算逻辑。

  3. 聚合总体统计 — 使用 Ray 原生聚合算子（Mean、Std、Min、Max）计算总体统计信息，无需 pandas 物化，适用于任意规模的数据集。

  4. 导出 — 通过 RayExporter 将带有统计列的数据集导出到磁盘。

与本地 Analyzer 的差异#

特性

本地 Analyzer

RayAnalyzer

总体统计（count/mean/std/min/max）

支持

支持

逐列分布图表

支持

不支持

相关性分析

支持

不支持

分位数

支持

不支持

可扩展性

单机

分布式（Ray 集群）

RayAnalyzer 专注于在大规模数据上计算总体统计信息。如果需要详细的可视化和相关性分析，建议对采样子集使用本地 Analyzer。

快速开始#

在开始前，你应该安装 Data-Juicer 以及它的 dist 依赖需求：

uv pip install -v -e .  # 安装 Data-Juicer 的最小依赖需求
uv pip install -v -e ".[dist]"  # 包括 Ray 以及其他分布式相关的依赖库

然后启动一个 Ray 集群（参考 Ray 文档 ）：

# 启动一个集群并作为头节点
ray start --head

# （可选）在其他节点或机器上连接集群
ray start --address='{head_ip}:6379'

我们在目录 demos/process_on_ray/ 中准备了简单的例子，包括 2 个配置文件和 2 个测试数据集。

demos/process_on_ray
├── configs
│   ├── demo.yaml
│   └── dedup.yaml
└── data
    ├── demo-dataset.json
    └── demo-dataset.jsonl

[!Important] 如果你要在多个节点上运行这些例子，你需要将示例数据集放置与一个共享磁盘（如 NAS）上，并且将结果数据集导出到那里。你可以通过修改配置文件中的 dataset_path 和 export_path 参数来实现。

运行 Ray 模式样例#

在配置文件 demo.yaml 中，我们将执行器类型设置为 "ray" 并且指定了自动的 Ray 地址。

...
dataset_path: './demos/process_on_ray/data/demo-dataset.jsonl'
export_path: './outputs/demo/demo-processed'

executor_type: 'ray'  # 将执行器类型设置为 "ray"
ray_address: 'auto'  # 设置为自动 Ray 地址
...

运行这个例子，以使用 12 个常规算子处理测试数据集：

# 从源码运行处理工具
python tools/process_data.py --config demos/process_on_ray/configs/demo.yaml

# 使用命令行工具
dj-process --config demos/process_on_ray/configs/demo.yaml

Data-Juicer 会使用示例配置文件处理示例数据集，并将结果数据集导出到配置文件中 export_path 参数指定的目录中。

运行分布式去重样例#

在配置文件 dedup.yaml 中，我们将执行器类型设置为 "ray" 并且指定了自动的 Ray 地址。我们使用了 MinHash 去重算子专门的分布式版本来对数据集去重。

project_name: 'demo-dedup'
dataset_path: './demos/process_on_ray/data/'
export_path: './outputs/demo-dedup/demo-ray-bts-dedup-processed'

executor_type: 'ray'  # 将执行器类型设置为 "ray"
ray_address: 'auto'  # 设置为自动 Ray 地址

# process schedule
# a list of several process operators with their arguments
process:
  - ray_bts_minhash_deduplicator:  # minhash 去重算子的分布式版本
      tokenization: 'character'

运行该实例来对数据集去重：

# 从源码运行处理工具
python tools/process_data.py --config demos/process_on_ray/configs/dedup.yaml

# 使用命令行工具
dj-process --config demos/process_on_ray/configs/dedup.yaml

Data-Juicer 会使用示例配置文件对示例数据集去重，并将结果数据集导出到配置文件中 export_path 参数指定的目录中。

运行 Ray 分析样例#

在 demos/analyze_simple/ray_analyzer.yaml 配置文件中，我们将执行器类型设置为 "ray" 来使用分布式的 RayAnalyzer。

project_name: 'demo-ray-analyzer'
dataset_path: './demos/data/demo-dataset.jsonl'

executor_type: 'ray'  # 将执行器类型设置为 "ray"
ray_address: 'auto'  # 设置为自动 Ray 地址

export_path: './outputs/demo-ray-analyzer'

process:
  - text_length_filter:
      min_len: 10
      max_len: 10000
  - words_num_filter:
      lang: en
      min_num: 10
      max_num: 10000
  - alphanumeric_filter:
      min_ratio: 0.25
      max_ratio: 0.9

运行该示例来分析数据集：

# 从源码运行分析工具
python tools/analyze_data.py --config demos/analyze_simple/ray_analyzer.yaml

# 使用命令行工具
dj-analyze --config demos/analyze_simple/ray_analyzer.yaml

RayAnalyzer 会计算所有数值统计列的总体统计信息（count、mean、std、min、max）并打印结果。带有计算好的统计列的数据集将导出到 export_path 指定的路径。