Criar e executar um job que usa GPUs

Este documento explica como criar e executar um job que usa uma de processamento gráfico (GPU). Para saber mais sobre os recursos restrições para GPUs, consulte Sobre GPUs na na documentação do Compute Engine.

Ao criar um job em lote, você pode usar GPUs para acelerar cargas de trabalho específicas. Casos de uso comuns para jobs que usam GPUs incluem processamento de dados intensivo e cargas de trabalho de inteligência artificial (IA), como machine learning (ML).

Antes de começar

  1. Se você nunca usou o Batch antes, revise Introdução ao Batch e ativar o Batch. Para isso, conclua pré-requisitos para projetos e usuários.

  2. Para receber as permissões necessárias para criar um job, peça ao administrador para conceder a você os seguintes papéis do IAM:

    Para mais informações sobre a concessão de papéis, consulte Gerenciar o acesso a projetos, pastas e organizações.

    Também é possível conseguir as permissões necessárias por meio de papéis personalizados ou de outros papéis predefinidos.

Criar um job que usa GPUs

Para criar um job que use GPUs, faça o seguinte:

  1. Planeje os requisitos para um job que usa GPUs.
  2. Crie um trabalho com os requisitos e métodos que você identificou. Para conferir exemplos de como criar um job usando os métodos recomendados, consulte Criar um job de exemplo que usa GPUs neste documento.

Planejar os requisitos para um job que usa GPUs

Antes de criar um job que use GPUs, planeje os requisitos dele, conforme explicado nas seguintes seções:

  1. Selecionar o tipo de máquina de GPU
  2. Instalar os drivers da GPU
  3. Definir recursos de VM compatíveis

Etapa 1: selecionar o tipo de máquina da GPU

Os tipos de máquina de GPU disponíveis (as combinações válidas de tipo de GPU, número de GPUs e tipo de máquina (vCPUs e memória)) e os casos de uso estão listados na página Tipos de máquina de GPU na documentação do Compute Engine.

Os campos necessários para um job especificar um tipo de máquina GPU variam de acordo com a categorias na tabela a seguir:

Tipos de máquina da GPU e requisitos de job

GPUs para VMs otimizadas para acelerador: VMs com um tipo de máquina da família de máquinas com otimização para acelerador têm um tipo e um número específicos dessas GPUs automaticamente anexadas.

Para usar GPUs em VMs otimizadas para aceleradores, recomendamos que você especificar o tipo de máquina. Cada tipo de máquina com otimização para acelerador oferece suporte apenas a um tipo e número específicos de GPUs, ou seja, equivalente, quer você especifique ou não esses valores, além dos tipo de máquina otimização para acelerador.

Especificamente, o Batch também permite especificar e o número de GPUs para VMs com otimização de acelerador, mas o resultado As opções de vCPU e memória costumam ser muito limitadas. Por isso, recomendamos verificar se as opções de vCPU e memória disponíveis são compatíveis com os requisitos da tarefa.

GPUs para VMs N1: essas GPUs exigem que você especifique o tipo. e o valor a ser anexado a cada VM e precisam ser anexados às VMs com uma na linha de comando Série de máquinas N1.

Para usar GPUs em VMs N1, recomendamos especificar pelo menos o tipo e número de GPUs. Verifique se a combinação de valores corresponde a uma das opções de GPU válidas para os tipos de máquina N1. As opções de vCPU e memória para VMs N1 que usam qualquer tipo específico e número de GPUs é bastante flexível, portanto, se preferir, você pode deixar o Batch selecionar uma tipo de máquina que atenda aos requisitos da tarefa.

Etapa 2: instalar os drivers da GPU

Para instalar os drivers de GPU necessários, selecione um dos seguintes métodos:

Etapa 3: definir recursos de VM compatíveis

Para saber mais sobre os requisitos e as opções para definir os recursos da VM para um vaga, consulte Recursos do job.

Em resumo, você precisa realizar todas as ações a seguir ao definir os recursos da VM para um job que usa GPUs:

  • Verifique se o tipo de máquina da GPU está disponível no local das VMs do job.

    Para saber onde os tipos de máquinas GPU estão disponíveis, consulte Disponibilidade de GPU por regiões e zonas na documentação do Compute Engine.

  • Se você especificar o tipo de máquina do job, verifique se ele tem vCPUs e memória para os requisitos de tarefa do job. A especificação do tipo de máquina do job é obrigatória ao usar GPUs para VMs otimizadas para acelerador e opcional ao usar GPUs para VMs N1.

  • Defina os recursos da VM para um job usando um método válido:

    • Defina os recursos da VM diretamente usando o campo instances[].policy (recomendado se possível). Esse método é mostrado nos exemplos.
    • Defina os recursos da VM com um modelo usando o campo instances[].instanceTemplate. Esse método é necessário para instalar manualmente drivers de GPU por meio de um imagem. Para mais informações, consulte Defina recursos de job usando um modelo de instância de VM.

