Azure spot instances in the instance pool

databricks · Mar 24, 2021 · 5a804bb · 5a804bb
1 parent cb1f4ea
commit 5a804bb
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)
+}