算子插件
September 8, 2026 · View on GitHub
概述
Data-Juicer 支持通过 Python 标准的 entry points 机制发现并加载来自外部包的算子。这使得算子可以作为独立的 Python 包进行分发、版本管理和安装(例如发布到 PyPI 或私有索引),同时仍然注册到全局 OPERATORS 注册表中,并可在任意 recipe 中通过名称引用。
工作原理
在 import data_juicer.ops 时,load_op_plugins() 会扫描所有已安装发行包中声明在 data_juicer.ops 分组下的 entry points,并导入其指向的模块。导入插件模块会执行其模块级的 @OPERATORS.register_module(...) 装饰器,在 init_configs() 读取注册表之前完成算子注册。
关键特性:
- 自动发现:插件包安装后,其算子即可使用,无需额外配置。
- 故障隔离:若某个插件导入失败(例如缺少依赖),仅记录警告并跳过,不影响其余流程的执行。
- 向后兼容:未安装任何插件时,行为保持不变。
插件 vs. custom_operator_paths
Data-Juicer 提供两种互补的方式来使用库外算子:
| 算子插件(entry points) | custom_operator_paths | |
|---|---|---|
| 分发方式 | 可安装的包(PyPI / 私有索引) | 本地 .py 文件或包目录 |
| 发现方式 | 安装后自动发现 | 在 CLI / YAML 中显式指定路径 |
| 适用场景 | 可复用、可版本化、需要共享的算子 | 快速的本地或一次性算子 |
| 配置 | 无需额外配置 | --custom-operator-paths 或 YAML 中的 custom_operator_paths: |
开发算子插件
1. 包结构
my-dj-ops/
├── pyproject.toml
└── my_dj_ops/
└── __init__.py # 定义并注册算子
2. 实现并注册算子
算子的编写方式与内置算子一致:继承基类(Mapper、Filter、Deduplicator 等),并使用 @OPERATORS.register_module(<op_name>) 装饰器注册。重型依赖应通过 LazyLoader 延迟加载,避免在模块导入时直接引入。
# my_dj_ops/__init__.py
from data_juicer.ops.base_op import OPERATORS, Mapper
@OPERATORS.register_module("my_upper_mapper")
class MyUpperMapper(Mapper):
"""将每个样本的文本转为大写。"""
_batched_op = True
def process_batched(self, samples):
samples[self.text_key] = [t.upper() for t in samples[self.text_key]]
return samples
3. 声明 Entry Point
在 pyproject.toml 中,将模块暴露到 data_juicer.ops 分组下。entry point 的取值需指向一个模块(或对象),其导入会触发 @OPERATORS.register_module 调用——指向包的 __init__ 是最简单的做法。
[project]
name = "my-dj-ops"
version = "0.1.0"
dependencies = ["py-data-juicer"]
[project.entry-points."data_juicer.ops"]
my_dj_ops = "my_dj_ops"
4. 安装并使用
pip install -e . # 或:pip install my-dj-ops
然后即可在任意 recipe 中按名称引用该算子,default 和 Ray 执行器均支持:
process:
- my_upper_mapper: {}
注意事项与最佳实践
- 算子名称唯一性:注册的算子名(如
my_upper_mapper)不得与内置算子或其他插件冲突。 - 声明
py-data-juicer依赖:在插件的dependencies中列出,以确保基类与注册表可用。 - 重型依赖保持延迟加载:将重型 ML 库写入插件的
dependencies,但在运行时通过LazyLoader加载,与核心算子保持相同的约定。 - GPU 与不可 fork 算子:插件中可照常使用
_accelerator = "cuda"或use_cuda(),执行器会根据这些属性选择相应的多进程上下文。