Criar um job de exemplo que usa GPUs

As seções a seguir explicam como criar um exemplo de vaga para diferentes tipos de máquina GPU. Todos os jobs de exemplo instalam drivers de GPU automaticamente e definir diretamente os recursos da VM.

Usar GPUs para VMs otimizadas para aceleradores

É possível criar um job que usa GPUs para VMs com otimização para acelerador usando o CLI gcloud, API Batch, Java ou Python.

gcloud

  1. Crie um arquivo JSON que instale drivers de GPU e defina Campo machineType com tipo de máquina da família de máquinas com otimização para aceleradores e usa um local que tenha o tipo especificado de GPUs.

    Por exemplo, para criar um job de script básico que usa GPUs para com as VMs otimizadas para aceleradores, crie um arquivo JSON com o seguinte conteúdo:

    {
    "taskGroups": [
        {
            "taskSpec": {
                "runnables": [
                    {
                        "script": {
                            "text": "echo Hello world from task ${BATCH_TASK_INDEX}."
                        }
                    }
                ]
            },
            "taskCount": 3,
            "parallelism": 1
        }
    ],
    "allocationPolicy": {
        "instances": [
            {
                "installGpuDrivers": INSTALL_GPU_DRIVERS,
                "policy": {
                    "machineType": "MACHINE_TYPE"
                }
            }
        ],
        "location": {
            "allowedLocations": [
                "ALLOWED_LOCATIONS"
            ]
        }
    }
}

    Substitua:

    • INSTALL_GPU_DRIVERS: quando definido como true, o Batch busca os drivers necessários para Tipo de GPU especificado no campo policy de um local de terceiros, e o Batch os instala em seu nome. Se você definir esse campo como false (padrão), será necessário instalar os drivers de GPU manualmente para usar GPUs para esse job.

    • MACHINE_TYPE: um tipo de máquina do família de máquinas com otimização para acelerador.

    • ALLOWED_LOCATIONS: o campo allowedLocations[] define uma região e, opcionalmente, uma ou mais zonas em que as instâncias de VM do seu job podem ser executadas. Por exemplo, regions/us-central1, zones/us-central1-a permite a zona us-central1-a. Especifique os locais que oferecem o tipo de máquina de GPU que você quer para esse job. Caso contrário, se você omitir esse campo, verifique se o local do job oferece o tipo de máquina de GPU.

  2. Para criar e executar o job, use o comando gcloud batch jobs submit:

    gcloud batch jobs submit JOB_NAME \
    --location LOCATION \
    --config JSON_CONFIGURATION_FILE

    Substitua:

    • JOB_NAME: o nome do job.

    • LOCATION: o local do job.

    • JSON_CONFIGURATION_FILE: o caminho de um arquivo JSON com os detalhes de configuração do job.

API

Faça uma solicitação POST ao Método jobs.create que instala drivers de GPU, define Campo machineType com tipo de máquina da família de máquinas com otimização para aceleradores e usa um local que tenha o tipo especificado de GPUs.

Por exemplo, para criar um job de script básico que usa GPUs para VMs otimizadas para acelerador, faça a seguinte solicitação:

POST https://batch.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/jobs?job_id=JOB_NAME

{
    "taskGroups": [
        {
            "taskSpec": {
                "runnables": [
                    {
                        "script": {
                            "text": "echo Hello world from task ${BATCH_TASK_INDEX}."
                        }
                    }
                ]
            },
            "taskCount": 3,
            "parallelism": 1
        }
    ],
    "allocationPolicy": {
        "instances": [
            {
                "installGpuDrivers": INSTALL_GPU_DRIVERS,
                "policy": {
                    "machineType": "MACHINE_TYPE"
                }
            }
        ],
        "location": {
            "allowedLocations": [
                "ALLOWED_LOCATIONS"
            ]
        }
    }
}

Substitua:

  • PROJECT_ID: o ID do projeto do seu projeto.

  • LOCATION: o local do trabalho.

  • JOB_NAME: o nome do job.

  • INSTALL_GPU_DRIVERS: quando definido como true, o lote busca os drivers necessários para o tipo de GPU especificado no campo policy de um local de terceiros e os instala em seu nome. Se você definir esse campo como false (padrão), será necessário para instalar drivers de GPU manualmente a fim de usar GPUs nesse job.

  • MACHINE_TYPE: um tipo de máquina do família de máquinas com otimização para acelerador.

  • ALLOWED_LOCATIONS: o campo allowedLocations[] define uma região e, opcionalmente, uma ou mais zonas em que as instâncias de VM do seu job podem ser executadas. Por exemplo, regions/us-central1, zones/us-central1-a permite a zona us-central1-a. Especifique os locais que oferecem o tipo de máquina de GPU que você quer para esse job. Caso contrário, se você omitir esse campo, o local do job mostra o tipo de máquina GPU.

Java 


