MLflow Pipelinesを試してみた

こちらの記事にあるように、今週発表されたMLflowの新機能であるMLflow Pipelinesがもう試せるということだったので、早速動かしてみました。

https://qiita.com/taka_yayoi/items/29a54ba9a5df4a38b3ad

こちらのリポジトリにサンプルが格納されています。

https://github.com/mlflow/mlp-regression-template

# MLflow Pipelinesとは

以下は私の理解です。

- [MLflow](https://mlflow.org/)自体、機械学習モデル本体、ハイパーパラメータ、メトリクスを記録することができ、モデルの再現性確保、可搬性の確保に役立つものです。
- しかし、機械学習モデルはモデル単体で役立つのではなく、トレーニングにおけるデータ取り込み、特徴量エンジニアリング、モデルデプロイ、推論の実施と一連のパイプラインを経ていくことで価値を生み出します。
- このパイプライン、問題のタイプにもよりますが定型化できる部分が多い(回帰問題なら、データ取り込み、データ分割、データ変換、トレーニング、評価、デプロイなど)のですが、これまでは毎回ファイルやノートブックでパイプラインの処理を記述(あるいは既存のものを流用)していました。
- パイプラインが定型化できるのであれば、これをテンプレートとして提供するとともに、パイプライン自体もMLflowでトラッキング、再活用できるようにすればいいのでは？という考えのもと生み出されたのが[MLflow Pipelines](https://mlflow.org/docs/latest/pipelines.html#)です。

# リポジトリのクローン

ローカルでも実行できるとのことですが、ここではDatabricksの[Repos](https://qiita.com/taka_yayoi/items/b89f199ff0d3a4c16140)を使って、ワークスペース上にクローンします。

![Screen Shot 2022-07-01 at 12.17.49.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/13225974-f92d-c22c-4712-580acbf43d2c.png)

これで全てクローンされました。
![Screen Shot 2022-07-01 at 12.19.11.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/5f27de33-758e-dc51-cce0-31d8ed753e6e.png)

# ノートブックの実行

`notebooks/databricks`を開きます。実行にはDatabricksランタイム11.0以降が動作しているクラスターが必要なので、クラスターを作成して起動しておき、ノートブックをアタッチします。

![Screen Shot 2022-07-01 at 12.20.20.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/befe8505-bf0c-cb88-be92-695ec5e9a212.png)

## ライブラリのインストール

以下を実行して、MLflow Pipelinesと依存ライブラリをインストールします。依存ライブラリは上のステップでリポジトリをクローンした際に含まれている`requirements.txt`を使用します。

```
%pip install mlflow[pipelines]
%pip install -r ../requirements.txt
```

## クライアントの作成

以下を実行することで、プロファイル`databricks`を用いてMLflow Pipelinesクライアントを作成します。

```py:Python
from mlflow.pipelines import Pipeline

p = Pipeline(profile="databricks")
```

現時点では回帰問題のパイプラインのみが提供されています。今後増えていくことになるのでしょう。

![Screen Shot 2022-07-01 at 12.24.29.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/f2139a6f-c3f3-c348-d5c4-82713bd1d1b6.png)

なお、プロファイルは`profiles`にYAMLで格納されています。いろいろなパラメーターを指定できます。

```yaml:databricks.yaml
# [Recommended] Uncomment fields below to set an MLflow experiment to track the pipeline execution.
# If unset, a default experiment based on runtime context will be created.
# experiment:
  # name: "/Shared/sklearn_regression_experiment"

# Use a section of the TLC Trip Record Dataset for model development
# (https://www1.nyc.gov/site/tlc/about/tlc-trip-record-data.page)
INGEST_DATA_LOCATION: dbfs:/databricks-datasets/nyctaxi-with-zipcodes/subsampled
# Specify the format of the dataset
INGEST_DATA_FORMAT: spark_sql
# Override the default train / validation / test dataset split ratios
SPLIT_RATIOS: [0.75, 0.125, 0.125]
```

パイプラインの実装はルートのpipeline.yamlに記述されています。

```yaml:pipeline.yaml
# `pipeline.yaml` is the main configuration file for an MLflow Pipeline.
# Required pipeline parameters should be defined in this file with either concrete values or
# variables such as {{ INGEST_DATA_LOCATION }}.

# Variables must be dereferenced in a profile YAML file, located under `profiles/`.
# See `profiles/local.yaml` for example usage. One may switch among profiles quickly by
# providing a profile name such as `local` in the Pipeline object constructor:
# `p = Pipeline(profile="local")`
#
# NOTE: YAML does not support tabs for indentation. Please use spaces and ensure that all YAML
# files are properly formatted.

template: "regression/v1"
# Specifies the dataset to use for model development
data:
  # Dataset locations on the local filesystem are supported, as well as HTTP(S) URLs and
  # any other remote locations resolvable by MLflow, such as those listed in
  # https://mlflow.org/docs/latest/tracking.html#artifact-stores
  location: {{INGEST_DATA_LOCATION}}
  # Beyond `parquet` datasets, the `spark_sql` and `delta` formats are also natively supported for
  # use with Spark
  format: {{INGEST_DATA_FORMAT|default('parquet')}}
  # Datasets with other formats, including `csv`, can be used by implementing and
  # specifying a `custom_loader_method`
  custom_loader_method: steps.ingest.load_file_as_dataframe
  # If the `spark_sql` `format` is specified, the `sql` entry is used to specify a SparkSQL
  # statement that identifies the dataset to use
  sql: SELECT * FROM delta.`{{INGEST_DATA_LOCATION}}`
  # If the `delta` `format` is specified, you can also configure the Delta table `version` to read
  # or the `timestamp` at which to read data
  # version: 2
  # timestamp: 2022-06-01T00:00:00.000Z
# Specifies the name of the column containing targets / labels for model training and evaluation
target_col: "fare_amount"
steps:
  split:
    # Train/validation/test split ratios
    split_ratios: {{SPLIT_RATIOS|default([0.75, 0.125, 0.125])}}
    # Specifies the method to use to perform additional processing and cleaning on split datasets
    post_split_method: steps.split.process_splits
  transform:
    # Specifies the method that defines the data transformations to apply during model inference
    transformer_method: steps.transform.transformer_fn
  train:
    # Specifies the method that defines the estimator type and parameters to use for model training
    estimator_method: steps.train.estimator_fn
  evaluate:
    # Sets performance thresholds that a trained model must meet in order to be eligible for
    # registration to the MLflow Model Registry
    validation_criteria:
      - metric: root_mean_squared_error
        threshold: 10
      - metric: mean_absolute_error
        threshold: 50
      - metric: weighted_mean_squared_error
        threshold: 20
  register:
    # Specifies the name of the Registered Model to use when registering a trained model to
    # the MLflow Model Registry
    model_name: "taxi_fare_regressor"
    # Indicates whether or not a model that fails to meet performance thresholds should still
    # be registered to the MLflow Model Registry
    allow_non_validated_model: true
metrics:
  # Defines custom performance metrics to compute during model training and evaluation
  custom:
    - name: weighted_mean_squared_error
      # Specifies the name of the function in `steps/custom_metrics.py` to use to compute the metric
      function: weighted_mean_squared_error
      greater_is_better: False
  # Sets the primary metric to use to evaluate model performance. This primary metric is used
  # to sort MLflow Runs corresponding to the pipeline in the MLflow Tracking UI
  primary: "root_mean_squared_error"
```

## パイプラインの確認

`inspect`でパイプラインを確認することができます。

```py:Python
p.inspect()
```

パイプラインのグラフ構造が表示されます。
![Screen Shot 2022-07-01 at 12.26.40.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/e2aefc86-b9e4-6ba6-1e01-b781cd178c02.png)

ステップ(縦二重線の四角)をクリックすると詳細情報が表示されます。
![Screen Shot 2022-07-01 at 12.27.37.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/247a012a-7415-6f8a-ef37-92430ebd2d40.png)

## パイプラインの実行

`run`でパイプラインを実行します。引数にステップ名を指定することで特定のステップを実行できます。引数無しの場合はパイプライン全体が実行されます。

以下を実行することで、データを取り込みます。

```py:Python
p.run("ingest")
```

実行結果には「カード」というものが表示され、処理結果をインタラクティブに確認することができます。
![Screen Shot 2022-07-01 at 12.34.41.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/4af81444-e85e-078c-30e1-f07f686205f6.png)

以降同じようにステップを順に実行していきます。

```py:Python
p.run("split")
```
![Screen Shot 2022-07-01 at 12.38.03.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/3088dae7-768d-3530-9b83-51f3b1e21fa9.png)

```py:Python
p.run("transform")
```
![Screen Shot 2022-07-01 at 12.38.33.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/e39f50fb-4059-4a29-67bf-db3cff191a0c.png)

トレーニングを行うと、当たり前ですがMLflowにトラッキングされます。

```py:Python
p.run("train")
```
![Screen Shot 2022-07-01 at 12.39.13.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/3e26d07a-0926-763f-9503-114bdf0ab028.png)

```py:Python
p.run("evaluate")
```
![Screen Shot 2022-07-01 at 12.39.59.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/409c43cb-1f8f-6a64-9a2b-c08c9e3ff974.png)

最後にモデルレジストリにモデルを登録します。

```py:Python
p.run("register")
```
![Screen Shot 2022-07-01 at 12.40.47.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/8786774c-80f4-16b2-1caf-f4b9a4a388dc.png)

`inspect`で各ステップを調査することもできます

```py:Python
p.inspect("train")
```

# パイプラインのトラッキング

個人的にすごいと思ったのはこちらです。これまではモデルにまつわる情報がトラッキングされていたのですが、Pipelinesを使うことでモデルを作り、利用するMLパイプラインをまるっとトラッキングできるようになります。
![Screen Shot 2022-07-01 at 12.43.50.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/fed25b4b-dba1-b3ff-5293-497561d334ae.png)

機械学習モデルがどのように構築されたのか、どのように使われるのかなど、モデルの「文脈」をこれまで以上に捕捉できるようになり、機械学習モデルを本格運用に持っていく際にも、既存のパイプラインをほぼほぼそのままに移行できるようになるMLflow Pipelines、機械学習のビジネスへの本格運用をさらに支援するテクノロジーだと思いました。

### Databricks 無料トライアル

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