Comment accélérer Python avec Rust et PyO3

Python est magnifique par sa concision et sa riche sélection de paquets tiers, mais lorsqu'il s'agit de calculs intensifs, de l'analyse de gigaoctets de données, ou de la validation de millions d'objets, l'interpréteur commence à ralentir. Auparavant, dans de telles situations, les développeurs devaient écrire manuellement des extensions de l'API C. Ce processus est laborieux : le débogage des fuites mémoire et la gestion manuelle des références éliminent rapidement tout le plaisir du développement.

Une voie plus sûre est apparue il y a plusieurs années. La bibliothèque PyO3 permet d'écrire des modules natifs en Rust tout en maintenant des garanties strictes de sécurité mémoire et des abstractions claires.

Si vous avez utilisé le validateur Pydantic v2, l'analyseur JSON orjson, le framework Polars, ou le tokeniseur Tiktoken d'OpenAI, vous avez déjà exécuté du code créé avec PyO3.

Pourquoi lier Rust et Python

Python excelle dans la logique métier et le prototypage rapide. Rust l'emporte là où une vitesse maximale, une faible consommation de ressources et la concurrence sont nécessaires.

PyO3 résout deux problèmes :

1. Délester les goulots d'étranglement des applications Python vers des fonctions ou structures Rust distinctes.
2. Intégrer l'interpréteur Python dans un binaire Rust pour l'exécution de scripts.

Le plus souvent, le projet est utilisé pour le premier scénario.

Comment écrire un module en pratique

Pour compiler du code Rust en paquet Python, les auteurs recommandent l'outil maturin. Il gère la compilation et crée les fichiers wheel.

Créons un module simple qui additionne deux nombres et renvoie le résultat sous forme de chaîne.

D'abord, configurons l'environnement virtuel :

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

La logique côté Rust dans le fichier 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(())
    }
}

Les macros #[pyfunction] et #[pymodule] génèrent automatiquement du code pour l'API C de Python. Les développeurs n'ont pas besoin de convertir manuellement les types PyObject en structures Rust et vice versa.

Démarrons la compilation :

maturin develop

Maintenant, la fonction peut être appelée directement depuis Python :

import string_sum

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

Si vous souhaitez mesurer les performances réelles en production, compilez simplement le module avec l'indicateur d'optimisation : maturin develop --release.

Cas d'utilisation réels

Dans Pydantic V2, toute la logique de vérification de type et de validation a été réécrite en Rust dans un paquet séparé pydantic-core. Les performances ont considérablement augmenté.

La bibliothèque Polars remplace Pandas lors du travail avec de grands ensembles de données grâce aux calculs parallèles en Rust. L'utilisateur écrit du code Python familier, tandis que le travail lourd se fait dans le module compilé.

En apprentissage automatique, PyO3 est utilisé par Hugging Face dans le paquet tokenizers et par OpenAI dans tiktoken. L'analyse rapide de texte avant l'envoi au réseau neuronal réduit considérablement la latence du service.

Outils complémentaires

Un ensemble d'utilitaires utiles s'est formé autour de PyO3 :

• maturin élimine le besoin d'écrire des configurations complexes pour setup.py.
• rust-numpy lie les tableaux NumPy avec Rust sans copier les données en mémoire.
• pyo3-async-runtimes connecte asyncio avec le runtime async Tokio.
• pythonize transforme les structures de données en utilisant Serde.

Conclusion

Si votre projet présente un goulot d'étranglement en termes de performances, PyO3 vous aidera à résoudre le problème sans passer à C ou C++. La principale difficulté réside dans l'apprentissage de Rust lui-même. La compréhension des concepts de possession et des types de données sera nécessaire, mais la bibliothèque gère entièrement le travail de routine avec CPython.