MLflowによるモデルのトラッキング

2024/4/12に翔泳社より**Apache Spark徹底入門**を出版します！

<a href="https://www.amazon.co.jp/dp/4798182281/"><img src="https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/87910d7c-8a63-c198-8e82-216ea7a8ee07.jpeg" width=300></a>

書籍の[サンプルノートブック](https://github.com/databricks/LearningSparkV2/blob/master/notebooks/LearningSparkv2.dbc)をウォークスルーしていきます。`Python/Chapter11/11-1 MLflow`となります。

翻訳ノートブックのリポジトリはこちら。

https://github.com/taka-yayoi/SparkSample/tree/main

ノートブックはこちら

https://github.com/taka-yayoi/SparkSample/blob/main/Chapter11/11-1%20MLflow.py

MLflowはDatabricks Runtime for MLにプレインストールされています。MLランタイムを使っていない場合には、MLflowをインストールする必要があります。

```py
from pyspark.ml import Pipeline
from pyspark.ml.feature import StringIndexer, VectorAssembler
from pyspark.ml.regression import RandomForestRegressor
from pyspark.ml.evaluation import RegressionEvaluator

filePath = "dbfs:/databricks-datasets/learning-spark-v2/sf-airbnb/sf-airbnb-clean.parquet"
airbnbDF = spark.read.parquet(filePath)
(trainDF, testDF) = airbnbDF.randomSplit([.8, .2], seed=42)

categoricalCols = [field for (field, dataType) in trainDF.dtypes 
                   if dataType == "string"]
indexOutputCols = [x + "Index" for x in categoricalCols]
stringIndexer = StringIndexer(inputCols=categoricalCols, 
                              outputCols=indexOutputCols, 
                              handleInvalid="skip")

numericCols = [field for (field, dataType) in trainDF.dtypes 
               if ((dataType == "double") & (field != "price"))]
assemblerInputs = indexOutputCols + numericCols
vecAssembler = VectorAssembler(inputCols=assemblerInputs, 
                               outputCol="features")

rf = RandomForestRegressor(labelCol="price", maxBins=40, maxDepth=5, 
                           numTrees=100, seed=42)

pipeline = Pipeline(stages=[stringIndexer, vecAssembler, rf])
```

# MLflowによるモデルのトラッキング

```py
import mlflow
import mlflow.spark
import pandas as pd

with mlflow.start_run(run_name="random-forest") as run:
  # パラメーターの記録: Num Trees と Max Depth
  mlflow.log_param("num_trees", rf.getNumTrees())
  mlflow.log_param("max_depth", rf.getMaxDepth())
 
  # モデルの記録
  pipelineModel = pipeline.fit(trainDF)
  mlflow.spark.log_model(pipelineModel, "model")

  # メトリクスの記録: RMSE と R2
  predDF = pipelineModel.transform(testDF)
  regressionEvaluator = RegressionEvaluator(predictionCol="prediction", 
                                            labelCol="price")
  rmse = regressionEvaluator.setMetricName("rmse").evaluate(predDF)
  r2 = regressionEvaluator.setMetricName("r2").evaluate(predDF)
  mlflow.log_metrics({"rmse": rmse, "r2": r2})

  # アーティファクトの記録: 特徴量の重要度スコア
  rfModel = pipelineModel.stages[-1]
  pandasDF = (pd.DataFrame(list(zip(vecAssembler.getInputCols(), 
                                    rfModel.featureImportances)), 
                          columns=["feature", "importance"])
              .sort_values(by="importance", ascending=False))
  # 最初にローカルファイルシステムに書き出し、MLflowにファイルの場所を通知
  pandasDF.to_csv("/tmp/feature-importance.csv", index=False)
  mlflow.log_artifact("/tmp/feature-importance.csv")
```

モデルがトラッキングされます。
![Screenshot 2024-03-27 at 19.02.58.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/6be771b1-0ccb-109e-08d6-3eb73129deb5.png)
![Screenshot 2024-03-27 at 19.03.59.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/a2ebb66e-6712-367e-27f8-c8dd6a0af10a.png)
![Screenshot 2024-03-27 at 19.04.09.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/bf540dc4-d120-491a-b5c0-4d6e6fe390b0.png)


# MLflowClient

[MLflowClient](https://mlflow.org/docs/latest/python_api/mlflow.client.html)は、MLflowエクスペリメント、ラン、モデルバージョン、登録モデルに対するPythonのCRUDインタフェースを提供します。

```py
from mlflow.tracking import MlflowClient

client = MlflowClient()
runs = client.search_runs(run.info.experiment_id,
                          order_by=["attributes.start_time desc"], 
                          max_results=1)
run_id = runs[0].info.run_id
runs[0].data.metrics
```
```
{'rmse': 211.5096898777315, 'r2': 0.22794251914574226}
```

# バッチ予測の生成

バッチ予測を生成するためにモデルをロードしなおしましょう。

```py
# MLflowで保存したモデルのロード
pipelineModel = mlflow.spark.load_model(f"runs:/{run_id}/model")

# 予測の生成
inputPath = "dbfs:/databricks-datasets/learning-spark-v2/sf-airbnb/sf-airbnb-clean.parquet"
inputDF = spark.read.parquet(inputPath)
predDF = pipelineModel.transform(inputDF)
```
```py
display(predDF)
```
![Screenshot 2024-03-27 at 19.02.03.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/5b069606-aed3-aa91-3f10-d09e429332e0.png)

# ストリーミング予測の生成

ストリーミングの予測結果を生成するために同じことを行えます。

```py
# MLflowで保存したモデルのロード
pipelineModel = mlflow.spark.load_model(f"runs:/{run_id}/model")

# シミュレーションされたストリーミングデータをセットアップ
repartitionedPath = "dbfs:/databricks-datasets/learning-spark-v2/sf-airbnb/sf-airbnb-clean-100p.parquet"
schema = spark.read.parquet(repartitionedPath).schema

streamingData = (spark
                 .readStream
                 .schema(schema) # このようにスキーマを設定可能
                 .option("maxFilesPerTrigger", 1)
                 .parquet(repartitionedPath))

# 予測の生成
streamPred = pipelineModel.transform(streamingData)
```
```
# ストリーミング予測結果の表示
display(streamPred)
```

### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

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