RustバインディングでPython Lambdaを高速化

AWS Lambdaを使ったPythonアプリケーションで、処理速度がボトルネックになることはありませんか？特に計算負荷の高い処理では、Pythonの柔軟性と簡潔さがパフォーマンスの犠牲となることがあります。このような場合、Rustを利用して一部の処理を高速化することが可能です。

本記事では、Rustバインディングを活用して、Python Lambda関数を高速化する手順をご紹介します。

# この記事で解決する課題
Pythonの処理をRustで置き換える方法を知りたい。
AWS LambdaでRustとPythonを組み合わせる手法を学びたい。
処理速度を計測し、改善効果を確認したい。

# Rustバインディングとは？
Rustバインディングは、Rustの関数やロジックを他の言語（本記事ではPython）から直接利用可能にする仕組みです。Rustはその高速性と安全性から、多くの場面で性能改善の目的で利用されています。
PythonとRustを組み合わせることで、Pythonの簡潔さとRustの高性能を同時に活かすことができます。

# 実施環境
・Mac OS Sonoma 14.6.1
・Python 3.12.7
・AWS SAM CLI 1.127.0
・AWS CLI 2.15.25
・Docker 27.2.0

# Rustの実行環境のセットアップ
Rustバインディングを実装するために、以下の環境をセットアップします。

## Rustのインストール
rustupを使ってRustをインストールします。
```
$ curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh
```
インストールを実行すると以下のようにインストール方法のオプション選択が表示されます。標準（default）インストールである`1`を入力してEnterキーを押下します。
```
Current installation options:


   default host triple: aarch64-apple-darwin
     default toolchain: stable (default)
               profile: default
  modify PATH variable: yes

1) Proceed with standard installation (default - just press enter)
2) Customize installation
3) Cancel installation
>1
```
インストール後、環境変数を更新します。
```
$ source $HOME/.cargo/env
```
インストールが成功したか確認します。
rustupのバージョンが表示されればインストールが完了です。
```
$ rustup -V
rustup 1.27.1 (54dd3d00f 2024-04-24)
info: This is the version for the rustup toolchain manager, not the rustc compiler.
info: The currently active `rustc` version is `rustc 1.83.0 (90b35a623 2024-11-26)`
```

デフォルトではビルドのターゲットプラットフォームが、インストール時のOS環境となっているので、必要なターゲットを追加します。(本記事ではLambdaのアーキテクチャはarm64を想定)
```
$ rustup target add aarch64-unknown-linux-gnu
```
:::note info
※アーキテクチャのx86_64を使用する場合は以下を追加します。
x86_64-unknown-linux-gnu
:::

## maturinのインストール
maturinはRustで書かれたコードをPythonから利用できるようにするためのツールです。
pipを使ってインストールします。
```
$ pip install maturin
```
インストールが成功したか確認します。
maturinのバージョンが表示されればインストールが完了です。
```
$ maturin --version
maturin 1.7.6
```
これでRustとmaturinの準備が整いました。

# Rustプロジェクトの作成
rustupと共にインストールされるcargoを使用してRustプロジェクトを初期化し、maturinでPythonから利用可能なRustモジュールを構築します。

## Rustプロジェクトの初期化
以下のコマンドでRustプロジェクトを作成します。
```
$ cargo new --lib rust_lib
```
これにより、以下のようなディレクトリ構成が生成されます。
```
my_project
└──rust_lib/
    ├── Cargo.toml
    └── src/
        └── lib.rs
```
## Cargo.tomlの編集
Cargo.tomlに以下の依存関係とプロジェクトがPython拡張モジュールであることを示す設定を追加します。
```rust:Cargo.toml
[package]
name = "rust_lib"
version = "0.1.0"
edition = "2021"

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

[dependencies]
pyo3 = { version = "0.18", features = ["extension-module"] }
```

## Rustコードの実装
`src/lib.rs`を編集して、Pythonで呼び出す関数を実装します。
本記事では1000万ループする処理を実施してみます。
```rust:lib.rs
use pyo3::prelude::*;

#[pyfunction]
fn execute_loop() -> i32 {
    let mut count = 0;
    for _ in 0..10_000_000 {
        count += 1;
    }
    count
}

#[pymodule]
fn rust_lib(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(execute_loop, m)?)?;
    Ok(())
}
```
## Rustコードのビルド
`rust_lib`ディレクトリ配下で以下のコマンドで実行して、RustコードをPythonから使用可能な形式（Python拡張モジュール）にビルドします。
```
$ maturin build --release --target aarch64-unknown-linux-gnu --zig -i python3.12
```
:::note info
各オプションの説明
`--release`：デバック情報や使用されていないデッドコードの除去などが行われ、最適化されたバイナリが作成されます。
`--target`：ビルドのターゲットプラットフォームを指定します。
`--zig`：Zigランタイムを使用してクロスコンパイル（例：macOSでLinux向けのビルド）を有効にします。
`-i`：ビルドターゲットとするPythonのバージョンを指定します。
:::
コマンドが成功すると、`rust_lib/target/wheels`配下にPythonのwheel形式（.whl）のバイナリが作成され、パッケージ化されます。
```
my_project
└──rust_lib/
    ├── Cargo.lock
    ├── Cargo.toml
    ├── __init__.py
    ├── src/
    │   └── lib.rs
    └── target/
        ├── aarch64-unknown-linux-gnu
        ├── maturin
        ├── release
        └── wheels/
            └── rust_lib-0.1.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
```

# ビルドとデプロイ
AWS SAMを使ってPython Lambdaをビルド、デプロイします。

