enechainにおけるMLOps基盤の技術選定

ogp

データプラットフォームデスクの佐崎です。今回は、enechain の MLOps 基盤の全体構成や、各コンポーネントの選定理由についてお話します。

背景

enechain はあらゆるエネルギーの価値を自由に取引できるプラットフォームを運営しており、電力などの先物商品を取り扱っています。

先物取引において、売り手も買い手もどこまで妥協するかを考えるのは非常に難しい問題です。この時、「今の市況ではどれくらいの価格が基準となるのか」という情報があれば、市場参加者はそれを参考にして売り/買いの価格を検討することができるため、取引量を増やす起爆剤となることが期待されます。

enechain では市場の活性化を目的として、電力などの先物商品に対して基準価格を算定する独自モデルを作成、運用しています。これらの一部で機械学習を利用していますが、以下のような課題により開発、運用コストが肥大化していました。

jupyter notebook ベースで実験コードを乱立させてしまい、その管理、共有に手間がかかる
複数エンジニアが実験するが、その結果の管理、比較に時間がかかる
実行パイプラインが Cloud Functions のつなぎ合わせで構成されており、プロセスのどの部分で失敗したかの特定など運用に工数がかかる

これらの課題の解決のため、今回新たに MLOps 基盤を整備しました。

MLOps 基盤で実現したかった要件

今回 MLOps 基盤を新規構築するにあたって、下記要件の実現を目指しました。

簡易に実験環境が構築でき、結果を共有しやすいこと
実験結果の管理を省力化し、比較可能な状態を保つこと
モデルとその実行環境をバージョン管理し、出力を再現可能にすること
機械学習パイプラインの状態把握を容易にすること
CI/CD を整備し、自動テストや自動デプロイを実現すること

また、enechain ではエネルギー・電力事業という社会的に非常に重要かつ市場規模の大きな領域でビジネスを行っています。そのため、セキュリティを強固にすること、今後のビジネス拡大によるデータ量の増大に対応して計算能力をスケールできることが非常に重要となります。これらの前提を守った上で、エンジニアの利便性を向上させる必要がありました。

システム構成

MLOps 基盤の構成は図のようになっており、主に下記コンポーネントで構成されています。

実験環境
実行環境管理
実験管理
ワークフローエンジン
CI/CD
通知・監視

マネージドサービスを中心とした構成となっていますが、いくつかのコンポーネントは自前の GKE 上にデプロイしています。

各コンポーネントの選定理由と使用感をお伝えするとともに、セキュリティ担保、スケーラビリティ、運用負荷の低減のために工夫したポイントもご紹介します。

1. 実験環境

実験環境はエンジニアが前処理、モデリング、評価等を試行錯誤するための jupyter notebook を実行する環境を提供します。

jupyter notebook の運用では下記の課題がありました。

実行環境がローカルに存在するため、機微な情報を扱う際の証跡を残すことが困難
ローカルマシンのスペック以上の計算処理ができない
複数名のエンジニアがローカル環境の notebook で機械学習モデルを検証しており、notebook ファイルの受け渡しで情報共有する必要があったため手間がかかる

今回、これらの対策として Vertex AI workbench, JupyterHub on GKE, JupyterLab on local machine について比較、検討し、最終的に Vertex AI workbench を利用することにしました。比較結果が以下になります。

証跡を残す観点ではクラウド上で環境を構築するのが現実的な対応となります。また、機微データを誤ってローカルに保持してしまうリスクを低減できることも大きなポイントです。

計算能力の拡張性という観点では JupyterHub on GKE が上回りますが、インフラ側で問題が発生したときの対応コストを鑑み、今回は採用を見送っています。

使用感

notebook については、今までファイルを共有していたところがリンクで共有できるようになったため、かなり手軽になった印象です。

また、Vertex AI workbench では BQ と GCS をブラウジングでき、探索的なデータ解析を手軽に実行できるようになりました。

一方、notebook を差分管理できておらず、過去バージョンを参照・共有したいときには別ファイルで保存するなど依然として手間がかかっています。2022 年に Fast.ai から発表された nbdev2 を利用して適切な Git 管理を可能にするなど、より簡易に実験結果を管理・共有することは今後改善が必要なポイントです。

2. 実行環境管理

実行環境管理は実行環境の container image を管理するレジストリサービスと、モデル保管機能を提供しています。

実行環境管理については下記の課題がありました。

機械学習では結果の再現性を担保し、事後検証できる状態を保つことが重要だが、実行環境とモデルのバージョンのそれぞれの組み合わせによって結果が変化する可能性があり、組み合わせの管理に大きな手間がかかる

この課題の対策として、コンテナレジストリサービスとの組み合わせで GCS, MLflow Model Registry, Vertex AI Metadata を比較し、最終的に GAR + GCS の構成を選定しています。比較結果は以下のとおりです。