import com.google.cloud.batch.v1.AllocationPolicy;
import com.google.cloud.batch.v1.AllocationPolicy.Accelerator;
import com.google.cloud.batch.v1.AllocationPolicy.InstancePolicy;
import com.google.cloud.batch.v1.AllocationPolicy.InstancePolicyOrTemplate;
import com.google.cloud.batch.v1.BatchServiceClient;
import com.google.cloud.batch.v1.CreateJobRequest;
import com.google.cloud.batch.v1.Job;
import com.google.cloud.batch.v1.LogsPolicy;
import com.google.cloud.batch.v1.Runnable;
import com.google.cloud.batch.v1.Runnable.Script;
import com.google.cloud.batch.v1.TaskGroup;
import com.google.cloud.batch.v1.TaskSpec;
import com.google.protobuf.Duration;
import java.io.IOException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

public class CreateGpuJob {

  public static void main(String[] args)
      throws IOException, ExecutionException, InterruptedException, TimeoutException {
    // TODO(developer): Replace these variables before running the sample.
    // Project ID or project number of the Google Cloud project you want to use.
    String projectId = "YOUR_PROJECT_ID";
    // Name of the region you want to use to run the job. Regions that are
    // available for Batch are listed on: https://cloud.google.com/batch/docs/get-started#locations
    String region = "europe-central2";
    // The name of the job that will be created.
    // It needs to be unique for each project and region pair.
    String jobName = "JOB_NAME";
    // Optional. When set to true, Batch fetches the drivers required for the GPU type
    // that you specify in the policy field from a third-party location,
    // and Batch installs them on your behalf. If you set this field to false (default),
    // you need to install GPU drivers manually to use any GPUs for this job.
    boolean installGpuDrivers = false;
    // Accelerator-optimized machine types are available to Batch jobs. See the list
    // of available types on: https://cloud.google.com/compute/docs/accelerator-optimized-machines
    String machineType = "g2-standard-4";

    createGpuJob(projectId, region, jobName, installGpuDrivers, machineType);
  }

  // Create a job that uses GPUs
  public static Job createGpuJob(String projectId, String region, String jobName,
                                  boolean installGpuDrivers, String machineType)
      throws IOException, ExecutionException, InterruptedException, TimeoutException {
    // Initialize client that will be used to send requests. This client only needs to be created
    // once, and can be reused for multiple requests.
    try (BatchServiceClient batchServiceClient = BatchServiceClient.create()) {
      // Define what will be done as part of the job.
      Runnable runnable =
          Runnable.newBuilder()
              .setScript(
                  Script.newBuilder()
                      .setText(
                          "echo Hello world! This is task ${BATCH_TASK_INDEX}. "
                                  + "This job has a total of ${BATCH_TASK_COUNT} tasks.")
                      // You can also run a script from a file. Just remember, that needs to be a
                      // script that's already on the VM that will be running the job.
                      // Using setText() and setPath() is mutually exclusive.
                      // .setPath("/tmp/test.sh")
                      .build())
              .build();

      TaskSpec task = TaskSpec.newBuilder()
                  // Jobs can be divided into tasks. In this case, we have only one task.
                  .addRunnables(runnable)
                  .setMaxRetryCount(2)
                  .setMaxRunDuration(Duration.newBuilder().setSeconds(3600).build())
                  .build();

      // Tasks are grouped inside a job using TaskGroups.
      // Currently, it's possible to have only one task group.
      TaskGroup taskGroup = TaskGroup.newBuilder()
          .setTaskCount(3)
          .setParallelism(1)
          .setTaskSpec(task)
          .build();

      // Policies are used to define on what kind of virtual machines the tasks will run.
      // Read more about machine types here: https://cloud.google.com/compute/docs/machine-types
      InstancePolicy instancePolicy =
          InstancePolicy.newBuilder().setMachineType(machineType).build();  

      // Policies are used to define on what kind of virtual machines the tasks will run on.
      AllocationPolicy allocationPolicy =
          AllocationPolicy.newBuilder()
              .addInstances(
                  InstancePolicyOrTemplate.newBuilder()
                      .setInstallGpuDrivers(installGpuDrivers)
                      .setPolicy(instancePolicy)
                      .build())
              .build();

      Job job =
          Job.newBuilder()
              .addTaskGroups(taskGroup)
              .setAllocationPolicy(allocationPolicy)
              .putLabels("env", "testing")
              .putLabels("type", "script")
              // We use Cloud Logging as it's an out of the box available option.
              .setLogsPolicy(
                  LogsPolicy.newBuilder().setDestination(LogsPolicy.Destination.CLOUD_LOGGING))
              .build();

      CreateJobRequest createJobRequest =
          CreateJobRequest.newBuilder()
              // The job's parent is the region in which the job will run.
              .setParent(String.format("projects/%s/locations/%s", projectId, region))
              .setJob(job)
              .setJobId(jobName)
              .build();

      Job result =
          batchServiceClient
              .createJobCallable()
              .futureCall(createJobRequest)
              .get(5, TimeUnit.MINUTES);

      System.out.printf("Successfully created the job: %s", result.getName());

      return result;
    }
  }
}

