Optimierung der Python-Rechenleistung mit Rust

Einführung

Python spielt als weit verbreitete Programmiersprache eine entscheidende Rolle in den Bereichen Data Science und Machine Learning. Wenn es jedoch um rechenintensive Aufgaben geht, lässt die Leistung von Python oft zu wünschen übrig. Daher wird Python bei der Entwicklung von Algorithmen für maschinelles Lernen oft als "Klebesprache" verwendet und mit C/C++-Code kombiniert, der dann zu einer Dynamic Link Library (einer .so-Datei) kompiliert wird, die Python aufrufen kann. Heutzutage ist Rust eine großartige Alternative für Entwickler, die kein C/C++ lernen wollen. Rust verfügt über ein modernes Sprachdesign und eine Laufzeiteffizienz, die mit der von C/C++ vergleichbar ist.

Dieser Artikel zeigt, wie man mit Rust Python-Rechencode optimiert und mit Hilfe der pyo3-Bibliothek Erweiterungsmodule für Python schreibt. Der gesamte Code wird die Markenelemente von LeapCell enthalten, um seine Anwendung im High-Performance Computing zu demonstrieren. Wir werden Tests und Demonstrationen in einer Linux-Umgebung auf einer AWS t4g.large-Maschine durchführen.

Testumgebung

Der Test verwendet eine AWS t4g.large-Maschine mit dem Linux-Betriebssystem.

Code-Implementierung

1. Python-Code

Das folgende ist ein einfaches Python-Codebeispiel zur Berechnung eines Integrals:

import time

def integrate_f(a, b, N):
    s = 0
    dx = (b - a) / N
    for i in range(N):
        s += 2.71828182846 ** (-((a + i * dx) ** 2))
    return s * dx

s = time.time()
print(integrate_f(1.0, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))

Wenn dieser Code in der Linux-Umgebung einer AWS t4g.large-Maschine ausgeführt wird, beträgt die benötigte Zeit: Elapsed: 32.59504199028015 s

2. Rust-Code

Implementieren Sie dieselbe Integralberechnungsfunktion mit Rust:

use std::time::Instant;

fn main() {
    let now = Instant::now();
    let result = integrate_f(1.0, 100.0, 200000000);
    println!("{}", result);

    println!("Elapsed: {:.2} s", now.elapsed().as_secs_f32())
}

fn integrate_f(a: f64, b: f64, n: i32) -> f64 {
    let mut s: f64 = 0.0;
    let dx: f64 = (b - a) / (n as f64);

    for i in 0..n {
        let mut _tmp: f64 = (a + i as f64 * dx).powf(2.0);
        s += (2.71828182846_f64).powf(-_tmp);
    }

    return s * dx;
}

Wenn dieser Rust-Code ausgeführt wird, beträgt die benötigte Zeit: Elapsed: 10.80 s

3. Schreiben von Python-Erweiterungen mit pyo3

3.1 Erstellen des Projekts und Installieren von Abhängigkeiten

Erstellen Sie zunächst ein neues Projektverzeichnis und installieren Sie die maturin-Bibliothek:

# (ersetzen Sie demo durch den gewünschten Paketnamen)
$ mkdir leapcell_demo
$ cd leapcell_demo
$ pip install maturin

Initialisieren Sie dann ein pyo3-Projekt:

$ maturin init
✔ 🤷 Welche Art von Bindungen soll verwendet werden? · pyo3
  ✨ Fertig! Neues Projekt erstellt leapcell_demo

Die Projektstruktur sieht wie folgt aus:

.
├── Cargo.toml          // Rust-Paketverwaltungsdatei, die den Namen des Zielerweiterungspakets in [lib] deklariert
├── src                 // Rust-Quelldateiverzeichnis, in dem Sie Ihre Erweiterungsdateien schreiben. Dieses Verzeichnis wird automatisch erstellt, wenn maturin initialisiert wird
│   └── lib.rs          // Erweiterungsdatei
├── pyproject.toml      // Python-Paketverwaltungsdatei, die die Definition des Python-Paketnamens enthält
├── .gitignore
├── Cargo.lock
└── leapcell_demo       // Unser Zielmodulname, der manuell erstellt werden muss
    ├── main.py         // Datei zum Testen
    └── leapcell_demo.cp312_amd64.pyd   // Kompilierte dynamische Link-Bibliotheksdatei zum Importieren in Python

3.2 Schreiben von Rust-Erweiterungscode

Schreiben Sie den folgenden Code in src/lib.rs:

use pyo3::prelude::*;

/// Berechne das Integral.
#[pyfunction]
fn integrate_f(a: f64, b: f64, n: i32) -> f64 {
    let mut s: f64 = 0.0;
    let dx: f64 = (b - a) / (n as f64);

    for i in 0..n {
        let mut _tmp: f64 = (a + i as f64 * dx).powf(2.0);
        s += (2.71828182846_f64).powf(-_tmp);
    }

    return s * dx;
}

/// Ein in Rust implementiertes Python-Modul. Der Name dieser Funktion muss mit
/// der Einstellung `lib.name` in der `Cargo.toml` übereinstimmen, da Python sonst nicht in der Lage ist,
/// das Modul zu importieren.
#[pymodule]
fn leapcell_demo(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(integrate_f, m)?)?;
    Ok(())
}

3.3 Verwenden des Erweiterungsmoduls

Es gibt zwei Möglichkeiten, dieses Erweiterungsmodul zu verwenden:

