こちらの続編です。

Pythonからマルチタスクジョブを構成できるのは便利ですが、GUIではクラスターを再利用することができます。では、Databricks SDK for Pythonからできるのでしょうか？

ということで調べました。結論できました。

まず、ジョブを作成するメソッドを確認すると、パラメータにjob_clustersがあります。

job_clusters – List[JobCluster] (optional) A list of job cluster specifications that can be shared and reused by tasks of this job. Libraries cannot be declared in a shared job cluster. You must declare dependent libraries in task settings.

このジョブのタスクで共有、再利用できるジョブクラスターの仕様のリストです。

まさにこれです。でも、どうやって指定するのか？よく見るとJobClusterのリストとなっています。こちらを確認します。

class databricks.sdk.service.jobs.JobCluster(job_cluster_key: 'str', new_cluster: 'Optional[compute.ClusterSpec]' = None)

引数にjob_cluster_keyとnew_clusterがありますのでピンときました。ジョブのタスクの仕様で、パラメータにjob_cluster_keyがあります。

job_cluster_key: str | None = None
If job_cluster_key, this task is executed reusing the cluster specified in job.settings.job_clusters.

job_cluster_keyが指定されている場合には、job.settings.job_clustersで指定されているクラスターを再利用してこのタスクが実行されます。

まさしくこれでした。job.settings.job_clustersがなんだろうと思いましたが、jobのインスタンス作成時にしている上記のパラメータjob_clustersということでした。

図にします。

上の構成を実装するコードが以下となります。

%pip install databricks-sdk --upgrade
dbutils.library.restartPython()

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.jobs import Task, NotebookTask, Source, TaskDependency
from databricks.sdk.service.compute import ClusterSpec, RuntimeEngine

w = WorkspaceClient()

こちらが肝になります。

from databricks.sdk import WorkspaceClient
from databricks.sdk.service.jobs import Task, NotebookTask, Source, JobSettings, JobCluster
from databricks.sdk.service.compute import ClusterSpec, RuntimeEngine

w = WorkspaceClient()

job_name = "cluster re-use job"
description = "タスクのクラスターを再利用"
notebook_path = (
    "/Users/takaaki.yayoi@databricks.com/20231210_workflow/ml-quickstart-training-jpn"
)
task_key = "first-task"

# 最新MLランタイム
latest = w.clusters.select_spark_version(ml=True)

print("ジョブの作成を試みます。お待ちください...\n")

j = w.jobs.create(
    name=job_name,
    # ジョブで共有するジョブクラスターのリスト
    job_clusters=[
        JobCluster(
            job_cluster_key="first_job_cluster",
            new_cluster = ClusterSpec(num_workers=1, spark_version=latest, node_type_id="i3.xlarge"),
        ),
        JobCluster(
            job_cluster_key="second_job_cluster",
            new_cluster = ClusterSpec(num_workers=1, spark_version=latest, node_type_id="i3.xlarge"),
        )
    ],
    # ジョブを構成するタスクのリスト、タスクから上記ジョブクラスターを job_cluster_key で指定
    tasks=[
        Task(
            description="first-task",
            job_cluster_key="first_job_cluster", # first_job_clusterを使用
            notebook_task=NotebookTask(
                base_parameters={"n_estimators": 50},
                notebook_path=notebook_path,
                source=Source("WORKSPACE"),
            ),
            task_key="first-task",
        ),
        Task(
            description="second-task",
            job_cluster_key="second_job_cluster", # second_job_clusterを使用
            notebook_task=NotebookTask(
                base_parameters={"n_estimators": 50},
                notebook_path=notebook_path,
                source=Source("WORKSPACE"),
            ),
            task_key="second-task",
        ),

        Task(
            description="third-task",
            job_cluster_key="first_job_cluster", # first_job_clusterを使用
            notebook_task=NotebookTask(
                base_parameters={"n_estimators": 50},
                notebook_path=notebook_path,
                source=Source("WORKSPACE"),
            ),
            task_key="third-task",
        ),
    ],
)

print(f"View the job at {w.config.host}/#job/{j.job_id}\n")

期待した通りの構成でジョブを作成することができました！

ご活用ください！

はじめてのDatabricks

はじめてのDatabricks

Databricks無料トライアル

Databricks無料トライアル