Python 

from google.cloud import batch_v1


def create_gpu_job(project_id: str, region: str, job_name: str) -> batch_v1.Job:
    """
    This method shows how to create a sample Batch Job that will run
    a simple command on Cloud Compute instances on GPU machines.

    Args:
        project_id: project ID or project number of the Cloud project you want to use.
        region: name of the region you want to use to run the job. Regions that are
            available for Batch are listed on: https://cloud.google.com/batch/docs/get-started#locations
        job_name: the name of the job that will be created.
            It needs to be unique for each project and region pair.

    Returns:
        A job object representing the job created.
    """
    client = batch_v1.BatchServiceClient()

    # Define what will be done as part of the job.
    task = batch_v1.TaskSpec()
    runnable = batch_v1.Runnable()
    runnable.script = batch_v1.Runnable.Script()
    runnable.script.text = "echo Hello world! This is task ${BATCH_TASK_INDEX}. This job has a total of ${BATCH_TASK_COUNT} tasks."
    # You can also run a script from a file. Just remember, that needs to be a script that's
    # already on the VM that will be running the job. Using runnable.script.text and runnable.script.path is mutually
    # exclusive.
    # runnable.script.path = '/tmp/test.sh'
    task.runnables = [runnable]

    # We can specify what resources are requested by each task.
    resources = batch_v1.ComputeResource()
    resources.cpu_milli = 2000  # in milliseconds per cpu-second. This means the task requires 2 whole CPUs.
    resources.memory_mib = 16  # in MiB
    task.compute_resource = resources

    task.max_retry_count = 2
    task.max_run_duration = "3600s"

    # Tasks are grouped inside a job using TaskGroups.
    # Currently, it's possible to have only one task group.
    group = batch_v1.TaskGroup()
    group.task_count = 4
    group.task_spec = task

    # Policies are used to define on what kind of virtual machines the tasks will run on.
    # In this case, we tell the system to use "g2-standard-4" machine type.
    # Read more about machine types here: https://cloud.google.com/compute/docs/machine-types
    policy = batch_v1.AllocationPolicy.InstancePolicy()
    policy.machine_type = "g2-standard-4"

    instances = batch_v1.AllocationPolicy.InstancePolicyOrTemplate()
    instances.policy = policy
    instances.install_gpu_drivers = True
    allocation_policy = batch_v1.AllocationPolicy()
    allocation_policy.instances = [instances]

    job = batch_v1.Job()
    job.task_groups = [group]
    job.allocation_policy = allocation_policy
    job.labels = {"env": "testing", "type": "container"}
    # We use Cloud Logging as it's an out of the box available option
    job.logs_policy = batch_v1.LogsPolicy()
    job.logs_policy.destination = batch_v1.LogsPolicy.Destination.CLOUD_LOGGING

    create_request = batch_v1.CreateJobRequest()
    create_request.job = job
    create_request.job_id = job_name
    # The job's parent is the region in which the job will run
    create_request.parent = f"projects/{project_id}/locations/{region}"

    return client.create_job(create_request)

Usar GPUs para VMs N1

É possível criar um job que usa GPUs para VMs N1 usando a CLI gcloud, a API Batch, Java, Node.js ou Python.

gcloud

  1. Crie um arquivo JSON que instale drivers de GPU, defina os subcampos type e count do campo accelerators[] e use um local com o tipo especificado de GPUs.

    Por exemplo, para criar um job de script básico que usa GPUs para VMs N1 e permite que o Batch selecione exatamente o tipo de máquina N1, um arquivo JSON com o seguinte conteúdo:

    {
    "taskGroups": [
        {
            "taskSpec": {
                "runnables": [
                    {
                        "script": {
                            "text": "echo Hello world from task ${BATCH_TASK_INDEX}."
                        }
                    }
                ]
            },
            "taskCount": 3,
            "parallelism": 1
        }
    ],
    "allocationPolicy": {
        "instances": [
            {
                "installGpuDrivers": INSTALL_GPU_DRIVERS,
                "policy": {
                    "accelerators": [
                        {
                            "type": "GPU_TYPE",
                            "count": GPU_COUNT
                        }
                    ]
                }
            }
        ],
        "location": {
            "allowedLocations": [
                "ALLOWED_LOCATIONS"
            ]
        }
    }
}

    Substitua:

    • INSTALL_GPU_DRIVERS: quando definido como true, o Batch busca os drivers necessários para Tipo de GPU especificado no campo policy de um local de terceiros, e o Batch os instala em seu nome. Se você definir esse campo como false (padrão), será necessário para instalar drivers de GPU manualmente a fim de usar GPUs nesse job.

    • GPU_TYPE: a GPU. não é válido. Você pode ver uma lista dos tipos de GPU disponíveis usando o comando gcloud compute accelerator-types list. Use este campo apenas para GPUs para VMs N1.

    • GPU_COUNT: o número de GPUs do tipo especificado. Para mais informações sobre as opções válidas, consulte os tipos de máquinas de GPU para a série de máquinas N1. Use este campo apenas para GPUs de VMs N1.

    • ALLOWED_LOCATIONS: o campo allowedLocations[] define uma região e, opcionalmente, uma ou mais zonas em que as instâncias de VM do seu job podem ser executadas. Por exemplo, regions/us-central1, zones/us-central1-a permite a zona us-central1-a. Certifique-se de especificar locais que oferecem o tipo de máquina GPU que você quer para este trabalho. Caso contrário, se você omitir esse campo, o local do job mostra o tipo de máquina GPU.

  2. Para criar e executar o job, use o Comando gcloud batch jobs submit:

    gcloud batch jobs submit JOB_NAME \
    --location LOCATION \
    --config JSON_CONFIGURATION_FILE

    Substitua:

    • JOB_NAME: o nome do job.

    • LOCATION: o local do trabalho.

    • JSON_CONFIGURATION_FILE: o caminho de uma Arquivo JSON com os detalhes de configuração do job.

