Databricks社内環境をTerraformで構築した話

はじめに

こんにちは、DATUM STUDIO株式会社の亀井です。弊社でもDatabricksを利用した業務が増加してきたため、Databricksの社内環境を構築しました。環境構築にはTerraformを採用しました。本記事では、その構築手順や工夫した点、苦労した点などを共有します。

なぜTerraformを採用したの？

Databricksの環境（ワークスペース）構築方法は、公式ドキュメントによると以下の選択肢があります：

• AWSクイックスタートを使用したワークスペースの作成
• ワークスペースを手動で作成する（新しいDatabricksアカウント）
• ワークスペースを手動で作成する（既存のDatabricksアカウント）
• アカウントAPIを使用したワークスペースの作成
• Terraformを用いたDatabricksワークスペースの作成

その中でTerraformを選択した理由は以下の通りです：

• サンドボックス環境運用において、面倒くさそうな頻発するユーザ追加・削除作業が簡潔に対応可能だから
• Databricksの裏側で動作しているAWSリソースをコードベースで確認したかったから
• Terraformの勉強になりそうだったから

構築したリソース

1. アカウント（手動作成）
2. ワークスペース
3. カタログ（ユーザごとに作成）
4. 共用のAll Purpose Cluster
5. 共用のSQL Cluster
6. コストダッシュボード

構築手順

1. アカウント作成

こちらのサイトからDatabricksアカウントを作成します。account_idは控えておきます。

2. Service Principalの作成

1. Account UI → User Management → Service Principals → 「Add Service Principal」から作成します。作成したService PrincipalにはAccount Adminを付与しておきます。その後、「Service Principal」→「作成したService Principal」→「Generate Secret」でシークレットを生成します。生成したシークレットは安全な場所に保存しておきます。

3. ワークスペース作成

こちらのガイドをもとにワークスペースを作成します。
Terraform Registryに、上記のコードがまるっと入ったaws-workspace-basicモジュールがあるので、こちらを使用しても良いです。（こちらのexampleに使用例があるので、そのまま引用してOK。）
terraform.tfvarsには、上記で作ったアカウントのaccount_idと、上記で作ったサービスプリンシパルのid, secretを追加します。

4. Unity Catalog Metastore作成 & ユーザ作成

Deploying pre-requisite resources and enabling Unity CatalogをもとにUnity Catalog Metastoreを作ります。
ただ、create-users-and-groupsの前後に不足している工程があるので補足しておきます。

1. databricks_storage_credentialにprovider = databricks.workspaceを追加

2. ユーザをワークスペースにアサイン

   • 「databricks_users」「databricks_metastore_admins」で追加したユーザをワークスペースにアサインする必要があります
     • これがないと、users, metastore_adminsで追加したユーザがワークスペースにログインできません
     • databricks_mws_permission_assignment Resourceを使って、ユーザをワークスペースに追加しましょう

     resource "databricks_mws_permission_assignment" "add_admin_group" {
     provider = databricks.mws
     for_each = toset(concat(var.databricks_users, var.databricks_metastore_admins))
     workspace_id = module.aws-workspace-basic.workspace_id
     principal_id = databricks_user.unity_users[each.value].id
     permissions = ["ADMIN"]
     }
     

3. Unity Catalog metastore 作成 & ユーザ作成公式ドキュメント通りに作業していると、Configure external locations and credentialsあたりでError: cannot create external location: User does not have CREATE EXTERNAL LOCATION on Metastoreエラーが発生します。

   • これは、2. Service Principalの作成 で作成したService Principalが、新規作成した Unity Catalog metasoreに対してCREATE EXTERNAL LOCATION権限が付与されていないことが原因で発生するエラーです。
   • 2. Service Principalの作成 で作成したService Principalをadmin_groupsに追加するコードを書きましょう

   data "databricks_service_principal" "this" {
     provider  = databricks.mws
     application_id = var.client_id
   }
   
   resource "databricks_group_member" "add_service_pricipal" {
     provider  = databricks.mws
     group_id  = databricks_group.admin_group.id
     member_id = data.databricks_service_principal.this.id
   }
   

5. ユーザごとのCatalog作成

↑で作成したものは、"sandbox"という一つのcatalogを作成します。
catalogはユーザごとに作成したいので、for_eachで書き換えます。

resource "databricks_catalog" "sandbox" {
  for_each     = toset(concat(var.databricks_users, var.databricks_metastore_admins))
  provider     = databricks.workspace
  storage_root = "s3://${aws_s3_bucket.external.id}/some"
  name         = each.key
  comment      = "this catalog is managed by terraform"
  properties = {
    purpose = "testing"
  }
  depends_on = [databricks_metastore_assignment.default_metastore]
}

resource "databricks_grants" "catalog_permissions" {
  for_each     = toset(concat(var.databricks_users, var.databricks_metastore_admins))
  provider = databricks.workspace
  catalog  = databricks_catalog.sandbox[each.key].name

  grant {
    principal  = each.value
    privileges = ["ALL_PRIVILEGES"]
  }
}

6. 共用コンピュートリソース作成

公式ドキュメントを参照し、共有クラスタやSQL Warehouseを構築します。以下は例です：

resource "databricks_cluster" "standard_cluster" {
  provider   = databricks.workspace
  cluster_name            = "Standard Cluster"
  spark_version           = data.databricks_spark_version.latest_lts.id
  autotermination_minutes = 45
  autoscale {
    min_workers = 1
    max_workers = 2
  }
  policy_id =  databricks_cluster_policy.fair_use.id
  apply_policy_default_values = true
}

resource "databricks_sql_endpoint" "standard_sql_warehouse" {
  provider   = databricks.workspace
  name             = "Standard SQL Warehouse"
  cluster_size     = "2X-Small"
  warehouse_type   = "PRO"
  enable_photon    = true
  enable_serverless_compute = false
  auto_stop_mins = 45
}

resource "databricks_sql_endpoint" "cost_report_sql_warehouse" {
  provider   = databricks.workspace
  name             = "Cost Report SQL Warehouse"
  cluster_size     = "2X-Small"
  warehouse_type   = "PRO"
  enable_photon    = true
  enable_serverless_compute = false
  auto_stop_mins = 10
}

7. コストダッシュボード作成

環境の公開後、無駄なクラスタ使用を監視するため、公式ガイドを参考にUsage Reportを管理者に毎日送信するよう設定しました。

工夫した点

• ユーザをメールアドレスで追加するだけで、ワークスペースにユーザ追加 + 個人用カタログ作成までを自動化できるように構成しました。
• コスト監視ダッシュボードを低コストで稼働できるように、SQL Warehouseを別で用意
  • 起動時間を最小限に抑えることで、定常監視によるコストを最小限に抑えました。

苦労した点

• ワークスペースデプロイからCatalog作成までの一貫したドキュメントがない
• Databricks Terraform Providerのドキュメントが微妙に間違えている
  • 4. Unity Catalog Metastore作成 & ユーザ作成で申し上げたように、ドキュメント通りだとうまくいかない点がところどころで見受けられます。ある程度Databricksの使用経験があり、必要なリソースの知識がないと構成が難しい可能性があります。

実際に得られた（or今後得られそうな）メリット

• ユーザ管理作業（追加、削除）が簡単
• Databricksインフラ構築の勉強になった
• 他プロジェクトにも転用可能

総括

Databricks社内Sandbox環境の構築について共有しました。
Databricksのインフラ構築に興味がある方はぜひ挑戦してみてください！