🦀

Rust・Python間でNumPyの配列を受け渡したい

2023/01/13に公開

4件

概要

Rustを使って機械学習アルゴリズムを実装してPythonから呼び出せるようにしたいのですが、言語間でのデータの受け渡しをどうすれば良いかが分からなかったので、幾つかの方法を実装して比較してみました。

Pythonで機械学習の実装を呼び出すということは、多くの場合NumPyの配列の受け渡しが出来なければいけないと思うので、今回は特にNumPyの配列の受け渡しの方法について考えます。

この記事で試した方法は以下の3つです。

ctypesを使って、配列をシリアライズして文字列として渡す
ctypesを使って、配列をdoubleのポインタとして渡す
PyO3とmaturinを使って、RustのコードからネイティブのPythonモジュールを作る

記事を読むにあたり注意してもらいたいこと

Rustに関しては勉強中なので、内容の正確性についてはあまり自信がありません。
あくまでも「こうやったら出来ました」という事例紹介だと思って読んで頂けるとありがたいです。

準備: 実験用Crateの作成

以下のコマンドにより、data_sharingという名前でライブラリクレートを作成します。

cargo new --lib data_sharing

次に、Cargo.tomlを以下のように編集します。

Cargo.toml

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

[lib]
crate-type = ["cdylib"]

[dependencies]
serde = { version = "1.0.152", features = ["derive"] }
serde_json = "1.0.91"
pyo3 = { version = "0.17", features = ["extension-module"] }
numpy = "0.17"

手法1: ctypesを使って、配列をシリアライズして文字列として渡す

まずは、手法1を実装してみます。おおまかに「シリアライズ」と言っていますが、データの受け渡しが出来れば何でも良いので今回はJSONを使うことにします。

Rust側のコードは以下の通りです。
何を解説するべきか自分でも分かりませんが、個人的にポイントだと思う点を列挙します。

Pythonから呼び出したい関数の頭には #[no_mangle] と pub unsafe extern "C" を付ける。
引数と返り値はC言語の文字列になるので、どちらも *const c_char にする (生ポインタ)。
生ポインタから値を取り出すのはRust的にはunsafeな処理なので、unsafeブロックで囲む必要がある。
CString::new() でCの文字列に変換し、.into_raw()で所有権を放棄しつつ生ポインタに変換して返す。^[1]
文字列の所有権を放棄したことでメモリリークが起きてしまうので、free_strという関数を作ってPython側から呼び出すことでメモリを解放する。

なお、Serdeの使用方法については解説がいくらでもありそうな感じがするのでここでは省略します。

use serde::{Deserialize, Serialize};
use serde_json::json;
use std::ffi::{c_char, CStr, CString};

#[derive(Serialize, Deserialize)]
struct JsonInput {
    arr: Vec<Vec<f64>>,
}

/// # Safety
/// PythonからctypesでJSON文字列を受け取ってJSON文字列を返す
#[no_mangle]
pub unsafe extern "C" fn json_to_json(json_str: *const c_char) -> *const c_char {
    // Pythonから受け取ったJSON文字列をrustの&strに変換
    let json_str = unsafe {
        assert!(!json_str.is_null());
        CStr::from_ptr(json_str).to_str().unwrap()
    };

    // JSON文字列をデシリアライズして配列を取り出す
    let input: JsonInput = serde_json::from_str(json_str).unwrap();
    let arr = input.arr;

    // 作成された配列にデータが入っているか念のため確認
    let first = arr[0][0];
    let last = arr.last().unwrap().last().unwrap();
    println!("arr[0][0] = {}, arr[-1][-1] = {}", first, last);

    // 配列をシリアライズしてJSON文字列に戻す
    let output = json!({
        "arr": arr,
    });
    let output = serde_json::to_string(&output).unwrap();

    // JSON文字列をCの文字列に変換して返す
    let c_str = CString::new(output).unwrap();
    c_str.into_raw()
}

/// Rust側で一度所有権を放棄した文字列をPython側から戻してもらって解放する
#[no_mangle]
pub extern "C" fn free_str(ptr: *const c_char) {
    unsafe {
        let _ = CString::from_raw(ptr as *mut c_char);
    }
}

コードが書けたらリリースモードでビルドして、共有ライブラリを作成します。
(速度の比較をするのが目的なのでリリースモードでビルドする必要があります。)

cargo build --release

OSによって出力されるファイル名が異なるので、環境に応じて以下のように読み替えてください。

Windowsの場合 target/release/data_sharing.dll
MacOSの場合 target/release/libdata_sharing.dylib
Linuxの場合 target/release/libdata_sharing.so

