Add ProvisioningModelMix to InstanceFlexibilityPolicy to support mixing of spot and on-demands instances for Dataproc cluster

Author: ahzaz

mmv1/third_party/terraform/services/dataproc/resource_dataproc_cluster.go

@@ -1227,6 +1227,35 @@ func ResourceDataprocCluster() *schema.Resource {
 														},
 													},
 												},
+												"provisioning_model_mix": {
+													Type:     schema.TypeList,
+													Optional: true,
+													Computed: true,
+													AtLeastOneOf: []string{
+														"cluster_config.0.preemptible_worker_config.0.instance_flexibility_policy.0.provisioning_model_mix.0.standard_capacity_base",
+														"cluster_config.0.preemptible_worker_config.0.instance_flexibility_policy.0.provisioning_model_mix.0.standard_capacity_percent_above_base",
+													},
+													MaxItems:    1,
+													Description: `Defines how Dataproc should create VMs with a mixture of provisioning models.`,
+													Elem: &schema.Resource{
+														Schema: map[string]*schema.Schema{
+															"standard_capacity_base": {
+																Type:         schema.TypeInt,
+																Optional:     true,
+																Description:  `The base capacity that will always use Standard VMs to avoid risk of more preemption than the minimum capacity you need.`,
+																ValidateFunc: validation.IntAtLeast(0),
+															},
+
+															"standard_capacity_percent_above_base": {
+																Type:         schema.TypeInt,
+																Optional:     true,
+																Description:  `The percentage of target capacity that should use Standard VM. The remaining percentage will use Spot VMs.`,
+																ForceNew:     true,
+																ValidateFunc: validation.IntBetween(0, 100),
+															},
+														},
+													},
+												},
 											},
 										},
 									},
@@ -2413,6 +2442,9 @@ func expandPreemptibleInstanceGroupConfig(cfg map[string]interface{}) *dataproc.
 			if v, ok := flexibilityPolicy["instance_selection_list"]; ok {
 				icg.InstanceFlexibilityPolicy.InstanceSelectionList = expandInstanceSelectionList(v)
 			}
+			if v, ok := flexibilityPolicy["provisioning_model_mix"]; ok {
+				icg.InstanceFlexibilityPolicy.ProvisioningModelMix = expandProvisioningModelMix(v)
+			}
 		}
 
 	}
@@ -2444,6 +2476,18 @@ func expandInstanceSelectionList(v interface{}) []*dataproc.InstanceSelection {
 	return instanceSelections
 }
 
+func expandProvisioningModelMix(v interface{}) *dataproc.ProvisioningModelMix {
+	pmm := v.([]interface{})
+	if len(pmm) > 0 {
+		provisioningModelMix := pmm[0].(map[string]interface{})
+		return &dataproc.ProvisioningModelMix{
+			StandardCapacityBase:             int64(provisioningModelMix["standard_capacity_base"].(int)),
+			StandardCapacityPercentAboveBase: int64(provisioningModelMix["standard_capacity_percent_above_base"].(int)),
+		}
+	}
+	return nil
+}
+
 func expandMasterInstanceGroupConfig(cfg map[string]interface{}) *dataproc.InstanceGroupConfig {
 	icg := &dataproc.InstanceGroupConfig{}
 
@@ -3184,8 +3228,13 @@ func flattenPreemptibleInstanceGroupConfig(d *schema.ResourceData, icg *dataproc
 			disk["local_ssd_interface"] = icg.DiskConfig.LocalSsdInterface
 		}
 		if icg.InstanceFlexibilityPolicy != nil {
-			instanceFlexibilityPolicy["instance_selection_list"] = flattenInstanceSelectionList(icg.InstanceFlexibilityPolicy.InstanceSelectionList)
-			instanceFlexibilityPolicy["instance_selection_results"] = flattenInstanceSelectionResults(icg.InstanceFlexibilityPolicy.InstanceSelectionResults)
+			if icg.InstanceFlexibilityPolicy.InstanceSelectionList != nil {
+				instanceFlexibilityPolicy["instance_selection_list"] = flattenInstanceSelectionList(icg.InstanceFlexibilityPolicy.InstanceSelectionList)
+				instanceFlexibilityPolicy["instance_selection_results"] = flattenInstanceSelectionResults(icg.InstanceFlexibilityPolicy.InstanceSelectionResults)
+			}
+			if icg.InstanceFlexibilityPolicy.ProvisioningModelMix != nil {
+				instanceFlexibilityPolicy["provisioning_model_mix"] = flattenProvisioningModelMix(icg.InstanceFlexibilityPolicy.ProvisioningModelMix)
+			}
 		}
 	}
 
