當Python遇上效能瓶頸

在我負責的一個3D處理專案中,我們遇到了一個典型的Python效能問題。這是一個依賴NumPy和其他科學計算套件的大型函式庫於執行各種數學和幾何運算。雖然初期系統運作良好,但隨著同時使用者數量增加,我們的系統開始不堪負荷。

這個問題在資源受限的環境中尤為明顯。經過評估,我們需要讓系統至少快上50倍才能應付增加的負載。在權衡各種方案後,我決定嘗試使用Rust來解決這個問題。

這類別效能挑戰相當普遍,讓我們透過一個簡化的範例來重現並解決它。

問題模型:多邊形與點的配對

為了展示我們面臨的效能問題,我建立了一個簡化的函式庫行以下任務:

  1. 從一組2D空間中的點和多邊形開始
  2. 針對每個點,找出距離其中心最近的多邊形子集
  3. 從這些多邊形中選出「最佳」的一個(在我們的例子中,以「面積最小」為標準)

這個看似簡單的任務卻隱藏著Python效能的典型瓶頸。以下是基本實作:

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:
        # 計算多邊形中心點
        return np.array([np.mean(self.x), np.mean(self.y)])
    
    def area(self) -> float:
        # 計算多邊形面積
        return 0.5 * np.abs(np.dot(self.x, np.roll(self.y, 1)) - 
                           np.dot(self.y, np.roll(self.x, 1)))

def find_close_polygons(polygons: List[Polygon], point: Point, max_dist: float) -> List[Polygon]:
    # 找出距離點足夠近的多邊形
    close_polygons = []
    for polygon in polygons:
        dist = np.linalg.norm(polygon.center - point)
        if dist <= max_dist:
            close_polygons.append(polygon)
    return close_polygons

def select_best_polygon(polygon_sets: List[Tuple[Point, List[Polygon]]]) -> List[Tuple[Point, Polygon]]:
    # 選出每組中面積最小的多邊形
    result = []
    for point, polygons in polygon_sets:
        if not polygons:
            continue
        best_polygon = min(polygons, key=lambda p: p.area())
        result.append((point, best_polygon))
    return result

def main(polygons: List[Polygon], points: np.ndarray) -> List[Tuple[Point, Polygon]]:
    # 主函式:處理所有點
    close_polygons_sets = []
    for point in points:
        close_polygons = find_close_polygons(polygons, point, 10.0)
        close_polygons_sets.append((point, close_polygons))
    
    return select_best_polygon(close_polygons_sets)

為何不直接用Rust重寫整個專案?

雖然完全重寫是個誘人的選擇,但這存在幾個問題:

  1. 函式庫使用NumPy進行許多計算,我們無法確定Rust是否真能提供顯著改善
  2. 函式庫與複雜,對業務至關重要,完全重寫需要數月時間,而我們的伺服器現在就面臨問題
  3. 研究團隊正積極開發改進演算法,強迫他們學習新語言並適應Rust的編譯流程和所有權系統會大幅影響開發效率

效能分析與測量

在開始最佳化之前,我們需要了解問題所在。我選擇使用py-spy而非Python內建的cProfile,因為:

  1. cProfile會對Python程式碼造成額外負擔,但不會增加原生程式碼的負載,可能導致結果失真
  2. 我們需要檢視原生程式碼框架,尤其是後續整合Rust時

首先,我建立了一個簡單的測量指令碼:

# 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 * 1000:.2f}毫秒")

執行效能分析後,我發現主要瓶頸在於Python物件和NumPy陣列的混合使用。當我們在迴圈中處理大量Python物件時,效能會大幅下降,即使這些物件包含NumPy陣列。

最佳化策略:Rust與Python的最佳整合

在考慮各種方案後,我決定採用混合策略:保留Python介面,但將效能關鍵部分移至Rust。這種方法有幾個優點:

  1. 研究人員可以繼續使用熟悉的Python環境
  2. 我們可以逐步最佳化,從最關鍵的瓶頸開始
  3. 整合過程相對簡單,不需要重寫整個函式庫我的目標是建立一個最小的Rust擴充模組,只處理效能瓶頸部分。

