作者 | Ohad Ravid
译者 | 平川
策划 | 刘燕
本文最初发布于 Ohad Ravid 的个人博客 Tea and Bits。
不久前，在 $work，我们的一个核心 Python 库遇到了性能问题。
这个库是我们 3D 处理管道的基础库。它相当大，也很复杂。它使用 NumPy 和其他科学相关的 Python 包进行广泛的数学和几何操作。
我们的系统还必须在 CPU 资源有限的环境中运行，虽然它一开始表现得不错，但随着实际并发用户数的增长，我们开始遇到问题，系统难以满足负载增长的需求。
我们的结论是，要处理增加的工作负载，系统至少要快 50 倍。我们认为，Rust 可以帮助我们实现这一目标。
因为我们遇到的性能问题很常见，所以我们要在这篇（并不算短的）文章中重现并解决它们。
泡壶茶或冲杯咖啡，我将介绍：
（a）基本问题； （b）用来解决这个问题的几次迭代优化。
你也可以直接跳到文章末尾，查看最终代码。
一个实用的例子 让我们创建一个小型库，用它重现我们最初遇到的性能问题（但其所完成的工作是任意的）。 假设你有一个多边形列表和一个点列表，都是二维的。出于业务原因，我们希望将每个点“匹配”到单个多边形。
我们假想的库将：
1、从点和多边形的初始列表开始（均为二维的）。
2、对于每个点，计算它到多边形中心的距离，找到一个离它最近的、小得多的多边形子集。
3、从这些多边形中，选择一个“最佳”的（我们将以“面积最小的”为“最佳”）。代码类似下面这样（这里 有完整代码）：
from typing import List, Tuple
import numpy as np
from dataclasses import dataclass
from functools import cached_property
Point = np.array
@dataclass
class Polygon:
x: np.array
y: np.array
@cached_property
def center(self) -> Point: ...
def area(self) -> float: ...
def find_close_polygons(polygon_subset: List[Polygon], point: Point, max_dist: float) -> List[Polygon]:
...
def select_best_polygon(polygon_sets: List[Tuple[Point, List[Polygon]]]) -> List[Tuple[Point, Polygon]]:
...
def main(polygons: List[Polygon], points: np.ndarray) -> List[Tuple[Point, Polygon]]:
主要的困难（性能方面）在于混合使用 Python 对象和 numpy 数组。
我们马上会深入分析这个问题。
值得注意的是，对于这个玩具库来说，把部分 / 所有内容转换为 向量化 的 numpy 或许是可行的，但对于真正的库来说，这几乎是不可能的，因为这会使代码的可读性和可修改性大大降低，并且收益非常有限（这里有一个部分向量化的版本，它的速度更快，但与我们想要实现的结果差很远）。
此外，任何基于 JIT 的技巧（PyPy/numba）所能带来的收益都非常小（保准起见，我们会进行测量）。
为什么不完全用 Rust™重写呢？
虽然完全重写很有吸引力，但存在几个问题：
该介绍我们的性能分析器朋友了。
Python 有一个内置的性能分析器（cProfile），但在这种情况下，它并不是真正适合这项工作的工具：
py-spy（GitHub）。py-spy是一个 采样分析器 #Statistical_profilers)，可以看到原生栈帧。
该项目还慷慨地将预先构建好的版本发布到了 PyPi，我们只需运行pip install py-spy 就可以开始工作了。
我们还需要一些东西来度量。
# measure.py
import time
import poly_match
import os
# 减少噪音，实际地改进性能
os.environ["OPENBLAS_NUM_THREADS"] = "1"
polygons, points = poly_match.generate_example()
# 随着代码的速度越来越快，我们会增加这个值
NUM_ITER = 10
t0 = time.perf_counter()
for _ in range(NUM_ITER):
poly_match.main(polygons, points)
t1 = time.perf_counter()
took = (t1 - t0) / NUM_ITER
print(f"Took and avg of {took * 1000:.2f}ms per iteration")
虽然不是很科学，但也可以让我们走很远。
要做好基准测试很难。话虽如此，不要因为追求完美的基准测试设置而倍感压力，特别是当你开始优化一个程序时。
——Nicholas Nethercote，《Rust 性能手册》运行这个脚本可以得出基线： $ python measure.py 平均每次迭代耗时 293.41 毫秒
对于原始库，为了涵盖了所有情况，我们用了 50 个不同的示例。
这与系统的整体性能是匹配的，这意味着我们可以着手降低这个数值了。
注：我们还可以使用 PyPy 进行度量（为了充分发挥 JIT 的能力，我们还在前面加了一个预热操作）。
$ conda create -n pypyenv -c conda-forge pypy numpy && conda activate pypyenv
$ pypy measure_with_warmup.py
平均每次迭代耗时 1495.81 毫秒
先测量一下
那么，首先让我们看看是哪里慢。
$ py-spy record --native -o profile.svg -- python measure.py
py-spy> 每秒采样 100 次。按 Control-C 退出。
平均每次迭代耗时 365.43 毫秒
py-spy> 进程退出，采样结束。
py-spy> 将火焰图数据写入'profile.svg'。样本数：391 错误数：0
可以看到，开销已经非常小。为了进行比较，我们使用cProfile得出了以下结果：
$ python -m cProfile measure.py
平均每次迭代耗时 546.47 毫秒
7.806 秒 7551778 次函数调用（7409483 次原始调用）
..我们得到了下面这个漂亮的淡红色 火焰图：
每个框都对应一个函数，我们可以看每个函数花费的相对时间，包括它正在调用的函数（沿着图 / 堆栈向下。注意： 原文中每个框都可以点击进入）。
这里，我们主要得出了以下结论：
find_close_polygons上。norm ，这是一个 numpy 函数。我们看一下 find_close_polygons：
def find_close_polygons(
polygon_subset: List[Polygon], point: np.array, max_dist: float
) -> List[Polygon]:
close_polygons = []
for poly in polygon_subset:
if np.linalg.norm(poly.center - point) < max_dist:
close_polygons.append(poly)
return close_polygons
我们将在 Rust 中重写这个函数。
在深入研究细节之前，有几点需要注意：
Polyon、np.array）。我的第一个 Rust 模块
pyo3是一个用于 Python 和 Rust 交互的 crate。它的文档非常完善。要了解基本设置，请点击 这里。
我们将把这个 crate 命名为poly_match_rs，并添加名为find_close_polygons的函数。
mkdir poly_match_rs && cd "$_"
pip install maturin
maturin init --bindings pyo3
maturin develop
一开始，这个 crate 看起来是下面这样：
use pyo3::prelude::*;
#[pyfunction]
fn find_close_polygons() -> PyResult {
Ok(())
}
#[pymodule]
fn poly_match_rs(_py: Python, m: &PyModule) -> PyResult {
m.add_function(wrap_pyfunction!(find_close_polygons, m)?)?;
Ok(())
}
我们还需要记住，每次修改 Rust 库时都要执行maturin develop。
就是这样！我们调用下新函数，看看会发生什么。
>>> poly_match_rs.find_close_polygons(polygons, point, max_dist)
E TypeError: poly_match_rs.poly_match_rs.find_close_polygons() takes no arguments (3 given)
v1 —— 简单的 Rust 译文
我们将从匹配预期的 API 开始。
PyO3 在从 Python 到 Rust 的转换方面非常聪明，所以这非常简单：
#[pyfunction]
fn find_close_polygons(polygons: Vec, point: PyObject, max_dist: f64) -> PyResult {
Ok(vec![])
}
顾名思义，PyObject是一个通用的 Python 对象。我们将尝试和它交互。
这样程序应该可以运行了（尽管不正确）。
我将复制粘贴原始的 Python 函数，然后仅仅修正下语法。
#[pyfunction]
fn find_close_polygons(polygons: Vec, point: PyObject, max_dist: f64) -> PyResult {
let mut close_polygons = vec![];
for poly in polygons {
if norm(poly.center - point) < max_dist {
close_polygons.push(poly)
}
}
Ok(close_polygons)
}
很酷，但编译不通过：
% maturin develop
...
error[E0609]: no field `center` on type `Py`
--> src/lib.rs:8:22
|
8 | if norm(poly.center - point) src/lib.rs:8:12
|
8 | if norm(poly.center - point) < max_dist {
| ^^^^ not found in this scope
error: aborting due to 2 previous errors ] 58/59: poly_match_rs
我们需要三个 crate 来实现我们的函数：
# 针对 Rust 原生数组操作
ndarray = "0.15"
# 面向数组的`norm`函数
ndarray-linalg = "0.16"
# 访问 numpy 基于`ndarray`创建的对象
numpy = "0.18"
首先，将不透明的通用对象point: PyObject 转换为我们可以使用的对象。
就像我们向 PyO3 请求“PyObjects的Vec”一样，我们可以请求 numpy-array，它会自动为我们转换参数。
use numpy::PyReadonlyArray1;
#[pyfunction]
fn find_close_polygons(
// 一个“拥有 GIL 锁”的对象，这样我们就可以访问 Python 管理的内存
py: Python,
polygons: Vec,
point: PyReadonlyArray1,
max_dist: f64,
) -> PyResult {
let mut close_polygons = vec![];
let point = point.as_array();
for poly in polygons {
let center = poly
.getattr(py, "center")?
.extract::(py)?
.as_array()
.to_owned();
if (center - point).norm() < max_dist {
close_polygons.push(poly)
}
}
Ok(close_polygons)
}
与原来的版本进行比较：
def find_close_polygons(
polygon_subset: List[Polygon], point: np.array, max_dist: float
) -> List[Polygon]:
close_polygons = []
for poly in polygon_subset:
if np.linalg.norm(poly.center - point) < max_dist:
close_polygons.append(poly)
return close_polygons
我们希望，与原来的函数相比，这个版本有一些优势，但是有多大的优势呢？
$ (cd ./poly_match_rs/ && maturin develop)
$ python measure.py
平均每次迭代耗时 609.46 毫秒
所以……Rust 真的很慢吗？不！我们只是忘了要求速度！如果我们使用maturin develop --release来运行，得到的结果就会好很多：
$ (cd ./poly_match_rs/ && maturin develop --release)
$ python measure.py
平均每次迭代耗时 23.44 毫秒
我们还想查看原生代码，因此，我们将在 release 配置中启用调试符号。既然这样，我们不妨要求最高速度。
# 添加到 Cargo.toml
[profile.release]
debug = true # 启用性能分析器的调试符号
lto = true # 链接时优化
codegen-units = 1 # 编译慢，但运行快
v2 —— 使用 Rust 进一步重写
现在，使用py-spy中的--native 标志，我们可以同时查看 Python 和新的原生代码。
再次运行py-spy ：
$ py-spy record --native -o profile.svg -- python measure.py
py-spy> 采样过程每秒 100 次。按 Ctrl+C 键退出。
我们得到了下面这个火焰图（非红色的部分也加进来了，以供参考）：
从分析器的输出中，我们可以看到一些有趣的东西：
find_close_polygons::...::trampoline （Python 直接调用的符号）和__pyfunction_find_close_polygons（我们的实际实现）的相对大小。 耗时占比 95% 对 88%，所以开销很小。If (center - point).norm() < max_dist{…｝）即lib_v1.rs: 22（右侧的小框）占总运行时间的大约 9%。 所以 10 倍的改进还是可能的！lib_v1.rs: 16，即poly.getattr(…).extract(…)，放大一下就可以看到getattr，它在获取使用了as_array的底层数组。 我们从中得出的结论是，我们需要专注于解决第 3 点，方法就是用 Rust 重写 Polygon。看下我们的目标：
@dataclass
class Polygon:
x: np.array
y: np.array
_area: float = None
@cached_property
def center(self) -> np.array:
centroid = np.array([self.x, self.y]).mean(axis=1)
return centroid
def area(self) -> float:
if self._area is None:
self._area = 0.5 * np.abs(
np.dot(self.x, np.roll(self.y, 1)) - np.dot(self.y, np.roll(self.x, 1))
)
return self._area
我们将尽可能地保留现有的 API，但我们真的不需要area 那么快（目前）。
实际的类可能会包含其他复杂的东西，比如merge方法，它使用了scipy.spatial中的ConvexHull。
为了降低成本（并限制这篇本已很长的文章的篇幅），我们只将Polygon的“核心”功能迁移到 Rust，并在 Python 中通过子类化来实现 API 的其余部分。
struct 看起来是这样的：
// `Array1`是一维数组， `numpy` crate 可以很方便地操作
use ndarray::Array1;
// `subclass` 告诉 PyO3 可以在 Python 中将其子类化
#[pyclass(subclass)]
struct Polygon {
x: Array1,
y: Array1,
center: Array1,
}
现在，我们需要实际地实现它。我们想把poly.{x, y, center} 暴露为：
我们还需要一个构造函数，以便 Python 可以创建新的 Polygon 。
use numpy::{PyArray1, PyReadonlyArray1, ToPyArray};
#[pymethods]
impl Polygon {
#[new]
fn new(x: PyReadonlyArray1, y: PyReadonlyArray1) -> Polygon {
let x = x.as_array();
let y = y.as_array();
let center = Array1::from_vec(vec![x.mean().unwrap(), y.mean().unwrap()]);
Polygon {
x: x.to_owned(),
y: y.to_owned(),
center,
}
}
// 返回类型中的`Py`是一种说明“Python 拥有一个对象”的方式。
#[getter]
fn x(&self, py: Python,
polygons: Vec, // 引用 Python 拥有的对象
point: PyReadonlyArray1,
max_dist: f64,
) -> PyResult { // 返回相同的`Py`引用，未经修改
let mut close_polygons = vec![];
let point = point.as_array();
for poly in polygons {
let center = poly.borrow(py).center // 需要使用 GIL (`py`) 借用底层的`Polygon`
.to_owned();
if (center - point).norm() < max_dist {
close_polygons.push(poly)
}
}
Ok(close_polygons)
}
让我们看看这段代码会带来什么：
$ python measure.py
平均每次迭代耗时 6.29 毫秒
就快达到目标了！再提升 2 倍就行了！
v3 —— 避免分配
让我们再一次启动剖析器。
select_best_polygon，它现在调用了一些 Rust 代码（当它获取x & y向量） 我们可以修复这个问题，但潜在的改善可能非常小（可能只有 10%）extract_argument（在lib_v2.rs: 48下）上，所以我们还是付出了不小的开销！ 但大部分时间都在PyIterator::next和PyTypeInfo::is_type_of中，要改进并不容易。lib_v2.rs: 58是我们的if，我们还看到了drop_in_place和to_owned。 这行代码大约了占总时间的 35%，远超我们的预期：所有数据都到位的话，这应该是”快速位（fast bit）“。以下是有问题的代码片段：
let center = poly.borrow(py).center
.to_owned();
if (center - point).norm() < max_dist { ... }
我们希望避免to_owned 。但norm 需要一个 owned 对象，所以我们必须手动实现。
（这里，我们之所以可以改进ndarray是因为我们知道数组实际上只包含 2 个f32）。
看起来就像下面这样：
use ndarray_linalg::Scalar;
let center = &poly.as_ref(py).borrow().center;
if ((center[0] - point[0]).square() + (center[1] - point[1]).square()).sqrt() < max_dist {
close_polygons.push(poly)
}
但是，借用检查器不允许我们这样做：
error[E0505]: cannot move out of `poly` because it is borrowed
--> src/lib.rs:58:33
|
55 | let center = &poly.as_ref(py).borrow().center;
| ------------------------
| |
| borrow of `poly` occurs here
| a temporary with access to the borrow is created here ...
...
58 | close_polygons.push(poly);
| ^^^^ move out of `poly` occurs here
59 | }
60 | }
| - ... and the borrow might be used here, when that temporary is dropped and runs the `Drop` code for type `PyRef`
像往常一样，借用检查器是正确的：我们正在进行内存犯罪。
更简单的修复方法是直接克隆，并编译close_polygons.push(poly.clone())。
实际上，这个克隆的成本非常低，因为我们只incrPython 对象的引用计数。
然而，在这种情况下，我们也可以通过做一个典型的 Rust 技巧来缩短借用：
let norm = {
let center = &poly.as_ref(py).borrow().center;
((center[0] - point[0]).square() + (center[1] - point[1]).square()).sqrt()
};
if norm < max_dist {
close_polygons.push(poly)
}
因为poly只在内部作用域内被借用，一旦到达close_polygons.push ，编译器就可以知道我们不再持有该引用，并将愉快地完成新版本的编译。
最后，我们得到如下结果：
$ python measure.py
平均每次迭代耗时 2.90 毫秒
比原来的代码快了 100 倍。
小结
我们是从下面的 Python 代码开始的：
@dataclass
class Polygon:
x: np.array
y: np.array
_area: float = None
@cached_property
def center(self) -> np.array:
centroid = np.array([self.x, self.y]).mean(axis=1)
return centroid
def area(self) -> float:
...
def find_close_polygons(
polygon_subset: List[Polygon], point: np.array, max_dist: float
) -> List[Polygon]:
close_polygons = []
for poly in polygon_subset:
if np.linalg.norm(poly.center - point) < max_dist:
close_polygons.append(poly)
return close_polygons
# Rest of file (main, select_best_polygon).
我们使用py-spy对它进行了性能分析，即使只是对 find_close_polyons 做 最简单的行对行翻译，性能提升也超过了 10 倍。
我们又做了多轮 profile-write-measure 迭代，直到运行速度获得了 100 倍的改进，而且 API 与原来的库相同。
最终的 Python 代码如下所示：
import poly_match_rs
from poly_match_rs import find_close_polygons
class Polygon(poly_match_rs.Polygon):
_area: float = None
def area(self) -> float:
...
# 文件的其余部分没有变化 (main, select_best_polygon)。
它调用以下 Rust 代码：
use pyo3::prelude::*;
use ndarray::Array1;
use ndarray_linalg::Scalar;
use numpy::{PyArray1, PyReadonlyArray1, ToPyArray};
#[pyclass(subclass)]
struct Polygon {
x: Array1,
y: Array1,
center: Array1,
}
#[pymethods]
impl Polygon {
#[new]
fn new(x: PyReadonlyArray1, y: PyReadonlyArray1) -> Polygon {
let x = x.as_array();
let y = y.as_array();
let center = Array1::from_vec(vec![x.mean().unwrap(), y.mean().unwrap()]);
Polygon {
x: x.to_owned(),
y: y.to_owned(),
center,
}
}
#[getter]
fn x(&self, py: Python,
polygons: Vec,
point: PyReadonlyArray1,
max_dist: f64,
) -> PyResult {
let mut close_polygons = vec![];
let point = point.as_array();
for poly in polygons {
let norm = {
let center = &poly.as_ref(py).borrow().center;
((center[0] - point[0]).square() + (center[1] - point[1]).square()).sqrt()
};
if norm PyResult {
m.add_class::()?;
m.add_function(wrap_pyfunction!(find_close_polygons, m)?)?;
Ok(())
}
本文要点
- Rust（在 pyo3 的帮助下）以最小的妥协为普通的 Python 代码解锁了真正的原生性能体验。
- 对于研究人员来说，Python 是一个极好的 API，和 Rust 快速构建块一起组成了一个极其强大的组合。
- 性能分析非常有趣，它能让你真正地了解代码中发生的一切。
最后：电脑的速度快得惊人。下次，当你等待某件事完成时，可以考虑启动一个分析器，你可能会了解到一些新东西。
原文链接：
https://ohadravid.github.io/posts/2023-03-rusty-python/