@@ -3222,6 +3271,14 @@ func flattenInstanceSelectionResults(isr []*dataproc.InstanceSelectionResult) []
 
 }
 
+func flattenProvisioningModelMix(pmm *dataproc.ProvisioningModelMix) []map[string]interface{} {
+	provisioningModelMix := map[string]interface{}{}
+	provisioningModelMix["standard_capacity_base"] = pmm.StandardCapacityBase
+	provisioningModelMix["standard_capacity_percent_above_base"] = pmm.StandardCapacityPercentAboveBase
+
+	return []map[string]interface{}{provisioningModelMix}
+}
+
 func flattenMasterInstanceGroupConfig(d *schema.ResourceData, icg *dataproc.InstanceGroupConfig) []map[string]interface{} {
 	disk := map[string]interface{}{}
 	data := map[string]interface{}{}

mmv1/third_party/terraform/services/dataproc/resource_dataproc_cluster_test.go.tmpl

@@ -534,6 +534,29 @@ func TestAccDataprocCluster_spotWithInstanceFlexibilityPolicy(t *testing.T) {
 	})
 }
 
+func TestAccDataprocCluster_spotOnDemandMixing(t *testing.T) {
+	t.Parallel()
+
+	rnd := acctest.RandString(t, 10)
+	var cluster dataproc.Cluster
+	acctest.VcrTest(t, resource.TestCase{
+		PreCheck:                 func() { acctest.AccTestPreCheck(t) },
+		ProtoV5ProviderFactories: acctest.ProtoV5ProviderFactories(t),
+		CheckDestroy:             testAccCheckDataprocClusterDestroy(t),
+		Steps: []resource.TestStep{
+			{
+				Config: testAccDataprocCluster_spotOnDemandMixing(rnd),
+				Check: resource.ComposeTestCheckFunc(
+					testAccCheckDataprocClusterExists(t, "google_dataproc_cluster.spot_mixing", &cluster),
+					resource.TestCheckResourceAttr("google_dataproc_cluster.spot_mixing", "cluster_config.0.preemptible_worker_config.0.preemptibility", "SPOT"),
+					resource.TestCheckResourceAttr("google_dataproc_cluster.spot_mixing", "cluster_config.0.preemptible_worker_config.0.instance_flexibility_policy.0.provisioning_model_mix.0.standard_capacity_base", "1"),
+					resource.TestCheckResourceAttr("google_dataproc_cluster.spot_mixing", "cluster_config.0.preemptible_worker_config.0.instance_flexibility_policy.0.provisioning_model_mix.0.standard_capacity_percent_above_base", "50"),
+				),
+			},
+		},
+	})
+}
+
 func TestAccDataprocCluster_spotWithAuxiliaryNodeGroups(t *testing.T) {
 	t.Parallel()
 
@@ -1955,6 +1978,50 @@ resource "google_dataproc_cluster" "spot_with_instance_flexibility_policy" {
 	`, rnd)
 }
 
+func testAccDataprocCluster_spotOnDemandMixing(rnd string) string {
+	return fmt.Sprintf(`
+resource "google_dataproc_cluster" "spot_mixing" {
+  name   = "tf-test-dproc-%s"
+  region = "us-central1"
+
+  cluster_config {
+    gce_cluster_config {
+      internal_ip_only = false
+    }
+    master_config {
+      num_instances = "1"
+      machine_type  = "e2-medium"
+      disk_config {
+        boot_disk_size_gb = 35
+      }
+    }
+
+    worker_config {
+      num_instances = "2"
+      machine_type  = "e2-medium"
+      disk_config {
+        boot_disk_size_gb = 35
+      }
+    }
+
+    preemptible_worker_config {
+      num_instances = "3"
+      preemptibility = "SPOT"
+      disk_config {
+        boot_disk_size_gb = 35
+      }
+	  instance_flexibility_policy {
+        provisioning_model_mix {
+          standard_capacity_base = 1
+          standard_capacity_percent_above_base = 50
+        }
+      }
+    }
+  }
+}
+	`, rnd)
+}
+
 func testAccDataprocCluster_withAuxiliaryNodeGroups(rnd string) string {
 	return fmt.Sprintf(`
 resource "google_dataproc_cluster" "with_auxiliary_node_groups" {