コンテナレジストリとしては、GCP 利用を前提とすると GAR (Google Artifact Registry)の利用が妥当となり、こちらを採用しています。GAR は GCR (Google Container Registry) の後継サービスであり、vulnerability scan を活用して container image の脆弱性検知を自動化でき、セキュリティを向上できるのが大きな魅力です。

モデル管理としては、モデル管理数の増大に対してスケーラビリティがあること、シンプルな構成で運用負荷が低いことを鑑みて、GCS を採用しました。

MLflow Model Registry や Vertex AI Metadata は GCS と比較して高機能ではありますが、現状 MLflow Tracking にてパラメータやメトリクスなどの実験メタデータは管理できていること、今後必要性があれば移行も可能なことから、今回は採用を見送っています。

使用感

実行環境とモデルをそれぞれバージョン管理して保持することで、過去の実験の再現にかかる工数を大きく削減できました。実装上の工夫としては、image name や model file name に実行者の氏名を含ませた上で管理しており、自分と他者の実行環境を取り違えを防いでいます。

また、invoke という python のタスクランナーを利用して image push や model training のコマンドを自作し、tag 名も実行ユーザごとに自動生成したことで、開発者による命名規則のブレを防いでいます。

切り戻しの柔軟性が向上したことも大きな利点です。モデル、実行環境のいずれの変更によっても想定外の結果となる可能性があるため、独立して柔軟にバージョンを変えて試行可能なことで、原因調査時の切り分けなど不具合対応を簡略化できるようになりました。

3. 実験管理

実験管理は、モデルの実行結果や実験パラメータを記録し、結果比較を可能にする機能を提供しています。

実験管理については下記の課題がありました。

実験結果をスプレッドシートで学習結果やパラメータを記録する必要があり、デグレしないようにメンテナンスすることに時間がかかる
実験試行ごとの精度比較を手作業で実施する必要があり、時間がかかる

この対策として MLflow Tracking, Vertex AI Metadata, Comet.ml などの SaaS について比較、検討し、最終的に MLflow Tracking を選定しています。比較結果は以下のとおりです。

多少の運用コストはあるものの、実験記録と比較の機能が充実していること、チームメンバーに利用経験があり学習コストを抑えられること、GUI など必要な機能が網羅されていることから MLflow Tracking を選定しています。

MLflow tracking は GKE 上にホスティングしており、アクセス可能な IP の制限、通信の SSL 化、Workload Identity を利用した権限管理などの工夫により、セキュリティを向上させています。また、ストレージなども自由に変更可能なため、今後の実験試行の増大にも耐えられる構成となっています。

使用感

一時期スプレッドシートを開くのも嫌になるくらい実験管理には苦労していましたが、モデルごと、実験ごとに自動で結果が記録され、GUI から気軽に比較できるようになったことで工数を大幅に削減できました。任意の実験同士の比較結果を URL を介して共有できることも非常に都合がよく、情報共有も効率化できています。

また、GKE 上にサーバ立ち上げが必要なため MLflow Tracking の構築は多少の時間はかかりましたが、運用開始後はトラブルもなく安定して利用できています。

4. ワークフローエンジン

ワークフローエンジンは、機械学習の処理をパイプライン化して管理する機能を提供します。

ワークフローエンジンについては下記の課題がありました。

既存システムでは Cloud Functions を DAG の一要素のように扱ってパイプラインを構成していたが、GUI などわかりやすく状態を可視化する仕組みが無く、試行錯誤がやりづらかったり、運用負荷が大きくなっている

解決策として、 AI platform pipelines と Composer を比較検討し、最終的に AI platform pipelines を選定しています。比較結果は以下のとおりです。

比較対象は両方 GCP マネージドサービスであるため、証跡管理といったセキュリティ面やスケーラビリティといった観点では決定的な差異はありませんでした。

今回は、パイプラインの途中結果の可視化機能など試行錯誤を効率化する仕組みが充実していることを重視し、Kubeflow ベースの AI platform pipelines を採用しています。

AI platform pipelines に対しても IP 制限や Workload Identity による権限管理を設定しており、セキュリティの向上を図っています。

使用感

パイプラインの実行状況を GUI で見れるようになり、状況把握が劇的にスムースになりました。特に、エラー発生時にひと目でどのコンポーネントで失敗したかがわかるため、対応工数を削減できてます。

また、パイプラインのバージョン管理も可能なため、問題発生時にパイプラインバージョンのみを切り戻して検証でき、障害対応などの運用コストも低減できています。

一方で、細かい話にはなりますが、exit handler の使い勝手が悪いことは改善したいポイントです。kubeflow の大本である Argo workflows では複数の exit handler を設定でき柔軟な終了処理が定義できますが、AI platform pipeline の kubeflow は v1 系のため単一の終了処理しか定義できません。そのため、エラー発生したコンポーネントごとに動的に通知先を変更する場合に手間がかかっています。

