【Databricks】1000人分のワークスペースとクラスターを1日で作る方法【SDK for Python】

はじめに

今年1月から2月にかけて、「第2回 金融データ活用チャレンジ」が一般社団法人金融データ活用推進協会 (FDUA) 様の主催で開催され、3月8日に表彰式が行われました。

このイベントは、金融業界において発展途上であるデータ活用の分野における人材を育成・発掘するための分析コンペティションであり、Databricks などのビッグデータ分析ツールを希望した参加者に提供いたしました。ありがたいことに、Databricks の利用を希望し、ワークスペースを提供させていただいたユーザー数は最終的に1000人を突破しました。

今回、Databricks 様、そして Microsoft 様のパートナーとして、Microsoft Azure 上の環境にて提供する、大規模なワークスペースの環境立ち上げに携わらせていただくことになり、何百人ものユーザーが利用する（かもしれない）環境で各個人用にリソースを用意することになりました。

通常、Databricks でユーザーを登録したり、クラスターを作成したりするには GUI からの作業で十分なのですが、それを100人、1000人単位のユーザーに対して行うのはかなりの作業量になります。この作業を自動化するための手段として、Databricks SDK for Python を利用することを決めました。

以前の記事では、 ARM テンプレートを用いた Azure リソース作成の自動化を行いましたが、今回はこの時できなかったワークスペース内部の構築も自動化を目指すということになります。

環境の要件

コンペ開始前に提示された要件をまとめると、以下のようになりました。

• 100人単位の規模のユーザーが利用するらしい（前年度の規模感から）
• Unity Catalog でデータベースのアクセス制御を行う（元データのテーブルは全員アクセス可、など）
• 各ユーザーにプライベートな、自作の加工データなどを保存する領域が必要（スキーマで分離することにした）

コンペが始まってしばらくすると、参加者の方からそれまで用意していた共用クラスターに普段使っている機械学習用のライブラリが入っていないなどの声が挙がりました。
そのため機械学習関連のライブラリが充実した ML Runtime を利用できるようにしたいのですが、現状では Unity Catalogを利用でき、なおかつ ML Runtime をインストールするには Personal Computing のポリシーでクラスターを作成しなければいけません。ということで以下の要件も加わりました。

• ML Runtime 入りのクラスターを、参加者と同数作成して各ユーザーに提供する

準備・前提条件

まず、Databricks のワークスペースを作成し、Unity Catalog を有効化する必要があります。以前とは異なり、ワークスペースは１つだけ作るので、今回はテンプレートを使わず手動で立ち上げました。ちなみに、最近になって自分で関連するリソースを作る必要もなく自動で Unity Catalog が有効化されるようになったらしいです（リージョンによる）。

Databricks は、2023 年 11 月 9 日、アカウント全体で順次ロールアウトする Unity Catalog の新しいワークスペースの自動有効化を開始しました。 自動的に有効になったワークスペースには、次のプロパティがあります。

• 自動的にプロビジョニングされた Unity Catalog メタストア (ワークスペース リージョンに Unity Catalog メタストアが既に存在する場合を除く)。
• カタログや外部データベース接続を作成する機能など、ワークスペース管理者の既定のアクセス許可。
• メタストア管理者がいない (既存の Unity Catalog メタストアが使用され、メタストア管理者が既に割り当てられている場合を除く)。
• マネージド テーブルとマネージド ボリュームにメタストア レベルのストレージがない (メタストア レベルのストレージを持つ既存の Unity Catalog メタストアが使用されている場合を除く)。

ワークスペース カタログは、最初にプロビジョニングされるときに、ワークスペースにちなんだ名前になります。

(Unity Catalog の設定と管理 | Mircosoft Learnより引用)

Azure Databricks に登録する際には Entra ID でユーザーを作成する必要がありますが、この際はAzure側の機能でcsvファイルを用いた一括作成が可能です。この時に使ったCSVファイルを取っておいたり、アカウント名(onmicrosoft.comで終わるメールアドレス)に命名規則を付けておくと後の自動化がしやすくなります。

SDK のクライアントを作成

今回 SDK for Python を利用する環境はワークスペース上のノートブック上に作ることにしました。 バージョン13.2以降のクラスターにはデフォルトでインストールされています。

