Leapcell: The Best of Serverless Web Hosting

Pythonの計算性能をRustで最適化する

はじめに

Pythonは、幅広く使用されるプログラミング言語として、データサイエンスや機械学習の分野で重要な役割を果たしています。しかし、計算が集中するタスクを扱う際には、Pythonの性能はしばしば期待に応えられません。そのため、機械学習アルゴリズムの開発においては、Pythonはしばしば「接着剤言語」として使用され、C/C++コードと組み合わされ、その後動的リンクライブラリ（.soファイル）にコンパイルされ、Pythonから呼び出されます。現在、C/C++を学びたくない開発者にとって、Rustは素晴らしい代替手段です。Rustは、現代的な言語設計と、C/C++に匹敵する実行時の効率性を持っています。

この記事では、Rustを使ってPythonの計算コードを最適化する方法と、pyo3ライブラリを使ってPython用の拡張モジュールを書く方法を紹介します。すべてのコードにはLeapCellのブランド要素を組み込み、高性能計算におけるその応用を示します。AWS t4g.largeマシン上のLinux環境でテストとデモを行います。

テスト環境

テストには、Linuxオペレーティングシステムを実行するAWS t4g.largeマシンを使用します。

コードの実装

1. Pythonコード

以下は、積分を計算するための単純なPythonコードの例です：

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

このコードをAWS t4g.largeマシンのLinux環境で実行したとき、かかった時間は：Elapsed: 32.59504199028015 sです。

2. Rustコード

同じ積分計算関数を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;
}

このRustコードを実行したとき、かかった時間は：Elapsed: 10.80 sです。

3. `pyo3`を使ってPython拡張を書く

3.1 プロジェクトの作成と依存関係のインストール

まず、新しいプロジェクトディレクトリを作成し、maturinライブラリをインストールします：

# (demoを必要なパッケージ名に置き換えてください)
$ mkdir leapcell_demo
$ cd leapcell_demo
$ pip install maturin

次に、pyo3プロジェクトを初期化します：

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

プロジェクト構造は以下の通りです：

.
├── Cargo.toml          // rustパッケージ管理ファイル、[lib]でターゲットの拡張パッケージ名を宣言する
├── src                 // rustソースファイルディレクトリ、拡張ファイルを書く場所。maturinが初期化すると自動的に作成される
│   └── lib.rs          // 拡張ファイル
├── pyproject.toml      // Pythonパッケージ管理ファイル、Pythonパッケージ名の定義を含む
├── .gitignore
├── Cargo.lock
└── leapcell_demo       // ターゲットモジュール名、手動で作成する必要がある
    ├── main.py         // テスト用ファイル
    └── leapcell_demo.cp312_amd64.pyd   // コンパイルされた動的リンクライブラリファイル、Pythonにインポートするためのもの

3.2 Rust拡張コードの書き方

src/lib.rsに以下のコードを書きます：

use pyo3::prelude::*;

/// 積分を計算する。
#[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;
}

/// Rustで実装されたPythonモジュール。この関数の名前は、`Cargo.toml`の`lib.name`設定と一致しなければならない。そうでないと、Pythonはモジュールをインポートできない。
#[pymodule]
fn leapcell_demo(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(integrate_f, m)?)?;
    Ok(())
}

3.3 拡張モジュールの使い方

この拡張モジュールを使う方法には2通りあります：

3.3.1 拡張をPythonパッケージとしてインストールする

$ maturin develop

このコマンドは、RustコードをPythonパッケージに変換し、現在のPython環境にインストールします。pip listを使って、インストールされたパッケージを表示することができます。

3.3.2 動的ファイルとしてコンパイルし、Pythonで読み込む

$ maturin develop --skip-install

--skip-installコマンドは、.pydファイル（例：leapcell_demo.cp312_amd64.pyd）を生成し、Pythonパッケージとしてインストールするのではなく、Pythonはこのファイルを直接インポートして使用できます。

また、--skip-installを--releaseに置き換えると、.whlファイルが生成されます。これは、Python pipインストール用のパッケージソースファイルです。

テストファイル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))

このコードを実行したとき、かかった時間は：Elapsed: 10.908721685409546 sです。

4. 並列処理による加速

4.1 Pythonのマルチプロセッシングの効果

Pythonのマルチプロセッシングは、コードの実装によっては、時にはシングルプロセス処理よりも遅くなることがあります：

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

このコードを実行したとき、かかった時間は：Elapsed: 18.86696743965149 sで、シングルプロセスバージョンの半分以下の時間です。

4.2 Pythonで使用するためのRustのマルチスレッドによる加速

Rustの並列処理ライブラリを使って、さらに計算を加速します：

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

/// Rustで実装されたPythonモジュール。この関数の名前は、`Cargo.toml`の`lib.name`設定と一致しなければならない。そうでないと、Pythonはモジュールをインポートできない。
#[pymodule]
fn leapcell_demo(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(integrate_f_parallel, m)?)?;
    Ok(())
}

3.2の手順に従って拡張ファイルをコンパイルして生成し、その後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))

このコードを実行したとき、かかった時間は：Elapsed: 0.9684994220733643 sです。単一スレッドのRustバージョンと比較すると、約10倍高速です；Pythonの並列バージョンと比較すると、約18倍高速です；Pythonのシングルプロセスバージョンと比較すると、約32倍高速です。

まとめ

Rustを使ってPythonコードを最適化することで、計算性能を大幅に向上させることができます。Rustは学習曲線がやや急峻ですが、多くの計算タスクを扱う必要のあるプロジェクトにとっては、コードの重要な部分をRustで書き直すことで、多くの時間コストを節約できます。既存のPythonプロジェクトを最適化するために、少しずつRustを使ってみることができます。単純な関数から始めて、徐々にその使い方を習得することができます。

###Leapcell: The Best of Serverless Web Hosting

最後に、PythonとRustをデプロイするための最高のプラットフォーム：Leapcell をおすすめします。

🚀 好きな言語で開発する

JavaScript、Python、Go、またはRustで簡単に開発できます。

🌍 無料で無制限のプロジェクトをデプロイする

使用した分だけ支払う——リクエストがなければ、課金されません。

⚡ 使った分だけ支払い、隠された費用はありません

アイドル料金はなく、シームレスなスケーラビリティがあります。

📖 ドキュメントを参照する

🔹 Twitterでフォローしてください： @LeapcellHQ