SageMaker Serverless Inference + PyTorch + Lambda + API Gateway

こんにちは!!!クライアントエンジニアの小林です。

今回はSageMaker Serverless Inferenceを使用してPyTorchな推論モデルをサーバーレスで動かしていきます。

作業環境

・windows 10
・visual studio code
・python 3.9.12

概要

SageMaker Serverless Inferenceを使用してPyTorchな推論モデルをサーバーレスで動かしていきます。

例としてCRAFTを扱います。
CRAFTはざっくりいうと文字領域検出モデルです。
あくまでサーバーレス推論をすることがメインなのでモデルの詳細は省きます。

アーキテクチャはこんな感じです。
あれだけサービスあるのになぜ公式で作図ツールを用意していないのだ。

1. 画像を送信
2. APIGatewayでLambdaで送りつける
3. Lambdaからサーバーレス推論の呼び出し
4. 推論結果をLambdaで受け取り
5. APIGatewayで画像の送信元に結果を投げつける

注意点として筆者はAWSと初対面な上にあまり仲良くありません。
サーバーレス推論をするために触った程度です。
そのため「動けばええやろ」精神で、各項目の設定は必要最低限となっています。

SageMaker Serverless Inferenceとは

SageMakerはAWSが提供している機械学習プラットフォームです。

その中のサービスの1つServerless Inferenceがサーバーレス推論というわけです。

SageMaker Serverless Inferenceの料金プラン

料金は推論にかかった時間とデータの送受信量に依存します。
おそらくモデルの起動時間も含まれる気がします。どうせケチだから。

流れ

こんな感じで進めます。
流石に機械学習専用サービスなだけあってフロー自体もシンプルです。

1. モデルと推論コードの準備
2. ファイルのアップロード
3. モデルのデプロイ
4. Lambdaの構築
5. API Gatewayの構築
6. 動作テスト

モデルと推論コードの準備

SageMakerで動作するコード一式を用意します。

モデル実装

基本的にはローカルで動かしているものを流用できます。
例としてCRAFTの実装コードを載せています。

注意点として2023/03/06時点のSageMaker for PyTorchのPythonバージョンは3.8です。
型ヒントに旧仕様であるtypingモジュールを使う必要があったり、dictやlistの結合にoperatorが使えなかったりと、3.8以降のバージョンを使っている場合はダウングレード対応が発生します。
その他にもGPU推論が非対応だったりします。

craft.py	HatenaBlogSample/blog_craft/code/craft.py at sagemaker-serverless-inference/main · spark-kobayashi-arata/HatenaBlogSample · GitHub
craft_base.py	HatenaBlogSample/blog_craft/code/craft_base.py at sagemaker-serverless-inference/main · spark-kobayashi-arata/HatenaBlogSample · GitHub
gaussian.py	HatenaBlogSample/blog_craft/code/gaussian.py at sagemaker-serverless-inference/main · spark-kobayashi-arata/HatenaBlogSample · GitHub
vgg16_bn.py	HatenaBlogSample/blog_craft/code/vgg16_bn.py at sagemaker-serverless-inference/main · spark-kobayashi-arata/HatenaBlogSample · GitHub

追加のパッケージ指定

SageMakerに標準インストールされていないパッケージはpip installでお馴染みなrequiments.txtに指定する必要があります。

CRAFTは標準モジュールで事足りるため、実際のデプロイには含めませんが、分かりやすいように画像では配置した状態で撮っています。

学習済みモデルの用意

サーバーレス推論のお試し用としてCRAFTの学習済みモデルを公開しています。
公開している学習済みモデルはCC BY-NC 4.0を適用しています。
ImageNetを元に学習したvgg16モデルを使用している関係で非営利にしています。

SageMaker PyTorch Model Server

HatenaBlogSample/blog_craft/code/inference.py at sagemaker-serverless-inference/main · spark-kobayashi-arata/HatenaBlogSample · GitHub

PyTorch LightningやHugging FaceのようにSageMakerにはSageMakerのフレームワークが存在します。

SageMakerのフレームワークに対応した関数を定義していればそれが実行され、未定義の場合はデフォルトの関数が実行されるといった、挙動的にはC系のオーバーライドと同様です。

詳しくはこちらのドキュメントに載っています。

ファイルを圧縮

先ほど用意した学習済みモデルですが、pth形式でのアップロードはできません。
これをtar.gz形式で圧縮する必要があります。

そして圧縮時に先ほど実装したinference.pyやcraft.py、requiments.txtなどを含めることでデプロイ時の作業が減ります。

作業自体は減るのですがコーディングミスや不具合修正をするたびに再圧縮をする必要が発生します。
そのためpython初心者にはむしろ悪手だったりします。

