最近，我們的一個(gè)核心 Python 庫遇到了性能問題。這是一個(gè)非常龐大且復(fù)雜的庫，是我們 3D 處理管道的支柱，使用了 NumPy 以及其他 Python 數(shù)據(jù)科學(xué)庫來執(zhí)行各種數(shù)學(xué)和幾何運(yùn)算。

具體來說，我們的系統(tǒng)必須在 CPU 資源有限的情況下在本地運(yùn)行，雖然起初的性能還不錯(cuò)，但隨著并發(fā)用戶數(shù)量的增長，我們開始遇到問題，系統(tǒng)也出現(xiàn)了超負(fù)載。

我們得出的結(jié)論是：系統(tǒng)至少需要再快 50 倍才能處理這些增加的工作負(fù)載——而我們認(rèn)為，Rust 可以幫助我們實(shí)現(xiàn)這一目標(biāo)。

因?yàn)槲覀冇龅降男阅軉栴}很常見，所以下面，我來簡單介紹一下解決過程：

（a）基本的潛在問題；

（b）我們可以通過哪些優(yōu)化來解決這個(gè)問題。

我們的運(yùn)行示例

首先，我們通過一個(gè)小型庫來展示最初的性能問題。

假設(shè)有一個(gè)多邊形列表和一個(gè)點(diǎn)列表，且都是二維的，出于業(yè)務(wù)需求，我們需要將每個(gè)點(diǎn)“匹配”到一個(gè)多邊形。

我們的庫需要完成下列任務(wù)：

? 從點(diǎn)和多邊形的初始列表（全部為 2D）著手。

? 對(duì)于每個(gè)點(diǎn)，根據(jù)與中心的距離，找到離點(diǎn)最近的多邊形的子集。

? 從這些多邊形中，選擇一個(gè)“最佳”多邊形。

代碼大致如下：

性能方面最主要的難點(diǎn)在于，Python 對(duì)象和 numpy 數(shù)組的混合。

下面，我們簡單地分析一下這個(gè)問題。

需要注意的是，對(duì)于上面這段代碼，我們當(dāng)然可以把一切都轉(zhuǎn)化成 numpy 的向量計(jì)算，但真正的庫不可能這么做，因?yàn)檫@會(huì)導(dǎo)致代碼的可讀性和可修改性大大降低，收益也非常有限。此外，使用任何基于 JIT 的技巧（PyPy / numba）產(chǎn)生的收益都非常小。

為什么不直接使用 Rust 重寫所有代碼？

雖然重寫所有代碼很誘人，但有一些問題：

? 該庫的大量計(jì)算使用了 numpy，Rust 也不一定能提高性能。

? 該庫龐大而復(fù)雜，關(guān)系到核心業(yè)務(wù)邏輯，而且高度算法化，因此重寫所有代碼需要付出幾個(gè)月的努力，而我們可憐的本地服務(wù)器眼看就要掛了。

? 一群好心的研究人員積極努力改進(jìn)這個(gè)庫，實(shí)現(xiàn)了更好的算法，并進(jìn)行了大量實(shí)驗(yàn)。他們不太愿意學(xué)習(xí)一門新的編程語言，而且還要等待編譯，還要研究復(fù)雜的借用檢查器——他們不希望離開舒適區(qū)太遠(yuǎn)。

小心探索

下面，我來介紹一下我們的分析器。

Python 有一個(gè)內(nèi)置的 Profiler (cProfile)，但對(duì)于我們來說，選擇這個(gè)工具不太合適：

? 它會(huì)為所有 Python 代碼引入大量開銷，卻不會(huì)給原生代碼帶來額外開銷，因此測試結(jié)果可能有偏差。

? 我們將無法查看原生代碼的調(diào)用幀，這意味著我們也無法查看 Rust 代碼。

