MLOps1年目 - SageMakerを使う中で苦労した課題解決とこれからの展望

はじめに

 テクノロジー戦略室MLOpsチームの古賀です。MLOpsチームでは、レコメンドエンジンのMLOps基盤の構築や、AWS PersonalizeなどのMLサービスやLLMなどML活用を、企画から構築まで進めています。本テックブログでは前者のMLOps基盤構築の取り組みを紹介します。

 以前投稿したテックブログにあるように、AWS SageMakerとStep Functions Data Science SDKを導入し、データサイエンティスト主導のMLモデル改善フローを構築できました。しかし、ゴールというわけではなく、まだまだ改善の余地がありました。

 そこで、本テックブログでは、課題の洗い出しや分析、解決のための取り組みを紹介します。そもそもMLOpsを行う背景は以前投稿したテックブログの「背景」に記載しているので、興味ある方はご確認ください。

以前のテックブログまでにやったことと課題

 まず前提として、今までのMLOpsシステム構成を知る必要があります。下記のような構成になっています。説明の都合上、図を書き直していますが、構成は一緒です。

 一連のMLプロセスがパイプライン化されているため、データサイエンティスト主導でモデル改善フローを実行できます。また、学習-推論時で前処理を共通化しているため、Training-Serving Skewを改善出来ています。

 一方で課題がいくつかありました。

  • 特徴量作成
    • 特徴量を再利用できない
  • 学習:開発時、学習パイプライン(前処理から評価まで)の実行時間が長い
  • デプロイ:単一モデルのデプロイフローしか整備されていない
  • 全体
    • パイプライン実行の複雑化
    • PoCからシステムへの初期導入のリードタイムが長い

各課題とその解決策の分析

 各課題を分析し、必要な解決策を説明します。

特徴量の再利用性の低さ

 現状S3に特徴量を出力しています。S3に特徴量を格納している場合、作成者以外の人が特徴量の利用を判断するのが難しくなります。理由は、S3には特徴量のメタデータ管理やバージョン管理の機能がないからです。これらの機能がない状態で作成者以外の人が利用判断するためには、その特徴量が信頼できるデータソースから正しい変換処理により作られているかをソースコードから読み解き確認し、その結果作られる特徴量が正しいかも確認する必要があります。

 そこで、特徴量のメタデータ管理機能やバージョン管理機能を持ったデータストアに格納する必要があります。

学習にかかる時間の長期化

 学習パイプラインの実行による動作確認に時間がかかっていました。理由は2つあります。

  • 不要なステップも実行されるため。例えば、評価の処理を変更した場合、評価ステップのみ実行すべきですが、前処理ステップから実行してしまいます。
  • SageMaker Processingの起動に時間がかかるため。1つあたり5分程度かかるため、全て実行すると起動だけで15分程度かかります。

 これらの課題を解決するためには、必要なステップだけ実行し、起動時間を減らす必要があります。

単一モデルデプロイフローしか整備されていない

 単一モデルのデプロイのみ実装しており、複数モデルのオンライン評価の仕組みが用意されておりませんでした。そのため、複数モデルのデプロイとそれらのオンライン評価の仕組みが必要です。

パイプライン実行の煩雑化

 デプロイパイプライン以外の動作確認手順が煩雑でした。例えば、ライブラリをインストールし特徴量変換処理を変更した場合、下記の作業が必要でした。

  • 動作確認用のAWSリソースを作成(データ取得パイプライン、S3バケット)
  • Dockerイメージのbuildとpush
  • 動作確認用のデータ取得パイプライン実行notebookを実行

変更するたびに、動作確認手順が変わるのは認知負荷が高い上、対応漏れも発生します。加えて、変更箇所によっては、AWSリソースは他のパイプラインやLambdaの作成も必要になります。これらの課題を解決するには、変更箇所によらない動作確認手段が必要になります。

PoCからシステムへの初期導入のリードタイムの長期化

 レコメンドプロジェクトを下記のように進めていました。

 MLOpsとして最優先で解決したいのは、SageMakerへの載せ替えの長期化です。長期化している理由は2つあります。

  • プロジェクトごとに一からインフラ構築しているため
  • PoC後にMLエンジニアが構築したパイプラインに載せ替えているため

 特に後者の工数が大きく、PoCではPoC後の移植を想定したコードを書いておらず、スパゲッティコードのような状態で移行が大変でした。加えて、載せ替えにあたり、データサイエンティストとMLエンジニアのコミュニケーションコストもかかっていました。

 これらの課題を解決するためには、データサイエンティストの開発スピードを落とさず、作ったモデルをそのままシステムに組み込めるような基盤が必要になります。

