使用PyO3/Rust 构建 Python 扩展：实现高性能的 GIL 释放与并发计算
大家好！今天我们来探讨一个重要的主题：如何利用 PyO3/Rust 构建高性能的 Python 扩展，并充分利用 GIL 释放和并发计算来提升性能。Python 的 GIL（Global Interpreter Lock）一直是其并发性能的一大瓶颈。虽然多线程在 I/O 密集型任务中能带来一些提升，但在 CPU 密集型任务中，由于 GIL 的存在，多线程并不能真正实现并行计算。Rust 语言以其安全性、高性能以及与 C 语言的良好互操作性，成为了解决这一问题的理想选择。通过 PyO3，我们可以轻松地将 Rust 代码集成到 Python 中，并利用 Rust 的线程模型来绕过 GIL 的限制。
1. GIL 的本质与限制
首先，我们需要理解 GIL 的作用。GIL 确保同一时刻只有一个线程可以执行 Python 字节码。这简化了 Python 解释器的设计，避免了复杂的线程安全问题。然而，这也意味着在 CPU 密集型任务中，即使我们使用多线程，也无法真正利用多核 CPU 的优势。
举个例子，假设我们有一个计算密集型的函数 calculate_sum：
def calculate_sum(n: int) -> int:
"""
计算 1 到 n 的累加和。
"""
total = 0
for i in range(1, n + 1):
total += i
return total
import time
import threading
def worker(n: int, result: list):
start_time = time.time()
sum_result = calculate_sum(n)
end_time = time.time()
result.append((sum_result, end_time - start_time))
print(f"Thread finished, sum: {sum_result}, time: {end_time - start_time:.4f}s")
if __name__ == "__main__":
n = 100000000
num_threads = 4
results = []
threads = []
for _ in range(num_threads):
t = threading.Thread(target=worker, args=(n // num_threads, results))
threads.append(t)
t.start()
for t in threads:
t.join()
total_time = sum([r[1] for r in results])
print(f"Total time taken: {total_time:.4f}s")
在这个例子中，我们创建了多个线程来并行计算累加和。然而，由于 GIL 的限制，这些线程实际上是交替执行的，并没有真正实现并行加速。
2. PyO3 简介与环境配置
PyO3 是一个 Rust 库，用于创建 Python 扩展模块。它提供了方便的 API，使得我们能够轻松地将 Rust 代码集成到 Python 中。
首先，我们需要安装 Rust 和 Cargo：
curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh然后，我们需要创建一个新的 Rust 项目：
cargo new pyo3_example --lib
cd pyo3_example接下来，我们需要在 Cargo.toml 文件中添加 PyO3 依赖：
[package]
name = "pyo3_example"
version = "0.1.0"
edition = "2021"
[lib]
name = "pyo3_example" # 模块名，Python 中 import 的名字
crate-type = ["cdylib"]
[dependencies]
pyo3 = { version = "0.20", features = ["extension-module"] }
rayon = "1.9" # 添加 rayon 作为并发库注意：crate-type = ["cdylib"] 指定了我们将构建一个动态链接库，这是 Python 扩展模块的必要条件。我们还添加了 rayon 依赖，这是一个 Rust 的数据并行库，我们将使用它来实现并发计算。
3. 使用 PyO3 构建基本的 Python 扩展
下面是一个简单的 PyO3 扩展模块的例子：
use pyo3::prelude::*;
#[pyfunction]
fn greet(name: &str) -> PyResult<String> {
Ok(format!("Hello, {}!", name))
}
#[pymodule]
fn pyo3_example(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(greet)?)?;
Ok(())
}在这个例子中，我们定义了一个名为 greet 的函数，它接收一个字符串参数 name，并返回一个包含问候语的字符串。#[pyfunction] 宏将 Rust 函数暴露给 Python。#[pymodule] 宏定义了 Python 模块，并将 greet 函数添加到模块中。
要构建这个扩展模块，我们需要运行以下命令：
cargo build --release构建成功后，会在 target/release 目录下生成一个名为 libpyo3_example.so（或者 libpyo3_example.dylib，具体取决于操作系统）的动态链接库。
现在，我们可以在 Python 中导入并使用这个模块：
import pyo3_example
print(pyo3_example.greet("World")) # 输出: Hello, World!4. GIL 释放：允许 Rust 代码并行执行
PyO3 提供了多种方式来释放 GIL，允许 Rust 代码在没有 GIL 的情况下执行。最常用的方式是使用 Python::allow_threads。
下面是一个使用 Python::allow_threads 释放 GIL 的例子：
use pyo3::prelude::*;
use std::time::Duration;
use std::thread;
#[pyfunction]
fn long_running_task(duration_secs: u64) -> PyResult<()> {
Python::with_gil(|py| {
println!("Starting long running task...");
// 释放 GIL
py.allow_threads(|| {
thread::sleep(Duration::from_secs(duration_secs));
println!("Long running task finished.");
});
println!("Back in Python GIL context.");
});
Ok(())
}
#[pymodule]
fn pyo3_example(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(long_running_task)?)?;
Ok(())
}在这个例子中，long_running_task 函数模拟了一个耗时的操作。Python::allow_threads 闭包中的代码将在没有 GIL 的情况下执行。这意味着在执行 thread::sleep 时，Python 解释器可以执行其他的 Python 代码。
5. 使用 Rayon 实现并发计算
Rayon 是一个 Rust 的数据并行库，它提供了方便的 API，使得我们能够轻松地将计算任务分解成多个子任务，并在多个线程上并行执行。
下面是一个使用 Rayon 并行计算累加和的例子：
use pyo3::prelude::*;
use rayon::prelude::*;
#[pyfunction]
fn parallel_calculate_sum(n: usize) -> PyResult<usize> {
let result = (1..=n).into_par_iter().sum::<usize>();
Ok(result)
}
#[pymodule]
fn pyo3_example(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(parallel_calculate_sum)?)?;
Ok(())
}在这个例子中，parallel_calculate_sum 函数使用 Rayon 的 into_par_iter() 方法将 1..=n 范围内的整数转换为并行迭代器。然后，sum() 方法将并行迭代器中的所有元素累加起来。Rayon 会自动将计算任务分解成多个子任务，并在多个线程上并行执行。
为了在没有 GIL 的情况下运行 Rayon 计算，我们需要结合 Python::allow_threads 使用：
use pyo3::prelude::*;
use rayon::prelude::*;
#[pyfunction]
fn parallel_calculate_sum_gil_released(py: Python, n: usize) -> PyResult<usize> {
let result = py.allow_threads(|| {
(1..=n).into_par_iter().sum::<usize>()
});
Ok(result)
}
#[pymodule]
fn pyo3_example(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(parallel_calculate_sum_gil_released)?)?;
Ok(())
}在这个例子中，我们将 Rayon 的计算任务放在 Python::allow_threads 闭包中执行，从而确保计算任务在没有 GIL 的情况下并行执行。
6. 高级技巧：处理复杂的数据结构
在实际应用中，我们经常需要处理复杂的数据结构，例如 NumPy 数组。PyO3 提供了方便的 API，使得我们能够轻松地在 Rust 代码中访问和操作 NumPy 数组。
首先，我们需要在 Cargo.toml 文件中添加 numpy 依赖：
[dependencies]
pyo3 = { version = "0.20", features = ["extension-module", "numpy"] }
numpy = "0.20" # 确保 numpy 版本与 pyo3 兼容
rayon = "1.9"然后，我们可以使用 numpy crate 中的 API 来访问和操作 NumPy 数组：
use pyo3::prelude::*;
use pyo3::types::PyArray;
use numpy::{PyArray1, ToPyArray};
use rayon::prelude::*;
#[pyfunction]
fn process_numpy_array(py: Python, arr: &PyArray1<f64>) -> PyResult<Py<PyArray1<f64>>> {
// 获取 NumPy 数组的切片
let data = arr.as_slice().unwrap();
// 并行处理 NumPy 数组
let result: Vec<f64> = py.allow_threads(|| {
data.par_iter().map(|x| x * 2.0).collect()
});
// 将结果转换为 NumPy 数组
let result_array = result.to_pyarray(py).to_owned();
Ok(result_array)
}
#[pymodule]
fn pyo3_example(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(process_numpy_array)?)?;
Ok(())
}在这个例子中，process_numpy_array 函数接收一个 NumPy 数组作为参数，并将数组中的每个元素乘以 2。我们首先使用 as_slice() 方法获取 NumPy 数组的切片，然后使用 Rayon 并行处理切片中的每个元素。最后，我们将处理后的结果转换为 NumPy 数组并返回。
7. 错误处理与性能优化
在编写 PyO3 扩展时，错误处理和性能优化是非常重要的。
- 错误处理： PyO3 提供了
PyResult类型，用于处理可能发生的错误。我们可以使用?运算符来传播错误，或者使用PyErr::set方法来设置 Python 异常。 - 性能优化： 在编写 Rust 代码时，我们需要注意避免不必要的内存分配和复制。我们可以使用引用和借用来避免内存复制，或者使用
unsafe代码来绕过 Rust 的安全检查。
8. 示例：图像处理加速
让我们以一个图像处理的例子来展示 PyO3 的威力。 假设我们有一个 Python 函数，需要对图像的每个像素进行处理。 这个函数用纯 Python 实现会非常慢。
import numpy as np
import time
def process_image_python(image: np.ndarray) -> np.ndarray:
"""
纯 Python 实现的图像处理函数。
"""
height, width, channels = image.shape
new_image = np.zeros_like(image)
for y in range(height):
for x in range(width):
for c in range(channels):
new_image[y, x, c] = image[y, x, c] * 2 # 简单的像素值翻倍
return new_image
if __name__ == "__main__":
image = np.random.randint(0, 256, size=(512, 512, 3), dtype=np.uint8)
start_time = time.time()
new_image = process_image_python(image)
end_time = time.time()
print(f"Python processing time: {end_time - start_time:.4f}s")现在我们用 PyO3 和 Rust 实现相同的功能，并使用 Rayon 进行并行处理：
use pyo3::prelude::*;
use pyo3::types::PyArray;
use numpy::{PyArray3, ToPyArray};
use rayon::prelude::*;
#[pyfunction]
fn process_image_rust(py: Python, image: &PyArray3<u8>) -> PyResult<Py<PyArray3<u8>>> {
let (height, width, channels) = (image.shape()[0], image.shape()[1], image.shape()[2]);
let data = image.as_slice().unwrap();
let processed_data: Vec<u8> = py.allow_threads(|| {
data.par_iter()
.map(|&pixel| pixel.wrapping_mul(2)) // 使用 wrapping_mul 防止溢出
.collect()
});
let processed_image = PyArray::from_shape_vec(py, (height, width, channels), processed_data).unwrap();
Ok(processed_image.to_owned())
}
#[pymodule]
fn pyo3_example(_py: Python, m: &PyModule) -> PyResult<()> {
m.add_function(wrap_pyfunction!(process_image_rust)?)?;
Ok(())
}在 Python 中调用 Rust 函数：
import numpy as np
import time
import pyo3_example
def process_image_python(image: np.ndarray) -> np.ndarray:
height, width, channels = image.shape
new_image = np.zeros_like(image)
for y in range(height):
for x in range(width):
for c in range(channels):
new_image[y, x, c] = image[y, x, c] * 2
return new_image
if __name__ == "__main__":
image = np.random.randint(0, 256, size=(512, 512, 3), dtype=np.uint8)
start_time = time.time()
new_image_python = process_image_python(image)
end_time = time.time()
print(f"Python processing time: {end_time - start_time:.4f}s")
start_time = time.time()
new_image_rust = pyo3_example.process_image_rust(image)
end_time = time.time()
print(f"Rust processing time: {end_time - start_time:.4f}s")
编译并运行这段代码，你会发现 Rust 版本的图像处理速度远快于 Python 版本，尤其是在多核 CPU 上。
9. 注意事项与最佳实践
- 数据所有权： 在 Rust 和 Python 之间传递数据时，需要注意数据的所有权。通常情况下，我们需要将数据复制到 Rust 中，或者使用借用来避免复制。
- 内存安全： Rust 是一门内存安全的语言，但在使用
unsafe代码时，我们需要格外小心，避免出现内存泄漏和悬垂指针等问题。 - 错误处理： 在编写 PyO3 扩展时，我们需要充分考虑可能发生的错误，并使用
PyResult类型来处理错误。 - 测试： 我们需要编写充分的测试用例来验证 PyO3 扩展的正确性和性能。
表格：性能对比
| 实现方式 | 是否释放 GIL | 是否并行 | 性能（相对Python） |
|---|---|---|---|
| 纯 Python | 否 | 否 | 1x |
| PyO3 | 否 | 否 | 2-5x |
| PyO3 + GIL释放 | 是 | 否 | 2-5x |
| PyO3 + GIL释放 + Rayon | 是 | 是 | 10-50x |
总结
通过 PyO3，我们可以将 Rust 的高性能特性带到 Python 中，并通过 GIL 释放和并发计算来突破 Python 的性能瓶颈。 Rust 带来的内存安全，编译时检查和无畏并发特性都使得使用 PyO3 构建高性能扩展成为可能。 通过仔细的设计和优化，我们可以利用 Rust 和 Python 的各自优势，构建出高效、可靠的应用程序。
更多IT精英技术系列讲座，到智猿学院