こんにちは。NewsPicksエンジニアの雲越です。

NewsPicksでは、記事のレコメンドを始めとする機械学習を組み込んだシステムがいくつか動いており、中にはAmazon SageMakerを使って実装されているものもあります。

今回はそんなシステムの一つである記事分類システムが、どのようにSageMakerを使っているかについて紹介します。

どんなシステム？
学習について
- Training Jobを利用している背景
推論について
- Serverless Inference導入
- CDKで推論エンドポイントを作るのをやめました
今回の対応まとめ
今後の改善
終わりに

どんなシステム？

NewsPicks編集部によってオリジナル記事が投稿されたり、URLピックという機能やRSSによって外部メディアの記事が取り込まれたりすると、NewsPicks内に記事が登録されます。

この際に、その記事のタイトルとサマリー(本文を要約した文章)からどのカテゴリ(金融やキャリアなど)に属するのかを分類して返しているのが、今回紹介する記事分類システムです。

モデル学習についてはSageMaker Training Jobを使っていて、記事のタイトルとサマリーを受け取って、推論して、結果を返すところをSageMaker Serverless Inferenceで行っています。

実はこのシステムは昨年末にリリースしたものなのですが、早く稼働させることを優先してメンテナンス性等について十分考慮されておらず、かなりの課題を抱えた状態でした。ずっと心残りになっていたのですが、今回改善の時間を取らせていただいて課題をある程度解消することができました。

以下では現状の構成と併せて、当時の課題とその解決のために行ったことについてもお伝えできればと思います。

学習について

訓練データをS3から、学習用コードを含むDocker imageをECRから取得し、Training Jobを実行することでモデルを学習します。ジョブが完了して作成されたモデルはS3に配置されます。

Training Jobの実行にあたってSageMakerで用意されたコンテナ等を利用する方法もありますが、このシステムの場合は形態素解析にMeCabを利用しているため、予めインストールしてあるimageを別途自前で用意する必要があります。独自のimageの作り方については公式のサンプルコードがあるので、こちらを参考に作成しました。(推論でも同様のことを行ってます。)

データやモデルを配置するS3バケットやDocker imageを管理するECRのリポジトリ、Training Jobの実行に必要な権限などはCDKで管理をしています。 Training Jobの実行はboto3を使ったPythonスクリプトによって行っています。

Training Jobを利用している背景

元々はSageMakerのノートブックでモデルの学習を行って本番環境にデプロイしていたのですが、以下を課題に感じていました。

ノートブックだとコードの管理がしづらいこと
デプロイ手順が煩雑になっていたこと
- AWSコンソールにログイン、SageMakerノートブックインスタンスの起動を待つ、ノートブックを開いてデプロイ処理を実行するといった手順が必要で、そのままだと効率化が難しそうでした。

これらを解決するためにTraining Jobでのモデル学習を行うようにしました。

Training Jobにしたことで、

学習処理をPythonスクリプトに書いてgitで管理
Training Jobの実行はPythonスクリプトで手元から実行

というようにどちらの課題も解決することができました。

注意点としては、一応モデルの出力先が指定できるのですが完全にはできません。例えば、出力先としてs3://trained-models/のように指定できますが、モデルが出力されるのはs3://trained-models/{Training Job Name}/output/model.tar.gzになります。

利用するモデルの配置場所を固定したい場合などはこれを補助する実装や手順が必要になります。

推論について

推論については、上記のTraining Jobで作成されたモデルをS3から、推論用コードを含むDocker imageをECRから取得し、サーバーレスの推論エンドポイントを立てています。

学習のときと同様にS3とECR、IAMについてはCDKで管理していますが、SageMakerについてはスクリプトによって設定の作成とエンドポイントの更新を行っています。

Serverless Inference導入

先日一般公開されたServerless Inferenceを本番環境で使っています。

以前まではリアルタイム推論でエンドポイントを作っていたのですが、このときは全開発環境にエンドポイントを用意していませんでした。リアルタイム推論ではインスタンスを立てる必要があるので、普段使わないときでもコストがかかってしまいます。弊社では開発環境が10面以上あるためそのコストを考慮して、限られた環境でだけ使えるようになっていました。そのため、他の環境で使う必要があれば立ち上げる作業を行う必要がありました。

