How the Integrations Between Ray & MLflow Aids Distributed ML Production - The Databricks Blogの翻訳です。

このゲスト記事はRay.ioのコントリビュータである、AnyscaleのソフトウェアエンジニアであるAmog KamsettyとArchit Kulkarniによるものです。

このブログ記事では、RayとMLflowの新たなインテグレーションを発表します。Ray TuneとMLflowトラッキング、Ray ServeとMLflowモデルを組み合わせることで、機械学習(ML)モデルの構築、本格運用が容易になります。

これらのインテグレーションは最新のRay wheelで利用できます。利用を始めるにはpipによる最新バージョンのRayのインストール手順を参考にしてください。これらのインテグレーションは、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は、様々な環境への配備をサポートする標準的なパッケージを提供します。

ライブラリのまとめ

Ray Tune、Ray Serve、MLflowを組み合わせることで、MLエンジニアの運用負荷を削減し、メインタスクであるMLモデル、アルゴリズムの構築に注力できます。

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

Ray Tune + MLflowトラッキング

分散チューニングから容易にMLflowのトラッキングサーバーに情報を記録するために、MLflow Tracking APIとRay Tuneを組み合わせます。

インテグレーションには二つのAPIを使用します。MLflowLoggerCallbackとmlflow_mixinです。

MLflowLoggerCallbackを用いることで、Ray Tuneは処理におけるハイパーパラメーターの設定、結果、モデルのチェックポイントをMLflowに自動的に記録するようになります。

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のトラッキングサーバーがリモートにある場合には、他の方が実験結果や成果物にアクセスすることも可能です。

MLflow UI Ray Tune Run from Databricks on Vimeo.

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

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

Python

from ray.tune.integration.mlflow import mlflow_mixin

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

試すことのできるサンプルや詳細情報に関しては、こちらのドキュメントを参照ください。

Ray Serve + MLflowモデル

MLflowモデルは、pythonの関数として容易にロードすることができ、Ray Serveを用いることで容易に配備を行うことができます。モデルのURIを指定することで、モデルのチェックポイントや、Mlflowのモデルレジストリから必要なバージョンのモデルをロードすることができます。

以下はサンプルです。

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トラッキングによって、迅速かつ管理の行き届いた開発、実験が可能になります。

サンプルノートブックをDatabricksコミュニティエディションで試してみませんか。

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

次のステップ

最新のRay nightly wheelsとpip install mlflowを試してみてください。あるいは、コミュニティエディションでサンプルノートブックを試してみてください。そして、Ray ServeとMLflowモデルのインテグレーションをさらに推し進めるプラグイン開発状況に注目してください。

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

Ray Tune + MLflowトラッキングのインテグレーションに関するドキュメント、サンプルをご確認下さい。
Pytorch Lightning modelの自動ロギング、チューニングを行う際にこのインテグレーションを活用したサンプルをご覧ください。

謝辞

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

Databricks 無料トライアル

Databricks 無料トライアル