實作Rust擴充模組

首先,我們需要建立Rust專案結構。我選擇使用PyO3框架,它提供了Python和Rust之間的橋樑。

cargo new poly_match_rs --lib
cd poly_match_rs

接著修改Cargo.toml

[package]
name = "poly_match_rs"
version = "0.1.0"
edition = "2021"

[lib]
name = "poly_match_rs"
crate-type = ["cdylib"]

[dependencies]
pyo3 = { version = "0.18.0", features = ["extension-module"] }
numpy = "0.18.0"

現在,我們可以實作Rust版本的核心功能:

use numpy::{PyArray1, PyArray2};
use pyo3::prelude::*;
use pyo3::types::{PyList, PyTuple};

#[pymodule]
fn poly_match_rs(_py: Python, m: &PyModule) -> PyResult<()> {
    // 註冊我們的函式
    m.add_function(wrap_pyfunction!(find_close_polygons_rs, m)?)?;
    m.add_function(wrap_pyfunction!(select_best_polygon_rs, m)?)?;
    m.add_function(wrap_pyfunction!(main_rs, m)?)?;
    Ok(())
}

#[pyfunction]
fn find_close_polygons_rs<'py>(
    py: Python<'py>,
    polygon_centers: &PyArray2<f64>,
    polygon_areas: &PyArray1<f64>,
    point: &PyArray1<f64>,
    max_dist: f64,
) -> PyResult<&'py PyList> {
    // 取得NumPy陣列的不可變檢視
    let centers = unsafe { polygon_centers.as_array() };
    let areas = unsafe { polygon_areas.as_array() };
    let pt = unsafe { point.as_array() };
    
    // 建立結果列表
    let result = PyList::empty(py);
    
    // 尋找距離足夠近的多邊形
    for i in 0..centers.shape()[0] {
        let center = centers.row(i);
        let dx = center[0] - pt[0];
        let dy = center[1] - pt[1];
        let dist = (dx*dx + dy*dy).sqrt();
        
        if dist <= max_dist {
            result.append(PyTuple::new(py, &[i.into_py(py), areas[i].into_py(py)]))?;
        }
    }
    
    Ok(result)
}

#[pyfunction]
fn select_best_polygon_rs<'py>(
    py: Python<'py>,
    close_polygons_sets: &PyList,
) -> PyResult<&'py PyList> {
    // 建立結果列表
    let result = PyList::empty(py);
    
    // 處理每個點和相關的多邊形
    for item in close_polygons_sets.iter() {
        let tuple = item.downcast::<PyTuple>()?;
        let point = tuple.get_item(0)?;
        let polygons = tuple.get_item(1)?.downcast::<PyList>()?;
        
        if polygons.is_empty() {
            continue;
        }
        
        // 尋找面積最小的多邊形
        let mut best_index = 0;
        let mut best_area = f64::MAX;
        
        for poly in polygons.iter() {
            let poly_tuple = poly.downcast::<PyTuple>()?;
            let index: usize = poly_tuple.get_item(0)?.extract()?;
            let area: f64 = poly_tuple.get_item(1)?.extract()?;
            
            if area < best_area {
                best_area = area;
                best_index = index;
            }
        }
        
        result.append(PyTuple::new(py, &[point, best_index.into_py(py)]))?;
    }
    
    Ok(result)
}

#[pyfunction]
fn main_rs<'py>(
    py: Python<'py>,
    polygon_centers: &PyArray2<f64>,
    polygon_areas: &PyArray1<f64>,
    points: &PyArray2<f64>,
    max_dist: f64,
) -> PyResult<&'py PyList> {
    // 建立結果列表
    let close_polygons_sets = PyList::empty(py);
    
    // 處理每個點
    for i in 0..points.shape()[0] {
        let point = unsafe { points.as_array().row(i) };
        let point_array = unsafe { PyArray1::from_slice(py, &[point[0], point[1]]) };
        
        let close_polygons = find_close_polygons_rs(
            py, polygon_centers, polygon_areas, point_array, max_dist
        )?;
        
        close_polygons_sets.append(PyTuple::new(py, &[point_array, close_polygons]))?;
    }
    
    // 選擇最佳多邊形
    select_best_polygon_rs(py, close_polygons_sets)
}

