使用Rust让python加速50倍

发表于 2023/09/28 更新于 2023/09/30

作者 jusbin 12 分钟阅读

python语法便利，可以让人以接近自然表达的方式编写代码，因此有着极高的开发效率。但有些时候，python代码的性能会很低下，虽然很多情况下python都会调用底层的C模块加快运算，但总有无法覆盖的情况，如果恰好程序对性能要求很高，那么这些未能优化的地方将成为运行的瓶颈。将运行缓慢的部分用底层语言（如C、C++、Rust）重写，python层调用重写后的方法，这就是很好的解决问题的方法。

问题场景

需要对视频流进行加密和解密操作：

视频流中包含一个又一个NALU，NALU使用0x000001作为起始标志。
封装NALU时为了避免内部出现0x000001，会对数据进行处理，如果出现了00000000 00000000 000000xx，那么中间添加0x03，变成00000000 00000000 00000011 000000xx。简而言之，需要对原始数据进行编码后才能放入到视频流中。
加密操作流程：获取NALU -> 加密NALU -> 编码加密后的NALU
解密操作流程：获取NALU -> 解码加密后的NALU -> 解密NALU

以加密过程为例，获取NALU的速度很快；加密使用cryptography提供的AES加密模块，速度同样很快；最后进行编码的速度却很慢。

encrypt cost 1.012859
encrypt time 0.10626899999999999
encode  time 0.8735130000000012

decrypt cost 0.888086
decrypt time 0.037574000000000024
decode  time 0.818685000000001

可以看到，AES加密仅仅耗费了0.1s，而编码加密后的数据却耗费了0.87秒。编码过程其实很简单，仅仅是遍历一遍数据，在某些情况下添加额外的0x03字节。

  
def nalu_encode(data):
  data_raw = data
  data_enc = bytearray()
  
  i = 0
  i_max = len(data_raw)
  while i < i_max:
    if (i + 2 < i_max) and (data_raw[i] == 0 and data_raw[i + 1] == 0 and data_raw[i + 2] < 4):
        data_enc.append(data_raw[i])
        data_enc.append(data_raw[i + 1])
        data_enc.append(3)
        data_enc.append(data_raw[i + 2])
        i += 2
    else:
        data_enc.append(data_raw[i])
    i += 1
  return data_enc

加速改造

知道了问题所在，就开始进行优化。可以使用C/C++实现这一函数，然后编译成静态库供python调用。Rust同样可以实现这一过程，同时Rust生态链里已经有脚手架，可以更方便地进行python/rust之间的相互调用。

使用pyo3将函数功能放到rust实现

pyo3就是一个这样的进行语言间绑定的工具，并且提供maturin作为管理工具，能一键构建、发布。

安装`maturin`

pip install maturin

初始化项目

由于我是在现有项目中初始化，先进入到项目根目录下

maturin init

初始化后会生成几个配置文件和代码目录

rust配置文件
- Cargo.toml
- pyproject.toml
GitHub actions配置文件，利用GitHub actions完成自动化的包发布
- .github/workflows/CI.yml
rust源代码
- src/lib.rs
  修改配置文件

修改Cargo.toml中[lib]的name字段为python调用的包名，maturin会将rust代码安装为当前python环境下的一个包。

  
[package]
name = "h26x-extractor"
version = "0.1.0"
edition = "2021"

# name修改为python调用的包名 构建后python就能 import rust_utils
[lib]
name = "rust_utils"
crate-type = ["cdylib"]

[dependencies]
pyo3 = "0.19.0"

同时修改pyproject.toml中[project]中name，保持二者一致。

  
[build-system]
requires = ["maturin>=1.2,<2.0"]
build-backend = "maturin"

# 和Cargo.toml中名称保持一致
[project]
name = "rust_utils"
requires-python = ">=3.7"
classifiers = [
  "Programming Language :: Rust",
  "Programming Language :: Python :: Implementation :: CPython",
  "Programming Language :: Python :: Implementation :: PyPy",
]


[tool.maturin]
features = ["pyo3/extension-module"]

编写rust功能函数

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

#[pyfunction]
fn nalu_decode(data: &[u8]) -> Vec<u8> {
  let mut res = Vec::with_capacity(data.len());
  let mut i = 0;
  let i_max = data.len();
  while i < i_max {
      if (i + 2 < i_max) && (data[i] == 0) && (data[i + 1] == 0) && (data[i + 2] == 3) {
          res.push(0);
          res.push(0);
          i += 3;
      } else {
          res.push(data[i]);
          i += 1;
      }
  }
  res
}

/// A Python module implemented in Rust.
#[pymodule]
fn rust_utils(_py: Python, m: &PyModule) -> PyResult<()> {
  m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
  m.add_function(wrap_pyfunction!(nalu_decode, m)?)?;
  Ok(())
}

lib.rs中已经提供了一个示例，sum_as_string方法就是一个暴露给python的方法，这里照葫芦画瓢：

