跳转至

使用 papermill 运行 Jupyter Notebook

当执行两个具有前后依赖关系的 Jupyter Notebook 时,我们需要等待第一个 Notebook 运行完成,再开始运行第二个 Notebook。

一个笨拙的方法是,先预估好第一个 Notebook 运行需要的时间(例如 1 小时),然后在第二个 Notebook 的第一个 Cell 添加:

Python
import time

time.sleep(1 * 60 * 60)

这将使第二个 Notebook 运行 1 小时后,再运行后续的 Cell。

上述方法需要事先预估第一个 Notebook 运行需要的时间,预估时间过短会导致第二个 Notebook 提前运行,预估时间过长又会导致第二个 Notebook 浪费很多时间在 time.sleep() 上。

本文介绍了如何使用 papermill 运行 Jupyter Notebook,实现控制先后运行两个 Notebook 的功能。

screen-capture

安装 papermill

首先,我们需要安装 papermill 库。可以使用以下命令通过 pip 安装:

Bash
pip install papermill

示例场景

假设我们有两个 Jupyter Notebook,第一个 Notebook 构造一个包含日期索引的 pandas DataFrame,并导出为 CSV 文件。第二个 Notebook 读取该 CSV 文件,并绘制相应的折线图,然后导出为 PNG 文件。为了实现这一点,我们可以在第一个 Notebook 中使用 Papermill 来调用第二个 Notebook。

Notebook 1: papermill-1.ipynb

papermill-1.ipynb
import time

import numpy as np
import pandas as pd
import papermill as pm

# 创建 DataFrame
dates = pd.date_range(start="2020-01-01", end="2020-12-31")
np.random.seed(42)
data = np.cumsum(np.random.randn(len(dates)) + 0.1)
df = pd.DataFrame(data, index=dates, columns=["Value"])

# 将 DataFrame 导出为 CSV 文件
df.to_csv("papermill-data.csv")
time.sleep(2)

# 使用 Papermill 运行第二个 Notebook
pm.execute_notebook(
    input_path="papermill-2.ipynb",
    output_path="papermill-2.ipynb",
)

print("第二个 Notebook 已经被运行")

在这个 Notebook 中,我们生成了一个包含从 2020-01-012020-12-31 的日期索引的 DataFrame,数据为逐渐递增的随机数。然后将这个 DataFrame 导出为 data.csv 文件,最后通过 Papermill 调用并运行第二个 Notebook。

Notebook 2: papermill-2.ipynb

papermill-2.ipynb
import matplotlib.pyplot as plt
import pandas as pd

# 读取 CSV 文件
df = pd.read_csv("papermill-data.csv", index_col=0, parse_dates=True)

# 绘制折线图
plt.figure(figsize=(10, 6))
plt.plot(df.index, df["Value"], label="Value", color="b")
plt.title("Value Over Time")
plt.xlabel("Date")
plt.ylabel("Value")
plt.legend()

# 保存折线图为 PNG 文件
png_file_path = "papermill-data.png"
plt.savefig(png_file_path)
plt.show()

print(f"折线图已导出到 {png_file_path}")

第二个 Notebook 读取第一个 Notebook 生成的 CSV 文件,并绘制一个折线图,展示数据随时间的变化趋势。最后将折线图保存为 papermill-data.png 文件。

运行 papermill-1.ipynb,就可以实现按顺序运行两个 Notebook。

screen-capture

papermill 的更多功能

papermill 的强大之处不仅在于能够运行 Jupyter Notebook,它还支持参数化运行。通过传递参数,papermill 可以在不同的运行中执行相同的 Notebook,但使用不同的输入数据。这对于批处理任务和需要重复分析的场景非常有用。

更多使用方法可以参考 papermill 官方文档

评论