FAAS in Python, WASM, WASI and Go

FAAS in Go with WASM, WASI and Rust - Eli Bendersky's website という投稿を読んで、Python で似たようなことをやるにはどうすれば良いか知るために試してみました。

この記事の内容は以下の通りです。

Python で WASI を使って WASM コードをロードし、Web サーバーに接続する方法。

WASI を使う Wasm の作成方法などは、ここでは説明しません（コードは載せておきます）。

Design

アプリケーションの基本デザインは FAAS in Go with WASM, WASI and Rust - Eli Bendersky's website と同じです。

以下のステップで動作します。

Python web server が HTTP GET リクエストを受け取る。
Python web server が HTTP GET リクエストのパスを解析し、対応する Wasm モジュールをロードし、実行する。
Wasm モジュールは、Python web server から環境変数でデータを受け取って、処理結果を stdout に出力する。
Python web server は、Wasm プラグインの処理結果を stdout から読み取って、HTTP レスポンスとして返す。

受け渡すデータについて

今回は Python と Wasm 間のデータのやり取りを JSON で行います。

サーバー実装

Python (FastAPI) でサーバーを実装します（全体のコードはこちら）。

app = FastAPI()

logging.basicConfig(level="INFO")
logger = logging.getLogger(__name__)


class Item(BaseModel):
    """Item."""

    name: str
    age: int
    description: str


@app.get("/{funcname}")
async def execute(funcname: str) -> Item:
    """Execute wasm function."""
    # refs: https://gist.github.com/pims/711549577759ad1341f1a90860f1f3a5

    wasm_path = f"functions/{funcname}.wasm"
    # なんらかの処理をして Wasm モジュールに渡すデータを作成
    # これはテスト用のデータ
    d = {"name": "zztkm", "age": 20}

    try:
        item_dict = await invoke_wasm_module(funcname, wasm_path, d)
        return Item(**item_dict)
    except Exception:
        logger.exception("error")
        raise HTTPException(status_code=400, detail="Bad Request") from None

このサーバーは、HTTP GET リクエストを受け取り、パスパラメーターを解析し、対応する Wasm モジュールパスを invoke_wasm_module 関数に渡します。そして、その返り値を Item オブジェクトに詰め替えて HTTP レスポンスとして返します。

Wasm を実行する関数の実装

サーバーの実装ができたので、次に Wasm モジュールを実行する関数を実装します。

async def invoke_wasm_module(modname: str, wasm_path: str, input_data: dict[str, Any]) -> dict[str, Any]:
    """Invoke wasm module."""
    # refs: https://gist.github.com/pims/711549577759ad1341f1a90860f1f3a5

    engine = Engine()
    linker = Linker(engine)
    linker.define_wasi()

    # functions/funcname.wasm が存在するか確認
    try:
        module = Module.from_file(linker.engine, wasm_path)
    except Exception:
        logger.exception("module not found %s", wasm_path)
        raise

    config = WasiConfig()

    # generate tenmporary file for stdout_file
    async with tempfile.NamedTemporaryFile() as f:
        config.stdout_file = f.name
        config.env = [["DATA", json.dumps(input_data)]]

        store = Store(linker.engine)
        store.set_wasi(config)

        instance = linker.instantiate(store, module)
        start = instance.exports(store)["_start"]

        logger.info("start wasm module %s", modname)
        try:
            start(store)
            out = await f.read()
            return json.loads(out.decode(encoding="utf-8"))
        except WasmtimeError as e:
            logger.debug(e)
            out = await f.read()
            return json.loads(out.decode(encoding="utf-8"))

関数の説明

wasmtime-py は WASI の実行をサポートしています。
- コードの大半は WASM モジュールをインスタンス化するためのものです。
Wasm モジュールに値を渡すために環境変数を使っています。
- これはいくつかの可能性を示すためであり、Wasm モジュールに値を渡す方法には stdin を使う方法もあります。
Wasm モジュールの出力は stdout ストリームから受け取ります
- WasiConfig.stdout_file にファイルパスを指定することで、Wasm モジュールの標準出力への書き込みが指定したファイルに書き込まれます。
  - 詳しくは、WasiConfig.stdout_file | wasmtime API documentation
最後に、ファイルの内容を読み取って、JSON を dict に変換して返します。

Wasm モジュールの実装

main.go

package main

import (
	"encoding/json"
	"os"
)

type Input struct {
	Name string `json:"name"`
	Age int `json:"age"`
}

type Output struct {
	Name string `json:"name"`
	Age int `json:"age"`
	Description string `json:"description"`
}

func main() {
	d := os.Getenv("DATA")
	var input Input
	json.Unmarshal([]byte(d), &input)

	output := Output{
		Name: input.Name,
		Age: input.Age,
		Description: "This is a description",
	}

	b, _ := json.Marshal(output)
	os.Stdout.Write(b)
}

そして、以下のようにビルドします。

GOOS=wasip1 GOARCH=wasm go build -o gojson.wasm main.go

Go の WASI サポートについては、公式ブログを読んでみてください。

ここで生成した gojson.wasm を functions ディレクトリに配置して、サーバーを起動します。起動したら curl で叩いてみます。

$ curl "http://127.0.0.1:8000/gojson"
{"name":"zztkm","age":20,"description":"This is a description"}

上記のような出力が得られたら成功です 😄

今回は、Wasm モジュールをローカルファイルシステムから取得しましたが、S3 や R2 などのオブジェクトストレージから読み込んでみると、サーバーを停止することなく、機能拡張をすることができるようになります（実際には毎回取得するわけにはいかないのでキャッシュするなど工夫が必要）！

まとめ

Python (ホスト)と Wasm モジュールでデータのやり取りをする方法を示しました。
今回は Wasm モジュールを Go で実装しましたが、Rust などの別言語でも同じように実装することができます。
Web サーバーだけでローカルの CLI ツールでも同様に、標準入出力を通してデータのやり取りをすることができます。
- 例: sqlc
このようなプラグインシステムは色々と応用が効きそうなので、今後もいろんなパターンで遊んでみたいと思います。

Veltioblog

Table of contents

Design

受け渡すデータについて

サーバー実装

Wasm を実行する関数の実装

Wasm モジュールの実装

まとめ