また、pod 上限数などのインフラ設定を適切に設計しないと backtest など大規模処理を実行した際の待ち時間が増大します。現状は、TL_SECONDS_AFTER_WORKFLOW_FINISH というオプションで pod の実行後の生存期間を運用の邪魔にならない範囲で短くして対応しています。また、連続的に大規模処理を実行する際には kubectl delete pod --field-selector=status.phase==Succeeded で成功済みの pod を削除して pod 上限に引っかからないように対応しています。

5. CI/CD

CI/CD においては、container image やパイプラインのデプロイなど、モデル更新以外で工数が多くかかっているという課題がありました。この課題の解決を目指し、以下のようなツールを組み合わせて CI/CD 環境を構築しました。

CI ツールとしては、Github を利用していること、enechain の他環境でも利用実績があることを考慮し、 Github Actions を採用しています。セキュリティ向上とスケーラビリティ確保のため、GKE 上にホスティングした Github Actions self-hosted runners を利用しており、mangaged な Actions と比較して一定管理コストは掛かりますが、サービス単位での計算資源管理や各種クラウドリソースへのアクセス権限を詳細に設定して安全性を高めることを重視しています。

CI では pytest を利用してパイプラインの自動テストと合わせて ML モデルの動作と精度検証を実行しています。mape, rmse などの設定した評価指標の値がしきい値以上になった場合には CI が失敗するように設定し、劣化したモデルがデプロイされないようにする仕組みを整えています。

CD ツールとしては build 速度向上のため Cloud Build と、k8s リソースのデプロイのため社内標準である Argo CD を利用しています。Argo CD についてもデプロイなどの権限を細かく設定し、安全性を高めています。

CD では Git の Tag と Release 作成、モデルの学習とパイプライン DAG の登録、image build とデプロイを実施しており、結果を Slack 通知する機能も実装しています。

使用感

CI/CD の構築にはかなりの時間がかかりましたが、結果的にエンジニアがモデルやパイプラインの作り込みだけに集中できるようになったため、早期に投資時間をペイできました。

一方、Cloud Build によって一定の時間短縮はできたものの、依然として image build で時間がかかっており、image 軽量化などの工夫を今後進める必要があります。

また、CI で実行しているモデルの精度検証のしきい値については事前調査した上で固定値を利用していますが、今後は過去実績から統計値を算出してしきい値を自動更新するなどの改善もしていきたいと考えています。

6. 通知・監視

通知・監視は、Datadog, Slackで構成しています。

ログ確認からパフォーマンスモニタリングまで監視に必要な機能を備えており、大規模サービスのログ管理・インシデント管理ツールとして利用実績も豊富なため、Datadog を採用しています。

また、社内コミュニケーションツールとして使っている Slack にアラートなど通知を集約し、情報共有の工数を下げることを意識しています。Slack にはパイプラインの実行ステータス等を通知しており、エンジニアがすぐに気づいて詳細を確認できるようにしています。

使用感

アラートや実験結果を Slack に通知することで、スムースに対応に移れるようになりました。具体例としては、下図のようにパイプラインの結果や実行ステータスを通知しています。これにより、実行待ち時間を有効に使うことができ、異常が起きたらすぐに気付ける状態を保っています。

今後の展望

以上のように、構築した MLOps 基盤をフル活用して開発効率化でき始めましたが、改善したいポイントも見えてきました

notebook の Git 管理

現状差分管理できておらず、どのように実験を変更していったかの記録や共有に手間がかかっています。nbdev2 などの notebook の差分管理を支援するライブラリを活用し、情報共有をより充実させることで開発スピードを向上させる必要があります。

特徴量のバージョン管理

今はすべての特徴量を BigQuery に単純に格納しており、バージョン管理はできていない状態です。そのため、過去データで実験、検証したいときには都度再計算する必要があります。バージョン管理できれば再計算の必要性がなくなり、検証高速化と計算コスト削減が実現できるため、Feature Store コンポーネントの取り込みは検討の価値があります。

SLO 管理や精度監視の高度化

現状は実行ステータス履歴やモデル精度推移などをトラックしたいときは都度 Kubeflow や MLflow を参照しています。DataDog の Dashboard 機能や Looker Studio などを活用し、より監視を充実させて運用品質を改善し続けたいと考えています。

上記の課題に対して改善を続けていき、開発者にとってより使いやすい MLOps 基盤を構築していきたいと考えています。

まとめ

今回は、enechain で利用している MLOps 基盤についてご紹介しました。

MLflow Tracking や AI platform pipelines など、ドキュメントが充実している一定成熟したサービスを採用したため、構築やその後の運用でも大きく困ることはありませんでした。

ML エンジニアとしては環境構築や実験リソース管理、実験管理から開放されたことで、検証そのものやパイプライン実装に集中できるようになり、開発工数を大幅に削減できました。また、モデルについてお客様からフィードバックを頂くこともあり、安定的に ML パイプラインを運用する責任感とやりがいを感じながら日々の業務に取り組むことができています。

機会があれば、human-in-the-loop 構成で構築した機械学習システムの事例なども紹介したいと思います。