システムをシンプルにするための攻めのリアルタイム推論導入 ―― SageMaker Endpointのscale to zero機能

はじめに

レバレジーズで機械学習エンジニアをしている稲垣です。

現在、とあるシステムにレコメンド機能を導入するプロジェクトを担当しています。
その中で、システムにSageMakerの新しめの機能をうまく活用できたので、その紹介をしたいと思います。

ぶつかった課題感

今回担当したレコメンドの主な要件は以下でした。

  • レコメンドの更新頻度は日次でOK
  • レコメンドを必要とするシステムAから、レコメンドサーバ(今回の開発対象)に日次でリクエストが来るので、そのリクエストに応じてレコメンドデータを返す

レコメンドサーバの実装方法の大きな分類としては「バッチ推論」と「リアルタイム推論」の 2 つがありますが、
今回の要件だと「バッチ推論」を採用するのが自然な選択肢となります。

レコメンドシステムに関する良書「推薦システム実践入門」の6章でも、分/秒レベルのレコメンドの内容更新が求められない場合はバッチ推論が適していると書かれており、我々のチームでも当初はバッチ推論を採用する予定でした。

しかし、バッチ推論を採用した場合、関連システム全体は次のような形になります。

※補足
システムA:ユーザーが閲覧するフロントエンドを持つシステム
システムB:社内の営業担当者向けシステム

また、システムの主な処理の流れを説明すると以下のようになります。

  1. システムBのRDBからBigQueryにデータ連携
  2. SageMaker Transform Jobでバッチ推論し結果をS3に出力
  3. S3に出力されたデータをECSタスクのジョブでElasticCacheに書き込み
  4. システムBからレコメンドAPIにリクエストが送られ、レコメンドデータを返す
  5. システムBからシステムAにレコメンドデータを連携
  6. ユーザーがレコメンドを閲覧できる

ポイントとしては、今回の関連システム全体ではバッチ処理によるデータ連携が4箇所もあり、新しいレコメンドをきちんとシステムAに渡すためにはそれぞれのバッチ処理の実施タイミングを考慮する必要があります。

また、バッチ推論の結果を保存するためのDB (図ではElasticCache) や、S3からDBにデータをロードする仕組み(図ではECSジョブ)も必要になります。

さらに、バッチ推論では「システムBから実際にどのレコメンドデータがリクエストされるか」はリクエストされるタイミングまで分からないため、候補となるすべてのレコメンドデータを事前計算してDBに貯めておく必要があり、計算時間やコスト面での無駄も多いという課題もありました。

説明しよう!SageMaker EndpointのScale to Zeroとは

そんな中、何か良い方法ないかなとネットを漁っていると、2024年のAWS re:InventでSageMaker EndpointのScale to Zeroという機能がリリースされていることを発見しました。

従来のリアルタイム推論では、インスタンスが常時起動しているためコストが高くなりがちでしたが、Scale to Zero機能を使用することで、リクエストがない時はインスタンスが自動的に停止し、コストを劇的に抑えることができます。
(似た機能にSageMaker Endpointのサーバレス推論もありますが、こちらはGPUが使えないため今回の候補からは外れました)

この機能を採用するデメリットとしては、インスタンスが停止した状態でリクエストを送るとインスタンスの初回立ち上げに7-8分程度かかるため、そこをシステムとして許容できるかは確認する必要があります。(今回はOKでした)

また、この機能を利用するにはSageMakerのInference Componentというこれまた比較的新しい機能を使う必要があり、実装スケジュールもそこそこタイトだったため、攻めた技術選定をした結果予期せぬ沼にハマリ結果実装が遅れてしまうというリスクもありました。

しかし、既存の問題をキレイに解決できそうという点と、開発メンバーに頼りになる方々が揃っていたため、「まあ何とかなるだろう!!」と採用することにしました。

改善された世界

Scale to Zero機能の実装では次のような点に苦労しました。

  • 実装の参考になるドキュメントが非常に少なく、我々がやりたいような実装をするにはPython SDKの実装を一部読む必要があった
  • Inference Componentへ配分するリソース調整が割とシビアで適正な値でないとエラーで落ちる
  • デプロイごとにImage Push、Instance立ち上げ、Inference Component立ち上げが発生するので1回のデバッグに30分ほどかかる(これが一番つらかった)

試行錯誤の末、無事機能の実装が完了しシステム構成は以下のようにシンプルにできました。

改善された世界における主な処理の流れを説明すると以下のようになります。

  1. システムBのRDBからBigQueryにデータ連携
  2. システムBからレコメンドAPIにリクエストが送られる
  3. レコメンドAPIからSageMaker Endpointにリクエストを送り、リアルタイム推論された結果を受け取り、システムBへ返す
  4. システムBからシステムAにレコメンドデータを連携
  5. ユーザーがレコメンドを閲覧できる

この構成では、システムBからリクエストされたレコメンドデータのみを計算すればよいため、無駄な計算やストレージコストも削減できます。

実際、バッチ推論を採用した場合よりリアルタイム推論(scale to zero)を採用した方が主に以下の観点でコストを抑えることができました。

  • ElasticCache自体が不要になる
    • ex. cache.m6g.large使用なら月80$程度の削減
  • バッチ推論の時間が短縮される
    • ex. 推論時間を3/4にできた場合、SageMakerのInstanceにml.g4dn.4xlargeを使っていたとすると$60/月程度の削減

おわりに

この記事では、SageMaker Endpointの新しめの機能を使ってレコメンドシステム全体をシンプルにする試みについて紹介しました。

レバレジーズの開発組織では、エンジニア一人一人の裁量が大きく、自由度の高い設計や攻めた技術選定ができる点が特徴の1つかと思います。
(もちろんオーバーエンジニアリングは危険なので関係者との十分な合意形成がある前提ですが)

レバレジーズでのMLエンジニアの働き方に興味をお持ちいただけた方は、ぜひこちらから求人/カジュアル面談へのご応募お待ちしております。

おまけ

当たり前ですが、auto scalingが正しく設定されていないとコスト削減にならないので慎重に確認しましょう。
私はauto scalingの設定をミスって笑えない金額の請求が発生し各方面に謝罪しました。(震え)




Source link

関連記事

コメント

この記事へのコメントはありません。