Как я проектирую и разрабатываю расширения Python на Rust

Вы наверняка видели множество статей на тему "Python, Rust - производительность, бла-бла-бла... Вот, реализуем foo2plus2". Вся беда в том, что все эти статьи демонстрируют очень простые примеры уровня "hello-world". Напротив, в этой статье я хочу рассказать о том, как я проектирую комплексные расширения и почему я принимаю те или иные проектные решения.

Статья является переводом и адаптацией моей же статьи на Medium. Возможно, кто-то предпочтет прочесть в оригинале - ссылка позволяет прочитать без подписки.

На данный момент я написал четыре библиотеки для Python на Rust (1, 2, 3, 4) и приобрел определенный опыт, но все еще не чувствую, что достиг той квалификации, которая позволяет утверждать, что правильно, а что нет. Некоторые из моих подходов вдохновлены другими людьми, другие являются результатом анализа и долгих попыток рефакторинга кода, и все же, я не уверен, что мои решения являются лучшими из возможных.

Кроме того, я не очень опытный разработчик на Rust, поэтому в моем коде могут встречаться забавные вещи - будьте готовы к фейспалмам, если вы сильно прошарены в Rust.

Благодарности

Если порыться в моих репозиториях, то можно обнаружить эволюцию подхода к разработке расширений Rust для Python: например, в Similari многие решения странные, но именно в ходе данного проекта у меня произошел существенный скачок вперед, когда я получил PR от Андрея Ткаченко.

Другая команда продемонстрировала, как использовать Sphinx для документирования расширений, и я написал об этом в соответствующей статье. Я провел несколько часов, погружаясь в среду Manylinux и особенности ее использование в Maturin GitHub Action, в результате чего появилась статья, демонстрирующая, как создавать переносимые расширения Rust для Python с помощью GitHub CI и сборке в Docker.

Хочу сказать "спасибо" всем тем, кто помог мне продвинуться в данной теме: все вы проделали огромную работу, создавая свой код без общепринятых подходов и обширной документации.

Благодарю следующие коллективы и людей за участие в моем прогрессе:

• PyO3 team — за их усилия по предоставлению разработчикам превосходного инструмента для создания расширений Python на Rust;

• Andrei Tkachenko — за демонстрацию возможностей использования transmute при разработке расширений Python на Rust и за огромный вклад в развитие Similari;

• Pola.rs project team — идеи по дизайну и фрагменты кода;

• Pometry/Raphtory team — за идеи по документированию расширений Python на Rust.

Предварительные условия

Это не вводная статья. Я предполагаю, что вы знаете Rust и можете написать расширение типа "hello world" для Python на Rust с помощью PyO3. Я предполагаю, что вы знаете, что такое GIL и как он работает.

Специфика модели управления памятью

Rust и Python подходят по разному к управлению памятью, что иногда превращает проектирование сложных структур данных и способов их взаимодействия между собой в нетривиальную задачу: в Rust модель памяти основана на владении, а в Python - на ссылках на объекты в куче.

Наиболее близкой структурой, представляющей модель памяти Python, в Rust является:

Arc<Mutex<T>>
// or
Arc<RwLock<T>>

Однако, в Python вышеуказанный мьютекс представлен GIL - глобальная блокировка, общая для всех объектов в рантайме, в то время как в Rust возможно осуществлять изолированную блокировку только нужных для работы объектов.

Одна из ключевых проблем при проектировании объектной модели совместимой с Python заключается в том, что передать ссылку (&mut или &) из Rust в Python невозможно: для этого необходимо использовать Arc<T>.

В PyO3 реализованы некоторые решения, позволяющие передавать аллоцированные в Python объекты в Rust-функции по ссылке, но обратная операция невозможна: вы не можете просто сохранить ссылку на Rust-объект в памяти Python. PyO3 поддерживает передачу объектов Rust в Python по значению, оборачивая их в объекты Python (PyObject).

Понимание моделей памяти и их особенностей их взаимодействия между собой важно для правильного проектирования кода, особенно если он освобождает GIL, поскольку освобождение GIL требует от объекта реализации трейта Ungil, который, по сути, является Send.

