自动保存、查询并版本化 Python 计算

mandala 通过两种通用工具，消除了机器学习实验跟踪（以及更多）所需的努力和代码开销：

1. @op 装饰器：
   • 捕获 Python 函数调用的输入、输出及代码（包括依赖关系）
   • 自动重用过往结果，绝不会对同一调用重复计算
   • 专为组合成端到端持久化程序而设计，从而在纯 Python 中实现高效的迭代开发，无需考虑存储后端。

ComputationFrame 数据结构： 自动将命令式代码的执行组织为变量和操作的高层计算图。能够检测反馈循环、分支/合并以及聚合/索引等模式。 查询变量之间的关系，通过提取一个数据框，其中列代表计算图中的变量和操作，每行包含计算图部分执行的值或调用。 自动化对异构的 `@op` 调用网络进行探索和高层操作。

视频演示

快速演示如何在 mandala 中运行计算，同时更新对应的 ComputationFrame 视图及其提取的数据框（代码可在这里找到）：

https://github.com/amakelov/mandala/assets/1467702/85185599-10fb-479e-bf02-442873732906

安装

pip install pymandala

或者

pip install git+https://github.com/amakelov/mandala

教程

• 快速入门： | 阅读文档
• ComputationFrames： | 阅读文档
• 玩具机器学习项目： | 阅读文档

博客与论文

• 整洁计算：介绍 ComputationFrame 数据结构及其应用
• Python 函数调用的实用依赖跟踪：描述 mandala 依赖跟踪系统的设计动机与实现细节
• 该论文即将发表于 SciPy 2024 会议论文集。
• 在Hacker News上的讨论

常见问题

这与其他实验跟踪框架有何不同？

与 W&B、MLFlow 或 Comet 等流行工具相比，mandala：

• 更细粒度地与实际 Python 代码执行集成
  • 函数调用是持久化、版本管理和查询的同步单元，而非整个脚本或笔记本，从而实现更高效的复用和增量开发。
  • 更进一步，Python 集合类型（如 list、dict）可以对存储系统透明化，使得单个元素被单独存储和跟踪，并可在不同集合和调用之间复用。
  • 由于基于记忆化而非日志记录，您无需考虑如何命名所记录的内容。
• 提供 ComputationFrame 数据结构，这是一种强大且简单的方法，用于表示、查询和操作复杂的已保存计算。
• 自动从当前代码库状态和调用输入中解析每个 @op 调用的版本号。

@op 缓存是如何失效的？

• 对于某个 @op 函数的调用，例如 f，系统会搜索过去使用相同内容输入（由哈希函数决定）的调用，且这些调用所访问的依赖项（包括 f 自身）的版本与当前状态兼容。
• 函数版本之间的兼容性由用户决定：您可以自由标记某些更改与过往结果兼容，但关于标记兼容性的限制，请参阅限制部分。
• 内部而言，mandala 使用略微修改后的 joblib 哈希算法来计算 Python 对象的内容哈希值。这在许多场景下非常实用，但并不完美，具体讨论请参见限制部分。

我可以修改 @op 的代码吗？如果修改了会发生什么？

• 常见用例：您一直在使用某个 @op，随后希望扩展其功能而不使过往结果失效。推荐的做法是添加一个新参数 a，并为其提供一个包裹在 NewArgDefault(x) 中的默认值。当传入等于 x 的值时，存储系统会回退到之前的调用。
• 除此之外，您可能需要利用版本控制系统来检测 @op 的依赖关系及其变化。详情请参阅文档。

是否已准备好投入生产环境？

• mandala 目前处于 Alpha 阶段，API 可能会发生变化。
• 此外，已知存在性能瓶颈，可能导致处理超过 1 万次调用的存储时速度较慢。

它有多自包含？

• mandala 的核心只有几千行代码，并且仅依赖于 pandas 和 joblib。
• 为了可视化 ComputationFrame，您需要在系统级别安装 dot 工具，和/或安装 Python 的 graphviz 库。

局限性