これまでの取り組み

 上記課題の中で実際に取り組んだことを紹介します。特徴量の再利用性についてはFeature Store + Dataflowの検証までになりましたが、他の課題はある程度解決できました。Feature Store + Dataflowの検証についても紹介します。

特徴量の再利用性の低さ

 AWS SageMaker Feature StoreとGoogle Cloud Dataflowを選択し、検証を進めています。AWS SageMaker Feature Storeを採用した理由は下記の通りです。

  • 特徴量のバージョン管理機能やメタデータ機能があり、特徴量の再利用性を高められるため

Google Cloud Dataflowを採用した理由は下記の通りです。

  • データ変更後、リアルタイムに特徴量に変換して格納するため
  • pandas互換のAPIをサポートしており、データサイエンティストも扱えるため
  • 他チームでRDSからBigQueryにETL処理するためにGoogle Cloud Dataflowを導入予定で技術を標準化するため

 現在は下記の構成の検証を進めています。

白抜き部分の実装は概ね完了したので、これから検証 / 本番環境にデプロイし問題がないか確認していきます。問題がなければ、稼働中のレコメンドプロジェクトにも導入していく予定です。

学習にかかる時間の長期化

 SageMaker Pipelinesを導入し解決しました。導入した理由は下記の通りです。

  • キャッシュモードが存在し、不要なステップの実行をスキップできるため
  • ローカルモードが存在し、SageMaker Processing Jobの起動時間を短縮できるため

処理を変更したステップのみ実行できるので、最小限の実行時間にできました。

単一モデルのデプロイフローしか整備されていない

 複数モデルをデプロイできるようにし、推論エンドポイントであるSageMaker Endpointの前段に配置しているキャッシュサーバーにトラフィックを振り分ける処理を実装しました。なお、トラフィックを振り分ける際は、ユーザーに対して推論するモデルを固定しました。

 まず、ユーザーごとにモデルを固定した理由は下記の通りです。

  • ユーザーを困惑させないため。レコメンド機能を使うたびに結果が変わると、ユーザーが混乱するため。
  • 分析を簡単にするため。もしトラフィックをランダムに振り分けると、どのモデルが推薦したか追う必要があります。この方法よりは、ABテスト期間内でユーザーごとにモデルを固定した方がシンプルで分析しやすいと判断しました。

 次に、キャッシュサーバーに実装した理由は、SageMaker Endpointのトラフィック振り分け機能はランダムな振り分けしかなく、ユーザーごとにモデルを固定できないためです。そのため、SageMaker Endpointの前段のサーバーに振り分けロジックを実装しました。

パイプライン実行の煩雑化

 どのファイルを変更しても、Github Actionsワークフローを動かせば動作確認できるようにしました。加えて、Model Registryに登録された評価結果を承認し、対象のモデルをデプロイできるようにしました。

 Github ActionsワークフローはDockerイメージのbuildとpush、Lambda関数の作成 / 更新、データ取得と学習パイプラインを作成 / 更新し実行します。どこを変更しても、このワークフローを動かせば確認できるため、認知負荷を減らせました。加えて、必要な時のみDockerイメージの構築やLambda関数の更新を実行しているため、不要な実行時間はありません。

 Model Registryに登録されたモデルを承認し、デプロイパイプラインを実行します。このようなフローにしてる理由は、人間が評価結果を確認してからデプロイしたいためです。ただし、この判断基準を明確なルールにできるならば、Github Actionsワークフローでデプロイまで実行しても良いと考えています。

これからやりたいこと

 短期的には、レコメンドプロジェクトの開発効率を高めるために、解決できてない課題を解決したいです。

  • 特徴量を他プロジェクトで再利用できない
    • Feature Store + Dataflowを本格導入したいです。パフォーマンスに問題が無いことを確認する、他の機能との優先度の兼ね合いなどありますが、実現していきたいです。
  • PoCからシステムへの初期導入のリードタイムの長期化
    • MLOpsフレームワークの構築、または、標準構成のテンプレート化などで短縮させたいです。PoCする段階からSageMaker上で開発することで、SageMaker移行の工数を無くすためです。

 長期的には、開発効率だけでなく、ビジネスサイドがレコメンドが事業に与える影響を直感的に分かるようにし、ビジネスサイドとデータサイエンティストのコミュニケーションも効率化させ、レコメンドの精度改善スピードを向上させたいです。

まとめ

 以前のテックブログ執筆時の課題と解決のために、提案し実行した取り組みを紹介しました。まだGoogleが提唱する MLOps レベル1に達してないですし、使い勝手も改善していく必要があります。加えて、今回は紹介できませんでしたが、ML活用も推進しています。レバレジーズではMLOps基盤やML活用を企画から考え構築までやりたいエンジニアを募集しています。興味のある方の応募をお待ちしています。