🦀

maturinを使ってRustのコードをpythonから呼び出す

2023/02/18に公開3件

Python

Rust

maturin

tech

はじめに

自分が機械学習エンジニアとして日々業務をしているので、Pythonを使ってサーバサイドを開発したり、機械学習モデルを構築したりしています。そんな中で、ふつふつと思うことがあります

Pythonって遅くない？！

そうです。どれだけきれいに書けたとしても、多少雑に書いたC++やRustのコードと比べても遅い場合があります。
でも、できるだけPythonで書きたいけど、アルゴリズムは高速化したい...なんてことありますよね？

そこで、本記事では、maturinというPythonバインディングツールを使ってPythonの自作ライブラリを作ろうと思います。

Maturin

matrinは、pyo3、rust-cpython、cffi、uniffiバインディングツールをPythonパッケージとして公開しやすいようにしてくれるツールで、rustバイナリでクレートを構築し公開します。

つまり、PyPIで公開したり、その他プラットフォームで公開するのが容易になります。

How to use

まずはmaturinを入手します。

pip install maturin

Rustのツールキットもインストールしておきましょう。公式では、以下のコマンドで入れることを推奨しています。

curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh

準備できれば、任意のディレクトリで以下のコマンド操作をします。今回は、Rustバインディングツールで一番メジャーpyo3を使ってみます。

$ maturin init
✔ 🤷 Which kind of bindings to use? · pyo3
  ✨ Done! Initialized project <path to project>/maturin_handson

階層構造やその他必要なモジュールを自動生成してくれて、以下の階層構造になるはずです。

❯ tree .
.
├── Cargo.toml
├── README.md
├── pyproject.toml
└── src
    └── lib.rs

スクリプトが4つ生成されており、一つづつ説明すると、

Cargo.toml
Rustの必要なモジュールを追加するところです。
pyproject.toml
パッケージとして公開する際に必要となるものです。
src/lib.rs
パッケージの処理を書くメインスクリプトです。
README.md
よく見るやつ

2つの引数を足す関数

まずは、自動生成される、2つの引数の数字を足して返す関数を実行してみます。

スクリプト

lib.rs

use pyo3::prelude::*;

/// Formats the sum of two numbers as string.
#[pyfunction]
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
    Ok((a + b).to_string())
}

/// A Python module implemented in Rust.
#[pymodule]
fn maturin_handson(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
    Ok(())
}

L4-L7 #[pyfunction]をつけると、Pythonバインディングできるようになります。

L10-L14 : #[pymodule]をつけることで、パッケージとして出すことができます。m.add_functionでバインディングする関数を定義することでpythonから使えるようになります。

その他にもここを見ると、PyClassやPython classの__init__()に対応するような#[new]などもあるので、試してみてください。

build

以下のコマンドでビルドできます。--releaseをつけると、実行速度が早くなるそうです。

maturin build -i python3 --release

Python パッケージ化

buildだけでは、pythonパッケージにならないので、以下のコマンドで、install までします。developを使う場合、pyenvやcondaなどで仮想環境を作っていないと動かないので注意です

# vertual env にパッケージを入れる
maturin develop --release

# vertual env なしで入れる方法
maturin build -i python3 --release
pip instlal .

Pythonから呼び出す

パッケージ化できたら、Pythonスクリプトから呼び出してみましょう。testディレクトリを作成します。

$ mkdir test
$ touch test/__init__.py
$ touch test/test_sum.py

$ tree .
.
├── Cargo.lock
├── Cargo.toml
├── README.md
├── pyproject.toml
├── src
│   └── lib.rs
└── test
    ├── __init__.py
    └── test_sum.py

test_sum.pyの中身は以下の通りです。今回はpytestを用いているので関数名の最初にtest_をつけるのを忘れずに

test_sum.py

from maturin_handson import sum_as_string

def test_add():
    a, b = 1, 2
    c = sum_as_string(a, b)
    assert c == "3"

それでは、pytestを使って、テストしていきましょう。

$ pip install pytest
$ pytest

$ pytest
========== test session starts ==========
platform darwin -- Python 3.10.9, pytest-7.2.1, pluggy-1.0.0
rootdir: <path to file from root>/maturin_handson
collected 1 item