• 目前的版本控制系统功能不够丰富，文档也不够完善，无法满足实际使用场景的需求。例如，它不支持以一致的方式删除旧版本，也无法根据函数版本限制 ComputationFrame。此外，许多错误信息不够清晰，也未提供解决方案建议。
• 在使用版本控制时，如果您将某次更改标记为与历史结果兼容，需格外小心：如果此次更改引入了未被 mandala 跟踪的新依赖项，这些“隐形”依赖的变化可能不会被存储系统察觉，从而导致您误以为某些结果是最新的，而实际上并非如此。
• 请参阅“注意事项”笔记本，了解需要避免的常见错误：

未来功能路线图

总体

• 支持 @op 中的命名输出
• 支持重命名 @op 及其输入/输出

记忆化

• 为选定对象添加自定义序列化
• 找到一种忽略内容哈希中微小数值误差的解决方案
• 改进集合相关的文档
• 通过例如 dask 或 ray 支持并行执行 @op
• 支持从存储中排除特定的输入/输出

计算框架

• 添加对计算图中循环的支持
• 改进 expand_... 方法的启发式算法
• 增加通过谓词将 CF 限制到变量值特定子集的工具
• 改进集合的使用支持及示例
• 添加对 CF 中节点合并、拆分等简化操作的支持

版本控制

• 支持以一致方式删除旧版本
• 改进文档和错误信息
• 对该系统进行更全面的测试
• 支持按函数版本限制 CF
• 支持手动为版本添加依赖项，以避免“隐形依赖”问题

性能

• 提升内存缓存的性能
• 提升 ComputationFrame 操作的性能

高远愿景

从理想角度来看，mandala 的意义远不止于机器学习实验跟踪。其主要目标是 让持久化逻辑与最佳实践成为 Python 的自然延伸。一旦实现这一点，您原本就必须编写的纯“计算”代码，同时也将成为存储接口。很难再找到比这更简单、更可靠的方式来管理计算产物了。

一种基于第一性原理的计算产物管理方法

我们期望存储系统能够提供以下功能：

• 用简短且无歧义的描述引用产物：“这是我计算出的[庞大而复杂的 Python 对象]，对我而言，它代表[人类可读的描述]”
• 保存产物：“保存[庞大而复杂的 Python 对象]”
• 在后续时间点引用并加载产物：“请给我之前计算出的[人类可读的描述]”
• 确认是否已计算过某产物：“我是否已经计算过[人类可读的描述]？”
• 以更复杂的方式查询结果：“请给我所有满足[更高层次的人类可读描述]的事物”，这在实践中意味着对多个产物组合施加某种谓词条件。
• 获取产物生成过程的报告：“用于生成[人类可读的描述]的代码是什么？”

关键的观察是，执行轨迹几乎可以回答上述所有问题。

相关工作

mandala 结合了多种技术的思想，并与之存在相似之处。 以下是一些有用的对比点：

• 记忆化：
  • provenance 库在记忆化方面与 mandala 非常相似，但缺少查询和依赖关系跟踪的功能。
  • 标准的 Python 记忆化解决方案包括 joblib.Memory 和 functools.lru_cache。mandala 在底层使用了 joblib 的序列化和哈希机制。
  • incpy 是一个将记忆化直接集成到 Python 解释器中的项目。
  • funsies 是一个基于记忆化的分布式工作流执行器，它使用与 mandala 类似的哈希机制来跟踪哪些计算已经完成。该工具作用于脚本级别（而非函数级别），并且缺乏可查询性和版本管理功能。
  • koji 是一种增量计算数据处理框架的设计，能够统一处理不同类型的资源（文件或服务）。它同样采用类似的哈希机制来跟踪计算过程。
• 计算框架：
  • 计算框架是 关系数据库 的特例：计算图中的每个函数节点都有一张调用表，其列对应所有连接到该函数的输入输出边标签。类似地，每个变量节点则是一张单列表，记录该变量中的所有 Ref。外键约束将函数的列与变量关联起来，而通过表之间的各种连接操作可以表达变量联合计算历史的不同概念。
  • 计算框架也与 图数据库 存在联系，因为对计算框架的一些相关查询，例如涉及沿 @op 边可达性的查询，实际上是图数据库查询的特例。此外，Storage 的内部表示形式更接近图数据库，而非关系数据库。
  • 计算框架还与应用 范畴论 中的一些思想相关，例如将有限范畴到集合范畴的函子（即“余层”）用作“通用”内存数据结构的蓝图，这种结构再次等价于关系数据库；参见例如 这篇论文，其中描述了这一范畴论构造。