ちなみに今回は Windows で実装しているので、data_sharing.dll というファイル名になります。

最後に、作成された共有ライブラリを Python から呼び出すコードを実装すれば完成です。
ここでのポイントは以下の通り。

入力と出力のデータ型は c_void_pにする。^[2]
json.dumpsでJSON文字列を作成し、.encode('utf-8')でバイト列に変換する。
Rustから返ってきたvoidポインタをcastして文字列を取り出す。
文字列を取り出したら、不要になったvoidポインタをfree_strに渡してメモリを解放する。

import ctypes
import json
import numpy as np

lib = ctypes.cdll.LoadLibrary('data_sharing/target/release/data_sharing.dll')
lib.json_to_json.argtypes = [ctypes.c_void_p]
lib.json_to_json.restype = ctypes.c_void_p
lib.free_str.argtypes = [ctypes.c_void_p]

def json_to_json(x):
    input_json = json.dumps({'arr': x.tolist()}).encode('utf-8') # 配列をJSON文字列に変換
    ret_void = lib.json_to_json(input_json) # Rust側の関数を呼び出してvoid型のポインタとして文字列を受け取る
    ret_str = ctypes.cast(ret_void, ctypes.c_char_p).value.decode('utf-8') # void型のポインタをキャストして文字列を取り出す
    lib.free_str(ret_void) # Rust側で確保したメモリを解放する
    ret = json.loads(ret_str) # JSON文字列をデシリアライズして辞書型に変換
    return np.array(ret['arr']) # 辞書型から配列を取り出して返す


# 利用例
input_array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float64)
output_array = json_to_json(input_array)

このやり方の長所は、誰にでも理解できそうな単純なやり方だということです。
JSONで表現できる内容であればどのようなデータの受け渡しも容易に出来るので、パフォーマンスを気にしないという場合は案外便利かもしれません。
一方で、シリアライズやデシリアライズにそれなりの時間を要するので、パフォーマンスを気にする場合には別の方法を検討するべきだと思います。

手法2: ctypesを使って、配列をdoubleのポインタとして渡す

次に、手法2として NumPy の配列を double のポインタとして渡す方法を紹介します。
とはいえ、ctypesを使うので手法1と比較してそんなに大きな違いがあるわけではありません。
実装の手間はちょっと増えますが、シリアライズやデシリアライズの処理が不要なのでこちらの方が手法1よりもかなり高速に動作します。

Rust側の実装は以下の通りです。
基本的には手法1と大差ないのですが、敢えてポイントを挙げると以下の2点になります。

*const c_charとは異なり、*const f64からデータを読み取る時にはデータサイズの指定が必要 ^[3] (from_raw_partsの第2引数)
手法1と同じようにデータの所有権を放棄した上で生ポインタを返さなければいけないが、手法1と違って as_ptr と std::mem::forget を使っている ^[4]
配列の所有権を放棄したことでメモリリークが起きてしまうので、free_arrayという関数を作ってPython側から呼び出すことでメモリを解放する。^[5]

/// # Safety
/// Pythonからctypesでf64の配列を受け取ってf64の配列を返す
#[no_mangle]
pub unsafe extern "C" fn array_to_array(
    data: *const f64,
    n_rows: usize,
    n_cols: usize,
) -> *const f64 {
    // 生ポインタから値を読み取ってVec<Vec<f64>>に変換
    let data = unsafe { std::slice::from_raw_parts(data, n_rows * n_cols) };
    let arr: Vec<Vec<f64>> = data.chunks(n_cols).map(|x| x.to_vec()).collect();

    // 作成された配列にデータが入っているか念のため確認
    let first = arr[0][0];
    let last = arr.last().unwrap().last().unwrap();
    println!("arr[0][0] = {}, arr[-1][-1] = {}", first, last);

    // Vec<Vec<f64>> を *const f64 に変換して、データの所有権を放棄してから返す
    let output = arr.into_iter().flatten().collect::<Vec<f64>>();
    let p = output.as_ptr();
    std::mem::forget(output);
    p
}

/// Rust側で一度所有権を放棄した配列をPython側から戻してもらって解放する
#[no_mangle]
pub extern "C" fn free_array(ptr: *const f64, len: usize) {
    unsafe {
        let _ = std::slice::from_raw_parts(ptr as *mut f64, len);
    }
}

手法1と同じ手順でビルドしたら、以下のようにPythonから呼び出すことができます。
ここでのポイントは以下の4点です。

