情報系院生のノート

情報技術系全般,自分用メモを公開してます。

セグメンテーションするpytorch機械学習モデルをAWSへデプロイ(API Gateway, Lambda, ECR)

はじめに

pythonを使った機械学習モデルは巷でよく見ます。今回はpytorchを使って学習させたモデルをAWSにサクッとデプロイしたいと思います。

この記事をめちゃめちゃ参考にしました。良記事です。

tech.aptpod.co.jp

今回は上の記事の丸パクリで、以下のような構成にします。

Frontend --- API Gateway --- Lambda --- ECR

独り言

多くの人がアクセスするような環境だと、EC2にWebサーバ建ててロードバランサー噛ませるのが一番良いと思いますが、 小さな規模やプロトタイピングなどではそのような環境を作るのは大変です。

サーバレスでAPIを作る手っ取り早い方法としてLambdaがあります。しかしパッケージ含めて250MB制約があり厳しい。

そんな中2020/12/04にLambdaのコンテナをサポートが発表されました。なんと10GBまでデプロイ可能です。 また、ローカルで実行できるLambdaのRuntime APIツールも提供されました。

Lambdaがローカルでデバッグできるなんて感激です!

ガートナーによれば、AIは幻滅期に入ったとされ、次に啓蒙期、生産性の安定期ときます。 つまり、これからは機械学習モデルの社会実装が進む頃合いです。

これからどんどん機械学習モデルのデプロイしやすい環境が整備されていくんでしょうね。

構成

今回は画像をセグメンテーションするモデルを動かします。

以下に今回構築するAWSの構成の詳細を示します。

f:id:tontainoti:20210302215341p:plain
Segmentation endpoint overview

コンテナイメージの作成

Lambda上で動かすには、 Lambda Runtime Interface Clientsを入れなければいけません。

こちら1の公式が提供しているイメージには既に必要なコンポーネントが含まれていると思われます。(要確認)

私はマルチステージビルドしたかったので使ってません。 Dockerfileのダイエットについてはこちら2を参考にしています。

フォルダ構成

├── app
│   ├── app.py
│   ├── modelとかcheckpointとか
├── Dockerfile
├── entry.sh
├── requirements.txt
  • ローカルの/appフォルダにapp.pyを作成し、そこのhandler関数が呼ばれるように書きました。
  • modelは別にフォルダを作って格納しておきます。
  • Dockerfileは次の節で説明します。
  • entry.shは、これは公式チュートリアル3にて掲載されていました。
    • ローカルとLambda上とで条件分岐しているみたいです。以下にentry.shを置いておきます。
#!/bin/sh
if [ -z "${AWS_LAMBDA_RUNTIME_API}" ]; then
    exec /usr/bin/aws-lambda-rie /usr/local/bin/python -m awslambdaric $1
else
    exec /usr/local/bin/python -m awslambdaric $1
fi
  • requrement.txtに必要なパッケージを記述しておきます。PipfileとかでもOKです。(その場合はDockerfileの書き換えが必要ですが)

Dockerfile

以下に私が構成したDockerfileを示します。 このサイト[^3]を参考にして、マルチステージビルドしています。

最初にFROM python:3.7.9 as buildとして、ビルド用のイメージを引いてきて、生成するイメージに使うコンテナは次のようにFROM python:3.7.9-slim-stretch as production slim-stretchにします。 また、COPYでbuildからpipでインストールしたパッケージを持って来ています。

# Define function directory
ARG FUNCTION_DIR="/function" 

FROM python:3.7.9 as build

