[Rust] 使用 Rust 极致提升 Python 性能:图表和绘图提升 24 倍,数据计算提升 10 倍
💥 内容涉及著作权,均归属作者本人。若非作者注明,默认欢迎转载:请注明出处,及相关链接。
Summary: 有时候,仅采用标准方法还不够好。本篇文章,是关于在重要的地方做最小的改变,从而达到最大的效果。处理来自船舶的 GPS 信号时,使用 Python 各种类库优化后,性能仍然不能满足。通过大约 300 行 Rust 代码,包含细节的整个实现,甚至包括 Rust 文档和单元测试!使得数据处理的速度大幅提升:图表和绘图提升 24 倍,数据计算提升 10 倍。生产环境中的体现,表明这是经过深思熟虑的、有针对性的优化。这次改进,不仅仅是学术上的,也不仅仅是为了降低工作成本。
Topics: rust python pyo3 rust-python-集成 python-性能改进
关于作者
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 的整体处理,花费更长的时间。
所以,我们有一个选择题。我们可以:
- 尝试将数据分块,然后使用多进程 multi-processing 模块处理(在 Python 中是不推荐的),从而利用更强大的云虚拟机,用来支撑 matplotlib 计算。
- 使用线程,编写一个非常小的本地自定义库,用来完成我们想要的数学运算。
第一种方法可以工作,但不太可能是非常经济高效的,因为我们只是并行地运行多个较慢代码的副本。于是,我决定试试第二种选择。
规划自定义本地库
考虑到在早期的 Java point-in-polygon 开发中,吸取到的一些经验教训,这次我们可以使用一些技巧。例如:
- 避免为每个多边形计算都进行库调用,为每个 dataframe 只进行一次调用,可以大量减少库调用的开销。
- 避免在实际问题非常简单的情况下,使用重量级几何计算库,否则开销会严重影响性能。
- 对每个多边形进行边界测试。
- 尽可能基于 32 位整数(比浮点更快)。
- 使用线程。
需要说明的是,Java 肯定不是这里的答案。Java 与 Python 的集成,真是太吓人了。
Rust
最近,我一直在使用 PyO3 做一些实验性的工作,它允许 Rust / Python 的双向集成。这里,我们将重点介绍 Python 导入和使用 Rust 实现的模块。
以下是实现的功能明细:
- 在 Rust 中实现 Python 类。
- 在构造函数中,存放 geojson 字符串数组,表示我们的多边形区域。
- 从船舶位置 dataframe,获取纬度/经度坐标,存入 numpy 数组。
- 返回结果为 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 倍。当我们的目标是更快的 QA 周期时,需要最终部署在 AWS 中的 Kubernetes 集群中。集群中运行的代码,将其计算核心数量增加到 4 个,是完全合理的。考虑到后续的过滤算法,Rust 处理时间约占任务总运行时间的 20%,因此添加更多线程几乎没有意义,除非任务的其他部分可以受益。
生产环境的提升
以上小修改的具体代码,已经部署在正式生产环境中。上文提到,数据量会扩大到 2500 倍。
“这个处理过程,过去需要 30 个小时,现在需要 6 个小时,速度提升 500%。”
这次改进,不仅仅是学术上的,也不仅仅是为了降低工作成本。
“我们为客户带来模型变更后的内部流程,包括 QA,现在比以前快了一天——每次都快。”
这是经过深思熟虑的、有针对性的优化。
我们必须考虑到,我们在这里添加了一项新技术,使代码复杂化了,并使维护源代码存储库变得更加困难。但是,通过限制新库的功能实现范围,具体地小改进,可以缓解这种情况。业务逻辑没有改变,但实现方式已经改变了,只要 point-in-polygon
“正常工作”——我们有单元测试来证明这一点——这次代码改进就不会造成任何伤害。
原文链接:Using Rust to corrode insane Python run-times
谢谢您的阅读!