API

Faça uma solicitação POST ao Método jobs.create que instala drivers de GPU, define Subcampos type e count do campo accelerators[] e usa um local que tem o e tipos específicos de GPUs.

Por exemplo, para criar um job de script básico que usa GPUs para VMs N1 e permite que o Batch selecione exatamente o tipo de máquina N1, seguinte solicitação:

POST https://batch.googleapis.com/v1/projects/PROJECT_ID/locations/LOCATION/jobs?job_id=JOB_NAME

{
    "taskGroups": [
        {
            "taskSpec": {
                "runnables": [
                    {
                        "script": {
                            "text": "echo Hello world from task ${BATCH_TASK_INDEX}."
                        }
                    }
                ]
            },
            "taskCount": 3,
            "parallelism": 1
        }
    ],
    "allocationPolicy": {
        "instances": [
            {
                "installGpuDrivers": INSTALL_GPU_DRIVERS,
                "policy": {
                    "accelerators": [
                        {
                            "type": "GPU_TYPE",
                            "count": GPU_COUNT
                        }
                    ]
                }
            }
        ],
        "location": {
            "allowedLocations": [
                "ALLOWED_LOCATIONS"
            ]
        }
    }
}

Substitua:

  • PROJECT_ID: o ID do projeto do seu projeto.

  • LOCATION: o local do trabalho.

  • JOB_NAME: o nome do job.

  • INSTALL_GPU_DRIVERS: quando definido como true, o lote busca os drivers necessários para o tipo de GPU especificado no campo policy de um local de terceiros e os instala em seu nome. Se você definir esse campo como false (padrão), será necessário para instalar drivers de GPU manualmente a fim de usar GPUs nesse job.

  • GPU_TYPE: o tipo de GPU. Você pode ver uma lista dos tipos de GPU disponíveis usando o comando gcloud compute accelerator-types list. Use este campo apenas para GPUs para VMs N1.

  • GPU_COUNT: o número de GPUs do tipo especificado. Para mais informações sobre as opções válidas, consulte Tipos de máquinas de GPU para séries de máquinas N1. Use este campo apenas para GPUs de VMs N1.

  • ALLOWED_LOCATIONS: o campo allowedLocations[] define uma região e, opcionalmente, uma ou mais zonas em que as instâncias de VM do seu job podem ser executadas. Por exemplo, regions/us-central1, zones/us-central1-a permite a zona us-central1-a. Especifique os locais que oferecem o tipo de máquina de GPU que você quer para esse job. Caso contrário, se você omitir esse campo, o local do job mostra o tipo de máquina GPU.

Java 


import com.google.cloud.batch.v1.AllocationPolicy;
import com.google.cloud.batch.v1.AllocationPolicy.Accelerator;
import com.google.cloud.batch.v1.AllocationPolicy.InstancePolicy;
import com.google.cloud.batch.v1.AllocationPolicy.InstancePolicyOrTemplate;
import com.google.cloud.batch.v1.BatchServiceClient;
import com.google.cloud.batch.v1.CreateJobRequest;
import com.google.cloud.batch.v1.Job;
import com.google.cloud.batch.v1.LogsPolicy;
import com.google.cloud.batch.v1.Runnable;
import com.google.cloud.batch.v1.Runnable.Script;
import com.google.cloud.batch.v1.TaskGroup;
import com.google.cloud.batch.v1.TaskSpec;
import com.google.protobuf.Duration;
import java.io.IOException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

public class CreateGpuJobN1 {

