将 Python 和 Rust 融合在一起,为 pyQuil® 4.0 带来和谐

2023-12-28 16:58:10 浏览数 (1)

前言

pyQuil 一直是在 Rigetti 量子处理单元(QPUs)上构建和运行量子程序的基石,通过我们的 Quantum Cloud Services(QCS™)平台提供服务。它是我们的一个重要客户端库。然而,随着 QCS 平台的发展,我们越来越倾向于使用 Rust,因为它具有出色的性能、类型系统和强调正确性。为了支持Rigetti 不断增长的 Rust 工具和服务生态系统,pyQuil 中的许多功能已被我们的 Rust 库取代。幸运的是,Rust 很适合用作外部函数接口(FFI)。这对我们来说是 Rust 的另一个重要优势,因为它是在我们的服务和高级语言(如 Python)或低级语言(如 C)之间架设桥梁的理想选择。

我们仍然致力于支持 Python 和 pyQuil,因此我们花了过去一年的时间用我们现代的 Rust SDKs 改装了 pyQuil。这对 pyQuil 进行了基础性的更改,以一种透明的方式为用户带来了 Rust 的好处,并为在 Rigetti 的第四代 QPUs 上编译和运行程序提供了所需的增强功能。您可以在我们的 “Introducing pyQuil v4” 指南中了解有关主要更改的详细信息。在本文的其余部分,我们将讨论在 Python 中集成 Rust 时遇到的一些挑战和突破。

设定方向

在继续之前,让我们明确集成我们的 Rust SDKs 与 pyQuil 所需的两个主要目标:

在我们现有的 Rust 库之上构建 Python 软件包,而不损害这些 Rust 库的设计或惯用“Rustiness”。

将这些软件包合并到 pyQuil 中,同时最小化对现有API和行为的破坏性更改。

从 Rust 库构建 Python 软件包

我们知道我们希望我们的 Rust 库保持纯粹的 Rust 库,不包含任何 Python 特定的代码或类型。相反,我们希望确保我们的 Python 软件包符合 Python 开发人员的期望。这些目标是相互冲突的,因此很明显前进的最有效方式是保持我们的 Rust crate 中的核心逻辑,并构建一个具有 Rust 绑定的 Rust 软件包的单独 crate。

我们决定使用 PyO3 crate 作为在 Rust 中构建 Python 软件包的首选框架。它被广泛使用并有很好的文档。pyo3 提供了许多宏,可以用于包装您的 Rust 代码并将其公开为 Python 对象。这些宏注释了类型和函数的定义,但在尝试从外部 crate 中的类型构建 Python 软件包时,它们的实用性受到限制。

典型的解决方法涉及在外部类型周围创建 newtype 包装器,但这会导致繁琐的样板代码。例如,newtype 包装器缺乏使用 pyo3 生成 getter 和 setter 属性的便利性。相反,使用 newtype 包装器需要手动实现。

quil-rs 中的这个例子说明了这个问题。在 Quil 中,一个 EXCHANGE a b 指令交换内存引用 a 和 b 中的值。这在 quil-rs 中使用 MemoryReference 和 Exchange 结构表示:

代码语言:rust复制
pub struct MemoryReference {
    pub name: String,
    pub index: u64
}

pub struct Exchange {
    pub left: MemoryReference,
    pub right: MemoryReference
}

如果我们直接用 PyO3 包装这个结构,我们将使用 pyclass 和 pyo3 属性将 ExchangeMemoryReference 分别包装为 Python 类,完全具有它们的字段的 gettersetter

代码语言:rust复制
use pyo3::pyclass;

#[pyclass(get_all, set_all)]
pub struct MemoryReference {
    pub name: String,
    pub index: u64
}

#[pyclass(get_all, set_all)]
pub struct Exchange {
    pub left: MemoryReference,
    pub right: MemoryReference
}

虽然方便,但这种方法需要将 Python 特定的代码和依赖项注入我们的 Rust库,从而破坏其纯度。但是,我们应该如何处理外部 crate 的代码呢?