Если вы хотите одновременно использовать объект в Python и Rust, необходимо использовать обертку для него на базе Arc<T> или Arc<Mutex<T>>.

Пример объекта, используемого только в Python:

use pyo3::prelude::*;

#[pyclass]
struct Object {
    x: i32,
}

#[pymethods]
impl Object {
    fn inc(&mut self) {
        self.x += 1;
    }
}

Ссылка на экземпляр данного класса может быть временно передана в функцию Rust, вызываемую из Python, например:

#[pyfunction]
fn callme(o: &mut Object) {
    o.inc();
}

Или вы можете передать объект в Rust как копию (Clone) - объект в пространстве памяти Python не пропадет:

#[pyfunction]
fn callme(mut o: Object) {
    o.inc();
}

Пример объекта, используемого одновременно из Rust и Python:

use parking_lot::Mutex;
use pyo3::prelude::*;
use std::sync::Arc;

struct Object {
    x: i32,
}

impl Object {
    fn inc(&mut self) {
        self.x += 1;
    }
}

#[pyclass]
#[pyo3(name = "Object")]
#[derive(Clone)]
struct ObjectWrapper(Arc<Mutex<Object>>);

#[pymethods]
impl ObjectWrapper {
    fn inc(&self) {
        let mut bound = self.0.lock();
        bound.inc();
    }
}

Теперь вы можете иметь ссылки на один и тот же объект как в Python, так и в Rust, например:

thread_local! {
    static OBJECT: ObjectWrapper = ObjectWrapper(Arc::new(Mutex::new(Object { x: 0 })))
}

#[pyfunction]
fn get_shared_object() -> ObjectWrapper {
    OBJECT.with(|o| o.clone())
}

#[pyfunction]
fn callme(o: &mut ObjectWrapper) {
    o.inc();
}

Его можно одновременно использовать из Python и Rust (в двух потоках, например). Mutex<T> защищает его от гонки, а Arc<T> обеспечивает механизм совместного использования. Вместо мьютекса можно использовать другие примитивы синхронизации, например RwLock<T>.

Нельзя безопасно написать функцию, возвращающую в Python ссылку на объект, выделенный Rust. Но можно вернуть Python ссылку на объект, выделенный в памяти Python:

fn callme<'p>(slf: PyRef<'p, Self>, _py: Python<'p>) -> PyResult<PyRef<'p, Self>> {
    let o = slf.deref();
    ...
    Ok(slf)
}

Сфера применения вышеуказанного кода ограничена, но можно, например, использовать для реализации билдеров или для реализации with.

Структура проекта

При работе с PyO3, структура проекта сильно влияет на вашу производительность и удобство работы с кодом . Приступая к работе над расширением, вооружившись руководством по PyO3, вы обычно начинаете с проекта, сочетающего в себе как Rust-код, так и Python-аннотации, что приводит к тесной связи с Python - вы не сможете компилировать код без Python и тестировать его без инициализации среды выполнения Python.

В принципе, у вас есть два подхода к структурированию кода:

• с использованием пространств имен и feature-флагов, размещая код, связанный с Python, за флагами (этот подход реализован в Similari);

• с использованием отдельных крейтов, объединенных в одно рабочее пространство (Savant-RS, RocksQ).

Изначально я придерживался первого подхода, но сейчас считаю, что второй подход гораздо лучше. Вот как выглядит проект RocksQ в рамках подхода с разными крейтами:

Итак, как вы видите, у меня два крейта: в queue_rs я реализовал логику и функциональность, он ничего не знает о Python. Во втором - только код, связанный с Python. Я считаю, что такая компоновка является оптимальной по ряду причин. Давайте обсудим их.

Продуктивность при работе над кодом с логикой на Rust. Я хочу, чтобы мой код быстро компилировался и тестировался без использования среды исполнения Python и ее типов данных. Такая компоновка проекта позволяет разрабатывать код с логикой с использованием только нативных типов.

Зависимость от PyO3 увеличивает время компиляции и часто делает невозможной реализацию тестов без инициализации среды Python, что неудобно, и может приводить к инетересным эффектам, поскольку GIL сериализует код, ваши тесты могут приобретать неожиданное изменчивое от запуска к запуску поведение.

