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

In this blog post, we’re announcing two new integrations with Ray and MLflow: Ray Tune+MLflow Tracking and Ray Serve+MLflow Models, which together make it much easier to build machine learning (ML) models and take them to production.

These integrations are available in the latest Ray wheels. You can follow the instructions here to pip install the nightly version of Ray and take a look at the documentation to get started. They will also be in the next Ray release — version 1.2

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

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

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

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

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

Ray Tune is a library for executing hyperparameter tuning experiments at any scale and can save you tens of hours in training time.

With Ray Tune you can:

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

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

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

Ray Serve is an easy-to-use scalable model serving library that:

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

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

Ray Tune and Ray Serve make it easy to distribute your ML development and deployment, but how do you manage this process? This is where MLflow comes in.

During experiment execution, you can leverage MLflow’s Tracking API to keep track of the hyperparameters, results, and model checkpoints of all your experiments, as well as easily visualize and share them with other team members. And when it comes to deployment, MLflow Models provides standardized packaging to support deployment in a variety of different environments.

主要なポイント

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 integrates with MLflow Tracking API to easily record information from your distributed tuning run to an MLflow server.

There are two APIs for this integration: an MLflowLoggerCallback and an mlflow_mixin.

With the MLflowLoggerCallback, Ray Tune will automatically log the hyperparameter configuration, results, and model checkpoints from each run in your experiment to MLflow.

from ray.tune.integration.mlflow import MLflowLoggerCallback tune.run( train_fn, config={ # define search space here "parameter_1": tune.choice([1, 2, 3]), "parameter_2": tune.choice([4, 5, 6]), }, callbacks=[MLflowLoggerCallback( experiment_name="experiment1", save_artifact=True)]) 

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

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

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

from ray.tune.integration.mlflow import mlflow_mixin @mlflow_mixin def train_fn(config): ... mlflow.log_metric(...) 

You can check out the documentation here for full runnable examples and more information.

Ray Serve + MLflow モデル

MLflow models can be conveniently loaded as python functions, which means that they can be served easily using Ray Serve. The desired version of your model can be loaded from a model checkpoint or from the MLflow Model Registry by specifying its Model URI. Here’s how this looks:

import ray from ray import serve import mlflow.pyfunc class MLflowBackend: def __init__(self, model_uri): self.model = mlflow.pyfunc.load_model(model_uri=model_uri) async def __call__(self, request): return self.model.predict(request.data) ray.init() client = serve.start() # This can be the same checkpoint that was saved by MLflow Tracking model_uri = "/Users/ray_user/my_mlflow_model" # Or you can load a model from the MLflow model registry model_uri = "models:/my_registered_model/Production" client.create_backend("mlflow_backend", MLflowBackend, model_uri) 

まとめ

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

Try running this example in the Databricks Community Edition (DCE) with this notebook. Note: This Ray Tune + MLflow extension has only been tested on DCE runtimes 7.5 and MLR 7.5.

次のステップ

Give this integration a try by pip install the latest Ray nightly wheels and pip install mlflow. Or try this notebook on DCE. Also, stay tuned for a future deployment plugin that further integrates Ray Serve and MLflow Models.

参考資料：

• Check out the documentation for the Ray Tune + MLflow Tracking integration and the runnable example
• See how you can use this integration to tune and autolog a Pytorch Lightning model. Example.

謝辞

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