use of com.google.api.services.serviceusage.v1beta1.model.QuotaBucket in project pipeline5 by hartwigmedical.
the class QuotaConstrainedComputeEngineTest method ensuresVmsDontExceedMaxCPUForRegion.
@Test
public void ensuresVmsDontExceedMaxCPUForRegion() throws Exception {
ComputeEngine decorated = mock(ComputeEngine.class);
ServiceUsage serviceUsage = mock(ServiceUsage.class);
ServiceUsage.Services services = mock(ServiceUsage.Services.class);
ServiceUsage.Services.ConsumerQuotaMetrics consumerQuotaMetrics = mock(ServiceUsage.Services.ConsumerQuotaMetrics.class);
ServiceUsage.Services.ConsumerQuotaMetrics.Limits limits = mock(ServiceUsage.Services.ConsumerQuotaMetrics.Limits.class);
ServiceUsage.Services.ConsumerQuotaMetrics.Limits.Get limitsGet = mock(ServiceUsage.Services.ConsumerQuotaMetrics.Limits.Get.class);
when(serviceUsage.services()).thenReturn(services);
when(services.consumerQuotaMetrics()).thenReturn(consumerQuotaMetrics);
when(consumerQuotaMetrics.limits()).thenReturn(limits);
ArgumentCaptor<String> quotaName = ArgumentCaptor.forClass(String.class);
when(limits.get(quotaName.capture())).thenReturn(limitsGet);
ConsumerQuotaLimit limit = new ConsumerQuotaLimit().setQuotaBuckets(List.of(new QuotaBucket().setEffectiveLimit(10L).setDimensions(Map.of(REGION, REGION))));
when(limitsGet.execute()).thenReturn(limit);
VirtualMachineJobDefinition jobDefinition = VirtualMachineJobDefinition.builder().name("test").namespacedResults(ResultsDirectory.defaultDirectory()).startupCommand(BashStartupScript.of("empty")).performanceProfile(VirtualMachinePerformanceProfile.custom(10, 10)).build();
ArgumentCaptor<VirtualMachineJobDefinition> constrained = ArgumentCaptor.forClass(VirtualMachineJobDefinition.class);
when(decorated.submit(any(), constrained.capture(), any())).thenReturn(PipelineStatus.SUCCESS);
QuotaConstrainedComputeEngine victim = new QuotaConstrainedComputeEngine(decorated, serviceUsage, REGION, PROJECT, 0.6);
PipelineStatus result = victim.submit(MockRuntimeBucket.test().getRuntimeBucket(), jobDefinition);
assertThat(result).isEqualTo(PipelineStatus.SUCCESS);
MachineType machineType = constrained.getValue().performanceProfile().machineType();
assertThat(machineType.cpus()).isEqualTo(6);
assertThat(machineType.memoryGB()).isEqualTo(6);
}
use of com.google.api.services.serviceusage.v1beta1.model.QuotaBucket in project pipeline5 by hartwigmedical.
the class QuotaConstrainedComputeEngine method submit.
@Override
public PipelineStatus submit(final RuntimeBucket bucket, final VirtualMachineJobDefinition jobDefinition, final String discriminator) {
VirtualMachineJobDefinition constrained = jobDefinition;
MachineType machineType = jobDefinition.performanceProfile().machineType();
try {
QuotaBucket regionalQuota = serviceUsage.services().consumerQuotaMetrics().limits().get(cpuQuotaName(project)).execute().getQuotaBuckets().stream().filter(b -> b.getDimensions() != null).filter(b -> region(b) != null).filter(b -> region(b).equals(region)).findFirst().orElseThrow();
int maxCPU = (int) (regionalQuota.getEffectiveLimit().intValue() * constrainByPercentage);
if (machineType.cpus() > maxCPU) {
double reductionRatio = (double) maxCPU / machineType.cpus();
constrained = VirtualMachineJobDefinition.builder().from(constrained).performanceProfile(VirtualMachinePerformanceProfile.custom(maxCPU, (int) (machineType.memoryGB() * reductionRatio))).build();
}
return decorated.submit(bucket, constrained, discriminator);
} catch (IOException e) {
throw new RuntimeException(e);
}
}
Aggregations