Например, при использовании PyBytes необходимо иметь рабочую среду Python для создания объекта данного типа. Я же не хочу иметь дело с этими лишними операциями при работе с логикой.

Следуя вышеизложенному, я пишу свой код с учетом особенностей модели памяти, но без типов, связанных с Python, и зависимости от PyO3.

Мои бенчмарки, размещенные в крейте с логикой, работают только с кодом на Rust, поэтому я уверен в том, что они проверяют скорость работы оптимизированного кода, а не транзакционные издержки, связанные с передачей данных между Python и Rust. В свою очередь, в крейте связи с Python вы можете размещать тесты, направленные на тестирование e2e поведения.

Однажды с помощью такого подхода я выяснил, что передача в Rust больших массивов Vec<u8> убивает производительность, и начал использовать PyBytes в своем коде, связанном с PyO3.

Продуктивность при работе над кодом, связанным с Python. Когда я работаю над кодом, связанным с Python, мне необходимо реализовывать только преобразование между типами Python и типами Rust, но не логику. Время компиляции уменьшается, поскольку крейте с логикой изменений не проиходит.

Особенно время компиляции беспокоит меня при работе над документацией - я очень часто меняю аннотации и перекомпилирую код для того, чтобы посмотреть как аннотации будут выглядеть после обработки Sphix. Мои соображения по созданию документации для расширений Rust читайте в соответствующей статье. Предложенная схема помогает мне значительно сократить время компиляции.

Структурирование кодовой базы. Я всегда знаю, где найти код - если он связан с Python, то в проекте *_py, в противном случае - в проекте *_rs. Это помогает уменьшить хаос, который возникает в случае, когда в одном файле встречается много кода, не относящегося к логике, но обеспечивающего взаимодействие с Python.

Именование объектов кода. Rust не поддерживает переопределения имен функций; обычно это означает, что если вам нужна одна и та же функция для внутренних целей (только для Rust) и для экспорта в Python, то одну из них требуется назвать как-то по-хитрому.

Раньше мне приходилось делать что-то в духе:

fn callme(x: i64) -> i64 {
...
}

#[pyfunction]
#[pyo3(name = "callme")]
fn callme_py(x: i64) -> i64 {
    callme(x)
}

Это работает, но порождает беспорядочный код. В предлагаемом подходе это выглядит следующим образом:

use lib_rs as rust;

#[pyfunction]
fn callme(x: i64) -> i64 {
    rust::callme(x)
}

Рано или поздно кто-то захочет получить сборку без Python. Так было с Similari, люди хотели получить код без зависимости от Python. Это стало проблемой. В результате Andrei Tkachenko предложил большой PR, поместив функциональность, связанную с Python, за feature-флаг.

Обертки Python для объектов

Теперь, зная мой подход, вы, наверное, начали думать, что у меня дублируется больше кода, чем требуется. Да, это так, вместо аннотирования одного объекта с #[pyclass] у меня две реализации:

pub struct Object {
    x: i32,
}

impl Object {
    pub fn inc(&mut self) {
        self.x += 1;
    }
}

use pyo3::prelude::*;
use lib_rs as rust;

#[pyclass]
struct Object(rust::Object);

#[pymethods]
impl Object {
    fn inc(&mut self) {
        self.0.inc()
    }
}

Тем не менее, я нахожу подход очень удобным, поскольку код в крейте *_py является абсолютно шаблонным и легко генерируется GitHub Copilot практически без доработок.

Обертка объекта в Python представляет собой кортежную структуру с единственным аргументом

Обычно я не использую традиционные структуры Rust в своих крейтах *_py : только кортежные структуры. Ранее я писал что-то вроде:

struct Object {
    inner: rust::Object
}

Однако теперь я вижу недостатки такой компоновки:

• Я обнаружил, что иногда я использую "inner", иногда "object", "internal" и т.д. Это вносит хаос.

• Код выглядит менее унифицированным, при кортежной компоновке обращение к внутреннему объекту всегда происходит через .0. .

• Вы можете легко использовать std::mem::transmute