  public static void main(String[] args)
      throws IOException, ExecutionException, InterruptedException, TimeoutException {
    // TODO(developer): Replace these variables before running the sample.
    // Project ID or project number of the Google Cloud project you want to use.
    String projectId = "YOUR_PROJECT_ID";
    // Name of the region you want to use to run the job. Regions that are
    // available for Batch are listed on: https://cloud.google.com/batch/docs/get-started#locations
    String region = "europe-central2";
    // The name of the job that will be created.
    // It needs to be unique for each project and region pair.
    String jobName = "JOB_NAME";
    // Optional. When set to true, Batch fetches the drivers required for the GPU type
    // that you specify in the policy field from a third-party location,
    // and Batch installs them on your behalf. If you set this field to false (default),
    // you need to install GPU drivers manually to use any GPUs for this job.
    boolean installGpuDrivers = false;
    // The GPU type. You can view a list of the available GPU types
    // by using the `gcloud compute accelerator-types list` command.
    String gpuType = "nvidia-tesla-t4";
    // The number of GPUs of the specified type.
    int gpuCount = 2;

    createGpuJob(projectId, region, jobName, installGpuDrivers, gpuType, gpuCount);
  }

  // Create a job that uses GPUs
  public static Job createGpuJob(String projectId, String region, String jobName,
                                  boolean installGpuDrivers, String gpuType, int gpuCount)
      throws IOException, ExecutionException, InterruptedException, TimeoutException {
    // Initialize client that will be used to send requests. This client only needs to be created
    // once, and can be reused for multiple requests.
    try (BatchServiceClient batchServiceClient = BatchServiceClient.create()) {
      // Define what will be done as part of the job.
      Runnable runnable =
          Runnable.newBuilder()
              .setScript(
                  Script.newBuilder()
                      .setText(
                          "echo Hello world! This is task ${BATCH_TASK_INDEX}. "
                                  + "This job has a total of ${BATCH_TASK_COUNT} tasks.")
                      // You can also run a script from a file. Just remember, that needs to be a
                      // script that's already on the VM that will be running the job.
                      // Using setText() and setPath() is mutually exclusive.
                      // .setPath("/tmp/test.sh")
                      .build())
              .build();

      TaskSpec task = TaskSpec.newBuilder()
                  // Jobs can be divided into tasks. In this case, we have only one task.
                  .addRunnables(runnable)
                  .setMaxRetryCount(2)
                  .setMaxRunDuration(Duration.newBuilder().setSeconds(3600).build())
                  .build();

      // Tasks are grouped inside a job using TaskGroups.
      // Currently, it's possible to have only one task group.
      TaskGroup taskGroup = TaskGroup.newBuilder()
          .setTaskCount(3)
          .setParallelism(1)
          .setTaskSpec(task)
          .build();

      // Accelerator describes Compute Engine accelerators to be attached to the VM.
      Accelerator accelerator = Accelerator.newBuilder()
          .setType(gpuType)
          .setCount(gpuCount)
          .build();

      // Policies are used to define on what kind of virtual machines the tasks will run on.
      AllocationPolicy allocationPolicy =
          AllocationPolicy.newBuilder()
              .addInstances(
                  InstancePolicyOrTemplate.newBuilder()
                      .setInstallGpuDrivers(installGpuDrivers)
                      .setPolicy(InstancePolicy.newBuilder().addAccelerators(accelerator))
                      .build())
              .build();

      Job job =
          Job.newBuilder()
              .addTaskGroups(taskGroup)
              .setAllocationPolicy(allocationPolicy)
              .putLabels("env", "testing")
              .putLabels("type", "script")
              // We use Cloud Logging as it's an out of the box available option.
              .setLogsPolicy(
                  LogsPolicy.newBuilder().setDestination(LogsPolicy.Destination.CLOUD_LOGGING))
              .build();

      CreateJobRequest createJobRequest =
          CreateJobRequest.newBuilder()
              // The job's parent is the region in which the job will run.
              .setParent(String.format("projects/%s/locations/%s", projectId, region))
              .setJob(job)
              .setJobId(jobName)
              .build();

      Job result =
          batchServiceClient
              .createJobCallable()
              .futureCall(createJobRequest)
              .get(5, TimeUnit.MINUTES);

      System.out.printf("Successfully created the job: %s", result.getName());

      return result;
    }
  }
}

Node.js

Para criar um job com GPUs usando Node.js, selecione uma das opções a seguir com base no tipo de máquina do seu modelo de GPU:

Criar um job que usa GPUs com VMs otimizadas para acelerador

Para usar GPUs com VMs otimizadas para acelerador, basta especificar o tipo de máquina que você quer para as VMs do job:

// Imports the Batch library
const batchLib = require('@google-cloud/batch');
const batch = batchLib.protos.google.cloud.batch.v1;

// Instantiates a client
const batchClient = new batchLib.v1.BatchServiceClient();

/**
 * TODO(developer): Update these variables before running the sample.
 */
