概要

Databricks SQL Serverless Warehouseを効率的に管理するために、Databricks SDKを使用して1分で自動停止するよう設定する方法をご紹介します。この@taka_yayoi(Takaaki Yayoi) さんの次の記事にて、1 分で自動停止するDatabricks SQL Serverless Warehouse を作成できることを知りました。Databricks SDK により、新規で作成する方法と既存の Databricks SQL Serverless Warehouse の設定を変更する方法を検証しました。

アイドル時間1分で停止するDatabricks SQLサーバレスウェアハウスの作成 #Databricks - Qiita

検証コードと実行結果

事前準備

まず、databricks-sdkライブラリを最新版に更新します。以下のコードを実行して、Python環境を再起動します。

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

1. 新規で1 分で自動停止するDatabricks SQL Serverless Warehouse を作成する方法

新規に1分で自動停止する Warehouse を作成するには、auto_stop_minsパラメータに1を設定します。以下のコードを実行して Warehouse を作成します。

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

created = w.warehouses.create(
    name=f"qiita-wh-01",
    cluster_size="2X-Small",
    max_num_clusters=1,
    auto_stop_mins=1,
    enable_serverless_compute=True,
).result()

Databricks Workspace にて SQL warehouses の設定を確認するとAuto Stopが 1 分となっていることを確認できました。

2. 既存の Databricks SQL Serverless Warehouse の設定を変更する方法

設定の変更を検証するために、auto_stop_minsに10に設定した Databricks SQL Serverless Warehouse を作成します。

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

created = w.warehouses.create(
    name=f"qiita-wh-02",
    cluster_size="2X-Small",
    max_num_clusters=1,
    auto_stop_mins=10,
    enable_serverless_compute=True,
).result()

設定を変更する場合には Databricks SQL Serverless Warehouse の ID が必要となるため取得します。

warehouses = w.warehouses.list()

# 特定のSQLウェアハウスの名称からIDを検索
warehouse_name = "qiita-wh-02"
warehouse_id = None
for warehouse in warehouses:
    if warehouse.name == warehouse_name:
        warehouse_id = warehouse.id
        break

# IDを表示
if warehouse_id:
    print(f"'ID: {warehouse_id}")
else:
    print(f"ID が見つかりませんでした。")

Databricks SQL Serverless Warehouse の自動停止時間を 1 分に設定します。

# 設定値を更新
w.warehouses.edit(
    warehouse_id,
    auto_stop_mins=1,
)

Databricks Workspace にて SQL warehouses の設定を確認するとAuto Stopが 1 分となっていることを確認できました。