test/test_sum.py .                  [100%]

========== 1 passed in 0.00s ============

うまくいってますね。これでpythonからrust コードを呼び出すことができました。

Makefileにまとめる

上記のコマンドをいちいち打つのは面倒なので、自分はMakefileにまとめました。

Makefile

clean:
	cargo clean
	rm -rf *~ dist *.egg-info build target

build:
	maturin build -i python3 --release
	
develop:
	maturin develop --release

fmt:
	rustfmt src/*
	black test/*

Numpyを渡して数値解析

それでは本格的に、Pythonからベクトルなどをrustに渡して数値解析してみましょう。

今回は以下の記事を参考に、rustを用いて、行列の基本的な計算をrustで行わせます。

スクリプトの処理としては、単純に以下の流れを作成します。

Pythonでnumpyオブジェクトを生成
rustに渡して
rustで処理して結果を返す

rustパッケージを持ってくる

基本的にCargo.tomlに記述していけば、buildしたときに取ってきてくれます。
しかし、以下のコマンドを使えば、latestを自動で取ってきてくれます。pythonでいうとpoetryと同じような使い方ができるのでめちゃくちゃわかりやすいですね！

$ cargo add ndarray numpy ndarray-linalg

Cargo.tomlを見てみると、ndarrayが追加されているのがわかります。

Cargo.toml

[package]
name = "maturin_handson"
version = "0.1.0"
edition = "2021"

# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[lib]
name = "maturin_handson"
crate-type = ["cdylib"]

[dependencies]
+ndarray = "0.15.6"
+numpy = "0.17.2"
+ndarray-linalg = "0.16.0"
pyo3 = { version = "0.17.3", features = ["extension-module"] }

main スクリプトを追加

数値解析をするスクリプトを作っていきます。

PyArrayモジュールを型として引数に指定することで、pythonで作ったnumpyモジュールをそのまま引数に渡して、扱うことができます。ただし、as_array()を使う際はunsafeを使わないといけません。

注意点として、一度関数を作って、pyfunctionでラップしてあげないと、色々とエラーがでました。
どなたか、この原因がわかる方がいればコメントいただけると幸いです。

src/lib.rs

+use ndarray::prelude::*;
+use numpy::{IntoPyArray, PyArrayDyn};
use pyo3::prelude::*;

/// Formats the sum of two numbers as string.
#[pyfunction]
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
    Ok((a + b).to_string())
}

+fn axpy(a: f64, x: ArrayViewD<f64>, y: ArrayViewD<f64>) -> ArrayD<f64> {
+    a * &x + &y
+}

+#[pyfunction]
+fn axpy_py(py: Python, a: f64, x: &PyArrayDyn<f64>, y: &PyArrayDyn<f64>) -> Py<PyArrayDyn<f64>> {
+    // axpyのラッパー
+    unsafe {
+        let x = x.as_array();
+        let y = y.as_array();
+        axpy(a, x, y).into_pyarray(py).to_owned()
+    }
+}

/// A Python module implemented in Rust.
#[pymodule]
fn maturin_handson(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
+    m.add_function(wrap_pyfunction!(axpy_py, m)?)?;
    Ok(())
}

テスト

同じくpythonのテストコードを作成します。

test_sum.py

+from maturin_handson import sum_as_string, axpy_py
+import numpy as np
def test_add():
   a, b = 1, 2
   c = sum_as_string(a, b)
   assert c == "3"

+def test_axpy():
+    a = 2
+    x = np.array([[1., 1., 1.],
+                  [1., 1., 1.],
+                  [1., 1., 1.]])
+
+    y = np.array([[1., 1., 1.],
+                  [1., 1., 1.],
+                  [1., 1., 1.]])
+
+    res = axpy_py(a, x, y)
+    assert (res == a*x+y).all()

テストを走らせてみましょう

$ pytest
============ test session starts =======
platform darwin -- Python 3.10.9, pytest-7.2.1, pluggy-1.0.0
rootdir: <path to project>/maturin_handson
collected 2 items

test/test_sum.py ..                 [100%]

============ 2 passed in 0.17s ==========