如何使用 Rust 加速 Python

最新推荐文章于 2024-05-02 15:36:55 发布

言蜕

最新推荐文章于 2024-05-02 15:36:55 发布

阅读量786

点赞数

文章标签： rust python

原文链接：https://towardsdatascience.com/how-to-make-your-python-packages-really-fast-with-rust-91a9bebacbc2

版权

本文翻译自原文地址，主要用于记录如何将 Rust 代码发布成 python 的包供 python 调用以提升性能。

引言

Python很慢，但这不仅仅是Python的问题，很多动态语言都会有这样的问题。因此性能至上的 Python 包的开发者都转向C语言，但C不够有趣，同时。下面会简单介绍一下Rust。

Rust 是一类内存高效的语言，没有运行时，也没有垃圾回收机制。Rust 难以置信地快，而且尤其可靠，同时还有一个很棒的社区。而且由于有了maturin 和 PyO3 等工具的存在，它使得嵌入你的Python代码变得尤其简单。

接下来将讲述如何一步一步地创建你的 Python 包，涉及到 Rust 的部分不会过于深入。

预备条件

开始之前，确保你的机器上安装了 Rust ，你可以直接从官网下载，按照引导安装，建议安装虚拟环境，后续测试 Rust 包的时候需要。

脚本预览

下面是一个脚本，给定一个数 $n$ ，将会计算斐波那契数列的第100个数，同时还会计算函数耗时多久。

import sys

from timeit import timeit

RUNS = 100

def fibonacci(n: int) -> int:

if n <= 1:

return n

return fibonacci(n - 1) + fibonacci(n - 2)

def main():

n = int(sys.argv[1])

print(f"{fibonacci(n) = }")

python_time_per_call = timeit(lambda: fibonacci(n), number=RUNS) / RUNS

print(f"\nPython μs per call: {python_time_per_call * 1_000_000:.2f} μs")

print(f"Python ms per call: {python_time_per_call * 1_000:.2f} ms")

if __name__ == "__main__":

main()

这是一个非常原始的，完全未经优化的函数，只使用 Python 也可以使它运行更快，但今天不打算从这些角度优化。我们将使用这段代码在 Rust 中创建一个 Python 包。

Maturin 的配置

第一步是安装 Maturin，这是一个构建系统，用于构建和发布 Rust 的包(Crate)作为 Python 包。你可以使用 pip install maturin 来实现。接下来，为你的包创建一个目录。最后一步是在新目录下运行 maturin init。此时，将提示您选择要使用的 Rust 绑定，选择 pyo3。

现在，如果你查看上一步构建的目录，会看到一些文件。Maturin 为我们创建了一些配置文件，即Cargo.toml和 pyproject.toml。Cargo.toml文件是 Rust 构建工具cargo的配置文件，它包含一些关于包的默认元数据、一些构建选项和 pyo3 的依赖项。 pyproject.toml文件是相当标准的，但它被设置为使用 maturin 作为构建后端。

Maturin还会创建一个GitHub Actions工作流来发布你的包。在你需要维护一个开源项目时，它会使工作更加轻松。我们最关心的文件是src目录下的 lib.rs 文件。

刚才设置之下的目录结构如下所示

fibbers/  
├── .github/  
│ └── workflows/  
│ └── CI.yml  
├── .gitignore  
├── Cargo.toml  
├── pyproject.toml  
└── src/  
└── lib.rs

Rust 部分

Maturin 已经使用我们前面提到的 PyO3 绑定为我们创建了Python模块的脚手架。

use pyo3::prelude::*;

/// Formats the sum of two numbers as string.

#[pyfunction]

fn sum_as_string(a: usize, b: usize) -> PyResult<String> {

Ok((a + b).to_string())

}

/// A Python module implemented in Rust.

#[pymodule]

fn fibbers(_py: Python, m: &PyModule) -> PyResult<()> {

m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;

Ok(())

}

这段代码的主要部分是sum_as_string函数，被标记为pyfunction属性，fibbers函数，代表我们的Python模块名。fibbers 函数所实现的唯一功能在于在 fibbers 模块中注册 sum_as_string函数。
现在进行安装，就可以从python中调用fibbers.sum_as_string()，并且得到的是预想的结果。

接下来我们将把函数 sum_as_string 替换为如下的 fib 函数

use pyo3::prelude::*;

/// Calculate the nth Fibonacci number.

#[pyfunction]

fn fib(n: u32) -> u32 {

if n <= 1 {

return n;

}

fib(n - 1) + fib(n - 2)

}

/// Fast Fibonacci number calculation.

#[pymodule]

fn fibbers(_py: Python, m: &PyModule) -> PyResult<()> {

m.add_function(wrap_pyfunction!(fib, m)?)?;

Ok(())

}

对比不同的实现

在安装 fibbers 包之前，我们所需做的是在终端运行 maturin develop 。这会下载并且编译我们的 Rust 包，并且安装到虚拟环境里面去。

接下来回到 fib.py 文件中去，我们可以导入fibbers ，输出 fibbers.fib() 并且加上一个 timeit 来衡量其性能。

import sys

from timeit import timeit

import fibbers

RUNS = 100

def fibonacci(n: int) -> int:

if n <= 1:

return n

return fibonacci(n - 1) + fibonacci(n - 2)

def main():

n = int(sys.argv[1])

print(f"{fibonacci(n) = }")

print(f"{fibbers.fib(n) = }")

python_time_per_call = timeit(lambda: fibonacci(n), number=RUNS) / RUNS

print(f"\nPython μs per call: {python_time_per_call * 1_000_000:.2f} μs")

print(f"Python ms per call: {python_time_per_call * 1_000:.2f} ms")

rust_time_per_call = timeit(lambda: fibbers.fib(n), number=RUNS) / RUNS

print(f"\nRust μs per call: {rust_time_per_call * 1_000_000:.2f} μs")

print(f"Rust ms per call: {rust_time_per_call * 1_000:.2f} ms")

if __name__ == "__main__":

main()

如果两个程序同时计算斐波那契数列的第十个数，可以看到 Rust 函数比 python 函数快 5 倍
计算第十个数

如果计算第20个和30个斐波那契数列，可以看到Rust比python快大约15倍。

计算第二十个数
计算第三十个数

但如果我告诉你们目前 Rust 还没有达到其速度的极限，可以看到默认情况下 maturin develop 会构建Rust包的Dev版本，该版本在构建的时候会放弃很多编译器的优化来减少编译时间，这意味着构建出的程序无法以其最高速度运行。此时你可以回到 fibbers 目录，再次执行 maturin develop 命令，不过这次加上 --release 选项，这次将能得到二进制包的最快版本。

如果我们现在获取斐波那契数列的第30个数，相对于python ，Rust 将给我们带来 40 倍速的提升。

优化之后的计算第三十个数