注意点としてSDKのバージョンが最新版（この作業をしていた時点ではver0.18、公開前に確認したら0.21になってた…）でなければドキュメント通りに動作しない機能が多かったです。最新版のクラスター（14.3）でもインストールされているのは0.1.6とかなり古いバージョンであり、ノートブックでpipを実行するなりクラスター設定をいじるなりして SDK をアップデートしないと動きません。

Python SDKを操作するためには、認証情報を含めてクライアントを作成する必要があります。クライアントには'WorkspaceClient'と'AccountClient'の２種類があり、主にユーザー(Unity Catalog)関連の操作はアカウントの、主にクラスター関連の操作はワークスペースのクライアントを使うことになります。
ノートブック上で実行した場合、ワークスペースの認証はデフォルトで行われているのですが、アカウント側については手動で設定を構成する必要があります。認証方法は複数ありますので、公式ドキュメントを参照してください。ここではAzureのサービスプリンシパルを利用する方法を使います。

from databricks.sdk import WorkspaceClient

w = WorkspaceClient()

from databricks.sdk import AccountClient

a = AccountClient(
    host          = "https://accounts.azuredatabricks.net",
    account_id    = os.getenv("ACCOUNT_ID"),
    client_id     = os.getenv("AZURE_CLIENT_ID"),
    client_secret = os.getenv("AZURE_CLIENT_SECRET"),
)

Databricks のクラスターに SDK がインストール済みという点からノートブックで実行することを選びましたが、アップデートの必要があることやクライアント情報を置かないといけないセキュリティ上の観点からローカルで書いた方がよかったかもしれないと終わってから思いました。

ユーザーを作成する

次に、 Entra ID 上のユーザーをワークスペースに割り当てる作業です。こちらは基本的にSCIMプロビジョニングを利用した自動化が推奨されていますが、アプリへのユーザーの割り当てのためにユーザー名を一つずつクリックしていくという作業があり、アカウントコンソールの GUI をポチポチしながら登録するのと労力があまり変わらないんじゃないか？という理由で、自動化することにしました。

Azure Portal 上で CSV ファイルを用いた一括作成を行っているなど、ユーザーのIDと表示名のリストがファイル化されていることを前提とします。そのリストをワークスペースにアップロードし、ノートブック上にデータとして読み込みます。

user_file = 'userlist.csv'

users = []
with open(user_file, mode='r') as f:
    reader = csv.DictReader(f)
    for row in reader:
        user = {
            'id': row['ユーザー名 [userPrincipalName] 必須'],
            'name': row['名前 [displayName] 必須']
        }
        users.append(user)

（列名は Entra の一括作成テンプレートの書式に従った物です）

必要なのは、Databricks アカウントにユーザーを作成することと、ワークスペースにユーザーを割り当てることです。
割り当てのためにはワークスペースIDが必要になりますが、この部分はhttps://adb-xxxxxxxxxxxxxx.yy.azuredatabricks.netのxの部分と同じです。
この部分をユーザー名とIDを入力したら実行する関数を作ります。

from databricks.sdk.service import iam

WORKSPACE_ID = 'xxxxxxxxxxxxxx'

def create_user(user_name,display_name):
    u = a.users.create(user_name=user_name,display_name=display_name) # ユーザーを作成する
    a.workspace_assignment.update(permissions = [iam.WorkspacePermission.USER], workspace_id=WORKSPACE_ID, principal_id=u.id) # ワークスペースにユーザーを割り当てる
    return u

for user in users:
    u = create_user(user_name=user['名前 [displayName] 必須'],display_name=user['ユーザー名 [userPrincipalName] 必須'])

ところで、ワークスペースへの登録はユーザーにワークスペースへの 'USER' としての権限を与えるという形式で行います。こういったロール関連の変数は SDK では enum 形式で定義されており、普通に文字列で 'USER' と渡したりなどするとエラーになります。

カタログスキーマを割り当てる

まず、事前にUnity Catalogを有効化し、全員がアクセス可能なカタログを一つ作成します。ここは GUI でも済ませられると思います。
この下に、各ユーザーのみが個別にアクセスできるようなスキーマを作ります。

スキーマの作成は SDK の機能に含まれていますが、スキーマへの権限の付与は SQL で行う必要があるので、spark.sql を呼び出します。

import re