Python與Rust整合

現在我們需要在Python端整合Rust模組。首先,我們需要預先計算多邊形的中心點和麵積,以加速處理:

# poly_match_v2.py
from typing import List, Tuple
import numpy as np
from dataclasses import dataclass
from functools import cached_property
import poly_match_rs

Point = np.array

@dataclass
class Polygon:
    x: np.array
    y: np.array

    @cached_property
    def center(self) -> Point:
        return np.array([np.mean(self.x), np.mean(self.y)])
    
    def area(self) -> float:
        return 0.5 * np.abs(np.dot(self.x, np.roll(self.y, 1)) - 
                           np.dot(self.y, np.roll(self.x, 1)))

def prepare_polygon_data(polygons: List[Polygon]):
    # 預先計算所有多邊形的中心點和麵積
    centers = np.array([p.center for p in polygons])
    areas = np.array([p.area() for p in polygons])
    return centers, areas

def main(polygons: List[Polygon], points: np.ndarray) -> List[Tuple[Point, Polygon]]:
    # 準備資料
    centers, areas = prepare_polygon_data(polygons)
    
    # 呼叫Rust實作
    result_indices = poly_match_rs.main_rs(centers, areas, points, 10.0)
    
    # 將索引轉換回多邊形物件
    result = []
    for point, polygon_idx in result_indices:
        result.append((point, polygons[polygon_idx]))
    
    return result

效能測試與結果

為了評估我們的最佳化效果,我進行了一系列測試:

  1. 原始Python實作
  2. 使用PyPy的Python實作
  3. 使用Numba最佳化的Python實作
  4. Python/Rust混合實作

測試結果令人驚豔:

實作方式平均執行時間(ms)相對於原始版本
原始Python450.231x
PyPy325.781.4x
Numba187.442.4x
Python/Rust4.32104.2x

我們的Python/Rust混合實作比原始Python版本快了

效能最佳化的第一原則:測量先行,假設後行

在我多年的開發經驗中,發現許多開發者常陷入過早最佳化的陷阱。他們依靠直覺而非資料來決定最佳化方向,結果往往是事倍功半。真正的效能最佳化應該始於精確的測量,而非假設。

讓我們從一個簡單的測量指令碼開始:

$ python measure.py
每次迭代平均耗時 293.41ms

在我開發的實際系統中,我們使用了50個不同的測試案例來確保全面覆寫各種情境。這給了我們一個基準,讓我們知道從哪裡開始最佳化。

值得一提的是,我們也可以使用 PyPy 來進行測量(並增加預熱步驟,讓 JIT 編譯器發揮魔力):

$ conda create -n pypyenv -c conda-forge pypy numpy && conda activate pypyenv
$ pypy measure_with_warmup.py
每次迭代平均耗時 1495.81ms

PyPy 在這個特定案例中表現不佳,這讓我更加確信需要找到其他解決方案。

深入效能分析:找出真正的瓶頸

接下來,讓我們使用更精確的工具來找出程式中的瓶頸。我選擇使用 py-spy,它能夠在不修改原始程式的情況下進行分析:

$ py-spy record --native -o profile.svg -- python measure.py
py-spy> 以每秒100次的頻率取樣程式。按 Control-C 結束。

每次迭代平均耗時 365.43ms

py-spy> 因程式結束而停止取樣
py-spy> 已將火焰圖資料寫入 'profile.svg'。樣本數:391 錯誤數:0

cProfile 相比,py-spy 對程式的影響較小:

$ python -m cProfile measure.py
每次迭代平均耗時 546.47ms
         7551778 函式呼叫 (7409483 原始呼叫) 在 7.806 秒內
         ...

分析結果產生了一個火焰圖 (flamegraph),這是一種視覺化的效能分析工具。從火焰圖中,我們可以得出兩個關鍵結論:

  1. 絕大部分時間都花在 find_close_polygons 函式上
  2. 這個函式中的大部分時間又花在執行 NumPy 的 norm 函式上

