From 242415e1c755588b7c596cda3ef9ef41ba463311 Mon Sep 17 00:00:00 2001
From: Tanya Borisova <tborisova@microsoft.com>
Date: Tue, 8 Aug 2023 22:57:31 +0000
Subject: [PATCH] Add autoscale config

---
 config.sample.yaml                     | 3 +++
 infrastructure/transform/databricks.tf | 4 ++++
 infrastructure/transform/variables.tf  | 4 ++++
 3 files changed, 11 insertions(+)

diff --git a/config.sample.yaml b/config.sample.yaml
index 9cf2e979..b9c28d71 100644
--- a/config.sample.yaml
+++ b/config.sample.yaml
@@ -62,6 +62,9 @@ transform:  # Optional
       local_disk_min_size: 600
       category: "Memory Optimised"
     autotermination_minutes: 120
+    autoscale:
+      min_workers: 1
+      max_workers: 3
     init_scripts:
       - /workspaces/FlowEHR/transform/sample_init_script.sh
 
diff --git a/infrastructure/transform/databricks.tf b/infrastructure/transform/databricks.tf
index 50cee19a..e314ded3 100644
--- a/infrastructure/transform/databricks.tf
+++ b/infrastructure/transform/databricks.tf
@@ -70,6 +70,10 @@ resource "databricks_cluster" "cluster" {
   spark_version           = data.databricks_spark_version.latest.id
   node_type_id            = data.databricks_node_type.node_type.id
   autotermination_minutes = var.transform.databricks_cluster.autotermination_minutes
+  autoscale {
+    min_workers = var.transform.databricks_cluster.autoscale.min_workers
+    max_workers = var.transform.databricks_cluster.autoscale.max_workers
+  }
 
   spark_conf = merge(
     # Secrets for SQL Feature store
diff --git a/infrastructure/transform/variables.tf b/infrastructure/transform/variables.tf
index abda900a..3a89bea4 100644
--- a/infrastructure/transform/variables.tf
+++ b/infrastructure/transform/variables.tf
@@ -146,6 +146,10 @@ variable "transform" {
       }), {}),
       autotermination_minutes = optional(number, 0),
       init_scripts            = optional(list(string), [])
+      autoscale = optional(object({
+        min_workers = optional(number, 0)
+        max_workers = optional(number, 0)
+      }), {})
     }), {})
   })
   default = {