// Project ID or project number of the Google Cloud project you want to use.
const projectId = await batchClient.getProjectId();
// Name of the region you want to use to run the job. Regions that are
// available for Batch are listed on: https://cloud.google.com/batch/docs/get-started#locations
const region = 'europe-central2';
// The name of the job that will be created.
// It needs to be unique for each project and region pair.
const jobName = 'batch-gpu-job';
// The GPU type. You can view a list of the available GPU types
// by using the `gcloud compute accelerator-types list` command.
const gpuType = 'nvidia-l4';
// The number of GPUs of the specified type.
const gpuCount = 1;
// Optional. When set to true, Batch fetches the drivers required for the GPU type
// that you specify in the policy field from a third-party location,
// and Batch installs them on your behalf. If you set this field to false (default),
// you need to install GPU drivers manually to use any GPUs for this job.
const installGpuDrivers = false;
// Accelerator-optimized machine types are available to Batch jobs. See the list
// of available types on: https://cloud.google.com/compute/docs/accelerator-optimized-machines
const machineType = 'g2-standard-4';

// Define what will be done as part of the job.
const runnable = new batch.Runnable({
  script: new batch.Runnable.Script({
    commands: ['-c', 'echo Hello world! This is task ${BATCH_TASK_INDEX}.'],
  }),
});

const task = new batch.TaskSpec({
  runnables: [runnable],
  maxRetryCount: 2,
  maxRunDuration: {seconds: 3600},
});

// Tasks are grouped inside a job using TaskGroups.
const group = new batch.TaskGroup({
  taskCount: 3,
  taskSpec: task,
});

// Policies are used to define on what kind of virtual machines the tasks will run on.
// In this case, we tell the system to use "g2-standard-4" machine type.
// Read more about machine types here: https://cloud.google.com/compute/docs/machine-types
const instancePolicy = new batch.AllocationPolicy.InstancePolicy({
  machineType,
  // Accelerator describes Compute Engine accelerators to be attached to the VM
  accelerators: [
    new batch.AllocationPolicy.Accelerator({
      type: gpuType,
      count: gpuCount,
      installGpuDrivers,
    }),
  ],
});

const allocationPolicy = new batch.AllocationPolicy.InstancePolicyOrTemplate({
  instances: [{installGpuDrivers, policy: instancePolicy}],
});

const job = new batch.Job({
  name: jobName,
  taskGroups: [group],
  labels: {env: 'testing', type: 'script'},
  allocationPolicy,
  // We use Cloud Logging as it's an option available out of the box
  logsPolicy: new batch.LogsPolicy({
    destination: batch.LogsPolicy.Destination.CLOUD_LOGGING,
  }),
});
// The job's parent is the project and region in which the job will run
const parent = `projects/${projectId}/locations/${region}`;

async function callCreateBatchGPUJob() {
  // Construct request
  const request = {
    parent,
    jobId: jobName,
    job,
  };

  // Run request
  const [response] = await batchClient.createJob(request);
  console.log(JSON.stringify(response));
}

await callCreateBatchGPUJob();

Criar um job que usa GPUs com VMs N1

Para usar GPUs com VMs N1, é necessário especificar o número e o tipo de GPUs que você quer para cada uma das VMs do job:

// Imports the Batch library
const batchLib = require('@google-cloud/batch');
const batch = batchLib.protos.google.cloud.batch.v1;

// Instantiates a client
const batchClient = new batchLib.v1.BatchServiceClient();

/**
 * TODO(developer): Update these variables before running the sample.
 */
// Project ID or project number of the Google Cloud project you want to use.
const projectId = await batchClient.getProjectId();
// Name of the region you want to use to run the job. Regions that are
// available for Batch are listed on: https://cloud.google.com/batch/docs/get-started#locations
const region = 'europe-central2';
// The name of the job that will be created.
// It needs to be unique for each project and region pair.
const jobName = 'batch-gpu-job-n1';
// The GPU type. You can view a list of the available GPU types
// by using the `gcloud compute accelerator-types list` command.
const gpuType = 'nvidia-tesla-t4';
// The number of GPUs of the specified type.
const gpuCount = 1;
// Optional. When set to true, Batch fetches the drivers required for the GPU type
// that you specify in the policy field from a third-party location,
// and Batch installs them on your behalf. If you set this field to false (default),
// you need to install GPU drivers manually to use any GPUs for this job.
const installGpuDrivers = false;
// Accelerator-optimized machine types are available to Batch jobs. See the list
// of available types on: https://cloud.google.com/compute/docs/accelerator-optimized-machines
const machineType = 'n1-standard-16';

// Define what will be done as part of the job.
const runnable = new batch.Runnable({
  script: new batch.Runnable.Script({
    commands: ['-c', 'echo Hello world! This is task ${BATCH_TASK_INDEX}.'],
  }),
});

const task = new batch.TaskSpec({
  runnables: [runnable],
  maxRetryCount: 2,
  maxRunDuration: {seconds: 3600},
});

// Tasks are grouped inside a job using TaskGroups.
const group = new batch.TaskGroup({
  taskCount: 3,
  taskSpec: task,
});