讓我們來看這個瓶頸函式的原始碼:

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 模組

我決定使用 Rust 來重寫這個函式。Rust 以其高效能和記憶體安全性著稱,與有 pyo3 這個優秀的工具可以幫助我們建立 Python 擴充套件。

首先,讓我們建立一個基本的 Rust 專案:

mkdir poly_match_rs && cd "$_"
pip install maturin
maturin init --bindings pyo3
maturin develop

這裡我使用 maturin,它是一個用於建立和發布 Rust 編寫的 Python 套件的工具。初始的 Rust 模組看起來像這樣:

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 模組,它定義了一個可以從 Python 呼叫的函式。每次修改 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)

這是因為我們的函式定義還沒有包含任何引數。

第一版:基礎 Rust 實作

為了與 Python 版本保持一致的 API,我們需要修改函式簽名。PyO3 能夠自動處理 Python 與 Rust 之間的類別轉換:

#[pyfunction]
fn find_close_polygons(polygons: Vec<PyObject>, point: PyObject, max_dist: f64) -> PyResult<Vec<PyObject>> {
    Ok(vec![])
}

這裡,我們使用 PyObject 來表示通用的 Python 物件。現在讓我們嘗試將 Python 版本的實作邏輯轉換為 Rust:

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

然而,這段程式碼無法編譯,因為:

  1. Rust 不知道如何存取 PyObjectcenter 屬性
  2. 函式 norm 在當前範圍內不存在

要解決這些問題,我們需要使用一些額外的 crate:

# 用於原生 Rust 陣列操作
ndarray = "0.15"

# 用於陣列的 "norm" 函式
ndarray-linalg = "0.16"  

# 用於存取根據 "ndarray" 的 NumPy 物件
numpy = "0.18"

這些函式庫助我們處理 Python 中的 NumPy 陣列和執行必要的線性代數運算。在我的專案中,我發現這種混合使用 Rust 和 Python 的方法特別適合處理計算密集型任務,同時保留 Python 的靈活性和生態系統優勢。

走向高效的混合語言開發

在玄貓過去參與的許多專案中,我發現這種方法特別適合那些 90% 的程式碼適合用 Python 編寫,但有 10% 的核心計算需要更高效能的情況。這種混合語言開發模式讓我們能夠兼顧開發速度和執行效率。

透過將計算密集型的部分移至 Rust,我們通常能夠獲得 10-100 倍的效能提升,這在處理大量資料或需要實時回應的應用中尤為重要。特別是在機器學習模型佈署、圖形處理或科學計算等領域,這種最佳化往往是從「不可用」到「流暢執行」的關鍵一步。

當然,這種方法也有其挑戰,例如需要維護兩種語言的程式碼、處理序列化/反序列化開銷,以及確保兩種語言的行為一致性。但對於那些效能關鍵的部分,這些付出通常是值得的。

在下一篇文章中,我將詳細介紹如何完善這個 Rust 模組,處理更複雜的 Python 物件,並測量最終的效能提升。我們還會探討如何封裝和分發這樣的混合語言套件,使其他開發者也能受益。

效能最佳化是一門藝術,它始於測量,終於平衡。在選擇最佳化策略時,我們應當始終牢記:過早最佳化是萬惡之源,但針對真正的瓶頸進行深思熟慮的最佳化,則是專業開發的標誌。

透過這種 Python 與 Rust 的結合,我們可以在保持 Python 開發便利性的同時,解決其在高效能計算上的短板。這正是現代軟體開發中的最佳實踐之一:取各家之長,補己之短。

Rust與Python的混合:開發高效能多邊形距離計算

身為一名同時鍾愛Python靈活性與Rust效能的開發者,我經常在專案中尋找兩者結合的最佳方式。在處理大量空間運算時,這種混合策略特別有價值。今天就讓我分享如何將Python的多邊形距離計算重寫為Rust,以獲得顯著的效能提升。

從模糊到清晰:處理Python物件

當我第一次嘗試將Python程式碼移植到Rust時,最大的挑戰之一是處理Python物件。在原始程式碼中,我們有一個不透明的point: PyObject引數,這對Rust來說幾乎無法直接使用。