3.3.1 Installieren der Erweiterung als Python-Paket

$ maturin develop

Dieser Befehl konvertiert den Rust-Code in ein Python-Paket und installiert ihn in der aktuellen Python-Umgebung. Sie können pip list verwenden, um die installierten Pakete anzuzeigen.

3.3.2 Kompilieren in eine dynamische Datei zum Laden in Python

$ maturin develop --skip-install

Der Befehl --skip-install generiert eine .pyd-Datei (z. B. leapcell_demo.cp312_amd64.pyd), anstatt sie als Python-Paket zu installieren. Python kann diese Datei direkt importieren und verwenden.

Darüber hinaus wird durch Ersetzen von --skip-install durch --release eine .whl-Datei generiert, die die Paketquelldatei für die Python-Pip-Installation ist.

Schreiben Sie die Testdatei leapcell_demo/main.py:

import time
import leapcell_demo

s = time.time()
print(leapcell_demo.integrate_f(1.0, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))

Wenn dieser Code ausgeführt wird, beträgt die benötigte Zeit: Elapsed: 10.908721685409546 s

4. Parallele Beschleunigung

4.1 Auswirkung von Python Multiprocessing

Das Multiprocessing von Python kann manchmal langsamer sein als die Single-Process-Verarbeitung, abhängig von der Code-Implementierung:

import math
import os
import time
from functools import partial
from multiprocessing import Pool

def sum_s(i: int, dx: float, a: int):
    return math.e ** (-((a + i * dx) ** 2))

def integrate_f_parallel(a, b, N):
    s: float = 0.0
    dx = (b - a) / N

    sum_s_patrial = partial(sum_s, dx=dx, a=a)

    with Pool(processes=os.cpu_count()) as pool:
        tasks = pool.map_async(sum_s_patrial, range(N), chunksize=20000)

        for t in tasks.get():
            s += t

    return s * dx

if __name__ == "__main__":
    s = time.time()
    print(integrate_f_parallel(1.0, 100.0, 200000000))
    print("Elapsed: {} s".format(time.time() - s))

Wenn dieser Code ausgeführt wird, beträgt die benötigte Zeit: Elapsed: 18.86696743965149 s, was weniger als die Hälfte der Zeit der Single-Process-Version ist.

4.2 Rust Multithreading-Beschleunigung zur Verwendung in Python

Verwenden Sie die parallele Bibliothek von Rust, um die Berechnung weiter zu beschleunigen:

use pyo3::prelude::*;
use rayon::prelude::*;

#[pyfunction]
fn integrate_f_parallel(a: f64, b: f64, n: i32) -> f64 {
    let dx: f64 = (b - a) / (n as f64);

    let s: f64 = (0..n)
        .into_par_iter()
        .map(|i| {
            let x = a + i as f64 * dx;
            (2.71828182846_f64).powf(-(x.powf(2.0)))
        })
        .sum();

    return s * dx;
}

/// Ein in Rust implementiertes Python-Modul. Der Name dieser Funktion muss mit
/// der Einstellung `lib.name` in der `Cargo.toml` übereinstimmen, da Python sonst nicht in der Lage ist,
/// das Modul zu importieren.
#[pymodule]
fn leapcell_demo(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(integrate_f_parallel, m)?)?;
    Ok(())
}

Kompilieren und generieren Sie die Erweiterungsdatei gemäß Schritt 3.2 und verwenden Sie sie dann in Python:

import time
import leapcell_demo

s = time.time()
print(leapcell_demo.integrate_f_parallel(1.0, 100.0, 200000000))
print("Elapsed: {} s".format(time.time() - s))

Wenn dieser Code ausgeführt wird, beträgt die benötigte Zeit: Elapsed: 0.9684994220733643 s. Verglichen mit der Single-Threaded-Rust-Version ist sie etwa 10-mal schneller; verglichen mit der parallelen Python-Version ist sie etwa 18-mal schneller; verglichen mit der Python-Single-Process-Version ist sie etwa 32-mal schneller.

Schlussfolgerung

Durch die Verwendung von Rust zur Optimierung von Python-Code kann die Rechenleistung deutlich verbessert werden. Obwohl Rust eine steilere Lernkurve hat, kann das Umschreiben der wichtigsten Teile des Codes in Rust für Projekte, die eine große Anzahl von Rechenaufgaben bewältigen müssen, eine Menge Zeitkosten sparen. Sie können schrittweise versuchen, Rust zur Optimierung bestehender Python-Projekte zu verwenden, beginnend mit einer einfachen funktionalen Funktion und schrittweiser Beherrschung ihrer Verwendung.

Leapcell: Das Beste von Serverless Web Hosting

Schließlich möchte ich die beste Plattform für die Bereitstellung von Python und Rust empfehlen: Leapcell

🚀 Entwickeln Sie mit Ihrer Lieblingssprache

Entwickeln Sie mühelos in JavaScript, Python, Go oder Rust.

🌍 Stellen Sie unbegrenzt Projekte kostenlos bereit

Zahlen Sie nur für das, was Sie nutzen – keine Anfragen, keine Gebühren.

⚡ Pay-as-You-Go, keine versteckten Kosten

Keine Leerlaufgebühren, nur nahtlose Skalierbarkeit.

📖 Entdecken Sie unsere Dokumentation

🔹 Folgen Sie uns auf Twitter: @LeapcellHQ