Databricksクラスターでアクセスモードを共有にすることで、複数人で一つのクラスターを共有することができますが、いくつかの制限を意識する必要があります。その中でも大きなものは機械学習ランタイムを使えないというものです。
これは、そのままではMLflowでモデルの記録ができないことを意味します。機械学習ランタイムにはMLflowがプレインストールされています。
しかし、Databricksランタイム13.0以降ではシングルノードで動作する機械学習ライブラリ(scikit-learnなど)が動作するようになりました。
この記事では、実際に共有クラスターでMLflowを活用してみます。ここでは、Databricksランタイム13.3LTSを使っています。
まず、手動でMLflowをインストールします。dbutils.library.restartPython()でPythonカーネルを再起動してください。これを行わないと後のステップでエラーになります。
%pip install mlflow
dbutils.library.restartPython()
ライブラリをインポートします。
import mlflow
import numpy as np
import pandas as pd
import sklearn.datasets
import sklearn.metrics
import sklearn.model_selection
import sklearn.ensemble
データをロードして準備します。
# データのロードと前処理
white_wine = pd.read_csv("/dbfs/databricks-datasets/wine-quality/winequality-white.csv", sep=';')
red_wine = pd.read_csv("/dbfs/databricks-datasets/wine-quality/winequality-red.csv", sep=';')
white_wine['is_red'] = 0.0
red_wine['is_red'] = 1.0
data_df = pd.concat([white_wine, red_wine], axis=0)
# ワイン品質に基づいた分類ラベルの定義
data_labels = data_df['quality'] >= 7
data_df = data_df.drop(['quality'], axis=1)
# 80/20でトレーニング/テストデータセットを分割
X_train, X_test, y_train, y_test = sklearn.model_selection.train_test_split(
data_df,
data_labels,
test_size=0.2,
random_state=1
)
MLflowのオートロギングを有効化します。
# このノートブックでのMLflow autologgingを有効化
mlflow.autolog()
モデルのトレーニングとMLflowによるトラッキングを行います。
with mlflow.start_run(run_name='gradient_boost') as run:
model = sklearn.ensemble.GradientBoostingClassifier(random_state=0)
# モデル、パラメータ、トレーニングメトリクスが自動でトラッキングされます
model.fit(X_train, y_train)
predicted_probs = model.predict_proba(X_test)
roc_auc = sklearn.metrics.roc_auc_score(y_test, predicted_probs[:,1])
# テストデータに対するAUCスコアは自動で記録されないので、手動で記録します
mlflow.log_metric("test_auc", roc_auc)
print("Test AUC of: {}".format(roc_auc))
2023/11/21 00:32:38 WARNING mlflow.utils.autologging_utils: MLflow autologging encountered a warning: "/databricks/python/lib/python3.10/site-packages/_distutils_hack/__init__.py:33: UserWarning: Setuptools is replacing distutils."
Test AUC of: 0.8834365701533531
トラッキングされました！
こちらの記事にありますように、共有クラスター自身、今後も改善してまいりますので、是非ご活用ください！