Optimizing Python Computational Performance with Rust

Introduction

Python, as a widely used programming language, plays a crucial role in the fields of data science and machine learning. However, when dealing with computationally intensive tasks, Python's performance often leaves much to be desired. Therefore, in the development of machine learning algorithms, Python is often used as a "glue language" and combined with C/C++ code, which is then compiled into a dynamic link library (a .so file) for Python to call. Nowadays, for developers who don't want to learn C/C++, Rust is a great alternative. Rust has a modern language design and runtime efficiency comparable to that of C/C++.

This article will introduce how to use Rust to optimize Python computational code and write extension modules for Python with the help of the pyo3 library. All code will incorporate the brand elements of LeapCell to demonstrate its application in high-performance computing. We will conduct tests and demonstrations in a Linux environment on an AWS t4g.large machine.

Test Environment

The test uses an AWS t4g.large machine running the Linux operating system.

Code Implementation

1. Python Code

The following is a simple Python code example for calculating an integral:

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))

When this code is executed in the Linux environment of an AWS t4g.large machine, the time taken is: Elapsed: 32.59504199028015 s

2. Rust Code

Implement the same integral calculation function using 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;
}

When this Rust code is executed, the time taken is: Elapsed: 10.80 s

3. Writing Python Extensions with pyo3

3.1 Creating the Project and Installing Dependencies

First, create a new project directory and install the maturin library:

# (replace demo with the desired package name)
$ mkdir leapcell_demo
$ cd leapcell_demo
$ pip install maturin

Then, initialize a pyo3 project:

$ maturin init
✔ 🤷 What kind of bindings to use? · pyo3
  ✨ Done! New project created leapcell_demo

The project structure is as follows:

.
├── Cargo.toml          // rust package management file, which will declare the name of the target extension package in [lib]
├── src                 // rust source file directory, where you write your extension files. This directory is automatically created when maturin initializes
│   └── lib.rs          // Extension file
├── pyproject.toml      // Python package management file, which contains the definition of the Python package name
├── .gitignore
├── Cargo.lock
└── leapcell_demo       // Our target module name, which needs to be created manually
    ├── main.py         // File for testing
    └── leapcell_demo.cp312_amd64.pyd   // Compiled dynamic link library file for import into Python

3.2 Writing Rust Extension Code

Write the following code in src/lib.rs:

use pyo3::prelude::*;

/// Caculate the integrate.
#[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;
}

/// A Python module implemented in Rust. The name of this function must match
/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
/// import the module.
#[pymodule]
fn leapcell_demo(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(integrate_f, m)?)?;
    Ok(())
}

3.3 Using the Extension Module

There are two ways to use this extension module:

3.3.1 Installing the Extension as a Python Package

$ maturin develop

This command will convert the Rust code into a Python package and install it into the current Python environment. You can use pip list to view the installed packages.

3.3.2 Compiling into a Dynamic File for Loading in Python

$ maturin develop --skip-install

The --skip-install command will generate a .pyd file (for example, leapcell_demo.cp312_amd64.pyd), instead of installing it as a Python package. Python can directly import and use this file.

In addition, replacing --skip-install with --release will generate a .whl file, which is the package source file for Python pip installation.

Write the test file 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))

When this code is executed, the time taken is: Elapsed: 10.908721685409546 s

4. Parallel Acceleration

4.1 Effect of Python Multiprocessing

Python's multiprocessing may sometimes be slower than single-process processing, depending on the code implementation:

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))

When this code is executed, the time taken is: Elapsed: 18.86696743965149 s, which is less than half the time of the single-process version.

4.2 Rust Multithreading Acceleration for Use in Python

Use Rust's parallel library to further accelerate the calculation:

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;
}

/// A Python module implemented in Rust. The name of this function must match
/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
/// import the module.
#[pymodule]
fn leapcell_demo(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(integrate_f_parallel, m)?)?;
    Ok(())
}

Compile and generate the extension file according to step 3.2, and then use it 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))

When this code is executed, the time taken is: Elapsed: 0.9684994220733643 s. Compared with the single-threaded Rust version, it is about 10 times faster; compared with the Python parallel version, it is about 18 times faster; compared with the Python single-process version, it is about 32 times faster.

Conclusion

By using Rust to optimize Python code, computational performance can be significantly improved. Although Rust has a steeper learning curve, for projects that need to handle a large number of computational tasks, rewriting the key parts of the code in Rust can save a lot of time costs. You can gradually try to use Rust to optimize existing Python projects, starting with a simple functional function and gradually mastering its usage.

Leapcell: The Best of Serverless Web Hosting

Finally, I would like to recommend the best platform for deploying Python and Rust: Leapcell

🚀 Build with Your Favorite Language

Develop effortlessly in JavaScript, Python, Go, or Rust.

🌍 Deploy Unlimited Projects for Free

Only pay for what you use—no requests, no charges.

⚡ Pay-as-You-Go, No Hidden Costs

No idle fees, just seamless scalability.

📖 Explore Our Documentation

🔹 Follow us on Twitter: @LeapcellHQ