# Install aws-lambda-cpp build dependencies
RUN apt-get update \
  && apt-get install -y \
  g++ \
  make \
  cmake \
  unzip \
  libcurl4-openssl-dev \
  libsm6 \
  libxrender1 \
  libxtst6 \
  && apt-get autoremove -y \
  && apt-get clean \
  && rm -rf /var/lib/apt/lists/*
 
# Include global arg in this stage of the build
ARG FUNCTION_DIR

# Install the runtime interface client & other python package
COPY requirements.txt  /
RUN pip install --upgrade pip \  
    && pip install awslambdaric \
    && pip --no-cache-dir install -r requirements.txt  \
    && rm -rf ~/.cache

# production stageの定義
FROM python:3.7.9-slim-stretch as production

# build stageでinstallされたpackage群を丸ごと持ってくる
COPY --from=build /usr/local/lib/python3.7/site-packages /usr/local/lib/python3.7/site-packages

ARG FUNCTION_DIR
# Create function directory
RUN mkdir -p ${FUNCTION_DIR} && mkdir -p ${FUNCTION_DIR}/model/

# Copy function code
COPY app/ ${FUNCTION_DIR}/

# Set working directory to function root directory
WORKDIR ${FUNCTION_DIR}

# (optional) for TEST
ADD https://github.com/aws/aws-lambda-runtime-interface-emulator/releases/latest/download/aws-lambda-rie /usr/bin/aws-lambda-rie
RUN chmod 755 /usr/bin/aws-lambda-rie
COPY entry.sh /
ENTRYPOINT [ "/entry.sh" ]

CMD [ "app.handler" ]

POST, Responseスキーム

POSTはmultipart/form-dataです。

{
    "img": base64 utf-8エンコード画像
}

ResponseははstatusCodeとbodyを含む必要があるみたいです。 私は2つの画像を返したかったので、bodyをjsonにして、その中の要素にbase64の画像を突っ込んでます。

Response

{
    "isBase64Encoded": false,
    "headers": {
        "Content-Type": "application/json"
    },
    "statusCode": 200,
    "body": "{
        "img1": base64 utf-8エンコード画像,
        "img2": base64 utf-8エンコード画像
    }"
}

app.py

以下にapp.pyの主要部を示します。

multipartのデコードに苦労しました…。

import io
import base64
import json

def b64toPIL(b64img):
    im_bytes = base64.b64decode(b64img)
    im_file = io.BytesIO(im_bytes)
    img = Image.open(im_file)
    return img


def PILtob64(img):
    im_file = io.BytesIO()
    img.save(im_file, format="PNG", quality=100)
    im_bytes = im_file.getvalue()
    im_b64 = base64.b64encode(im_bytes).decode('utf-8')
    return im_b64


def predict(img):
# 機械学習モデルのロード
# 推論
# return 画像
...

def parse_multipart_from_api_gateway(event):
    c_type, c_data = parse_header(event['headers']['Content-Type'])
    encoded_string = event['body'].encode('utf-8')
    c_data['boundary'] = bytes(c_data['boundary'], "utf-8")
    data_dict = parse_multipart(io.BytesIO(encoded_string), c_data)

    # 整形
    formatted_dict = {}
    for k, v in data_dict.items():
        formatted_dict[k] = v[0]

    return formatted_dict


def handler(event, context):

    response = {
        "isBase64Encoded": False,
        "headers": {
            'Content-Type': 'application/json'
        },
        "statusCode": 200,
        "body": ""
    }

    # 'multipart/form-data'をデコード
    data_dict = parse_multipart_from_api_gateway(event)

    img = b64toPIL(data_dict['img'])

    # 推論
    pred_img1, pred_img2 = predict(img)

    body_dict = {
        "img1": PILtob64(pred_img1),
        "img2": PILtob64(pred_img2)
    }

    response["body"] = json.dumps(body_dict)

    return response

ローカルでテスト

早速ローカルで動作確認します。

  • build
docker build -t segmentation_model:latest .
  • 実行
docker run  -p 9000:8080 --entrypoint /usr/bin/aws-lambda-rie  --name serverless  --rm segmentation_model:latest  /usr/local/bin/python -m awslambdaric app.handler
  • POSTしてみる

ローカルでPOSTしたいのですが、multipart/form-dataをうまく送ることができませんでした。

以下の感じで送ればevent['headers']でヘッダ情報を取得できますが、boundary属性がありません。

endpoint = "http://localhost:9000/2015-03-31/functions/function/invocations"
response = requests.post(endpoint, json={'body': multipart_string, 'headers': {'Content-Type': content_type}})

ここはnc -l ポート番号の部分にPOSTして、その内容をコピペしてPOSTするとかしかなさそうですね…

よいデバッグ方法があれば知りたいです。

Lmabdaのロギング

Lambdaから詳細なエラーログがほしいときがあると思います。 app.pyに以下のコードを追加すると、詳細なログを出してくれます。

import logging

logger = logging.getLogger()

formatter = logging.Formatter(
    '[%(levelname)s]\t%(asctime)s.%(msecs)dZ\t%(aws_request_id)s\t%(filename)s\t%(funcName)s\t%(lineno)d\t%(message)s\n',
    '%Y-%m-%dT%H:%M:%S')
for handler in logger.handlers:
    handler.setFormatter(formatter)

任意のログを出したいとき

logger.info("any log")

ECRにプッシュ

AWSコンソールに入り、ECRでプライベートレジストリを作成します。

以下記事を参考にaws-cliでECRにログインします。-v ~/.aws:/root/.awsでマウントしていることに注意です。

zenn.dev

docker imageのtag名を変更します

docker tag segmentation_model:latest   {AWS_ACCOUNT_NO}.dkr.ecr.ap-northeast-1.amazonaws.com/{REPO_NAME}:latest

pushします

docker push {AWS_ACCOUNT_NO}.dkr.ecr.ap-northeast-1.amazonaws.com/{REPO_NAME}:latest

Lambda関数作成

Lambda関数を作成します。 関数の作成時、「コンテナイメージ」を選択するとECRにコミットされているコンテナイメージを選択します。

API Gatewayの作成

Lambdaのデザイナーからトリガーを追加でAPI Gatewayを追加します。 ここではHTTPのAPI Gatewayの作り方を紹介します。(HTTPの方が低コスト)

HTTPではJWT認証が必須なので、Cognitoを使います。 面倒だな…と思った方はRESTでAPIキー認証のAPI Gatewayを建ててください。すぐできます。

ほぼこの良記事を参考にします。

qiita.com

注意点として、ユーザープールを作成するときに、シークレットキーのチェックボックスは外してください。 aws-cliからアクセス出来なくなります。

これでAPI Gatewayも建てることができました。

POSTテスト

headersにaws cognito-idp admin-initiate-authコマンドで取得したIdTokenを入れます。

"headers": {
        "Authorization": IdToken
    },

あとエンドポイントの部分をAWSに変えれば、ローカルでテストしたソースがそのまま使えます!

感想

僕はいろいろ手間取って構築に2日くらいかかったので全然サクッとは行きませんでしたが、 こんだけでサーバレスなAPIが完成します。

素晴らしいですね。

参考