この記事では、NTT Communications Advent Calendar 2024 16 日目の記事です。本記事では MLflow という実験管理 OSS を Google Cloud の Vertex AI Experiments に置き換えを検討してみた話について記載しています。
はじめに
こんにちは、イノベーションセンターの林です。普段はノーコード AI 開発ツールである「Node-AI」というプロダクトでソフトウェアエンジニアとして開発に携わっています。また、SRE・オブザーバビリティ・Google Cloud といった領域も好きで趣味で勉強していたりします。
Node-AI は AI 開発ツールということもあり、アプリケーションの構成の中に機械学習系のタスクを担うコンポーネントも存在しています。その中で MLflow と連携しているエンドポイントがあり、その部分において運用続ける中で課題が浮き彫りになりました。
本記事では、この課題に対するアプローチとして同じく実験管理機能を提供する Google Cloud の Vertex AI Experiments に置き換える検討をしたことについてまとめていきます。
結論
簡単に本記事の結論をまとめると、現時点での MLflow の利用方法では Vertex AI Experiments に置き換えるのは厳しいという結果でした。
Google Cloud のリソースはものによって上限が割り当てられていてサービスの過度な負荷を防いでいます。Vertex AI Experiments というサービスにも実験情報の記録や取得について上限が割り当てられています。
検証の中で単純なユースケースにも関わらず、この上限に当たってしまいエラーが発生する事態に陥りました。上限は申請によって緩和できますが、実際のユースケースを想定したときに多少緩和しただけでは同様のエラーに陥ることが考えられました。
これは Node-AI の MLflow の利用の仕方が起因しており、本来の想定されている利用の仕方から逸脱していることが根本原因だと考えています。まずは、その点の見直しから整理していく必要があると結論に至りました。
以降では、実験管理機能の説明、 MLflow・Vertex AI Experiments の紹介、検証する中でぶち当たった壁、その壁に取り組んでみての考察と今後という流れで執筆しています。
話題の中心となる実験管理機能
前提知識として、実験管理について説明します。実験管理の目的は、実験 = モデルの評価状況を記録することで後から比較・モデル選定を可能とすることです。一般的に管理対象となるのは、下記の要素が挙げられるようです。
- コードバージョン：どのようなコードで実験をしたか
- データバージョン：どのようなデータで学習・テストをしたか
- ハイパーパラメータ：ハイパーパラメータとして何を使用したか
- メトリクス：実験の結果、どのような結果が得られたか
- 環境：実際に動作したときのマシンとしてどのようなものを使用していたか
つまるところ、実験を再現するための情報とその結果の記録が主要な機能と言えそうです。
この実験管理機能を提供している代表的な OSS が MLflow となります。MLflow にはこの実験管理機能に加えて、記録した精度を比較しやすくする可視化機能、記録した情報に紐づくモデル保存機能なども提供しています。
引用：https://mlflow.org/docs/latest/getting-started/quickstart-2/index.html#compare-the-results
浮き彫りになった課題
Node-AI では上述した通り、機械学習系のタスクを担うコンポーネントが存在しており、一部のエンドポイントで MLflow と連携しています。連携方法としては、データベースに対する取得・保存のためのインターフェースとして MLflow の実験管理に関する API を利用しています。
この MLflow は導入当初問題なく利用できていたのですが、運用を続ける中でユーザー数の増加や学習規模の増大によって以下のような課題が顕著となりました。ここでは 2 つ取り上げたいと思います。
パフォーマンス面
Node-AI ではモデル作成時に複数回の学習を実行します。この複数回の学習結果をニアリアルタイムで取得し UI 上に描画したいがために、学習量ごとの保存に加えて 1 秒間に 3~4 回の取得の API コールをしています。
この利用の仕方では書き込み及び読み取りの処理が遅く、学習開始から完了までの一連のリクエストのボトルネックになっています。これは同時に学習を開始するユーザー数が多いほど、負荷は増大してパフォーマンスの劣化を引き起こします。
セキュリティ面
クリティカルなものとして、定期的に検出される脆弱性によって緊急のアップデート作業が必要となる点です。下記は、昨年末に検出された脆弱性の例です。
機械学習のライフサイクルを管理するプラットフォーム「MLflow」に3件の脆弱性が明らかとなった。いずれも「クリティカル（Critical）」とレーティングされている。
サーバ上にファイルを配置することが可能となるパストラバーサルの脆弱性「CVE-2023-6015」が判明。また認証なしにファイルを上書きし、リモートよりコマンドを実行してデータやモデルにアクセスが可能となる「CVE-2023-6018」が明らかとなった。
さらに「REST API」の認証をバイパスしてアカウントの作成が可能となる「CVE-2023-6014」が判明している。
CVE番号を採番した脆弱性情報報告サイトのhuntrでは、共通脆弱性評価システム「CVSSv3.1」のベーススコアを「CVE-2023-6015」「CVE-2023-6018」のいずれも最高値である「10」とした。また「CVE-2023-6014」は「9.1」と評価している。重要度はいずれも「クリティカル（Critical）」。
こういった側面に加えて、OSS 特有のセルフでのメンテナンスコストの高さも相まって運用の負荷を増大させる要因となっています。
Vertex AI Experiments によるアプローチ
私自身が Google Cloud を好んでいることと実験管理機能を提供しているという 2 点から Vertex AI Experiments を置き換えの検討候補としました。
Vertex AI Experiments に関する説明は、公式ドキュメントより引用します。
Vertex AI Experiments は、さまざまなモデル アーキテクチャ、ハイパーパラメータ、トレーニング環境を追跡および分析し、テストの実行の手順、入力、出力を追跡する際に役立つツールです。Vertex AI Experiments では、テストデータセットに対して、トレーニングの実行中に、モデルの集計のパフォーマンスを評価することもできます。この情報を使用して、特定のユースケースに最適なモデルを選択できます。
引用：https://cloud.google.com/vertex-ai/docs/experiments/intro-vertex-ai-experiments?hl=ja
上述した一般的な要素に加えて実行順序や関連したオブジェクトの紐付けといったところも可能になっているようです。
Vertex AI Experiments に期待した点としては、Google マネージドなサービスというところでのセキュリティ面の課題からの解放およびコンピューティングスペックのオートスケーリングによるパフォーマンス面の課題からの解放という点でした。（検討開始当初は、前者は間違いなく実現できると思っていましたが、後者は要検証項目となっていました。）
検討の中でぶち当たった壁
検討は大きく下記の手順で進めました。
- Vertex AI Experiments の理解
- 現状の実装における MLflow の利用の仕方の理解
- 現状の実装の MLflow を Vertex AI Experiments で置き換えられるかの検証
また、検討過程は Architecture Decision Record（ADR）として GitHub Issue にまとめました。こちらの記事執筆にも非常に役立っています。
Vertex AI Experiments の理解では、主に公式ドキュメントの読み込みと用意されているサンプルコードを触ったりでどういった特性があるのか、利用の仕方ができるのかなどを確認しました。サービスによるのですが、公式ドキュメントに Colab や Colab Enterprise ですぐに環境を立ち上げてサンプルコードを実行できるリンクが用意されていたりします。
引用：https://cloud.google.com/vertex-ai/docs/experiments/intro-vertex-ai-experiments?hl=ja
上図内の「Colab Enterprise で開く」をクリックすると、Google Cloud Console に遷移して Vertex AI の Colab Enterprise が開きます。
現状の実装における MLflow の利用の仕方の理解では、自分自身今までに触ったことがないコンポーネントだっためチームのテックリードにコードを解説してもらいながら理解を進めました。
それらの手順を経て、現状の実装の MLflow を Vertex AI Experiments で置き換えられるかの検証に進みました。Vertex AI Experiments は MLflow と一部互換性があり、どの程度現状のコードを置き換えられるかという点にフォーカスして進めました。
そこでぶち当たった壁 = 課題について 2 点紹介します。
前提知識の整理
実際の課題について説明する前に単語の認識を合わせたいと思います。Vertex AI Experiments の公式ドキュメントを参考にします。
1 回の学習記録を 「テスト実行」 と呼びます。テスト実行にはその学習で設定したハイパーパラメータ・学習結果のメトリクス・その他の環境値などが記録されます。コード上では ExperimentRun というオブジェクトで表現されます。
テスト実行には、ユーザー定義の指標、パラメータ、実行、アーティファクト、Vertex リソース（たとえば、PipelineJob）を含めることができます。
引用：https://cloud.google.com/vertex-ai/docs/experiments/intro-vertex-ai-experiments?hl=ja#experimentrun
また、上記のテスト実行の集合を 「テスト」 と呼びます。コード上では Experiment というオブジェクトで表現されます。
テストは、ユーザーが入力アーティファクトやハイパーパラメータなどのさまざまな構成をグループとして調査できるパイプライン実行に加えて、一連の n 個のテスト実行を含むことができるコンテキストです。 https://cloud.google.com/vertex-ai/docs/experiments/intro-vertex-ai-experiments?hl=ja#experiment
テスト実行のメタデータの扱い方
最初の大きな課題となったのは、テスト実験のメタデータの扱い方です。Node-AI の実装の中では、特定のタイミングのテスト実験結果を取得して描画する画面が存在します。
学習方法はクロスバリデーションの有無やパラメータ探索の有無などいくつかの組み合わせが存在する上に、学習完了後にまとめて描画するのではなくニアリアルタイムに描画していることから特定のタイミングの結果を取得する必要がありました。これには、テスト実行にメタデータを付与して高速に検索することで対応しています。
これを実現するために、MLflow の API の中でも set_tags() を利用しています。その他にも利用しているものもありますが、主だったものとして log_metrics() と log_params() というものがあります。それぞれの簡単な用途が下記になります。
- log_metrics()：学習毎の評価制度を記録する
- log_params()：学習毎のパラメーターを記録する
- set_tags()：学習毎のメタデータを記録する
上述した通り、Vertex AI Experiments は MLflow に一部互換性があり、これらの API を置き換えられると検証がスムーズに進むため大きなポイントでした。しかし、そんなに甘いことはなくテスト実行のメタデータを記録する set_tags() のみ実装されていませんでした。
この課題は、Vertex AI Experiments を扱う Python ライブラリである google-cloud-aiplatform がどのような構造になっているかを読み解くことで解決しました。
こちらにまとめるとボリュームが 1.5 倍になりそうなので、詳細は「Vertex AI Experiments の実態 - コードを辿った先にあったもの -」という資料と「Vertex AI Experiments をコードから読み解いてみた」というブログにまとめています。前者の資料はここで述べている課題にフォーカスした内容で登壇した際のものです。後者のブログは実際にどのようにコードを読み解いたのかという内容になっていますので、興味ある方は両方を合わせて読んでいただけると幸いです。
ここでは、その結論のみ述べます。コードを読み解くと Vertex AI Experiments は Vertex ML Metadata というメタデータ管理サービスの実験管理観点でのラッパーであることがわかりました。Vertex AI Experiments の Experiment オブジェクトや ExperimentRun オブジェクトの実態は、Vertex ML Metadata の Context というオブジェクトに紐づいていました。
このことをきっかけに Context オブジェクトを読み解くと、 Context オブジェクトは自らのメタデータを更新するメソッドが実装されており、こちらを利用することでテスト実行のメタデータを扱えました。結果として課題だった、set_tags() の置き換えを実現できました。
API コールの上限
上述の一部互換課題を解決したことで、MLflow の置き換え検証を進めることができました。置き換え最中は何度も実行と修正を繰り替えすので、実験管理機能を使用する操作の規模は最小限のものとしていました。
一通りコードの置き換えが完了したので、実験管理機能を使用する操作の規模を徐々に大きくしていたところ下記のエラーが発生しました。
ERROR 429 Quota exceeded for quota metric 'Resource management (CRUD) requests' and limit 'Resource management (CRUD) requests per minute per region' of service 'aiplatform.googleapis.com' for consumer 'project_number:***'. reason: "RATE_LIMIT_EXCEEDED".
少し解説すると、 aiplatform.googleapis.com とは Vertex AI 周りの API エンドポイントのドメインを指します。AI Platform は Vertex AI の前身のサービス名だったと思います。
Quota exceeded for quota metric 'Resource management (CRUD) requests' ということなので、
リソースに対する操作の指標が上限に達したと言われています。
Vertex AI に限らずですが、Google Cloud のリソースにはものによっては上限が割り当てられています。Vertex AI に関しては、こちらのページから各種上限を確認できます。この中でもリソース管理（CRUD）リクエストの 1 分間あたりのリクエスト数が 600 というものが上述のエラーに該当します。
Google Cloud では、割り当てを使用して公平性を確保し、リソースの使用量と可用性の急増を抑えます。割り当ては、Google Cloud プロジェクトで使用できる Google Cloud リソースの量を制限します。割り当ては、ハードウェア、ソフトウェア、ネットワーク コンポーネントなど、さまざまなリソースタイプに適用されます。
今回でいえば、テスト実行を検索して取得する部分でこちらのエラーが発生していることがわかりました。上述した通り学習を実行した際、ニアアリルタイムにテスト実行の結果を取得するために秒間 3~4 回ほどの取得リクエストを投げています。この操作時間が長ければそれだけリクエスト数が増加することとなります。
結果として、この操作が Vertex AI Experiments におけるテスト実行情報の取得について想定を超えるような使い方となっており、API コールの上限に達したということでした。
こちらの上限は申請すれば緩和させることができるのですが、今回のような検証レベルのユースケースでこの上限に達してエラーが出たということは、複数人が同時に該当の操作をするような実際のユースケース場合では上限を緩和させてもすぐに達してしてしまうことが考えられました。
考察と今後
上述した検証レベルのユースケースで API コールの上限に達したという点について、現状の Node-AI の MLflow への API コールの仕方は想定外の使い方であることが考えられました。だからこそ、MLflow がパフォーマンスがボトルネックになっているとこいうことにつながっているのかもしれません。
こちらの検討で現状の Node-AI の該当の実装があまりイケていないことが改めてわかりました。今後取り得る選択肢としては、API コールを減らすような実装の見直ししたのちに MLflow を Vertex AI Experiments に置き換えるといったところでしょうか。
バックエンドの API コールの減少は、フロントの振る舞いも加味して調整する必要があるので、慎重に進める必要がありそうです。
さいごに
今回 Node-AI で利用している MLflow のパフォーマンス面とセキュリティ面の課題に対して Google Cloud の Vertex AI Experiments というサービスの置き換えを検討しました。
検討の中で、テスト実行のメタデータの扱いと API コールの上限の 2 つの課題にぶつかり、前者は解決できたものの後者はクリティカルなものでした。この検討の考察から取り得る選択肢を洗い出すといったとこまでを本記事でまとめました。
かなり長くなってしまいました。ここまでお読みいただいた方、ありがとうございました！ 明日以降のブログにもご期待ください！
