Papermill扩展指南：通过Entry Points实现自定义I/O和引擎

Papermill扩展指南：通过Entry Points实现自定义I/O和引擎

papermill 📚 Parameterize, execute, and analyze notebooks 项目地址: https://gitcode.com/gh_mirrors/pa/papermill

什么是Entry Points？

在Python生态系统中，Entry Points是一种强大的插件机制，允许已安装的包向系统注册其提供的功能组件。对于Papermill项目而言，Entry Points机制使得开发者能够扩展其核心功能，特别是：

1. 自定义I/O处理器：支持从不同存储系统读写笔记本文件
2. 自定义执行引擎：改变笔记本的执行方式和行为

这种设计遵循了"开放-封闭"原则，使得Papermill核心保持稳定，同时通过插件机制实现功能扩展。

开发自定义I/O处理器

I/O处理器的作用

Papermill执行流程的第一步就是读取输入笔记本文件，最后一步则是写入执行后的笔记本。I/O处理器抽象了这些操作，使得Papermill不仅支持本地文件系统，还能对接各种存储服务。

实现规范

一个完整的I/O处理器需要实现以下四个核心方法：

class CustomIO:
    @classmethod
    def read(cls, file_path):
        """读取指定路径的文件内容并返回"""
        pass
    
    @classmethod
    def write(cls, file_content, file_path):
        """将内容写入指定路径"""
        pass
    
    @classmethod
    def pretty_path(cls, path):
        """返回美化后的路径表示"""
        return path
    
    @classmethod
    def listdir(cls, path):
        """列出指定路径下的内容"""
        pass

注意事项：

• 如果某些操作不需要支持（如只读系统），相应方法应抛出明确异常
• 路径处理应采用统一约定，通常使用类似URL的格式（如s3://bucket/path）

注册处理器

开发完成后，需要在包的setup.py中声明entry point：

setup(
    # ...其他配置
    entry_points={
        "papermill.io": [
            "customproto://=mypackage.module:CustomIOHandler"
        ]
    }
)

其中customproto://是协议前缀，当Papermill遇到以此开头的路径时，就会使用对应的处理器类。

实战：SFTP处理器示例

下面我们实现一个完整的SFTP处理器，支持通过SFTP协议访问远程笔记本文件。

项目结构

papermill_sftp/
├── setup.py
└── src/
    └── papermill_sftp/
        └── __init__.py

核心实现

# src/papermill_sftp/__init__.py
import os
import tempfile
import pathlib
import urllib.parse
import pysftp

class SFTPHandler:
    @classmethod
    def read(cls, path):
        """从SFTP服务器读取笔记本文件"""
        parsed = urllib.parse.urlparse(path)
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp_path = pathlib.Path(tmpdir) / pathlib.Path(parsed.path).name
            with cls._create_sftp_connection(parsed) as sftp:
                sftp.get(parsed.path, str(tmp_path))
            return tmp_path.read_text()

    @classmethod
    def write(cls, content, path):
        """将笔记本写入SFTP服务器"""
        parsed = urllib.parse.urlparse(path)
        with tempfile.TemporaryDirectory() as tmpdir:
            tmp_path = pathlib.Path(tmpdir) / "output.ipynb"
            tmp_path.write_text(content)
            with cls._create_sftp_connection(parsed) as sftp:
                sftp.put(str(tmp_path), parsed.path)

    @classmethod
    def _create_sftp_connection(cls, parsed_url):
        """创建SFTP连接"""
        return pysftp.Connection(
            host=parsed_url.hostname,
            username=os.getenv('SFTP_USERNAME'),
            password=os.getenv('SFTP_PASSWORD'),
            port=parsed_url.port or 22
        )

    @classmethod
    def pretty_path(cls, path):
        return path

    @classmethod
    def listdir(cls, path):
        raise NotImplementedError("SFTP目录列表暂未实现")

打包配置

# setup.py
from setuptools import setup, find_packages

setup(
    name="papermill_sftp",
    version="0.1",
    packages=find_packages("./src"),
    package_dir={"": "src"},
    install_requires=["pysftp"],
    entry_points={
        "papermill.io": ["sftp://=papermill_sftp:SFTPHandler"]
    }
)

安装后，用户即可使用sftp://协议前缀访问SFTP服务器上的笔记本文件。

开发自定义执行引擎

引擎的作用

执行引擎负责实际运行笔记本代码。默认引擎使用本地内核执行，但通过自定义引擎可以实现：

• 远程执行（集群、云环境）
• 执行前后处理（如数据注入、结果分析）
• 特殊执行模式（如参数化多次运行）

实现规范

自定义引擎需继承papermill.engines.Engine基类，并实现核心方法：

from papermill.engines import Engine

class CustomEngine(Engine):
    @classmethod
    def execute_managed_notebook(cls, nb_man, kernel_name, **kwargs):
        """
        nb_man: 笔记本管理器对象
        kernel_name: 内核名称
        kwargs: 其他执行参数
        """
        # 自定义执行逻辑

实战：执行时间记录引擎

下面实现一个在单元格输出中添加执行时间的引擎。

项目结构

papermill_timing/
├── setup.py
└── src/
    └── papermill_timing/
        └── __init__.py

核心实现

# src/papermill_timing/__init__.py
from datetime import datetime
from papermill.engines import NBClientEngine
from nbformat.v4 import new_output

class TimingEngine(NBClientEngine):
    @classmethod
    def execute_managed_notebook(cls, nb_man, kernel_name, **kwargs):
        # 先执行原始引擎逻辑
        super().execute_managed_notebook(nb_man, kernel_name, **kwargs)
        
        # 后处理：为每个代码单元格添加执行时间输出
        for cell in nb_man.nb.cells:
            if cell.cell_type == "code" and hasattr(cell.metadata, 'papermill'):
                start = datetime.fromisoformat(cell.metadata.papermill.start_time)
                end = datetime.fromisoformat(cell.metadata.papermill.end_time)
                duration = (end - start).total_seconds()
                
                time_output = new_output(
                    "display_data",
                    data={"text/plain": [f"执行耗时: {duration:.3f}秒"]}
                )
                cell.outputs.insert(0, time_output)

打包配置

# setup.py
from setuptools import setup, find_packages

setup(
    name="papermill_timing",
    version="0.1",
    packages=find_packages("./src"),
    package_dir={"": "src"},
    install_requires=["papermill"],
    entry_points={
        "papermill.engine": ["timing_engine=papermill_timing:TimingEngine"]
    }
)

用户可通过--engine timing_engine参数使用此引擎，每个代码单元格的输出将包含执行时间信息。

最佳实践建议

1. 错误处理：自定义组件应提供清晰的错误信息
2. 性能考虑：远程I/O操作应考虑缓存和批量处理
3. 兼容性：保持与Papermill核心版本的兼容
4. 文档完善：为自定义组件编写详细的使用说明
5. 测试覆盖：确保自定义功能的稳定性和可靠性

通过合理利用Entry Points机制，开发者可以极大地扩展Papermill的能力边界，使其适应各种复杂场景的需求。

papermill 📚 Parameterize, execute, and analyze notebooks 项目地址: https://gitcode.com/gh_mirrors/pa/papermill