機械学習実験管理ツールの移行検討

こんにちは、技術本部 研究開発部 ML Platformチームでインターンをしていた岡部です。ML Platformチームでは、主に社内の研究員の成果を高速にリリースするための基盤であるCircuitの開発・運用を行っています。 CircuitはEKSをベースとしたシステムで、シングルクラスタの運用を行っています。Circuitに関する詳細は以下のスライドを参照してください。

speakerdeck.com

本記事は Sansan Advent Calendar 2024、16日目の記事です。

背景・課題

研究開発部では、機械学習の実験管理ツールとしてClearML（旧Trains）を利用していました。これは約5年前に導入され、社内の研究員に広く利用されていました。詳しくは以下の記事を参照してください。

buildersbox.corp-sansan.com

実はこのClearMLサーバーはAmazon EC2上で単純な docker compose up -d によって立ち上げられており、長期にわたりアラートやソフトウェアアップデートなどの運用が行われていませんでした。またEC2インスタンス自体もTerraformのようなIaCツールを使っていないため、どのような設定で立ち上げられているのかも不透明でした。その結果、インフラがレガシー化し、メンテナンス性や拡張性に課題がありました。

先述したように、社内にはCircuitと呼ばれるEKSベースのモダンなアプリケーション基盤が存在しており、さまざまなアプリケーションを統合的に管理・運用する仕組みが整いつつありました。この新基盤にClearMLを移行することで、運用負荷や技術的負債を軽減し、より効率的な運用を実現できるのではないかという期待がありました。

ClearMLのCircuitへの移行

ClearMLをCircuitに移行するにあたり、以下の2つを中心に調査を進めました。

• Helm Chartの活用：ClearMLは公式にHelm Chartが提供されているため、Kubernetes上へのデプロイ自体は比較的容易でした。
• 内部コンポーネントのデータ永続化：ClearMLは内部的にRedis、Elasticsearch、MongoDBといったデータストアを利用しています。これらの永続化方法（PV/PVCの使用、AWSマネージドサービスの利用）を検討しました。

ところで、2024年11月時点ではCircuitのクラスターの更新はブルーグリーンデプロイによって行われています。これは、クラスターの内部でデータを永続化していないために採用されています。 つまり、ClearMLのデータをクラスター内で永続化することは、クラスターの更新時の運用を複雑化させることになります。

当初は、この運用負荷を抑える方法を模索していました。EKSクラスター内部でRedisやElasticsearch、MongoDBなどを運用しつつ、クラスター移行時にデータをスムーズに移行する仕組みを検討していました。しかし、これにはいくつかの困難が伴いました。

• ログ基盤であるElasticsearchのデータ移行：Elasticsearchの保持するデータ量が膨大であり、データ移行には時間がかかる。
• データ移行オペレーションの煩雑さ: 社内固有の事情により、データ移行を行うオペレーションが大変になる。

このような事情を考慮すると、Circuitに移行するのではなく、現在のEC2上でのClearML運用を継続する方が実用的ではないか、という結論に近づきました。ホストボリュームでデータを保持する現行方式であれば、データ移行の煩雑さから解放され、現状維持のまま運用が続けられます。

現状の課題の見直しとManaged MLflowへの移行

改めて課題を整理した結果、社内のユーザー（研究員）が求めているものは、ClearMLという特定のソフトウェアではなく、機械学習実験管理そのものであることを再認識しました。つまり、ClearMLに限らず、機械学習の実験管理ツールである他のソフトウェアも検討の対象となります。実現可能性や社内固有の要件を考慮した結果、AWSのManaged MLflowが有力な代替案として浮上しました。

これは2024年6月に一般公開されたサービスで、SageMaker上のコンポーネントとして、フルマネージドのMLflow環境が提供されています。プレスリリースは以下のリンクから参照できます。

Announcing the general availability of fully managed MLflow on Amazon SageMaker | AWS News Blog

プレスリリースにも記載されている通り、S3バケットを用意するだけで、すぐにManaged MLflowを利用できます。また、SageMaker Studioとの連携も容易であり、既存のSageMakerベースのML Pipelineと統合することも可能です。 MLflowはML実験管理ツールとして現在広く普及しており、既に社内ではSageMakerベースのML Pipelineが活用されていることも導入する上での大きなメリットとなりました。

よってClearMLからManaged MLflowへの移行が、技術的にも運用的にも、研究員とインフラエンジニアの双方にとって有益であると判断しました。

まとめ

すでにレガシー化していたClearMLサーバーを、Circuit基盤へ移行する実現可能性を検討しましたが、データの永続化やクラスター更新時の運用負荷が大きいことから断念しました。 改めて要件を整理した結果、ClearMLに限らず機械学習実験管理ツールの選択肢を広げることが重要であると判断し、Managed MLflowへの移行を無事に決定することができました。

今後はManaged MLflowの有効活用方法を整えつつ、ClearML環境を段階的に縮退することで、社内ML実験管理体制のシンプル化と効率化を実現していく予定です。