所以，我們計(jì)劃使用 py-spy，它是一個(gè)采樣分析器，可以查看原生幀。他們還將預(yù)構(gòu)建的輪子發(fā)布到了 pypi，因此我們只需運(yùn)行 pip install py-spy 即可。

此外，我們還需要一些測量指標(biāo)。

# measure.pyimport timeimport poly_matchimport os
# Reduce noise, actually improve perf in our case.os.environ['OPENBLAS_NUM_THREADS'] = '1'polygons, points = poly_match.generate_example()# We are going to increase this as the code gets faster and faster.NUM_ITER = 10t0 = time.perf_counter()for _ in range(NUM_ITER):poly_match.main(polygons, points)t1 = time.perf_counter()took = (t1 - t0) / NUM_ITERprint(f'Took and avg of {took * 1000:.2f}ms per iteration')

這些測量指標(biāo)雖然不是很科學(xué)，但可以幫助我們優(yōu)化性能。

“我們很難找到合適的測量基準(zhǔn)。但請(qǐng)不要過分強(qiáng)調(diào)擁有完美的基準(zhǔn)測試設(shè)置，尤其是當(dāng)你優(yōu)化某個(gè)程序時(shí)?！?/span>

—— Nicholas Nethercote，《The Rust Performance Book》

運(yùn)行該腳本，我們就可以獲得測量基準(zhǔn)：

對(duì)于原來的庫，我們使用了 50 個(gè)不同的樣本來確保涵蓋所有情況。

這個(gè)測量結(jié)果與實(shí)際的系統(tǒng)性能相符，這意味著，我們的工作就是突破這個(gè)數(shù)字。

我們還可以使用 PyPy 進(jìn)行測量：

$ conda create -n pypyenv -c conda-forge pypy numpy && conda activate pypyenv$ pypy measure_with_warmup.pyTook an avg of 1495.81ms per iteration

先測量

首先，我們來找出什么地方如此之慢。

我們可以看到開銷非常小。相較而言，使用 cProfile 得到的數(shù)據(jù)如下：

$ python -m cProfile measure.pyTook an avg of 546.47ms per iteration7551778 function calls (7409483 primitive calls) in 7.806 seconds…

下面是我們獲得的火焰圖：

每個(gè)方框都是一個(gè)函數(shù)，我們可以看到每個(gè)函數(shù)花費(fèi)的相對(duì)時(shí)間，包括它正在調(diào)用的函數(shù)（沿著圖形/棧向下）。

要點(diǎn)總結(jié)：

? 絕大部分時(shí)間花在 find_close_polygons 上。

? 大部分時(shí)間都花在執(zhí)行 norm，這是一個(gè) numpy 函數(shù)。

下面，我們來仔細(xì)看看 find_close_polygons：

我們打算用 Rust 重寫這個(gè)函數(shù)。

在深入細(xì)節(jié)之前，請(qǐng)務(wù)必注意以下幾點(diǎn)：

? 此函數(shù)接受并返回復(fù)雜對(duì)象（Polygon、np.array）。

? 對(duì)象的大小非常重要（因此復(fù)制需要一定的開銷）。

? 這個(gè)函數(shù)被調(diào)用了很多次（所以我們引入的開銷可能會(huì)引發(fā)問題）。

我的第一個(gè) Rust 模塊

PyO3 是一個(gè)用于 Python 和 Rust 之間交互的 crate ，擁有非常好的文檔。

我們將調(diào)用自己的 poly_match_rs，并添加一個(gè)名為 find_close_polygons 的函數(shù)。

mkdir poly_match_rs && cd '$_'pip install maturinmaturin init --bindings pyo3maturin develop

剛開始的時(shí)候，我們的 crate 大致如下：

我們還需要記住，每次修改 Rust 庫時(shí)都需要執(zhí)行 maturin develop。

改動(dòng)就這么多。下面，我們來調(diào)用新函數(shù)，看看情況會(huì)怎樣。

>>> 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)

第一版：Rust 轉(zhuǎn)換