首先,我们必须围绕外部类型创建 newtype 包装器,以将 #[pyclass] 属性应用于它们:

代码语言:rust复制
use quil_rs::instruction::{Exchange, MemoryReference};
use pyo3::prelude::*;

#[pyclass(name = "MemoryReference")]
pub struct PyMemoryReference(MemoryReference);

#[pyclass(name = "Exchange")]
pub struct PyExchange(Exchange)

接下来,由于我们不能在新类型包装器上使用 get_all 和 set_all 访问 MemoryReferenceExchange 的内部字段,我们必须为内部类型的每个字段手动实现 getter 和 setter:

代码语言:rust复制
#[pymethods]
impl PyMemoryReference {
    #[getter]
    fn get_name(self) -> String { ... }
    #[setter]
    fn set_name(self, name: String) -> PyResult<()> { ... }
    #[getter]
    fn get_index(self) -> u64 { ... }
    #[setter]
    fn set_index(self, index: u64) -> PyResult<()> { ... }
}

#[pymethods]
impl PyExchange {
    #[getter]
    fn get_left(self) -> MemoryReference { ... }
    #[setter]
    fn set_left(self, memory_reference: PyMemoryReference) -> PyMemoryReference { ... }
    #[getter]
    fn get_right(self) -> MemoryReference { ... }
    #[setter]
    fn set_right(self, memory_reference: PyMemoryReference) -> PyMemoryReference { ... }
}

这种方法牺牲了 PyO3 提供的许多便利性,容易出错,并且显著增加了维护构建在外部 Rust crate 上的 Python 软件包所需的样板代码。对于我们来说,这是一个重大问题,特别是因为 quil-rs 在很大程度上依赖于 Rust 的类型系统来表示 Quil 程序。

如果我们能够同时拥有两个世界的最佳优势呢?这就是 rigetti-pyo3 的目标,这是我们构建的一个开源库,通过引入 traits 和宏,大大减少了构建围绕外部 Rust 类型的 Python 软件包所需的样板代码。使用 rigetti-pyo3,我们可以使用 py_wrap_data_struct! 宏生成 newtype 包装器,包含每个字段的 getter 和 setter。我们所需做的就是指定字段、预期的 Rust 类型以及用于转换的 Python 兼容类型:

代码语言:rust复制
py_wrap_data_struct! {
    PyMemoryReference(MemoryReference) as "MemoryReference" {
        name: String => Py<PyString>,
        index: u64 => Py<PyInt>
    }
}

py_wrap_data_struct! {
    PyExchange(Exchange) as "Exchange" {
        left: MemoryReference => PyMemoryReference,
        right: MemoryReference => PyMemoryReference
    }
}

“rigetti-pyo3”包含一系列宏,使得利用基本类型的 trait 实现变得轻而易举,从而实现 Python 方法。例如,impl_hash! 宏利用包装的 Rust 类型上的 Hash 实现,在包装类型上实现了 Python 的 __hash__ 方法。

这些宏的存在不仅减少了样板代码,而且通过确保每个绑定都以相同的方式实现常见功能,使得 Python API 更加一致。py_wrap_union_enum! 宏就是一个很好的例子,它用简单的 API 包装了一个带标签的联合(或 Rust 枚举的变体),用于构造和与 Rust 枚举交互的 Python 类。

“rigetti-pyo3”已经被证明是在外部 Rust crate 上构建 Python 软件包的宝贵框架。它使我们能够在 Rust 库和相应的 Python 库之间建立无缝的集成,而无需在任一设计中进行妥协。

改装 pyQuil

尽管 pyQuil 和我们的 Rust 库解决了一些共同的问题,但它们的解决方案在许多情况下是非常不同的。它们的方法在许多情况下相似,但也存在很大的灵活性。总的来说,从我们的 Rust 库中添加新功能到 pyQuil 并不是一个挑战,因为我们可以自由选择如何将它们整合。然而,在 pyQuil 具有更多功能的情况下,我们通常不得不将其迁移到我们的 Rust 库中。在这里需要谨慎决策,我们希望回溯任何必要的功能以提供完整而一致的 API,但与此同时,我们不希望过多地将 pyQuil 特定的功能移植回我们的 Rust SDKs。

