原文链接: http://saidvandeklundert.net/learn/2021-11-06-calling-rust-from-python/
翻译:ChenYe
选题:minikiller
从python调用rust-使用rust加速你的python代码
Python 是一门很棒的语言!至少,在我看来是这样。它有着丰富的第三方库,可以非常好的完成很多工作。但作为解释型语言,Python 并不是最快的。如果要加速代码运行速度,你可以使用 C/C++ 。这种编程语言上的扩展使得你可以在 Python 中调用这些其他语言写的函数。这种方式提供了更快的执行速度,同时你也不会被全局解释器锁困扰。
除了使用 C 或 C++ 之外,我们还可以使用 Rust 。Rust 提供了 FFI(Foreign Function Interface 外部函数接口),它允许你将 Rust 函数导出到 C 语言。通过让 Python 使用这些导出的 C 函数,你可以将 Python 和 Rust 缝合在一起。如此一来就可以在需要的地方使用外部的 Rust 程序。
在本文中,我将介绍一些关于如何从 Python 调用多个 Rust 函数的基本示例。在 Rust 一侧,我将使用 std 中的 ffi
,在 Python 一侧,我将仍然使用 ctypes
:
在 Python 中调用 Rust 函数打印一个字符串.
首先,我们将编写一个打印字符串的 Rust 函数。下图说明具体发生了什么:
在 Python 侧我们做了以下几件事:
- Rust 导出为 C 编译生成的库文件被导入
- 将 string 转换为 bytes
- 调用从 Rust 中导入的函数
- Python 侧的 Rust 函数的参数是 UTF-8 编码的 bytes
在 Rust 侧我们做了以下几件事:
- 创建一个新的库
- 通过 Cargo.toml 指明我们要构建一个 Rust 动态库
- 用 Rust 编写与 C 兼容的外部函数接口
- 读取经由 C 传递来的 Python 侧的输入,类型是
Char *
- 编译 Rust 库
完成之后,我们就可以运行代码了!
Python 侧.
import ctypes
rust = ctypes.CDLL("target/release/librust_lib.so")
if __name__ == "__main__":
SOME_BYTES = "Python says hi inside Rust!".encode("utf-8")
rust.print_string(SOME_BYTES)
首先,我们导入 ctypes ,然后,我们指定 .so 文件的问题。在我这里,命名为 print_string.py 的 Python 脚本被放在构建 rust 库的目录下。在我运行完 cargo build --release
后,会产生以下目录结构:
rust_lib/
├── target/
| ├── lib.rs
├── target/
│ ├── release/
│ ├── librust_lib.so
├── print_string.py
在 Python 代码的 __main__
部分,我先通过调用 .encode("utf-8")
将字符串转换为 bytes 。这一步将得到 UTF-8 编码的 bytes 。然后,我通过 rust.print_string(SOME_BYTES)
调用 Rust 函数,rust
是对我们之前加载的库的引用。print_string
是导出的 Rust 函数名。SOME_BYTES
作为参数传递给目标函数。
我们还可以以另一种方式为 Rust 函数提供参数。试试将以下内容添加到 print_string_script.py :
rust.print_string(
ctypes.c_char_p("Another way of sending strings to Rust via C.".encode("utf-8"))
)
在这种情况下,我们使用 ctypes.c_char_p
将值传递给 rust 函数。 c_char_p
是一个指向字符串的指针。
Rust 侧.
在 Rust 侧,我们先使用 cargo new --lib
创建一个库。然后我们编辑 Cargo.toml 文件:
[package]
name = "rust_from_python"
version = "0.1.0"
authors = ["Said van de Klundert"]
[lib]
name = "rust_lib"
crate-type = ["dylib"]
这表明我们正在构建一个 Rust 动态库,并且我们将使用 libc 。
然后,我们编写 lib.rs 文件:
use std::ffi::CStr;
use std::os::raw::c_char;
use std::str;
/// Turn a C-string into a string slice and print to console:
#[no_mangle]
pub extern "C" fn print_string(c_string_ptr: *const c_char) {
let bytes = unsafe { CStr::from_ptr(c_string_ptr).to_bytes() };
let str_slice = str::from_utf8(bytes).unwrap();
println!("{}", str_slice);
}
第一行将 CStr 引入,这是一个表示 C 字符串的借用的结构体,它的文档是这样说的:
This type represents a borrowed reference to a nul-terminated array of bytes. It can be constructed safely from a &[u8] slice, or unsafely from a raw *const c_char. It can then be converted to a Rust &str by performing UTF-8 validation, or into an owned CString.
我们将用它,将 CStr 转换成 Rust 中的 &str 。
第二行 use std::os::raw::c_char;
,将 c_char 类型引入。
第三行 use std::os::raw::c_char;
,使得我们能访问 from_utf8 ,使用这个方法将 bytes 转换为 Rust 的 &str 。
现在,我们看看这个函数:
#[no_mangle]
pub extern "C" fn print_string(c_string_ptr: *const c_char) {
let bytes = unsafe { CStr::from_ptr(c_string_ptr).to_bytes() };
let str_slice = str::from_utf8(bytes).unwrap();
println!("{}", str_slice);
}
extern
关键字用于创建 FFI(Foreign Function Interface)。它可用于调用其他语言的函数或创建允许其他语言调用 Rust 函数接口。
以下引用来自 book of Rust:
We also need to add a #[no_mangle] annotation to tell the Rust compiler not to mangle the name of this function. Mangling is when a compiler changes the name we’ve given a function to a different name that contains more information for other parts of the compilation process to consume but is less human readable. Every programming language compiler mangles names slightly differently, so for a Rust function to be nameable by other languages, we must disable the Rust compiler’s name mangling.
函数参数类型指定为 c_string_ptr: *const c_char
。其中,*const
表示原始指针,c_char
是 C 语言 char 类型。送一组合起来,就是一个指向 c_char
的原始指针。
为了解引用这个原始指针,我们使用 unsafe
代码块。在这个代码块中,我们使用 CStr::from_ptr
,它会使用 CStr 包装器包装由参数传进来的指针。之后,我们调用 to_bytes()
将 C 字符串转换为 byte 切片。这些字符都将存储在 bytes
变量中。
将字符切片传给 str::from_utf8()
并解包返回值,就会得到一个 &str
类型的值,最后将其打印出来。
当上述一切都写好之后,使用命令 cargo build --release
就可以得到以下目录结构:
rust_lib/
├── target/
| ├── lib.rs
├── target/
│ ├── release/
│ ├── librust_lib.so
├── print_string.py
从 rust_lib 目录,我们运行 Python 脚本:
root@rust:/python/rust/rust_lib# ./print_string.py
Python says hi inside Rust!
此时,我们就已经从 Python 调用了一个 Rust 函数并将一个字符串打印到控制台。
在 Python 中调用一个 Rust 函数打印一个整数.
当从 Python 向 Rust 传递参数时,我们需要考虑 Python、C 和 Rust 中使用的类型。现在我们试试在屏幕上打印一个整数。首先,我们创建 Python 脚本:
import ctypes
rust_lib = ctypes.CDLL("target/release/librust_lib.so")
if __name__ == "__main__":
SOME_BYTES = (2).to_bytes(4, byteorder="little")
rust_lib.print_int(SOME_BYTES)
接下来,我们在 Rust 端并将以下内容添加到我们的 lib.rs 中:
use std::os::raw::c_int;
// Turn a C-int into a &i32 and print to console:
#[no_mangle]
pub extern "C" fn print_int(c_int_ptr: *const c_int) {
let int_ptr = unsafe { c_int_ptr.as_ref().unwrap() };
println!("Python gave us number {}", int_ptr);
}
我们将另一种数据类型 c_int
引入 。该函数将指向 C 整型的指针作为参数。在我们执行命令 cargo build --release
之后,我们可以运行 Python 脚本:
root@rust:/python/rust/rust_lib# python3 print_number.py
Python gave us number 2
Rust、Python 和 C 中有许多类型。这些类型之间的对应可能会把人搞的很头疼!
从 Python 调用多类型的 Rust 函数
现在,我们使用 Python 调用一个 start_procedure
。为了专注研究跨语言调用,该函数仅仅获取一个结构并返回另一个结构。在 Python 侧,我们使用 Pydantic basemodel 来创建 Rust 函数所需的输入。Pydantic basemodel 将具有与 Rust 结构相同的字段。我们对 Rust 的返回值做同样的事情。我们创建了一个 Pydantic basemodel ,它是 Rust struct 在 Python 侧的镜像。 Rust struct 和 Pydantic basemodel 将包含多种不同类型的字段。这是我们将以最简单的方式(至少在我看来)处理的这件事:使用 C 语言中的 Char *
。
上图说明了该过程会发生什么。在 Rust 和 Python 之间,JSON 字符串用于传递值。在 Rust 侧,我们将 JSON 解析到相应的结构中。在 Python 侧,我们将 JSON 加载到相应的 Pydantic basemodel 中。这样做的优点是我们只需要在 C 中使用 Char *
,而不必使用 C 结构体或 C 中的其他类型。
Python 侧.
我们将使用的 Python 脚本如下:
import ctypes
from pydantic import BaseModel
from typing import List
rust = ctypes.CDLL("target/release/librust_lib2.so")
class ProcedureInput(BaseModel):
timeout: int
retries: int
host_list: List[str]
action: str
job_id: int
class ProcedureOutput(BaseModel):
job_id: int
result: str
message: str
failed_hosts: List[str]
if __name__ == "__main__":
procedure_input = ProcedureInput(
timeout=10,
retries=3,
action="reboot",
host_list=["server1", "server2"],
job_id=1,
)
ptr = rust.start_procedure(procedure_input.json().encode("utf-8"))
returned_bytes = ctypes.c_char_p(ptr).value
procedure_output = ProcedureOutput.parse_raw(returned_bytes)
print(procedure_output.json(indent=2))
我们先加载了目标库,然后定义了两个类。这脸各类是 Pydantic basemodels ,它在运行时强制执行类型提示。ProcedureInput 类是 Rust 函数的参数,ProcedureOutput 类是我们希望从 Rust 函数中返回的东西。
完成定义后,我们实例化一个 ProcedureInput 对象。我们使用 rust.start_procedure
调用 Rust 函数。当我们进行 Rust 函数调用时,procedure_input.json().encode("utf-8")
这个表达式将会以 JSON 字符串的形式输出类实例的字段,并将字符串转成 bytes 。
我们会收到从 Rust 中返回的 ptr 。接下来,将其转换为 bytes,并且使用 parse_raw 方法将使我们能够从这些 bytes 构建 ProcedureOutput 实例。
当我们运行 Python 代码时,将得到以下输出:
root@rust:/python/rust/rust_lib# python3 call_rust_function.py
{
"job_id": 1,
"result": "success",
"message": "1 host failed",
"failed_hosts": [
"server1"
]
}
Rust 侧.
以下是 Rust 侧的代码:
extern crate serde;
extern crate serde_json;
use serde::{Deserialize, Serialize};
use std::ffi::CStr;
use std::ffi::CString;
use std::os::raw::c_char;
use std::str;
#[derive(Debug, Serialize, Deserialize)]
struct ProcedureInput {
timeout: u8,
retries: u8,
host_list: Vec<String>,
action: String,
job_id: i32,
}
#[derive(Serialize, Deserialize)]
struct ProcedureOutput {
result: String,
message: String,
failed_hosts: Vec<String>,
job_id: i32,
}
#[no_mangle]
pub extern "C" fn start_procedure(c_string_ptr: *const c_char) -> *mut c_char {
let bytes = unsafe { CStr::from_ptr(c_string_ptr).to_bytes() };
let string = str::from_utf8(bytes).unwrap();
let model: ProcedureInput = serde_json::from_str(string).unwrap();
let result = long_running_task(model);
let result_json = serde_json::to_string(&result).unwrap();
let c_string = CString::new(result_json).unwrap();
c_string.into_raw()
}
fn long_running_task(model: ProcedureInput) -> ProcedureOutput {
let result = ProcedureOutput {
result: "success".to_string(),
message: "1 host failed".to_string(),
failed_hosts: vec!["server1".to_string()],
job_id: model.job_id,
};
return result;
}
在这段代码中,我们先定义了两个结构体作为 Pydantic basemodels 在 Rust 侧的镜像:
#[derive(Debug, Serialize, Deserialize)]
struct ProcedureInput {
timeout: u8,
retries: u8,
host_list: Vec<String>,
action: String,
job_id: i32,
}
#[derive(Serialize, Deserialize)]
struct ProcedureOutput {
result: String,
message: String,
failed_hosts: Vec<String>,
job_id: i32,
}
然后,定义了 start_procedure 函数:
#[no_mangle]
pub extern "C" fn start_procedure(c_string_ptr: *const c_char) -> *mut c_char {
let bytes = unsafe { CStr::from_ptr(c_string_ptr).to_bytes() };
let string = str::from_utf8(bytes).unwrap();
let model: ProcedureInput = serde_json::from_str(string).unwrap();
let result = long_running_task(model);
let result_json = serde_json::to_string(&result).unwrap();
let c_string = CString::new(result_json).unwrap();
c_string.into_raw()
}
与之前相似,我们先从原始指针构建一个字符串。然后我们得到了从 Python 传过来的 JSON 字符串,我们将其解析为 ProcedureInput 结构体。我们将该结构传递给 long_running_task
方法,它会为我们返回一个 ProcedureOutput 对象。使用 serde_json::to_string
可以将结构体转为 JSON 字符串。我们使用 CString::new
从字节容器创建新的兼容 C 的字符串。使用 into_raw()
方法将 c_string
的所有权转移给 C 中的调用者。into_raw()
方法返回一个指针,我们之前在 Python 侧用它来读取返回值。
从 Python 调用带有内存泄漏的 Rust 函数
如果我们把 start_procedure
放在一个 while 循环里,让它运行一段时间,我们可以看到进程会逐渐开始消耗越来越多的内存。是因为从 Rust 返回的值没有被清理。
通过对 Python 脚本进行以下更改,我们将持续调用 start_procedure
:
while True:
ptr = rust.start_procedure(procedure_input.json().encode("utf-8"))
returned_bytes = ctypes.c_char_p(ptr).value
procedure_output = ProcedureOutput.parse_raw(returned_bytes)
print(procedure_output.json(indent=2))
调用脚本,我们可以看到 Python 脚本的内存使用量在慢慢增加。
修复内存泄露问题
为了解决这个内存泄漏问题,我们首先需要在 Rust 中创建一个清理内存的函数:
#[no_mangle]
pub extern "C" fn free_mem(c_string_ptr: *mut c_char) {
unsafe { CString::from_raw(c_string_ptr) };
}
此函数将使用 from_raw
获取已转移到 C 的 CString
的所有权。当函数结束时,变量将不被任何所有者拥有,值被删除,释放内存。
我们需要在 Python 侧调用这个函数。 free_mem
函数的输入应该是 Rust 返回给 Python 的值。
以下 Python 可以连续运行而不会泄漏内存:
while True:
ptr = rust.start_procedure(procedure_input.json().encode("utf-8"))
returned_bytes = ctypes.c_char_p(ptr).value
procedure_output = ProcedureOutput.parse_raw(returned_bytes)
print(procedure_output)
rust.free_mem(ptr)
start_procedure
返回的值就是我们传递给 free_mem 的值。另一件需要注意的事情是我们在处理完值后调用 free_mem
。
在下面的示例中,我们在 Python 代码中使用它之前释放该值:
ptr = rust.start_procedure(procedure_input.json().encode("utf-8"))
returned_bytes = ctypes.c_char_p(ptr).value
rust.free_mem(ptr)
procedure_output = ProcedureOutput.parse_raw(returned_bytes)
在让 Rust 释放内存之后,我们尝试在 Python 侧读取相同的字节。当我们运行这段代码时,我们完成了一次 double free :
root@rust:/python/rust/rust_lib# python3 call_rust_continuously_free_mem.py
job_id=1 result='success' message='1 host failed' failed_hosts=['server1']
free(): double free detected in tcache 2
Aborted
总结
从 Python 中使用 Rust 是我这段时间一直想尝试的事情。即使我已经研究 Rust 一段时间了,完全转向 Rust 对我现在参与的任何项目都没有任何意义。在某些情况下,我使用的库在 Rust 中不可用,而在其他情况下,我正在工作的项目太大而无法在 Rust 中重写。另外,其实我对于 Rust 还处于学习阶段。
像这样从 Python 调用 Rust 为我铺平了道路:
- 将 Rust 合并到现有的 Python 项目中并从小处着手
- 对现有的 Python 项目更有信心,知道如果速度真的成为问题,我可以用 Rust 来加速
- 同时使用 Python 和 Rust!
Rust 社区的人们在使用 Rust 来加速 Python 方面做出了很多努力。在很多情况下,你会看到这些项目使用 pyo3。这个库提供了 “Python 的 Rust 绑定,包括用于创建原生 Python 扩展库的工具” 。使用 pyo3 的一个例子是 polars 。
这篇文章中的例子可以在这里找到。
注意:写这篇文章时,我使用了 CPython 3.9.2、Pydantic 版本 1.8.2 和 Rust 1.56.0。