应用-使用Rust加速Python的最简单方法

目录

使用Rust加速Python的最简单方法

如果你想加速现有的Python代码，使用Rust编写编译扩展可能是一个非常好的选择：

• 在许多情况下，Rust代码的运行速度比Python快得多。
• Rust可以防止C、C++和Cython代码中常见的内存管理错误。
• Rust有越来越多的第三方包生态系统，并且与C和C++不同，它还有一个内置的包管理器和构建系统。

然而，如果你只是想快速尝试Rust扩展，打包和集成样板代码可能会成为障碍：每增加一点摩擦都会阻碍你进行实验，看看Rust是否真的能带来帮助。

这就是rustimport的用武之地，它是一个库，使得独立的Rust文件可以轻松地在Python中导入（目前仅支持Linux和macOS）。在本文中，我们将介绍：

• 如何使用rustimport快速尝试你的Rust代码。
• Rust初学者最常见的性能错误，以及如何避免它。
• 使用rustimport时的一些注意事项。

注意： 与典型的Rust扩展一样，你仍然会使用PyO3来桥接Python和Rust，但正如我们将看到的，rustimport使得使用PyO3比替代的打包方法更加容易。

前提条件

rustimport会即时编译你的Rust代码，因此你需要在计算机上安装Rust编译器。这不是你最终打包库或应用程序时所需的，但我们的用例是原型设计，无论如何你都需要一个编译器。

你还需要安装rustimport，例如在虚拟环境或Conda环境中运行pip install rustimport。

示例：计算斐波那契数列

我们将用Rust实现斐波那契数列，但首先用Python实现，以便我们可以比较性能：

def fibonacci(number: int) -> int:
    if number == 0:
        return 0
    if number == 1:
        return 1
    prevprev = 0
    prev = 1
    current = 1
    for _ in range(number - 1):
        current = prevprev + prev
        prevprev = prev
        prev = current
    return current

然后我们可以使用IPython交互式提示的%timeit魔法来计时这个实现：

In [1]: from pyfib import fibonacci

In [2]: %timeit fibonacci(50)
954 ns ± 1.66 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

在Rust中实现斐波那契数列

让我们在Rust中实现相同的函数，并将其暴露给Python。我们将创建一个新的独立文件rustfib.rs，内容如下：

// rustimport:pyo3

// PyO3是一个用于编写Python扩展的Rust库；
// 我们将导入其最常用的API。
use pyo3::prelude::*;

#[pyfunction]  // ← 将函数暴露给Python
fn fibonacci(number: u64) -> u64 {
    if number == 0 {
        return 0;
    }
    if number == 1 {
        return 1;
    }
    let mut prevprev = 0;
    let mut prev = 1;
    let mut current = 1;
    for _ in 0..(number - 1) {
        current = prevprev + prev;
        prevprev = prev;
        prev = current;
    }
    current
}

在解释这段代码的含义之前，我们先看看它的运行情况。

默认情况下，如果我们尝试在Python提示符中导入Rust文件，它不会工作。Python不知道如何导入以.rs结尾的文件：

In [1]: from rustfib import fibonacci
ModuleNotFoundError: No module named 'rustfib'

这就是rustimport的用武之地。顾名思义，它允许你导入Rust文件；你只需要先import rustimport.import_hook。

In [2]: import rustimport.import_hook

In [3]: from rustfib import fibonacci
    Updating crates.io index
  Downloaded proc-macro2 v1.0.64
  Downloaded 1 crate (44.8 KB) in 0.32s
   Compiling target-lexicon v0.12.8
   Compiling autocfg v1.1.0
...
   Compiling pyo3-macros v0.18.3
   Compiling rustfib v0.1.0 (/tmp/rustimport/rustfib-7eb3578b36e7d1e44917eae13823c3f8/rustfib)
    Finished dev [unoptimized + debuginfo] target(s) in 10.13s
In [4]: fibonacci(50)
Out[4]: 12586269025

我们刚刚导入了一个Rust文件，它被自动编译，然后我们运行了代码！

如果我们查看目录中的文件，可以看到一个新的编译好的Python扩展被创建：

$ ls
rustfib.cpython-311-x86_64-linux-gnu.so  rustfib.rs

此时我们不需要再导入rustimport.import_hook：我们可以直接导入扩展，因为它已经编译好了：

$ python -c "import rustfib; print(rustfib.fibonacci(10))"
55

在讨论性能之前，让我们先解释一下我们编写的代码以及rustimport为我们做了什么。

使用rustimport编写扩展

rustimport在多个方面减少了你的工作量。

首先，rustimport.import_hook会钩入Python的导入系统；当它看到具有相关名称的Rust文件时，它会将其编译为Python扩展并导入。为了防止你随机导入任何Rust文件，你需要在文件顶部添加一个特殊注释：

// rustimport:pyo3

其次，通常使用Rust编写Python扩展的方式是使用一个名为PyO3的第三方库。在正常的Rust设置中，你需要将pyo3添加为依赖项；上面的// rustimport: pyo3注释告诉rustimport自动为你完成这项工作。

第三，一个正常的PyO3扩展需要你有一个模块初始化函数，你会在其中注册函数。在我们的例子中，它可能看起来像这样：

