分散型 ML の生産性を高める Ray と MLflow の統合

Community Editionに代わり、Free Editionでは無料でより充実した機能をご利用いただけます。ぜひ今日からぜひFree Editionをお試しください。
 

このブログ投稿では、RayとMLflowの新たな2つの統合を発表します：Ray Tune+MLflow TrackingとRay Serve+MLflow Modelsです。これらは、機械学習（ML）モデルの構築とプロダクションへの移行を大幅に容易にします。

これらの統合機能は、最新の Ray ホイールで利用可能です。Ray のナイトリーバージョンを pip でインストールする方法、およびその他の参考資料は、こちらから参照できます。これらは、Ray の次期リリース（バージョン 1.2）にも搭載される予定です。

私たちの目標は、Ray の分散型ライブラリによるトレーニングやサービスの拡張性と、MLflow のエンドツーエンドのモデルライフサイクル管理という、それぞれのプロジェクトの強みを活かすことです。

これらのツールで解決できること

新しい統合機能の説明に入る前に、これらのライブラリで何ができるかを見てみましょう。

Ray Tune によるハイパーパラメータチューニングの拡張

ML モデルのサイズとトレーニング時間の増加により、大規模な ML 実験を単一のマシンで実行することはほぼ不可能です。そのため、実験を複数のマシンに分散させることが必須となっています。

Ray Tune は、ハイパーパラメータのチューニング実験を任意の規模で実行するためのライブラリで、トレーニング時間を大幅に短縮することができます。

Ray Tune を利用することで、以下のことが可能になります。

• 10 行以下のコードでマルチノードのハイパーパラメータスイープを開始
• Pytorch、Tensorflow、MXNet、Keras など、任意の ML フレームワークの使用
• Population Based Training、HyperBand、 Asynchronous Successive Halving（ASHA）など、最先端のハイパーパラメータ最適化アルゴリズムの活用

Ray Serve によるモデルサービングの拡張

機械学習モデルを開発した後、実際に予測リクエストを処理するためにモデルをデプロイすることがよくあります。しかし、ML モデルはコンピューティング負荷が高いことが多いため、実際にデプロイするには分散システムにスケールアウトする必要があります。

Ray Serve は、使いやすいスケーラブルなモデルサービングライブラリで、以下のことを可能にします。

• 複数のマシン間で GPU を使ったモデルサービングを簡素化し、本番環境の稼働率とパフォーマンスの要件を満たす
• Pytorch、Tensorflow、MXNet、Kerasなど、あらゆる ML フレームワークで動作
• プログラムによる設定インターフェースを提供（YAML や JSON は不要）

MLflow によるエンドツーエンドの
モデルライフサイクル管理

Ray Tune と Ray Serve は、ML の開発・デプロイを簡単に分散させることができますが、このプロセスの管理はどのようにすればよいでしょうか ?そこで必要になるのが MLflow です。

実験の実行中には、MLflow 追跡 API により、全ての実験のハイパーパラメータ、結果、モデルのチェックポイントを記録し、それらを簡単に可視化して他のチームメンバーと共有することができます。また、MLflow モデルが提供する標準化されたパッケージにより、さまざまな異なる環境でのデプロイがサポートされます。

主要なポイント

Together, Ray Tune, Ray Serve, and MLflow remove the scaling and managing burden from ML Engineers, allowing them to focus on the main task– building ML models and algorithms.

では、これらのライブラリを活用する方法を見てみましょう。

Ray Tune + MLflow 追跡

Ray Tune と MLflow 追跡 APIとの統合により、分散チューニングの実行情報を容易に MLflow サーバーに記録することができます。

この統合には2つのAPIがあります：MLflowLoggerCallbackとmlflow_mixinです。

MLflowLoggerCallbackを使用すると、Ray Tuneは自動的に各実験のハイパーパラメータ設定、結果、モデルのチェックポイントをMLflowにログします。

以下の図は、Ray Tune が、複数の異なるトレーニングを、それぞれ異なるハイパーパラメータ設定で、全て並行して実行している状態を示しています。これらの実行状況は全て MLflow の UI 上で確認することができ、この UI 上では、ログに記録されたメトリクスを視覚化することができます。MLflow の追跡サーバーがリモートである場合も、他のユーザーが実験結果やアーティファクトにアクセスすることも可能です。

どの情報を記録するかを Ray Tune に頼らずに自分で管理する場合は、mlflow_mixin API を使います。

トレーニングの関数にデコレータを追加し、関数内で MLflow のメソッドを呼び出すことができます。

完全に実行可能な例と詳細な情報については、ここでドキュメンテーションをご覧いただけます。

Ray Serve + MLflow モデル

MLflow モデル は、Python 関数として読み込めるため、Ray Serve を使って簡単に提供することができます。利用するモデルのバージョンは、モデルのチェックポイントから、または、MLflow モデルレジストリからモデルの URI を指定して読み込むことができます。以下に例を示します。

まとめ

Ray と MLflow を共に使用することで、分散型 ML アプリケーションの構築や本番環境への導入が非常に容易になります。Ray Tune + MLflow 追跡は、開発や実験をより迅速かつ管理しやすくし、Ray Serve + MLflow モデルは、大規模なモデルのデプロイを容易にします。

Databricks Community Edition (DCE) での実行例をこちらのNotebookで紹介しています。ぜひご参照ください。
注：この Ray Tune + MLflow の拡張機能は、DCE ランタイム 7.5 および MLR7.5 でのみテストを行っています。

次のステップ

この統合機能は、最新の Ray ナイトリーホイール を pip インストールし、pip install mlflow を実行することで利用することができます。DCE で利用する方法については、こちらのNotebookを参照してください。将来的に、Ray Serve と MLflow モデルのさらなる統合を目的としたプラグインの提供も予定しています。

参考資料：

• Ray Tune + MLflow Tracking統合のドキュメンテーションをご覧ください。
• この統合を使用してPytorch Lightningモデルを調整し、自動ログを取る方法をご覧ください。

謝辞

Anyscale社および Databricks で、Ray とMLflow の開発に携わる以下チームメンバーのご協力に感謝します。
チームメンバー：Richard Liaw、Kai Fricke、Eric Liang、Simon Mo、Edward Oakes、Michael Galarnyk、Jules Damji、Sid Murching、Ankit Mathur