首先，我們來定義 API。

PyO3 可以幫助我們將 Python 轉(zhuǎn)換成 Rust：

PyObject （顧名思義）是一個(gè)通用、“一切皆有可能”的 Python 對(duì)象。稍后，我們將嘗試與它進(jìn)行交互。

這樣程序應(yīng)該就可以運(yùn)行了（盡管不正確）。

我直接把原來的 Python 函數(shù)復(fù)制粘帖進(jìn)去，并修復(fù)了語法問題。

#[pyfunction]fn find_close_polygons(polygons: Vec<PyObject>, point: PyObject, max_dist: f64) -> PyResult<Vec<PyObject>> {let mut close_polygons = vec![];
for poly in polygons {if norm(poly.center - point) < max_dist {close_polygons.push(poly)}}
Ok(close_polygons)}

可惜未能通過編譯：

我們需要 3 個(gè) crate 才能實(shí)現(xiàn)函數(shù)：

# For Rust-native array operations.ndarray = '0.15'# For a `norm` function for arrays.ndarray-linalg = '0.16'# For accessing numpy-created objects, based on `ndarray`.numpy = '0.18'

首先，我們將 point: PyObject 轉(zhuǎn)換成可以使用的東西。

我們可以利用 PyO3 來轉(zhuǎn)換 numpy 數(shù)組：

現(xiàn)在 point 變成了 ArrayView1，我們可以直接使用了。例如：

// Make the `norm` function available.use ndarray_linalg::Norm;assert_eq!((point.to_owned() - point).norm(), 0.);

接下來，我們需要獲取每個(gè)多邊形的中心，然后將其轉(zhuǎn)換成 ArrayView1。

雖然信息量有點(diǎn)大，但總的來說，結(jié)果就是逐行轉(zhuǎn)換原來的代碼：

use pyo3::prelude::*;use ndarray_linalg::Norm;use numpy::PyReadonlyArray1;#[pyfunction]fn find_close_polygons(py: Python<'_>,polygons: Vec<PyObject>,point: PyReadonlyArray1<f64>,max_dist: f64,) -> PyResult<Vec<PyObject>> {let mut close_polygons = vec![];let point = point.as_array();for poly in polygons {let center = poly.getattr(py, 'center')?.extract::<PyReadonlyArray1<f64>>(py)?.as_array().to_owned();if (center - point).norm() < max_dist {close_polygons.push(poly)}}Ok(close_polygons)}

對(duì)比一下原來的代碼：

我們希望這個(gè)版本優(yōu)于原來的函數(shù)，但究竟有多少提升呢？

$ (cd ./poly_match_rs/ && maturin develop)$ python measure.pyTook an avg of 609.46ms per iteration

看起來 Rust 非常慢？實(shí)則不然，使用maturin develop --release運(yùn)行，就能獲得更好的結(jié)果：

這個(gè)速度提升很不錯(cuò)啊。

我們還想查看我們的原生代碼，因此發(fā)布時(shí)需要啟用調(diào)試符號(hào)。即便啟用了調(diào)試，我們也希望看到最大速度。

# added to Cargo.toml[profile.release]debug = true # Debug symbols for our profiler.lto = true # Link-time optimization.codegen-units = 1 # Slower compilation but faster code.

第二版：用 Rust 重寫更多代碼

接下來，在 py-spy 中通過 --native 標(biāo)志，查看 Python 代碼與新版的原生代碼。

再次運(yùn)行 py-spy：

這次得到的火焰圖如下所示（添加紅色之外的顏色，以方便參考）：

看看分析器的輸出，我們發(fā)現(xiàn)了一些有趣的事情：

1.find_close_polygons::...::trampoline（Python 直接調(diào)用的符號(hào)）和__pyfunction_find_close_polygons（我們的實(shí)現(xiàn)）的相對(duì)大小。

? 可以看到二者分別占據(jù)了樣本的 95% 和 88%，因此額外開銷非常小。

