Databricks Jobsの実行結果を保存したMLflowからWorkspaceにノートブックとしてインポートする方法 #Python

概要

Databricks Jobsの実行結果を保存したMLflowからWorkspaceにノートブックとしてインポートする方法を紹介します。

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

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

https://github.com/manabian-/databricks_tecks_for_qiita/blob/main/tecks/get_notebook_from_mlflow_artifacts/dbc/get_notebook_from_mlflow_artifacts.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/get_notebook_from_mlflow_artifacts" # 実行結果の保存先を指定

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

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

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

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

%run ./includes/put_notebookresult_to_mlflow_artifact

3. MLflow ArtifactsからHTMLファイルを取得

import json 
# 現在のノートブックのディレクトリ（ノートブック名を含む）を設定
notebook_current_path = json.loads(dbutils.notebook.entry_point
                                      .getDbutils()
                                      .notebook()
                                      .getContext()
                                  .toJson()
                                 )['extraContext']['notebook_path']

# ノートブック名を除外した現在のノートブックのディレクトリを設定
notebook_current_dir = notebook_current_path[:notebook_current_path.rfind('/')]

# ダウンロード先のディレクトリを設定
local_path = '/dbfs/FileStore/qiita/get_notebook_from_mlflow_artifacts'

# MLflow Artifactから取得対象のディレクトリを設定
mlflow_artifacts_folder = 'notebok_results'

# インポート先のノートブックのディレクトリを設定
notebook_dir = f'{notebook_current_dir}/notebok_results' # フォルダを事前に作成する必要あり

import mlflow
from mlflow.tracking import MlflowClient

## ノートブックを格納するフォルダを作成
dbutils.fs.mkdirs(local_path.replace('/dbfs', 'dbfs:'))

## MLflowからArtifactsをダウンロード
client = MlflowClient()
downloaded_path = client.download_artifacts(mlflow_run_id, 
                                            mlflow_artifacts_folder, 
                                            local_path
                                           )

# ダウンロードしたファイル名を取得
import glob
html_file_path = glob.glob(f'{downloaded_path}/*')[0]

4. htmlファイルの内容のbase64文字列を取得

# htmlファイルの内容のbase64文字列を取得
import base64
import os

# htmlファイルの名前を取得
html_file_name = os.path.splitext(os.path.basename(html_file_path))[0]

# インポート先のパスとファイル名をセット
notebook_import_path = f'{notebook_dir}/runid_{html_file_name}'

# htmlファイルの内容を取得
with open(html_file_path) as f:
    file_value = f.read()

html_file_encode_base64 = base64.b64encode(file_value.encode())

5. REST API経由でDatabricks Workspaceにノートブックインポート

# REST API経由でDatabricks Workspaceにノートブックインポート
import requests

response = requests.post(
    f'{db_url}/api/2.0/workspace/import',
    headers={
        "Authorization": f'Bearer {token}'},
    json={
        "content": html_file_encode_base64,
        "path": notebook_import_path,
        "overwrite": 'true',
        "format": 'HTML',
    }
)

if response.status_code == 200:
    print("success")
elif response.status_code == 403:
    print(response.text)
else:
    print("Error geting the job: {0}: {1}".format(response.json()["error_code"],response.json()["message"]))

6. 実行結果を確認

実行結果のノートブックがインポートされていることを確認。

Databricks Jobsの実行結果を保存したMLflowからWorkspaceにノートブックとしてインポートする方法

概要