Search in sources :

Example 1 with QuotaBucket

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);
}
Also used : PipelineStatus(com.hartwig.pipeline.execution.PipelineStatus) MachineType(com.hartwig.pipeline.execution.MachineType) ServiceUsage(com.google.api.services.serviceusage.v1beta1.ServiceUsage) ConsumerQuotaLimit(com.google.api.services.serviceusage.v1beta1.model.ConsumerQuotaLimit) QuotaBucket(com.google.api.services.serviceusage.v1beta1.model.QuotaBucket) Test(org.junit.Test)

Example 2 with QuotaBucket

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);
    }
}
Also used : ServiceUsage(com.google.api.services.serviceusage.v1beta1.ServiceUsage) MachineType(com.hartwig.pipeline.execution.MachineType) QuotaBucket(com.google.api.services.serviceusage.v1beta1.model.QuotaBucket) RuntimeBucket(com.hartwig.pipeline.storage.RuntimeBucket) IOException(java.io.IOException) PipelineStatus(com.hartwig.pipeline.execution.PipelineStatus) MachineType(com.hartwig.pipeline.execution.MachineType) IOException(java.io.IOException) QuotaBucket(com.google.api.services.serviceusage.v1beta1.model.QuotaBucket)

Aggregations

ServiceUsage (com.google.api.services.serviceusage.v1beta1.ServiceUsage)2 QuotaBucket (com.google.api.services.serviceusage.v1beta1.model.QuotaBucket)2 MachineType (com.hartwig.pipeline.execution.MachineType)2 PipelineStatus (com.hartwig.pipeline.execution.PipelineStatus)2 ConsumerQuotaLimit (com.google.api.services.serviceusage.v1beta1.model.ConsumerQuotaLimit)1 RuntimeBucket (com.hartwig.pipeline.storage.RuntimeBucket)1 IOException (java.io.IOException)1 Test (org.junit.Test)1