Cómo acelerar Python con Rust y PyO3

Python es hermoso en su concisión y su amplia selección de paquetes de terceros, pero cuando se trata de cálculos pesados, analizar gigabytes de datos o validar millones de objetos, el intérprete comienza a ralentizarse. Anteriormente, en estas situaciones, los desarrolladores tenían que escribir extensiones de la API de C manualmente. Este proceso era laborioso: depurar fugas de memoria y la gestión manual de referencias eliminaba rápidamente toda la alegría del desarrollo.

Hace varios años surgió un camino más seguro. La biblioteca PyO3 permite escribir módulos nativos en Rust manteniendo estrictas garantías de seguridad de memoria y abstracciones claras.

Si has usado el validador de Pydantic v2, el parser JSON orjson, el framework Polars o el tokenizador Tiktoken de OpenAI, ya has ejecutado código creado con PyO3.

Por qué vincular Rust y Python

Python destaca en lógica de negocio y prototipado rápido. Rust gana donde se necesita máxima velocidad, bajo consumo de recursos y concurrencia.

PyO3 resuelve dos problemas:

1. Transferir los cuellos de botella de aplicaciones Python a funciones o estructuras separadas en Rust.
2. Integrar el intérprete de Python en un binario de Rust para la ejecución de scripts.

La mayoría de las veces, el proyecto se utiliza para el primer escenario.

Cómo escribir un módulo en la práctica

Para compilar código Rust en un paquete de Python, los autores recomiendan la herramienta maturin. Se encarga de la compilación y crea los archivos wheel.

Creemos un módulo simple que sume dos números y devuelva el resultado como una cadena.

Primero, configuremos el entorno virtual:

mkdir string_sum && cd string_sum
python -m venv .env
source .env/bin/activate
pip install maturin
maturin init --bindings pyo3

La lógica del lado de Rust en el archivo src/lib.rs:

use pyo3::prelude::*;

#[pyfunction]
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
    Ok((a + b).to_string())
}

#[pymodule]
mod string_sum {
    use super::*;

    #[pymodule_export]
    fn init(m: &Bound<'_, PyModule>) -> PyResult<()> {
        m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
        Ok(())
    }
}

Las macros #[pyfunction] y #[pymodule] generan automáticamente código para la API de C de Python. Los desarrolladores no necesitan convertir manualmente tipos de PyObject a estructuras de Rust y viceversa.

Iniciemos la compilación:

maturin develop

Ahora la función se puede llamar directamente desde Python:

import string_sum

result = string_sum.sum_as_string(5, 20)
print(result)  # '25'

Si quieres medir el rendimiento real en producción, simplemente compila el módulo con el flag de optimización: maturin develop --release.

Casos de uso en el mundo real

En Pydantic V2, toda la lógica de verificación de tipos y validación se reescribió en Rust en un paquete separado pydantic-core. El rendimiento aumentó drásticamente.

La biblioteca Polars está desplazando a Pandas cuando se trabaja con grandes conjuntos de datos gracias a los cálculos paralelos en Rust. El usuario escribe código Python familiar, mientras que el trabajo pesado se realiza en el módulo compilado.

En aprendizaje automático, PyO3 es utilizado por Hugging Face en el paquete tokenizers y por OpenAI en tiktoken. El análisis rápido de texto antes de enviarlo a la red neuronal reduce significativamente la latencia del servicio.

Herramientas de apoyo

Se ha formado un conjunto de utilidades útiles alrededor de PyO3:

• maturin elimina la necesidad de escribir configuraciones complejas para setup.py.
• rust-numpy vincula arrays de NumPy con Rust sin copiar datos en memoria.
• pyo3-async-runtimes conecta asyncio con el runtime asíncrono Tokio.
• pythonize transforma estructuras de datos usando Serde.

Conclusión

Si tu proyecto tiene un cuello de botella de rendimiento, PyO3 ayudará a resolver el problema sin cambiar a C o C++. La principal dificultad aquí radica en aprender Rust en sí. Se requerirá comprender los conceptos de propiedad y los tipos de datos, pero la biblioteca maneja completamente el trabajo rutinario con CPython.