Amazon SageMaker BYOCとAWS Step Functionsで推論エンドポイントを構成する

はじめに

この記事はAWS Summit Online 2020で登壇した内容ですが、テックブログのため再度まとめ直した記事です。

目的と課題

Baseconnectでは500万件ほどの企業情報を保持しており、その内140万件ほどをサービスで提供しております。 その中から利用者が欲しい情報を探し出すために、検索軸や検索性能を改善し続けています。 ただ、検索で探し出すのではなく、利用者自身が保持しているリストと突合させて情報を補完（名寄せ）したいという要望があり、それに応えるためにAIを補助的に入れることを目的としました。

進めるにあたり、以下のことを意識して進めました。

  • ゴールはどのような状態か
  • 精度の基準はどうするか
  • どのような手法で機械学習を行うか

取り組み

ゴールはどのような状態か

今回は既にサービスで提供している部分をAIでリプレースすることを目指しました。

これまでに、利用者の手持ちの営業先リストがアップロードされた際に140万件の企業情報との名寄せを実装済みでした。
そのため、ゴールとしては利用者としては意識すること無く、バックエンドをAIへリプレースすることで利便性の向上を目指しました。
精度の基準はどうするか

今後入力項目や名寄せ対象が増加されることも考えられましたが、まずは既存の仕組みでの精度よりも向上させることを基準としました。

※入力データの状態により異なるため、サンプルデータを用意 ※取り組み前の時点で既存の機能でサンプルデータを用いた精度は名寄せ率が55%ほど

どのような手法で機械学習を行うか

機械学習として教師あり、教師なし、強化学習とありますが、大きくはどれで（もしくはどの組み合わせで）取り組むかについて検討しました。（以下、検討の一部）

  • 教師あり
    • 名寄せ先となるデータが140万件あり、それらは日々増加や更新が行われている
    • 正解が140万件となるのは現実的では無さそう
  • 教師なし
    • クラスタリングを活用するのは一定の効果がありそう
    • if文などのロジックが増えてきたら、最初からロジックで作るのとあまり変わらない
  • 強化学習
    • 名寄せした結果が求めているデータに近いかどうか判断をすることで、名寄せ先のデータに変動があっても耐えられる仕組みに出来そう

ただ、ブレストの中で教師なし学習も今回の目的の中で活用は出来そうと話はしていました。 クラスタリングなども間に含めて多段的にAIを使うことも考えましたが、まずはシンプルに強化学習でどこまで精度が出せるのか取り組んでみました。

取り組み内容

具体的に取り組んだ内容です。

全体構成

Amazon SageMakerをコアに使うことを考え、そのフローをどうするかとして全体を設計しました。

当初推論はLambdaで構成しようとしてましたが、メモリの制限や起動に時間がかかる（プロビジョニングにすると高くなる）ため、Amazon SageMakerの利用を選択しました。 Amazon SageMakerはモデルの学習や推論エンドポイントの提供、またマルチエンドポイントで利用者毎のモデル提供など将来的に展開を考えられそうでした。

また、名寄せ先となるデータは変動が考えられるため、今回はデータベースを直接参照するという手段を取りました。 なお、学習・推論と検証する中でデータベースのパフォーマンスも負荷がかかりすぎないか監視しながら実施しました。

学習とDocker Imageの作成

Amazon SageMakerで推論エンドポイントを利用するために、そのコンテナ構造に合わせています。

◆ディレクトリ構成
project-root/
  ├ opt/
  │  └ ml/
  │      ├ code/
  │      │  ├ agent.py
  │      │  ├ environment.py
  │      │  ├ train.py
  │      │  ├ eval.py
  │      │  ├ handler.py … Amazon SageMakerの推論API用
  │      │  ├ nginx.conf … Amazon SageMakerの推論API用
  │      │  └ serve … Amazon SageMakerの推論API用
  │      │
  │      ├ failure/
  │      ├ input/
  │      │  ├ config/
  │      │  ├ data/ … 学習用データ
  │      │  └ eval/ … 評価用データ
  │      ├ model/ … 学習済みモデル
  │      └ output/ … 評価結果
  ├ build_and_push.sh
  ├ buildspec.yml
  ├ Dockerfile
  └ requirements.txt

Amazon SageMakerにもNotebookがあり、そこでもコーディングと実行は出来るのですが、今回は先に開発端末で試していたこともあり開発端末にて学習／評価を実行し、GitHubのPullRequestでAWS CodePipelineを利用してDocker Imageを作成（AWS CodeBuildでイメージを作成し、Amazon Elastic Container Repositoryへプッシュ）するようにしました。

Docker Imageの作成はサンプルを参考にしましたが、一部変更しました。

  • DockerHubの tensorflow/tensorflow:2.0.0-py3 を利用
    • その当時の新しいバージョンを利用したかったため
  • 学習済みモデルもDocker Imageに含める
    • パイプラインにAmazon SageMakerでトレーニングも含めることも出来たのですが、ひとまず開発端末でモデル作成まで行っていたため
  • Docker Image内にDBの接続情報も環境変数で定義
    • 秘匿情報を限定的に参照としたく、暫定的に進めたかったため
  • Docker Image内でもPythonライブラリをセットアップ
    • 後述のhttp serverもサンプルとは異なる方式を取ったため
  • tensorflow-model-server は利用しない
    • 深層強化学習で作成したモデルがtensorflow-model-serverでそのまま使えなかったため
    • 学習済みモデルのServingにはPythonのHTTServerを利用
