Fix Google Batch default instance family types (#3960) [ci fast]

Signed-off-by: Paolo Di Tommaso <[email protected]>
nextflow-io · May 19, 2023 · b5257cd · b5257cd
1 parent f48a473
commit b5257cd
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 11 deletions.
diff --git a/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchMachineTypeSelector.groovy b/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchMachineTypeSelector.groovy
@@ -69,7 +69,9 @@ class GoogleBatchMachineTypeSelector {
      * https://cloud.google.com/compute/docs/disks#local_ssd_machine_type_restrictions
      * LAST UPDATE 2023-02-17
      */
-    private static final List<String> DEFAULT_FAMILIES_WITH_SSD = ['n1-*', 'n2-*', 'n2d-*', 'c2-*', 'c2d-*', 'm3-*']
+    private static final List<String> DEFAULT_FAMILIES_FOR_FUSION = ['n1-*', 'n2-*', 'n2d-*', 'c2-*', 'c2d-*', 'm3-*']
+
+    private static final List<String> DEFAULT_FAMILIES = ['n1-*', 'n2-*', 'n2d-*', 'c2-*', 'c2d-*', 'm1-*', 'm2-*', 'm3-*', 'e2-*']
 
     @Immutable
     static class MachineType {
@@ -81,10 +83,10 @@ class GoogleBatchMachineTypeSelector {
         int memPerVm
     }
 
-    String bestMachineType(int cpus, int memoryMB, String region, boolean spot, boolean localSSD, List<String> families) {
+    String bestMachineType(int cpus, int memoryMB, String region, boolean spot, boolean fusionEnabled, List<String> families) {
         final machineTypes = getAvailableMachineTypes(region)
         if (families == null)
-            families = []
+            families = Collections.<String>emptyList()
 
         // Check if a specific machine type was defined
         if (families.size() == 1) {
@@ -98,12 +100,14 @@ class GoogleBatchMachineTypeSelector {
 
         final memoryGB = Math.ceil(memoryMB / 1024.0 as float) as int
 
-        // Use only families that can have a local SSD
-        if (!families && localSSD)
-            families = DEFAULT_FAMILIES_WITH_SSD
+        if (!families ) {
+            families = fusionEnabled
+                    ? DEFAULT_FAMILIES_FOR_FUSION
+                    : DEFAULT_FAMILIES
+        }
 
         // All types are valid if no families are defined, otherwise at least it has to start with one of the given values
-        final matchMachineType = {String t -> !families || families.find { matchType(it, t) }}
+        final matchMachineType = {String type -> !families || families.find { matchType(it, type) }}
 
         // find machines with enough resources and SSD local disk
         final validMachineTypes = machineTypes.findAll {

diff --git a/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy b/plugins/nf-google/src/main/nextflow/cloud/google/batch/GoogleBatchTaskHandler.groovy
@@ -406,13 +406,12 @@ class GoogleBatchTaskHandler extends TaskHandler implements FusionAwareTask {
         final cpus = config.getCpus()
         final memory = config.getMemory() ? config.getMemory().toMega().toInteger() : 1024
         final spot = executor.config.spot ?: executor.config.preemptible
-        final useSSD = fusionEnabled()
         final families = config.getMachineType() ? config.getMachineType().tokenize(',') : []
         final priceModel = spot ? PriceModel.spot : PriceModel.standard
 
         try {
             return new CloudMachineInfo(
-                    type: GoogleBatchMachineTypeSelector.INSTANCE.bestMachineType(cpus, memory, location, spot, useSSD, families),
+                    type: GoogleBatchMachineTypeSelector.INSTANCE.bestMachineType(cpus, memory, location, spot, fusionEnabled(), families),
                     zone: location,
                     priceModel: priceModel
             )

diff --git a/.../nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchMachineTypeSelectorTest.groovy b/.../nf-google/src/test/nextflow/cloud/google/batch/GoogleBatchMachineTypeSelectorTest.groovy
@@ -24,10 +24,10 @@ class GoogleBatchMachineTypeSelectorTest extends Specification {
             getAvailableMachineTypes(REGION) >> MACHINE_TYPES
         }
         expect:
-        selector.bestMachineType(CPUS, MEM, REGION, SPOT, SSD, FAMILIES) == EXPECTED
+        selector.bestMachineType(CPUS, MEM, REGION, SPOT, FUSION, FAMILIES) == EXPECTED
 
         where:
-        CPUS | MEM  | REGION | SPOT  | SSD   | FAMILIES                   | EXPECTED
+        CPUS | MEM  | REGION | SPOT  | FUSION | FAMILIES                  | EXPECTED
         1    | 1000 | 'reg'  | true  | false | null                       | 'e2-type01'
         1    | 1000 | 'reg'  | false | true  | null                       | 'n1-type02'
         4    | 4000 | 'reg'  | false | false | []                         | 'e2-type03'