## lambda_handlerの実装
作成したRustコードを参照するlambda_handlerを実装します。
time関数を使って、参照したRust関数の処理時間を計測します。
```python:handler.py
import rust_lib
import time

def lambda_handler(event, context):

    start_time = time.time()

    # Rust関数の実行
    count = rust_lib.execute_loop()

    elapsed_time = time.time() - start_time

    print(f'{count} 回, {elapsed_time} 秒')

    return
```
`src/functions/python_binding_sample`ディレクトリを作成し、`handler.py`ファイルを格納します。
```
my_project
├── rust_lib/
│   ├── Cargo.lock
│   ├── Cargo.toml
│   ├── src/
│   │   └── lib.rs
│   └── target/
│       ├── aarch64-unknown-linux-gnu
│       ├── maturin
│       ├── release
│       └── wheels/
│           └── rust_lib-0.1.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
└── src/
    └── functions/
        └── python_binding_sample/
            └── handler.py
```

## Lambda Layerの作成
`src/layers/rust_lib`ディレクトリを作成し、
先ほど作成した、Rustコードのパッケージをコピーします。
```
$ cp rust_lib/target/wheels/rust_lib-0.1.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl src/layers/r
ust_lib 
```
:::note info
Rustのビルド時にそのまま配置したい場合は、以下のコマンド。
```
$ maturin build -m rust_lib/Cargo.toml --release --target aarch64-unknown-linux-gnu --zig -i python3.12 -o src/layers/rust_lib
```
:::
ビルド時にパッケージをインストールするため、`requirements.txt`を作成します。
```txt:requirements.txt
./rust_lib-0.1.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
```
作成した`requirements.txt`ファイルを`src/layers/rust_lib`ディレクトリに格納します。
```
my_project
├── rust_lib/
│   ├── Cargo.lock
│   ├── Cargo.toml
│   ├── src/
│   │   └── lib.rs
│   └── target/
│       ├── aarch64-unknown-linux-gnu
│       ├── maturin
│       ├── release
│       └── wheels/
│           └── rust_lib-0.1.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
└── src/
    ├── functions/
    │   └── python_binding_sample/
    │       └── handler.py
    └── layers/
        └── rust_lib/
            ├── requirements.txt
            └── rust_lib-0.1.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
```
## SAMテンプレート作成
`template.yaml`を作成します。
```yaml:template.yaml
AWSTemplateFormatVersion: '2010-09-09'
Transform: AWS::Serverless-2016-10-31
Description: >
  my_project
  
  Sample SAM Template for my_project

Globals:
  Function:
    Timeout: 15

    Tracing: Active
    LoggingConfig:
      LogFormat: JSON

Resources:
  RustLibLayer:
    Type: AWS::Serverless::LayerVersion
    Properties:
      Description: Layer description
      LayerName: rust-lib-layer
      ContentUri: src/layers/rust_lib
      CompatibleRuntimes:
        - python3.12
      CompatibleArchitectures:
        - arm64
    Metadata:
      BuildMethod: python3.12
      BuildArchitecture: arm64

  PythonBindingSampleFunction:
    Type: AWS::Serverless::Function
    Properties:
      CodeUri: src/functions/python_binding_sample
      Handler: handler.lambda_handler
      Runtime: python3.12
      Architectures:
      - arm64
      MemorySize: 256
      Layers:
        - !Ref RustLibLayer

```
プロジェクト直下に格納します。

```
my_project
├── rust_lib/
│   ├── Cargo.lock
│   ├── Cargo.toml
│   ├── src/
│   │   └── lib.rs
│   └── target/
│       ├── aarch64-unknown-linux-gnu
│       ├── maturin
│       ├── release
│       └── wheels/
│           └── rust_lib-0.1.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
├── src/
│   ├── functions/
│   │   └── python_binding_sample/
│   │       └── handler.py
│   └── layers/
│       └── rust_lib/
│           ├── requirements.txt
│           └── rust_lib-0.1.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl
└── template.yaml
```

## デプロイの実行
Rustコードのパッケージは、Lambda(AmazonLinux2)環境をターゲットとしてビルドされているため、`--use-container`オプションでコンテナを使ってビルドします。
```
$ sam build --use-container
```
Lambdaをデプロイします。
```
$ sam deploy  --stack-name my-sam-stack --region ap-northeast-1 --capabilities CAPABILITY_IAM --s3-bucket my-sam-stack-bucket
```

# パフォーマンス比較
Pythonで同様に1000万回ループする処理を実行した際の処理時間を比較します。
```python:python_lib.py
def execute_loop():
    count = 0
    for i in range(10000000):
        count += 1

    return count
```

## 実行結果
処理時間を比較したところ、Pythonの関数では3秒程かかっていた処理が、Rustの関数では約3マイクロ秒と単純に見ると10万倍のパフォーマンス向上が確認できました。
```
Python実装: 10000000 回, 3.527838945388794 秒
Rust実装  : 10000000 回, 3.5762786865234375e-06 秒
```


# おわりに
本記事では、Rustバインディングを用いてPython Lambda関数を高速化する方法を解説しました。
Rustを使うことで、Pythonでは時間がかかっていた処理のパフォーマンス向上が期待できます。

最後までお読みいただきありがとうございました！
この記事の情報が、少しでも読んでくださった方々のお役にたてれば嬉しいです。


# 参考ドキュメント

https://community.aws/content/2ZSJAWJt2VbYrGOXPTv1rfmmLSP/rustifying-serverless-boost-aws-lambda-performance-with-rust

https://gihyo.jp/article/2023/07/monthly-python-2307

https://blog.agatan.dev/posts/maturin-pyo3/