#[pymodule]是python中module的入口，在这里添加一个nalu_decode方法
在nalu_decode方法上添加#[pyfunction]，表明这是和python相关的函数
参数类型可以直接使用rust语言的类型，pyo3会自动进行转换

在python版本的函数中，函数参数是一个bytes，输出也是一个bytes，从参数类型表中可看出对应的rust类型为Vec<u8>, &[u8], Cow<[u8]>。因此rust版本的nalu_decode中输入类型为&[u8]，输出类型为Vec<u8>（三种类型都可以作为输入或者输出类型）。代码的内部逻辑和rust完全一致，甚至代码都是copilot自动补全的。

构建rust代码

maturin develop

运行构建指令后，构建成功

📦 Including license file "/Users/jusbin/Code/h26x-extractor/LICENSE"
🔗 Found pyo3 bindings
🐍 Found CPython 3.9 at /Users/jusbin/Code/h26x-extractor/venv/bin/python
📡 Using build options features from pyproject.toml
💻 Using `MACOSX_DEPLOYMENT_TARGET=11.0` for aarch64-apple-darwin by default
 Compiling h26x-extractor v0.1.0 (/Users/jusbin/Code/h26x-extractor)
  Finished dev [unoptimized + debuginfo] target(s) in 0.36s
📦 Built wheel for CPython 3.9 to /var/folders/yq/h7g177zx3fn_yvwy6zyk8blm0000gn/T/.tmpp3LddZ/rust_utils-0.1.0-cp39-cp39-macosx_11_0_arm64.whl
🛠 Installed rust_utils-0.1.0

此时rust_utils应该已经成功安装至当前的虚拟环境，测试一下

  
>>> import rust_utils
>>> rust_utils.sum_as_string(1, 2)
'3'

说明构建成功，#[pymodule]中添加的rust函数能成功地在python侧运行。

这时通过rust中实现的nalu_encode以及nalu_decode方法进行编码/解码，看一下效果

encrypt cost 0.467197
encrypt time 0.103564
encode  time 0.2867820000000001

decrypt cost 0.392586
decrypt time 0.038244000000000035
decode  time 0.32320400000000016

编码时间由之前的0.87s缩短为0.28s，性能提升了200%。

使用release模式编译rust代码

虽然提升的幅度很大，但编码时间仍远比AES加密时间长，肯定是不合理的。

但rust没有发挥出全部的能力，看一下release模式下的性能。加上--release参数进行构建

maturin develop --release

性能得到飞跃，从0.28s缩短为0.038s

encrypt cost 0.18675
encrypt time 0.0748939999999999
encode  time 0.03856800000000004

decrypt cost 0.13551
decrypt time 0.03571200000000001
decode  time 0.07222100000000002

使用零开销返回类型

只是简单地使用rust进行重写，并开启release模式，就能使效率提升20倍。

进一步的，在将值返回python中有提到，返回rust中的python本地类型(&PyAny, &PyDict…)是零开销（zero cost）的，而前面的代码里为了方便返回的是rust语言中的类型Vec<u8>。

  
use pyo3::prelude::*;
use pyo3::types::PyBytes;

#[pyfunction]
fn nalu_decode(py: Python, data: &PyBytes) -> Py<PyAny> {
  let data = data.as_bytes();
  let mut res = Vec::with_capacity(data.len());
  let mut i = 0;
  let i_max = data.len();
  while i < i_max {
      if (i + 2 < i_max) && (data[i] == 0) && (data[i + 1] == 0) && (data[i + 2] == 3) {
          res.push(0);
          res.push(0);
          i += 3;
      } else {
          res.push(data[i]);
          i += 1;
      }
  }
  PyBytes::new(py, &res).into()
}  

参数中添加py: Python用于创建python中arrays对应的零开销类型PyBytes

返回类型不能是&PyBytes，会出现生命周期问题，需要显示声明生命周期。可以通过into()接管引用并返回Py<PyAny>类型，而Py<T>都是零开销的。

encrypt cost 0.117545
encrypt time 0.07031599999999986
encode  time 0.016480000000000022

decrypt cost 0.082091
decrypt time 0.03545599999999999
decode  time 0.01656300000000001

可以看到，效率又提升了一倍。

总结

	原本时长	使用rust重写	重写+relase	重写+release+零开销返回
encode方法	0.87s	0.28s	0.038s	0.017s
decode方法	0.81s	0.32s	0.072s	0.016s
效率(对比python)	1	3.1	22.8	51.1

将python中性能较差的方法使用rust进行重写，并简单地进行返回类型的优化后就能获得50倍的性能提升。

编程语言总是要在人与机器之间进行取舍，要么接近人的思维模式，如python，那么开发效率高，运行效率低；要么接近机器指令，如C，那么开发效率低，运行效率高。但语言总是为程序员服务的，在不同的应用场景下使用最为适宜的语言，在一个项目使用多种语言进行开发，都是为了能更好地完成项目，使得最终的成品在开发效率和运行效率之间获得平衡。

文章, 计算机

本文由作者按照 CC BY 4.0 进行授权