今回はデプロイ設定を最低限で行いたいので圧縮時に必要なファイルをすべて突っ込んでいますが、慣れない方はデプロイ時に行う方法も試してみてください。
今回は紹介しませんが先ほどのドキュメントになんとなく載っているので、たぶんそれで理解できます。

前置き長くなりました。
rootフォルダに学習済みモデルを配置、rootフォルダにcodeフォルダを作成します。
codeフォルダにはモデル実装ファイルやinference.py、requiments.txtを突っ込みます。

rootフォルダを右クリックで圧縮なのですが、操作手順が面倒なので適当にスクリプトを作ります。
メモ帳にコードを貼り付けて、いくつかの変数を変更した後に*.bat拡張子で保存、実行することで動作するはずです。
HatenaBlogSample/pack-blog-craft.bat at sagemaker-serverless-inference/main · spark-kobayashi-arata/HatenaBlogSample · GitHub

実行するとROOT_DIRにtar.gz形式のファイルができていると思います。

ファイルのアップロード

tar.gz形式に圧縮したファイルをS3にアップロードします。

バケットの作成

S3を開きます。

『バケットを作成』を選択します。

バケット名の入力と、リージョン選択をして作成します。

作成が完了すると、一覧に追加されます。

アップロード

作成したバケットを開いて、『アップロード』を選択します。

『ファイルを追加』から先ほど圧縮したファイルを追加します。

『アップロード』を選択します。

数秒待ちます。

無事にアップロードが完了しました。

モデルのデプロイ

S3にアップロードしたファイルを元にSageMakerノートブックインスタンスからサーバーレス推論モデルをデプロイします。
デプロイが完了したらそのモデルが正常に動くか軽く動作テストをします。

ノートブックインスタンスの作成

SageMakerを開きます。

『ノートブックインスタンスの作成』を選択します。

インスタンス名を入力します。

『ノートブックインスタンスの作成』を選択します。

作成に成功したらステータスが緑色になるまで数分待ちます。

ステータスが緑色になったら『JupyterLabを開く』を選択します。

モデルのデプロイ

https://github.com/spark-kobayashi-arata/HatenaBlogSample/blob/sagemaker-serverless-inference/main/aws/deply_model.ipynb
上記のコードを貼り付けて実行します。
model_data引数にはS3にアップロードしたファイルのS3 URIを指定する必要があります。

S3 URIはアップロードしたファイルのプロパティから確認できます。

実行したら数分待ちます。
！マークが末尾に出たらデプロイ完了の合図です。

Notebookで動作テスト

https://github.com/spark-kobayashi-arata/HatenaBlogSample/blob/sagemaker-serverless-inference/main/aws/test_inference.ipynb
上記のコードを新しいセルに貼り付けます。
CRAFT_ENDPOINTにはデプロイしたモデルのエンドポイント名を指定します。

エンドポイント名は『推論』>『エンドポイント』から確認できます。
これをCRAFT_ENDPOINTに指定します。

実行したら数分後に結果が返ってきます。
画像は公開している学習済みモデルに含まれるsample2.pngを使用しています。

デプロイが完了したらインスタンスを停止します。

Lambdaの構築

主にサーバーレス推論の呼び出しを行います。

関数の作成

Lambdaを開きます。

『関数の作成』を選択します。

関数名を入力、ランタイムは『Python 3.9』を選択します。

『関数の作成』を選択します。

モジュールの追加

素のLambda for Pythonでは主要モジュールであるnumpyやboto3などがインストールされていません。
そのため不足しているモジュールをインストールするためにレイヤーを使用します。

リポジトリで公開されているものを使用しました。感謝。
なんでこういうの公式で用意しないんでしょうか。Amazonさんの怠慢ですね。

『レイヤーの追加』を選択します。

『ARNを指定』を選択します。

arn:aws:lambda:ap-northeast-1:770693421928:layer:Klayers-p39-numpy:10
arn:aws:lambda:ap-northeast-1:770693421928:layer:Klayers-p39-boto3:11
『ARNを入力』>『検証』>『追加』の順番です。

これでnumpyとboto3が使えるようになりました。

コードの実装

https://github.com/spark-kobayashi-arata/HatenaBlogSample/blob/sagemaker-serverless-inference/main/aws/lambda_function.py
上記のコードをlambda_function.pyに貼り付けます。

コードの変更をしたら『Deploy』を選択します。

タイムアウトの変更

先ほども少し触れましたがサーバーレス推論はGPU非対応です。
それ故に推論には数秒要します。

Lambdaのデフォルトタイムアウトは3秒です。
つまり余裕でタイムアウトします。

そのためタイムアウトを30秒に変更します。
Lambda自体は30秒以上待ってくれるのですが、後述するAPIGatewayが最大で29秒しか待ってくれないため30秒です。