# カタログスキーマの自動作成と権限付与
def create_schema(id):
    usrn = re.match(r'(.+)@dbworkspace.onmicrosoft.com', id) # ユーザー名をアカウント名から抽出
    w.schemas.create(name=usrn,catalog_name='competition_data') # スキーマを作成する
    spark.sql(f'GRANT ALL PRIVILEGES ON SCHEMA competition_data.{usrn} TO `{id}`')


for user in users:
    create_schema(user['id'])

クラスターを作成する

ML Runtime は共用クラスターでは利用できないので、Personal Compute ポリシーを適用したクラスターを人数分作成し、各ユーザーがそれらにアクセスできるようにする必要があります。

また、重要な点として 作成時にユーザーを指定する Personal クラスターでも、該当ユーザーにアクセス権限を与えないとその人がアクセスできません。 この点に気づかず、提供直後になってクラスターにアクセスできないとクレームを貰ってしまいました・・・

改めて要件を確認すると、

• 組み込みの「Personal Compute」ポリシーを利用
• ランタイムは ML の最新版
• ノードタイプは最小の Standard_DS3_v2 を選択、シングルノード構成
• シングルユーザーに対して「再起動可」の権限を与える（これより高いとユーザーに構成を変えられる可能性があり、低いとクラスターが自由に起動出来ない）

まずはユーザー名とクラスター名を指定してクラスターの作成と権限付与を行う関数を定義します。

from databricks.sdk.service import compute
from databricks.sdk.service.compute import ClusterAccessControlRequest

def create_cluster(cluster_name,user_name,node = 'Standard_DS3_v2'):

    latest = w.clusters.select_spark_version(latest=True,ml=True) 

    clstr = w.clusters.create(cluster_name=cluster_name,
                            spark_version=latest,
                            autotermination_minutes=15,
                            single_user_name= user_name,
                            spark_conf={
                                "spark.databricks.cluster.profile": "singleNode",
                                "spark.master": "local[*, 4]"
                            },
                            node_type_id=node,
                            custom_tags={
                                "ResourceClass": "SingleNode"
                            },
                            policy_id="XXXXXXXXXXXXXXXX", # Personal Compute のポリシーID。ワークスペース毎に異なる
                            data_security_mode=compute.DataSecurityMode.SINGLE_USER
                            )
    w.clusters.set_permissions(cluster_id=clstr.cluster_id, access_control_list=[ClusterAccessControlRequest(user_name=user_name, permission_level=compute.ClusterPermissionLevel.CAN_RESTART)])
    w.clusters.delete(cluster_id=clstr.cluster_id)

クラスターへの権限付与の記法はやや複雑です。権限の種類を示す変数にClusterAccessControlRequestというオブジェクトを要求されており、その情報にユーザーIDと、権限の区分としてClusterPermissionLevel.CAN_RESTARTを指定しなければいけません。

この部分はclusterモジュールの公式ドキュメントに具体例が書かれておらず、パラメータ指定の方法（具体的にはページ内に説明のないClusterAccessControlRequestの扱い方）が分かりづらく、個人的にハマりポイントでした。

クラスターの作成と割り当てを、各ユーザーに対して実行します。
ユーザーIDに命名規則を定めていれば、もう少しシンプルな形のループで書くこともできます。

for user in users:
    clstrn = re.match(r'(.+)@dbworkspace.onmicrosoft.com', user['id']) +"'s cluster" # ユーザー名を抽出し、クラスター名を指定
    create_cluster(clstrn, user['id'])     

実行すると、UIから作成したのと同じようにクラスターが自動的に起動します。そのままだとコストが余計にかかってしまうため、w.clusters.delete()でクラスターを自動的に停止する処理を入れます。ややこしいですがこれはクラスターをただ停止するだけの関数で、削除したい場合はw.clusters.permanent_delete()を使うらしいです。

おわりに

以上が、今回のコンペを通して作り上げた自動化の手法になります。

実際の作業では、新規ユーザーの情報が毎日更新される中で、ドキュメントを調べながら手探りでのスクリプト作成に取り組んでいました。（1000人分のワークスペースとクラスターを1日で作った訳じゃないです）設定したリソースに不備があって参加者の皆様にご不便をかけてしまった場面もあり、もっと余裕をもってリサーチしておけばよかったと思います。

今後また Databricks の環境構築に携わった際は、今回の経験を生かしてより迅速にセキュアな環境を提供できるように頑張ります。