LayerXにおけるVertex AI Pipelinesの導入と活用

ツール導入前の課題

既存の機械学習開発では、時間のかかる部分やGPUを使う部分だけ、Vertex AI CustomJobを使って、実験を行っていました。その際は、CustomJob内で作成したモデルや予測結果などはGCSやS3にアップロードし、CustomJobが終わったら手元に予測結果を持ってきて評価を行う、といった流れで行なっていました。このような開発を行っていた中で以下のような課題が出てきました。

1. わざわざVertex AI CustomJobに渡す必要のない前処理・評価などの作業が、CustomJobの実行終了を待ってから手元で行うため、軽い作業の割にインスタンスなどを立ち上げる必要があり、心理的コストになっていた。
2. 一度のCustomJobで訓練から予測まで一気に実行すると、途中で落ちた時に訓練からやり直す必要があり、開発速度の低下を招いていた。
3. 依存関係のある複雑な処理を行おうとするとバグが起こりやすかったりと、複雑度を上げるようなモデル改善を行うことに対する心理的コストが高くなっていた。
4. チームメンバーが改善を行った評価結果などがスプレッドシートで管理されており、手動での転記作業が面倒で行われなくなったり、コードレベルでの実験への紐付けなどが行えていなかったので、再現性が怪しくなっていた。

どのような状態を目指していたか

実験に対しての心理的障壁をなくし、改善を高速に回せるようになり、その結果が自動的に溜まってチームメンバーに知見として広まっていくような状態を目指しました。 また、今まで手動でモデルの再学習していたような作業も、この際定期実行できるようにしたいなと考えました。

• Amazon SageMaker Pipelines
• Apache Airflow (Cloud Composer, AWS MWAA)
• gokart
• Prefect

コスト

少人数チームで専任のMLOpsエンジニアがいないため、インフラなどの管理コストなどのかからないツールであることを重視していました。また、フルマネージドのサービスであっても固定でインフラコストがかかってしまうようなツールやSaaSのようなツールは今のチーム規模ではまだ必要ではないと考えていました。

開発の活発度 (コミュニティの発達具合)

開発が活発であったり、技術発信が良くされているようなプロダクトであれば、詰まった時に問題が解決しやすく、すぐに解決できなくてもいずれ解消されるであろう、という考えから検討軸に入れていました。

機能面での充実

ざっくり普段の開発や今後の開発への拡張性として、以下のような機能が欲しいなと考えていました。

• 定期実行を行うことができるか。
• キャッシュ機能があるか。
• パイプラインの途中から再実行できるか。
• コンテナベースで実行ができ、かつk8sを持つ必要がない。
• 各プロジェクトのパイプラインをUI上で一元管理できる。

私も全てのツールについて完全に調査できているわけではありませんが、ざっくり比較ポイントを表に整理しました。

こう見るとどのツールも一長一短で、悩ましいと考えていました。

そこで、チームに使い方を伝える際に一定肌感のあるものを使った方が利用を促進しやすいので、まず自分自身が使ったことのあるAirflowやVertex AI Pipelinesに絞りました。この二つだと、Airflowはチームで管理するプロジェクトが増えてきたときにUI上で色々なパイプラインを一元的に管理できたり、同じパイプラインの過去の実行履歴などに対しての横軸比較が簡単にできたりする部分が魅力的ではあったのですが、チーム規模やコストを考えてVertex AI Pipelinesを選択しました。

コスト面以外でのAirflowと比べたVertex AI Pipelinesの良い点としましては、GCSとのインテグレーションが強く、コンポーネントで生成したデータを特に意識せずにGCSに配置してくれる点です。Airflowだとデータの受け渡しをクラウドストレージ経由などで自前で実装する必要があったのですが、その部分を意識せず開発できる点は強みでした。 また、Airflowの場合は運用フェーズに載っている静的なパイプラインの運用を意識しているため、パラメータや構造などを頻繁に変えるような実験フェイズのパイプラインだと開発しにくいというデメリットもありました。

Prefectも元Airflow開発者によってAirflowにあった課題を解決したとされるツールなため、試したい気持ちはありましたが、今回は導入をスムーズにするため使い慣れているツールを選択しました。

改善したかった課題はどれくらい解決されたか

まず、重い処理も軽い処理も全てVertex AI Pipelinesに移行したことで、前処理、訓練、予測、後処理、評価まで全て自動化することができました。複雑な依存関係もDAGとして整理することで見通しが良くなりました。

評価結果はVertex AI Experimentに保存することで、チーム間での実験結果の共有も楽になりました。 一方で、パイプラインの途中で落ちてしまった場合にエラーが出て再度途中から実行したい場合や、他のメンバーの実験をパイプラインの途中から再現したい、といった場合には後述するような工夫が必要でした。

どのような成果が得られたか

精度改善の実験を手軽に行えるようになったことで、現状では平均にチームメンバー一人当たり週に1回以上の実験を回すことができるようになりました。その際に、実験結果とそれに対応するパイプライン実行URLも自動で貯まるようになり、結果の共有や実験の信頼性、再現性が高まりました。

また、単純な機械学習モデルだけでなく、弊社が最近リリースした「パーソナライズドAI-OCR」などの複雑な依存関係のある機械学習モデルも開発しやすくなりました。

Vertex AI Pipelinesで提供されているScheduler APIを使うことで、定期実行もCloud SchedulerやCloud Functionを使わずとも手軽にできるようになり、今まで手動で再学習や特徴更新していた部分を自動化することができるようになりました。

副次的な効果ではありますが、並列実行などもしやすくなったため、時間のかかるコンポーネントを並列化することによって高速に実験を回すことができるようになりました。

Vertex AI Pipelines自体の持っているキャッシュ機能が、チーム開発やデバッグなどを想定しているものではなく、キャッシュをかけたいのにかからない時があった点です。 例えば、他のチームメンバーが行った実験の後処理部分だけ修正したい、となったときに後処理以前の部分はキャッシュを効かせたいといった要件があります。Vertex AI Pipelinesのキャッシュ機能はinputのデータなどだけではなく、実行に使ったDockerのimageや環境変数などもみた上でキャッシュを判断するため、新しくimageをpushしたり、環境変数にユーザなどを使っていたりすると簡単にキャッシュが掛からなくなってしまいます。実験を効率化するために、途中から再実行させるといった機能はmustだと思っていたので、デフォルトのキャッシュ機能に振り回されないようにするような工夫が必要でした。この工夫については、Vertex AI PIpelinesでの実験を加速させるためのTIPSという記事で紹介しています。

k8sなどのインフラ管理のできない小規模なチームであれば、Vertex AI Pipelinesの導入は第一選択肢になると思います。特にKubeflow Pipelines v2が整備されたことで、以前のようにyamlでComponentを定義する必要がなくなり、PythonicにComponentを定義できるようになり、使い勝手が良くなったと思います。また、Scheduler APIは現状だと更新周りが不十分なので、現状だとCloud SchedulerとCloud Functionを使う構成も一つの手になると思います。

あまり依存関係のない単発の実験を動かすのであれば、Vertex AI CustomJobをとりあえず使うのも手です。

今後、機械学習プロジェクトが増えていく想定なので、パイプラインのパフォーマンス管理をどのようにしていくかが課題です。自作での監視ツールを作成か、Prefectなどの複数パイプラインを管理できるツールに移行するかを進めていく予定です。