Docker Imageの動作確認

まず開発端末で動作を確認したかったため、以下のコマンドでDocker Imageの作成とコンテナの起動を行いました。

$ docker build -t demo . --build-arg DB_HOST=”xxxxxx” --build-arg DB_USER=”xxxxxx” --build-arg DB_PASS=”xxxxxx”
$ docker run -p 8009:8009 -it demo serve

ECRの設定まで済んでいる場合は以下でも出来出来ます。

$ ./build_and_push.sh
$ docker run -it 8009:8009 -it {account}.dkr.ecr.{region}.amazonaws.com/demo:latest serve

動作イメージは以下のようになっています。サンプルで作成するとエンドポイントは決まるので、そこに対してPostmanなどで確認しました。

AWS Step Functions Data Science SDKの利用

Amazon SageMakerで推論エンドポイントを更新するまでを自動化するためにAWS CodeBuildからAWS Step Functionsをトリガーするようにパイプラインを構成しました。 そして、AWS Step Functionsを構成するためにAWS Step Functions Data Science SDKを利用しました。 SDKはPythonで利用でき、notebook上で記述と実行を試しながら進めました。

AWS Step Functionsの構成としては以下となっています。

各ステップで気にした点は以下です。

[モデル] ・独自Docker Imageでは学習済みなのでトレーニングは行いませんでした ・データベースへの接続が必要なのでVPC内で構成する必要がありました ・上記２点を考慮して構成する場合はSDKではダメでAWS Lambdaを構成に含める必要がありました

[エンドポイント設定] ・モデルの指定 ・インスタンスタイプを ml.t2.medium（最小）で構成

[エンドポイント] ・既存のエンドポイントを新しいエンドポイント設定で更新するように構成

[その他] ・SDKではAWS Step FunctionsのCreate or Updateは出来なかったので別々に構成しました ・AWS Step FunctionsではエンドポイントのCreate or Updateは出来なかったので別々に構成しました ・SDKのドキュメントはstableだと詳細の記載が無かったのでlatestで見ていました

推論エンドポイントの確認

まず、Amazon SageMakerのエンドポイントはsig-v4で保護されており、通常のREST APIのように呼び出すためには手順が必要なため、boto3でクライアントを構成して推論エンドポイントの確認をするようにしました。 https://docs.aws.amazon.com/sagemaker/latest/APIReference/API_runtime_InvokeEndpoint.html https://docs.aws.amazon.com/AmazonS3/latest/API/sig-v4-authenticating-requests.html

aws cliに同梱されているboto3で動作確認した サービスから利用する場合も各言語で提供されているaws-sdkのboto3で実装

呼び出しサンプル（Python）
import json 
import boto3 

sagemaker_runtime = boto3.client("sagemaker-runtime") 

data = dataset[start : start + SPLIT_SIZE] 
headers = { "Content-Type": "application/json" }
payload = { "mode": "bulk_linkage", "clients": data }

# 推論エンドポイントにリクエスト 
response = sagemaker_runtime.invoke_endpoint( 
                        EndpointName="demo-endpoint", 
                        Body=json.dumps(payload), 
                        ContentType="application/json", 
                        Accept="application/json" 
) 

decoded_response = response["Body"].read().decode() 
# bodyがstringになっているのでjson化 
result = json.loads(decoded_response)

結果

評価と改善を繰り返すことで既存の名寄せ率が23%向上しました。（サンプルでの結果） 　既存：55% 　新方式：78%

その後一部改善などを行い、以下の精度になっています。
　名寄せ率：77%
　正解率：86% （後から確認したところ、取り組み時の正解率はおおよそ65%ほどでした）

また、リポジトリを分離でき、マネージドサービス上にDevOpsを構築できました。 既存の処理はプロダクトと同一リポジトリ内に実装されていたため、変更をリリースする際にもプロダクト全体のリリースとなっていた点が解消されています。 リリースの運用についても、分離する際にAWS Step Functionsなどを利用することで軽減した状態で構成することが出来ました。

さいごに

今回の取り組みでAmazon SageMakerを用いた本番環境までの流れを把握出来たので、次は監視やスケール、Amazon SageMaker上での学習もDevOpsに組み込んで最適化していくとさらに他の領域も展開しやすくなると考えています。

また、SageMakerのマルチエンドポイントは現在のサービスにもフィット出来そうなので、自社で作り上げる膨大なバックエンドのデータを用いてエンドユーザー毎に最適なAIを届けることも取り組んでいければさらに価値が出せるのではないかと思います。

名寄せの精度という点では、今回は深層強化学習を利用したところですが、他にも辞書の改善や多段的な推論など含めてアプローチできることはあるので、そういうところにも手を入れることによって向上出来るのではないでしょうか。

Baseconnectではエンジニアメンバーを絶賛募集しております！ もしご興味をお持ちいただけた方がいらっしゃいましたら、下記のリンクよりぜひカジュアルにお話しさせていただけましたら嬉しく思います。

herp.careers