2.實(shí)際邏輯(if (center - point).norm() < max_dist { ... }) 是 lib_v1.rs:22（右側(cè)非常小的框），大約占總運(yùn)行時(shí)間的 9%。

? 所以應(yīng)該可以實(shí)現(xiàn) 10 倍的提升。

3.大部分時(shí)間花在 lib_v1.rs:16 上，它是 poly.getattr(...).extract(...)，可以看到實(shí)際上只是 getattr 以及使用 as_array 獲取底層數(shù)組。

也就是說，我們需要專心解決第 3 點(diǎn)，而解決方法是用 Rust 重寫 Polygon。

我們來看看目標(biāo)類：

@dataclassclass Polygon:x: np.arrayy: np.array_area: float = None@cached_propertydef center(self) -> np.array:centroid = np.array([self.x, self.y]).mean(axis=1)return centroiddef 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

我們希望盡可能保留現(xiàn)有的 API，但我們不需要 area 的速度大幅提升。

實(shí)際的類可能有其他復(fù)雜的東西，比如 merge 方法——使用了 scipy.spatial 中的 ConvexHull。

為了降低成本，我們只將 Polygon 的“核心”功能移至 Rust，然后從 Python 中繼承該類來實(shí)現(xiàn) API 的其余部分。

我們的 struct 如下所示：

下面，我們需要實(shí)現(xiàn)這個(gè) struct。我們先公開 poly.{x, y, center}，作為：

? 屬性

? numpy 數(shù)組

我們還需要一個(gè) constructor，以便 Python 創(chuàng)建新的 Polygon：

use numpy::{PyArray1, PyReadonlyArray1, ToPyArray};#[pymethods]impl Polygon {#[new]fn new(x: PyReadonlyArray1<f64>, y: PyReadonlyArray1<f64>) -> 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,}}
// the `Py<..>` in the return type is a way of saying 'an Object owned by Python'.#[getter]               fn x(&self, py: Python<'_>) -> PyResult<Py<PyArray1<f64>>> {Ok(self.x.to_pyarray(py).to_owned()) // Create a Python-owned, numpy version of `x`.}// Same for `y` and `center`.}

我們需要將這個(gè)新的 struct 作為類添加到模塊中：

然后更新 Python 代碼：

class Polygon(poly_match_rs.Polygon):_area: float = Nonedef area(self) -> float:...

下面，編譯代碼——雖然可以運(yùn)行，但速度非常慢！

為了提高性能，我們需要從 Python 的 Polygon 列表中提取基于 Rust 的 Polygon。

PyO3 可以非常靈活地處理這類操作，所以我們可以通過幾種方法來完成。我們有一個(gè)限制是我們還需要返回 Python 的 Polygon，而且我們不想克隆任何實(shí)際數(shù)據(jù)。

我們可以針對(duì)每個(gè) PyObject 調(diào)用 .extract::<Polygon>(py)?，但也可以要求 PyO3 直接給我們 Py<Polygon>。

這是對(duì) Python 擁有的對(duì)象的引用，我們希望它包含原生 pyclass 結(jié)構(gòu)的實(shí)例（或子類，在我們的例子中）。

下面，我們來看看使用這些代碼的效果如何：

$ python measure.pyTook an avg of 6.29ms per iteration

我們快要成功了，只需再提升一倍的速度即可。

第三版：避免內(nèi)存分配

我們?cè)賮砜匆豢捶治銎?/span>的結(jié)果。

1.首先，我們看到 select_best_polygon，現(xiàn)在它調(diào)用的是一些 Rust 代碼（在獲取 x 和 y 向量時(shí)）。

? 我們可以解決這個(gè)問題，但這是一個(gè)非常小的提升（大約為 10%）。

2.我們看到 extract_argument 花費(fèi)了大約 20% 的時(shí)間（在 lib_v2.rs:48 下），這個(gè)開銷相對(duì)比較大。

