FastAPI+MessagePackで画像とメタデータを一緒にPOSTする
FastAPIで画像とメタデータ（テキストやリストなど）を一緒に、1つのオブジェクトとしてPOSTする方法を解説します。MessagePackで全部丸々固めてしまうやり方で、サクッと作りたいときに便利ではないかと思います。SageMaker Serverless Inferenceでの利用を想定しています。
目次
前提
※前置きが長いので、タイトルの内容知りたい人は、「MessagePackのバイナリをFastAPIで読む」までスキップしてください
前提としては、SageMaker Serverless InferenceをFastAPIで実装したいというのがあります。この記事を拡張する形です。
https://tech.fusic.co.jp/posts/2021-12-03-trying-sagemaker-serverless-inference/
なぜSageMaker Serverless Inferenceかというと、この記事に載ってたやり方がとても楽そうだったからです。ちょっとつらい体験しても、Lambdaのほうが拡張性はあると思います。
ただ、SageMaker Serverless Inferenceはにはかなり制約があり、
- 設置できるAPIがヘルスチェックの「ping」と推論の「invocations」だけ
- もともとTraiinng Jobなどで訓練したモデルをサクッとデプロイする想定で作られたものなのでここは致し方ない
- AWSが想定しているCommon Data Formats for InferenceのContent-Typeが限られ、FastAPIのContent-Typeの仕様と絶妙にマッチしない
- boto3のランタイムでは、invoke_endpointで叩くのだが、Request Bodyがバイナリをそのまま送るので、「Content-Type=multipart/form-data」のような凝ったことをすると、Boundaryから書く必要が出てくる（仕様）
- 最悪クライアント側はBoundaryから書いてもよくても、サーバー側でパースする際に、バイナリが入れ子になったJSONを扱う必要が出てきて非常に面倒くさい
ここらへんの問題点はうまくやれば解決するかもしれませんが、自分はあまりに面倒なので投げました。ただ、今回紹介する方法だと比較的サクッとできたので、覚え書きとして記録しておきます。
Common Data Formats for Inference
2点目と3点目について、AWSの公式にあるCommon Data Formats for Inferenceを見てみましょう。
JSONやCSVが使えるのはいいですが、微妙に痒いところに手が届きません。例えば、画像なら「application/x-image」ですが、「画像＋メタデータ（ハイパーパラメータやテキスト設定など）」となるとどうでしょう？ 動画ならどうでしょう？（ペイロードサイズの4MB制約があるため、大きなデータを推論したい場合はS3に置くことを考えたほうがいいかもしれません）
具体的にはこのようなデータです。
{
"message": "Some text",
"number": 123,
"image": <image_file_binary>
}
このように画像データ＋なにかのメタデータというのが実践的な形でしょう。公式で推奨されている「RecordIO」（ContentTypeでは「application/x-recordio-protobuf」）を使えば一見できるように思えます。
RecordIOの罠
Python用のRecordIOはライブラリとして用意されています。pipからもインストールできます。
https://github.com/PaddlePaddle/recordio
ただよくソースを読むと、このRecordIOはProtocol Bufferを備えていません。シリアライズの部分はPickleを使っていました。一方でAWSが推奨しているContent-Typeは「application/x-recordio-protobuf」なので、RecordIOにProtocol Bufが入ったものと考えられます。
AWSが推奨しているRecordIO（RecordIO＋ProtoBuf）への変換は、SageMaker SDKで公開されています。どんな形式かはこの記事が参考になるでしょう。
Amazon SageMakerにおけるRecordIO形式のデータの作成と読み込み
シリアライザの部分はこのソースファイルです。シリアライズの部分に使われるのが「write_numpy_to_dense_tensor」という関数ですが、この関数よく見ると、行列のレコードとベクトルのラベルしか対応していません（2022年11月時点）。あくまで構造化データでの利用想定で、画像みたいな3階テンソルは使い物にならないという歯がゆい事情があります。
シリアライザを独自に設定できないだろうか
ここで考えるのは、AWSが推奨している形以外で、以下のようなデータをさっくりAPIに渡せる方法がないかということです。
{
"message": "Some text",
"number": 123,
"image": <image_file_binary>
}
まず思いつくのは、Dict全体をこちらが用意したシリアライザでバイナリに変換して、バイナリをまるごとAPIのInputにしてしまおうということです。
ここではMessage Packを使ってみました（さすがにpickleはセキュリティ的にアレなので）。以前使ったときは複雑なNumPy配列だとうまくいかなかったこともありましたが、画像ファイルのバイナリやメタデータのリストぐらいなら全然いけます。JSONではなく、BSONを使うとかもありだと思います。
MessagePackのバイナリをFastAPIで読む
ここからようやく本題に入ります。クライアント側でMessagePackで上記のようなデータをエンコードしてPOSTすることは簡単なので、問題はサーバー側の設計になります。受け取ったバイナリをどうやってデコードするのでしょうか？
こちらの記事が大変参考になりました。
この記事だとクラスを用意してがっつりスキーマを定義しているのですが、正直JSONに毛が生えた程度の感覚で使いたいので、今回のケースだとここまでガッツリやらなくてもいいかなという気持ちはありました。もう少し簡単なやり方を探してみます。
要するにリクエストボディをダイレクトにFastAPIから取れればいいのです。これはFastAPIの公式ドキュメントにありました。
APIの引数をRequestにしてしまえば、リクエストの内容そのものにアクセスできます。例えば、ここをUploadFileとしてしまうとformとして解釈されしまい、Content-Typeの問題が発生しますが、リクエストのままアクセスしてしまえば、Content-Typeは関係なくなるそうです。
from fastapi import FastAPI, Request
import msgpack
app = FastAPI()
@app.get("/ping")
def ping():
return {"Hello": "World"}
@app.post("/invocations")
async def transformation(request: Request):
raw_binary = await request.body()
data = msgpack.unpackb(raw_binary, raw=False)
したがって、上記のようにすればいいことになります。
デモプログラム
以下のようなディレクトリ構成です。
+ src
- main.py
- Dockerfile
- requirements.txt
Dockerfile。今回はローカルでFastAPIを動かしますが、SageMakerに合わせてポートを8080に変えています。
FROM ubuntu:20.04
RUN apt-get update
ENV TZ=Asia/Tokyo
RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone
RUN apt-get install -yq --no-install-recommends python3-pip \
python3-dev \
wget \
git \
tzdata && apt-get upgrade -y && apt-get clean
RUN ln -s /usr/bin/python3 /usr/bin/python
COPY requirements.txt .
RUN pip install -U pip &&\
pip install --no-cache-dir python-multipart &&\
pip install --no-cache-dir -r requirements.txt
EXPOSE 8080
COPY ./src .
ENTRYPOINT [ "/usr/bin/python", "main.py"]
requirements.txt。FastAPIに必要なものと、MessagePack、画像を読み込むのでPILを使っています。
fastapi
uvicorn
msgpack
pillow
main.py。サーバーのメイン部分でAPIを定義しているコードです。
from fastapi import FastAPI, Request
import uvicorn
import io
import msgpack
from PIL import Image
app = FastAPI()
@app.get("/ping")
def ping():
return {"Hello": "World"}
@app.post("/invocations")
async def transformation(request: Request):
raw_binary = await request.body()
data = msgpack.unpackb(raw_binary, raw=False)
with Image.open(io.BytesIO(data["d"])) as img:
width = img.width
height = img.height
return {
"a":data["a"],
"b":data["b"],
"c":data["c"],
"width": width,
"height": height
}
if __name__ == "__main__":
uvicorn.run(
app,
host="0.0.0.0",
port=8080
)
今回は簡単に、MessagePackから展開したキーの値と、バイナリ画像（d）の解像度を表示しています。入力データのバリデーション？ナニソレオイシイノ……ってぐらい適当な実装です。
Dockerをビルド＋起動
想定はSageMaker Serverless Inferenceですが、簡単にするためにローカルのDockerでFastAPIを動かします。ローカルだともしかしたらポートを8080以外したほうがいいかもしれません。
適当にDockerをビルドして起動します。
docker build -t msgpack_fastapi .
docker run --rm -d --name mycontainer -p 80:8080 msgpack_fastapi
クライアント側での動作確認
実際にローカルにデプロイしたAPIを叩いてみましょう。ただ、ここではrequestライブラリを使わずに、urllibで叩きます。これはSageMaker SDKでBodyをバイナリから直書きしないといけないため、urllibのほうが検証には都合がいいからです。
フリー素材から「cat.jpg」としました。これをメタデータと一緒にPOSTします。
import urllib
import msgpack
data = {
"a": 123,
"b": "Hey guys, ",
"c": [
"we",
"have",
"a",
"gift",
"for",
"you"
],
}
with open("cat.jpg", "rb") as fp:
data["d"] = fp.read()
こんな感じに適当にDictを用意して、画像のバイナリを貼り付けます。MessagePackでシリアライズしてペイロードデータを作ります。
payload = msgpack.packb(data, use_bin_type=True)
あとはRequestを投げてあげると、
headers = {
"Content-Type": "application/x-msgpack",
}
req = urllib.request.Request(
"http://127.0.0.1/invocations",
payload, headers
)
with urllib.request.urlopen(req) as res:
print(res.read())
じゃーん。このとおりうまくいきました！ PILによる画像の解像度解析もうまく行っていますね。
b'{"a":123,"b":"Hey guys, ","c":["we","have","a","gift","for","you"],"width":2067,"height":1163}'
当初計画していたように「とりあえず何でもかんでも1個のMessagePackのバイナリに打ち込んで、サーバー側で展開する」というのができました。MessagePackの仕様上、複雑なNumPy配列だとハマるかもしれませんが、デプロイの制約がある環境下でもFastAPIの汎用性が上がるのではないかと思います。
（バリデーションってなんだっけ…）
Shikoan's ML Blogの中の人が運営しているサークル「じゅ～しぃ～すくりぷと」の本のご案内
技術書コーナー
北海道の駅巡りコーナー