使いたいときに色々準備が必要というのは開発者体験としても良くないのでなんとかしたいと悩んでいたところ、ちょうどServerless Inferenceが一般公開されたと聞き、即座に導入を進めることにしました。

この記事分類システム自体、新たに記事が登録されたときだけ動けば良く常時立っている必要がないという点でサーバーレスとは相性が良かったことや、現状Serverless Infereneceで対応していないGPUを使っていないシステムだったことなどが幸いし、すぐに導入することができました。実装についても、InstanceType、AcceleratorTypeなどのリアルタイム推論で必要なパラメータを削除してServerless Configを追加するくらいで特に難しいところはありませんでした。

以下のようなイメージです。

{
  "VariantName": "AllTrafic",
  "ModelName": "ModelName",
  "InitialVariantWeight": 1,
  "ServerlessConfig": {
    "MemorySizeInMB": 4096,
    "MaxConcurrency": 5,
  }
}

結果として、コストがほぼかからなくなりました。 Cost Explorerにてサービスを「SageMaker」で絞り込み、使用タイプごとに表示して「{リージョンコード}-ServerlessInf-{メモリ数}」というタイプを確認することで、Serverless Inferenceでかかっているコストが確認できます。

Serverless化によるコスト削減(5/12にServerlessに置き換え)

上図は本番環境におけるコスト削減のグラフで、それだけだと元々利用していたインスタンスサイズも小さかったため削減した金額は大きくないですが、開発環境も含めるとそれなりのインパクトがある改善になりました。また、Serverless Inferenceであれば呼ばれない限りはコストがかからないので、すべての開発環境に設定だけ作っておいて使いたいときはすぐに使えるという開発者体験の改善もできました。

一点、最大同時接続数(MaxConcurrency)はデフォルトが5に設定されていますが、多めにしてあげないとリクエストを捌ききれずにエラーを返してしまうことがあったのでそこだけ注意が必要です。こちらによると、MaxConcurrencyはアカウントのリージョンごとに合計200まで設定できるようです。(申請すれば上限を引き上げることも可能みたいです。)

CDKで推論エンドポイントを作るのをやめました

つい先日まで推論エンドポイントの作成もCDKで行うようにしていたのですが、boto3で推論エンドポイントを作成するようにしました。

なぜかというと、推論エンドポイントを作成するために必要なSageMaker ModelやSageMaker Endpoint Configは一度作ると更新ができず、また名前も一意である必要があるため、CDKで変更をデプロイしようとしても名前を変えないとエラーが発生して更新ができないという課題があったからです。そのため、当時はcdk destroyしてからcdk deployしなおすという心理的負荷がかなり高いデプロイ作業を行っており、早急に解決したいと思っていました。

当初はCDKのままでなんとかする方法はないかと考え、名前の末尾にgitのcommit-idやcdk deployの実行日時をつけることで差分を作ることを検討していました。ですが、このシステムでは学習用データはS3で管理しているのでそこだけ変えた場合はcommit-idは変わらないので差分ができませんし、実行日時をつけるとなると何も変更がないのに差分がでるような状態になってしまいます。

そこで、CDKでは設定の作成に必要な権限だけ作成するようにして、それ以降のModelやEndpoint Config, Endpointはboto3を使ったPythonスクリプトによってデプロイするように切り離しました。

下図のようにModelやEndpoint Configについては末尾に実行日時を付与して作成し、それらを使ってEndpointを更新するようにしています。

今回の対応まとめ

改めて改善したことをまとめると以下のようになります。

SageMaker Training Jobを使うことで、従来はAWSコンソールにログインして手作業で行っていたモデルの学習がスクリプト一発でできるようになり、手順が効率化してメンテナンスコストが下がりました。
SageMaker Serverless Inferenceを使うことで、我々のユースケースではコストが激減しました。
推論エンドポイントの作成/更新処理をAWS CDKによるIaCからデプロイ用スクリプトに変更することで、学習用データの更新に伴うデプロイにも柔軟に効率良く対応できるようになりました。