PyO3: когда Rust в Python окупается, а когда только мешает

Профилировщик показал, что 80% времени обработки батча уходит на одну функцию: парсинг и нормализацию 2 млн строк. Чистый Python делал это за 14 секунд на батч. Переписали функцию на Rust и подключили через PyO3 — стало 0,7 секунды. Двадцатикратный выигрыш при том, что вокруг остался тот же Python-код. Но соседняя попытка — ускорить функцию, которая тысячу раз в секунду дёргается из горячего цикла и возвращает маленький словарь, — дала замедление на 15%. Разница между этими двумя случаями и есть тема разбора.

Речь про PyO3 0.22.x и сборку через maturin — это актуальная связка на конец 2024 года. API PyO3 заметно менялся между мажорными версиями (например, устройство Bound<'py, T> пришло в 0.21), так что примеры из старых статей могут не собраться. Проверяйте версию в Cargo.toml.

Где Rust реально окупается

PyO3 не ускоряет Python. Он позволяет вынести кусок работы в скомпилированный код, у которого нет GIL внутри и нет накладных расходов интерпретатора на каждой операции. Выигрыш появляется, когда выполняется одно из условий:

• CPU-bound работа над большим объёмом данных за один вызов. Парсинг, числодробилка, обход дерева, сжатие. Вы платите за переход границы Python→Rust один раз, а внутри Rust перемалывает миллионы элементов.
• Освобождение GIL на время работы. Rust-код может отпустить GIL через Python::allow_threads, и тогда несколько Python-потоков реально пойдут параллельно по ядрам — то, чего чистый Python не умеет.
• Нужен код без сборщика мусора и с предсказуемой памятью. Обработка потока, где важна латентность хвоста, а не средняя.

Обратная сторона: каждый вызов из Python в Rust и каждое возвращаемое значение проходят конвертацию между типами. Для list из миллиона int это делается один раз оптом и почти бесплатно на элемент. Для функции, которую зовут в тесном цикле и которая возвращает маленькую структуру, стоимость конвертации превышает выигрыш от Rust.

PyO3 окупается там, где вы пересекаете границу редко, а работаете за ней много. Как только граница пересекается часто, а работа мелкая — вы платите за FFI больше, чем экономите на вычислениях.

Минимальный рабочий пример

Возьмём функцию, которая суммирует квадраты — не потому, что это полезно, а потому что она показывает границу конвертации. Структура проекта после maturin new: Cargo.toml, pyproject.toml, src/lib.rs.

// src/lib.rs — PyO3 0.22.x
use pyo3::prelude::*;

// Функция получает срез f64 напрямую, без поэлементной конвертации в цикле
#[pyfunction]
fn sum_of_squares(values: Vec<f64>) -> f64 {
    values.iter().map(|x| x * x).sum()
}

// CPU-bound работа с освобождением GIL: пока считаем, другие потоки живут
#[pyfunction]
fn heavy_sum(py: Python<'_>, values: Vec<f64>) -> f64 {
    py.allow_threads(|| {
        values.iter().map(|x| x.sqrt() * x).sum()
    })
}

#[pymodule]
fn fastmath(m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_of_squares, m)?)?;
    m.add_function(wrap_pyfunction!(heavy_sum, m)?)?;
    Ok(())
}

Cargo.toml должен объявить тип библиотеки как расширение Python:

[lib]
name = "fastmath"
crate-type = ["cdylib"]

[dependencies]
pyo3 = { version = "0.22", features = ["extension-module"] }

Сборка и установка в текущее окружение — одна команда:

maturin develop --release

Флаг --release здесь не косметика. Без него Rust собирается без оптимизаций, и вы легко получите код медленнее Python. Все замеры делайте только на release-сборке.

После этого в Python:

from fastmath import sum_of_squares, heavy_sum

result = sum_of_squares([1.0, 2.0, 3.0])  # 14.0

Цена перехода границы: почему мелкие вызовы проигрывают

