Examples with AcceleratorConfig com.google.api.services.notebooks.v1.model.AcceleratorConfig used on opensource projects

Example 1 with AcceleratorConfig

use of com.google.api.services.notebooks.v1.model.AcceleratorConfig in project terra-workspace-manager by DataBiosphere.

the class CreateAiNotebookInstanceStepTest method setFields.

@Test
public void setFields() {
    var creationParameters = new ApiGcpAiNotebookInstanceCreationParameters().postStartupScript("script.sh").machineType("machine-type").installGpuDriver(true).customGpuDriverPath("custom-path").bootDiskType("boot-disk-type").bootDiskSizeGb(111L).dataDiskType("data-disk-type").dataDiskSizeGb(222L).metadata(Map.of("metadata-key", "metadata-value")).acceleratorConfig(new ApiGcpAiNotebookInstanceAcceleratorConfig().coreCount(4L).type("accelerator-type")).vmImage(new ApiGcpAiNotebookInstanceVmImage().projectId("project-id").imageFamily("image-family").imageName("image-name")).containerImage(new ApiGcpAiNotebookInstanceContainerImage().repository("repository").tag("tag"));
    Instance instance = CreateAiNotebookInstanceStep.setFields(creationParameters, "foo@bar.com", new Instance());
    assertEquals("script.sh", instance.getPostStartupScript());
    assertEquals(true, instance.getInstallGpuDriver());
    assertEquals("custom-path", instance.getCustomGpuDriverPath());
    assertEquals("boot-disk-type", instance.getBootDiskType());
    assertEquals(111L, instance.getBootDiskSizeGb());
    assertEquals("data-disk-type", instance.getDataDiskType());
    assertEquals(222L, instance.getDataDiskSizeGb());
    assertThat(instance.getMetadata(), Matchers.aMapWithSize(2));
    // git secrets gets a false positive if 'service_account' is double quoted.
    assertThat(instance.getMetadata(), Matchers.hasEntry("proxy-mode", "service_" + "account"));
    assertThat(instance.getMetadata(), Matchers.hasEntry("metadata-key", "metadata-value"));
    assertEquals("foo@bar.com", instance.getServiceAccount());
    assertEquals(sa_scopes, instance.getServiceAccountScopes());
    assertEquals(4L, instance.getAcceleratorConfig().getCoreCount());
    assertEquals("accelerator-type", instance.getAcceleratorConfig().getType());
    assertEquals("project-id", instance.getVmImage().getProject());
    assertEquals("image-family", instance.getVmImage().getImageFamily());
    assertEquals("image-name", instance.getVmImage().getImageName());
    assertEquals("repository", instance.getContainerImage().getRepository());
    assertEquals("tag", instance.getContainerImage().getTag());
}

Example 2 with AcceleratorConfig

use of com.google.api.services.notebooks.v1.model.AcceleratorConfig in project terra-workspace-manager by DataBiosphere.

the class CreateAiNotebookInstanceStep method setFields.

@VisibleForTesting
static Instance setFields(ApiGcpAiNotebookInstanceCreationParameters creationParameters, String serviceAccountEmail, Instance instance) {
    instance.setPostStartupScript(Optional.ofNullable(creationParameters.getPostStartupScript()).orElse(DEFAULT_POST_STARTUP_SCRIPT)).setMachineType(creationParameters.getMachineType()).setInstallGpuDriver(creationParameters.isInstallGpuDriver()).setCustomGpuDriverPath(creationParameters.getCustomGpuDriverPath()).setBootDiskType(creationParameters.getBootDiskType()).setBootDiskSizeGb(creationParameters.getBootDiskSizeGb()).setDataDiskType(creationParameters.getDataDiskType()).setDataDiskSizeGb(creationParameters.getDataDiskSizeGb());
    Map<String, String> metadata = new HashMap<>();
    Optional.ofNullable(creationParameters.getMetadata()).ifPresent(metadata::putAll);
    // https://cloud.google.com/ai-platform/notebooks/docs/troubleshooting#opening_a_notebook_results_in_a_403_forbidden_error
    if (metadata.put(PROXY_MODE_METADATA_KEY, PROXY_MODE_SA_VALUE) != null) {
        throw new BadRequestException("proxy-mode metadata is reserved for Terra.");
    }
    instance.setMetadata(metadata);
    instance.setServiceAccount(serviceAccountEmail);
    instance.setServiceAccountScopes(SERVICE_ACCOUNT_SCOPES);
    ApiGcpAiNotebookInstanceAcceleratorConfig acceleratorConfig = creationParameters.getAcceleratorConfig();
    if (acceleratorConfig != null) {
        instance.setAcceleratorConfig(new AcceleratorConfig().setType(acceleratorConfig.getType()).setCoreCount(acceleratorConfig.getCoreCount()));
    }
    ApiGcpAiNotebookInstanceVmImage vmImageParameters = creationParameters.getVmImage();
    if (vmImageParameters != null) {
        instance.setVmImage(new VmImage().setProject(vmImageParameters.getProjectId()).setImageFamily(vmImageParameters.getImageFamily()).setImageName(vmImageParameters.getImageName()));
    }
    ApiGcpAiNotebookInstanceContainerImage containerImageParameters = creationParameters.getContainerImage();
    if (containerImageParameters != null) {
        instance.setContainerImage(new ContainerImage().setRepository(containerImageParameters.getRepository()).setTag(containerImageParameters.getTag()));
    }
    return instance;
}