Databricks SDK for Pythonでジョブを作成する際のクラスターの再利用

こちらの続編です。

https://qiita.com/taka_yayoi/items/79e83aa1124226504892

Pythonからマルチタスクジョブを構成できるのは便利ですが、GUIでは[クラスターを再利用](https://docs.databricks.com/ja/workflows/jobs/use-compute.html#use-shared-job-clusters)することができます。では、Databricks SDK for Pythonからできるのでしょうか？

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

まず、ジョブを作成する[メソッド](https://databricks-sdk-py.readthedocs.io/en/latest/autogen/jobs.html#databricks.sdk.service.jobs.JobsAPI.create)を確認すると、パラメータに`job_clusters`があります。

> job_clusters – List[[JobCluster](https://databricks-sdk-py.readthedocs.io/en/latest/autogen/jobs.html#databricks.sdk.service.jobs.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`のリストとなっています。[こちら](https://databricks-sdk-py.readthedocs.io/en/latest/autogen/jobs.html#databricks.sdk.service.jobs.JobCluster)を確認します。

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

引数に`job_cluster_key`と`new_cluster`がありますのでピンときました。ジョブの[タスクの仕様](https://databricks-sdk-py.readthedocs.io/en/latest/autogen/jobs.html#databricks.sdk.service.jobs.Task)で、パラメータに`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`ということでした。

図にします。

![Screenshot 2024-01-19 at 10.29.47.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/d953795e-e3a9-994c-ce98-99eb19f8bdb2.png)

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

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

```py
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()
```

こちらが肝になります。

```py
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")
```

期待した通りの構成でジョブを作成することができました！
![Screenshot 2024-01-19 at 10.19.15.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/458be2e3-1662-fa33-20fc-531121d87bb5.png)

ご活用ください！

### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)