[2023] Tornando o Python 100× mais rápido com PyO3

Recentemente, ao estudar Python com free-threading, passei a me interessar por PyO3, então estou compartilhando este texto, embora ele já tenha 2 anos.

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

Contexto

• A biblioteca Python central do pipeline interno de processamento 3-D passou a gerar gargalo com o aumento de usuários simultâneos.
• Reescrever tudo em Rust seria arriscado e demorado, então a escolha foi pela otimização parcial.

Abordagem

1. Começar medindo: identificação dos gargalos com o profiler por amostragem py-spy.
2. Introdução gradual de Rust
   • Conexão Python ↔ Rust com PyO3 + maturin.
   • Primeiro, migração apenas da função find_close_polygons para Rust.
   • Depois, até a estrutura de dados Polygon foi movida para Rust, com subclassing a partir do Python.
3. Profiling e melhorias iterativas
   • Minimização de conversões desnecessárias de NumPy → Rust.
   • Redução de alocações e cópias, com micro-otimizações por cálculo direto de distância.

Mudanças de desempenho

Etapa	Tempo médio de execução (ms)	Fator de melhoria
Python puro inicial	293.41	1×
v1 – só a função em Rust (--release)	23.44	12.5×
v2 – Polygon também em Rust	6.29	46.5×
v3 – remoção de alocações e cálculo direto	2.90	101×

Tecnologias principais

• PyO3 : FFI segura entre Python ↔ Rust.
• maturin : automação de build e distribuição.
• ndarray / numpy crate : arrays e álgebra linear no lado Rust.
• py-spy : profiler que mostra até a pilha nativa.

Lições

• Fazer profiling primeiro permite obter grandes ganhos com pequenas mudanças de código.
• Mesmo mantendo a API Python, só trocar o módulo em Rust já permite aplicação imediata em serviços reais.
• Rust é bastante eficaz mesmo quando introduzido de forma enxuta apenas na “camada de desempenho”.

[allinux]

Criar extensões Python com C/C++ reduz demais a produtividade, mas com o PyO3 é muito mais conveniente porque já tem maturin e cargo.
Além disso, como para módulos Python a compilação cruzada também é essencial, Rust também facilita bastante a compilação cruzada.

[lamanus]

maturin... sofrimento...

[aer0700]

Eu aguento o máximo possível com vetorização do NumPy; se não dá, encaixo uma GPU e mudo para CuPy ou Torch; e, se ainda não der, escrevo código nativo com Cython... mas parece melhor evitar código nativo sempre que possível. É difícil.