Каждый вызов Rust-функции из Python — это не прыжок, а протокол. Аргументы проверяются на тип, распаковываются, значение результата упаковывается обратно в Python-объект и получает счётчик ссылок. Для Vec<f64> из аргумента PyO3 обходит список целиком и конвертирует каждый элемент — на входе. Это линейно от длины, и на миллионе элементов оно тонет в полезной работе.

А теперь антипример: функция принимает два int, складывает и возвращает int, и вы зовёте её 5 млн раз в цикле. Здесь на каждый вызов приходится фиксированная стоимость FFI-границы, а полезной работы — одно сложение. Такой цикл почти всегда быстрее оставить в Python или, если он критичен, целиком перенести в Rust, а не дёргать Rust из Python-цикла.

Правило, которое экономит недели: переносите в Rust цикл целиком, а не тело цикла. Граница должна быть снаружи горячего цикла, а не внутри него.

Как мерить, а не гадать

Прежде чем переписывать, снимите профиль. Синтетические цифры вроде «Rust в 40 раз быстрее» бесполезны — выигрыш полностью зависит от соотношения полезной работы к стоимости границы на ваших данных. Меряйте на реальном объёме:

1. Профилируйте Python-версию (cProfile или py-spy record), найдите функцию, которая действительно ест время. Оптимизировать функцию с 2% времени бессмысленно.
2. Оцените, крупная ли она по объёму работы за вызов. Обрабатывает массив/файл/дерево за раз — кандидат. Зовётся часто и возвращает мелочь — плохой кандидат.
3. Напишите Rust-версию, соберите --release, замерьте ту же нагрузку через timeit или perf_counter на одинаковых входных данных.
4. Сравните с учётом времени конвертации входа/выхода, а не только внутреннего цикла.

Что вы платите за интеграцию

Ускорение — не единственная строка в счёте. PyO3 добавляет постоянные издержки, которые не видны в бенчмарке, но бьют по команде.

Если ускоряемый участок — 5% времени всего сервиса, эти издержки почти наверняка не окупятся. Если это ядро продукта, через которое проходит весь трафик, — окупятся быстро.

Где подход ломается

• Паника в Rust вместо исключения. Если Rust-код паникует (например, unwrap() по None), PyO3 превратит это в pyo3_runtime.PanicException, но состояние может остаться неопределённым. Возвращайте PyResult и явные ошибки, а не паникуйте на невалидном входе.
• Забыли allow_threads. Долгая CPU-функция без освобождения GIL держит весь интерпретатор — другие потоки стоят. В однопоточном коде это незаметно, под нагрузкой веб-сервера ловится как рост латентности.
• Конвертация больших вложенных структур. Возврат Vec<HashMap<String, Vec<...>>> в Python строит гору Python-объектов, и время конвертации может съесть весь выигрыш. Иногда дешевле вернуть bytes и распарсить на стороне Python или отдать numpy-массив через rust-numpy.
• Debug-сборка в замере. Классическая ошибка: maturin develop без --release, вывод «Rust медленнее», отказ от идеи. Замеряйте только release.

Что делать дальше

Начните не с переписывания, а с профиля: найдите одну функцию, которая ест заметную долю времени и делает много работы за вызов. Соберите для неё Rust-версию через maturin develop --release, замерьте на реальных данных с учётом конвертации границы. Если выигрыш меньше двух-трёх раз — почти наверняка не стоит постоянных издержек на toolchain и колёса.

Для работы с числовыми массивами смотрите rust-numpy: он даёт доступ к буферу numpy без копирования, что убирает главный источник стоимости границы. Для распространения библиотеки настройте сборку колёс под все платформы через cibuildwheel — иначе интеграция, которая ускорила код на вашей машине, сломает установку у пользователей. Асинхронный мост Python↔Rust (pyo3-async-runtimes) — отдельная тема со своими ловушками, здесь мы её не трогаем.