Skip to content

Commit 2c42799

Browse files
javiergaitanmcovarr
authored and
SophiaPerzan-DG
committed
feat: [GCP Batch] Support passing standard machine types to the Google backend (#1)
* WX-1810 WX-1830 n1/n2/n2d machine types, cpuPlatform on GCPBATCH (broadinstitute#7518) * feat: [GCP Batch] Support passing standard machine types to the Google backend --------- Co-authored-by: Miguel Covarrubias <mcovarr@users.noreply.github.com> @ Conflicts: @ CHANGELOG.md @ supportedBackends/google/batch/src/main/scala/cromwell/backend/google/batch/api/GcpBatchRequestFactoryImpl.scala @ supportedBackends/google/batch/src/main/scala/cromwell/backend/google/batch/models/GcpBatchRuntimeAttributes.scala @ supportedBackends/google/batch/src/main/scala/cromwell/backend/google/batch/util/GcpBatchMachineConstraints.scala @ supportedBackends/google/batch/src/test/scala/cromwell/backend/google/batch/models/GcpBatchRuntimeAttributesSpec.scala @ supportedBackends/google/batch/src/test/scala/cromwell/backend/google/batch/util/GcpBatchMachineConstraintsSpec.scala
1 parent 648f65a commit 2c42799

File tree

6 files changed

+135
-58
lines changed

6 files changed

+135
-58
lines changed

‎supportedBackends/google/batch/src/main/scala/cromwell/backend/google/batch/api/GcpBatchRequestFactoryImpl.scala‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -262,9 +262,11 @@ class GcpBatchRequestFactoryImpl()(implicit gcsTransferConfiguration: GcsTransfe
262262
val machineType = GcpBatchMachineConstraints.machineType(runtimeAttributes.memory,
263263
runtimeAttributes.cpu,
264264
cpuPlatformOption = runtimeAttributes.cpuPlatform,
265+
standardMachineTypeOption = runtimeAttributes.standardMachineType,
266+
googleLegacyMachineSelection = false,
265267
jobLogger = jobLogger
266268
)
267-
val instancePolicy =
269+
val instancePolicy =
268270
createInstancePolicy(cpuPlatform = cpuPlatform, spotModel, accelerators, allDisks, machineType = machineType)
269271
val locationPolicy = LocationPolicy.newBuilder.addAllAllowedLocations(zones.asJava).build
270272
val allocationPolicy =

‎supportedBackends/google/batch/src/main/scala/cromwell/backend/google/batch/models/GcpBatchCustomMachineType.scala‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ import wom.format.MemorySize
1111

1212
import scala.math.{log, pow}
1313

14+
case class StandardMachineType(machineType: String) {}
15+
1416
/**
1517
* Adjusts memory and cpu for custom machine types.
1618
*

‎supportedBackends/google/batch/src/main/scala/cromwell/backend/google/batch/models/GcpBatchRuntimeAttributes.scala‎

Lines changed: 25 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,9 @@ final case class GcpBatchRuntimeAttributes(cpu: Int Refined Positive,
4545
failOnStderr: Boolean,
4646
continueOnReturnCode: ContinueOnReturnCode,
4747
noAddress: Boolean,
48-
checkpointFilename: Option[String]
48+
useDockerImageCache: Option[Boolean],
49+
checkpointFilename: Option[String],
50+
standardMachineType: Option[String]
4951
)
5052

5153
object GcpBatchRuntimeAttributes {
@@ -75,6 +77,13 @@ object GcpBatchRuntimeAttributes {
7577
val CpuPlatformIntelIceLakeValue = "Intel Ice Lake"
7678
val CpuPlatformAMDRomeValue = "AMD Rome"
7779

80+
val UseDockerImageCacheKey = "useDockerImageCache"
81+
private val useDockerImageCacheValidationInstance = new BooleanRuntimeAttributesValidation(
82+
UseDockerImageCacheKey
83+
).optional
84+
85+
val StandardMachineTypeKey = "standardMachineType"
86+
7887
val CheckpointFileKey = "checkpointFile"
7988
private val checkpointFileValidationInstance = new StringRuntimeAttributesValidation(CheckpointFileKey).optional
8089

@@ -88,6 +97,8 @@ object GcpBatchRuntimeAttributes {
8897
)
8998
private def cpuPlatformValidation(runtimeConfig: Option[Config]): OptionalRuntimeAttributesValidation[String] =
9099
cpuPlatformValidationInstance
100+
private def standardMachineTypeValidation(runtimeConfig: Option[Config]): OptionalRuntimeAttributesValidation[String] =
101+
new StringRuntimeAttributesValidation(StandardMachineTypeKey).optional
91102
private def gpuTypeValidation(runtimeConfig: Option[Config]): OptionalRuntimeAttributesValidation[GpuType] =
92103
GpuTypeValidation.optional
93104

@@ -150,7 +161,8 @@ object GcpBatchRuntimeAttributes {
150161
memoryValidation(runtimeConfig),
151162
bootDiskSizeValidation(runtimeConfig),
152163
checkpointFileValidationInstance,
153-
dockerValidation
164+
dockerValidation,
165+
standardMachineTypeValidation(runtimeConfig)
154166
)
155167
}
156168

@@ -201,6 +213,14 @@ object GcpBatchRuntimeAttributes {
201213
RuntimeAttributesValidation.extract(memoryValidation(runtimeAttrsConfig), validatedRuntimeAttributes)
202214
val disks: Seq[GcpBatchAttachedDisk] =
203215
RuntimeAttributesValidation.extract(disksValidation(runtimeAttrsConfig), validatedRuntimeAttributes)
216+
val useDockerImageCache: Option[Boolean] = RuntimeAttributesValidation.extractOption(
217+
useDockerImageCacheValidation(runtimeAttrsConfig).key,
218+
validatedRuntimeAttributes
219+
)
220+
val standardMachineType: Option[String] = RuntimeAttributesValidation.extractOption(
221+
standardMachineTypeValidation(runtimeAttrsConfig).key,
222+
validatedRuntimeAttributes
223+
)
204224

205225
new GcpBatchRuntimeAttributes(
206226
cpu = cpu,
@@ -215,7 +235,9 @@ object GcpBatchRuntimeAttributes {
215235
failOnStderr = failOnStderr,
216236
continueOnReturnCode = continueOnReturnCode,
217237
noAddress = noAddress,
218-
checkpointFilename = checkpointFileName
238+
useDockerImageCache = useDockerImageCache,
239+
checkpointFilename = checkpointFileName,
240+
standardMachineType = standardMachineType
219241
)
220242
}
221243

‎supportedBackends/google/batch/src/main/scala/cromwell/backend/google/batch/util/GcpBatchMachineConstraints.scala‎

Lines changed: 22 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -4,28 +4,37 @@ import cromwell.backend.google.batch.models.{
44
GcpBatchRuntimeAttributes,
55
N1CustomMachineType,
66
N2CustomMachineType,
7-
N2DCustomMachineType
7+
N2DCustomMachineType,
8+
StandardMachineType
89
}
910
import cromwell.core.logging.JobLogger
1011
import eu.timepit.refined.api.Refined
1112
import eu.timepit.refined.numeric.Positive
13+
import wdl4s.parser.MemoryUnit
1214
import wom.format.MemorySize
1315

1416
object GcpBatchMachineConstraints {
1517
def machineType(memory: MemorySize,
1618
cpu: Int Refined Positive,
1719
cpuPlatformOption: Option[String],
20+
standardMachineTypeOption: Option[String],
21+
googleLegacyMachineSelection: Boolean,
1822
jobLogger: JobLogger
19-
): String = {
20-
// If someone requests Intel Cascade Lake or Intel Ice Lake as their CPU platform then switch the machine type to n2.
21-
// Similarly, CPU platform of AMD Rome corresponds to the machine type n2d.
22-
val customMachineType =
23-
cpuPlatformOption match {
24-
case Some(GcpBatchRuntimeAttributes.CpuPlatformIntelCascadeLakeValue) => N2CustomMachineType
25-
case Some(GcpBatchRuntimeAttributes.CpuPlatformIntelIceLakeValue) => N2CustomMachineType
26-
case Some(GcpBatchRuntimeAttributes.CpuPlatformAMDRomeValue) => N2DCustomMachineType
27-
case _ => N1CustomMachineType
28-
}
29-
customMachineType.machineType(memory, cpu, jobLogger)
30-
}
23+
): String =
24+
if (standardMachineTypeOption.exists(_.trim.nonEmpty)) {
25+
StandardMachineType(standardMachineTypeOption.get).machineType
26+
} else if (googleLegacyMachineSelection) {
27+
s"predefined-$cpu-${memory.to(MemoryUnit.MB).amount.intValue()}"
28+
} else {
29+
// If someone requests Intel Cascade Lake or Intel Ice Lake as their CPU platform then switch the machine type to n2.
30+
// Similarly, CPU platform of AMD Rome corresponds to the machine type n2d.
31+
val customMachineType =
32+
cpuPlatformOption match {
33+
case Some(GcpBatchRuntimeAttributes.CpuPlatformIntelCascadeLakeValue) => N2CustomMachineType
34+
case Some(GcpBatchRuntimeAttributes.CpuPlatformIntelIceLakeValue) => N2CustomMachineType
35+
case Some(GcpBatchRuntimeAttributes.CpuPlatformAMDRomeValue) => N2DCustomMachineType
36+
case _ => N1CustomMachineType
37+
}
38+
customMachineType.machineType(memory, cpu, jobLogger)
39+
}
3140
}

‎supportedBackends/google/batch/src/test/scala/cromwell/backend/google/batch/models/GcpBatchRuntimeAttributesSpec.scala‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,9 @@ trait GcpBatchRuntimeAttributesSpecsMixin {
285285
failOnStderr = false,
286286
continueOnReturnCode = ContinueOnReturnCodeSet(Set(0)),
287287
noAddress = false,
288-
checkpointFilename = None
288+
useDockerImageCache = None,
289+
checkpointFilename = None,
290+
standardMachineType = None
289291
)
290292

291293
def assertBatchRuntimeAttributesSuccessfulCreation(runtimeAttributes: Map[String, WomValue],

‎supportedBackends/google/batch/src/test/scala/cromwell/backend/google/batch/util/GcpBatchMachineConstraintsSpec.scala‎

Lines changed: 80 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -24,65 +24,105 @@ class GcpBatchMachineConstraintsSpec extends AnyFlatSpec with CromwellTimeoutSpe
2424

2525
it should "generate valid machine types" in {
2626
val validTypes = Table(
27-
("memory", "cpu", "cpuPlatformOption", "machineTypeString"),
27+
("memory", "cpu", "cpuPlatformOption", "standardMachineTypeOption", "googleLegacyMachineSelection", "machineTypeString"),
2828
// Already ok tuple
29-
(MemorySize(1024, MemoryUnit.MB), refineMV[Positive](1), None, "custom-1-1024"),
29+
(MemorySize(1024, MemoryUnit.MB), refineMV[Positive](1), None, None, false, "custom-1-1024"),
3030
// CPU must be even (except if it's 1)
31-
(MemorySize(4, MemoryUnit.GB), refineMV[Positive](3), None, "custom-4-4096"),
31+
(MemorySize(4, MemoryUnit.GB), refineMV[Positive](3), None, None, false, "custom-4-4096"),
3232
// Memory must be a multiple of 256
33-
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](1), None, "custom-1-1024"),
33+
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](1), None, None, false, "custom-1-1024"),
3434
// Memory / cpu ratio must be > 0.9GB, increase memory
35-
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](4), None, "custom-4-3840"),
36-
(MemorySize(14, MemoryUnit.GB), refineMV[Positive](16), None, "custom-16-14848"),
35+
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](4), None, None, false, "custom-4-3840"),
36+
(MemorySize(14, MemoryUnit.GB), refineMV[Positive](16), None, None, false, "custom-16-14848"),
3737
// Memory / cpu ratio must be < 6.5GB, increase CPU
38-
(MemorySize(13.65, MemoryUnit.GB), refineMV[Positive](1), None, "custom-4-14080"),
38+
(MemorySize(13.65, MemoryUnit.GB), refineMV[Positive](1), None, None, false, "custom-4-14080"),
3939
// Memory should be an int
40-
(MemorySize(1520.96, MemoryUnit.MB), refineMV[Positive](1), None, "custom-1-1536"),
41-
(MemorySize(1024.0, MemoryUnit.MB), refineMV[Positive](1), None, "custom-1-1024"),
40+
(MemorySize(1520.96, MemoryUnit.MB), refineMV[Positive](1), None, None, false, "custom-1-1536"),
41+
(MemorySize(1024.0, MemoryUnit.MB), refineMV[Positive](1), None, None, false, "custom-1-1024"),
4242
// Increase to a cpu selection not valid for n2 below
43-
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, "custom-34-31488"),
43+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, None, false, "custom-34-31488"),
44+
45+
// Same tests as above but with legacy machine type selection (cpu and memory as specified. No 'custom machine
46+
// requirement' adjustments are expected this time, except float->int)
47+
(MemorySize(1024, MemoryUnit.MB), refineMV[Positive](1), None, None, true, "predefined-1-1024"),
48+
(MemorySize(4, MemoryUnit.GB), refineMV[Positive](3), None, None, true, "predefined-3-4096"),
49+
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](1), None, None, true, "predefined-1-1024"),
50+
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](4), None, None, true, "predefined-4-1024"),
51+
(MemorySize(14, MemoryUnit.GB), refineMV[Positive](16), None, None, true, "predefined-16-14336"),
52+
(MemorySize(13.65, MemoryUnit.GB), refineMV[Positive](1), None, None, true, "predefined-1-13977"),
53+
(MemorySize(1520.96, MemoryUnit.MB), refineMV[Positive](1), None, None, true, "predefined-1-1520"),
54+
(MemorySize(1024.0, MemoryUnit.MB), refineMV[Positive](1), None, None, true, "predefined-1-1024"),
55+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, None, true, "predefined-33-2048"),
4456

4557
// Same tests but with cascade lake (n2)
46-
(MemorySize(1024, MemoryUnit.MB), refineMV[Positive](1), n2OptionCascadeLake, "n2-custom-2-2048"),
47-
(MemorySize(4, MemoryUnit.GB), refineMV[Positive](3), n2OptionCascadeLake, "n2-custom-4-4096"),
48-
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](1), n2OptionCascadeLake, "n2-custom-2-2048"),
49-
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](4), n2OptionCascadeLake, "n2-custom-4-4096"),
50-
(MemorySize(14, MemoryUnit.GB), refineMV[Positive](16), n2OptionCascadeLake, "n2-custom-16-16384"),
51-
(MemorySize(13.65, MemoryUnit.GB), refineMV[Positive](1), n2OptionCascadeLake, "n2-custom-2-14080"),
52-
(MemorySize(1520.96, MemoryUnit.MB), refineMV[Positive](1), n2OptionCascadeLake, "n2-custom-2-2048"),
53-
(MemorySize(1024.0, MemoryUnit.MB), refineMV[Positive](1), n2OptionCascadeLake, "n2-custom-2-2048"),
54-
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), n2OptionCascadeLake, "n2-custom-36-36864"),
58+
(MemorySize(1024, MemoryUnit.MB), refineMV[Positive](1), n2OptionCascadeLake, None, false, "n2-custom-2-2048"),
59+
(MemorySize(4, MemoryUnit.GB), refineMV[Positive](3), n2OptionCascadeLake, None, false, "n2-custom-4-4096"),
60+
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](1), n2OptionCascadeLake, None, false, "n2-custom-2-2048"),
61+
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](4), n2OptionCascadeLake, None, false, "n2-custom-4-4096"),
62+
(MemorySize(14, MemoryUnit.GB), refineMV[Positive](16), n2OptionCascadeLake, None, false, "n2-custom-16-16384"),
63+
(MemorySize(13.65, MemoryUnit.GB), refineMV[Positive](1), n2OptionCascadeLake, None, false, "n2-custom-2-14080"),
64+
(MemorySize(1520.96, MemoryUnit.MB), refineMV[Positive](1), n2OptionCascadeLake, None, false, "n2-custom-2-2048"),
65+
(MemorySize(1024.0, MemoryUnit.MB), refineMV[Positive](1), n2OptionCascadeLake, None, false, "n2-custom-2-2048"),
66+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), n2OptionCascadeLake, None, false, "n2-custom-36-36864"),
5567

5668
// Same tests, but with ice lake. Should produce same results as cascade lake since they're both n2.
57-
(MemorySize(1024, MemoryUnit.MB), refineMV[Positive](1), n2OptionIceLake, "n2-custom-2-2048"),
58-
(MemorySize(4, MemoryUnit.GB), refineMV[Positive](3), n2OptionIceLake, "n2-custom-4-4096"),
59-
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](1), n2OptionIceLake, "n2-custom-2-2048"),
60-
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](4), n2OptionIceLake, "n2-custom-4-4096"),
61-
(MemorySize(14, MemoryUnit.GB), refineMV[Positive](16), n2OptionIceLake, "n2-custom-16-16384"),
62-
(MemorySize(13.65, MemoryUnit.GB), refineMV[Positive](1), n2OptionIceLake, "n2-custom-2-14080"),
63-
(MemorySize(1520.96, MemoryUnit.MB), refineMV[Positive](1), n2OptionIceLake, "n2-custom-2-2048"),
64-
(MemorySize(1024.0, MemoryUnit.MB), refineMV[Positive](1), n2OptionIceLake, "n2-custom-2-2048"),
65-
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), n2OptionIceLake, "n2-custom-36-36864"),
69+
(MemorySize(1024, MemoryUnit.MB), refineMV[Positive](1), n2OptionIceLake, None, false, "n2-custom-2-2048"),
70+
(MemorySize(4, MemoryUnit.GB), refineMV[Positive](3), n2OptionIceLake, None, false, "n2-custom-4-4096"),
71+
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](1), n2OptionIceLake, None, false, "n2-custom-2-2048"),
72+
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](4), n2OptionIceLake, None, false, "n2-custom-4-4096"),
73+
(MemorySize(14, MemoryUnit.GB), refineMV[Positive](16), n2OptionIceLake, None, false, "n2-custom-16-16384"),
74+
(MemorySize(13.65, MemoryUnit.GB), refineMV[Positive](1), n2OptionIceLake, None, false, "n2-custom-2-14080"),
75+
(MemorySize(1520.96, MemoryUnit.MB), refineMV[Positive](1), n2OptionIceLake, None, false, "n2-custom-2-2048"),
76+
(MemorySize(1024.0, MemoryUnit.MB), refineMV[Positive](1), n2OptionIceLake, None, false, "n2-custom-2-2048"),
77+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), n2OptionIceLake, None, false, "n2-custom-36-36864"),
6678

6779
// Same tests but with AMD Rome (n2d) #cpu > 16 are in increments of 16
68-
(MemorySize(1024, MemoryUnit.MB), refineMV[Positive](1), n2dOption, "n2d-custom-2-1024"),
69-
(MemorySize(4, MemoryUnit.GB), refineMV[Positive](3), n2dOption, "n2d-custom-4-4096"),
70-
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](1), n2dOption, "n2d-custom-2-1024"),
71-
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](4), n2dOption, "n2d-custom-4-2048"),
72-
(MemorySize(14, MemoryUnit.GB), refineMV[Positive](16), n2dOption, "n2d-custom-16-14336"),
73-
(MemorySize(13.65, MemoryUnit.GB), refineMV[Positive](1), n2dOption, "n2d-custom-2-14080"),
74-
(MemorySize(1520.96, MemoryUnit.MB), refineMV[Positive](1), n2dOption, "n2d-custom-2-1536"),
75-
(MemorySize(1024.0, MemoryUnit.MB), refineMV[Positive](1), n2dOption, "n2d-custom-2-1024"),
76-
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), n2dOption, "n2d-custom-48-24576"),
77-
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](81), n2dOption, "n2d-custom-96-49152"),
78-
(MemorySize(256, MemoryUnit.GB), refineMV[Positive](128), n2dOption, "n2d-custom-96-262144")
80+
(MemorySize(1024, MemoryUnit.MB), refineMV[Positive](1), n2dOption, None, false, "n2d-custom-2-1024"),
81+
(MemorySize(4, MemoryUnit.GB), refineMV[Positive](3), n2dOption, None, false, "n2d-custom-4-4096"),
82+
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](1), n2dOption, None, false, "n2d-custom-2-1024"),
83+
(MemorySize(1, MemoryUnit.GB), refineMV[Positive](4), n2dOption, None, false, "n2d-custom-4-2048"),
84+
(MemorySize(14, MemoryUnit.GB), refineMV[Positive](16), n2dOption, None, false, "n2d-custom-16-14336"),
85+
(MemorySize(13.65, MemoryUnit.GB), refineMV[Positive](1), n2dOption, None, false, "n2d-custom-2-14080"),
86+
(MemorySize(1520.96, MemoryUnit.MB), refineMV[Positive](1), n2dOption, None, false, "n2d-custom-2-1536"),
87+
(MemorySize(1024.0, MemoryUnit.MB), refineMV[Positive](1), n2dOption, None, false, "n2d-custom-2-1024"),
88+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), n2dOption, None, false, "n2d-custom-48-24576"),
89+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](81), n2dOption, None, false, "n2d-custom-96-49152"),
90+
(MemorySize(256, MemoryUnit.GB), refineMV[Positive](128), n2dOption, None, false, "n2d-custom-96-262144"),
91+
92+
// Test Standard Machine types
93+
// General-purpose machine family
94+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, Option("n1-standard-2"), false, "n1-standard-2"),
95+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, Option("n1-highmem-2"), false, "n1-highmem-2"),
96+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, Option("n1-highcpu-4"), false, "n1-highcpu-4"),
97+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, Option("f1-micro"), false, "f1-micro"),
98+
99+
// Accelerator-optimized machine family
100+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, Option("a2-highgpu-1g"), false, "a2-highgpu-1g"),
101+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, Option("a3-megagpu-8g"), false, "a3-megagpu-8g"),
102+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, Option("g2-standard-4"), false, "g2-standard-4"),
103+
104+
// Other machine families
105+
// Storage-optimized
106+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, Option("z3-highmem-88"), false, "z3-highmem-88"),
107+
// Compute-optimized
108+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, Option("h3-standard-88"), false, "h3-standard-88"),
109+
// Memory-optimized
110+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, Option("m3-ultramem-128"), false, "m3-ultramem-128"),
111+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, Option("a2-highgpu-1g"), false, "a2-highgpu-1g"),
112+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, Option("a2-highgpu-1g"), false, "a2-highgpu-1g"),
113+
114+
// Standard machine type overrides legacy selection
115+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, Option("a2-highgpu-1g"), true, "a2-highgpu-1g"),
116+
(MemorySize(2, MemoryUnit.GB), refineMV[Positive](33), None, Option("a2-highgpu-1g"), false, "a2-highgpu-1g")
79117
)
80118

81-
forAll(validTypes) { (memory, cpu, cpuPlatformOption, expected) =>
119+
forAll(validTypes) { (memory, cpu, cpuPlatformOption, standardMachineTypeOption, googleLegacyMachineSelection, expected) =>
82120
GcpBatchMachineConstraints.machineType(
83121
memory = memory,
84122
cpu = cpu,
85123
cpuPlatformOption = cpuPlatformOption,
124+
standardMachineTypeOption = standardMachineTypeOption,
125+
googleLegacyMachineSelection = googleLegacyMachineSelection,
86126
jobLogger = mock[JobLogger]
87127
) shouldBe expected
88128
}

0 commit comments

Comments
 (0)