double型のポインタを引数や返り値にする場合は ctypes.POINTER(ctypes.c_double) を使う。
Numpyの配列xをdouble型のポインタに変換するには x.ctypes.data_as(ctypes.POINTER(ctypes.c_double)) とする。
double型のポインタをNumpyの配列に変換するには np.ctypeslib.as_array(ret, shape=(n_rows, n_cols)) とする。
メモリリークを防ぐために.copy()でコピーした配列を使い、Rustから受け取ったメモリはすぐに解放してしまう。^[6]

import ctypes
import numpy as np

lib = ctypes.cdll.LoadLibrary('data_sharing/target/release/data_sharing.dll')
lib.array_to_array.argtypes = [ctypes.POINTER(ctypes.c_double), ctypes.c_size_t, ctypes.c_size_t]
lib.array_to_array.restype = ctypes.POINTER(ctypes.c_double)
lib.free_array.argtypes = [ctypes.POINTER(ctypes.c_double), ctypes.c_size_t]

def array_to_array(x):
    n_rows, n_cols = x.shape
    ret = lib.array_to_array(x.ctypes.data_as(ctypes.POINTER(ctypes.c_double)), n_rows, n_cols)
    ret_arr = np.ctypeslib.as_array(ret, shape=(n_rows, n_cols)).copy()
    lib.free_array(ret, ret_arr.size)
    return ret_arr


# 利用例
input_array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float64)
output_array = array_to_array(input_array)

こちらの手法は、単純で分かりやすい上にパフォーマンス的にも悪くないので、面倒でなければ手法1よりはこちらの手法を使うほうがいいでしょう。

手法3: PyO3とmaturinを使って、RustのコードからネイティブのPythonモジュールを作る

最後に、RustをPythonと連携させるのに便利なクレートであるPyO3を使って、Pythonのネイティブモジュールを作る方法を紹介します。
まず、以下のようなRustのコードを書きます。

use numpy::{convert::IntoPyArray, PyArray2, PyReadonlyArray2};
use pyo3::{pymodule, types::PyModule, PyResult, Python};

#[pymodule]
fn data_sharing(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    /// 二次元のndarrayを受け取り、二次元のndarrayを返す
    #[pyfn(m)]
    #[pyo3(name = "ndarray_to_ndarray")]
    fn ndarray_to_ndarray_py<'py>(
        py: Python<'py>,
        arr: PyReadonlyArray2<f64>,
    ) -> &'py PyArray2<f64> {
        // input
        let arr = arr.as_array();

        // 作成された配列にデータが入っているか念のため確認
        let first = arr[[0, 0]];
        let last = arr[[arr.nrows() - 1, arr.ncols() - 1]];
        println!("arr[0][0] = {}, arr[-1][-1] = {}", first, last);

        // output
        let output = arr.to_owned();
        output.into_pyarray(py)
    }

    Ok(())
}

これを maturin というツールでビルドすると、data_sharingというPythonのモジュールが出来て、その中にndarray_to_ndarrayという関数がある状態になります。
そんなに解説できるほど理解出来ていませんが、numpyというクレートを使って PythonのNumPyみたいな操作が出来るのは Python ユーザー的には嬉しい点ですね。

Rustのコードが書けたら、Cargo.toml があるディレクトリに移動して以下のコマンドを実行します。^[7]

maturin develop --release

すると、以下のようにPythonから呼び出すことができるようになります。^[8]

import data_sharing

input_array = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float64)
output_array = data_sharing.ndarray_to_ndarray(input_array)

実験

実験環境

OS: Windows11
CPU: Core i7-12700
rustc 1.65.0
Python 3.10.8
Numpy 1.23.5

各手法の速度比較

様々なサイズのデータに対して、各手法の実行にどれくらい時間がかかるかを計測してみました。実験に使ったコードは以下の通りです。

import ctypes
import json
import time
import numpy as np
import pandas as pd
import data_sharing

lib = ctypes.cdll.LoadLibrary('data_sharing/target/release/data_sharing.dll')

# json_to_json
lib.json_to_json.argtypes = [ctypes.c_void_p]
lib.json_to_json.restype = ctypes.c_void_p
lib.free_str.argtypes = [ctypes.c_void_p]

# array_to_array
lib.array_to_array.argtypes = [ctypes.POINTER(ctypes.c_double), ctypes.c_size_t, ctypes.c_size_t]
lib.array_to_array.restype = ctypes.POINTER(ctypes.c_double)
lib.free_array.argtypes = [ctypes.POINTER(ctypes.c_double), ctypes.c_size_t]


def benchmark(f):
    n_trials = 10
    d = 10
    times = []
    for n in num_samples:
        tmp = []
        for _ in range(n_trials):
            x = np.arange(n * d, dtype=np.float64).reshape(n, d)
            start = time.time()
            f(x)
            end = time.time()
            tmp.append(end - start)
        times.append(tmp)
    return np.array(times)