? 但大部分時(shí)間都花在了 PyIterator::next 和 PyTypeInfo::is_type_of 中，這可不容易修復(fù)。

3.我們看到大量時(shí)間花在了內(nèi)存分配上。

? lib_v2.rs:58 是我們的 if 語句，我們還看到了drop_in_place和to_owned。

? 實(shí)際的代碼大約占總時(shí)間的 35%，遠(yuǎn)超我們的預(yù)期。所有數(shù)據(jù)都已存在，所以這一段本應(yīng)非常快。

下面，我們來解決最后一點(diǎn)。

有問題的代碼如下：

我們希望避免 to_owned。但是，我們需要一個(gè)已擁有的 norm 對(duì)象，所以我們必須手動(dòng)實(shí)現(xiàn)。

具體的寫法如下：

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)}

然而，借用檢查器報(bào)錯(cuò)了：

借用檢查器是正確的，我們使用內(nèi)存的方式不正確。

更簡單的修復(fù)方法是直接克隆，然后 close_polygons.push(poly.clone()) 就可以通過編譯了。

這實(shí)際上是一個(gè)開銷很低的克隆，因?yàn)槲覀冎辉黾恿?Python 對(duì)象的引用計(jì)數(shù)。

然而，在這個(gè)例子中，我們也可以通過一個(gè) 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 只在內(nèi)部范圍內(nèi)被借用，如果我們接近 close_polygons.pus，編譯器就可以知道我們不再持有引用，因此就可以通過編譯。

最后的結(jié)果：

相較于原來的代碼，整體性能得到了 100 倍的提升。

總結(jié)

我們?cè)瓉淼?Python 代碼如下：

@dataclassclass Polygon:x: np.arrayy: np.array_area: float = None@cached_propertydef center(self) -> np.array:centroid = np.array([self.x, self.y]).mean(axis=1)return centroiddef 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 對(duì)其進(jìn)行了分析，即便用最簡單的、逐行轉(zhuǎn)換的 find_close_polygons，也可以獲得 10 倍的性能提升。

我們反復(fù)進(jìn)行分析-修改代碼-測量結(jié)果，并最終獲得了 100 倍的性能提升，同時(shí) API 仍然保持與原來的庫相同。

最終得到的 Python 代碼如下：

調(diào)用的 Rust 代碼如下：

use pyo3::prelude::*;use ndarray::Array1;use ndarray_linalg::Scalar;use numpy::{PyArray1, PyReadonlyArray1, ToPyArray};#[pyclass(subclass)]struct Polygon {x: Array1<f64>,y: Array1<f64>,center: Array1<f64>,}#[pymethods]impl Polygon {#[new]fn new(x: PyReadonlyArray1<f64>, y: PyReadonlyArray1<f64>) -> 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<'_>) -> PyResult<Py<PyArray1<f64>>> {Ok(self.x.to_pyarray(py).to_owned())}// Same for `y` and `center`.}#[pyfunction]fn find_close_polygons(py: Python<'_>,polygons: Vec<Py<Polygon>>,point: PyReadonlyArray1<f64>,max_dist: f64,) -> PyResult<Vec<Py<Polygon>>> {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 < max_dist {close_polygons.push(poly)}}Ok(close_polygons)}#[pymodule]fn poly_match_rs(_py: Python, m: &PyModule) -> PyResult<()> {m.add_class::<Polygon>()?;m.add_function(wrap_pyfunction!(find_close_polygons, m)?)?;Ok(())}

要點(diǎn)總結(jié)

? Rust（在 PyO3 的幫助下）能夠以非常小的代價(jià)換取 Python 代碼性能的大幅提升。

? 對(duì)于研究人員來說，Python API 非常優(yōu)秀，同時(shí)使用 Rust 快速構(gòu)建基本功能是一個(gè)非常強(qiáng)大的組合。

? 分析非常有趣，可以幫助你了解代碼中的一切。