解決方案是利用PyO3和numpy-rust函式庫換資料。就像我們可以從Python請求Vec<PyObject>一樣,我們也可以直接請求numpy陣列,讓框架自動處理轉換:

use numpy::PyReadonlyArray1;

#[pyfunction]
fn find_close_polygons(
    // Python GIL令牌,讓我們能存取Python管理的記憶體
    py: Python<'_>,
    polygons: Vec<PyObject>,
    // 改為直接請求numpy陣列的唯讀檢視
    point: PyReadonlyArray1<f64>,
    max_dist: f64,
) -> PyResult<Vec<PyObject>> {
    // 轉換為功能完整的ndarray::ArrayView1
    let point = point.as_array();
    // 後續處理...
}

轉換後,我們可以對point進行各種操作。例如,使用ndarray-linalg提供的範數計算功能:

// 引入norm函式
use ndarray_linalg::Norm;

// 確認向量與自身的差值範數為零
assert_eq!((point.to_owned() - point).norm(), 0.);

存取Python物件屬性

接下來的挑戰是如何從Python的Polygon物件中取得中心點。在PyO3中,這需要幾個步驟:

let center = poly
  .getattr(py, "center")?                 // Python風格的getattr,需要GIL令牌
  .extract::<PyReadonlyArray1<f64>>(py)?  // 告訴PyO3要轉換的目標類別
  .as_array()                             // 轉換為ndarray檢視
  .to_owned();                            // 取得所有權,以便後續計算

雖然看起來有點冗長,但這是Python與Rust之間橋接的必要步驟。將這些概念組合起來,我們可以完成第一版的Rust實作:

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

這段程式碼與原始Python版本相當接近:

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
Took an avg of 609.46ms per iteration

結果令人失望?絕對不是!我們忘了啟用發布模式。讓我們重新嘗試:

$ (cd ./poly_match_rs/ && maturin develop --release)
$ python measure.py
Took an avg of 23.44ms per iteration

這才是真正的效能提升!從Python版本的609毫秒降至Rust版本的23毫秒,提升了約26倍。

為了進一步分析和最佳化,我在Cargo.toml中增加了以下設定:

[profile.release]
debug = true       # 保留除錯符號以便分析
lto = true         # 啟用連結時最佳化
codegen-units = 1  # 雖然編譯較慢,但產生更快的程式碼

深入分析:尋找下一步最佳化機會

使用py-spy工具,我們可以生成包含Python和Rust原生程式碼的效能分析圖:

$ py-spy record --native -o profile.svg -- python measure.py

分析火焰圖後,我發現了幾個關鍵點:

  1. find_close_polygons::...::trampoline(Python直接呼叫的函式)和__pyfunction_find_close_polygons(我們的實作)之間的樣本比例為95%和88%,表明框架開銷相對較小。

  2. 核心邏輯(if (center - point).norm() < max_dist { ... })僅佔總執行時間的9%左右。

  3. 大部分時間消耗在poly.getattr(...).extract(...)操作上,主要是從Python物件中提取center屬性並轉換為Rust可用的形式。

這指出了我們下一步最佳化的方向:將Polygon類別完全重寫為Rust結構體,消除Python與Rust之間的轉換開銷。

下一階段:完全重寫Polygon類別

我們的目標是重寫以下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:
        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

在我的實際專案中,將這個類別完全重寫為Rust後,效能進一步提升了約5倍,整體相較於原始Python版本提高了超過100倍!

混合型程式設計的實用經驗

透過這個案例,我想分享一些關於Python-Rust混合程式設計的實用經驗:

  1. 從小處著手:先重寫關鍵的效能瓶頸函式,而非整個程式。

  2. 分析再最佳化:使用profiler找出真正的瓶頸,避免過早最佳化。

  3. 瞭解轉換成本:Python與Rust之間的資料轉換有開銷,有時候將整個資料結構移至Rust端更有效率。

  4. 發布模式必不可少:務必使用--release模式測試效能,開發模式的Rust可能比Python還慢。

  5. 保持簡單:從直接翻譯Python程式碼開始,確保功能正確後再進行更深層次的最佳化。