def json_to_json(x):
    input_json = json.dumps({'arr': x.tolist()}).encode('utf-8')
    ret_void = lib.json_to_json(input_json)
    ret_str = ctypes.cast(ret_void, ctypes.c_char_p).value.decode('utf-8')
    lib.free_str(ret_void)
    ret = json.loads(ret_str)
    return np.array(ret['arr'])


def array_to_array(x):
    n_rows, n_cols = x.shape
    ret = lib.array_to_array(x.ctypes.data_as(ctypes.POINTER(ctypes.c_double)), n_rows, n_cols)
    ret_arr = np.ctypeslib.as_array(ret, shape=(n_rows, n_cols)).copy()
    lib.free_array(ret, ret_arr.size)
    return ret_arr


if __name__ == '__main__':
    num_samples = [1000, 10000, 100000, 1000000]
    df = pd.DataFrame({
            'json_to_json': benchmark(json_to_json).mean(axis=1),
            'array_to_array': benchmark(array_to_array).mean(axis=1),
            'ndarray_to_ndarray': benchmark(data_sharing.ndarray_to_ndarray).mean(axis=1),
        },
        index=num_samples,
    )
    print(df)

実験結果はこのようになりました。

         json_to_json  array_to_array  ndarray_to_ndarray
1000         0.003131        0.000000            0.000000
10000        0.042302        0.001562            0.000000
100000       0.459450        0.009969            0.001562
1000000      5.151885        0.117958            0.019355

各行はデータサイズ、各列はそれぞれの手法に対応しています。書かれている数字は平均実行時間 (秒)です。

予想通り、手法1 (json_to_json) は相対的にかなり遅く、手法2 (array_to_array) と手法3 (ndarray_to_ndarray) はそれに比べるとかなり早いです。
手法2と手法3を比較すると、特にデータが大きい時には手法3の方が早い傾向にあるようです。^[9]

結局どの手法を使っていくか

速度的には手法3が一番早く、便利そうでもあるのですが、以下のような理由でひとまずは手法2を使っていくことにしたいと思います。

速度面で差があるといっても自分の目的からするとあまり問題にならないくらいの差
Python や Numpy との連携に特化するよりは、より一般的な手法の方に興味がある
PyO3のドキュメントを読むのが億劫 (さっさと機械学習の話にいきたいので...)
Python のことは忘れて Rust のプログラミングを楽しみたい気分だが、numpy クレートを使うと NumPy を使っているような気持ちになりそう

適当な理由も混ざっていますが、趣味プログラミングなので勘弁してください。

脚注

所有権が残っていると関数が終了したときにRustによってメモリが解放されてPython側で受け取れない。 ↩︎
c_char_pにしていたら自動で文字列に変換されてしまいPython側からメモリを解放できなくなってしまったのでc_void_pにしています。 ↩︎
文字列と違って終端にNULLがないのが理由だと思われます。 ↩︎
何か深い理由があってこうしているわけではなく、単純に色々試してこれで上手くいったからこうなっています。 ↩︎
termoshttさん、ご指摘頂きありがとうございました。 ↩︎
これにより余計なコピーが発生してしまうわけですが、単純さと安全性重視でこうすることにしました。 ↩︎
事前にPython側でmaturinをインストールしておく必要があります　(pip install maturin) ↩︎
パスが通っているどこかに wheel 形式のファイルを作ってくれるみたいですが、あまり細かいことは調べていません。 ↩︎
ただ、手法2は実装次第でもう少し早くできるだろうという感じもしています。 ↩︎

Discussion

termoshtt

手法1と同じようにデータの所有権を放棄した上で生ポインタを返さなければいけないが、手法1と違って as_ptr と std::mem::forge を使っている

std::mem::forgetした以上、このメモリはリークします。Rust側で確保したメモリをPythonに返す場合、メモリリークさせない為には

もう一度Python側からRust側にメモリを返し、Rustでポインタから構造体を復元した上でDropする
Python側のデストラクタが実行された際に一緒に上手く解放されるようにフックを差し込む

の二通りが考えられ、PyO3は後者を自動的にやってくれます。自分でこの機能を実装する場合は

が参考になるでしょう。

skwbc

ありがとうございます。この辺ちょっとあやふやだったので大変勉強になりました。

ご提案頂いた修正方法についてなのですが、とりあえず手法1と手法2についてはRustにメモリを返す方法で対応しようと思います。
Python の C API を使うとなると個人的にはPyO3を使うよりもハードルが高いので、今後後者の方法を採用する場合は素直にPyO3を使うことになりそうです。

