Azure spot instances in the instance pool

databricks · Mar 23, 2021 · 2e2a2f6 · 2e2a2f6
1 parent 81fa607
commit 2e2a2f6
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 2 deletions.
diff --git a/compute/resource_cluster.go b/compute/resource_cluster.go
@@ -328,7 +328,10 @@ func modifyClusterRequest(clusterModel *Cluster) {
 		clusterModel.AzureAttributes = nil
 	}
 	if clusterModel.GcpAttributes != nil {
-		clusterModel.GcpAttributes = nil
+		gcpAttributes := GcpAttributes{
+			GoogleServiceAccount: clusterModel.GcpAttributes.GoogleServiceAccount,
+		}
+		clusterModel.GcpAttributes = &gcpAttributes
 	}
 	clusterModel.EnableElasticDisk = false
 	clusterModel.NodeTypeID = ""

diff --git a/compute/resource_cluster_test.go b/compute/resource_cluster_test.go
@@ -1055,7 +1055,7 @@ func TestResourceClusterUpdate_FailNumWorkersZero(t *testing.T) {
 	require.Equal(t, true, strings.Contains(err.Error(), "NumWorkers could be 0 only for SingleNode clusters"))
 }
 
-func TestModifyClusterRequest(t *testing.T) {
+func TestModifyClusterRequestAws(t *testing.T) {
 	c := Cluster{
 		InstancePoolID: "a",
 		AwsAttributes: &AwsAttributes{
@@ -1072,3 +1072,37 @@ func TestModifyClusterRequest(t *testing.T) {
 	assert.Equal(t, "", c.DriverNodeTypeID)
 	assert.Equal(t, false, c.EnableElasticDisk)
 }
+
+func TestModifyClusterRequestAzure(t *testing.T) {
+	c := Cluster{
+		InstancePoolID: "a",
+		AzureAttributes: &AzureAttributes{
+			FirstOnDemand: 1,
+		},
+		EnableElasticDisk: true,
+		NodeTypeID:        "d",
+		DriverNodeTypeID:  "e",
+	}
+	modifyClusterRequest(&c)
+	assert.Nil(t, c.AzureAttributes)
+	assert.Equal(t, "", c.NodeTypeID)
+	assert.Equal(t, "", c.DriverNodeTypeID)
+	assert.Equal(t, false, c.EnableElasticDisk)
+}
+
+func TestModifyClusterRequestGcp(t *testing.T) {
+	c := Cluster{
+		InstancePoolID: "a",
+		GcpAttributes: &GcpAttributes{
+			UsePreemptibleExecutors: true,
+		},
+		EnableElasticDisk: true,
+		NodeTypeID:        "d",
+		DriverNodeTypeID:  "e",
+	}
+	modifyClusterRequest(&c)
+	assert.Equal(t, false, c.GcpAttributes.UsePreemptibleExecutors)
+	assert.Equal(t, "", c.NodeTypeID)
+	assert.Equal(t, "", c.DriverNodeTypeID)
+	assert.Equal(t, false, c.EnableElasticDisk)
+}