Examples with MachineType com.hartwig.pipeline.execution.MachineType used on opensource projects

Example 1 with MachineType

use of com.hartwig.pipeline.execution.MachineType 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);
}

Example 2 with MachineType

use of com.hartwig.pipeline.execution.MachineType 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);
    }
}