[CDAP-21269] Support machine-specific disk types in Dataproc FlexVM

cdap-runtime-ext-dataproc/pom.xml

@@ -31,7 +31,7 @@
 
   <properties>
     <google.api.client.version>1.34.0</google.api.client.version>
-    <google.cloud.libraries.bom.version>26.34.0</google.cloud.libraries.bom.version>
+    <google.cloud.libraries.bom.version>26.88.1</google.cloud.libraries.bom.version>
   </properties>
 
 
@@ -127,14 +127,6 @@
       <groupId>org.mockito</groupId>
       <artifactId>mockito-core</artifactId>
     </dependency>
-    <dependency>
-      <groupId>org.powermock</groupId>
-      <artifactId>powermock-module-junit4</artifactId>
-    </dependency>
-    <dependency>
-      <groupId>org.powermock</groupId>
-      <artifactId>powermock-api-mockito2</artifactId>
-    </dependency>
   </dependencies>
 
   <dependencyManagement>

cdap-runtime-ext-dataproc/src/main/java/io/cdap/cdap/runtime/spi/provisioner/dataproc/DataprocClient.java

@@ -76,6 +76,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
@@ -229,41 +230,22 @@ ClusterOperationMetadata createCluster(String name, String imageVersion,
       clusterProperties.put("dataproc:dataproc.monitoring.stackdriver.enable",
           Boolean.toString(conf.isStackdriverMonitoringEnabled()));
 
-      DiskConfig workerDiskConfig = DiskConfig.newBuilder()
-          .setBootDiskSizeGb(conf.getWorkerDiskGb())
-          .setBootDiskType(conf.getWorkerDiskType())
-          .setNumLocalSsds(0)
-          .build();
       InstanceGroupConfig.Builder primaryWorkerConfig = InstanceGroupConfig.newBuilder()
           .setNumInstances(conf.getWorkerNumNodes())
-          .setMachineTypeUri(conf.getWorkerMachineType())
-          .setDiskConfig(workerDiskConfig);
+          .setMachineTypeUri(conf.getWorkerMachineType());
       InstanceGroupConfig.Builder secondaryWorkerConfig = InstanceGroupConfig.newBuilder()
           .setNumInstances(conf.getSecondaryWorkerNumNodes())
           .setMachineTypeUri(conf.getWorkerMachineType())
-          .setPreemptibility(InstanceGroupConfig.Preemptibility.NON_PREEMPTIBLE)
-          .setDiskConfig(workerDiskConfig);
-
-      if (!conf.getWorkerFlexVmMachineTypes().isEmpty()) {
-        InstanceFlexibilityPolicy workerFlexPolicy =
-          createInstanceFlexibilityPolicy(conf.getWorkerFlexVmMachineTypes());
-        primaryWorkerConfig.setInstanceFlexibilityPolicy(workerFlexPolicy);
-        secondaryWorkerConfig.setInstanceFlexibilityPolicy(workerFlexPolicy);
-      }
+          .setPreemptibility(InstanceGroupConfig.Preemptibility.NON_PREEMPTIBLE);
+      setDiskAndFlexVmConfigs(conf.getWorkerFlexVmMachineTypes(), conf.getWorkerFlexVmDiskTypes(),
+          conf.getWorkerDiskType(), conf.getWorkerDiskGb(),
+          primaryWorkerConfig, secondaryWorkerConfig);
 
       InstanceGroupConfig.Builder masterConfig = InstanceGroupConfig.newBuilder()
           .setNumInstances(conf.getMasterNumNodes())
-          .setMachineTypeUri(conf.getMasterMachineType())
-          .setDiskConfig(DiskConfig.newBuilder()
-                         .setBootDiskType(conf.getMasterDiskType())
-                         .setBootDiskSizeGb(conf.getMasterDiskGb())
-                         .setNumLocalSsds(0)
-                         .build());
-      if (!conf.getMasterFlexVmMachineTypes().isEmpty()) {
-        InstanceFlexibilityPolicy masterFlexPolicy =
-          createInstanceFlexibilityPolicy(conf.getMasterFlexVmMachineTypes());
-        masterConfig.setInstanceFlexibilityPolicy(masterFlexPolicy);
-      }
+          .setMachineTypeUri(conf.getMasterMachineType());
+      setDiskAndFlexVmConfigs(conf.getMasterFlexVmMachineTypes(), conf.getMasterFlexVmDiskTypes(),
+          conf.getMasterDiskType(), conf.getMasterDiskGb(), masterConfig);
 
       //Set default concurrency settings for fixed cluster
       if (Strings.isNullOrEmpty(conf.getAutoScalingPolicy())) {
@@ -376,13 +358,90 @@ ClusterOperationMetadata createCluster(String name, String imageVersion,
     }
   }
 
-  private InstanceFlexibilityPolicy createInstanceFlexibilityPolicy(List<String> machineTypes) {
-    return InstanceFlexibilityPolicy.newBuilder()
-      .addInstanceSelectionList(
-        InstanceSelection.newBuilder().addAllMachineTypes(machineTypes).build())
+  /**
+   * Sets the boot disk and Flex VM configs on the given instance groups. Flex VM disk types are
+   * applied per instance selection and take the place of the group level disk config.
+   */
+  private void setDiskAndFlexVmConfigs(List<String> flexVmMachineTypes,
+      List<String> flexVmDiskTypes, String diskType, int diskSizeGb,
+      InstanceGroupConfig.Builder... groups) {
+    for (InstanceGroupConfig.Builder group : groups) {
+      if (flexVmDiskTypes.isEmpty()) {
+        group.setDiskConfig(createDiskConfig(diskType, diskSizeGb));
+      }
+      if (!flexVmMachineTypes.isEmpty()) {
+        group.setInstanceFlexibilityPolicy(
+            createInstanceFlexibilityPolicy(flexVmMachineTypes, flexVmDiskTypes, diskSizeGb));
+      }
+    }
+  }
+
+  /**
+   * Validates that boot disk types pair positionally one-to-one with machine types.
+   */
+  private static void validateFlexVmDiskTypeCount(List<String> machineTypes,
+                                                  List<String> diskTypes) {
+    if (diskTypes.isEmpty()) {
+      return;
+    }
+
+    if (machineTypes.size() != diskTypes.size()) {
+      throw configurationError(machineTypes.size(), diskTypes.size());
+    }
+  }
+
+  private static DataprocRuntimeException configurationError(Object... args) {
+    String errorMessage = String.format("Invalid config. It must list exactly one boot disk type per machine type "
+                                          + "should be in  the same order. Found %d machine type(s) but %d "
+                                          + "disk type(s).",
+                                        args);
+    return new DataprocRuntimeException.Builder()
+      .withErrorCategory(DataprocRuntimeException.ERROR_CATEGORY_PROVISIONING_CONFIGURATION)
+      .withErrorReason(errorMessage)
+      .withErrorMessage(errorMessage)
+      .withErrorType(ErrorType.USER)
       .build();
   }
 
+  /**
+   * Creates one instance selection per distinct boot disk type, where the Nth machine type uses the
+   * Nth disk type. An empty {@code diskTypes} puts every machine type in a single selection that
+   * relies on the group level disk config.
+   */
+  private InstanceFlexibilityPolicy createInstanceFlexibilityPolicy(List<String> machineTypes,
+      List<String> diskTypes, int diskSizeGb) {
+    if (diskTypes.isEmpty()) {
+      return InstanceFlexibilityPolicy.newBuilder()
+          .addInstanceSelectionList(
+              InstanceSelection.newBuilder().addAllMachineTypes(machineTypes).build())
+          .build();
+    }
+
+    validateFlexVmDiskTypeCount(machineTypes, diskTypes);
+
+    Map<String, List<String>> machineTypesByDiskType = new LinkedHashMap<>();
+    for (int i = 0; i < machineTypes.size(); i++) {
+      machineTypesByDiskType.computeIfAbsent(diskTypes.get(i), k -> new ArrayList<>())
+          .add(machineTypes.get(i));
+    }
+
+    InstanceFlexibilityPolicy.Builder policy = InstanceFlexibilityPolicy.newBuilder();
+    machineTypesByDiskType.forEach((diskType, types) ->
+        policy.addInstanceSelectionList(InstanceSelection.newBuilder()
+            .addAllMachineTypes(types)
+            .setDiskConfig(createDiskConfig(diskType, diskSizeGb))
+            .build()));
+    return policy.build();
+  }
+
+  private static DiskConfig createDiskConfig(String diskType, int diskSizeGb) {
+    return DiskConfig.newBuilder()
+        .setBootDiskType(diskType)
+        .setBootDiskSizeGb(diskSizeGb)
+        .setNumLocalSsds(0)
+        .build();
+  }
+
   protected void setNetworkConfigs(Compute compute, GceClusterConfig.Builder clusterConfig,
       boolean privateInstance) throws RetryableProvisionException, IOException {
     String network = conf.getNetwork();

cdap-runtime-ext-dataproc/src/main/java/io/cdap/cdap/runtime/spi/provisioner/dataproc/DataprocConf.java

@@ -117,6 +117,8 @@ final class DataprocConf {
 
   public static final String MASTER_FLEX_VM_MACHINE_TYPES = "masterFlexVmMachineTypes";
   public static final String WORKER_FLEX_VM_MACHINE_TYPES = "workerFlexVmMachineTypes";
+  public static final String MASTER_FLEX_VM_DISK_TYPES = "masterFlexVmDiskTypes";
+  public static final String WORKER_FLEX_VM_DISK_TYPES = "workerFlexVmDiskTypes";
 
   private static final Splitter COMMA_SPLITTER =
     Splitter.on(',').trimResults().omitEmptyStrings();
@@ -139,6 +141,7 @@ final class DataprocConf {
   private final String masterDiskType;
   private final String masterMachineType;
   private final List<String> masterFlexVmMachineTypes;
+  private final List<String> masterFlexVmDiskTypes;
 
   private final int workerNumNodes;
   private final int secondaryWorkerNumNodes;
@@ -148,6 +151,7 @@ final class DataprocConf {
   private final String workerDiskType;
   private final String workerMachineType;
   private final List<String> workerFlexVmMachineTypes;
+  private final List<String> workerFlexVmDiskTypes;
 
   private final long pollCreateDelay;
   private final long pollCreateJitter;
@@ -201,10 +205,10 @@ private DataprocConf(@Nullable String accountKey, String region, String zone, St
       @Nullable String networkHostProjectId, @Nullable String network, @Nullable String subnet,
       int masterNumNodes, int masterCpus, int masterMemoryMb,
       int masterDiskGb, String masterDiskType, @Nullable String masterMachineType,
-      List<String> masterFlexVmMachineTypes,
+      List<String> masterFlexVmMachineTypes, List<String> masterFlexVmDiskTypes,
       int workerNumNodes, int secondaryWorkerNumNodes, int workerCpus, int workerMemoryMb,
       int workerDiskGb, String workerDiskType, @Nullable String workerMachineType,
-      List<String> workerFlexVmMachineTypes,
+      List<String> workerFlexVmMachineTypes, List<String> workerFlexVmDiskTypes,
       long pollCreateDelay, long pollCreateJitter, long pollDeleteDelay, long pollInterval,
       @Nullable String encryptionKeyName, @Nullable String gcsBucket,
       @Nullable String tempBucket, @Nullable String serviceAccount, boolean preferExternalIp,
@@ -245,6 +249,7 @@ private DataprocConf(@Nullable String accountKey, String region, String zone, St
     this.masterDiskType = masterDiskType;
     this.masterMachineType = masterMachineType;
     this.masterFlexVmMachineTypes = masterFlexVmMachineTypes;
+    this.masterFlexVmDiskTypes = masterFlexVmDiskTypes;
     this.workerNumNodes = workerNumNodes;
     this.secondaryWorkerNumNodes = secondaryWorkerNumNodes;
     this.workerCpus = workerCpus;
@@ -253,6 +258,7 @@ private DataprocConf(@Nullable String accountKey, String region, String zone, St
     this.workerDiskType = workerDiskType;
     this.workerMachineType = workerMachineType;
     this.workerFlexVmMachineTypes = workerFlexVmMachineTypes;
+    this.workerFlexVmDiskTypes = workerFlexVmDiskTypes;
     this.pollCreateDelay = pollCreateDelay;
     this.pollCreateJitter = pollCreateJitter;
     this.pollDeleteDelay = pollDeleteDelay;
@@ -363,6 +369,14 @@ public List<String> getWorkerFlexVmMachineTypes() {
     return formatMachineType(workerFlexVmMachineTypes, workerCpus, workerMemoryMb);
   }
 
+  public List<String> getMasterFlexVmDiskTypes() {
+    return masterFlexVmDiskTypes;
+  }
+
+  public List<String> getWorkerFlexVmDiskTypes() {
+    return workerFlexVmDiskTypes;
+  }
+
   int getTotalWorkerCpus() {
     if (enablePredefinedAutoScaling) {
       return workerCpus
@@ -692,13 +706,15 @@ static DataprocConf create(Map<String, String> properties) {
       masterDiskType = "pd-standard";
     }
     final List<String> masterFlexVmMachineTypes = getStringList(properties, MASTER_FLEX_VM_MACHINE_TYPES);
+    final List<String> masterFlexVmDiskTypes = getDiskTypeList(properties, MASTER_FLEX_VM_DISK_TYPES);
     final int workerDiskGb = getInt(properties, "workerDiskGB", 1000);
     String workerDiskType = getString(properties, "workerDiskType");
     final String workerMachineType = getString(properties, "workerMachineType");
     if (workerDiskType == null) {
       workerDiskType = "pd-standard";
     }
     final List<String> workerFlexVmMachineTypes = getStringList(properties, WORKER_FLEX_VM_MACHINE_TYPES);
+    final List<String> workerFlexVmDiskTypes = getDiskTypeList(properties, WORKER_FLEX_VM_DISK_TYPES);
 
     final long pollCreateDelay = getLong(properties, "pollCreateDelay", 60);
     final long pollCreateJitter = getLong(properties, "pollCreateJitter", 20);
@@ -807,9 +823,9 @@ static DataprocConf create(Map<String, String> properties) {
     return new DataprocConf(accountKey, region, zone, projectId, networkHostProjectId, network,
         subnet,
         masterNumNodes, masterCpus, masterMemoryMb, masterDiskGb,
-        masterDiskType, masterMachineType, masterFlexVmMachineTypes,
+        masterDiskType, masterMachineType, masterFlexVmMachineTypes, masterFlexVmDiskTypes,
         workerNumNodes, secondaryWorkerNumNodes, workerCpus, workerMemoryMb, workerDiskGb,
-        workerDiskType, workerMachineType, workerFlexVmMachineTypes,
+        workerDiskType, workerMachineType, workerFlexVmMachineTypes, workerFlexVmDiskTypes,
         pollCreateDelay, pollCreateJitter, pollDeleteDelay, pollInterval,
         gcpCmekKeyName, gcpCmekBucket, tempBucket, serviceAccount, preferExternalIp,
         stackdriverLoggingEnabled, stackdriverMonitoringEnabled,
@@ -889,4 +905,14 @@ private static List<String> getStringList(Map<String, String> properties, String
       ? Collections.emptyList()
       : COMMA_SPLITTER.splitToList(val);
   }
-}
+
+  /**
+   * Parses a comma-separated list of boot disk types, normalized to lower case.
+   */
+  private static List<String> getDiskTypeList(Map<String, String> properties, String key) {
+    List<String> diskTypes = getStringList(properties, key).stream()
+      .map(String::toLowerCase)
+      .collect(Collectors.toList());
+    return Collections.unmodifiableList(diskTypes);
+  }
+}

cdap-runtime-ext-dataproc/src/main/java/io/cdap/cdap/runtime/spi/provisioner/dataproc/DataprocProvisioner.java

@@ -17,7 +17,6 @@
 package io.cdap.cdap.runtime.spi.provisioner.dataproc;
 
 import com.google.api.gax.grpc.GrpcStatusCode;
-import com.google.api.gax.rpc.ApiException;
 import com.google.api.gax.rpc.StatusCode;
 import com.google.cloud.dataproc.v1.ClusterOperationMetadata;
 import com.google.cloud.dataproc.v1.ClusterStatus.State;
@@ -45,7 +44,6 @@
 import java.util.Collections;
 import java.util.EnumSet;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -86,9 +84,6 @@ public class DataprocProvisioner extends AbstractDataprocProvisioner {
   private static final Set<ClusterStatus> TERMINAL_STATES =
     EnumSet.of(ClusterStatus.RUNNING, ClusterStatus.FAILED, ClusterStatus.NOT_EXISTS);
 
-  private static final Pattern MACHINE_TYPE_PATTERN =
-    Pattern.compile("^[a-z\\d]+(-[a-z\\d]+)*$");
-
   private final DataprocClientFactory clientFactory;
 
   @SuppressWarnings("WeakerAccess")
@@ -107,24 +102,6 @@ public void validateProperties(Map<String, String> properties) {
     DataprocConf conf = DataprocConf.create(properties);
     boolean privateInstance = Boolean.parseBoolean(
         getSystemContext().getProperties().get(PRIVATE_INSTANCE));
-    Set<String> allFlexTypes = new HashSet<>();
-    allFlexTypes.addAll(conf.getMasterFlexVmMachineTypes());
-    allFlexTypes.addAll(conf.getWorkerFlexVmMachineTypes());
-
-    for (String machineType : allFlexTypes) {
-      if (!MACHINE_TYPE_PATTERN.matcher(machineType).matches()) {
-        String errorMessage = String.format(
-          "Invalid flexible VM machine type '%s'. "
-            + "Machine types should follow standard GCP format.",
-          machineType);
-        throw new DataprocRuntimeException.Builder()
-          .withErrorCategory(DataprocRuntimeException.ERROR_CATEGORY_PROVISIONING_CONFIGURATION)
-          .withErrorReason(errorMessage)
-          .withErrorMessage(errorMessage)
-          .withErrorType(ErrorType.USER)
-          .build();
-      }
-  }
 
     if (privateInstance && conf.isPreferExternalIp()) {
       // When prefer external IP is set to true it means only Dataproc external ip can be used for

cdap-runtime-ext-dataproc/src/main/resources/gcp-dataproc.json

@@ -174,7 +174,7 @@
           "widget-type": "multi-select",
           "label": "Master Flexible Machine Types",
           "name": "masterFlexVmMachineTypes",
-          "description": "Optional comma-separated list of fallback machine types in priority order if the primary master machine type is unavailable (e.g. n1, n2, n2d, e2).",
+          "description": "Optional comma-separated list of alternate machine types Dataproc can use if the primary master machine type is unavailable (e.g. n1, n2, n2d, e2). All listed types are given equal preference.",
           "widget-attributes": {
             "options": [
               "n1",
@@ -272,7 +272,7 @@
           "widget-type": "multi-select",
           "label": "Worker Flexible Machine Types",
           "name": "workerFlexVmMachineTypes",
-          "description": "Optional comma-separated list of fallback machine types in priority order if the primary worker machine type is unavailable (e.g. n1, n2, n2d, e2).",
+          "description": "Optional comma-separated list of alternate machine types Dataproc can use if the primary worker machine type is unavailable (e.g. n1, n2, n2d, e2). All listed types are given equal preference.",
           "widget-attributes": {
             "options": [
               "n1",