メモリリークが本当に起きているかと、Rustにメモリを返してメモリリークを解消できるかについても一応試してみたのでご興味がありましたらご確認ください。

実験

メモリリークを起こすRustのコードを書く。

lib.rs

use std::ffi::{c_char, CString};

/// forgetを使ってf64のポインタを作ってPythonに渡す
#[no_mangle]
pub extern "C" fn make_array_with_forget() -> *const f64 {
    let mut v = Vec::new();
    for i in 0..10000 {
        v.push(i as f64);
    }
    let ptr = v.as_ptr();
    std::mem::forget(v);
    ptr
}

/// f64のポインタをPythonから受け取ってメモリを解放する
/// (from_raw_partsを使うのが適切なのかは不明)
#[no_mangle]
pub extern "C" fn free_array(ptr: *const f64) {
    unsafe {
        let _ = std::slice::from_raw_parts(ptr as *mut f64, 10000);
    }
}

/// 手法1と同じ方法でc_charの生ポインタを作ってPythonに渡す
#[no_mangle]
pub extern "C" fn make_str() -> *const c_char {
    let s = "a".repeat(80000);
    CString::new(s).unwrap().into_raw()
}

/// c_charのポインタをPythonから受け取ってメモリを開放する
/// (from_raw を使うのが適切なのかは不明)
#[no_mangle]
pub extern "C" fn free_str(ptr: *const c_char) {
    unsafe {
        let _ = CString::from_raw(ptr as *mut c_char);
    }
}

メモリリークするコードをPythonから呼び出して本当にリークするか確認する。

check_memory_leak.py

import ctypes
import psutil
import matplotlib.pyplot as plt
import numpy as np

plt.style.use('ggplot')

N = 10000

lib = ctypes.cdll.LoadLibrary('memory_leak/target/release/memory_leak.dll')
lib.make_array_with_forget.restype = ctypes.POINTER(ctypes.c_double)
lib.free_array.argtypes = [ctypes.POINTER(ctypes.c_double)]
lib.make_str.restype = ctypes.c_void_p
lib.free_str.argtypes = [ctypes.c_void_p]

def check_memory_leak(f, num_iter=10000):
    before = psutil.virtual_memory().used
    track = []
    for i in range(num_iter):
        f()
        track.append(psutil.virtual_memory().used - before)
    return track


def make_array_by_numpy() -> None:
    np.zeros((N,))


def make_array_with_forget() -> None:
    np.ctypeslib.as_array(lib.make_array_with_forget(), shape=(N,))


def make_array_with_forget_and_free() -> None:
    arr = lib.make_array_with_forget()
    lib.free_array(arr)


def make_str() -> None:
    lib.make_str()


def make_str_with_free() -> None:
    ptr = lib.make_str()
    s = ctypes.cast(ptr, ctypes.c_char_p).value.decode('utf-8')
    lib.free_str(ptr)


f_list = [
    make_array_by_numpy,
    make_array_with_forget,
    make_array_with_forget_and_free,
    make_str,
    make_str_with_free,
]
for f in f_list:
    track = check_memory_leak(f)
    print(f'{f.__name__}: {track[0]}, {track[-1]}')
    plt.plot(track, label=f.__name__)

plt.legend()
plt.savefig('memory_leak.png')

結果、手法1でも手法2でもメモリリークしていることが確認できた。
また、Rustにメモリを返す方法でメモリリークを防げているらしいことも確認できた。

シャクシャイヌ

PythonとRustの間をどうするか悩んでいたので大変参考にさせていただきました。sliceは所有権を持たないので、sliceをfrom_rawしてもメモリは解放されないと思います。Vecに戻せば解放されるはずですが、capacityとlengthがないとVecは復元できないので多少面倒かもしれません。

配列の生ポインタからメモリを解放するのはやたらと難しいようなので、基本的には配列を含む自作structを作り、それをBoxに入れて、into_raw()で所有権を放棄してポインタをpythonに渡す。pythonからstructのメモリを返してもらったらBox::from_raw()でBoxに戻して所有権を取り戻すことで、簡単にメモリを解放できると思います。配列の生ポインタはその自作structにVecを定義して、そのVecをas_ptr()することで取り出すのが良いのではないかと思います。

シャクシャイヌ

よく見るとmake_array_with_forget_and_freeでメモリが解放されてるように見えますね・・・

ただスライスに所有権がないというのはThe bookに書いてありますので、sliceを作って破棄してメモリが解放されるというのはかなり不思議（危険？）なことが起きているように思います。