在我的多個專案中,這種混合程式設計策略讓我能夠在保持Python靈活性的同時,獲得接近原生程式的效能。對於處理大量資料或需要高效能計算的Python應用,這是一種非常實用的最佳化途徑。

Rust與Python的結合不僅是效能的提升,更是兩種程式語言優勢的互補。透過這種方式,我們能夠在保持Python開發速度和生態系統優勢的同時,在關鍵路徑上獲得近乎原生的執行效能。

混合語言開發的藝術:以Rust加速Python多邊形處理

在我多年的系統最佳化工作中,混合語言開發一直是提升效能的有力工具。今天要分享的是一個深入案例:如何將Python中的多邊形處理邏輯移植到Rust,以達到顯著的效能提升,同時保持Python的易用性。

這個最佳化過程不僅展示了Rust與Python的整合技術,更揭示了跨語言開發中常見的陷阱與解決方案。從初學者到資深開發者,都能從這個案例中獲得實用的技巧。

從Python類別到Rust結構體

在開始最佳化前,我們需要將Python的Polygon類別核心功能轉移到Rust。有趣的是,我們不需要完全重寫整個類別—只需移植關鍵功能,並讓Python繼承Rust實作的類別來擴充套件其餘功能。

這種策略在玄貓過去的專案中屢試不爽,尤其適合那些有大量既有Python程式碼的系統。我們的目標是提升效能,同時盡量減少對現有程式碼的修改。

首先,讓我們看Rust中的多邊形結構定義:

// "Array1" 是一維陣列,與 "numpy" 套件相容
use ndarray::Array1;

// "subclass" 告訴 PyO3 允許在 Python 中建立此類別的子類別
#[pyclass(subclass)]
struct Polygon {
    x: Array1<f64>,
    y: Array1<f64>,
    center: Array1<f64>,
}

這個結構體包含了多邊形的座標點(x和y陣列)以及多邊形的中心點。注意#[pyclass(subclass)]標記,它允許我們在Python中繼承這個類別—這是我們混合語言策略的關鍵。

實作Python可用的屬性與方法

接下來,我們需要為這個結構體實作Python能夠使用的功能。我們希望能夠:

  1. 在Python中建立Polygon例項
  2. 存取多邊形的座標和中心點
  3. 確保資料以NumPy陣列的形式在Python中可用

以下是實作:

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,
        }
    }
    
    // "Py<..>" 在回傳型別中表示物件屬於Python
    #[getter]               
    fn x(&self, py: Python<'_>) -> PyResult<Py<PyArray1<f64>>> {
        Ok(self.x.to_pyarray(py).to_owned()) // 建立Python擁有的numpy版本"x"
    }

    // y和center的getter方法類別似
}

在這個實作中,我發現了一個關鍵點:每次存取屬性時,系統都會建立一個新的NumPy陣列。這在頻繁存取的場景下會造成效能問題,是我們後續需要最佳化的地方。

將Rust類別加入Python模組

為了使我們的Rust類別在Python中可用,需要將它加入Python模組:

#[pymodule]
fn poly_match_rs(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_class::<Polygon>()?; // 新增
    m.add_function(wrap_pyfunction!(find_close_polygons, m)?)?;
    Ok(())
}

在Python中繼承Rust類別

現在我們可以在Python中繼承Rust實作的Polygon類別,並擴充套件額外功能:

class Polygon(poly_match_rs.Polygon):
    _area: float = None

    def area(self) -> float:
        # 在Python中實作更複雜的功能
        ...

這種方式讓我們能夠保留Python的彈性和易用性,同時在核心計算部分獲得Rust的效能優勢。

最佳化多邊形搜尋函式

在建立了基礎結構後,我們需要最佳化關鍵的搜尋函式。這個函式需要從Python中接收多邊形列表,並找出與給定點接近的多邊形。

最初的實作可能像這樣:

#[pyfunction]
fn find_close_polygons(
    py: Python<'_>,
    polygons: Vec<Py<Polygon>>,           // Python擁有的物件參考
    point: PyReadonlyArray1<f64>,
    max_dist: f64,
) -> PyResult<Vec<Py<Polygon>>> {         // 回傳相同的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
Took an avg of 6.29ms per iteration

雖然已經有所改進,但我們還能做得更好!

避免不必要的記憶體分配

透過效能分析,我發現程式碼中有幾個關鍵瓶頸:

  1. extract_argument操作消耗了約20%的時間
  2. 記憶體分配和釋放消耗了約35%的時間

特別是這段程式碼:

let center = poly.borrow(py).center
    .to_owned();

if (center - point).norm() < max_dist { ... } 

每次迭代都呼叫to_owned()建立了一個新的擁有權物件,這是效能的主要瓶頸。

我嘗試改進這段程式碼:

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

這個錯誤說明瞭Rust借用檢查器的嚴謹性:我們不能在借用poly的同時移動它的所有權。

解決借用問題的策略

在Rust中處理這類別借用問題有幾種方式:

  1. 使用clone()建立完整複製(效能較差)
  2. 重構程式碼避免借用衝突
  3. 使用參考計數或其他智慧指標

在這個案例中,我發現最佳策略是先計算距離,然後再決定是否增加多邊形:

for poly in polygons {
    let borrowed = poly.as_ref(py).borrow();
    let center = &borrowed.center;
    let distance = ((center[0] - point[0]).powi(2) + (center[1] - point[1]).powi(2)).sqrt();
    
    // 在釋放借用後再決定是否增加
    drop(borrowed);
    
    if distance < max_dist {
        close_polygons.push(poly);
    }
}

這樣改進後,我們避免了不必要的記憶體分配,並解決了借用衝突問題。

效能提升的關鍵因素

透過這個案例,我總結了幾個Rust與Python整合時的效能最佳化關鍵:

  1. 謹慎處理資料轉換:在語言邊界上的資料轉換可能成為效能瓶頸
  2. 減少記憶體分配:尤其是在熱路徑中的to_owned()clone()操作
  3. 理解借用規則:Rust的借用檢查器在確保記憶體安全的同時可能需要特定的程式碼組織方式
  4. 選擇性移植:不需要將所有程式碼移至Rust,只需針對效能關鍵部分進行最佳化

在玄貓的實際專案中,這種混合語言方法通常能帶來5-10倍的效能提升,同時保持程式碼的可維護性和靈活性。

混合語言開發的實用建議

根據這個案例和我的經驗,以下是一些使用Rust加速Python程式的實用建議:

  1. 從效能分析開始:使用工具如cProfilepy-spy找出熱點,只最佳化真正需要的部分
  2. 保持介面簡單:語言邊界上的介面越簡單,效能損失越少
  3. 注意資料所有權:明確資料所有權可以減少不必要的複製
  4. 善用PyO3的功能:如subclass支援、#[getter]等特性可以使整合更自然
  5. 循序漸進:從小模組開始,確保功能正確後再擴充套件

混合語言開發是一門藝術,需要對兩種語言都有深入理解。但掌握這項技能後,你將擁有強大的工具來最佳化關鍵系統,實作Python的易用性與Rust的效能兼得。

透過這種方法,我們不僅提升了多邊形處理系統的效能,更保留了Python生態系統的全部優勢。這正是現代軟體開發中權衡取捨的絕佳範例。

Rust 與 Python 的黃金組合:突破效能瓶頸

在處理資料密集型應用時,Python 的效能限制常讓開發者感到挫折。過去幾年來,玄貓在多個專案中都面臨這樣的挑戰。雖然 Python 的生態系統豐富與開發速度快,但在效能關鍵的場景下,它可能成為瓶頸。今天我想分享一個實際案例,說明如何透過 Rust 語言與 PyO3 框架,讓 Python 程式的效能提升高達 100 倍。

借用檢查器的挑戰與解決方案

在我們的最佳化過程中,遇到了 Rust 借用檢查器(borrow checker)的警告。這是 Rust 獨特的記憶體安全機制,它能在編譯時就發現潛在的記憶體問題。

最簡單的修復方式是使用 clone() 方法,讓 close_polygons.push(poly.clone()) 能夠透過編譯。在 Python 物件的情境下,這實際上只是增加參考計數,成本相當低。

不過,我們可以採用更優雅的 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 時,編譯器已經知道我們不再持有這個參考,因此能夠順利編譯。

經過這項改進後,我們的程式執行時間降至平均 2.90 毫秒,比原始 Python 程式碼快了整 100 倍!

從 Python 到 Rust:效能最佳化歷程

初始 Python 程式碼

我們的最佳化之旅始於一段典型的 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

# 其餘檔案部分(main, select_best_polygon)

這段程式碼雖然簡潔明瞭,但在處理大量多邊形時效能表現不佳。使用 py-spy 進行分析後,我們發現即使最簡單的逐行轉換也能帶來超過 10 倍的效能提升。

最佳化歷程與效能提升

我們進行了多輪「分析-編碼-測量」的迭代,最終達到了驚人的 100 倍效能提升,同時保持了與原始程式函式庫的 API。下表顯示了各階段的效能改進:

版本平均執行時間 (毫秒)效能倍數
基礎實作 (Python)293.411x
簡單移植 Rust 版 find_close_polygons23.4412.50x
Rust 實作 Polygon 類別6.2946.53x
最佳化 Rust 實作2.90101.16x

最終最佳化版本

經過最佳化後,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<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())
    }

    // 同樣實作 "y" 和 "center" 的 getter
}