另一个挑战是如何在不破坏我们的 Rust SDKs API 的情况下满足 pyQuil 现有 API 的期望。其中之一涉及 asyncio 和 pyQuil 不支持 asyncio 的问题。

异步困境

我们的 Rust API 的大部分涉及与外部服务进行网络交互,这些任务自然适合异步 Rust。虽然 pyo3 本身不直接支持异步函数,但出色的 pyo3-asyncio 使将异步 Rust 函数公开为 Python asyncio 函数变得轻而易举。然而,pyQuil 在其自己的 API 中不使用 asyncio,并且使用这些 asyncio 函数的原样本需要在 pyQuil 的许多核心方法上引入 async 关键字。这将要求用户也采用 asyncio,这是我们不愿意做出的重大更改。

起初,我们尝试通过手动调用 asyncio 事件循环 API 以同步函数中运行将异步 Rust 绑定导出到 Python 中。这条路没有走得很远,对这个想法的所有变体都是可疑的。最终,没有一个在同步和异步上下文中都表现良好。

相反,如果我们将所有异步机制推到 Rust 运行时中会怎么样?这也带来了一系列挑战。首先,我们想确保我们适当地处理操作系统信号。用户经常希望通过按 Ctrl-C 来中止运行时间较长的函数,这会向运行中的程序发送 SIGINT 信号。在 Python 程序的情况下,运行中的 Python 解释器需要处理这些信号,这意味着在 Rust 掌控时,信号不会被处理。pyo3 文档记录了这个陷阱,这是我们在试图将潜在的长时间运行的异步函数变为同步函数时需要注意的事项。在所有这一切中,还有一个复杂的问题是 Python API 函数 PyErr_CheckSignals() 必须在主线程上调用,否则调用将是一个空操作。

总的来说,我们需要包装一个异步 Rust 函数,使其在 Python 中呈现为同步函数,同时确保在主线程上处理信号,以便尊重操作系统信号。

让我们来做吧。给定一个虚构的异步 Rust 函数 foo

代码语言:rust复制
async fn foo() -> String {
    tokio::time::sleep(Duration::from_secs(3));
    "hello".to_string()
}

使用 pyo3_asyncio,我们可以将其导出为一个 asyncio 函数:

代码语言:rust复制
#[pyfunction]
fn py_foo_async(py: Python<'_>) -> PyResult<&PyAny> {
    pyo3_asyncio::tokio::future_into_py(py, async { Ok(foo().await) })
}

但是,我们如何将其包装成同步 API 呢?首先,我们获取当前的运行时,然后将我们的异步函数作为任务在该运行时上启动。然后,我们可以使用 tokio::select! 来管理从我们的任务返回的结果,或从信号处理程序返回的结果,以先返回的为准。将所有这些都包装在当前运行时中,然后,大功告成!我们有一个在幕后使用 Rust 的异步运行时的同步 Python 函数:

代码语言:rust复制
#[pyfunction]
fn py_foo_sync() -> PyResult<String> {
    let runtime = pyo3_asyncio::tokio::get_runtime();
    let handle = runtime.spawn(foo());

    runtime.block_on(async {
        tokio::select! {
            result = handle => result.map_err(|err| pyo3::exceptions::PyRuntimeError::new_err(err.to_string())),
            signal_err = async {
                let delay = std::time::Duration::from_millis(100);
                loop {
                    Python::with_gil(|py| {
                        py.check_signals()
                    })?;
                    tokio::time::sleep(delay).await;
                }
            } => signal_err
        }
    })
}

这很好,但对于每个异步函数都做这么多事情太多了。为了每个异步函数在我们的 API 中都重复这个设置,我们可以使用一个宏。

