使用 Rust 极致提升 Python 性能:图表和绘图提升 24 倍,数据计算提升 10 倍

关于作者

Edward Wright

Vortexa 公司的首席 GIS 工程师。不写代码的时候,他忙着跑步机、山地自行车、建筑、修理东西,以及油画。

有时候,仅采用标准方法还不够好。本篇文章,是关于在重要的地方做最小的改变,从而达到最大的效果。

问题的边界

在 vortex 公司,我们广泛使用 Python。Python 非常适合于原型设计,也非常适合于数据的科学计算。虽然 Python 不是最快的语言,但它通常是非常棒的。

然而,最近我们发现一个特定的 Python 任务,需要 30 小时才能运行完毕。由于一些模型的变更,当我们想对一些业务调用重新计算时,这个运行时间真的影响了我们的 QA 反馈周期,使得将更新的模型引入到生产环境,变得更加困难。如果我们能够解决这个问题,将会加速模型的改进,为团队和我们的客户带来真正的好处。

在没有太多无关细节的情况下,我们的任务是处理来自船舶的 GPS 信号,并在应用其它算法之前,通过一组多边形算法,对信号进行过滤。

为什么这段代码如此慢?

无需做假设,我们的出发点必须是先测量这段代码。

我创建了代码的一个副本(复制/粘贴即可),但对其进行了修改,以便于可以处理一个小数据集。并在将来,对不同的技术进行比较。这段测试的代码,仍然忠实地再现了生产环境中所部署代码的运行负载。我使用优秀的 pyinstrument 模块,深入了解了 Python 代码中正在发生的事情。为了防止由于运行时间过短而扭曲结果,在所有初始化工作完成后,我才开始分析。

结果如下:

时间单位为秒。

main 方法,代表了算法在完成整个初始化之后的处理过程。test_python 方法,正在测试我认为很慢的部分代码的逻辑。当然,所有其它代码的逻辑仍然是存在的。

所以在 34.3 秒的运行时间中,29.8 秒花在了我前面提到的过滤逻辑中,25.1 秒消耗在 matplotlib 处理中,主要是做多边形绘图运算。

哪儿有问题?

我进行的测试数据,使用了近 8 米的船舶定位。我们正在研究全世界的数百个区域,数百个实现过滤功能的多边形算法要运行。我们使用的是 pandas,船舶的位置存储在 dataframe,但是我们需要将这个 dataframe 传递给 matplotlib,用于我们要测试的每个多边形区域。

我们对一个库进行了数百次调用,每次都要传递数百万条记录。在生产环境中,我们处理的数据可能要增加到 2500 倍,因此使用者才能看到 30 小时内,船舶的位置数据来自何处。

如何处理?

或许,在生产环境中进行繁重的任务处理,matplotlib 不是合适的工具?既然代码中已经在使用 pandas 了,为什么不试试 geopandas 呢?然后,我们可以在一个库调用中,计算所有多边形区域。

然而,这是一个灾难,我们增加了 10 倍的运行时间!Geopandas(以及它依次调用的其它库)使用了 423 个堆栈帧,而 matplotlib 只使用了 5 个堆栈帧,我觉得这非常惊人。测试跟踪还显示,即使创建 GeoDataFrames,也要比基于 matplotlib 的整体处理,花费更长的时间。

所以,我们有一个选择题。我们可以:

  1. 尝试将数据分块,然后使用多进程 multi-processing 模块处理(在 Python 中是不推荐的),从而利用更强大的云虚拟机,用来支撑 matplotlib 计算。
  2. 使用线程,编写一个非常小的本地自定义库,用来完成我们想要的数学运算。

第一种方法可以工作,但不太可能是非常经济高效的,因为我们只是并行地运行多个较慢代码的副本。于是,我决定试试第二种选择。

规划自定义本地库

考虑到在早期的 Java point-in-polygon 开发中,吸取到的一些经验教训,这次我们可以使用一些技巧。例如:

  1. 避免为每个多边形计算都进行库调用,为每个 dataframe 只进行一次调用,可以大量减少库调用的开销。
  2. 避免在实际问题非常简单的情况下,使用重量级几何计算库,否则开销会严重影响性能。
  3. 对每个多边形进行边界测试。
  4. 尽可能基于 32 位整数(比浮点更快)。
  5. 使用线程。

需要说明的是,Java 肯定不是这里的答案。Java 与 Python 的集成,真是太吓人了。

Rust

最近,我一直在使用 PyO3 做一些实验性的工作,它允许 Rust / Python 的双向集成。这里,我们将重点介绍 Python 导入和使用 Rust 实现的模块。

以下是实现的功能明细:

  1. 在 Rust 中实现 Python 类。
  2. 在构造函数中,存放 geojson 字符串数组,表示我们的多边形区域。
  3. 从船舶位置 dataframe,获取纬度/经度坐标,存入 numpy 数组。
  4. 返回结果为 numpy 数组(便于与 Python pandas 集成),表示每个坐标集对应的多边形(如果有的话)。

包含细节的整个实现,需要大约 300 行 Rust 代码,甚至包括 Rust 文档和单元测试!并且,还替换了大约 30 行 Python 代码(增加对 matplotlib 的调用)。PyO3 可以很好地与 numpy 和 ndarray crate(Rust 库)配合使用,允许其轻松地与 pandas 以及 numpy array 集成。并行处理方面,我们使用了 rayon。

有用吗?

当然有用。否则,这篇博文会很无聊的……

测试数据是完全相同的。

“使用 Rust,我们已经将 matplotlib 的处理时间,从 29.8 秒减少到 2.9 秒。”

Python 只使用一个线程,而 Rust 使用了 8 个线程(intel i7,超线程 4 核,所以称之为 4-5 倍的有效计算)。这还包括 Python 将结果集插回 pandas dataframe 的时间消耗。将实际的 matplotlib 与 Rust 库调用进行比较,可以得到 24 倍的改进。输出数据已经检查过,结果显示完全相同。

我们的新解决方案(在功能级别,即 dataframe 输入/输出),速度提高了 10 倍。集群中运行的代码,将其计算核心数量增加到 4 个,是完全合理的。考虑到后续的过滤算法,Rust 处理时间约占任务总运行时间的 20%,因此添加更多线程几乎没有意义,除非任务的其他部分可以受益。

生产环境的提升

以上小修改的具体代码,已经部署在正式生产环境中。上文提到,数据量会扩大到 2500 倍。

“这个处理过程,过去需要 30 个小时,现在需要 6 个小时,速度提升 500%。”

这次改进,不仅仅是学术上的,也不仅仅是为了降低工作成本。

“我们为客户带来模型变更后的内部流程,包括 QA,现在比以前快了一天——每次都快。”

这是经过深思熟虑的、有针对性的优化。

我们必须考虑到,我们在这里添加了一项新技术,使代码复杂化了,并使维护源代码存储库变得更加困难。但是,通过限制新库的功能实现范围,具体地小改进,可以缓解这种情况。业务逻辑没有改变,但实现方式已经改变了,只要 point-in-polygon “正常工作”——我们有单元测试来证明这一点——这次代码改进就不会造成任何伤害。

原文链接:Using Rust to corrode insane Python run-times

谢谢您的阅读!