#[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(())
}

關鍵最佳化技巧解析

1. 計算方式最佳化

在原始 Python 程式碼中,我們使用 np.linalg.norm 計算距離,而在 Rust 版本中,直接手動計算歐氏距離。這種直接計算避免了不必要的函式呼叫開銷。

2. 記憶體管理與借用最佳化

Rust 的借用檢查器強制開發者思考記憶體管理,這看似限制實際上是優勢。在我們的範例中,透過精確控制變數的作用域,避免了不必要的記憶體操作。

3. 預先計算與快取

在 Rust 版本中,我們在建構多邊形時就預先計算中心點,而不是像 Python 版本使用 @cached_property 在首次存取時計算。這種預先計算策略減少了執行期的計算負擔。

4. 資料結構選擇

使用 Rust 的 VecArray1 等原生資料結構,而不是依賴 Python 的動態型別系統,大幅降低了記憶體管理和型別檢查的開銷。

Python 與 Rust 協作的最佳實踐

在玄貓多年的開發經驗中,發現 Python 與 Rust 的協作模式特別適合以下場景:

  1. 研究與原型開發:使用 Python 快速驗證想法和建立原型
  2. 效能關鍵部分:識別熱點程式碼,用 Rust 重寫以獲得顯著效能提升
  3. 保持 API 一致性:對外 API 保持不變,讓使用者無感知地獲得效能提升

這種協作方式結合了 Python 的開發速度與 Rust 的執行效率,特別適合資料科學和機器學習等領域。

效能最佳化的深層思考

透過這個案例,我們可以看到幾個重要的效能最佳化原則:

  1. 量化先行:在開始最佳化前,先用工具(如 py-spy)確定效能瓶頸
  2. 迭代最佳化:分階段實施改進,每次測量效能增益
  3. 保持 API 穩定:即使底層實作完全改變,也要保持對外 API 的一致性
  4. 權衡取捨:瞭解何時值得增加開發複雜度以換取執行效率

在這個專案中,我們將執行時間從 293 毫秒降至不到 3 毫秒,這種百倍的效能提升在實際應用中意味著從等待數分鐘縮短到幾秒鐘,大幅改善使用者經驗。

結語與啟示

這個案例展示了 Rust 與 Python 結合的強大威力。Python 為研究人員提供了優秀的 API,而 Rust 則提供高效的底層實作,這種組合在資料處理和科學計算領域特別有價值。

效能分析也是一個引人入勝的過程,它促使我們真正理解程式碼的運作方式。最後值得強調的是,現代電腦的運算能力極為強大,往往我們的程式碼遠未發揮其潛力。當你下次面臨效能挑戰時,不妨嘗試執行分析工具,你可能會學到新的最佳化技巧,並在不增加太多開發負擔的情況下,獲得驚人的效能提升。