// Policies are used to define on what kind of virtual machines the tasks will run on.
// In this case, we tell the system to use "g2-standard-4" machine type.
// Read more about machine types here: https://cloud.google.com/compute/docs/machine-types
const instancePolicy = new batch.AllocationPolicy.InstancePolicy({
  machineType,
  // Accelerator describes Compute Engine accelerators to be attached to the VM
  accelerators: [
    new batch.AllocationPolicy.Accelerator({
      type: gpuType,
      count: gpuCount,
      installGpuDrivers,
    }),
  ],
});

const allocationPolicy = new batch.AllocationPolicy.InstancePolicyOrTemplate({
  instances: [{installGpuDrivers, policy: instancePolicy}],
});

const job = new batch.Job({
  name: jobName,
  taskGroups: [group],
  labels: {env: 'testing', type: 'script'},
  allocationPolicy,
  // We use Cloud Logging as it's an option available out of the box
  logsPolicy: new batch.LogsPolicy({
    destination: batch.LogsPolicy.Destination.CLOUD_LOGGING,
  }),
});
// The job's parent is the project and region in which the job will run
const parent = `projects/${projectId}/locations/${region}`;

async function callCreateBatchGPUJobN1() {
  // Construct request
  const request = {
    parent,
    jobId: jobName,
    job,
  };

  // Run request
  const [response] = await batchClient.createJob(request);
  console.log(JSON.stringify(response));
}

await callCreateBatchGPUJobN1();

Python 

from google.cloud import batch_v1


def create_gpu_job(
    project_id: str, region: str, zone: str, job_name: str
) -> batch_v1.Job:
    """
    This method shows how to create a sample Batch Job that will run
    a simple command on Cloud Compute instances on GPU machines.

    Args:
        project_id: project ID or project number of the Cloud project you want to use.
        region: name of the region you want to use to run the job. Regions that are
            available for Batch are listed on: https://cloud.google.com/batch/docs/get-started#locations
        zone: name of the zone you want to use to run the job. Important in regard to GPUs availability.
            GPUs availability can be found here: https://cloud.google.com/compute/docs/gpus/gpu-regions-zones
        job_name: the name of the job that will be created.
            It needs to be unique for each project and region pair.

    Returns:
        A job object representing the job created.
    """
    client = batch_v1.BatchServiceClient()

    # Define what will be done as part of the job.
    task = batch_v1.TaskSpec()
    runnable = batch_v1.Runnable()
    runnable.script = batch_v1.Runnable.Script()
    runnable.script.text = "echo Hello world! This is task ${BATCH_TASK_INDEX}. This job has a total of ${BATCH_TASK_COUNT} tasks."
    # You can also run a script from a file. Just remember, that needs to be a script that's
    # already on the VM that will be running the job. Using runnable.script.text and runnable.script.path is mutually
    # exclusive.
    # runnable.script.path = '/tmp/test.sh'
    task.runnables = [runnable]

    # We can specify what resources are requested by each task.
    resources = batch_v1.ComputeResource()
    resources.cpu_milli = 2000  # in milliseconds per cpu-second. This means the task requires 2 whole CPUs.
    resources.memory_mib = 16  # in MiB
    task.compute_resource = resources

    task.max_retry_count = 2
    task.max_run_duration = "3600s"

    # Tasks are grouped inside a job using TaskGroups.
    # Currently, it's possible to have only one task group.
    group = batch_v1.TaskGroup()
    group.task_count = 4
    group.task_spec = task

    # Policies are used to define on what kind of virtual machines the tasks will run on.
    # Read more about machine types here: https://cloud.google.com/compute/docs/machine-types
    policy = batch_v1.AllocationPolicy.InstancePolicy()
    policy.machine_type = "n1-standard-16"

    accelerator = batch_v1.AllocationPolicy.Accelerator()
    # Note: not every accelerator is compatible with instance type
    # Read more here: https://cloud.google.com/compute/docs/gpus#t4-gpus
    accelerator.type_ = "nvidia-tesla-t4"
    accelerator.count = 1

    policy.accelerators = [accelerator]
    instances = batch_v1.AllocationPolicy.InstancePolicyOrTemplate()
    instances.policy = policy
    instances.install_gpu_drivers = True
    allocation_policy = batch_v1.AllocationPolicy()
    allocation_policy.instances = [instances]

    location = batch_v1.AllocationPolicy.LocationPolicy()
    location.allowed_locations = ["zones/us-central1-b"]
    allocation_policy.location = location

    job = batch_v1.Job()
    job.task_groups = [group]
    job.allocation_policy = allocation_policy
    job.labels = {"env": "testing", "type": "container"}
    # We use Cloud Logging as it's an out of the box available option
    job.logs_policy = batch_v1.LogsPolicy()
    job.logs_policy.destination = batch_v1.LogsPolicy.Destination.CLOUD_LOGGING

    create_request = batch_v1.CreateJobRequest()
    create_request.job = job
    create_request.job_id = job_name
    # The job's parent is the region in which the job will run
    create_request.parent = f"projects/{project_id}/locations/{region}"

    return client.create_job(create_request)

A seguir