fix: personal instance pools var for databricks compute policies (#543)

chanzuckerberg · Nov 16, 2023 · 541f8d3 · 541f8d3
1 parent 074b18c
commit 541f8d3
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 6 deletions.
diff --git a/databricks-default-cluster-policies/main.tf b/databricks-default-cluster-policies/main.tf
@@ -36,6 +36,18 @@ locals {
       "defaultValue" : local.default_cluster_instance_profile_arn
     }
   }
+  personal_instance_pools = var.personal_compute_pool_ids != [] ? {
+    "instance_pool_id" : {
+      "type" : "allowlist",
+      "isOptional": true,
+      "values" : var.personal_compute_pool_ids
+    },
+    "driver_node_type_id": {
+      "type" : "allowlist",
+      "isOptional": true,
+      "values" : var.personal_compute_pool_ids
+    }
+  } : {}
 }
 
 resource "databricks_group" "power_user_group" {
@@ -68,7 +80,7 @@ module "personal_compute_cluster_policy" {
   databricks_workspace_id = var.databricks_workspace_id
   policy_name             = "${var.policy_name_prefix}Personal Compute"
   policy_family_id        = local.default_policy_family_ids["personal_compute"]
-  policy_overrides = merge(local.logging_override, {
+  policy_overrides = merge(local.logging_override, local.personal_instance_pools, {
     "autotermination_minutes" : {
       "type" : "fixed",
       "value" : 120
@@ -83,10 +95,19 @@ module "personal_compute_cluster_policy" {
       "pattern" : "([rcip]+[3-5]+[d]*\\.[0-1]{0,1}xlarge)",
       "hidden" : false
     },
-    "instance_pool_id" : {
-      type: "allowlist",
-      values: ["i3-xlarge-pool"]
-    }
+    "aws_attributes.availability": {
+      "type": "allowlist",
+      "values": [
+        "ON_DEMAND",
+        "SPOT_WITH_FALLBACK"
+      ],
+      "hidden": false
+    },
+    "runtime_engine": {
+      "type": "unlimited",
+      "defaultValue": "STANDARD",
+      "hidden": false
+    },
   })
   grantees = [local.all_users_group_name]
 }

diff --git a/databricks-default-cluster-policies/variables.tf b/databricks-default-cluster-policies/variables.tf
@@ -23,4 +23,10 @@ variable "policy_name_prefix" {
   description = "Prefix for policy names"
   type = string
   default = ""
-}
+}
+
+variable "personal_compute_pool_ids" {
+  description = "List of personal compute pool ids allowed"
+  type = list(string)
+  default = []
+}