モデル読込を含むと30秒超えることもあるので地味に厳しい縛りです。

『設定』>『一般設定』>『編集』を選択します。

タイムアウトを『30秒』に変更後『保存』を選択します。

アクセス権限の付与

LambdaからSageMakerにアクセスするためには権限付与が必要です。
付与する権限は『推論モデルの実行のみ』でいいのですが、面倒なので適当にフルアクセス権限を付与しちゃいます。

『設定』>『アクセス権限』からロール名のURLを開きます。

『許可を追加』>『ポリシーをアタッチ』を選択します。

『sagemaker』と検索します。

『AmazonSageMakerFullAccess』にチェックを入れます。

『許可を追加』を選択します。

アクセス権限を付与できました。

アクセス権限からSageMakerに関することが追記されているのが確認できます。

アクセス権限の付与（SageMaker）

Lambdaと同様にSageMakerにもアクセス権限を付与する必要があります。
たぶん。。。
というのも元から追加されていたのか、自分で追加したのか、もはや覚えてないのです。

『推論』>『モデル』からデプロイしたモデルを選択します。

IAM ロール ARNのURLを選択します。

Lambdaと同様な手順で『AmazonS3FullAccess』『AWSLambdaExecute』『AWSLambdaRole』を追加します。

API Gatewayの構築

リクエストとレスポンスを可能にするREST APIを構築します。

REST APIの作成

API Gatewayを開きます。

『APIを作成』を選択します。

REST APIの『構築』を選択します。

API nameを入力後、『Create API』を選択します。

Lambda関数と紐づけ

『Actions』>『Create Method』を選択します。

『POST』を選択します。

チェックマークをクリックします。

lambda関数名を入力します。

『Save』を選択します。

Lambdaにアクセス権限を付与してもいいか聞かれるので『OK』を選択します。

これでAPIを叩くとLambda関数が実行されるようになりました。

APIキーの作成と適用

アクセス制限を設けるためにAPIキーによる認証を設定します。

『API Keys』を選択します。

『Actions』>『Create API key』を選択します。

Nameを入力後、『Save』を選択します。

作成が完了しました。
API keyはコピーしておいてください。

POST設定に戻ります。
『Method Request』を選択します。

API Key Requiredの横の鉛筆マークを選択します。

『true』を選択します。

チェックマークを選択します。

これでリクエスト時にAPIキーが要求されるようになりました。

デプロイ

『Actions』>『Deploy API』を選択します。

『Deployment stage』>『[New Stage]』を選択します。

Stage nameを入力後、『Deploy』を選択します。

デプロイが完了しました。
Invoke URLはコピーしておいてください。

APIキーの紐づけ

『Usage Plans』を選択します。

『Create』を選択します。
Nameを入力後、『Enable throttling』と『Enable quota』のチェックを外して制限を解除します。
注意点として不特定多数の利用者を想定した場合は外さずに設定しないと危ないです。

『Next』を選択します。

『Add API Stage』を選択します。

APIキーと紐づけするAPIを選択します。

APIデプロイ時に作成したStageを選択します。

チェックボタンを選択します。

『Done』を選択します。

『Add API Key to Usage Plan』を選択します。

紐づけするAPIキーを選択します。

チェックマークを選択します。

『Done』を選択します。

紐づけが完了しました。

ローカルから動作テスト

HatenaBlogSample/request.py at sagemaker-serverless-inference/main · spark-kobayashi-arata/HatenaBlogSample · GitHub
上記のコードを実行します。

URLとAPIキーは先ほどコピーしたものを指定してください。

できた。ｖいえーいｖ

おわり!!!

お疲れさまでした!!!

一番疲れたのは筆者ですが。
数ヶ月後の自分用に画像マシマシで作ったはいいものの後半になるにつれて飽きがヤバかったです。

Serverless Inferenceはサーバーを知らぬ人でも構築できる良きサービスです。
ただ、GPU非対応だったり、複数モデルを同時に扱えなかったり、すぐにスリープ状態に入ったりと実運用するには不向きなんだろうなと感じました。
試験的なモデルの共有や簡易的なサンプル、推論頻度が低いモデルなどには最適なのでしょう。

とはいえ最近(2023/03/12)のAI需要的にAmazonさんも優先的に改善していくべきサービスなのではないかなと予想していたりします。
せめてGPU推論はできるようにしてほしいです。

参考サイト

API Gatewayとlambdaを使ってみる - Qiita
pythonのrequestsでAPIGateWayをAPIKeyを使って呼び出す - Qiita
API Gateway と Lambda で Hello World! - WhiteNote
【小ネタ】Lambda関数に画像データを渡す | DevelopersIO