Skip to content

Commit

Permalink
Include GPU info in instantiated VM info
Browse files Browse the repository at this point in the history
  • Loading branch information
@sam-schu
sam-schu committed Dec 13, 2024
1 parent f0924b8 commit 7233923
Show file tree
Hide file tree
Showing 4 changed files with 52 additions and 48 deletions.
2 changes: 1 addition & 1 deletion services/src/main/scala/cromwell/services/cost/GcpCostCatalogTypes.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ case class GpuInfo(count: Long, gpuType: String)
/*
* Case class that contains information retrieved from Google about a VM that cromwell has started
*/
case class InstantiatedVmInfo(region: String, machineType: String, preemptible: Boolean)
case class InstantiatedVmInfo(region: String, machineType: String, gpuInfo: Option[GpuInfo], preemptible: Boolean)

/*
* These types reflect hardcoded strings found in a google cost catalog.
Expand Down
87 changes: 45 additions & 42 deletions services/src/test/scala/cromwell/services/cost/GcpCostCatalogServiceSpec.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -146,64 +146,67 @@ class GcpCostCatalogServiceSpec
it should "find the skus for a VM when appropriate" in {
val lookupRows = Table(
("instantiatedVmInfo", "resource", "skuDescription"),
(InstantiatedVmInfo("europe-west9", "custom-16-32768", false),
(InstantiatedVmInfo("europe-west9", "custom-16-32768", None, false),
Cpu,
"N1 Predefined Instance Core running in Paris"
),
(InstantiatedVmInfo("europe-west9", "custom-16-32768", false),
(InstantiatedVmInfo("europe-west9", "custom-16-32768", None, false),
Ram,
"N1 Predefined Instance Ram running in Paris"
),
(InstantiatedVmInfo("us-central1", "custom-4-4096", true),
(InstantiatedVmInfo("us-central1", "custom-4-4096", None, true),
Cpu,
"Spot Preemptible N1 Predefined Instance Core running in Americas"
),
(InstantiatedVmInfo("us-central1", "custom-4-4096", true),
(InstantiatedVmInfo("us-central1", "custom-4-4096", None, true),
Ram,
"Spot Preemptible N1 Predefined Instance Ram running in Americas"
),
(InstantiatedVmInfo("europe-west9", "n1-custom-16-32768", false),
(InstantiatedVmInfo("europe-west9", "n1-custom-16-32768", None, false),
Cpu,
"N1 Predefined Instance Core running in Paris"
),
(InstantiatedVmInfo("europe-west9", "n1-custom-16-32768", false),
(InstantiatedVmInfo("europe-west9", "n1-custom-16-32768", None, false),
Ram,
"N1 Predefined Instance Ram running in Paris"
),
(InstantiatedVmInfo("us-central1", "n1-custom-4-4096", true),
(InstantiatedVmInfo("us-central1", "n1-custom-4-4096", None, true),
Cpu,
"Spot Preemptible N1 Predefined Instance Core running in Americas"
),
(InstantiatedVmInfo("us-central1", "n1-custom-4-4096", true),
(InstantiatedVmInfo("us-central1", "n1-custom-4-4096", None, true),
Ram,
"Spot Preemptible N1 Predefined Instance Ram running in Americas"
),
(InstantiatedVmInfo("us-central1", "n2-custom-4-4096", true),
(InstantiatedVmInfo("us-central1", "n2-custom-4-4096", None, true),
Cpu,
"Spot Preemptible N2 Custom Instance Core running in Americas"
),
(InstantiatedVmInfo("us-central1", "n2-custom-4-4096", true),
(InstantiatedVmInfo("us-central1", "n2-custom-4-4096", None, true),
Ram,
"Spot Preemptible N2 Custom Instance Ram running in Americas"
),
(InstantiatedVmInfo("us-central1", "n2-custom-4-4096", false),
(InstantiatedVmInfo("us-central1", "n2-custom-4-4096", None, false),
Cpu,
"N2 Custom Instance Core running in Americas"
),
(InstantiatedVmInfo("us-central1", "n2-custom-4-4096", false), Ram, "N2 Custom Instance Ram running in Americas"),
(InstantiatedVmInfo("us-central1", "n2d-custom-4-4096", true),
(InstantiatedVmInfo("us-central1", "n2-custom-4-4096", None, false),
Ram,
"N2 Custom Instance Ram running in Americas"
),
(InstantiatedVmInfo("us-central1", "n2d-custom-4-4096", None, true),
Cpu,
"Spot Preemptible N2D AMD Custom Instance Core running in Americas"
),
(InstantiatedVmInfo("us-central1", "n2d-custom-4-4096", true),
(InstantiatedVmInfo("us-central1", "n2d-custom-4-4096", None, true),
Ram,
"Spot Preemptible N2D AMD Custom Instance Ram running in Americas"
),
(InstantiatedVmInfo("us-central1", "n2d-custom-4-4096", false),
(InstantiatedVmInfo("us-central1", "n2d-custom-4-4096", None, false),
Cpu,
"N2D AMD Custom Instance Core running in Americas"
),
(InstantiatedVmInfo("us-central1", "n2d-custom-4-4096", false),
(InstantiatedVmInfo("us-central1", "n2d-custom-4-4096", None, false),
Ram,
"N2D AMD Custom Instance Ram running in Americas"
)
Expand All @@ -219,21 +222,21 @@ class GcpCostCatalogServiceSpec
it should "fail to find the skus for a VM when appropriate" in {
val lookupRows = Table(
("instantiatedVmInfo", "resource", "errors"),
(InstantiatedVmInfo("us-central1", "custooooooom-4-4096", true),
(InstantiatedVmInfo("us-central1", "custooooooom-4-4096", None, true),
Cpu,
List("Unrecognized machine type: custooooooom-4-4096")
),
(InstantiatedVmInfo("us-central1", "n2custom-4-4096", true),
(InstantiatedVmInfo("us-central1", "n2custom-4-4096", None, true),
Cpu,
List("Unrecognized machine type: n2custom-4-4096")
),
(InstantiatedVmInfo("us-central1", "standard-4-4096", true),
(InstantiatedVmInfo("us-central1", "standard-4-4096", None, true),
Cpu,
List("Unrecognized machine type: standard-4-4096")
),
(InstantiatedVmInfo("planet-mars1", "custom-4-4096", true),
(InstantiatedVmInfo("planet-mars1", "custom-4-4096", None, true),
Cpu,
List("Failed to look up Cpu SKU for InstantiatedVmInfo(planet-mars1,custom-4-4096,true)")
List("Failed to look up Cpu SKU for InstantiatedVmInfo(planet-mars1,custom-4-4096,None,true)")
)
)

Expand All @@ -249,18 +252,18 @@ class GcpCostCatalogServiceSpec
// Create BigDecimals from strings to avoid inequality due to floating point shenanigans
val lookupRows = Table(
("instantiatedVmInfo", "costPerHour"),
(InstantiatedVmInfo("us-central1", "custom-4-4096", true), BigDecimal(".0361")),
(InstantiatedVmInfo("us-central1", "n2-custom-4-4096", true), BigDecimal(".04254400000000000480")),
(InstantiatedVmInfo("us-central1", "n2d-custom-4-4096", true), BigDecimal(".02371600000000000040")),
(InstantiatedVmInfo("us-central1", "custom-4-4096", false), BigDecimal(".143392")),
(InstantiatedVmInfo("us-central1", "n2-custom-4-4096", false), BigDecimal(".150561600")),
(InstantiatedVmInfo("us-central1", "n2d-custom-4-4096", false), BigDecimal(".130989600000000012")),
(InstantiatedVmInfo("europe-west9", "custom-4-4096", true), BigDecimal(".035018080000000004")),
(InstantiatedVmInfo("europe-west9", "n2-custom-4-4096", true), BigDecimal("0.049532000000000004")),
(InstantiatedVmInfo("europe-west9", "n2d-custom-4-4096", true), BigDecimal("0.030608000000000004")),
(InstantiatedVmInfo("europe-west9", "custom-4-4096", false), BigDecimal(".1663347200000000040")),
(InstantiatedVmInfo("europe-west9", "n2-custom-4-4352", false), BigDecimal(".17594163050")),
(InstantiatedVmInfo("europe-west9", "n2d-custom-4-4096", false), BigDecimal(".151947952"))
(InstantiatedVmInfo("us-central1", "custom-4-4096", None, true), BigDecimal(".0361")),
(InstantiatedVmInfo("us-central1", "n2-custom-4-4096", None, true), BigDecimal(".04254400000000000480")),
(InstantiatedVmInfo("us-central1", "n2d-custom-4-4096", None, true), BigDecimal(".02371600000000000040")),
(InstantiatedVmInfo("us-central1", "custom-4-4096", None, false), BigDecimal(".143392")),
(InstantiatedVmInfo("us-central1", "n2-custom-4-4096", None, false), BigDecimal(".150561600")),
(InstantiatedVmInfo("us-central1", "n2d-custom-4-4096", None, false), BigDecimal(".130989600000000012")),
(InstantiatedVmInfo("europe-west9", "custom-4-4096", None, true), BigDecimal(".035018080000000004")),
(InstantiatedVmInfo("europe-west9", "n2-custom-4-4096", None, true), BigDecimal("0.049532000000000004")),
(InstantiatedVmInfo("europe-west9", "n2d-custom-4-4096", None, true), BigDecimal("0.030608000000000004")),
(InstantiatedVmInfo("europe-west9", "custom-4-4096", None, false), BigDecimal(".1663347200000000040")),
(InstantiatedVmInfo("europe-west9", "n2-custom-4-4352", None, false), BigDecimal(".17594163050")),
(InstantiatedVmInfo("europe-west9", "n2d-custom-4-4096", None, false), BigDecimal(".151947952"))
)

forAll(lookupRows) { case (instantiatedVmInfo: InstantiatedVmInfo, expectedCostPerHour: BigDecimal) =>
Expand All @@ -274,24 +277,24 @@ class GcpCostCatalogServiceSpec

val lookupRows = Table(
("instantiatedVmInfo", "errors"),
(InstantiatedVmInfo("us-central1", "custooooooom-4-4096", true),
(InstantiatedVmInfo("us-central1", "custooooooom-4-4096", None, true),
List("Unrecognized machine type: custooooooom-4-4096")
),
(InstantiatedVmInfo("us-central1", "n2_custom_4_4096", true),
(InstantiatedVmInfo("us-central1", "n2_custom_4_4096", None, true),
List("Unrecognized machine type: n2_custom_4_4096")
),
(InstantiatedVmInfo("us-central1", "custom-foo-4096", true),
(InstantiatedVmInfo("us-central1", "custom-foo-4096", None, true),
List("Could not extract core count from custom-foo-4096")
),
(InstantiatedVmInfo("us-central1", "custom-16-bar", true),
(InstantiatedVmInfo("us-central1", "custom-16-bar", None, true),
List("Could not extract Ram MB count from custom-16-bar")
),
(InstantiatedVmInfo("us-central1", "123-456-789", true), List("Unrecognized machine type: 123-456-789")),
(InstantiatedVmInfo("us-central1", "n2-16-4096", true),
List("Failed to look up Cpu SKU for InstantiatedVmInfo(us-central1,n2-16-4096,true)")
(InstantiatedVmInfo("us-central1", "123-456-789", None, true), List("Unrecognized machine type: 123-456-789")),
(InstantiatedVmInfo("us-central1", "n2-16-4096", None, true),
List("Failed to look up Cpu SKU for InstantiatedVmInfo(us-central1,n2-16-4096,None,true)")
),
(InstantiatedVmInfo("planet-mars1", "n2-custom-4-4096", true),
List("Failed to look up Cpu SKU for InstantiatedVmInfo(planet-mars1,n2-custom-4-4096,true)")
(InstantiatedVmInfo("planet-mars1", "n2-custom-4-4096", None, true),
List("Failed to look up Cpu SKU for InstantiatedVmInfo(planet-mars1,n2-custom-4-4096,None,true)")
)
)

Expand Down
5 changes: 3 additions & 2 deletions ...batch/src/main/scala/cromwell/backend/google/batch/api/request/BatchRequestExecutor.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -145,7 +145,7 @@ object BatchRequestExecutor {
// Get instances that can be created with this AllocationPolicy, only instances[0] is supported
val instancePolicy = allocationPolicy.getInstances(0).getPolicy
val machineType = instancePolicy.getMachineType
val preemtible = instancePolicy.getProvisioningModelValue == ProvisioningModel.PREEMPTIBLE.getNumber
val preemptible = instancePolicy.getProvisioningModelValue == ProvisioningModel.PREEMPTIBLE.getNumber

// location list = [regions/us-central1, zones/us-central1-b], region is the first element
val location = allocationPolicy.getLocation.getAllowedLocationsList.get(0)
Expand All @@ -155,7 +155,8 @@ object BatchRequestExecutor {
else
location.split("/").last

val instantiatedVmInfo = Some(InstantiatedVmInfo(region, machineType, preemtible))
// TODO: include GPU info
val instantiatedVmInfo = Some(InstantiatedVmInfo(region, machineType, None, preemptible))

if (job.getStatus.getState == JobStatus.State.SUCCEEDED) {
RunStatus.Success(events, instantiatedVmInfo)
Expand Down
6 changes: 3 additions & 3 deletions ...c/main/scala/cromwell/backend/google/pipelines/v2beta/api/request/GetRequestHandler.scala
View file
Original file line number Diff line number Diff line change
Expand Up @@ -129,12 +129,12 @@ trait GetRequestHandler { this: RequestHandler =>
gpusList.asScala.headOption
}
} yield GpuInfo(gpus.getCount, gpus.getType)
// temp
logger.error("GPUs option: " + gpuInfo.toString)

// Unlike with region and machineType, gpuInfo's being None does not indicate an invalid
// result - it just means no GPUs are being used by the VM
val instantiatedVmInfo: Option[InstantiatedVmInfo] = (region, machineType) match {
case (Some(instantiatedRegion), Some(instantiatedMachineType)) =>
Option(InstantiatedVmInfo(instantiatedRegion, instantiatedMachineType, preemptible))
Option(InstantiatedVmInfo(instantiatedRegion, instantiatedMachineType, gpuInfo, preemptible))
case _ => Option.empty
}
if (operation.getDone) {
Expand Down

0 comments on commit 7233923

Please sign in to comment.