Databricks Jobsの実行結果であるHTMLファイルをMLflowに登録する方法

概要

Databricks Jobsの実行結果であるHTMLファイルをMLflowに登録する方法を紹介します。

本手順をMLOpsにおけるモデル学習のパイプラインに組み込むことで、MLモデルの品質向上に繋がります。
たとえば、下記画像のようにMicrosoft社のInterpretMLの実行結果をMLモデルと共に管理することであり、機械学習モデルの解釈可能性が高くなります。

詳細は下記のGithub pagesのページをご確認ください。

コードを実行したい方は、下記のdbcファイルを取り込んでください。

https://github.com/manabian-/databricks_tecks_for_qiita/blob/main/tecks/put_notebookresult_to_mlflow_artifact/dbc/put_notebookresult_to_mlflow_artifact.dbc

検証環境

databricks runtime: 8.1.x-scala2.12
Python version: 3.8.8
pyspark version: 3.1.1.dev0

手順

1. 変数へ基本的な情報をセット

import json

# DatabricksのワークスペースURLを設定
browserHostName = json.loads(dbutils.notebook.entry_point
                                      .getDbutils()
                                      .notebook()
                                      .getContext()
                                  .toJson()
                                 )['tags']['browserHostName']

# プロトコルを追加
db_url = f"https://{browserHostName}" # DatabricksのURL 

# トークン。本来は、dbutils.secrets.get(scope = <SOCPE_NAME>, key = <TOKEN>)
token = "dapia5ac9cadd5988cd3874cef7113a92b41" 

notebook_name = "code" # 実行するノートブック名を指定
notebook_path = "./includes/" + notebook_name # 実行するノートブックの相対パスと名称を指定
output_dir_notebookresult = "/dbfs/FileStore/qiita/put_notebookresult_to_mlflow_artifact" # 実行結果の保存先を指定

# MLflowのArtifactts内で保存先となっているディレクトリを設定
mlflow_html_path = "notebok_results"

2. Databricks REST APIによりJobsの実行結果をHTMLファイルとして出力する方法を実行

投稿済みの下記記事と同様の手順を実施します。

• Databricks REST APIによりJobsの実行結果をHTMLファイルとして出力する方法 - Qiita

%run ./includes/get_notebook_results_by_restapi

3. MLflowにDatabricksノートブックの実行結果を登録

# MLflowのArtifactts内で保存先とするディレクトリを設定
mlflow_html_path = "notebok_results" 

import mlflow

with mlflow.start_run(nested=True):
    # MLflowの登録する際のファイル名を設定
    mlflow_html_name = run_id + ".html"

    with open(output_file_path) as f:
        html_code = f.read()

    with open(mlflow_html_name, 'w') as f:
        f.write(html_code)

    # MLflow TrackingによりArtifactsにDatabricks実行結果のHTMLファイルを登録
    mlflow.log_artifact(mlflow_html_name, mlflow_html_path)

3. 実行結果を確認

MLflowにおけるRUNの実行結果のUIにて、{jobs_run_id}.htmlのファイルが格納されていることを確認。