DatabricksにおけるMLflow 3.0のサポート - デプロイメントジョブ編

Posted at 2025-04-10

こちらをウォークスルーします。

こちらのマニュアルに沿ってサンプルノートブックを動かします。

注意
執筆時点ではベータ版です。

MLflow 3.0

こちらにもまとめていますが、エクスペリメントページにおけるモデルタブの追加、Logged Modelエンティティの追加、モデルレジストリUIの改善などが含まれています。特に大きな変更点は、情報アーキテクチャが見直されたことだと思います。これまでは、MLflowラン(トレーニングの実行単位)にモデルやアーティファクト、メトリクス、パラメータなどが紐づけられていました。生成AIのように一つのモデルに繰り返し評価が行われる場合や、ディープラーニングのように複数のモデルチェックポイントができる場合には、この情報アーキテクチャでは困難さが生じていたため、今回の変更でモデルに複数のメトリクス、アーティファクトが紐づけられるようになっています。

有効化

ワークスペースのプレビューメニューにアクセスし、MLflow 3を有効化します。

デプロイメントジョブとは？

デプロイジョブを使用すると、新しいモデルバージョンが作成されるたびに評価、承認、デプロイなどのタスクを自動化し、Unity Catalog モデルや Databricks ジョブとシームレスに統合できます。これらのジョブにより、モデルデプロイパイプラインの設定が簡素化され、ヒューマンインザループの承認が組み込まれ、各モデルバージョンの進行状況と履歴コンテキストを明確に可視化できる管理されたワークフローが提供されます。

こちらの機能、懐かしいです。モデル管理がUnity Catalogに統合される前のワークスペースモデルレジストリで提供されていた承認フローの代替です。ワークスペースモデルレジストリでは、モデルのステージ分けができ、StagingからProductionに移行する際に承認者による承認を挟むことができました。Unity Catalogのモデル管理ではこの機能がオミットされていたので、悲しんでいたところです。

早速動かします。

ノートブックの準備

こちらを実行してモデルを作成しておきます。

こちらにある4つのノートブック(評価ノートブック、承認ノートブック、デプロイノートブック、デプロイジョブ作成ノートブック)をワークスペースにインポートします。

デプロイジョブ作成ノートブック

create-deployment-jobは環境に合わせた変更が必要です。

このノートブックを実行すると、MLflow 3.0 デプロイメントジョブとしてテンプレート化された Databricks ジョブが作成されます。このジョブには、評価、承認チェック、およびデプロイメントの3つのタスクがあります。評価タスクはデータセット上でモデルを評価し、承認チェックタスクは UC タグと UC モデル UI の承認ボタンを使用してモデルがデプロイメントのために承認されているかどうかを確認し、デプロイメントタスクはモデルをサービングエンドポイントにデプロイします。

デプロイメントジョブのテンプレートノートブックの例を Databricks ワークスペースにコピーします (AWS | Azure | GCP)。
UC モデルを作成するか、既存のものを使用します。例については、MLflow 3 の例を参照してください (AWS | Azure | GCP)。
ノートブックを実行する前に、次のセルの TODO/値を更新します。
ノートブックを実行した後、作成されたジョブは UC モデルに接続されません。ドキュメントに示されているように、UC モデル UI でジョブを UC モデルに接続する必要があります (AWS | Azure | GCP)。

# TODO: 必要に応じてこれらの値を更新してください
model_name = "takaakiyayoi_catalog.mlflow.ml_model" # 既に作成されたUCモデルの名前
job_name = "example_deployment_job" # デプロイメントジョブの希望する名前

# TODO: 各タスクのノートブックを作成し、以下の無効なパスを置き換えてノートブックパスを入力してください。
# これらのパスは、Databricksワークスペース内の例のデプロイメントジョブテンプレートノートブックからテンプレート化されたノートブックを配置した場所に対応する必要があります。
evaluation_notebook_path = "/Workspace/Users/takaaki.yayoi@databricks.com/20250409_mlflow3.0/evaluation"
approval_notebook_path = "/Workspace/Users/takaaki.yayoi@databricks.com/20250409_mlflow3.0/approval"
deployment_notebook_path = "/Workspace/Users/takaaki.yayoi@databricks.com/20250409_mlflow3.0/deployment"

# 必要な構成を持つジョブを作成して、デプロイメントジョブとしてモデルに接続します
from databricks.sdk import WorkspaceClient
from databricks.sdk.service import jobs

w = WorkspaceClient()
job_settings = jobs.JobSettings(
    name=job_name,
    tasks=[
        jobs.Task(
            task_key="Evaluation",
            notebook_task=jobs.NotebookTask(notebook_path=evaluation_notebook_path),
            max_retries=0,
        ),
        jobs.Task(
            task_key="Approval_Check",
            notebook_task=jobs.NotebookTask(
                notebook_path=approval_notebook_path,
                base_parameters={"approval_tag_name": "{{task.name}}"}
            ),
            depends_on=[jobs.TaskDependency(task_key="Evaluation")],
            max_retries=0,
        ),
        jobs.Task(
            task_key="Deployment",
            notebook_task=jobs.NotebookTask(notebook_path=deployment_notebook_path),
            depends_on=[jobs.TaskDependency(task_key="Approval_Check")],
            max_retries=0,
        ),
    ],
    parameters=[
        jobs.JobParameter(name="model_name", default=model_name),
        jobs.JobParameter(name="model_version", default=""),
    ],
    queue=jobs.QueueSettings(enabled=True),
    max_concurrent_runs=1,
)

created_job = w.jobs.create(**job_settings.__dict__)
print("ジョブ名 " + job_name + " を使用して、UCモデル " + model_name + " にデプロイメントジョブを接続します。UCモデルUIで確認してください。")
print("\n参考のため、ジョブIDは: " + str(created_job.job_id))
print("\nドキュメント: \nAWS: https://docs.databricks.com/aws/mlflow/deployment-job#connect \nAzure: https://learn.microsoft.com/azure/databricks/mlflow/deployment-job#connect \nGCP: https://docs.databricks.com/gcp/mlflow/deployment-job#connect")

実行するとジョブが作成されます。

ジョブ名 example_deployment_job を使用して、UCモデル takaakiyayoi_catalog.mlflow.ml_model にデプロイメントジョブを接続します。UCモデルUIで確認してください。

参考のため、ジョブIDは: 663294281963588

ドキュメント: 
AWS: https://docs.databricks.com/aws/mlflow/deployment-job#connect 
Azure: https://learn.microsoft.com/azure/databricks/mlflow/deployment-job#connect 
GCP: https://docs.databricks.com/gcp/mlflow/deployment-job#connect