feat: Support passing standard machine types to Google Batch by reusi…

…ng the cpuPlatform field in the RuntimeAttributes
broadinstitute · Sep 17, 2024 · 8b20d18 · 8b20d18
1 parent 081318d
commit 8b20d18
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 3 deletions.
diff --git a/...e/batch/src/main/scala/cromwell/backend/google/batch/api/GcpBatchRequestFactoryImpl.scala b/...e/batch/src/main/scala/cromwell/backend/google/batch/api/GcpBatchRequestFactoryImpl.scala
@@ -234,8 +234,10 @@ class GcpBatchRequestFactoryImpl()(implicit gcsTransferConfiguration: GcsTransfe
                                                              googleLegacyMachineSelection = false,
                                                              jobLogger = jobLogger
     )
+    // Don't set cpuPlatform if the field was used to specify standard machine type when submitting the GCP Batch request
+    val cpuPlatformBatchRequest = if (GcpBatchMachineConstraints.isStandardMachineType(cpuPlatform)) "" else cpuPlatform
     val instancePolicy =
-      createInstancePolicy(cpuPlatform = cpuPlatform, spotModel, accelerators, allDisks, machineType = machineType)
+      createInstancePolicy(cpuPlatform = cpuPlatformBatchRequest, spotModel, accelerators, allDisks, machineType = machineType)
     val locationPolicy = LocationPolicy.newBuilder.addAllowedLocations(zones).build
     val allocationPolicy =
       createAllocationPolicy(data, locationPolicy, instancePolicy.build, networkPolicy, gcpSa, accelerators)

diff --git a/...batch/src/main/scala/cromwell/backend/google/batch/models/GcpBatchCustomMachineType.scala b/...batch/src/main/scala/cromwell/backend/google/batch/models/GcpBatchCustomMachineType.scala
@@ -11,6 +11,8 @@ import wom.format.MemorySize
 
 import scala.math.{log, pow}
 
+case class StandardMachineType(machineType: String) {}
+
 /**
   * Adjusts memory and cpu for custom machine types.
   *

diff --git a/.../batch/src/main/scala/cromwell/backend/google/batch/util/GcpBatchMachineConstraints.scala b/.../batch/src/main/scala/cromwell/backend/google/batch/util/GcpBatchMachineConstraints.scala
@@ -4,22 +4,28 @@ import cromwell.backend.google.batch.models.{
   GcpBatchRuntimeAttributes,
   N1CustomMachineType,
   N2CustomMachineType,
-  N2DCustomMachineType
+  N2DCustomMachineType,
+  StandardMachineType
 }
 import cromwell.core.logging.JobLogger
 import eu.timepit.refined.api.Refined
 import eu.timepit.refined.numeric.Positive
+import scala.util.matching.Regex
 import wdl4s.parser.MemoryUnit
 import wom.format.MemorySize
 
 object GcpBatchMachineConstraints {
+  private val machineTypePattern: Regex = """^\w{2}\w?-\w+-\w+$""".r
+
   def machineType(memory: MemorySize,
                   cpu: Int Refined Positive,
                   cpuPlatformOption: Option[String],
                   googleLegacyMachineSelection: Boolean,
                   jobLogger: JobLogger
   ): String =
-    if (googleLegacyMachineSelection) {
+    if (isStandardMachineType(cpuPlatformOption.getOrElse(""))) {
+      StandardMachineType(cpuPlatformOption.getOrElse("")).machineType
+    } else if (googleLegacyMachineSelection) {
       s"predefined-$cpu-${memory.to(MemoryUnit.MB).amount.intValue()}"
     } else {
       // If someone requests Intel Cascade Lake or Intel Ice Lake as their CPU platform then switch the machine type to n2.
@@ -33,4 +39,6 @@ object GcpBatchMachineConstraints {
         }
       customMachineType.machineType(memory, cpu, jobLogger)
     }
+
+  def isStandardMachineType(machineType: String): Boolean = machineTypePattern.matches(machineType)
 }