Skip to content

Commit

Permalink
fix: Create a new standardMachineType runtime attribute instead of re…
Browse files Browse the repository at this point in the history
…using cpuPlatform
  • Loading branch information
javiergaitan committed Sep 17, 2024
1 parent 8b20d18 commit 366115d
Showing 1 changed file with 14 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,8 @@ final case class GcpBatchRuntimeAttributes(cpu: Int Refined Positive,
continueOnReturnCode: ContinueOnReturnCode,
noAddress: Boolean,
useDockerImageCache: Option[Boolean],
checkpointFilename: Option[String]
checkpointFilename: Option[String],
standardMachineType: Option[String]
)

object GcpBatchRuntimeAttributes {
Expand Down Expand Up @@ -85,6 +86,8 @@ object GcpBatchRuntimeAttributes {
UseDockerImageCacheKey
).optional

val StandardMachineTypeKey = "StandardMachineType"

val CheckpointFileKey = "checkpointFile"
private val checkpointFileValidationInstance = new StringRuntimeAttributesValidation(CheckpointFileKey).optional

Expand All @@ -98,6 +101,8 @@ object GcpBatchRuntimeAttributes {
)
private def cpuPlatformValidation(runtimeConfig: Option[Config]): OptionalRuntimeAttributesValidation[String] =
cpuPlatformValidationInstance
private def standardMachineTypeValidation(runtimeConfig: Option[Config]): OptionalRuntimeAttributesValidation[String] =
new StringRuntimeAttributesValidation(StandardMachineTypeKey).optional
private def gpuTypeValidation(runtimeConfig: Option[Config]): OptionalRuntimeAttributesValidation[GpuType] =
GpuTypeValidation.optional

Expand Down Expand Up @@ -171,7 +176,8 @@ object GcpBatchRuntimeAttributes {
bootDiskSizeValidation(runtimeConfig),
useDockerImageCacheValidation(runtimeConfig),
checkpointFileValidationInstance,
dockerValidation
dockerValidation,
standardMachineTypeValidation(runtimeConfig)
)
}

Expand Down Expand Up @@ -228,6 +234,10 @@ object GcpBatchRuntimeAttributes {
useDockerImageCacheValidation(runtimeAttrsConfig).key,
validatedRuntimeAttributes
)
val standardMachineType: Option[String] = RuntimeAttributesValidation.extractOption(
standardMachineTypeValidation(runtimeAttrsConfig).key,
validatedRuntimeAttributes
)

new GcpBatchRuntimeAttributes(
cpu = cpu,
Expand All @@ -243,7 +253,8 @@ object GcpBatchRuntimeAttributes {
continueOnReturnCode = continueOnReturnCode,
noAddress = noAddress,
useDockerImageCache = useDockerImageCache,
checkpointFilename = checkpointFileName
checkpointFilename = checkpointFileName,
standardMachineType = standardMachineType
)
}

Expand Down

0 comments on commit 366115d

Please sign in to comment.