はじめに

注意：この記事は2019/11/22時点での情報です。

LIFULLでは機械学習のサービング基盤にSageMakerを利用しています。
一方、LIFULLでは、データがBigQueryに集約される方針で、実際に集約されつつあるので、BQのデータを手軽に渡せるAutoML Tablesは簡単に試せて便利です。

AutoML Tablesにモデルのエクスポート機能がついたので、SageMakerのような別の環境で動かすできるようになりました。
エクスポート機能の使い方は公式に書いてあります。

LIFULLではメインサービスがAWSで動いているためデプロイはAWSである方が都合が良いです。
ですので今回は、AutoML Tables公式のDockerイメージをSageMakerで動くようにしてみました。

やったこと

要件の確認

SageMakerで動くDockerイメージの要件は最低以下を満たす必要があります。

エンドポイント
- /pingにGETされた時、ステータス200を返すこと
- /invocationsにPOSTされた時にモデルへの入力を受け取り60秒以内に予測結果を返すこと
モデルの読み込み場所
- /opt/ml/model以下にモデルが配置されるのでそこからモデルを読み取れるようにすること
port
- 8080番でリクエストを受け付けること

以上を満たせば、とりあえず動くようになります。

イメージの使い方の確認

AutoML Tablesで公式が提供するイメージの使い方を確認します。
先ほどの公式を参考にします。

イメージはgcr.io/cloud-automl-tables-public/model_serverを使うようなのでとりあえずpullします。

ボリュームは-v $(pwd)/model-export/tbl/new_folder_name:/models/default/0000001とあるので、エクスポートしてダウンロードしたモデルをmodels/default/000001にセットすればいいようです。
(あとでわかりますが、ここの000001は数字であればこの数字でなくても動きます。)

ポートは-p 8080:8080なのでそのまま利用できそうです。

リクエストはcurl -X POST --data @/tmp/request.json http://localhost:8080/predictとあるように
/predictに対して送るようです。

修正する点の確認

ポートはそのままで良さそう。
モデルの読み込み場所は変える必要がある。
エンドポイントは/predictを/invocationsに変更し/pingを生やす必要がある。

具体的に調べる

やることがわかったので一つずつ対処していきます。
イメージのソースを見つけられなかったのでとりあえず動かして挙動を確認していきます。

とりあえず、公式通りに起動してentrypointを確認すると

/usr/bin/tf_serving_entrypoint.sh

を起動していることがわかります。(起動せずに確認する方法はあるのだろうか？)

中に入って確認していきます。

docker run --rm -it --entrypoint bash gcr.io/cloud-automl-tables-public/model_server

/usr/bin/tf_serving_entrypoint.shから追っていった結果、
/prediction_server.pyや/translate.pyが重要であることがわかりました。

/usr/bin/tf_serving_entrypoint.shではtensorflow_model_serverを起動し/prediction_server.pyを実行していることがわかります。

tensorflow_model_server --port=8500 --rest_api_port=8501 --model_name=${MODEL_NAME} --model_base_path=${MODEL_BASE_PATH}/${MODEL_NAME} "$@"

ここでtensorflow_model_serverにモデルのパスを指定しています。
都合よく環境変数で指定されているのでここは弄らず、あとで環境変数で場所を/opt/ml/modelになるように変更します。

/prediction_server.pyではtornadeを起動してリクエストとレスポンスを加工する処理を挟んでいるようです。これも都合がいいのでここで/pingと/invocationsを作っていきます。

変更点した点を書いておきます。

prediction_server.py


class PingHandler(tornado.web.RequestHandler):

    def get(self):
        self.set_header("Content-Type", "application/json")
        self.set_status(200)
        self.write({'status': 'ok'})

def make_app():
    return tornado.web.Application([
        (r"/invocations", PredictHandler),
        (r"/ping", PingHandler),
    ])

これでエンドポイントは要件を満たしました。🎉

あと、私が作業した時点のイメージ(tag:11_21_2019)では/translate.pyにバグがあったので修正した点を書いておきます。

translate.py

import os
...

def user_request_to_tf_serving_request(body):
    ...

    model_spec = model_pb2.ModelSpec(
        name=os.environ['MODEL_NAME'], signature_name="serving_default")
    ...

def tf_serving_response_to_user_response(resp, batch_size):
    ...

    values = [{} for _ in range(batch_size)]
    for key, value in outputs.items():
        for i, v in enumerate(value.tolist()):
            values[i][key] = v

name=os.environ['MODEL_NAME']の部分はモデルの場所を変更する都合で元のname="default"だと読み込めなくなるためです。
あと、元々はvalues = [{} for _ in range(batch_size)]の部分はvalues = [{}] * batch_sizeと書いてあり、複数個のデータをリクエストした場合に最後の結果で全て上書きされてしまう状態でしたので修正しました。

Dockerイメージのビルド

修正を反映したイメージを作成します。

作業ディレクトリ

./
├ /override
│ ├ prediction_server.py
│ └ translate.py
└ Dockerfile

/overrideの中に先ほど修正した内容のファイルを入れました。

FROM gcr.io/cloud-automl-tables-public/model_server:11_21_2019

ENV LC_CTYPE="C.UTF-8"

ENV MODEL_BASE_PATH="/opt/ml" \
    MODEL_NAME="model" \
    PORT="8080"

COPY ./override/ /

ENTRYPOINT [ "/usr/bin/tf_serving_entrypoint.sh" ]

ここで先ほど保留にした、モデルディレクトリの要件も満たしました。🎉
イメージのタグは以下で確認し固定しました。
https://console.cloud.google.com/gcr/images/cloud-automl-tables-public/GLOBAL/model_server

デプロイして確認

作ったイメージをECRにpushしてSageMakerで確認していきます。

ここで重要になるのはs3に置くモデルの圧縮ファイルです。
/opt/ml/model直下に置くモデルディレクトリは数字でないといけないので、AutoML Tablesからダウンロードしたモデルをリネームする必要があります。

具体的に言うと、以下の感じ。

/1
├ /assets
├ /variables
└ save_mode.pb

あとはこれをtar zcvf ./1 model.tar.gzしてs3に置きそれをSageMakerでモデルに指定してエンドポイントを立てます。

エンドポイントを正常に起動できればミッションコンプリートです。お疲れ様でした。