こんにちは。AnyMind で機械学習エンジニアをしている河本直起です。
AnyMind では 0 から MLOps 環境を作成しており、前回の記事では機械学習導入初期フェーズにおいてとっていた構成について紹介させていただきました。今回は、その続編として Vertex AI を用いて実装している現状のバッチ推論基盤についてご紹介できればと思います。
従来の構成
こちらの記事でご紹介させていただいた通り、既存の仕組みでは以下のようにプロダクトアプリケーションが特徴量を取得し、それを推論 API にリクエストするという形になっていました。また、学習に使用するデータもプロダクトアプリケーション側の RDB から直接取得されていました。
課題
推論 API という観点では、プロダクトアプリケーション側の実装が機械学習モデルの使用する特徴量に依存するという点が問題となります。推論 API の入出力が特徴量の変更ごとに変わってしまうことになり、その度にプロダクトアプリケーション側で開発が必要になります。同時に、不要なコミュニケーションコストやミスコミュニケーションによるバグの発生にも繋がります。
次に、推論時にデータセットと接続しているのがプロダクトアプリケーションであるということから、推論結果の蓄積や再利用をプロダクトアプリケーション側で実装する必要があるという点も問題となります。推論結果を他のモデルへの特徴量として利用したい場合や一度推論されたレコードに対してキャッシュされた結果を返したいなど、予測結果が蓄積されている必要があるケースにおいて、都度プロダクトアプリケーション側と蓄積方法含めてコミュニケーションを取り実装してもらう必要がありました。ここで、モデルのバージョニングや推論結果の検証に向けた蓄積など、機械学習で必要なデータの蓄積の要件とプロダクトアプリケーションでの要件が大きく異なるため、コミュニケーションの難易度が高くなります。また、そもそも RDB はプロダクトアプリケーション側の要件に合わせて用意されたものであり、機械学習側のデータ蓄積要件には合わないという問題も発生しました。
最後に、推論 API においては A/B テストやカナリアリリースに向けてリクエストをモデル単位で振り分ける必要がありますが、その振り分けをプロダクトアプリケーション側で行う必要があるという点も問題となります。Vertex AI Endpoints には複数の Vertex AI Models がデプロイ可能であり、割合を定めてリクエストを振り分けることが可能です。しかし、推論に必要な特徴量がモデルごとに異なる一方で、その特徴量がクライアントアプリケーション側で取得されているためその機能を活用することもできませんでした。
大きく言えば、従来の構成では本来あるべき責任範囲の切り分けと実際に開発を行う範囲が一致していませんでした。機械学習エンジニアチームで実装すべきところをプロダクトアプリケーション側に任せることになってしまい、その結果コミュニケーションコストが大きくなり、ミスコミュニケーションやそもそもの要件との食い違いが発生するようになっていました。また、プロダクトアプリケーションから見たインターフェースが定まらないことから、リリースごとにプロダクト側でも開発が必要となりモデルの改善を頻繁に行うことが難しくなっていました。
要件
ここでは、上記課題を解決するためにどのような要件が必要なのかについて説明させていただきます。
機械学習用データソースの作成
まず、機械学習向けのデータソースがプロダクトアプリケーションから分離されている必要があります。そのため、推論処理とモデル学習処理が共通して参照するデータソースを新しく BigQuery に作成し、RDB から切り替えました。こちらの取り組みの詳細に関しては、以下の記事をご参照ください。
Cloud Composer (Airflow) を用いた機械学習向けデータ基盤の作成
この取り組みの結果、データソースに関しては以下のように BigQuery に移行できました。
推論処理
上記課題から、以下の二つを機械学習アプリケーション側で担当する機能として含め、切り分けることが要件となります。
- モデルへのリクエストの振り分け
- 推論結果の蓄積
今回の取り組みにおいては、現状はバッチ推論でビジネス要件が満たせるため、まず上記要件を満たしつつバッチ推論を行うことを目標としています。
新規構成
以降では、上記要件を満たした新規構成について紹介させていただきます。
推論フロー
モデル生成・バッチ推論に関しては Vertex Pipelines (Kubeflow) を利用しています。詳細に関しては以下の記事をご参照ください。
Vertex Pipelines (Kubeflow) の機械学習システムへの導入
推論結果が格納されるまでのフローは以下のようになっています。
まず、モデル学習パイプラインがモデルを生成します。後続のバッチ推論パイプラインではその生成されたモデルを用いて推論を行い、推論結果を BigQuery と Firestore (Datastore mode) に格納します。BigQuery は主に推論結果の確認や検証、後述する機械学習間での推論結果の再利用のため、Firestore (Datastore mode) は推論結果の取得のために使われます。
Vertex Featurestore や Cloud Spanner ではなく Firestore (Datastore mode) を使用している理由は、必要な要件を満たしていると同時に試算した時のコストが最も低かったためです。
推論結果のサービング
推論結果をサービングする際の構成は以下のようになっています。
プロダクトアプリケーションは推論結果が必要なキー（ユーザー単位の推論であればユーザー ID など）とメタデータ（レコメンデーションであれば上位何件が必要かなど）を新しく追加したサービング API にリクエストします。そして、サービング API はリクエストを元に Firestore (Datastore mode) から推論結果を取得し、返すという形になっています。サービング API は Cloud Run 上で走るようになっています。
こうすることで、推論結果の蓄積を機械学習アプリケーション側で行うことができ、プロダクトアプリケーションから見たインターフェースは一定となります。そのため、特徴量やモデルの変更時にプロダクトアプリケーション側で開発を行う必要がなくなります。
モデル単位のリクエストの振り分け
サービング API はプロダクトアプリケーションから見たインターフェースとしてリクエストに応じた推論結果の取得を行うと同時に、内部で A/B テストやカナリアリリースに向けたモデル単位でのリクエストの振り分けを行います。
現状、機械学習モデルは {プロジェクト}_{モデルのバージョン}_{モデルの生成された時間} を最小単位として保存しています。{モデルのバージョン} はモデルの内部処理が変更された場合に切り替えられる単位です。そして、推論結果は {プロジェクト}_{モデルのバージョン} という単位で切り分けて保存しています。
サービング API はどのバージョンのモデルをどの割合で使用するかを設定として保持しており、クライアント側からモデルのバージョンが指定されなかった場合は、定義された割合に応じたリクエストの振り分けを行います。推論結果におけるモデルのバージョンの切り替えは BigQuery ではテーブル単位、Firestore (Datastore mode) では Kind 単位で行なっています。
モデル単位での推論結果の振り分けは以下のように行っています。
デプロイ単位の切り分け
機械学習処理全体としては以下のようにデプロイ単位を切り分け、それぞれでバージョンを持たせています。
- 推論結果のサービング
- モデル学習 / モデルに依存したデータ処理 / バッチ推論
- 特徴量生成
こうすることによって、モデル変更ごとに特徴量生成処理を作成するなどの実装の重複を防ぐと同時に、将来的な汎用化、他チームへの引き継ぎや分業が容易になるようにしています。
機械学習アプリケーション間での推論結果の利用
例えば推論された KOL のカテゴリーを KOL の類似度算出の特徴量として利用する、モデルによって生成された潜在表現をモデル間で共通利用するなど、あるモデルの推論結果を他の機械学習モデルの特徴量として利用したい場合があります。
バッチ推論やモデル学習においては、BigQuery を共通のインターフェースとして推論結果を取得し特徴量テーブルを作成するようにしています。ここで、異なるモデルによる推論結果の混在を防ぐため、使用するモデルのバージョンを固定して推論結果を取得する形にしています。
まだ実装できていませんが、オンライン推論においては上記サービング API を共通のインターフェースとして利用することを考えています。
また、テキストのトークン化など機械学習モデルを使用しない処理においても、モデルという枠組みで生成し同じように共通利用しています。推論結果を共通利用することで、あるモデルの改善が他モデルの精度改善に繋がるようになっています。
終わりに
今回は機械学習モデルを用いたバッチ推論結果のサービングについて紹介させていただきました。参考になれば幸いです。