// 如果你使用rustimport，这是可选的：
#[pymodule]
fn rustfib(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(fibonacci, m)?)?;
    Ok(())
}

rustimport会自动为你生成这段代码。然而，如果默认的#[pyfunction]和#[pyclass]支持不够，你可以自己编写。

第四，你不需要编写或生成Cargo.toml文件，大多数Rust包都需要这个文件。

修复Rust初学者最常见的性能错误

那么我们的新扩展有多快，与我们Python实现的954纳秒相比如何？让我们来测量一下：

In [1]: import rustimport.import_hook

In [2]: from rustfib import fibonacci

In [3]: %timeit fibonacci(50)
920 ns ± 0.974 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

Rust不是应该很快吗？

不一定。如果你回到我们最初编译扩展的地方，它说Finished dev [unoptimized + debuginfo] target(s) in 10.13s。unoptimized提示了为什么这段代码运行得很慢。

当你编译Rust代码时，有不同的编译配置文件，它们以不同的方式编译。对我们来说，有两个有趣的配置文件：

• dev：有额外的调试断言和信息，但没有进行速度优化。在某些情况下，编译可能会更快。
• release：没有调试断言或信息，进行了速度优化。

如果你在基准测试你的Rust代码，或者只是希望它以全速运行，你需要使用release配置文件进行编译。 而且很容易忘记这样做！

• 如果你使用cargo，请使用cargo build --release。
• 如果你使用setuptools-rust或maturin，这是你用来为Python构建Rust扩展的正常打包工具，develop模式（pip install -e或maturin develop）将使用dev配置文件安装。你需要正常pip install以获得发布模式。

在rustimport的情况下，我们可以通过几种不同的方式使用发布配置文件进行编译（参见README了解详细信息）。我们将使用在导入Rust模块之前设置选项的方法。首先，我们删除现有的编译扩展。然后：

In [1]: import rustimport.import_hook

In [2]: import rustimport.settings

In [3]: rustimport.settings.compile_release_binaries = True

In [4]: from rustfib import fibonacci
   Compiling target-lexicon v0.12.8
   ...
   Compiling pyo3-macros v0.18.3
   Compiling rustfib v0.1.0 (/tmp/rustimport/rustfib-7eb3578b36e7d1e44917eae13823c3f8/rustfib)
    Finished release [optimized] target(s) in 5.77s

注意，它说Finished release [optimized] target(s)；我们现在有了一个release构建。这意味着我们可以对代码进行基准测试，以获得更真实的Rust性能表现：

In [5]: %timeit fibonacci(50)
44.7 ns ± 0.0438 ns per loop (mean ± std. dev. of 7 runs, 10,000,000 loops each)

比较所有三个版本：

版本	速度（越低越好）
Python	955 ns
Rust (dev profile)	920 ns
Rust (release profile)	45 ns

使用依赖项

如果你的原型想要使用第三方Rust库怎么办？通常，使用Rust程序时，你会编辑Cargo.toml，或者在命令行上运行cargo add thedependency来自动添加它。但使用rustimport时，你没有Cargo.toml。

相反，你可以将这些子句添加到你的可导入Rust文件中，使用特殊语法：

// rustimport:pyo3

// 嵌入的Cargo.toml：
//: [dependencies]
//: fibext = "0.2"

use fibext::Fibonacci;
use pyo3::prelude::*;

#[pyfunction]
fn fibonacci(number: u64) -> u64 {
    Fibonacci::new().nth(number as usize).unwrap()
}

在这个例子中，我们没有自己计算斐波那契数列，而是使用了已经实现它的fibext包。

避免误用旧代码

rustimport会注意到Rust文件和编译扩展是否已经分叉，并在必要时重新编译。然而，如果你不小心，很容易在编辑Rust文件后仍然运行旧代码：

1. 如果你忘记导入rustimport.import_hook，将使用文件系统中现有的编译扩展。
2. 如果你不重启Python，它将坚持使用你导入的第一个版本；它不会更新已经导入的模块。

何时使用rustimport

rustimport的亮点在于原型设计：它使得在现有代码库中尝试一些Rust代码并查看其运行情况变得非常简单。Python打包的两个主要Rust集成工具在现有代码库中进行原型设计时不如rustimport合适：

• setuptools-rust需要更新你的setup.py，并有一个额外的Cargo.toml，并为Rust设置目录结构。这相当快，但仍然需要额外的工作。
• maturin是一个优秀的工具，可以轻松创建新的基于Rust的Python库，但并不真正设计用于向现有项目添加Rust。

在这两种情况下，你都需要在每次更改Rust代码时手动安装或重新构建；工作量不大，但仍然增加了摩擦。

然而，当涉及到生产环境时，我可能会从rustimport切换到其他工具。例如，通过rustimport导入的单个Rust文件缺少普通Rust项目中的标准Cargo.lock文件，该文件包含传递性锁定的依赖项。这意味着你不会每次都获得具有相同依赖项的可重现构建。虽然rustimport确实可以扩展到完整的Rust项目，但在这一点上，它的优势下降，切换到更明确的构建系统可能更好。