RayとMLflow: 分散機械学習アプリケーションの本格運用

[How the Integrations Between Ray & MLflow Aids Distributed ML Production \- The Databricks Blog](https://databricks.com/blog/2021/02/03/ray-mlflow-taking-distributed-machine-learning-applications-to-production.html)の翻訳です。

> このゲスト記事はRay.ioのコントリビュータである、Anyscaleのソフトウェアエンジニアである[Amog Kamsetty](https://www.linkedin.com/in/amogkamsetty/)と[Archit Kulkarni](https://www.linkedin.com/in/archit-kulkarni-a0136a90/)によるものです。

このブログ記事では、[Ray](https://ray.io/)と[MLflow](https://mlflow.org/)の新たなインテグレーションを発表します。[Ray TuneとMLflowトラッキング](https://docs.ray.io/en/master/tune/tutorials/tune-mlflow.html)、[Ray Serve](https://docs.ray.io/en/master/serve/)と[MLflowモデル](https://www.mlflow.org/docs/latest/models.html)を組み合わせることで、機械学習(ML)モデルの構築、本格運用が容易になります。

これらのインテグレーションは最新のRay wheelで利用できます。利用を始めるには[pipによる最新バージョンのRayのインストール手順](https://docs.ray.io/en/master/installation.html#daily-releases-nightlies)を参考にしてください。これらのインテグレーションは、Rayの次期リリース(バージョン 1.2)にも搭載される予定です。

![](https://databricks.com/wp-content/uploads/2021/02/ray-logo2.png)

我々が目指しているのは、二つのプロジェクトの強みを活用するというものです。これらは、Rayのトレーニング、サービングをスケーリングするための分散ライブラリ、MLflowのエンドツーエンドのモデルライフサイクル管理機能です。

# これらのツールが解決する課題

この新たなインテグレーションの詳細に入る前に、これらのライブラリで何ができるのかを見てみましょう。

## Ray Tuneによる大規模ハイパーパラメーターチューニング

MLモデルのサイズやトレーニング回数の増加によって、単一のマシンで大規模ML実験を行うことは現実的ではなくなっています。たくさんのマシンにまたがって実験を分散させる必要があります。

[Ray Tune](https://docs.ray.io/en/latest/tune/)は、大規模なハイパーパラメーターチューニングを実行するためのライブラリであり、トレーニングに要する時間を大幅に削減することが可能です。

[Ray Tune](https://docs.ray.io/en/latest/tune/)を使用することで、以下のことが可能になります。

- 10行未満のコードでマルチノードのハイパーパラメータ探索を起動できます。
- Pytorch、Tensorflow、MXNet、KerasのようなあらゆるMLフレームワークを使用できます。
- Population Based Training、HyperBand、Asynchronous Successive Halving(ASHA)のような最先端のハイパーパラメータ最適化アルゴリズムを活用できます。

## Ray Serveによる大規模モデルサービング

機械学習モデルを開発した後、多くの場合、予測リクエストに対応するためにモデルを配備することになるかと思います。しかし、MLモデルは多くの計算資源を必要とし、実際の配備においては分散システムによるスケーラビリティが必要となります。

[Ray Serve](https://docs.ray.io/en/master/serve/)は簡単に利用できるスケーラブルなモデルサービングライブラリです。

- 複数マシンのGPUを活用してモデルサービングを簡素化しますので、実環境で必要となる稼働時間、性能要件に応えることができます。
- Pytorch、Tensorflow、MXNet、KerasのようなあらゆるMLフレームワークと連携できます。
- プログラム可能な設定インタフェースを提供します(YAMLやJSONは不要です！)。

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

![](https://databricks.com/wp-content/uploads/2021/02/mlflow-components2.png)
Ray TuneとRay Serveを活用することで、MLの開発、配備が容易にになりますが、このプロセスをどのように管理したらいいのでしょうか？ここで[MLflow](https://mlflow.org/)の出番となります。

実験を行なっている際、[MLflowのトラッキングAPI](https://www.mlflow.org/docs/latest/tracking.html)を活用して、実験における全てのハイパーパラメーター、結果、モデルのチェックポイントを記録することができ、可視化を行うことで他のチームメンバーに共有することも可能です。そして、モデルを配備する際にMLflowは、様々な環境への配備をサポートする標準的なパッケージを提供します。

## ライブラリのまとめ

[Ray Tune](https://docs.ray.io/en/latest/tune/)、[Ray Serve](https://docs.ray.io/en/master/serve/)、[MLflow](https://mlflow.org/)を組み合わせることで、MLエンジニアの運用負荷を削減し、メインタスクであるMLモデル、アルゴリズムの構築に注力できます。

それでは、どのようにこれらをインテグレーションするのかを見ていきましょう。

# Ray Tune + MLflowトラッキング

分散チューニングから容易にMLflowのトラッキングサーバーに情報を記録するために、[MLflow Tracking API](https://www.mlflow.org/docs/latest/tracking.html)と[Ray Tune](https://docs.ray.io/en/latest/tune/)を組み合わせます。

インテグレーションには二つのAPIを使用します。[MLflowLoggerCallback](https://docs.ray.io/en/master/tune/tutorials/tune-mlflow.html)と[mlflow\_mixin](https://docs.ray.io/en/master/tune/tutorials/tune-mlflow.html)です。

[MLflowLoggerCallback](https://docs.ray.io/en/master/tune/tutorials/tune-mlflow.html)を用いることで、[Ray Tune](https://docs.ray.io/en/latest/tune/)は処理におけるハイパーパラメーターの設定、結果、モデルのチェックポイントを[MLflow](https://mlflow.org/)に自動的に記録するようになります。

**Python**

```python
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が、異なるハイパーパラメーターの設定で様々なトレーニング(Run)を並列に実行している様子を見ることができます。これらの実行結果はMLflowのUIでも確認することができます。MLflow UIでは記録されたメトリクスを可視化することも可能です。MLflowのトラッキングサーバーがリモートにある場合には、他の方が実験結果や成果物にアクセスすることも可能です。

<iframe src="https://player.vimeo.com/video/507304409?title=0&byline=0&portrait=0" width="640" height="360" frameborder="0" allow="autoplay; fullscreen; picture-in-picture" allowfullscreen></iframe>
<p><a href="https://vimeo.com/507304409">MLflow UI Ray Tune Run</a> from <a href="https://vimeo.com/databricks">Databricks</a> on <a href="https://vimeo.com">Vimeo</a>.</p>

Ray Tuneにお任せするのではなく、記録したい情報を自分で定義したいのであれば、mlflow_mixin APIを使用することができます。

MLflowメソッドを呼び出すトレーニング関数にデコレーターを追加します。

**Python**

```python
from ray.tune.integration.mlflow import mlflow_mixin

@mlflow_mixin
def train_fn(config):
    ...
    mlflow.log_metric(...)
```

試すことのできるサンプルや詳細情報に関しては、こちらの[ドキュメント](https://docs.ray.io/en/master/tune/tutorials/tune-mlflow.html)を参照ください。

# Ray Serve + MLflowモデル

[MLflowモデル](https://www.mlflow.org/docs/latest/models.html)は、pythonの関数として容易にロードすることができ、[Ray Serve](https://docs.ray.io/en/master/serve/)を用いることで容易に配備を行うことができます。モデルのURIを指定することで、モデルのチェックポイントや、Mlflowのモデルレジストリから必要なバージョンのモデルをロードすることができます。

以下はサンプルです。

**Python**

```python
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を組み合わせることで、分散機械学習アプリケーションの構築と本格運用が容易になります。Ray Serve+MLflowモデルでモデルの大規模配備をシンプルにして、Ray Tune+MLflowトラッキングによって、迅速かつ管理の行き届いた開発、実験が可能になります。

[サンプルノートブック](https://github.com/taka-yayoi/public_repo/tree/main/ray_mlflow)を[Databricksコミュニティエディション](https://community.cloud.databricks.com/)で試してみませんか。

> **注意**
このRay Tune + MLflowの組み合わせは、コミュニティエディションのランタイム7.5、7.5MLでテストされています。

# 次のステップ

最新の[Ray nightly wheels](https://docs.ray.io/en/master/installation.html#daily-releases-nightlies)と`pip install mlflow`を試してみてください。あるいは、コミュニティエディションで[サンプルノートブック](https://github.com/taka-yayoi/public_repo/tree/main/ray_mlflow)を試してみてください。そして、Ray ServeとMLflowモデルのインテグレーションをさらに推し進めるプラグイン開発状況に注目してください。

現時点では以下の情報が利用可能です。

- Ray Tune + MLflowトラッキングのインテグレーションに関する[ドキュメント](https://docs.ray.io/en/master/tune/tutorials/tune-mlflow.html)、[サンプル](https://docs.ray.io/en/master/tune/examples/mlflow_example.html)をご確認下さい。
- **Pytorch Lightning model**の自動ロギング、チューニングを行う際にこのインテグレーションを活用した[サンプル](https://docs.ray.io/en/master/tune/examples/mlflow_ptl_example.html)をご覧ください。

# 謝辞

AnyscaleとDatabricksからRayとMLflowの開発に参画しているメンバーに感謝します。Richard Liaw、Kai Fricke、Eric Liang、Simon Mo、Edward Oakes、Michael Galarnyk、Jules Damji、Sid Murching、Ankit Mathur

### Databricks 無料トライアル

[Databricks 無料トライアル](https://databricks.com/jp/try-databricks)