[2023] Cómo hacer Python 100 veces más rápido con PyO3

Últimamente, mientras estudiaba Python con free-threading, me interesó PyO3, así que comparto este artículo aunque ya tiene 2 años.

Making Python 100× Faster with <100 Lines of Rust – resumen

Contexto

• La librería principal de Python de un pipeline interno de procesamiento 3-D se convirtió en un cuello de botella al aumentar los usuarios concurrentes.
• Reescribir todo en Rust implicaba demasiado riesgo y tiempo, así que se eligió una optimización parcial.

Enfoque

1. Primero medir: identificar el cuello de botella con el profiler de muestreo py-spy.
2. Introducción gradual de Rust
   • Conectar Python ↔ Rust con PyO3 + maturin.
   • Primero, portar solo la función find_close_polygons a Rust.
   • Después, mover también la estructura de datos Polygon a Rust y hacer subclassing desde Python.
3. Perfilado y mejora iterativos
   • Minimizar conversiones innecesarias de NumPy → Rust.
   • Reducir asignaciones y copias, y aplicar microoptimizaciones con cálculo directo de distancias.

Cambios de rendimiento

Etapa	Tiempo promedio de ejecución (ms)	Multiplicador de mejora
Python puro inicial	293.41	1×
v1 – solo la función en Rust (--release)	23.44	12.5×
v2 – Polygon también en Rust	6.29	46.5×
v3 – eliminación de asignaciones y cálculo directo	2.90	101×

Tecnologías clave

• PyO3 : FFI segura entre Python ↔ Rust.
• maturin : automatización de build y distribución.
• ndarray / numpy crate : arreglos y álgebra lineal del lado de Rust.
• py-spy : profiler que permite ver incluso el stack nativo.

Lecciones

• Si primero haces profiling, puedes obtener grandes ganancias con cambios pequeños en el código.
• Incluso manteniendo la API de Python, reemplazar solo el módulo en Rust permite aplicarlo de inmediato a un servicio en producción.
• Rust es suficientemente efectivo incluso si se introduce de forma acotada solo en las “zonas de rendimiento”.

[allinux]

Crear extensiones de Python con C/C++ reduce demasiado la productividad, pero PyO3 es muy cómodo porque de entrada cuenta con maturin y cargo.
Además, en los módulos de Python la compilación cruzada también es indispensable, y con Rust incluso eso es sencillo.

[lamanus]

maturin... dolor...

[aer0700]

Aguanta lo más posible con la vectorización de NumPy; si no alcanza, mete una GPU y pásate a CuPy o Torch, y si aun así no basta, escribe código nativo con Cython o algo por el estilo... pero parece que, en lo posible, es mejor evitar lo nativo. Es pesado.