代码语言:rust复制
macro_rules! py_sync {
    ($py: ident, $body: expr) => {{
        $py.allow_threads(|| {
            let runtime = ::pyo3_asyncio::tokio::get_runtime();
            let handle = runtime.spawn($body);

            runtime.block_on(async {
                tokio::select! {
                    result = handle => result.map_err(|err| ::pyo3::exceptions::PyRuntimeError::new_err(err.to_string()))?,
                    signal_err = async {
                        let delay = ::std::time::Duration::from_millis(100);
                        loop {
                            ::pyo3::Python::with_gil(|py| {py.check_signals()})?;
                            ::tokio::time::sleep(delay).await;
                        }
                    } => signal_err,
                }
            })
        })
    }};
}

我们宏的一个补充是我们如何将所有东西都包装在 py.allow_threads 中。这释放了全局解释器锁(GIL),以便在进行纯 Rust 工作时其他 Python 线程可以运行。我们只有在需要使用 Python::with_gil 检查 OS 信号时才重新获取 GIL。

现在,对于任何异步函数,我们只需写:

代码语言:rust复制
#[pyfunction]
fn py_foo(py: Python<'_>) -> PyResult<String> {
    py_sync!(py, async { Ok(foo().await) })
}

这也很好,但我们可以走得更远。这些同步函数对于兼容性来说是很好的,但一些用户可能会喜欢一个真正的 asyncio API。这就是为什么我们建立了另一个建立在上一个基础上的宏,用于提供单个 async 函数的同步和异步变体。这让我们在其自然的 async 形式中编写函数一次,并免费获得同步和异步变体。

代码语言:rust复制
// 这会生成两个Python函数:
//  def foo(): ...
//  async def foo(): ...
py_sync::py_function_sync_async! {
    #[pyfunction]
    async fn foo() -> PyResult<String> {
        Ok(foo().await)
    }
}

能够继续支持同步 API,同时不错过提供异步 API 的机会,对我们来说是一个巨大的胜利,也是将 Rust 与 Python 结合在一起能够带来的不易通过单独使用 Python 实现的好处的一个很好的例子。

回报:功能和性能

我们已经确定了在以不妥协任一库的质量或用户体验为代价的方式下,将现有的 Python 和 Rust 库之间的差距缩小的挑战。那么这给我们带来了什么?

如前所述,我们的 Rust 库已经开始在功能上超越 pyQuil。最重要的是,它们带来了在 Rigetti 的下一代 Ankaa 系统上编译和运行程序所需的增强功能。

此外,通过将解析和序列化 Quil 程序的逻辑、以编程方式构建它们以及执行和检索作业结果的逻辑集中到我们的 Rust 库中,我们已经为 pyQuil 现在和将来构建了一个坚实的基础。在我们的服务和客户端库中使用相同的逻辑,使我们更容易维护和扩展 pyQuil,同时为用户提供更一致的体验。

最后,我们不能结束一篇关于 Python 和 Rust 的博客文章,而不提到性能。通过将核心逻辑移植到 Rust,我们在许多方面看到了显著的性能提升,比如解析和序列化 Quil 程序。这是至关重要的,因为解析和序列化是 pyQuil 中常见的编译和执行工作流程中的关键步骤。

方法论:所有基准测试都使用 Python 3.8 在装有 M1 Max 的 2021 年 MacBook Pro 上执行。测试加载了一个大型的 Quil 程序文件,并对逐渐增大的程序块进行解析的基准测试。数据使用 pytest-benchmark 进行收集。

结论

将 Python 和 Rust 组合到 pyQuil v4 中提出了许多挑战。从构建在我们现有的 Rust 库之上而不妥协其设计的初步决策,到在不引入破坏性变更的情况下满足长时间 pyQuil 用户的期望,我们走过了一条复杂的道路。通过这些努力,我们现代化了 pyQuil,为用户提供了 Rust 的性能和类型安全性的好处,同时保持了 Python 的熟悉性和易用性。

这不仅仅是将两种语言结合在一起的技术问题。它还涉及到在两者之间找到平衡,以提供一致的用户体验,并为库的未来扩展奠定基础。通过解决这些问题,我们为 pyQuil 带来了一种令人满意的融合,展示了 Python 和 Rust 之间合作的潜力,以解决量子计算领域的挑战。

0 人点赞