• 版本管理：
  • 代码库中每个函数的修订历史被组织在一个“迷你 git 仓库”中，该仓库仅保留了 git 的最基本功能：它是一个 内容寻址存储 树，每条边记录从一端内容到另一端内容的差异。附加元数据用于标识语义等价的内容类。
  • 语义化版本控制 是另一种流行的代码版本管理系统。mandala 与 semver 的相似之处在于，它允许对依赖项的接口和逻辑进行向后兼容的更改。二者的不同之处在于，mandala 仍然以内容本身来标记版本，而不是使用“非规范”的数字。
  • Unison 编程语言 则通过函数内容的哈希值（确切地说是语法树的哈希值）来表示函数。

Mandala 快速上手指南

Mandala 是一个用于自动保存、查询和版本化 Python 计算的开源工具。它通过 @op 装饰器实现细粒度的函数级缓存与复用，并利用 ComputationFrame 数据结构将命令式代码执行自动组织为可查询的计算图，特别适合机器学习实验追踪及迭代开发。

环境准备

• 系统要求：支持 Linux、macOS 和 Windows。
• Python 版本：建议 Python 3.8 及以上。
• 核心依赖：安装后会自动包含 pandas 和 joblib。
• 可选依赖（可视化）：若需可视化 ComputationFrame 计算图，请确保系统已安装 Graphviz (dot 命令) 或 Python graphviz 库。
  • Ubuntu/Debian: sudo apt-get install graphviz
  • macOS: brew install graphviz
  • Windows: 下载并安装 Graphviz 安装包并将 bin 目录加入环境变量。

安装步骤

推荐使用 pip 进行安装。国内用户可使用清华或阿里镜像源加速下载。

方式一：从 PyPI 安装（推荐）

pip install pymandala -i https://pypi.tuna.tsinghua.edu.cn/simple

方式二：从 GitHub 源码安装（获取最新特性）

pip install git+https://github.com/amakelov/mandala -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

Mandala 的核心在于使用 @op 装饰器标记函数。一旦标记，该函数的输入、输出及代码依赖将被自动捕获。相同的输入和代码状态不会重复计算，而是直接返回缓存结果。

1. 最简单的示例

以下示例展示了如何定义一个带缓存的计算函数，并演示第二次调用时如何直接复用结果。

from mandala.core import OpStore

# 初始化存储后端 (默认使用本地临时目录)
store = OpStore()

# 使用 @op 装饰器标记函数
@store.op
def expensive_computation(x, y):
    print(f"正在计算 {x} + {y}...")  # 仅在首次执行或缓存失效时打印
    return x + y

# 第一次调用：实际执行计算
result1 = expensive_computation(2, 3)
print(f"结果 1: {result1}") 
# 输出:
# 正在计算 2 + 3...
# 结果 1: 5

# 第二次调用：相同参数，直接复用缓存，不执行函数体
result2 = expensive_computation(2, 3)
print(f"结果 2: {result2}")
# 输出:
# 结果 2: 5 (注意：没有打印 "正在计算..."，说明未重复执行)

# 参数变化：重新执行计算
result3 = expensive_computation(2, 4)
print(f"结果 3: {result3}")
# 输出:
# 正在计算 2 + 4...
# 结果 3: 6

2. 查询计算历史 (ComputationFrame)

Mandala 允许你将一系列 @op 调用组织成数据框（DataFrame）进行查询和分析。

# 假设已经运行了多次 expensive_computation
# 提取计算帧 (ComputationFrame)
cf = store.get_computation_frame()

# 查看包含所有变量和操作的数据框
df = cf.to_dataframe()
print(df.head())

# 你可以像操作 pandas DataFrame 一样筛选特定的计算结果
# 例如：找出所有第一个参数为 2 的计算记录
filtered = df[df['arg_x'] == 2]
print(filtered)

3. 开始你的第一个教程

想要更深入地了解如何在 ML 项目中使用或查看交互式演示，可以访问官方 Colab 教程：

• 快速入门: Open in Colab
• ML 项目实战: Open in Colab