Come velocizzare Python con Rust e PyO3

Python è bellissimo nella sua sintesi e nell'ampia selezione di pacchetti di terze parti, ma quando si tratta di calcoli intensivi, analizzare gigabyte di dati o validare milioni di oggetti, l'interprete inizia a rallentare. In precedenza, in queste situazioni, gli sviluppatori dovevano scrivere manualmente le estensioni dell'API C. Questo processo richiede molto lavoro: il debug delle perdite di memoria e la gestione manuale dei riferimenti eliminano rapidamente tutta la gioia dello sviluppo.

Alcuni anni fa è apparsa una strada più sicura. La libreria PyO3 permette di scrivere moduli nativi in Rust mantenendo rigorose garanzie di sicurezza della memoria e astrazioni chiare.

Se hai usato il validatore Pydantic v2, il parser JSON orjson, il framework Polars o il tokenizer Tiktoken di OpenAI, hai già eseguito codice creato con PyO3.

Perché collegare Rust e Python

Python eccelle nella logica di business e nel prototyping rapido. Rust vince dove servono massima velocità, basso consumo di risorse e concorrenza.

PyO3 risolve due problemi:

1. Spostare i colli di bottiglia delle applicazioni Python in funzioni o strutture Rust separate.
2. Incorporare l'interprete Python in un binario Rust per l'esecuzione di script.

Più spesso, il progetto viene usato per il primo scenario.

Come scrivere un modulo nella pratica

Per compilare codice Rust in un pacchetto Python, gli autori consigliano lo strumento maturin. Gestisce la compilazione e crea i file wheel.

Creiamo un semplice modulo che somma due numeri e restituisce il risultato come stringa.

Prima, impostiamo l'ambiente virtuale:

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

La logica lato Rust nel file 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(())
    }
}

Le macro #[pyfunction] e #[pymodule] generano automaticamente il codice per l'API C di Python. Gli sviluppatori non devono convertire manualmente i tipi PyObject in strutture Rust e viceversa.

Avviamo la build:

maturin develop

Ora la funzione può essere chiamata direttamente da Python:

import string_sum

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

Se vuoi misurare le prestazioni reali in produzione, compila semplicemente il modulo con il flag di ottimizzazione: maturin develop --release.

Casi d'uso reali

In Pydantic V2, tutta la logica di type checking e validazione è stata riscritta in Rust in un pacchetto separato pydantic-core. Le prestazioni sono aumentate drasticamente.

La libreria Polars sta sostituendo Pandas quando si lavora con grandi dataset grazie alle computazioni parallele in Rust. L'utente scrive codice Python familiare, mentre il lavoro pesante va nel modulo compilato.

Nel machine learning, PyO3 viene usato da Hugging Face nel pacchetto tokenizers e da OpenAI in tiktoken. Il parsing veloce del testo prima dell'invio alla rete neurale riduce significativamente la latenza del servizio.

Strumenti di supporto

Si è formato un insieme di utility utili attorno a PyO3:

• maturin elimina la necessità di scrivere configurazioni complesse per setup.py.
• rust-numpy collega gli array NumPy con Rust senza copiare i dati in memoria.
• pyo3-async-runtimes connette asyncio con il runtime asincrono Tokio.
• pythonize trasforma le strutture dati usando Serde.

Conclusione

Se il tuo progetto ha un collo di bottiglia nelle prestazioni, PyO3 ti aiuterà a risolvere il problema senza passare a C o C++. La principale difficoltà risiede nell'apprendere lo stesso Rust. Sarà necessario comprendere i concetti di ownership e i tipi di dati, ma la libreria gestisce completamente il lavoro di routine con CPython.